wgengine/magicsock: apply cryptorouting knob to periodic/init msg evaljwhited/lazy-endpoint-control-knobs

For consistency sake with the branch that can return lazyEndpoint when
no peerMap entry is found.

Updates tailscale/corp#31083

Signed-off-by: Jordan Whited <jordan@tailscale.com>

2 files changed, 113 insertions, 7 deletions

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index a4ba090ef..f7da2f820 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -1767,6 +1767,12 @@ func looksLikeInitiationMsg(b []byte) bool {
 		binary.LittleEndian.Uint32(b) == device.MessageInitiationType
 }
 
+// disableCryptoRouting returns true if controlKnobs are non-nil with
+// DisableCryptorouting set, otherwise it returns false.
+func (c *Conn) disableCryptoRouting() bool {
+	return c.controlKnobs != nil && c.controlKnobs.DisableCryptorouting.Load()
+}
+
 // receiveIP is the shared bits of ReceiveIPv4 and ReceiveIPv6.
 //
 // size is the length of 'b' to report up to wireguard-go (only relevant if
@@ -1842,7 +1848,7 @@ func (c *Conn) receiveIP(b []byte, ipp netip.AddrPort, cache *epAddrEndpointCach
 		de, ok := c.peerMap.endpointForEpAddr(src)
 		c.mu.Unlock()
 		if !ok {
-			if c.controlKnobs != nil && c.controlKnobs.DisableCryptorouting.Load() {
+			if c.disableCryptoRouting() {
 				// Note: UDP relay is dependent on cryptorouting enablement. We
 				// only update Geneve-encapsulated [epAddr]s in the [peerMap]
 				// via [lazyEndpoint].
@@ -1863,7 +1869,7 @@ func (c *Conn) receiveIP(b []byte, ipp netip.AddrPort, cache *epAddrEndpointCach
 	if stats := c.stats.Load(); stats != nil {
 		stats.UpdateRxPhysical(ep.nodeAddr, ipp, 1, geneveInclusivePacketLen)
 	}
-	if src.vni.isSet() && (connNoted || looksLikeInitiationMsg(b)) {
+	if src.vni.isSet() && (connNoted || looksLikeInitiationMsg(b)) && !c.disableCryptoRouting() {
 		// 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.
diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go
index 9399dab32..dacea6b17 100644
--- a/wgengine/magicsock/magicsock_test.go
+++ b/wgengine/magicsock/magicsock_test.go
@@ -3818,6 +3818,12 @@ func TestConn_receiveIP(t *testing.T) {
 		return ep
 	}
 
+	disableCryptoRoutingKnobs := func() *controlknobs.Knobs {
+		knobs := &controlknobs.Knobs{}
+		knobs.DisableCryptorouting.Store(true)
+		return knobs
+	}
+
 	tests := []struct {
 		name string
 		// A copy of b is used as input, tests may re-use the same value.
@@ -3832,6 +3838,8 @@ func TestConn_receiveIP(t *testing.T) {
 		// If insertWantEndpointTypeInPeerMap is true, use this [epAddr] for it
 		// in the [peerMap.setNodeKeyForEpAddr] call.
 		peerMapEpAddr epAddr
+		// controlKnobs may be nil
+		controlKnobs *controlknobs.Knobs
 		// If [*endpoint] then we expect 'got' to be the same [*endpoint]. If
 		// [*lazyEndpoint] and [*lazyEndpoint.maybeEP] is non-nil, we expect
 		// got.maybeEP to also be non-nil. Must not be reused across tests.
@@ -3891,6 +3899,19 @@ func TestConn_receiveIP(t *testing.T) {
 			wantNoteRecvActivityCalled: false,
 		},
 		{
+			name:                       "naked WireGuard init cryptorouting disabled empty peerMap",
+			b:                          looksLikeNakedWireGuardInit,
+			ipp:                        netip.MustParseAddrPort("127.0.0.1:7777"),
+			cache:                      &epAddrEndpointCache{},
+			controlKnobs:               disableCryptoRoutingKnobs(),
+			wantEndpointType:           nil,
+			wantSize:                   0,
+			wantIsGeneveEncap:          false,
+			wantOk:                     false,
+			wantMetricInc:              nil,
+			wantNoteRecvActivityCalled: false,
+		},
+		{
 			name:                            "naked WireGuard init endpoint matching peerMap entry",
 			b:                               looksLikeNakedWireGuardInit,
 			ipp:                             netip.MustParseAddrPort("127.0.0.1:7777"),
@@ -3917,6 +3938,19 @@ func TestConn_receiveIP(t *testing.T) {
 			wantNoteRecvActivityCalled: false,
 		},
 		{
+			name:                       "geneve WireGuard init cryptorouting disabled empty peerMap",
+			b:                          looksLikeGeneveWireGuardInit,
+			ipp:                        netip.MustParseAddrPort("127.0.0.1:7777"),
+			cache:                      &epAddrEndpointCache{},
+			controlKnobs:               disableCryptoRoutingKnobs(),
+			wantEndpointType:           nil,
+			wantSize:                   0,
+			wantIsGeneveEncap:          false,
+			wantOk:                     false,
+			wantMetricInc:              nil,
+			wantNoteRecvActivityCalled: false,
+		},
+		{
 			name:                            "geneve WireGuard init lazyEndpoint matching peerMap activity noted",
 			b:                               looksLikeGeneveWireGuardInit,
 			ipp:                             netip.MustParseAddrPort("127.0.0.1:7777"),
@@ -3933,6 +3967,21 @@ func TestConn_receiveIP(t *testing.T) {
 			wantNoteRecvActivityCalled: true,
 		},
 		{
+			name:                            "geneve WireGuard init cryptorouting disabled matching peerMap activity noted",
+			b:                               looksLikeGeneveWireGuardInit,
+			ipp:                             netip.MustParseAddrPort("127.0.0.1:7777"),
+			cache:                           &epAddrEndpointCache{},
+			insertWantEndpointTypeInPeerMap: true,
+			peerMapEpAddr:                   epAddr{ap: netip.MustParseAddrPort("127.0.0.1:7777"), vni: vni},
+			controlKnobs:                    disableCryptoRoutingKnobs(),
+			wantEndpointType:                newPeerMapInsertableEndpoint(0),
+			wantSize:                        len(looksLikeGeneveWireGuardInit) - packet.GeneveFixedHeaderLength,
+			wantIsGeneveEncap:               true,
+			wantOk:                          true,
+			wantMetricInc:                   nil,
+			wantNoteRecvActivityCalled:      true,
+		},
+		{
 			name:                            "geneve WireGuard init lazyEndpoint matching peerMap no activity noted",
 			b:                               looksLikeGeneveWireGuardInit,
 			ipp:                             netip.MustParseAddrPort("127.0.0.1:7777"),
@@ -3948,6 +3997,21 @@ func TestConn_receiveIP(t *testing.T) {
 			wantMetricInc:              nil,
 			wantNoteRecvActivityCalled: false,
 		},
+		{
+			name:                            "geneve WireGuard init cryptorouting disabled matching peerMap no activity noted",
+			b:                               looksLikeGeneveWireGuardInit,
+			ipp:                             netip.MustParseAddrPort("127.0.0.1:7777"),
+			cache:                           &epAddrEndpointCache{},
+			insertWantEndpointTypeInPeerMap: true,
+			peerMapEpAddr:                   epAddr{ap: netip.MustParseAddrPort("127.0.0.1:7777"), vni: vni},
+			controlKnobs:                    disableCryptoRoutingKnobs(),
+			wantEndpointType:                newPeerMapInsertableEndpoint(mono.Now().Add(time.Hour * 24)),
+			wantSize:                        len(looksLikeGeneveWireGuardInit) - packet.GeneveFixedHeaderLength,
+			wantIsGeneveEncap:               true,
+			wantOk:                          true,
+			wantMetricInc:                   nil,
+			wantNoteRecvActivityCalled:      false,
+		},
 		// TODO(jwhited): verify cache.de is used when conditions permit
 	}
 	for _, tt := range tests {
@@ -3959,15 +4023,14 @@ func TestConn_receiveIP(t *testing.T) {
 			}
 
 			// Init Conn.
-			c := &Conn{
-				privateKey: key.NewNode(),
-				netChecker: &netcheck.Client{},
-				peerMap:    newPeerMap(),
-			}
+			c := newConn(t.Logf)
+			c.privateKey = key.NewNode()
 			c.havePrivateKey.Store(true)
+			c.netChecker = &netcheck.Client{}
 			c.noteRecvActivity = func(public key.NodePublic) {
 				noteRecvActivityCalled = true
 			}
+			c.controlKnobs = tt.controlKnobs
 			c.SetStatistics(connstats.NewStatistics(0, 0, nil))
 
 			if tt.insertWantEndpointTypeInPeerMap {
@@ -4227,3 +4290,40 @@ func Test_lazyEndpoint_FromPeer(t *testing.T) {
 		})
 	}
 }
+
+func TestConn_disableCryptoRouting(t *testing.T) {
+	disableCryptoroutingSet := &controlknobs.Knobs{}
+	disableCryptoroutingSet.DisableCryptorouting.Store(true)
+
+	tests := []struct {
+		name  string
+		knobs *controlknobs.Knobs
+		want  bool
+	}{
+		{
+			name:  "nil",
+			knobs: nil,
+			want:  false,
+		},
+		{
+			name:  "DisableCryptorouting set",
+			knobs: disableCryptoroutingSet,
+			want:  true,
+		},
+		{
+			name:  "DisableCryptorouting unset",
+			knobs: &controlknobs.Knobs{},
+			want:  false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := &Conn{
+				controlKnobs: tt.knobs,
+			}
+			if got := c.disableCryptoRouting(); got != tt.want {
+				t.Errorf("disableCryptoRouting() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}