diff --git a/client/internal/routemanager/dnsinterceptor/handler.go b/client/internal/routemanager/dnsinterceptor/handler.go index b784cc274..f92300bfd 100644 --- a/client/internal/routemanager/dnsinterceptor/handler.go +++ b/client/internal/routemanager/dnsinterceptor/handler.go @@ -95,7 +95,7 @@ func (d *DnsInterceptor) RemoveRoute() error { // AllowedIPs should use real IPs if d.currentPeerKey != "" { - if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil { + if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil { merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err)) } } @@ -172,7 +172,7 @@ func (d *DnsInterceptor) removeAllowedIP(realPrefix netip.Prefix) error { } // AllowedIPs use real IPs - if _, err := d.allowedIPsRefcounter.Decrement(realPrefix); err != nil { + if _, err := d.allowedIPsRefcounter.Decrement(realPrefix, d.currentPeerKey); err != nil { return fmt.Errorf("remove allowed IP %s: %v", realPrefix, err) } @@ -205,7 +205,7 @@ func (d *DnsInterceptor) RemoveAllowedIPs() error { for _, prefixes := range d.interceptedDomains { for _, prefix := range prefixes { // AllowedIPs use real IPs - if _, err := d.allowedIPsRefcounter.Decrement(prefix); err != nil { + if _, err := d.allowedIPsRefcounter.Decrement(prefix, d.currentPeerKey); err != nil { merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %v", prefix, err)) } } diff --git a/client/internal/routemanager/dynamic/route.go b/client/internal/routemanager/dynamic/route.go index f0efd7b22..88f67a6f4 100644 --- a/client/internal/routemanager/dynamic/route.go +++ b/client/internal/routemanager/dynamic/route.go @@ -135,7 +135,7 @@ func (r *Route) RemoveAllowedIPs() error { var merr *multierror.Error for _, domainPrefixes := range r.dynamicDomains { for _, prefix := range domainPrefixes { - if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil { + if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil { merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err)) } } @@ -320,7 +320,7 @@ func (r *Route) removeRoutes(prefixes []netip.Prefix) ([]netip.Prefix, error) { merr = multierror.Append(merr, fmt.Errorf("remove dynamic route for IP %s: %w", prefix, err)) } if r.currentPeerKey != "" { - if _, err := r.allowedIPsRefcounter.Decrement(prefix); err != nil { + if _, err := r.allowedIPsRefcounter.Decrement(prefix, r.currentPeerKey); err != nil { merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %s: %w", prefix, err)) } } diff --git a/client/internal/routemanager/manager.go b/client/internal/routemanager/manager.go index 66b24cc5a..a784f90cc 100644 --- a/client/internal/routemanager/manager.go +++ b/client/internal/routemanager/manager.go @@ -215,7 +215,7 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) { ) } - m.allowedIPsRefCounter = refcounter.New( + m.allowedIPsRefCounter = refcounter.NewAllowedIPs( func(prefix netip.Prefix, peerKey string) (string, error) { // save peerKey to use it in the remove function return peerKey, m.wgInterface.AddAllowedIP(peerKey, prefix) diff --git a/client/internal/routemanager/refcounter/allowedips.go b/client/internal/routemanager/refcounter/allowedips.go new file mode 100644 index 000000000..3b92c2018 --- /dev/null +++ b/client/internal/routemanager/refcounter/allowedips.go @@ -0,0 +1,185 @@ +package refcounter + +import ( + "errors" + "fmt" + "net/netip" + "sort" + "sync" + + "github.com/hashicorp/go-multierror" + + nberrors "github.com/netbirdio/netbird/client/errors" +) + +// allowedIPsEntry holds the per-peer reference counts for a single prefix and which peer is +// currently installed in WireGuard. WireGuard allows a prefix on exactly one peer, so at most +// one peer is active at a time even when several peers reference the prefix. +type allowedIPsEntry struct { + // peers maps a peerKey to the number of references holding the prefix for that peer. + peers map[string]int + // active is the peerKey currently installed in WireGuard for this prefix ("" if none). + active string + // total is the sum of all per-peer reference counts (kept in sync with peers). + total int +} + +// AllowedIPsRefCounter is a peer-aware reference counter for WireGuard AllowedIPs. +// +// The generic Counter keys only by prefix and remembers a single Out value set by the first +// caller, which it never changes. That is wrong for AllowedIPs: two independent watchers (or +// multiple resolved domains) can reference the same prefix through different peers, and when the +// peer currently installed in WireGuard releases its last reference the prefix must be handed over +// to a surviving peer instead of being left pointing at the released one. +// +// It calls add/remove (which program WireGuard) only on the transitions that matter: +// - add on the first reference for a prefix, or when swapping the active peer; +// - remove on the last reference for a prefix, or on the old peer during a swap. +type AllowedIPsRefCounter struct { + mu sync.Mutex + entries map[netip.Prefix]*allowedIPsEntry + add AddFunc[netip.Prefix, string, string] + remove RemoveFunc[netip.Prefix, string] +} + +// NewAllowedIPs creates a new peer-aware AllowedIPs reference counter. +// add programs a prefix on a peer in WireGuard and returns the peerKey to store as the active peer. +// remove unprograms the prefix from the given peer. +func NewAllowedIPs(add AddFunc[netip.Prefix, string, string], remove RemoveFunc[netip.Prefix, string]) *AllowedIPsRefCounter { + return &AllowedIPsRefCounter{ + entries: map[netip.Prefix]*allowedIPsEntry{}, + add: add, + remove: remove, + } +} + +// Increment adds a reference to prefix for peerKey. WireGuard is programmed only for the first +// reference to a prefix; while a different peer is already installed the prefix is left with it +// (first peer wins, HA at the WireGuard layer is not possible) and only the reference count is kept. +func (rm *AllowedIPsRefCounter) Increment(prefix netip.Prefix, peerKey string) (Ref[string], error) { + rm.mu.Lock() + defer rm.mu.Unlock() + + e, ok := rm.entries[prefix] + if !ok { + e = &allowedIPsEntry{peers: map[string]int{}} + rm.entries[prefix] = e + } + + logCallerF("Increasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]", + prefix, peerKey, e.peers[peerKey], e.peers[peerKey]+1, e.total, e.total+1, e.active) + + // Program WireGuard only when nothing is installed yet for this prefix. + if e.active == "" { + out, err := rm.add(prefix, peerKey) + if errors.Is(err, ErrIgnore) { + if e.total == 0 { + delete(rm.entries, prefix) + } + return Ref[string]{Count: e.total, Out: e.active}, nil + } + if err != nil { + if e.total == 0 { + delete(rm.entries, prefix) + } + return Ref[string]{}, fmt.Errorf("failed to add allowed IP %v for peer %s: %w", prefix, peerKey, err) + } + e.active = out + } + + e.peers[peerKey]++ + e.total++ + + return Ref[string]{Count: e.total, Out: e.active}, nil +} + +// Decrement removes a reference to prefix for peerKey. When the peer currently installed in +// WireGuard releases its last reference, the prefix is swapped to a surviving peer if one exists, +// otherwise it is removed from WireGuard. +func (rm *AllowedIPsRefCounter) Decrement(prefix netip.Prefix, peerKey string) (Ref[string], error) { + rm.mu.Lock() + defer rm.mu.Unlock() + + e, ok := rm.entries[prefix] + if !ok { + logCallerF("No allowed IP reference found for prefix %v", prefix) + return Ref[string]{}, nil + } + + if e.peers[peerKey] == 0 { + logCallerF("No allowed IP reference found for prefix %v peer %s", prefix, peerKey) + return Ref[string]{Count: e.total, Out: e.active}, nil + } + + logCallerF("Decreasing allowed IP ref count for prefix %v peer %s [peer %d -> %d, total %d -> %d, active %q]", + prefix, peerKey, e.peers[peerKey], e.peers[peerKey]-1, e.total, e.total-1, e.active) + + e.peers[peerKey]-- + e.total-- + if e.peers[peerKey] == 0 { + delete(e.peers, peerKey) + } + + // Nothing to reprogram unless the peer we just released is the one installed in WireGuard + // and it no longer holds any reference. + if peerKey != e.active || e.peers[peerKey] > 0 { + return Ref[string]{Count: e.total, Out: e.active}, nil + } + + // The active peer is gone. Hand the prefix over to a surviving peer, or remove it entirely. + if survivor, ok := pickSurvivor(e.peers); ok { + if err := rm.remove(prefix, e.active); err != nil { + return Ref[string]{Count: e.total, Out: e.active}, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err) + } + out, err := rm.add(prefix, survivor) + if err != nil { + e.active = "" + return Ref[string]{Count: e.total, Out: e.active}, fmt.Errorf("swap allowed IP %v to peer %s: %w", prefix, survivor, err) + } + e.active = out + return Ref[string]{Count: e.total, Out: e.active}, nil + } + + if err := rm.remove(prefix, e.active); err != nil { + return Ref[string]{}, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err) + } + delete(rm.entries, prefix) + + return Ref[string]{Count: 0, Out: ""}, nil +} + +// Flush removes all prefixes from WireGuard and clears the counter. +func (rm *AllowedIPsRefCounter) Flush() error { + rm.mu.Lock() + defer rm.mu.Unlock() + + var merr *multierror.Error + for prefix, e := range rm.entries { + if e.active == "" { + continue + } + logCallerF("Flushing allowed IP for prefix %v peer %s", prefix, e.active) + if err := rm.remove(prefix, e.active); err != nil { + merr = multierror.Append(merr, fmt.Errorf("remove allowed IP %v for peer %s: %w", prefix, e.active, err)) + } + } + + clear(rm.entries) + + return nberrors.FormatErrorOrNil(merr) +} + +// pickSurvivor deterministically selects a peer still referencing the prefix. WireGuard cannot do +// multipath for a single prefix, so any surviving peer is a valid winner; the choice is made stable +// (lowest peerKey) for predictable behavior and testability. +func pickSurvivor(peers map[string]int) (string, bool) { + if len(peers) == 0 { + return "", false + } + keys := make([]string, 0, len(peers)) + for k := range peers { + keys = append(keys, k) + } + sort.Strings(keys) + return keys[0], true +} diff --git a/client/internal/routemanager/refcounter/allowedips_test.go b/client/internal/routemanager/refcounter/allowedips_test.go new file mode 100644 index 000000000..1d4b58143 --- /dev/null +++ b/client/internal/routemanager/refcounter/allowedips_test.go @@ -0,0 +1,171 @@ +package refcounter + +import ( + "net/netip" + "testing" +) + +// fakeWG models WireGuard's cryptokey routing: a prefix can be installed on exactly one peer. +type fakeWG struct { + installed map[netip.Prefix]string + adds int + removes int +} + +func newFakeWG() *fakeWG { + return &fakeWG{installed: map[netip.Prefix]string{}} +} + +func (f *fakeWG) counter() *AllowedIPsRefCounter { + return NewAllowedIPs( + func(prefix netip.Prefix, peerKey string) (string, error) { + f.adds++ + f.installed[prefix] = peerKey + return peerKey, nil + }, + func(prefix netip.Prefix, peerKey string) error { + f.removes++ + // only clear if this peer is the one installed, mirroring wg semantics + if f.installed[prefix] == peerKey { + delete(f.installed, prefix) + } + return nil + }, + ) +} + +func mustPrefix(t *testing.T, s string) netip.Prefix { + t.Helper() + p, err := netip.ParsePrefix(s) + if err != nil { + t.Fatalf("parse prefix %q: %v", s, err) + } + return p +} + +// TestAllowedIPs_SwapOnActivePeerRemoval reproduces the reported bug: two networks with the same +// prefix routed by different peers. Removing the network whose peer is installed must hand the +// prefix over to the surviving peer instead of leaving it on the removed one. +func TestAllowedIPs_SwapOnActivePeerRemoval(t *testing.T) { + f := newFakeWG() + c := f.counter() + p := mustPrefix(t, "10.44.8.0/24") + + if _, err := c.Increment(p, "peerA"); err != nil { + t.Fatal(err) + } + if _, err := c.Increment(p, "peerB"); err != nil { + t.Fatal(err) + } + // First peer wins while both are present. + if got := f.installed[p]; got != "peerA" { + t.Fatalf("expected peerA installed, got %q", got) + } + + // Remove the active peer's network -> must swap to peerB. + if _, err := c.Decrement(p, "peerA"); err != nil { + t.Fatal(err) + } + if got := f.installed[p]; got != "peerB" { + t.Fatalf("BUG: prefix stuck on removed peer, want peerB got %q", got) + } + + // Remove the last one -> prefix gone. + if _, err := c.Decrement(p, "peerB"); err != nil { + t.Fatal(err) + } + if _, ok := f.installed[p]; ok { + t.Fatalf("expected prefix removed, still installed on %q", f.installed[p]) + } +} + +// TestAllowedIPs_RemoveNonActivePeer removing a non-installed peer must not touch WireGuard. +func TestAllowedIPs_RemoveNonActivePeer(t *testing.T) { + f := newFakeWG() + c := f.counter() + p := mustPrefix(t, "10.44.8.0/24") + + _, _ = c.Increment(p, "peerA") + _, _ = c.Increment(p, "peerB") + removesBefore := f.removes + + if _, err := c.Decrement(p, "peerB"); err != nil { + t.Fatal(err) + } + if f.installed[p] != "peerA" { + t.Fatalf("active peer must stay peerA, got %q", f.installed[p]) + } + if f.removes != removesBefore { + t.Fatalf("removing a non-active peer must not call wg remove") + } +} + +// TestAllowedIPs_SamePeerMultipleRefs two references via the same peer must keep the prefix until +// the last reference is released (the reason the per-peer count must be an int, not a set). +func TestAllowedIPs_SamePeerMultipleRefs(t *testing.T) { + f := newFakeWG() + c := f.counter() + p := mustPrefix(t, "10.44.8.0/24") + + _, _ = c.Increment(p, "peerA") + _, _ = c.Increment(p, "peerA") + if f.adds != 1 { + t.Fatalf("expected a single wg add for the same peer, got %d", f.adds) + } + + if _, err := c.Decrement(p, "peerA"); err != nil { + t.Fatal(err) + } + if f.installed[p] != "peerA" { + t.Fatalf("prefix must stay while a reference remains, got %q", f.installed[p]) + } + if f.removes != 0 { + t.Fatalf("no wg remove expected while a reference remains, got %d", f.removes) + } + + if _, err := c.Decrement(p, "peerA"); err != nil { + t.Fatal(err) + } + if _, ok := f.installed[p]; ok { + t.Fatalf("prefix must be removed after last reference") + } +} + +// TestAllowedIPs_RefCountAndActive checks the Ref returned to callers (used for the HA-disabled log). +func TestAllowedIPs_RefCountAndActive(t *testing.T) { + f := newFakeWG() + c := f.counter() + p := mustPrefix(t, "10.44.8.0/24") + + ref, _ := c.Increment(p, "peerA") + if ref.Count != 1 || ref.Out != "peerA" { + t.Fatalf("want {1, peerA}, got {%d, %q}", ref.Count, ref.Out) + } + ref, _ = c.Increment(p, "peerB") + if ref.Count != 2 || ref.Out != "peerA" { + t.Fatalf("want {2, peerA}, got {%d, %q}", ref.Count, ref.Out) + } +} + +// TestAllowedIPs_Flush removes everything installed and clears the counter. +func TestAllowedIPs_Flush(t *testing.T) { + f := newFakeWG() + c := f.counter() + p1 := mustPrefix(t, "10.44.8.0/24") + p2 := mustPrefix(t, "10.44.9.0/24") + + _, _ = c.Increment(p1, "peerA") + _, _ = c.Increment(p2, "peerB") + + if err := c.Flush(); err != nil { + t.Fatal(err) + } + if len(f.installed) != 0 { + t.Fatalf("expected all prefixes removed, got %v", f.installed) + } + // After flush, a fresh increment must add again. + _, _ = c.Increment(p1, "peerC") + if f.installed[p1] != "peerC" { + t.Fatalf("counter not reset after flush") + } +} diff --git a/client/internal/routemanager/refcounter/types.go b/client/internal/routemanager/refcounter/types.go index aadac3e25..7da0e17e3 100644 --- a/client/internal/routemanager/refcounter/types.go +++ b/client/internal/routemanager/refcounter/types.go @@ -5,5 +5,7 @@ import "net/netip" // RouteRefCounter is a Counter for Route, it doesn't take any input on Increment and doesn't use any output on Decrement type RouteRefCounter = Counter[netip.Prefix, struct{}, struct{}] -// AllowedIPsRefCounter is a Counter for AllowedIPs, it takes a peer key on Increment and passes it back to Decrement -type AllowedIPsRefCounter = Counter[netip.Prefix, string, string] +// AllowedIPsRefCounter tracks WireGuard AllowedIPs per prefix. Unlike the generic Counter it is peer-aware: +// a prefix can be claimed by several peers at once and WireGuard allows a given prefix on exactly one peer, +// so the counter records the per-peer reference count and swaps the installed peer when the active one is released. +// See allowedips.go. diff --git a/client/internal/routemanager/static/route.go b/client/internal/routemanager/static/route.go index d480fdf00..7c45c3efa 100644 --- a/client/internal/routemanager/static/route.go +++ b/client/internal/routemanager/static/route.go @@ -15,6 +15,11 @@ type Route struct { route *route.Route routeRefCounter *refcounter.RouteRefCounter allowedIPsRefcounter *refcounter.AllowedIPsRefCounter + // currentPeerKey is the routing peer this watcher currently has the prefix installed on + // (the HA winner elected by the watcher). It can differ from route.Peer and change on + // failover, so it is recorded on AddAllowedIPs and used on RemoveAllowedIPs to decrement + // the exact peer that was incremented. + currentPeerKey string } func NewRoute(params common.HandlerParams) *Route { @@ -52,12 +57,14 @@ func (r *Route) AddAllowedIPs(peerKey string) error { ref.Out, ) } + r.currentPeerKey = peerKey return nil } func (r *Route) RemoveAllowedIPs() error { - if _, err := r.allowedIPsRefcounter.Decrement(r.route.Network); err != nil { + if _, err := r.allowedIPsRefcounter.Decrement(r.route.Network, r.currentPeerKey); err != nil { return err } + r.currentPeerKey = "" return nil }