From 13433af4a50a2b77a9f50562af05b88c31428ee8 Mon Sep 17 00:00:00 2001 From: Zoltan Papp Date: Sun, 19 Jul 2026 21:37:15 +0200 Subject: [PATCH] [client] Preserve routed allowed IPs across lazy connection idle transitions When a peer went idle under lazy connections, the WireGuard peer was removed and the wake endpoint re-armed with only the overlay allowed IPs, dropping routed prefixes installed by the route manager. The asynchronous route-watcher re-add uses update_only, which is a silent no-op while the peer is absent, so routed subnets stayed black-holed until the peer was woken by other means. Split the idle transition from the full close: Conn.Idle tears down transports, cancels pending delayed endpoint updates and updates the status without touching the WireGuard peer. The activity listener then arms the wake endpoint via the new IdlePeerEndpoint operation, which removes and re-creates the peer in a single transaction: handshake state is dropped (keeping the handshake-first wake semantics of packet staging and first-packet capture/reinjection) while the currently installed allowed IPs are preserved. The read-modify-write runs under the interface mutex, serializing it against concurrent allowed IP updates from the route manager. Conn.Close no longer takes signalToRemote: the idle transition is the only path that signals GOAWAY to the remote peer. --- client/iface/configurer/common.go | 38 +++++ client/iface/configurer/kernel_unix.go | 34 ++++ client/iface/configurer/usp.go | 58 +++++++ client/iface/configurer/usp_idle_test.go | 158 ++++++++++++++++++ client/iface/device/interface.go | 1 + client/iface/iface.go | 13 ++ client/internal/conn_mgr.go | 2 +- client/internal/engine.go | 2 +- client/internal/engine_test.go | 8 + client/internal/iface_common.go | 1 + .../lazyconn/activity/listener_bind.go | 2 +- .../lazyconn/activity/listener_bind_test.go | 3 +- .../lazyconn/activity/listener_udp.go | 2 +- client/internal/lazyconn/activity/manager.go | 4 +- .../lazyconn/activity/manager_test.go | 3 +- client/internal/lazyconn/manager/manager.go | 4 +- client/internal/lazyconn/wgiface.go | 5 +- client/internal/peer/conn.go | 24 ++- client/internal/peer/endpoint.go | 10 ++ client/internal/peer/endpoint_test.go | 99 +++++++++++ client/internal/peerstore/store.go | 18 +- 21 files changed, 455 insertions(+), 34 deletions(-) create mode 100644 client/iface/configurer/usp_idle_test.go create mode 100644 client/internal/peer/endpoint_test.go diff --git a/client/iface/configurer/common.go b/client/iface/configurer/common.go index 10162d703..1cdf9ab96 100644 --- a/client/iface/configurer/common.go +++ b/client/iface/configurer/common.go @@ -3,6 +3,7 @@ package configurer import ( "net" "net/netip" + "time" "golang.zx2c4.com/wireguard/wgctrl/wgtypes" ) @@ -19,6 +20,43 @@ func buildPresharedKeyConfig(peerKey wgtypes.Key, psk wgtypes.Key, updateOnly bo } } +// buildIdlePeerEndpointConfig creates a config that removes and re-creates a peer in a +// single transaction with the given allowed IPs, endpoint and disabled keepalive. +func buildIdlePeerEndpointConfig(peerKey wgtypes.Key, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) wgtypes.Config { + keepAlive := time.Duration(0) + return wgtypes.Config{ + Peers: []wgtypes.PeerConfig{ + { + PublicKey: peerKey, + Remove: true, + }, + { + PublicKey: peerKey, + AllowedIPs: prefixesToIPNets(allowedIPs), + Endpoint: endpoint, + PersistentKeepaliveInterval: &keepAlive, + }, + }, + } +} + +// mergePrefixes returns the union of the two prefix lists, keeping the original order and +// dropping duplicates. +func mergePrefixes(current, base []netip.Prefix) []netip.Prefix { + merged := make([]netip.Prefix, 0, len(current)+len(base)) + seen := make(map[netip.Prefix]struct{}, len(current)+len(base)) + for _, group := range [][]netip.Prefix{current, base} { + for _, prefix := range group { + if _, ok := seen[prefix]; ok { + continue + } + seen[prefix] = struct{}{} + merged = append(merged, prefix) + } + } + return merged +} + func prefixesToIPNets(prefixes []netip.Prefix) []net.IPNet { ipNets := make([]net.IPNet, len(prefixes)) for i, prefix := range prefixes { diff --git a/client/iface/configurer/kernel_unix.go b/client/iface/configurer/kernel_unix.go index da69c2a35..f335b5006 100644 --- a/client/iface/configurer/kernel_unix.go +++ b/client/iface/configurer/kernel_unix.go @@ -3,6 +3,7 @@ package configurer import ( + "errors" "fmt" "net" "net/netip" @@ -145,6 +146,39 @@ func (c *KernelConfigurer) RemovePeer(peerKey string) error { return nil } +// IdlePeerEndpoint re-creates the peer pointing at the lazy wake endpoint, dropping +// handshake state while preserving the peer's currently installed allowed IPs. +func (c *KernelConfigurer) IdlePeerEndpoint(peerKey string, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) error { + peerKeyParsed, err := wgtypes.ParseKey(peerKey) + if err != nil { + return err + } + + var current []netip.Prefix + existing, err := c.getPeer(c.deviceName, peerKey) + switch { + case errors.Is(err, ErrPeerNotFound): + case err != nil: + return fmt.Errorf("get peer: %w", err) + default: + for _, ipNet := range existing.AllowedIPs { + addr, ok := netip.AddrFromSlice(ipNet.IP) + if !ok { + log.Warnf("failed to convert allowed IP %s of peer %s", ipNet.String(), peerKey) + continue + } + ones, _ := ipNet.Mask.Size() + current = append(current, netip.PrefixFrom(addr.Unmap(), ones)) + } + } + + config := buildIdlePeerEndpointConfig(peerKeyParsed, mergePrefixes(current, allowedIPs), endpoint) + if err := c.configure(config); err != nil { + return fmt.Errorf(`received error "%w" while setting idle endpoint for peer %s on interface %s`, err, peerKey, c.deviceName) + } + return nil +} + func (c *KernelConfigurer) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error { ipNet := net.IPNet{ IP: allowedIP.Addr().AsSlice(), diff --git a/client/iface/configurer/usp.go b/client/iface/configurer/usp.go index 0a25c55bc..e7b10903f 100644 --- a/client/iface/configurer/usp.go +++ b/client/iface/configurer/usp.go @@ -210,6 +210,64 @@ func (c *WGUSPConfigurer) RemovePeer(peerKey string) error { return ipcErr } +// IdlePeerEndpoint re-creates the peer pointing at the lazy wake endpoint, dropping +// handshake state while preserving the peer's currently installed allowed IPs. +func (c *WGUSPConfigurer) IdlePeerEndpoint(peerKey string, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) error { + peerKeyParsed, err := wgtypes.ParseKey(peerKey) + if err != nil { + return err + } + + current, err := c.currentAllowedIPs(hex.EncodeToString(peerKeyParsed[:])) + if err != nil { + return fmt.Errorf("get current allowed IPs: %w", err) + } + + config := buildIdlePeerEndpointConfig(peerKeyParsed, mergePrefixes(current, allowedIPs), endpoint) + if err := c.device.IpcSet(toWgUserspaceString(config)); err != nil { + return fmt.Errorf("set idle peer endpoint: %w", err) + } + + if endpoint != nil { + addr, err := netip.ParseAddr(endpoint.IP.String()) + if err != nil { + return fmt.Errorf("parse endpoint address: %w", err) + } + c.activityRecorder.UpsertAddress(peerKey, netip.AddrPortFrom(addr.Unmap(), uint16(endpoint.Port))) + } + return nil +} + +// currentAllowedIPs returns the allowed IPs currently installed for the peer identified by +// its hex-encoded public key. It returns an empty list when the peer does not exist. +func (c *WGUSPConfigurer) currentAllowedIPs(hexKey string) ([]netip.Prefix, error) { + ipc, err := c.device.IpcGet() + if err != nil { + return nil, err + } + + var prefixes []netip.Prefix + foundPeer := false + for _, line := range strings.Split(ipc, "\n") { + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "public_key=") { + foundPeer = line == "public_key="+hexKey + continue + } + if !foundPeer || !strings.HasPrefix(line, "allowed_ip=") { + continue + } + + prefix, err := netip.ParsePrefix(strings.TrimPrefix(line, "allowed_ip=")) + if err != nil { + log.Warnf("failed to parse allowed IP %q: %v", line, err) + continue + } + prefixes = append(prefixes, prefix) + } + return prefixes, nil +} + func (c *WGUSPConfigurer) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error { ipNet := net.IPNet{ IP: allowedIP.Addr().AsSlice(), diff --git a/client/iface/configurer/usp_idle_test.go b/client/iface/configurer/usp_idle_test.go new file mode 100644 index 000000000..6ed1dbc9a --- /dev/null +++ b/client/iface/configurer/usp_idle_test.go @@ -0,0 +1,158 @@ +package configurer + +import ( + "encoding/hex" + "net" + "net/netip" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + wgconn "golang.zx2c4.com/wireguard/conn" + "golang.zx2c4.com/wireguard/device" + "golang.zx2c4.com/wireguard/tun/netstack" + "golang.zx2c4.com/wireguard/wgctrl/wgtypes" + + "github.com/netbirdio/netbird/client/iface/bind" +) + +// newTestConfigurer creates a configurer backed by an in-memory wireguard-go device. +func newTestConfigurer(t *testing.T) *WGUSPConfigurer { + t.Helper() + + tunDev, _, err := netstack.CreateNetTUN([]netip.Addr{netip.MustParseAddr("100.64.0.1")}, []netip.Addr{}, 1280) + require.NoError(t, err) + + wgDev := device.NewDevice(tunDev, wgconn.NewDefaultBind(), device.NewLogger(device.LogLevelSilent, "[test] ")) + t.Cleanup(wgDev.Close) + + c := NewUSPConfigurerNoUAPI(wgDev, "utun-test", bind.NewActivityRecorder()) + + key, err := wgtypes.GeneratePrivateKey() + require.NoError(t, err) + require.NoError(t, c.ConfigureInterface(key.String(), 0)) + + return c +} + +// peerAllowedIPs returns the allowed IPs currently installed for the given peer. +func peerAllowedIPs(t *testing.T, c *WGUSPConfigurer, pubKey string) []string { + t.Helper() + + var ips []string + for _, line := range peerIpcLines(t, c, pubKey) { + if strings.HasPrefix(line, "allowed_ip=") { + ips = append(ips, strings.TrimPrefix(line, "allowed_ip=")) + } + } + return ips +} + +// peerEndpoint returns the endpoint currently installed for the given peer. +func peerEndpoint(t *testing.T, c *WGUSPConfigurer, pubKey string) string { + t.Helper() + + for _, line := range peerIpcLines(t, c, pubKey) { + if strings.HasPrefix(line, "endpoint=") { + return strings.TrimPrefix(line, "endpoint=") + } + } + return "" +} + +// peerIpcLines returns the uapi config lines belonging to the given peer. +func peerIpcLines(t *testing.T, c *WGUSPConfigurer, pubKey string) []string { + t.Helper() + + key, err := wgtypes.ParseKey(pubKey) + require.NoError(t, err) + hexKey := hex.EncodeToString(key[:]) + + ipc, err := c.device.IpcGet() + require.NoError(t, err) + + var lines []string + inPeer := false + for _, line := range strings.Split(ipc, "\n") { + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "public_key=") { + inPeer = line == "public_key="+hexKey + continue + } + if inPeer && line != "" { + lines = append(lines, line) + } + } + return lines +} + +// TestUSPConfigurer_IdlePeerEndpointPreservesAllowedIPs verifies the invariant the lazy idle +// transition relies on: IdlePeerEndpoint re-creates the peer (dropping handshake state via the +// remove+add transaction) while keeping every installed allowed IP, including routed +// prefixes added later by the route manager, and points the endpoint at the wake listener. +func TestUSPConfigurer_IdlePeerEndpointPreservesAllowedIPs(t *testing.T) { + c := newTestConfigurer(t) + + peerKey, err := wgtypes.GeneratePrivateKey() + require.NoError(t, err) + pubKey := peerKey.PublicKey().String() + + overlay := netip.MustParsePrefix("100.64.0.5/32") + routed := netip.MustParsePrefix("10.99.0.0/24") + + realEndpoint := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 51821} + require.NoError(t, c.UpdatePeer(pubKey, []netip.Prefix{overlay}, 25*time.Second, realEndpoint, nil)) + require.NoError(t, c.AddAllowedIP(pubKey, routed)) + + ips := peerAllowedIPs(t, c, pubKey) + require.Contains(t, ips, overlay.String(), "overlay prefix must be installed before the idle endpoint swap") + require.Contains(t, ips, routed.String(), "routed prefix must be installed before the idle endpoint swap") + + wakeEndpoint := &net.UDPAddr{IP: net.ParseIP("127.2.0.5"), Port: 17473} + require.NoError(t, c.IdlePeerEndpoint(pubKey, []netip.Prefix{overlay}, wakeEndpoint)) + + ips = peerAllowedIPs(t, c, pubKey) + assert.Contains(t, ips, routed.String(), "routed prefix must survive the idle endpoint swap") + assert.Contains(t, ips, overlay.String(), "overlay prefix must survive the idle endpoint swap") + assert.Equal(t, "127.2.0.5:17473", peerEndpoint(t, c, pubKey), "endpoint must point at the wake listener after the idle endpoint swap") +} + +// TestUSPConfigurer_IdlePeerEndpointCreatesMissingPeer verifies the cold-start arm path: when the +// peer does not exist yet (never connected), IdlePeerEndpoint creates it with the given base +// allowed IPs and the wake endpoint. +func TestUSPConfigurer_IdlePeerEndpointCreatesMissingPeer(t *testing.T) { + c := newTestConfigurer(t) + + peerKey, err := wgtypes.GeneratePrivateKey() + require.NoError(t, err) + pubKey := peerKey.PublicKey().String() + + overlay := netip.MustParsePrefix("100.64.0.5/32") + wakeEndpoint := &net.UDPAddr{IP: net.ParseIP("127.2.0.5"), Port: 17473} + + require.NoError(t, c.IdlePeerEndpoint(pubKey, []netip.Prefix{overlay}, wakeEndpoint)) + + assert.Equal(t, []string{overlay.String()}, peerAllowedIPs(t, c, pubKey), "missing peer must be created with the base allowed IPs") + assert.Equal(t, "127.2.0.5:17473", peerEndpoint(t, c, pubKey), "missing peer must be created with the wake endpoint") +} + +// TestUSPConfigurer_AddAllowedIPOnMissingPeerIsSilentNoOp documents the wireguard-go +// behavior the removed-peer idle flow raced against: AddAllowedIP uses update_only, +// which is a silent no-op when the peer does not exist. The idle transition must +// therefore keep the WireGuard peer (Conn.Idle + IdlePeerEndpoint) instead of leaving a +// window where the peer is absent. +func TestUSPConfigurer_AddAllowedIPOnMissingPeerIsSilentNoOp(t *testing.T) { + c := newTestConfigurer(t) + + peerKey, err := wgtypes.GeneratePrivateKey() + require.NoError(t, err) + pubKey := peerKey.PublicKey().String() + + routed := netip.MustParsePrefix("10.99.0.0/24") + + require.NoError(t, c.AddAllowedIP(pubKey, routed), "update-only on a missing peer must not return an error") + + assert.Empty(t, peerIpcLines(t, c, pubKey), "update-only call must not create the peer") +} diff --git a/client/iface/device/interface.go b/client/iface/device/interface.go index 7bab7b757..4861662a9 100644 --- a/client/iface/device/interface.go +++ b/client/iface/device/interface.go @@ -15,6 +15,7 @@ type WGConfigurer interface { ConfigureInterface(privateKey string, port int) error UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error RemovePeer(peerKey string) error + IdlePeerEndpoint(peerKey string, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) error AddAllowedIP(peerKey string, allowedIP netip.Prefix) error RemoveAllowedIP(peerKey string, allowedIP netip.Prefix) error SetPresharedKey(peerKey string, psk wgtypes.Key, updateOnly bool) error diff --git a/client/iface/iface.go b/client/iface/iface.go index 247f421a2..0381531b3 100644 --- a/client/iface/iface.go +++ b/client/iface/iface.go @@ -184,6 +184,19 @@ func (w *WGIface) RemovePeer(peerKey string) error { return w.configurer.RemovePeer(peerKey) } +// IdlePeerEndpoint re-creates the peer pointing at the lazy wake endpoint, dropping +// handshake state while preserving the peer's currently installed allowed IPs. +func (w *WGIface) IdlePeerEndpoint(peerKey string, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) error { + w.mu.Lock() + defer w.mu.Unlock() + if w.configurer == nil { + return ErrIfaceNotFound + } + + log.Debugf("Resetting peer on interface %s: %s, endpoint %s", w.tun.DeviceName(), peerKey, endpoint) + return w.configurer.IdlePeerEndpoint(peerKey, allowedIPs, endpoint) +} + // AddAllowedIP adds a prefix to the allowed IPs list of peer func (w *WGIface) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error { w.mu.Lock() diff --git a/client/internal/conn_mgr.go b/client/internal/conn_mgr.go index 77d1e6ca5..e1df336f8 100644 --- a/client/internal/conn_mgr.go +++ b/client/internal/conn_mgr.go @@ -228,7 +228,7 @@ func (e *ConnMgr) RemovePeerConn(peerKey string) { if !ok { return } - defer conn.Close(false) + defer conn.Close() if !e.isStartedWithLazyMgr() { return diff --git a/client/internal/engine.go b/client/internal/engine.go index 1d00ed0d2..f689dd2a6 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -1778,7 +1778,7 @@ func (e *Engine) addNewPeer(peerConfig *mgmProto.RemotePeerConfig) error { } if exists := e.connMgr.AddPeerConn(e.ctx, peerKey, conn); exists { - conn.Close(false) + conn.Close() return fmt.Errorf("peer already exists: %s", peerKey) } diff --git a/client/internal/engine_test.go b/client/internal/engine_test.go index fbd47ed74..2989033e4 100644 --- a/client/internal/engine_test.go +++ b/client/internal/engine_test.go @@ -53,6 +53,7 @@ type MockWGIface struct { UpdateAddrFunc func(newAddr wgaddr.Address) error UpdatePeerFunc func(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error RemovePeerFunc func(peerKey string) error + IdlePeerEndpointFunc func(peerKey string, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) error AddAllowedIPFunc func(peerKey string, allowedIP netip.Prefix) error RemoveAllowedIPFunc func(peerKey string, allowedIP netip.Prefix) error CloseFunc func() error @@ -124,6 +125,13 @@ func (m *MockWGIface) RemovePeer(peerKey string) error { return m.RemovePeerFunc(peerKey) } +func (m *MockWGIface) IdlePeerEndpoint(peerKey string, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) error { + if m.IdlePeerEndpointFunc == nil { + return nil + } + return m.IdlePeerEndpointFunc(peerKey, allowedIPs, endpoint) +} + func (m *MockWGIface) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error { return m.AddAllowedIPFunc(peerKey, allowedIP) } diff --git a/client/internal/iface_common.go b/client/internal/iface_common.go index 8ffa0b102..1d4d87832 100644 --- a/client/internal/iface_common.go +++ b/client/internal/iface_common.go @@ -32,6 +32,7 @@ type wgIfaceBase interface { UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error RemoveEndpointAddress(key string) error RemovePeer(peerKey string) error + IdlePeerEndpoint(peerKey string, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) error AddAllowedIP(peerKey string, allowedIP netip.Prefix) error RemoveAllowedIP(peerKey string, allowedIP netip.Prefix) error Close() error diff --git a/client/internal/lazyconn/activity/listener_bind.go b/client/internal/lazyconn/activity/listener_bind.go index 72a0cfc76..942e1fa84 100644 --- a/client/internal/lazyconn/activity/listener_bind.go +++ b/client/internal/lazyconn/activity/listener_bind.go @@ -106,7 +106,7 @@ func (d *BindListener) setupLazyConn() error { IP: d.fakeIP.AsSlice(), Port: lazyBindPort, } - return d.wgIface.UpdatePeer(d.peerCfg.PublicKey, d.peerCfg.AllowedIPs, 0, endpoint, nil) + return d.wgIface.IdlePeerEndpoint(d.peerCfg.PublicKey, d.peerCfg.AllowedIPs, endpoint) } // ReadPackets blocks until activity is detected on the LazyConn or the listener is closed. diff --git a/client/internal/lazyconn/activity/listener_bind_test.go b/client/internal/lazyconn/activity/listener_bind_test.go index 7026a9c97..f56b752c9 100644 --- a/client/internal/lazyconn/activity/listener_bind_test.go +++ b/client/internal/lazyconn/activity/listener_bind_test.go @@ -9,7 +9,6 @@ import ( log "github.com/sirupsen/logrus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "golang.zx2c4.com/wireguard/wgctrl/wgtypes" "github.com/netbirdio/netbird/client/iface/device" "github.com/netbirdio/netbird/client/iface/wgaddr" @@ -45,7 +44,7 @@ type MockWGIfaceBind struct { endpointMgr *mockEndpointManager } -func (m *MockWGIfaceBind) UpdatePeer(string, []netip.Prefix, time.Duration, *net.UDPAddr, *wgtypes.Key) error { +func (m *MockWGIfaceBind) IdlePeerEndpoint(string, []netip.Prefix, *net.UDPAddr) error { return nil } diff --git a/client/internal/lazyconn/activity/listener_udp.go b/client/internal/lazyconn/activity/listener_udp.go index 4b7e0ddf7..bda14e89a 100644 --- a/client/internal/lazyconn/activity/listener_udp.go +++ b/client/internal/lazyconn/activity/listener_udp.go @@ -101,7 +101,7 @@ func (d *UDPListener) Close() { func (d *UDPListener) createEndpoint() error { d.peerCfg.Log.Debugf("creating lazy endpoint: %s", d.endpoint.String()) - return d.wgIface.UpdatePeer(d.peerCfg.PublicKey, d.peerCfg.AllowedIPs, 0, d.endpoint, nil) + return d.wgIface.IdlePeerEndpoint(d.peerCfg.PublicKey, d.peerCfg.AllowedIPs, d.endpoint) } func (d *UDPListener) newConn() (*net.UDPConn, error) { diff --git a/client/internal/lazyconn/activity/manager.go b/client/internal/lazyconn/activity/manager.go index 9de8c0fa7..4666f689f 100644 --- a/client/internal/lazyconn/activity/manager.go +++ b/client/internal/lazyconn/activity/manager.go @@ -5,10 +5,8 @@ import ( "net" "net/netip" "sync" - "time" log "github.com/sirupsen/logrus" - "golang.zx2c4.com/wireguard/wgctrl/wgtypes" "github.com/netbirdio/netbird/client/iface/wgaddr" "github.com/netbirdio/netbird/client/internal/lazyconn" @@ -30,7 +28,7 @@ type Event struct { } type WgInterface interface { - UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error + IdlePeerEndpoint(peerKey string, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) error IsUserspaceBind() bool Address() wgaddr.Address MTU() uint16 diff --git a/client/internal/lazyconn/activity/manager_test.go b/client/internal/lazyconn/activity/manager_test.go index 07dd8d84c..c8416a257 100644 --- a/client/internal/lazyconn/activity/manager_test.go +++ b/client/internal/lazyconn/activity/manager_test.go @@ -8,7 +8,6 @@ import ( "time" log "github.com/sirupsen/logrus" - "golang.zx2c4.com/wireguard/wgctrl/wgtypes" "github.com/netbirdio/netbird/client/iface/wgaddr" "github.com/netbirdio/netbird/client/internal/lazyconn" @@ -26,7 +25,7 @@ func (m *MocPeer) ConnID() peerid.ConnID { type MocWGIface struct { } -func (m MocWGIface) UpdatePeer(string, []netip.Prefix, time.Duration, *net.UDPAddr, *wgtypes.Key) error { +func (m MocWGIface) IdlePeerEndpoint(string, []netip.Prefix, *net.UDPAddr) error { return nil } diff --git a/client/internal/lazyconn/manager/manager.go b/client/internal/lazyconn/manager/manager.go index 3868e37e8..5808260cf 100644 --- a/client/internal/lazyconn/manager/manager.go +++ b/client/internal/lazyconn/manager/manager.go @@ -280,7 +280,7 @@ func (m *Manager) DeactivatePeer(peerID peerid.ConnID) { return } - m.peerStore.PeerConnClose(mp.peerCfg.PublicKey) + m.peerStore.PeerConnIdle(mp.peerCfg.PublicKey, false) mp.peerCfg.Log.Infof("start activity monitor") @@ -569,7 +569,7 @@ func (m *Manager) onPeerInactivityTimedOut(peerIDs map[string]struct{}) { mp.peerCfg.Log.Infof("connection timed out") // this is blocking operation, potentially can be optimized - m.peerStore.PeerConnIdle(mp.peerCfg.PublicKey) + m.peerStore.PeerConnIdle(mp.peerCfg.PublicKey, true) mp.expectedWatcher = watcherActivity diff --git a/client/internal/lazyconn/wgiface.go b/client/internal/lazyconn/wgiface.go index f003ab3cf..bfdef6dd4 100644 --- a/client/internal/lazyconn/wgiface.go +++ b/client/internal/lazyconn/wgiface.go @@ -3,9 +3,6 @@ package lazyconn import ( "net" "net/netip" - "time" - - "golang.zx2c4.com/wireguard/wgctrl/wgtypes" "github.com/netbirdio/netbird/client/iface/wgaddr" "github.com/netbirdio/netbird/monotime" @@ -13,7 +10,7 @@ import ( type WGIface interface { RemovePeer(peerKey string) error - UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error + IdlePeerEndpoint(peerKey string, allowedIPs []netip.Prefix, endpoint *net.UDPAddr) error IsUserspaceBind() bool Address() wgaddr.Address LastActivities() map[string]monotime.Time diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go index 09a4e8b02..9e7b1de13 100644 --- a/client/internal/peer/conn.go +++ b/client/internal/peer/conn.go @@ -285,8 +285,20 @@ func (conn *Conn) open(engineCtx context.Context, firstPacket []byte) error { return nil } -// Close closes this peer Conn issuing a close event to the Conn closeCh -func (conn *Conn) Close(signalToRemote bool) { +// Close closes this peer Conn issuing a close event to the Conn closeCh and removes the WireGuard peer +func (conn *Conn) Close() { + conn.close(false, true) +} + +// Idle tears down the connection for the lazy idle state: transports and proxies are +// closed but the WireGuard peer is kept, so its AllowedIPs (including routed prefixes +// installed by the route manager) survive until the activity listener re-points the +// endpoint at the wake listener. +func (conn *Conn) Idle(signalToRemote bool) { + conn.close(signalToRemote, false) +} + +func (conn *Conn) close(signalToRemote bool, removeWgPeer bool) { conn.mu.Lock() defer conn.wgWatcherWg.Wait() defer conn.mu.Unlock() @@ -329,8 +341,12 @@ func (conn *Conn) Close(signalToRemote bool) { conn.wgProxyICE = nil } - if err := conn.endpointUpdater.RemoveWgPeer(); err != nil { - conn.Log.Errorf("failed to remove wg endpoint: %v", err) + if removeWgPeer { + if err := conn.endpointUpdater.RemoveWgPeer(); err != nil { + conn.Log.Errorf("failed to remove wg endpoint: %v", err) + } + } else { + conn.endpointUpdater.CancelPendingUpdates() } if conn.evalStatus() == StatusConnected && conn.onDisconnected != nil { diff --git a/client/internal/peer/endpoint.go b/client/internal/peer/endpoint.go index 9ba1efb6e..424b72d73 100644 --- a/client/internal/peer/endpoint.go +++ b/client/internal/peer/endpoint.go @@ -65,6 +65,16 @@ func (e *EndpointUpdater) RemoveWgPeer() error { return e.wgConfig.WgInterface.RemovePeer(e.wgConfig.RemoteKey) } +// CancelPendingUpdates stops a scheduled delayed endpoint update without touching the +// WireGuard peer. Used on the idle transition where the peer is kept so a pending +// responder-side update cannot overwrite the wake endpoint later. +func (e *EndpointUpdater) CancelPendingUpdates() { + e.mu.Lock() + defer e.mu.Unlock() + + e.waitForCloseTheDelayedUpdate() +} + func (e *EndpointUpdater) RemoveEndpointAddress() error { e.mu.Lock() defer e.mu.Unlock() diff --git a/client/internal/peer/endpoint_test.go b/client/internal/peer/endpoint_test.go new file mode 100644 index 000000000..f1affe88c --- /dev/null +++ b/client/internal/peer/endpoint_test.go @@ -0,0 +1,99 @@ +package peer + +import ( + "net" + "net/netip" + "sync" + "testing" + "time" + + log "github.com/sirupsen/logrus" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.zx2c4.com/wireguard/wgctrl/wgtypes" + + "github.com/netbirdio/netbird/client/iface/configurer" + "github.com/netbirdio/netbird/client/iface/wgaddr" + "github.com/netbirdio/netbird/client/iface/wgproxy" +) + +type endpointTestWGIface struct { + mu sync.Mutex + updateCalls int + removeCalls int +} + +func (m *endpointTestWGIface) UpdatePeer(string, []netip.Prefix, time.Duration, *net.UDPAddr, *wgtypes.Key) error { + m.mu.Lock() + defer m.mu.Unlock() + m.updateCalls++ + return nil +} + +func (m *endpointTestWGIface) RemovePeer(string) error { + m.mu.Lock() + defer m.mu.Unlock() + m.removeCalls++ + return nil +} + +func (m *endpointTestWGIface) GetStats() (map[string]configurer.WGStats, error) { return nil, nil } + +func (m *endpointTestWGIface) GetProxy() wgproxy.Proxy { return nil } + +func (m *endpointTestWGIface) Address() wgaddr.Address { return wgaddr.Address{} } + +func (m *endpointTestWGIface) RemoveEndpointAddress(string) error { return nil } + +func (m *endpointTestWGIface) counts() (updates, removes int) { + m.mu.Lock() + defer m.mu.Unlock() + return m.updateCalls, m.removeCalls +} + +func newTestEndpointUpdater(iface *endpointTestWGIface, initiator bool) *EndpointUpdater { + cfg := WgConfig{ + RemoteKey: "remoteKey", + WgInterface: iface, + AllowedIps: []netip.Prefix{netip.MustParsePrefix("100.64.0.5/32")}, + } + return NewEndpointUpdater(log.WithField("peer", "test"), cfg, initiator) +} + +// TestEndpointUpdater_CancelPendingUpdates ensures a scheduled responder-side delayed +// update is stopped without removing the WireGuard peer, so an idle transition cannot +// be overwritten later by a stale endpoint update. +func TestEndpointUpdater_CancelPendingUpdates(t *testing.T) { + iface := &endpointTestWGIface{} + e := newTestEndpointUpdater(iface, false) + + addr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 51820} + require.NoError(t, e.ConfigureWGEndpoint(addr, nil)) + + updates, _ := iface.counts() + require.Equal(t, 1, updates, "responder must apply the immediate nil-endpoint update") + + // CancelPendingUpdates waits for the delayed-update goroutine to exit, so the + // call counts below are final: the 5s fallback update can never fire anymore. + e.CancelPendingUpdates() + + updates, removes := iface.counts() + assert.Equal(t, 1, updates, "delayed endpoint update must not fire after cancellation") + assert.Equal(t, 0, removes, "cancellation must not remove the WireGuard peer") +} + +// TestEndpointUpdater_CancelPendingUpdatesNoPending ensures cancellation is a safe no-op +// when no delayed update is scheduled (initiator path). +func TestEndpointUpdater_CancelPendingUpdatesNoPending(t *testing.T) { + iface := &endpointTestWGIface{} + e := newTestEndpointUpdater(iface, true) + + addr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 51820} + require.NoError(t, e.ConfigureWGEndpoint(addr, nil)) + + e.CancelPendingUpdates() + + updates, removes := iface.counts() + assert.Equal(t, 1, updates, "initiator applies exactly one direct update") + assert.Equal(t, 0, removes, "cancellation must not remove the WireGuard peer") +} diff --git a/client/internal/peerstore/store.go b/client/internal/peerstore/store.go index 112caa101..b9745ffd8 100644 --- a/client/internal/peerstore/store.go +++ b/client/internal/peerstore/store.go @@ -108,7 +108,10 @@ func (s *Store) PeerConnOpenWithFirstPacket(ctx context.Context, pubKey string, } } -func (s *Store) PeerConnIdle(pubKey string) { +// PeerConnIdle transitions the peer connection to the lazy idle state, keeping the +// WireGuard peer and its AllowedIPs in place. signalToRemote indicates whether the +// remote peer should be notified (false when the remote initiated the idle via GOAWAY). +func (s *Store) PeerConnIdle(pubKey string, signalToRemote bool) { s.peerConnsMu.RLock() defer s.peerConnsMu.RUnlock() @@ -116,18 +119,7 @@ func (s *Store) PeerConnIdle(pubKey string) { if !ok { return } - p.Close(true) -} - -func (s *Store) PeerConnClose(pubKey string) { - s.peerConnsMu.RLock() - defer s.peerConnsMu.RUnlock() - - p, ok := s.peerConns[pubKey] - if !ok { - return - } - p.Close(false) + p.Idle(signalToRemote) } func (s *Store) PeersPubKey() []string {