[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.
2026-10-26 22:39:07 +01:00 · 2026-07-19 21:37:15 +02:00
parent 92a5ed19d3
commit 13433af4a5
21 changed files with 455 additions and 34 deletions
@@ -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 {
@@ -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(),
@@ -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(),
@@ -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")
+}
@@ -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
@@ -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()