[client] Index peer tunnel IPs for faster PeerStateByIP lookup (#6412)

* [client] Index peer tunnel IPs for O(1) PeerStateByIP lookup

Replace the linear scan over all peers with an ipToKey map maintained
by AddPeer/RemovePeer, covering both IPv4 and IPv6 tunnel addresses.

Offline peers are intentionally no longer resolvable by IP: only active
peers can carry traffic, so IdentityForIP and the DNS disconnected-peer
filter now treat them as unknown, same as foreign IPs.

Skip the DNS answer filter for single-record responses; dropping the
only answer was always restored by the empty-answer escape hatch, so
the fast path is behavior-neutral.

* Ensure `ipToKey` entries are only removed if they match the peer being deleted, preventing accidental removal of unrelated mappings.

client/embed/embed.go

@@ -446,8 +446,8 @@ func (c *Client) Expose(ctx context.Context, req ExposeRequest) (*ExposeSession,
 
 // IdentityForIP looks up a remote peer by its tunnel IP using the
 // embedded client's status recorder. Returns the peer's WireGuard public
-// key and FQDN. ok=false means the IP isn't in this client's peer
-// roster — callers should treat that as "unknown peer".
+// key and FQDN. ok=false means the IP doesn't belong to an active peer
+// — offline roster peers are treated as unknown, same as foreign IPs.
 func (c *Client) IdentityForIP(ip netip.Addr) (pubKey, fqdn string, ok bool) {
 	if !ip.IsValid() || c.recorder == nil {
 		return "", "", false

client/internal/dns/local/local.go

@@ -482,7 +482,7 @@ func (d *Resolver) logDNSError(logger *log.Entry, hostname string, qtype uint16,
 // completely when every proxy peer is offline (the upstream may still
 // be reachable some other way, or the peerstore may be stale).
 func (d *Resolver) filterDisconnectedPeerAnswers(logger *log.Entry, question dns.Question, records []dns.RR) []dns.RR {
-	if len(records) == 0 {
+	if len(records) < 2 {
 		return records
 	}
 	d.mu.RLock()

client/internal/dns/local/local_test.go

@@ -2738,6 +2738,17 @@ func TestLocalResolver_FilterDisconnectedPeerAnswers(t *testing.T) {
 			connByIP:    nil,
 			wantInOrder: []string{"100.64.0.10", "100.64.0.11"},
 		},
+		{
+			// A single answer is never filtered: dropping it would only
+			// trigger the empty-answer escape hatch, so the fast path
+			// returns it untouched.
+			name:    "single disconnected answer passes through",
+			records: []nbdns.SimpleRecord{disconnectedRec},
+			connByIP: map[string]ipState{
+				"100.64.0.11": {known: true, connected: false},
+			},
+			wantInOrder: []string{"100.64.0.11"},
+		},
 	}
 
 	for _, tc := range tests {

client/internal/peer/conn_status.go

@@ -26,7 +26,6 @@ type connStatusInputs struct {
 	iceInProgress       bool // a negotiation is currently in flight
 }
 
-
 // ConnStatus describe the status of a peer's connection
 type ConnStatus int32
 

client/internal/peer/status.go

@@ -193,6 +193,7 @@ func (s *StatusChangeSubscription) Events() chan map[string]RouterState {
 type Status struct {
 	mux                   sync.RWMutex
 	peers                 map[string]State
+	ipToKey               map[string]string
 	changeNotify          map[string]map[string]*StatusChangeSubscription // map[peerID]map[subscriptionID]*StatusChangeSubscription
 	signalState           bool
 	signalError           error
@@ -231,6 +232,7 @@ type Status struct {
 func NewRecorder(mgmAddress string) *Status {
 	return &Status{
 		peers:                 make(map[string]State),
+		ipToKey:               make(map[string]string),
 		changeNotify:          make(map[string]map[string]*StatusChangeSubscription),
 		eventStreams:          make(map[string]chan *proto.SystemEvent),
 		eventQueue:            NewEventQueue(eventQueueSize),
@@ -282,6 +284,12 @@ func (d *Status) AddPeer(peerPubKey string, fqdn string, ip string, ipv6 string)
 		Mux:        new(sync.RWMutex),
 	}
 	d.peerListChangedForNotification = true
+	if ipv6 != "" {
+		d.ipToKey[ipv6] = peerPubKey
+	}
+	if ip != "" {
+		d.ipToKey[ip] = peerPubKey
+	}
 	return nil
 }
 
@@ -311,28 +319,22 @@ func (d *Status) PeerByIP(ip string) (string, bool) {
 
 // PeerStateByIP returns the full peer State for the given tunnel IP.
 // Matches against either the IPv4 (State.IP) or IPv6 (State.IPv6) tunnel
-// address so dual-stack peers are reachable on either family. Searches
-// both d.peers and d.offlinePeers — peers that have been moved into
-// the offline slice by ReplaceOfflinePeers are still part of the
-// account's roster and callers (DNS filter, embed.Client.IdentityForIP)
-// need to recognise them rather than treating them as unknown. Returns
-// the zero State and false when no peer matches or the input is empty.
+// address so dual-stack peers are reachable on either family. Only
+// active peers are matched; peers moved into the offline slice by
+// ReplaceOfflinePeers are intentionally treated as unknown.
 func (d *Status) PeerStateByIP(ip string) (State, bool) {
 	if ip == "" {
 		return State{}, false
 	}
 	d.mux.RLock()
 	defer d.mux.RUnlock()
-
-	for _, state := range d.peers {
-		if (state.IP != "" && state.IP == ip) || (state.IPv6 != "" && state.IPv6 == ip) {
-			return state, true
-		}
+	key, ok := d.ipToKey[ip]
+	if !ok {
+		return State{}, false
 	}
-	for _, state := range d.offlinePeers {
-		if (state.IP != "" && state.IP == ip) || (state.IPv6 != "" && state.IPv6 == ip) {
-			return state, true
-		}
+	state, ok := d.peers[key]
+	if ok {
+		return state, true
 	}
 	return State{}, false
 }
@@ -342,12 +344,18 @@ func (d *Status) RemovePeer(peerPubKey string) error {
 	d.mux.Lock()
 	defer d.mux.Unlock()
 
-	_, ok := d.peers[peerPubKey]
+	p, ok := d.peers[peerPubKey]
 	if !ok {
 		return errors.New("no peer with to remove")
 	}
 
 	delete(d.peers, peerPubKey)
+	if mappedKey, exists := d.ipToKey[p.IP]; exists && mappedKey == peerPubKey {
+		delete(d.ipToKey, p.IP)
+	}
+	if mappedKey, exists := d.ipToKey[p.IPv6]; exists && mappedKey == peerPubKey {
+		delete(d.ipToKey, p.IPv6)
+	}
 	d.peerListChangedForNotification = true
 	return nil
 }

client/internal/peer/status_test.go

@@ -90,12 +90,11 @@ func TestStatus_PeerStateByIP_MatchesIPv6(t *testing.T) {
 	req.Equal("pk-1", state.PubKey, "matching state must carry the right pub key")
 }
 
-// TestStatus_PeerStateByIP_MatchesOfflinePeers covers peers that have
-// been moved into the offline slice via ReplaceOfflinePeers. Callers
-// (DNS filter, embed.Client.IdentityForIP) need to treat them as known
-// rather than unknown — otherwise authentication / DNS filtering treats
-// known-but-offline peers as foreign IPs.
-func TestStatus_PeerStateByIP_MatchesOfflinePeers(t *testing.T) {
+// TestStatus_PeerStateByIP_IgnoresOfflinePeers documents that peers
+// moved into the offline slice via ReplaceOfflinePeers are intentionally
+// not resolvable by IP: only active peers can carry traffic, so callers
+// (DNS filter, embed.Client.IdentityForIP) treat them as unknown.
+func TestStatus_PeerStateByIP_IgnoresOfflinePeers(t *testing.T) {
 	status := NewRecorder("https://mgm")
 	req := require.New(t)
 
@@ -103,13 +102,31 @@ func TestStatus_PeerStateByIP_MatchesOfflinePeers(t *testing.T) {
 		{PubKey: "pk-offline", FQDN: "offline.netbird", IP: "100.64.0.20", IPv6: "fd00::20"},
 	})
 
-	state, ok := status.PeerStateByIP("100.64.0.20")
-	req.True(ok, "offline peer must resolve by IPv4 tunnel address")
-	req.Equal("pk-offline", state.PubKey, "matching state must carry the offline peer's pub key")
+	_, ok := status.PeerStateByIP("100.64.0.20")
+	req.False(ok, "offline peer must not resolve by IPv4 tunnel address")
 
-	state, ok = status.PeerStateByIP("fd00::20")
-	req.True(ok, "offline peer must resolve by IPv6 tunnel address")
-	req.Equal("pk-offline", state.PubKey, "IPv6 match must carry the offline peer's pub key")
+	_, ok = status.PeerStateByIP("fd00::20")
+	req.False(ok, "offline peer must not resolve by IPv6 tunnel address")
+}
+
+// TestStatus_PeerStateByIP_RemovedPeer verifies RemovePeer drops the
+// IP index entries for both address families.
+func TestStatus_PeerStateByIP_RemovedPeer(t *testing.T) {
+	status := NewRecorder("https://mgm")
+	req := require.New(t)
+
+	req.NoError(status.AddPeer("pk-1", "peer-1.netbird", "100.64.0.10", "fd00::1"))
+
+	_, ok := status.PeerStateByIP("100.64.0.10")
+	req.True(ok, "active peer must resolve before removal")
+
+	req.NoError(status.RemovePeer("pk-1"))
+
+	_, ok = status.PeerStateByIP("100.64.0.10")
+	req.False(ok, "removed peer must not resolve by IPv4 tunnel address")
+
+	_, ok = status.PeerStateByIP("fd00::1")
+	req.False(ok, "removed peer must not resolve by IPv6 tunnel address")
 }
 
 func TestStatus_UpdatePeerFQDN(t *testing.T) {