diff --git a/client/internal/conn_mgr.go b/client/internal/conn_mgr.go index 754ce37a3..7a591f60c 100644 --- a/client/internal/conn_mgr.go +++ b/client/internal/conn_mgr.go @@ -49,11 +49,21 @@ type ConnMgr struct { // engine.syncMsgMux; all other reads stay under engine.syncMsgMux only. lazyConnMgrMu sync.RWMutex + // 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, @@ -291,6 +301,7 @@ func (e *ConnMgr) Close() { func (e *ConnMgr) initLazyManager(engineCtx context.Context) { cfg := manager.Config{ InactivityThreshold: inactivityThresholdEnv(), + ReconcileAllowedIPs: e.reconcileRoutedIPs, } e.lazyConnMgrMu.Lock() diff --git a/client/internal/engine.go b/client/internal/engine.go index e1b03e878..617892e43 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -663,6 +663,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) // Wire DNS-time lazy-connection warm-up now that the connection manager 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..ef69b81a4 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,30 @@ 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 + } + + return m.allowedIPsRefCounter.ReapplyMatching( + func(out string) bool { return out == peerKey }, + func(prefix netip.Prefix) error { + if err := m.wgInterface.AddAllowedIP(peerKey, prefix); err != nil { + return fmt.Errorf("add allowed IP %s for peer %s: %w", prefix, peerKey, err) + } + return nil + }, + ) +} + // 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..917120275 100644 --- a/client/internal/routemanager/refcounter/refcounter.go +++ b/client/internal/routemanager/refcounter/refcounter.go @@ -94,6 +94,26 @@ func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) { return ref, ok } +// ReapplyMatching calls apply for every key whose stored Out satisfies pred, holding the +// counter lock for the whole pass. Running apply under the lock keeps it atomic with respect +// to Increment/Decrement: a prefix dropped to zero is removed from the map (and had its +// RemoveFunc called) before this pass observes it, so a stale key can never be re-applied. +// pred and apply are invoked under the lock, so they must not call back into the counter. +func (rm *Counter[Key, I, O]) ReapplyMatching(pred func(out O) bool, apply func(key Key) error) error { + rm.mu.Lock() + defer rm.mu.Unlock() + + var merr *multierror.Error + for key, ref := range rm.refCountMap { + if pred(ref.Out) { + if err := apply(key); err != nil { + merr = multierror.Append(merr, err) + } + } + } + return nberrors.FormatErrorOrNil(merr) +} + // 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..79a99c388 --- /dev/null +++ b/client/internal/routemanager/refcounter/refcounter_test.go @@ -0,0 +1,47 @@ +package refcounter + +import ( + "net/netip" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestReapplyMatching verifies ReapplyMatching invokes apply for exactly the keys whose stored +// Out satisfies the predicate (no duplicates for multiply-referenced keys) — the primitive +// ReconcilePeerAllowedIPs relies on to re-apply a single peer's routed prefixes. +func TestReapplyMatching(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") + + for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} { + _, err := rc.Increment(prefix, peer) + require.NoError(t, err) + } + // a second reference must not make the key applied twice + _, err := rc.Increment(peerA1, "peerA") + require.NoError(t, err) + + var applied []netip.Prefix + err = rc.ReapplyMatching( + func(out string) bool { return out == "peerA" }, + func(key netip.Prefix) error { applied = append(applied, key); return nil }, + ) + require.NoError(t, err) + assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, applied) + + var none []netip.Prefix + err = rc.ReapplyMatching( + func(out string) bool { return out == "missing" }, + func(key netip.Prefix) error { none = append(none, key); return nil }, + ) + require.NoError(t, err) + assert.Empty(t, none) +}