diff --git a/client/internal/conn_mgr.go b/client/internal/conn_mgr.go index 77d1e6ca5..bb5095c6b 100644 --- a/client/internal/conn_mgr.go +++ b/client/internal/conn_mgr.go @@ -43,11 +43,21 @@ type ConnMgr struct { lazyConnMgr *manager.Manager + // reconcileRoutedIPs re-applies a peer's routed allowed IPs after its lazy wake endpoint is + // (re)armed (Mode A at arm time). Injected by the engine; nil disables the reconcile. + reconcileRoutedIPs func(peerKey string) error + wg sync.WaitGroup lazyCtx context.Context lazyCtxCancel context.CancelFunc } +// SetRoutedIPsReconciler injects the callback used to re-apply a peer's routed allowed IPs when +// its lazy wake endpoint is (re)armed. Must be called before the lazy manager starts. +func (e *ConnMgr) SetRoutedIPsReconciler(fn func(peerKey string) error) { + e.reconcileRoutedIPs = fn +} + func NewConnMgr(engineConfig *EngineConfig, statusRecorder *peer.Status, peerStore *peerstore.Store, iface lazyconn.WGIface) *ConnMgr { e := &ConnMgr{ peerStore: peerStore, @@ -274,6 +284,7 @@ func (e *ConnMgr) Close() { func (e *ConnMgr) initLazyManager(engineCtx context.Context) { cfg := manager.Config{ InactivityThreshold: inactivityThresholdEnv(), + ReconcileAllowedIPs: e.reconcileRoutedIPs, } e.lazyConnMgr = manager.NewManager(cfg, engineCtx, e.peerStore, e.iface) diff --git a/client/internal/engine.go b/client/internal/engine.go index 1d00ed0d2..b07247b8f 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -652,6 +652,12 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL) iceCfg := e.createICEConfig() e.connMgr = NewConnMgr(e.config, e.statusRecorder, e.peerStore, wgIface) + e.connMgr.SetRoutedIPsReconciler(func(peerKey string) error { + if e.routeManager == nil { + return nil + } + return e.routeManager.ReconcilePeerAllowedIPs(peerKey) + }) e.connMgr.Start(e.ctx) e.srWatcher = guard.NewSRWatcher(e.signal, e.relayManager, e.mobileDep.IFaceDiscover, iceCfg) diff --git a/client/internal/lazyconn/manager/manager.go b/client/internal/lazyconn/manager/manager.go index 3868e37e8..b7424bb2f 100644 --- a/client/internal/lazyconn/manager/manager.go +++ b/client/internal/lazyconn/manager/manager.go @@ -29,6 +29,11 @@ type managedPeer struct { type Config struct { InactivityThreshold *time.Duration + // ReconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is + // armed. The activity listener creates the wake peer with the overlay /32 only; without the + // routed prefixes WireGuard would not steer subnet-bound traffic to the wake endpoint, so an + // idle routing peer could never be woken by that traffic. Optional; nil disables the reconcile. + ReconcileAllowedIPs func(peerKey string) error } // Manager manages lazy connections @@ -56,6 +61,9 @@ type Manager struct { peerToHAGroups map[string][]route.HAUniqueID // peer ID -> HA groups they belong to haGroupToPeers map[route.HAUniqueID][]string // HA group -> peer IDs in the group routesMu sync.RWMutex + + // reconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is armed. + reconcileAllowedIPs func(peerKey string) error } // NewManager creates a new lazy connection manager @@ -73,6 +81,7 @@ func NewManager(config Config, engineCtx context.Context, peerStore *peerstore.S activityManager: activity.NewManager(wgIface), peerToHAGroups: make(map[string][]route.HAUniqueID), haGroupToPeers: make(map[route.HAUniqueID][]string), + reconcileAllowedIPs: config.ReconcileAllowedIPs, } if wgIface.IsUserspaceBind() { @@ -201,7 +210,7 @@ func (m *Manager) AddPeer(peerCfg lazyconn.PeerConfig) (bool, error) { return false, nil } - if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil { + if err := m.armActivityListener(peerCfg); err != nil { return false, err } @@ -288,7 +297,7 @@ func (m *Manager) DeactivatePeer(peerID peerid.ConnID) { m.inactivityManager.RemovePeer(mp.peerCfg.PublicKey) - if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil { + if err := m.armActivityListener(*mp.peerCfg); err != nil { mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err) return } @@ -465,6 +474,31 @@ func (m *Manager) close() { } // shouldDeferIdleForHA checks if peer should stay connected due to HA group requirements +// armRoutedAllowedIPs re-applies the peer's routed allowed IPs onto its freshly armed wake +// endpoint. The activity listener creates the wake peer with the overlay /32 only, so without +// this the routed prefixes would be missing and traffic to a routed subnet could not wake the +// idle routing peer. It is a no-op when no reconciler is configured. +// armActivityListener (re)arms the peer's wake endpoint via the activity manager and then +// re-applies its routed allowed IPs, so traffic to a routed subnet can wake an idle routing +// peer. The routed prefixes must be re-applied after the wake endpoint exists because the +// listener creates it with the overlay /32 only. +func (m *Manager) armActivityListener(peerCfg lazyconn.PeerConfig) error { + if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil { + return err + } + m.armRoutedAllowedIPs(&peerCfg) + return nil +} + +func (m *Manager) armRoutedAllowedIPs(peerCfg *lazyconn.PeerConfig) { + if m.reconcileAllowedIPs == nil { + return + } + if err := m.reconcileAllowedIPs(peerCfg.PublicKey); err != nil { + peerCfg.Log.Errorf("failed to reconcile routed allowed IPs on wake endpoint: %v", err) + } +} + func (m *Manager) shouldDeferIdleForHA(inactivePeers map[string]struct{}, peerID string) bool { m.routesMu.RLock() defer m.routesMu.RUnlock() @@ -577,7 +611,7 @@ func (m *Manager) onPeerInactivityTimedOut(peerIDs map[string]struct{}) { mp.peerCfg.Log.Infof("start activity monitor") - if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil { + if err := m.armActivityListener(*mp.peerCfg); err != nil { mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err) continue } diff --git a/client/internal/routemanager/manager.go b/client/internal/routemanager/manager.go index 66b24cc5a..40bcce84f 100644 --- a/client/internal/routemanager/manager.go +++ b/client/internal/routemanager/manager.go @@ -61,6 +61,7 @@ type Manager interface { InitialRouteRange() []string SetFirewall(firewall.Manager) error SetDNSForwarderPort(port uint16) + ReconcilePeerAllowedIPs(peerKey string) error Stop(stateManager *statemanager.Manager) } @@ -232,6 +233,32 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) { ) } +// ReconcilePeerAllowedIPs re-applies every routed allowed IP currently tracked for the peer +// onto the WireGuard device. The allowed-IP refcounter only calls its AddFunc (which pushes to +// the device) on a prefix's 0->1 transition, so a peer whose device entry was rebuilt without a +// matching refcounter change — e.g. a lazy connection cycling through idle->wake, which recreates +// the WireGuard peer with the overlay /32 only — ends up missing routed prefixes the refcounter +// still considers installed, and nothing retries. Calling this when the peer's WireGuard entry is +// (re)created restores convergence. It is add-only and idempotent: AddAllowedIP is update-only, so +// prefixes are re-added to an existing peer and an absent peer is left untouched. +func (m *DefaultManager) ReconcilePeerAllowedIPs(peerKey string) error { + if m.allowedIPsRefCounter == nil { + return nil + } + + prefixes := m.allowedIPsRefCounter.KeysMatching(func(out string) bool { + return out == peerKey + }) + + var merr *multierror.Error + for _, prefix := range prefixes { + if err := m.wgInterface.AddAllowedIP(peerKey, prefix); err != nil { + merr = multierror.Append(merr, fmt.Errorf("add allowed IP %s for peer %s: %w", prefix, peerKey, err)) + } + } + return nberrors.FormatErrorOrNil(merr) +} + // Init sets up the routing func (m *DefaultManager) Init() error { m.routeSelector = m.initSelector() diff --git a/client/internal/routemanager/mock.go b/client/internal/routemanager/mock.go index 937314995..c1620b24c 100644 --- a/client/internal/routemanager/mock.go +++ b/client/internal/routemanager/mock.go @@ -112,6 +112,11 @@ func (m *MockManager) SetFirewall(firewall.Manager) error { func (m *MockManager) SetDNSForwarderPort(port uint16) { } +// ReconcilePeerAllowedIPs mock implementation of ReconcilePeerAllowedIPs from Manager interface +func (m *MockManager) ReconcilePeerAllowedIPs(peerKey string) error { + return nil +} + // Stop mock implementation of Stop from Manager interface func (m *MockManager) Stop(stateManager *statemanager.Manager) { if m.StopFunc != nil { diff --git a/client/internal/routemanager/reconcile_test.go b/client/internal/routemanager/reconcile_test.go new file mode 100644 index 000000000..2a8a4dc10 --- /dev/null +++ b/client/internal/routemanager/reconcile_test.go @@ -0,0 +1,90 @@ +//go:build !windows + +package routemanager + +import ( + "net" + "net/netip" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.zx2c4.com/wireguard/tun/netstack" + + "github.com/netbirdio/netbird/client/iface/device" + "github.com/netbirdio/netbird/client/iface/wgaddr" + "github.com/netbirdio/netbird/client/internal/routemanager/refcounter" +) + +// reconcileWGMock is a minimal iface.WGIface that only records AddAllowedIP calls; every other +// method is an inert stub because ReconcilePeerAllowedIPs exercises none of them. +type reconcileWGMock struct { + mu sync.Mutex + adds map[string][]netip.Prefix +} + +func (m *reconcileWGMock) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error { + m.mu.Lock() + defer m.mu.Unlock() + if m.adds == nil { + m.adds = map[string][]netip.Prefix{} + } + m.adds[peerKey] = append(m.adds[peerKey], allowedIP) + return nil +} + +func (m *reconcileWGMock) added(peerKey string) []netip.Prefix { + m.mu.Lock() + defer m.mu.Unlock() + return m.adds[peerKey] +} + +func (m *reconcileWGMock) RemoveAllowedIP(string, netip.Prefix) error { return nil } +func (m *reconcileWGMock) Name() string { return "utun-test" } +func (m *reconcileWGMock) Address() wgaddr.Address { return wgaddr.Address{} } +func (m *reconcileWGMock) ToInterface() *net.Interface { return nil } +func (m *reconcileWGMock) IsUserspaceBind() bool { return false } +func (m *reconcileWGMock) GetFilter() device.PacketFilter { return nil } +func (m *reconcileWGMock) GetDevice() *device.FilteredDevice { return nil } +func (m *reconcileWGMock) GetNet() *netstack.Net { return nil } + +// TestReconcilePeerAllowedIPs verifies the declarative reconcile re-applies every routed prefix +// tracked for the peer (self-heal, independent of refcount level) and stays scoped to that peer. +func TestReconcilePeerAllowedIPs(t *testing.T) { + wg := &reconcileWGMock{} + m := &DefaultManager{wgInterface: wg} + m.allowedIPsRefCounter = refcounter.New[netip.Prefix, string, string]( + func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil }, + func(netip.Prefix, string) error { return nil }, + ) + + peerA1 := netip.MustParsePrefix("10.0.0.0/24") + peerA2 := netip.MustParsePrefix("10.1.0.0/24") + peerB1 := netip.MustParsePrefix("10.2.0.0/24") + + for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} { + _, err := m.allowedIPsRefCounter.Increment(prefix, peer) + require.NoError(t, err) + } + // Extra reference: reconcile must still re-apply the prefix even though its refcount never + // hit 0 again (the exact case the plain incremental path skips). + _, err := m.allowedIPsRefCounter.Increment(peerA1, "peerA") + require.NoError(t, err) + + require.NoError(t, m.ReconcilePeerAllowedIPs("peerA")) + + assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, wg.added("peerA"), + "reconcile must re-apply all routed prefixes of the peer") + assert.Empty(t, wg.added("peerB"), "reconcile must not touch another peer's prefixes") +} + +// TestReconcilePeerAllowedIPsNoCounter verifies reconcile is a safe no-op before the refcounter is +// set up. +func TestReconcilePeerAllowedIPsNoCounter(t *testing.T) { + wg := &reconcileWGMock{} + m := &DefaultManager{wgInterface: wg} + + require.NoError(t, m.ReconcilePeerAllowedIPs("peerA")) + assert.Empty(t, wg.added("peerA")) +} diff --git a/client/internal/routemanager/refcounter/refcounter.go b/client/internal/routemanager/refcounter/refcounter.go index 27a724f50..02fc835cb 100644 --- a/client/internal/routemanager/refcounter/refcounter.go +++ b/client/internal/routemanager/refcounter/refcounter.go @@ -94,6 +94,23 @@ func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) { return ref, ok } +// KeysMatching returns every key whose stored Out satisfies the predicate. The predicate is +// evaluated under the counter's lock, so keep it cheap and side-effect free. It is used to +// enumerate the keys (e.g. allowed-IP prefixes) associated with a given Out value (e.g. a peer +// key) without exposing the internal map. +func (rm *Counter[Key, I, O]) KeysMatching(pred func(out O) bool) []Key { + rm.mu.Lock() + defer rm.mu.Unlock() + + var keys []Key + for key, ref := range rm.refCountMap { + if pred(ref.Out) { + keys = append(keys, key) + } + } + return keys +} + // Increment increments the reference count for the given key. // If this is the first reference to the key, the AddFunc is called. func (rm *Counter[Key, I, O]) Increment(key Key, in I) (Ref[O], error) { diff --git a/client/internal/routemanager/refcounter/refcounter_test.go b/client/internal/routemanager/refcounter/refcounter_test.go new file mode 100644 index 000000000..686782fd5 --- /dev/null +++ b/client/internal/routemanager/refcounter/refcounter_test.go @@ -0,0 +1,43 @@ +package refcounter + +import ( + "net/netip" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestKeysMatching verifies KeysMatching returns exactly the keys whose stored Out satisfies the +// predicate — the primitive ReconcilePeerAllowedIPs relies on to enumerate a single peer's routed +// prefixes without touching the others. +func TestKeysMatching(t *testing.T) { + rc := New[netip.Prefix, string, string]( + func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil }, + func(netip.Prefix, string) error { return nil }, + ) + + peerA1 := netip.MustParsePrefix("10.0.0.0/24") + peerA2 := netip.MustParsePrefix("10.1.0.0/24") + peerB1 := netip.MustParsePrefix("10.2.0.0/24") + + _, err := rc.Increment(peerA1, "peerA") + require.NoError(t, err) + _, err = rc.Increment(peerA2, "peerA") + require.NoError(t, err) + _, err = rc.Increment(peerB1, "peerB") + require.NoError(t, err) + + // a second reference to peerA1 must not duplicate it in the result + _, err = rc.Increment(peerA1, "peerA") + require.NoError(t, err) + + keysA := rc.KeysMatching(func(out string) bool { return out == "peerA" }) + assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, keysA) + + keysB := rc.KeysMatching(func(out string) bool { return out == "peerB" }) + assert.ElementsMatch(t, []netip.Prefix{peerB1}, keysB) + + keysNone := rc.KeysMatching(func(out string) bool { return out == "missing" }) + assert.Empty(t, keysNone) +}