fix: block incoming endpoints from call-me-maybe

wgengine/magicsock/magicsock.go

@@ -217,6 +217,9 @@ 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. Endpoints manually added (e.g. via coordination) will not be
+	// blocked.
 	blockEndpoints bool
 	// endpointsUpdateActive indicates that updateEndpoints is
 	// currently running. It's used to deduplicate concurrent endpoint
@@ -856,6 +859,7 @@ func (c *Conn) SetBlockEndpoints(block bool) {
 		return
 	}
 
+	// Re-gather local endpoints.
 	const why = "SetBlockEndpoints"
 	if c.endpointsUpdateActive {
 		if c.wantEndpointsUpdate != why {
@@ -866,6 +870,17 @@ 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) {
+			// Simulate receiving a call-me-maybe disco packet with no
+			// endpoints, which will cause all endpoints to be removed.
+			ep.handleCallMeMaybe(&disco.CallMeMaybe{
+				MyNumber: []netip.AddrPort{},
+			})
+		})
+	}
 }
 
 // determineEndpoints returns the machine's endpoint addresses. It
@@ -1523,6 +1538,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()),

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,25 @@ 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()
 	ms2 := newMagicStack(t, logger.WithPrefix(logf, "conn2: "), d.m2, derpMap)
 	defer ms2.Close()
 
-	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 +3092,56 @@ 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))
+	// meshStacks doesn't use call-me-maybe packets to update endpoints, which
+	// makes them not get cleared when we call SetBlockEndpoints.
+	ep2, ok := ms1.conn.peerMap.endpointForNodeKey(ms2.Public())
+	if !ok {
+		t.Fatalf("endpoint not found for ms2")
+	}
+	ep2.mu.Lock()
+	for k := range ep2.endpointState {
+		ep2.deleteEndpointLocked("TestBlockEndpointsDERPOK", k)
+	}
+	ep2.mu.Unlock()
+
+	// Wait for the call-me-maybe packet to arrive.
+	for i := 0; i < 50; i++ {
+		time.Sleep(100 * time.Millisecond)
+		ep2.mu.Lock()
+		if len(ep2.endpointState) != 0 {
+			ep2.mu.Unlock()
+			break
+		}
+		ep2.mu.Unlock()
 	}
-	if len(ms2.conn.activeDerp) == 0 {
-		t.Errorf("unexpected DERP empty got: %v want: >0", len(ms2.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 call-me-maybe endpoints should've been immediately removed from ms1.
+	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.
+	time.Sleep(time.Second)
+
+	ep2.mu.Lock()
+	defer ep2.mu.Unlock()
+	for i := range ep2.endpointState {
+		t.Fatalf("endpoint %q not missing", i.String())
 	}
 }
 
@@ -3129,3 +3165,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")
+}