fix: block incoming endpoints from call-me-maybe

wgengine/magicsock/endpoint.go

@@ -207,6 +207,26 @@ func (de *endpoint) deleteEndpointLocked(why string, ep netip.AddrPort) {
 	}
 }
 
+func (de *endpoint) clearEndpoints(why string) {
+	de.mu.Lock()
+	defer de.mu.Unlock()
+	de.debugUpdates.Add(EndpointChange{
+		When: time.Now(),
+		What: "clearEndpoints-" + why,
+	})
+	de.endpointState = map[netip.AddrPort]*endpointState{}
+	if de.bestAddr.AddrPort.IsValid() {
+		de.debugUpdates.Add(EndpointChange{
+			When: time.Now(),
+			What: "clearEndpoints-bestAddr-" + why,
+			From: de.bestAddr,
+		})
+		de.bestAddr = addrLatency{}
+	}
+	de.c.logf("magicsock: disco: node %s %s now using DERP only (all endpoints deleted)",
+		de.publicKey.ShortString(), de.discoShort())
+}
+
 // initFakeUDPAddr populates fakeWGAddr with a globally unique fake UDPAddr.
 // The current implementation just uses the pointer value of de jammed into an IPv6
 // address, but it could also be, say, a counter.

wgengine/magicsock/magicsock.go

@@ -217,6 +217,8 @@ type Conn struct {
 	// blockEndpoints is whether to avoid capturing, storing and sending
 	// endpoints gathered from local interfaces or STUN. Only DERP endpoints
 	// will be sent.
+	// This will also block incoming endpoints received via call-me-maybe disco
+	// packets.
 	blockEndpoints bool
 	// endpointsUpdateActive indicates that updateEndpoints is
 	// currently running. It's used to deduplicate concurrent endpoint
@@ -843,10 +845,10 @@ func (c *Conn) DiscoPublicKey() key.DiscoPublic {
 	return c.discoPublic
 }
 
-// SetBlockEndpoints sets the blockEndpoints field. If changed, endpoints will
-// be updated to apply the new settings. Existing connections may continue to
-// use the old setting until they are reestablished. Disabling endpoints does
-// not affect the UDP socket or portmapper.
+// SetBlockEndpoints sets the blockEndpoints field. If enabled, all peer
+// endpoints will be cleared from the peer map and every connection will
+// immediately switch to DERP. Disabling endpoints does not affect the UDP
+// socket or portmapper.
 func (c *Conn) SetBlockEndpoints(block bool) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -856,6 +858,7 @@ func (c *Conn) SetBlockEndpoints(block bool) {
 		return
 	}
 
+	// Re-gather local endpoints.
 	const why = "SetBlockEndpoints"
 	if c.endpointsUpdateActive {
 		if c.wantEndpointsUpdate != why {
@@ -866,6 +869,13 @@ func (c *Conn) SetBlockEndpoints(block bool) {
 		c.endpointsUpdateActive = true
 		go c.updateEndpoints(why)
 	}
+
+	// Clear any endpoints from the peerMap if block is true.
+	if block {
+		c.peerMap.forEachEndpoint(func(ep *endpoint) {
+			ep.clearEndpoints("SetBlockEndpoints")
+		})
+	}
 }
 
 // determineEndpoints returns the machine's endpoint addresses. It
@@ -1523,6 +1533,14 @@ func (c *Conn) handleDiscoMessage(msg []byte, src netip.AddrPort, derpNodeSrc ke
 			c.logf("[unexpected] CallMeMaybe from peer via DERP whose netmap discokey != disco source")
 			return
 		}
+		if c.blockEndpoints {
+			c.dlogf("[v1] magicsock: disco: %v<-%v (%v, %v)  got call-me-maybe with %d endpoints, but endpoints blocked",
+				c.discoShort, epDisco.short,
+				ep.publicKey.ShortString(), derpStr(src.String()),
+				len(dm.MyNumber),
+			)
+			dm.MyNumber = []netip.AddrPort{}
+		}
 		c.dlogf("[v1] magicsock: disco: %v<-%v (%v, %v)  got call-me-maybe, %d endpoints",
 			c.discoShort, epDisco.short,
 			ep.publicKey.ShortString(), derpStr(src.String()),
@@ -1921,6 +1939,13 @@ func (c *Conn) SetNetworkMap(nm *netmap.NetworkMap) {
 	// we'll fall through to the next pass, which allocates but can
 	// handle full set updates.
 	for _, n := range nm.Peers {
+		if c.blockEndpoints {
+			c.dlogf("[v1] magicsock: SetNetworkingMap: %v (%v, %v) received %d endpoints, but endpoints blocked",
+				n.DiscoKey.ShortString(), n.Key.ShortString(),
+				len(n.Endpoints),
+			)
+			n.Endpoints = []string{}
+		}
 		if ep, ok := c.peerMap.endpointForNodeKey(n.Key); ok {
 			if n.DiscoKey.IsZero() && !n.IsWireGuardOnly {
 				// Discokey transitioned from non-zero to zero? This should not

wgengine/magicsock/magicsock_test.go

@@ -3022,7 +3022,6 @@ func TestBlockEndpoints(t *testing.T) {
 			t.Fatal("DERP IP in endpoints list?", ep.Addr)
 		}
 		haveEndpoint = true
-		break
 	}
 	ms.conn.mu.Unlock()
 	if !haveEndpoint {
@@ -3060,30 +3059,27 @@ func TestBlockEndpointsDERPOK(t *testing.T) {
 	logf, closeLogf := logger.LogfCloser(t.Logf)
 	defer closeLogf()
 
-	derpMap, cleanup := runDERPAndStun(t, t.Logf, localhostListener{}, netaddr.IPv4(127, 0, 0, 1))
-	defer cleanup()
+	derpMap, cleanupDerp := runDERPAndStun(t, t.Logf, localhostListener{}, netaddr.IPv4(127, 0, 0, 1))
+	defer cleanupDerp()
 
 	ms1 := newMagicStack(t, logger.WithPrefix(logf, "conn1: "), d.m1, derpMap)
 	defer ms1.Close()
+	ms1.conn.SetDebugLoggingEnabled(true)
 	ms2 := newMagicStack(t, logger.WithPrefix(logf, "conn2: "), d.m2, derpMap)
 	defer ms2.Close()
+	ms2.conn.SetDebugLoggingEnabled(true)
 
-	cleanup = meshStacks(logf, nil, ms1, ms2)
-	defer cleanup()
+	cleanupMesh := meshStacks(logf, nil, ms1, ms2)
+	defer cleanupMesh()
 
 	m1IP := ms1.IP()
 	m2IP := ms2.IP()
 	logf("IPs: %s %s", m1IP, m2IP)
 
-	// SetBlockEndpoints is called later since it's incompatible with the test
-	// meshStacks implementations.
-	ms1.conn.SetBlockEndpoints(true)
-	ms2.conn.SetBlockEndpoints(true)
-	waitForNoEndpoints(t, ms1.conn)
-	waitForNoEndpoints(t, ms2.conn)
-
-	cleanup = newPinger(t, logf, ms1, ms2)
-	defer cleanup()
+	cleanupPinger1 := newPinger(t, logf, ms1, ms2)
+	defer cleanupPinger1()
+	cleanupPinger2 := newPinger(t, logf, ms2, ms1)
+	defer cleanupPinger2()
 
 	// Wait for both peers to know about each other.
 	for {
@@ -3098,14 +3094,39 @@ func TestBlockEndpointsDERPOK(t *testing.T) {
 		break
 	}
 
-	cleanup = newPinger(t, t.Logf, ms1, ms2)
-	defer cleanup()
+	waitForEndpoints(t, ms1.conn)
+	waitForEndpoints(t, ms2.conn)
 
-	if len(ms1.conn.activeDerp) == 0 {
-		t.Errorf("unexpected DERP empty got: %v want: >0", len(ms1.conn.activeDerp))
+	// SetBlockEndpoints is called later since it's incompatible with the test
+	// meshStacks implementations.
+	// We only set it on ms1, since ms2's endpoints should be ignored by ms1.
+	ms1.conn.SetBlockEndpoints(true)
+
+	// All endpoints should've been immediately removed from ms1.
+	ep2, ok := ms1.conn.peerMap.endpointForNodeKey(ms2.Public())
+	if !ok {
+		t.Fatalf("endpoint not found for ms2")
 	}
-	if len(ms2.conn.activeDerp) == 0 {
-		t.Errorf("unexpected DERP empty got: %v want: >0", len(ms2.conn.activeDerp))
+	ep2.mu.Lock()
+	if len(ep2.endpointState) != 0 {
+		ep2.mu.Unlock()
+		t.Fatalf("endpoints not removed from ms1")
+	}
+	ep2.mu.Unlock()
+
+	// Wait for endpoints to finish updating.
+	waitForNoEndpoints(t, ms1.conn)
+
+	// Give time for another call-me-maybe packet to arrive. I couldn't think of
+	// a better way than sleeping without making a bunch of changes.
+	t.Logf("sleeping for call-me-maybe packet to be received and ignored")
+	time.Sleep(time.Second)
+	t.Logf("done sleeping")
+
+	ep2.mu.Lock()
+	defer ep2.mu.Unlock()
+	for i := range ep2.endpointState {
+		t.Fatalf("endpoint %q not missing", i.String())
 	}
 }
 
@@ -3129,3 +3150,20 @@ func waitForNoEndpoints(t *testing.T, ms *Conn) {
 	}
 	t.Log("endpoints are blocked")
 }
+
+func waitForEndpoints(t *testing.T, ms *Conn) {
+	t.Helper()
+	for i := 0; i < 50; i++ {
+		time.Sleep(100 * time.Millisecond)
+		ms.mu.Lock()
+		for _, ep := range ms.lastEndpoints {
+			if ep.Addr.Addr() != tailcfg.DerpMagicIPAddr {
+				t.Log("endpoint found")
+				ms.mu.Unlock()
+				return
+			}
+		}
+		ms.mu.Unlock()
+	}
+	t.Fatal("endpoint was not found after 50 attempts")
+}