From a6559a892470de9d0cb7bb751698b41fb0b8aca6 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Thu, 16 Jul 2020 12:58:35 -0700
Subject: [PATCH] wgengine/magicsock: run test DERP in mode where only disco
 packets allowed

So we don't accidentally pass a NAT traversal test by having DERP pick up our slack
when we really just wanted DERP as an OOB messaging channel.
---
 derp/derp_server.go                  | 12 ++++++++++++
 disco/disco.go                       | 11 +++++++++++
 wgengine/magicsock/magicsock_test.go |  9 ++++++---
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/derp/derp_server.go b/derp/derp_server.go
index 36002ddca..3615bf77f 100644
--- a/derp/derp_server.go
+++ b/derp/derp_server.go
@@ -26,6 +26,7 @@ import (
 
 	"golang.org/x/crypto/nacl/box"
 	"golang.org/x/sync/errgroup"
+	"tailscale.com/disco"
 	"tailscale.com/metrics"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
@@ -47,6 +48,12 @@ type Server struct {
 	// before failing when writing to a client.
 	WriteTimeout time.Duration
 
+	// OnlyDisco controls whether, for tests, non-discovery packets
+	// are dropped. This is used by magicsock tests to verify that
+	// NAT traversal works (using DERP for out-of-band messaging)
+	// but the packets themselves aren't going via DERP.
+	OnlyDisco bool
+
 	privateKey key.Private
 	publicKey  key.Public
 	logf       logger.Logf
@@ -547,6 +554,11 @@ func (c *sclient) handleFrameSendPacket(ft frameType, fl uint32) error {
 		return fmt.Errorf("client %x: recvPacket: %v", c.key, err)
 	}
 
+	if s.OnlyDisco && !disco.LooksLikeDiscoWrapper(contents) {
+		s.packetsDropped.Add(1)
+		return nil
+	}
+
 	var fwd PacketForwarder
 	s.mu.Lock()
 	dst := s.clients[dstKey]
diff --git a/disco/disco.go b/disco/disco.go
index 8fa63cfae..b4d0ea2a1 100644
--- a/disco/disco.go
+++ b/disco/disco.go
@@ -31,6 +31,8 @@ import (
 // Magic is the 6 byte header of all discovery messages.
 const Magic = "TS💬" // 6 bytes: 0x54 53 f0 9f 92 ac
 
+const keyLen = 32
+
 // NonceLen is the length of the nonces used by nacl secretboxes.
 const NonceLen = 24
 
@@ -46,6 +48,15 @@ const v0 = byte(0)
 
 var errShort = errors.New("short message")
 
+// LooksLikeDiscoWrapper reports whether p looks like it's a packet
+// containing an encrypted disco message.
+func LooksLikeDiscoWrapper(p []byte) bool {
+	if len(p) < len(Magic)+keyLen+NonceLen {
+		return false
+	}
+	return string(p[:len(Magic)]) == Magic
+}
+
 // Parse parses the encrypted part of the message from inside the
 // nacl secretbox.
 func Parse(p []byte) (Message, error) {
diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go
index d8193d17d..30628bac6 100644
--- a/wgengine/magicsock/magicsock_test.go
+++ b/wgengine/magicsock/magicsock_test.go
@@ -243,19 +243,20 @@ func parseCIDR(t *testing.T, addr string) wgcfg.CIDR {
 	return cidr
 }
 
-func runDERP(t *testing.T, logf logger.Logf) (s *derp.Server, addr *net.TCPAddr, cleanupFn func()) {
+func runDERP(t *testing.T, logf logger.Logf, onlyDisco bool) (s *derp.Server, addr *net.TCPAddr, cleanupFn func()) {
 	var serverPrivateKey key.Private
 	if _, err := crand.Read(serverPrivateKey[:]); err != nil {
 		t.Fatal(err)
 	}
 
 	s = derp.NewServer(serverPrivateKey, logf)
+	s.OnlyDisco = onlyDisco
 
 	httpsrv := httptest.NewUnstartedServer(derphttp.Handler(s))
 	httpsrv.Config.ErrorLog = logger.StdLogger(logf)
 	httpsrv.Config.TLSNextProto = make(map[string]func(*http.Server, *tls.Conn, http.Handler))
 	httpsrv.StartTLS()
-	logf("DERP server URL: %s", httpsrv.URL)
+	logf("DERP server URL: %s  (onlyDisco=%v)", httpsrv.URL, onlyDisco)
 
 	cleanupFn = func() {
 		httpsrv.CloseClientConnections()
@@ -413,11 +414,13 @@ func testTwoDevicePing(t *testing.T, d *devices) {
 	rc := tstest.NewResourceCheck()
 	defer rc.Assert(t)
 
+	usingNatLab := d.m1 != (nettype.Std{})
+
 	// This gets reassigned inside every test, so that the connections
 	// all log using the "current" t.Logf function. Sigh.
 	logf, setT := makeNestable(t)
 
-	derpServer, derpAddr, derpCleanupFn := runDERP(t, logf)
+	derpServer, derpAddr, derpCleanupFn := runDERP(t, logf, usingNatLab)
 	defer derpCleanupFn()
 
 	stunAddr, stunCleanupFn := stuntest.ServeWithPacketListener(t, d.stun)