wgengine: ensure pingers are gone before returning from Close

We canceled the pingers in Close, but didn't wait around for their
goroutines to be cleaned up. This caused the ipn/e2e_test to catch
pingers in its resource leak check.

This commit introduces an object, but also simplifies the semantics
around the pinger's cancel functions. They no longer need to be called
while holding the mutex.

Signed-off-by: David Crawshaw <crawshaw@tailscale.com>

wgengine/userspace.go

@@ -66,7 +66,7 @@ type userspaceEngine struct {
 	statusCallback StatusCallback
 	peerSequence   []wgcfg.Key
 	endpoints      []string
-	pingers        map[wgcfg.Key]context.CancelFunc // mu must be held to call CancelFunc
+	pingers        map[wgcfg.Key]*pinger
 	linkState      *interfaces.State
 
 	// Lock ordering: wgLock, then mu.
@@ -128,7 +128,7 @@ func newUserspaceEngineAdvanced(logf logger.Logf, tundev *tstun.TUN, routerGen R
 		reqCh:   make(chan struct{}, 1),
 		waitCh:  make(chan struct{}),
 		tundev:  tundev,
-		pingers: make(map[wgcfg.Key]context.CancelFunc),
+		pingers: make(map[wgcfg.Key]*pinger),
 	}
 	e.linkState, _ = getLinkState()
 
@@ -259,28 +259,29 @@ func newUserspaceEngineAdvanced(logf logger.Logf, tundev *tstun.TUN, routerGen R
 //
 // These generated packets are used to ensure we trigger the spray logic in
 // the magicsock package for NAT traversal.
-func (e *userspaceEngine) pinger(peerKey wgcfg.Key, ips []wgcfg.IP) {
-	e.logf("generating initial ping traffic to %s (%v)", peerKey.ShortString(), ips)
-	var srcIP packet.IP
+type pinger struct {
+	e      *userspaceEngine
+	done   chan struct{} // closed after shutdown (not the ctx.Done() chan)
+	cancel context.CancelFunc
+}
 
-	e.wgLock.Lock()
-	if len(e.lastCfg.Addresses) > 0 {
-		srcIP = packet.NewIP(e.lastCfg.Addresses[0].IP.IP())
-	}
-	e.wgLock.Unlock()
+// close cleans up pinger and removes it from the userspaceEngine.pingers map.
+// It cannot be called while p.e.mu is held.
+func (p *pinger) close() {
+	p.cancel()
+	<-p.done
+}
 
-	if srcIP == 0 {
-		e.logf("generating initial ping traffic: no source IP")
-		return
-	}
+func (p *pinger) run(ctx context.Context, peerKey wgcfg.Key, ips []wgcfg.IP, srcIP packet.IP) {
+	defer func() {
+		p.e.mu.Lock()
+		if p.e.pingers[peerKey] == p {
+			delete(p.e.pingers, peerKey)
+		}
+		p.e.mu.Unlock()
 
-	e.mu.Lock()
-	if cancel := e.pingers[peerKey]; cancel != nil {
-		cancel()
-	}
-	ctx, cancel := context.WithCancel(context.Background())
-	e.pingers[peerKey] = cancel
-	e.mu.Unlock()
+		close(p.done)
+	}()
 
 	// sendFreq is slightly longer than sprayFreq in magicsock to ensure
 	// that if these ping packets are the only source of early packets
@@ -296,19 +297,6 @@ func (e *userspaceEngine) pinger(peerKey wgcfg.Key, ips []wgcfg.IP) {
 
 	payload := []byte("magicsock_spray") // no meaning
 
-	defer func() {
-		e.mu.Lock()
-		defer e.mu.Unlock()
-		select {
-		case <-ctx.Done():
-			return
-		default:
-		}
-		// If the pinger context is not done, then the
-		// CancelFunc is still in the pingers map.
-		delete(e.pingers, peerKey)
-	}()
-
 	ipid := uint16(1)
 	t := time.NewTicker(sendFreq)
 	defer t.Stop()
@@ -323,10 +311,52 @@ func (e *userspaceEngine) pinger(peerKey wgcfg.Key, ips []wgcfg.IP) {
 		}
 		for _, dstIP := range dstIPs {
 			b := packet.GenICMP(srcIP, dstIP, ipid, packet.ICMPEchoRequest, 0, payload)
-			e.tundev.InjectOutbound(b)
+			p.e.tundev.InjectOutbound(b)
 		}
 		ipid++
 	}
+
+}
+
+// pinger sends ping packets for a few seconds.
+//
+// These generated packets are used to ensure we trigger the spray logic in
+// the magicsock package for NAT traversal.
+func (e *userspaceEngine) pinger(peerKey wgcfg.Key, ips []wgcfg.IP) {
+	e.logf("generating initial ping traffic to %s (%v)", peerKey.ShortString(), ips)
+	var srcIP packet.IP
+
+	e.wgLock.Lock()
+	if len(e.lastCfg.Addresses) > 0 {
+		srcIP = packet.NewIP(e.lastCfg.Addresses[0].IP.IP())
+	}
+	e.wgLock.Unlock()
+
+	if srcIP == 0 {
+		e.logf("generating initial ping traffic: no source IP")
+		return
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	p := &pinger{
+		e:      e,
+		done:   make(chan struct{}),
+		cancel: cancel,
+	}
+
+	e.mu.Lock()
+	if e.closing {
+		e.mu.Unlock()
+		return
+	}
+	oldPinger := e.pingers[peerKey]
+	e.pingers[peerKey] = p
+	e.mu.Unlock()
+
+	if oldPinger != nil {
+		oldPinger.close()
+	}
+	p.run(ctx, peerKey, ips, srcIP)
 }
 
 func configSignature(cfg *wgcfg.Config, routerCfg *router.Config) (string, error) {
@@ -565,15 +595,16 @@ func (e *userspaceEngine) RequestStatus() {
 }
 
 func (e *userspaceEngine) Close() {
+	var pingers []*pinger
+
 	e.mu.Lock()
 	if e.closing {
 		e.mu.Unlock()
 		return
 	}
 	e.closing = true
-	for key, cancel := range e.pingers {
-		delete(e.pingers, key)
-		cancel()
+	for _, pinger := range e.pingers {
+		pingers = append(pingers, pinger)
 	}
 	e.mu.Unlock()
 
@@ -583,6 +614,13 @@ func (e *userspaceEngine) Close() {
 	e.linkMon.Close()
 	e.router.Close()
 	e.magicConn.Close()
+
+	// Shut down pingers after tundev is closed (by e.wgdev.Close) so the
+	// synchronous close does not get stuck on InjectOutbound.
+	for _, pinger := range pingers {
+		pinger.close()
+	}
+
 	close(e.waitCh)
 }