wgengine/wglog: stop using netmap.NetworkMap here too

This applies the same treatment from 8f210454dd (netlog) to wglog,
ending use of netmap.NetworkMap and instead getting the canonical data
from LocalBackend/nodeBackend.

This is a dependency to removing the netmap.NetworkMap from
upstream callers, like wgengine.Engine in general.

Updates #12542

Change-Id: Icb5af0799322def048a6f594b49f7d11273f025d
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

ipn/ipnlocal/local.go

@@ -646,6 +646,7 @@ func NewLocalBackend(logf logger.Logf, logID logid.PublicID, sys *tsd.System, lo
 	e.SetPeerByIPPacketFunc(b.lookupPeerByIP)
 	e.SetPeerSessionStateFunc(b.onPeerWireGuardState)
 	e.SetNetLogNodeSource(netLogNodeSource{b})
+	e.SetWGPeerLookup(b.lookupPeerWireGuardString)
 
 	if sys.InitialConfig != nil {
 		if err := b.initPrefsFromConfig(sys.InitialConfig); err != nil {
@@ -8090,6 +8091,25 @@ func (s netLogNodeSource) NodeByAddr(addr netip.Addr) (_ tailcfg.NodeView, _ tai
 // Compile-time assertion that netLogNodeSource implements [netlog.NodeSource].
 var _ netlog.NodeSource = netLogNodeSource{}
 
+// lookupPeerWireGuardString returns the Tailscale-conventional short string
+// (e.g. "[IMTBr]") for the peer whose wireguard-go-formatted public key
+// string is wgString (e.g. "peer(IMTB…r7lM)"), or "", false if no current
+// peer matches. It is installed on the engine via [Engine.SetWGPeerLookup]
+// in [NewLocalBackend] so that [wglog.Logger] can rewrite peer references
+// in wireguard-go log lines without any per-Reconfig denormalization.
+func (b *LocalBackend) lookupPeerWireGuardString(wgString string) (tsString string, ok bool) {
+	nb := b.currentNode()
+	nid, ok := nb.NodeByWireGuardString(wgString)
+	if !ok {
+		return "", false
+	}
+	nv, ok := nb.NodeByID(nid)
+	if !ok {
+		return "", false
+	}
+	return nv.Key().ShortString(), true
+}
+
 // ActiveSSHConns returns the number of active SSH connections,
 // or 0 if SSH is not linked into the binary or available on the platform.
 func (b *LocalBackend) ActiveSSHConns() int {

ipn/ipnlocal/node_backend.go

@@ -116,6 +116,15 @@ type nodeBackend struct {
 	// It is mutated in place (with mu held) and must not escape the [nodeBackend].
 	nodeByKey map[key.NodePublic]tailcfg.NodeID
 
+	// nodeByWGString indexes wireguard-go's truncated peer-string form
+	// (see [key.NodePublic.WireGuardGoString]) to node ID. It mirrors
+	// nodeByKey and lets the wireguard-go log path resolve
+	// "peer(XXXX…YYYY)" references in O(1), without scanning every
+	// peer, while still tolerating the fact that the wireguard-go form
+	// is lossy and can't be inverted to a [key.NodePublic].
+	// It is mutated in place (with mu held) and must not escape the [nodeBackend].
+	nodeByWGString map[string]tailcfg.NodeID
+
 	// userProfiles is the live set of user profiles, updated incrementally
 	// by mergeUserProfiles as deltas arrive. It parallels the peers map:
 	// netMap.UserProfiles is the frozen snapshot from the last full install,
@@ -229,6 +238,16 @@ func (nb *nodeBackend) NodeByKey(k key.NodePublic) (_ tailcfg.NodeID, ok bool) {
 	return nid, ok
 }
 
+// NodeByWireGuardString returns the node ID of the peer whose
+// [key.NodePublic.WireGuardGoString] form is s (e.g. "peer(IMTB…r7lM)").
+// ok is false if no current peer matches.
+func (nb *nodeBackend) NodeByWireGuardString(s string) (_ tailcfg.NodeID, ok bool) {
+	nb.mu.Lock()
+	defer nb.mu.Unlock()
+	nid, ok := nb.nodeByWGString[s]
+	return nid, ok
+}
+
 func (nb *nodeBackend) NodeByID(id tailcfg.NodeID) (_ tailcfg.NodeView, ok bool) {
 	nb.mu.Lock()
 	defer nb.mu.Unlock()
@@ -580,18 +599,26 @@ func (nb *nodeBackend) updateNodeByKeyLocked() {
 	nm := nb.netMap
 	if nm == nil {
 		nb.nodeByKey = nil
+		nb.nodeByWGString = nil
 		return
 	}
 
 	if nb.nodeByKey == nil {
 		nb.nodeByKey = map[key.NodePublic]tailcfg.NodeID{}
 	}
+	if nb.nodeByWGString == nil {
+		nb.nodeByWGString = map[string]tailcfg.NodeID{}
+	}
 	// First pass, mark everything unwanted.
 	for k := range nb.nodeByKey {
 		nb.nodeByKey[k] = 0
 	}
+	for k := range nb.nodeByWGString {
+		nb.nodeByWGString[k] = 0
+	}
 	addNode := func(n tailcfg.NodeView) {
 		nb.nodeByKey[n.Key()] = n.ID()
+		nb.nodeByWGString[n.Key().WireGuardGoString()] = n.ID()
 	}
 	if nm.SelfNode.Valid() {
 		addNode(nm.SelfNode)
@@ -605,6 +632,11 @@ func (nb *nodeBackend) updateNodeByKeyLocked() {
 			delete(nb.nodeByKey, k)
 		}
 	}
+	for k, v := range nb.nodeByWGString {
+		if v == 0 {
+			delete(nb.nodeByWGString, k)
+		}
+	}
 }
 
 func (nb *nodeBackend) updatePeersLocked() {
@@ -684,6 +716,7 @@ func (nb *nodeBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (handled bo
 				}
 			}
 			mak.Set(&nb.nodeByKey, m.Node.Key(), nid)
+			mak.Set(&nb.nodeByWGString, m.Node.Key().WireGuardGoString(), nid)
 			continue
 		case netmap.NodeMutationRemove:
 			nid := m.NodeIDBeingMutated()
@@ -694,6 +727,7 @@ func (nb *nodeBackend) UpdateNetmapDelta(muts []netmap.NodeMutation) (handled bo
 					}
 				}
 				delete(nb.nodeByKey, old.Key())
+				delete(nb.nodeByWGString, old.Key().WireGuardGoString())
 				delete(nb.peers, nid)
 			}
 			continue

ipn/ipnlocal/state_test.go

@@ -1986,6 +1986,7 @@ func (e *mockEngine) SetPeerByIPPacketFunc(func(netip.Addr) (_ key.NodePublic, o
 func (e *mockEngine) SetPeerSessionStateFunc(func(key.NodePublic, wgengine.PeerWireGuardState)) {
 }
 func (e *mockEngine) SetNetLogNodeSource(netlog.NodeSource) {}
+func (e *mockEngine) SetWGPeerLookup(func(wgString string) (tsString string, ok bool)) {}
 
 func (e *mockEngine) Close() {
 	e.mu.Lock()

wgengine/magicsock/magicsock_test.go

@@ -216,7 +216,7 @@ func newMagicStackWithKey(t testing.TB, logf logger.Logf, ln nettype.PacketListe
 	tsTun.SetFilter(filter.NewAllowAllForTest(logf))
 	tsTun.Start()
 
-	wgLogger := wglog.NewLogger(logf)
+	wgLogger := wglog.NewLogger(logf, nil)
 	dev := wgcfg.NewDevice(tsTun, conn.Bind(), wgLogger.DeviceLogger)
 	dev.Up()
 
@@ -244,7 +244,6 @@ func newMagicStackWithKey(t testing.TB, logf logger.Logf, ln nettype.PacketListe
 
 func (s *magicStack) Reconfig(cfg *wgcfg.Config) error {
 	s.tsTun.SetWGConfig(cfg)
-	s.wgLogger.SetPeers(cfg.Peers)
 
 	// In production, LocalBackend installs a PeerByIPPacketFunc via
 	// Engine.SetPeerByIPPacketFunc. Tests that bypass LocalBackend need
@@ -687,7 +686,7 @@ func TestDeviceStartStop(t *testing.T) {
 	defer conn.Close()
 
 	tun := tuntest.NewChannelTUN()
-	wgLogger := wglog.NewLogger(t.Logf)
+	wgLogger := wglog.NewLogger(t.Logf, nil)
 	dev := wgcfg.NewDevice(tun.TUN(), conn.Bind(), wgLogger.DeviceLogger)
 	dev.Up()
 	dev.Close()

wgengine/userspace.go

@@ -162,6 +162,13 @@ type userspaceEngine struct {
 	// logger from inside Reconfig. It may be nil if no source was installed.
 	netLogSource syncs.AtomicValue[netlog.NodeSource]
 
+	// wgPeerLookup is the lookup function installed via
+	// [Engine.SetWGPeerLookup]; it is consulted by [wgLogger] to translate
+	// wireguard-go peer references in log lines. It may be nil if no
+	// lookup was installed, in which case peer references are not
+	// rewritten.
+	wgPeerLookup syncs.AtomicValue[func(wgString string) (tsString string, ok bool)]
+
 	// tsmpLearnedDisco tracks per node key if a peer disco key was learned via TSMP.
 	// wgLock must be held when using this map.
 	tsmpLearnedDisco map[key.NodePublic]key.DiscoPublic
@@ -463,7 +470,13 @@ func NewUserspaceEngine(logf logger.Logf, conf Config) (_ Engine, reterr error)
 		e.tundev.PostFilterPacketOutboundToWireGuard = e.trackOpenPostFilterOut
 	}
 
-	e.wgLogger = wglog.NewLogger(logf)
+	e.wgLogger = wglog.NewLogger(logf, func(wgString string) (tsString string, ok bool) {
+		fn := e.wgPeerLookup.Load()
+		if fn == nil {
+			return "", false
+		}
+		return fn(wgString)
+	})
 	e.tundev.OnTSMPPongReceived = func(pong packet.TSMPPongReply) {
 		e.mu.Lock()
 		defer e.mu.Unlock()
@@ -685,7 +698,10 @@ func (e *userspaceEngine) maybeReconfigWireguardLocked() error {
 	}
 
 	full := e.lastCfgFull
-	e.wgLogger.SetPeers(full.Peers)
+	// The wireguard-go peer set may have changed; drop the cached
+	// peer-string rewrites so the next log line re-resolves them
+	// against the current lookup.
+	e.wgLogger.Invalidate()
 
 	// Rebuild the prefix-match peer routing table from the current
 	// (wireguard-filtered) peer list and publish it atomically.
@@ -746,6 +762,12 @@ func (e *userspaceEngine) SetNetLogNodeSource(src netlog.NodeSource) {
 	e.netLogSource.Store(src)
 }
 
+// SetWGPeerLookup installs the lookup function used by the engine's
+// wireguard-go log wrapper to rewrite peer references in log lines.
+func (e *userspaceEngine) SetWGPeerLookup(fn func(wgString string) (tsString string, ok bool)) {
+	e.wgPeerLookup.Store(fn)
+}
+
 func peerWireGuardStateFromDevice(state device.PeerSessionState) PeerWireGuardState {
 	switch state {
 	case device.PeerSessionNone:

wgengine/wgengine.go

@@ -172,6 +172,16 @@ type Engine interface {
 	// before any [Engine.Reconfig] call that starts up the network logger.
 	SetNetLogNodeSource(netlog.NodeSource)
 
+	// SetWGPeerLookup installs the function used by the engine's
+	// wireguard-go log wrapper to rewrite peer references in log lines
+	// (mapping wireguard-go's "peer(XXXX…YYYY)" form to the
+	// Tailscale-conventional short string form).
+	//
+	// It is expected to be called once during LocalBackend construction.
+	// The function is called concurrently and must be safe to call with
+	// no Engine locks held.
+	SetWGPeerLookup(func(wgString string) (tsString string, ok bool))
+
 	// SetPeerSessionStateFunc installs a callback used to observe WireGuard
 	// peer session state transitions.
 	//

wgengine/wglog/wglog.go

@@ -11,33 +11,42 @@
 
 	"github.com/tailscale/wireguard-go/device"
 	"tailscale.com/envknob"
-	"tailscale.com/syncs"
-	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
-	"tailscale.com/wgengine/wgcfg"
+	"tailscale.com/util/mak"
 )
 
 // A Logger is a wireguard-go log wrapper that cleans up and rewrites log lines.
-// It can be modified at run time to adjust to new wireguard-go configurations.
+// It rewrites wireguard-go peer references like "peer(XXXX…YYYY)" into the
+// Tailscale-conventional short string form like "[XXXXX]".
 type Logger struct {
 	DeviceLogger *device.Logger
-	replace      syncs.AtomicValue[map[string]string]
-	mu           sync.Mutex                   // protects strs
-	strs         map[key.NodePublic]*strCache // cached strs used to populate replace
-}
 
-// strCache holds a wireguard-go and a Tailscale style peer string.
-type strCache struct {
-	wg, ts string
-	used   bool // track whether this strCache was used in a particular round
+	// lookup, when non-nil, returns the Tailscale-conventional short
+	// string for a wireguard-go-formatted peer string (e.g. it maps
+	// "peer(IMTB…r7lM)" to "[IMTBr]"), or "", false if no peer matches.
+	// It is set once at construction and not changed afterwards.
+	lookup func(wgString string) (tsString string, ok bool)
+
+	mu sync.Mutex // protects cache
+
+	// cache memoizes lookup results. The key is the wireguard-go
+	// peer-string form ("peer(XXXX…YYYY)") and the value is the
+	// Tailscale-conventional short-string form ("[XXXXX]").
+	// It is cleared in bulk by Invalidate when the underlying peer set
+	// may have changed.
+	cache map[string]string
 }
 
 // NewLogger creates a new logger for use with wireguard-go.
 // This logger silences repetitive/unhelpful noisy log lines
 // and rewrites peer keys from wireguard-go into Tailscale format.
-func NewLogger(logf logger.Logf) *Logger {
+//
+// lookup, if non-nil, is consulted on cache misses to translate
+// wireguard-go peer references in log lines. If lookup is nil,
+// peer references are passed through unchanged.
+func NewLogger(logf logger.Logf, lookup func(wgString string) (tsString string, ok bool)) *Logger {
 	const prefix = "wg: "
-	ret := new(Logger)
+	ret := &Logger{lookup: lookup}
 	wrapper := func(format string, args ...any) {
 		if strings.Contains(format, "Routine:") && !strings.Contains(format, "receive incoming") {
 			// wireguard-go logs as it starts and stops routines.
@@ -60,32 +69,37 @@ func NewLogger(logf logger.Logf) *Logger {
 			// useful.
 			return
 		}
-		replace := ret.replace.Load()
-		if replace == nil {
-			// No replacements specified; log as originally planned.
+		if ret.lookup == nil {
+			// No lookup; log as originally planned.
 			logf(format, args...)
 			return
 		}
-		// Duplicate the args slice so that we can modify it.
-		// This is not always required, but the code required to avoid it is not worth the complexity.
-		newargs := make([]any, len(args))
-		copy(newargs, args)
-		for i, arg := range newargs {
-			// We want to replace *device.Peer args with the Tailscale-formatted version of themselves.
-			// Using *device.Peer directly makes this hard to test, so we string any fmt.Stringers,
-			// and if the string ends up looking exactly like a known Peer, we replace it.
-			// This is slightly imprecise, in that we don't check the formatting verb. Oh well.
+		// Replace any *device.Peer-shaped fmt.Stringer args with the
+		// Tailscale-formatted version of themselves. Using *device.Peer
+		// directly makes this hard to test, so we string any fmt.Stringers,
+		// and if the string ends up matching a known peer, we substitute.
+		// This is slightly imprecise, in that we don't check the formatting
+		// verb. Oh well.
+		var newargs []any
+		for i, arg := range args {
 			s, ok := arg.(fmt.Stringer)
 			if !ok {
 				continue
 			}
-			wgStr := s.String()
-			tsStr, ok := replace[wgStr]
+			tsStr, ok := ret.peerStringFor(s.String())
 			if !ok {
 				continue
 			}
+			if newargs == nil {
+				newargs = make([]any, len(args))
+				copy(newargs, args)
+			}
 			newargs[i] = tsStr
 		}
+		if newargs == nil {
+			logf(format, args...)
+			return
+		}
 		logf(format, newargs...)
 	}
 	if envknob.Bool("TS_DEBUG_RAW_WGLOG") {
@@ -95,36 +109,51 @@ func NewLogger(logf logger.Logf) *Logger {
 		Verbosef: logger.WithPrefix(wrapper, prefix+"[v2] "),
 		Errorf:   logger.WithPrefix(wrapper, prefix),
 	}
-	ret.strs = make(map[key.NodePublic]*strCache)
 	return ret
 }
 
-// SetPeers adjusts x to rewrite the peer public keys found in peers.
-// SetPeers is safe for concurrent use.
-func (x *Logger) SetPeers(peers []wgcfg.Peer) {
+// peerStringFor returns the Tailscale-conventional short string for
+// wgString, if wgString is a wireguard-go-formatted peer string for a
+// peer known to x.lookup. Results are memoized in x.cache for
+// subsequent log lines.
+func (x *Logger) peerStringFor(wgString string) (tsString string, ok bool) {
+	// Fast path: only strings shaped like wireguard-go's
+	// "peer(XXXX…YYYY)" output can be peer references; skip the cache
+	// and lookup call for anything else (e.g. arbitrary stringers used
+	// in unrelated log args).
+	if !strings.HasPrefix(wgString, "peer(") || !strings.HasSuffix(wgString, ")") {
+		return "", false
+	}
 	x.mu.Lock()
-	defer x.mu.Unlock()
-	// Construct a new peer public key log rewriter.
-	replace := make(map[string]string)
-	for _, peer := range peers {
-		c, ok := x.strs[peer.PublicKey] // look up cached strs
-		if !ok {
-			wg := peer.PublicKey.WireGuardGoString()
-			ts := peer.PublicKey.ShortString()
-			c = &strCache{wg: wg, ts: ts}
-			x.strs[peer.PublicKey] = c
-		}
-		c.used = true
-		replace[c.wg] = c.ts
+	if v, ok := x.cache[wgString]; ok {
+		x.mu.Unlock()
+		return v, true
 	}
-	// Remove any unused cached strs.
-	for k, c := range x.strs {
-		if !c.used {
-			delete(x.strs, k)
-			continue
-		}
-		// Mark c as unused for next round.
-		c.used = false
+	x.mu.Unlock()
+
+	// Call lookup without x.mu held so we don't have to reason about
+	// deadlocks between x.mu and any locks lookup acquires.
+	// In the worst case, two goroutines concurrently miss the cache
+	// for the same wgString and both call lookup; the redundant call
+	// is wasted but harmless.
+	tsString, ok = x.lookup(wgString)
+	if !ok {
+		return "", false
 	}
-	x.replace.Store(replace)
+	x.mu.Lock()
+	mak.Set(&x.cache, wgString, tsString)
+	x.mu.Unlock()
+	return tsString, true
+}
+
+// Invalidate clears the peer-string rewrite cache.
+//
+// Callers should invoke Invalidate when the set of wireguard-go peers
+// has changed so that the next log line re-resolves peer references
+// against the current lookup. It is safe to call concurrently with
+// logging.
+func (x *Logger) Invalidate() {
+	x.mu.Lock()
+	clear(x.cache)
+	x.mu.Unlock()
 }

wgengine/wglog/wglog_test.go

@@ -7,10 +7,11 @@
 	"fmt"
 	"testing"
 
+	extwgconn "github.com/tailscale/wireguard-go/conn"
+	extwgdevice "github.com/tailscale/wireguard-go/device"
+	"github.com/tailscale/wireguard-go/tun/tuntest"
 	"go4.org/mem"
 	"tailscale.com/types/key"
-	"tailscale.com/types/logger"
-	"tailscale.com/wgengine/wgcfg"
 	"tailscale.com/wgengine/wglog"
 )
 
@@ -40,12 +41,19 @@ type log struct {
 		}
 	}
 
-	x := wglog.NewLogger(logf)
-	key, err := key.ParseNodePublicUntyped(mem.S("20c4c1ae54e1fd37cab6e9a532ca20646aff496796cc41d4519560e5e82bee53"))
+	k, err := key.ParseNodePublicUntyped(mem.S("20c4c1ae54e1fd37cab6e9a532ca20646aff496796cc41d4519560e5e82bee53"))
 	if err != nil {
 		t.Fatal(err)
 	}
-	x.SetPeers([]wgcfg.Peer{{PublicKey: key}})
+	wantWG := k.WireGuardGoString()
+	wantTS := k.ShortString()
+	lookup := func(s string) (string, bool) {
+		if s == wantWG {
+			return wantTS, true
+		}
+		return "", false
+	}
+	x := wglog.NewLogger(logf, lookup)
 
 	for _, tt := range tests {
 		if tt.omit {
@@ -69,7 +77,7 @@ func TestSuppressLogs(t *testing.T) {
 	logf := func(format string, args ...any) {
 		logs = append(logs, fmt.Sprintf(format, args...))
 	}
-	x := wglog.NewLogger(logf)
+	x := wglog.NewLogger(logf, nil)
 	x.DeviceLogger.Verbosef("pass")
 	x.DeviceLogger.Verbosef("UAPI: Adding allowedip")
 
@@ -81,6 +89,35 @@ func TestSuppressLogs(t *testing.T) {
 	}
 }
 
+// TestWireGuardGoStringMatchesWireGuardGo guards against a wireguard-go bump
+// silently changing the wireguard-go peer-string format from under us. The
+// LocalBackend's nodeByWGString index is built using
+// [key.NodePublic.WireGuardGoString]; if wireguard-go's *device.Peer.String
+// were to drift, the index would quietly stop matching and wglog would no
+// longer translate peer references in log lines.
+func TestWireGuardGoStringMatchesWireGuardGo(t *testing.T) {
+	var raw [32]byte
+	for i := range raw {
+		raw[i] = byte(i + 1)
+	}
+	nodeKey := key.NodePublicFromRaw32(mem.B(raw[:]))
+
+	dev := extwgdevice.NewDevice(
+		tuntest.NewChannelTUN().TUN(),
+		extwgconn.NewDefaultBind(),
+		extwgdevice.NewLogger(extwgdevice.LogLevelError, ""),
+	)
+	t.Cleanup(dev.Close)
+	peer, err := dev.NewPeer(extwgdevice.NoisePublicKey(raw))
+	if err != nil {
+		t.Fatalf("NewPeer: %v", err)
+	}
+
+	if got, want := nodeKey.WireGuardGoString(), peer.String(); got != want {
+		t.Errorf("NodePublic.WireGuardGoString() = %q, want wireguard-go *device.Peer.String() = %q", got, want)
+	}
+}
+
 func stringer(s string) stringerString {
 	return stringerString(s)
 }
@@ -88,30 +125,3 @@ func stringer(s string) stringerString {
 type stringerString string
 
 func (s stringerString) String() string { return string(s) }
-
-func BenchmarkSetPeers(b *testing.B) {
-	b.ReportAllocs()
-	x := wglog.NewLogger(logger.Discard)
-	peers := [][]wgcfg.Peer{genPeers(0), genPeers(15), genPeers(16), genPeers(15)}
-	for range b.N {
-		for _, p := range peers {
-			x.SetPeers(p)
-		}
-	}
-}
-
-func genPeers(n int) []wgcfg.Peer {
-	if n > 32 {
-		panic("too many peers")
-	}
-	if n == 0 {
-		return nil
-	}
-	peers := make([]wgcfg.Peer, n)
-	for i := range peers {
-		var k [32]byte
-		k[n] = byte(n)
-		peers[i].PublicKey = key.NodePublicFromRaw32(mem.B(k[:]))
-	}
-	return peers
-}