fix code review issues

Signed-off-by: Naman Sood <mail@nsood.in>

wgengine/netstack/netstack.go

@@ -89,17 +89,17 @@ func Impl(logf logger.Logf, tundev *tstun.TUN, e wgengine.Engine, mc *magicsock.
 		for ip := range ipsToBeRemoved {
 			err := ipstack.RemoveAddress(nicID, ip)
 			if err != nil {
-				logf("netstack: could not deregister IP %v: %v", ip.String(), err)
+				logf("netstack: could not deregister IP %s: %v", ip, err)
 			} else {
-				logf("netstack: deregistered IP %v", ip.String())
+				logf("netstack: deregistered IP %s", ip)
 			}
 		}
 		for ip := range ipsToBeAdded {
 			err := ipstack.AddAddress(nicID, ipv4.ProtocolNumber, ip)
 			if err != nil {
-				logf("netstack: could not register IP %v: %v", ip.String(), err)
+				logf("netstack: could not register IP %s: %v", ip, err)
 			} else {
-				logf("netstack: registered IP %v", ip.String())
+				logf("netstack: registered IP %s", ip)
 			}
 		}
 	})

wgengine/userspace.go

@@ -109,7 +109,6 @@ type userspaceEngine struct {
 	trimmedDisco        map[tailcfg.DiscoKey]bool // set of disco keys of peers currently excluded from wireguard config
 	sentActivityAt      map[netaddr.IP]*int64     // value is atomic int64 of unixtime
 	destIPActivityFuncs map[netaddr.IP]func()
-	networkMapCallbacks map[*networkMapCallbackHandle]NetworkMapCallback
 
 	mu                  sync.Mutex // guards following; see lock order comment below
 	closing             bool       // Close was called (even if we're still closing)
@@ -119,6 +118,7 @@ type userspaceEngine struct {
 	endpoints           []string
 	pingers             map[wgkey.Key]*pinger // legacy pingers for pre-discovery peers
 	linkState           *interfaces.State
+	networkMapCallbacks map[*networkMapCallbackHandle]NetworkMapCallback
 
 	// Lock ordering: magicsock.Conn.mu, wgLock, then mu.
 }
@@ -216,7 +216,6 @@ func newUserspaceEngineAdvanced(conf EngineConfig) (_ Engine, reterr error) {
 		tundev:   tstun.WrapTUN(logf, conf.TUN),
 		resolver: tsdns.NewResolver(rconf),
 		pingers:  make(map[wgkey.Key]*pinger),
-		networkMapCallbacks: make(map[*networkMapCallbackHandle]NetworkMapCallback),
 	}
 	e.localAddrs.Store(map[netaddr.IP]bool{})
 	e.linkState, _ = getLinkState()
@@ -1266,9 +1265,16 @@ func (e *userspaceEngine) SetLinkChangeCallback(cb func(major bool, newState *in
 }
 
 func (e *userspaceEngine) AddNetworkMapCallback(cb NetworkMapCallback) func() {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	if e.networkMapCallbacks == nil {
+		e.networkMapCallbacks = make(map[*networkMapCallbackHandle]NetworkMapCallback)
+	}
 	handle := new(networkMapCallbackHandle)
 	e.networkMapCallbacks[handle] = cb
 	return func() {
+		e.mu.Lock()
+		defer e.mu.Unlock()
 		delete(e.networkMapCallbacks, handle)
 	}
 }
@@ -1291,6 +1297,8 @@ func (e *userspaceEngine) SetDERPMap(dm *tailcfg.DERPMap) {
 
 func (e *userspaceEngine) SetNetworkMap(nm *controlclient.NetworkMap) {
 	e.magicConn.SetNetworkMap(nm)
+	e.mu.Lock()
+	defer e.mu.Unlock()
 	for _, fn := range e.networkMapCallbacks {
 		fn(nm)
 	}

wgengine/wgengine.go

@@ -123,7 +123,7 @@ type Engine interface {
 	// that are called when the network map updates. It returns a
 	// function that when called would remove the function from the
 	// list of callbacks.
-	AddNetworkMapCallback(NetworkMapCallback) func()
+	AddNetworkMapCallback(NetworkMapCallback) (removeCallback func())
 
 	// SetNetInfoCallback sets the function to call when a
 	// new NetInfo summary is available.