From c8c493f3d9bf925e9459236bf1ecea823be6f825 Mon Sep 17 00:00:00 2001
From: David Anderson <danderson@tailscale.com>
Date: Fri, 18 Dec 2020 00:49:50 -0800
Subject: [PATCH] wgengine/magicsock: make ReceiveIPv4 a little easier to
 follow.

The previous code used a lot of whole-function variables and shared
behavior that only triggered based on prior action from a single codepath.
Instead of that, move the small amounts of "shared" code into each switch
case.

Signed-off-by: David Anderson <danderson@tailscale.com>
---
 wgengine/magicsock/magicsock.go | 45 ++++++++++++++++++---------------
 1 file changed, 25 insertions(+), 20 deletions(-)

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 9529a3701..185c4ee8a 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -1489,10 +1489,7 @@ func (c *Conn) ReceiveIPv4(b []byte) (n int, ep conn.Endpoint, addr *net.UDPAddr
 	// to udpRecvCh. The code below must not access b until it's
 	// completed a successful receive on udpRecvCh.
 
-	var addrSet *addrSet
-	var discoEp *discoEndpoint
 	var ipp netaddr.IPPort
-	var didNoteRecvActivity bool
 
 	select {
 	case dm := <-c.derpRecvCh:
@@ -1531,6 +1528,11 @@ func (c *Conn) ReceiveIPv4(b []byte) (n int, ep conn.Endpoint, addr *net.UDPAddr
 			goto Top
 		}
 
+		var (
+			didNoteRecvActivity bool
+			discoEp             *discoEndpoint
+			asEp                *addrSet
+		)
 		c.mu.Lock()
 		if dk, ok := c.discoOfNode[tailcfg.NodeKey(dm.src)]; ok {
 			discoEp = c.endpointOfDisco[dk]
@@ -1550,24 +1552,39 @@ func (c *Conn) ReceiveIPv4(b []byte) (n int, ep conn.Endpoint, addr *net.UDPAddr
 				c.mu.Lock()
 
 				discoEp = c.endpointOfDisco[dk]
-				c.logf("magicsock: DERP packet received from idle peer %v; created=%v", dm.src.ShortString(), discoEp != nil)
+				c.logf("magicsock: DERP packet received from idle peer %v; created=%v", dm.src.ShortString(), ep != nil)
 			}
 		}
-		if discoEp == nil {
-			addrSet = c.addrsByKey[dm.src]
-		}
+		asEp = c.addrsByKey[dm.src]
 		c.mu.Unlock()
 
-		if addrSet == nil && discoEp == nil {
+		if discoEp != nil {
+			ep = discoEp
+		} else if asEp != nil {
+			ep = asEp
+		} else {
 			key := wgcfg.Key(dm.src)
 			c.logf("magicsock: DERP packet from unknown key: %s", key.ShortString())
+			// TODO(danderson): after we fail to find a DERP endpoint, we
+			// seem to be falling through to passing the packet to
+			// wireguard with a garbage singleEndpoint. This feels wrong,
+			// should we goto Top above?
+			ep = c.findEndpoint(ipp, addr)
 		}
 
+		if !didNoteRecvActivity {
+			c.noteRecvActivityFromEndpoint(ep)
+		}
+		return n, ep, wgRecvAddr(ep, ipp, addr), nil
+
 	case um := <-c.udpRecvCh:
 		if um.err != nil {
 			return 0, nil, nil, err
 		}
 		n, addr, ipp = um.n, um.addr, um.ipp
+		ep = c.findEndpoint(ipp, addr)
+		c.noteRecvActivityFromEndpoint(ep)
+		return n, ep, wgRecvAddr(ep, ipp, addr), nil
 
 	case <-c.donec():
 		// Socket has been shut down. All the producers of packets
@@ -1582,18 +1599,6 @@ func (c *Conn) ReceiveIPv4(b []byte) (n int, ep conn.Endpoint, addr *net.UDPAddr
 		// with an error so it can clean up.
 		return 0, nil, nil, errors.New("socket closed")
 	}
-
-	if addrSet != nil {
-		ep = addrSet
-	} else if discoEp != nil {
-		ep = discoEp
-	} else {
-		ep = c.findEndpoint(ipp, addr)
-	}
-	if !didNoteRecvActivity {
-		c.noteRecvActivityFromEndpoint(ep)
-	}
-	return n, ep, wgRecvAddr(ep, ipp, addr), nil
 }
 
 func (c *Conn) ReceiveIPv6(b []byte) (int, conn.Endpoint, *net.UDPAddr, error) {