3 files changed, 126 insertions, 4 deletions
diff --git a/wgengine/magicsock/endpoint.go b/wgengine/magicsock/endpoint.go
index 4780c7f37..491c069a2 100644
--- a/wgengine/magicsock/endpoint.go
+++ b/wgengine/magicsock/endpoint.go
@@ -21,6 +21,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"go4.org/mem"
 	"golang.org/x/net/ipv4"
 	"golang.org/x/net/ipv6"
 	"tailscale.com/disco"
@@ -499,8 +500,9 @@ func (de *endpoint) initFakeUDPAddr() {
 }
 
 // noteRecvActivity records receive activity on de, and invokes
-// Conn.noteRecvActivity no more than once every 10s.
-func (de *endpoint) noteRecvActivity(src epAddr, now mono.Time) {
+// [Conn.noteRecvActivity] no more than once every 10s, returning true if it
+// was called, otherwise false.
+func (de *endpoint) noteRecvActivity(src epAddr, now mono.Time) bool {
 	if de.isWireguardOnly {
 		de.mu.Lock()
 		de.bestAddr.ap = src.ap
@@ -524,10 +526,12 @@ func (de *endpoint) noteRecvActivity(src epAddr, now mono.Time) {
 		de.lastRecvWG.StoreAtomic(now)
 
 		if de.c.noteRecvActivity == nil {
-			return
+			return false
 		}
 		de.c.noteRecvActivity(de.publicKey)
+		return true
 	}
+	return false
 }
 
 func (de *endpoint) discoShort() string {
@@ -2024,3 +2028,62 @@ func (de *endpoint) setDERPHome(regionID uint16) {
 	defer de.mu.Unlock()
 	de.derpAddr = netip.AddrPortFrom(tailcfg.DerpMagicIPAddr, uint16(regionID))
 }
+
+// rxPacketVerifyEndpoint wraps an [*endpoint] and implements
+// [conn.PeerAwareEndpoint] & [conn.InitiationAwareEndpoint]. We return an
+// [rxPacketVerifyEndpoint] in our [conn.ReceiveFunc]s periodically and for
+// all WireGuard initiation messages to resolve any potential [epAddr]
+// collisions.
+//
+// [epAddr] collisions have a higher chance of occurrence for packets received
+// over peer relays versus direct connections, as peer relay connections do not
+// upsert into [peerMap] around disco packet reception, but direct connections
+// do.
+//
+// We believe the packet from src was from the wrapped [*endpoint], but want to
+// verify with a Cryptokey Routing outcome from wireguard-go, and/or
+// just-in-time configure the potentially differing peer before wireguard-go
+// tries to decrypt.
+type rxPacketVerifyEndpoint struct {
+	*endpoint
+	src epAddr
+}
+
+// InitiationMessagePublicKey implements [conn.InitiationAwareEndpoint].
+func (r rxPacketVerifyEndpoint) InitiationMessagePublicKey(peerPublicKey [32]byte) {
+	pubKey := key.NodePublicFromRaw32(mem.B(peerPublicKey[:]))
+	if pubKey.Compare(r.endpoint.publicKey) == 0 {
+		return
+	}
+	r.endpoint.c.mu.Lock()
+	defer r.endpoint.c.mu.Unlock()
+	ep, ok := r.endpoint.c.peerMap.endpointForNodeKey(pubKey)
+	if !ok {
+		return
+	}
+	now := mono.Now()
+	ep.lastRecvUDPAny.StoreAtomic(now)
+	ep.noteRecvActivity(r.src, now)
+	// [ep.noteRecvActivity] may end up JIT configuring the peer, but we don't
+	// update [peerMap] as wireguard-go hasn't decrypted the initiation
+	// message yet. wireguard-go will call us below in
+	// [rxPacketVerifyEndpoint.FromPeer] if it successfully decrypts the
+	// message, at which point it's safe to insert r.src into the [peerMap]
+	// for ep, which will resolve the [epAddr] collision.
+}
+
+// FromPeer implements [conn.PeerAwareEndpoint].
+func (r rxPacketVerifyEndpoint) FromPeer(peerPublicKey [32]byte) {
+	pubKey := key.NodePublicFromRaw32(mem.B(peerPublicKey[:]))
+	if pubKey.Compare(r.endpoint.publicKey) == 0 {
+		return
+	}
+	r.endpoint.c.mu.Lock()
+	defer r.endpoint.c.mu.Unlock()
+	ep, ok := r.endpoint.c.peerMap.endpointForNodeKey(pubKey)
+	if !ok {
+		return
+	}
+	r.endpoint.c.peerMap.setNodeKeyForEpAddr(r.src, pubKey)
+	r.endpoint.c.logf("magicsock: rxPacketVerifyEndpoint.FromPeer(%v) setting epAddr(%v) in peerMap for node(%v)", pubKey.ShortString(), r.src, ep.nodeAddr)
+}
diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 1978867fa..f74802973 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -27,6 +27,7 @@ import (
 	"time"
 
 	"github.com/tailscale/wireguard-go/conn"
+	"github.com/tailscale/wireguard-go/device"
 	"go4.org/mem"
 	"golang.org/x/net/ipv6"
 
@@ -1632,6 +1633,16 @@ func (c *Conn) mkReceiveFunc(ruc *RebindingUDPConn, healthItem *health.ReceiveFu
 	}
 }
 
+// looksLikeInitiationMsg returns true if b looks like a WireGuard initiation
+// message, otherwise it returns false.
+func looksLikeInitiationMsg(b []byte) bool {
+	if len(b) == device.MessageInitiationSize &&
+		binary.BigEndian.Uint32(b) == device.MessageInitiationType {
+		return true
+	}
+	return false
+}
+
 // receiveIP is the shared bits of ReceiveIPv4 and ReceiveIPv6.
 //
 // size is the length of 'b' to report up to wireguard-go (only relevant if
@@ -1717,10 +1728,18 @@ func (c *Conn) receiveIP(b []byte, ipp netip.AddrPort, cache *epAddrEndpointCach
 	}
 	now := mono.Now()
 	ep.lastRecvUDPAny.StoreAtomic(now)
-	ep.noteRecvActivity(src, now)
+	connNoted := ep.noteRecvActivity(src, now)
 	if stats := c.stats.Load(); stats != nil {
 		stats.UpdateRxPhysical(ep.nodeAddr, ipp, 1, len(b))
 	}
+	if src.vni.isSet() && (connNoted || looksLikeInitiationMsg(b)) {
+		// connNoted is periodic, but we also want to verify if the peer is who
+		// we believe for all initiation messages, otherwise we could get
+		// unlucky and fail to JIT configure the "correct" peer.
+		// TODO(jwhited): relax this to include direct connections
+		//  See http://go/corp/29422 & http://go/corp/30042
+		return rxPacketVerifyEndpoint{endpoint: ep, src: src}, size, true
+	}
 	return ep, size, true
 }
 
diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go
index aea2de17d..0515162c7 100644
--- a/wgengine/magicsock/magicsock_test.go
+++ b/wgengine/magicsock/magicsock_test.go
@@ -3611,3 +3611,43 @@ func Test_peerAPIIfCandidateRelayServer(t *testing.T) {
 		})
 	}
 }
+
+func Test_looksLikeInitiationMsg(t *testing.T) {
+	initMsg := make([]byte, device.MessageInitiationSize)
+	binary.BigEndian.PutUint32(initMsg, device.MessageInitiationType)
+	initMsgSizeTransportType := make([]byte, device.MessageInitiationSize)
+	binary.BigEndian.PutUint32(initMsgSizeTransportType, device.MessageTransportType)
+	tests := []struct {
+		name string
+		b    []byte
+		want bool
+	}{
+		{
+			name: "valid initiation",
+			b:    initMsg,
+			want: true,
+		},
+		{
+			name: "invalid message type field",
+			b:    initMsgSizeTransportType,
+			want: false,
+		},
+		{
+			name: "too small",
+			b:    initMsg[:device.MessageInitiationSize-1],
+			want: false,
+		},
+		{
+			name: "too big",
+			b:    append(initMsg, 0),
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := looksLikeInitiationMsg(tt.b); got != tt.want {
+				t.Errorf("looksLikeInitiationMsg() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}