From f849018e3c1c0aecd1e786f098a4e32f8d00832a Mon Sep 17 00:00:00 2001
From: Naman Sood <mail@nsood.in>
Date: Thu, 14 Jan 2021 16:03:14 -0500
Subject: [PATCH] fix code review issues

Signed-off-by: Naman Sood <mail@nsood.in>
---
 wgengine/netstack/netstack.go |  8 +++----
 wgengine/userspace.go         | 42 +++++++++++++++++++++--------------
 wgengine/wgengine.go          |  2 +-
 3 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/wgengine/netstack/netstack.go b/wgengine/netstack/netstack.go
index 5c05ea1cb..72651bd41 100644
--- a/wgengine/netstack/netstack.go
+++ b/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)
 			}
 		}
 	})
diff --git a/wgengine/userspace.go b/wgengine/userspace.go
index 810c612c8..681ceace0 100644
--- a/wgengine/userspace.go
+++ b/wgengine/userspace.go
@@ -109,16 +109,16 @@ 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)
-	statusCallback     StatusCallback
-	linkChangeCallback func(major bool, newState *interfaces.State)
-	peerSequence       []wgkey.Key
-	endpoints          []string
-	pingers            map[wgkey.Key]*pinger // legacy pingers for pre-discovery peers
-	linkState          *interfaces.State
+	mu                  sync.Mutex // guards following; see lock order comment below
+	closing             bool       // Close was called (even if we're still closing)
+	statusCallback      StatusCallback
+	linkChangeCallback  func(major bool, newState *interfaces.State)
+	peerSequence        []wgkey.Key
+	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.
 }
@@ -209,14 +209,13 @@ func newUserspaceEngineAdvanced(conf EngineConfig) (_ Engine, reterr error) {
 		Forward: true,
 	}
 	e := &userspaceEngine{
-		timeNow:             time.Now,
-		logf:                logf,
-		reqCh:               make(chan struct{}, 1),
-		waitCh:              make(chan struct{}),
-		tundev:              tstun.WrapTUN(logf, conf.TUN),
-		resolver:            tsdns.NewResolver(rconf),
-		pingers:             make(map[wgkey.Key]*pinger),
-		networkMapCallbacks: make(map[*networkMapCallbackHandle]NetworkMapCallback),
+		timeNow:  time.Now,
+		logf:     logf,
+		reqCh:    make(chan struct{}, 1),
+		waitCh:   make(chan struct{}),
+		tundev:   tstun.WrapTUN(logf, conf.TUN),
+		resolver: tsdns.NewResolver(rconf),
+		pingers:  make(map[wgkey.Key]*pinger),
 	}
 	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)
 	}
diff --git a/wgengine/wgengine.go b/wgengine/wgengine.go
index 56871107d..5bd5e2e80 100644
--- a/wgengine/wgengine.go
+++ b/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.