From e6d0c92b1dd3ef9c6a50a722cb4474e24d4d8532 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Thu, 14 May 2020 10:01:48 -0700
Subject: [PATCH] wgengine/magicsock: clean up earlier fix a bit

Move WaitReady from fc88e34f42604bdf3c6c16200aa8051648890c2d into the
test code, and keep the derp-reading goroutine named for debugging.
---
 wgengine/magicsock/magicsock.go      | 28 ++++++++++------------------
 wgengine/magicsock/magicsock_test.go | 24 ++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 66b55f1e2..d9e3a4800 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -106,8 +106,8 @@ type Conn struct {
 
 	wantDerp      bool
 	privateKey    key.Private
-	myDerp        int // nearest DERP server; 0 means none/unknown
-	derpStarted   chan struct{}
+	myDerp        int           // nearest DERP server; 0 means none/unknown
+	derpStarted   chan struct{} // closed on first connection to DERP; for tests
 	activeDerp    map[int]activeDerp
 	prevDerp      map[int]*syncs.WaitGroupChan
 	derpTLSConfig *tls.Config // normally nil; used by tests
@@ -273,16 +273,6 @@ func Listen(opts Options) (*Conn, error) {
 
 func (c *Conn) donec() <-chan struct{} { return c.connCtx.Done() }
 
-// WaitReady waits until the magicsock is entirely initialized and connected
-// to its home DERP server. This is normally not necessary, since magicsock
-// is intended to be entirely asynchronous, but it helps eliminate race
-// conditions in tests. In particular, you can't expect two test magicsocks
-// to be able to connect to each other through a test DERP unless they are
-// both fully initialized before you try.
-func (c *Conn) WaitReady() {
-	<-c.derpStarted
-}
-
 // ignoreSTUNPackets sets a STUN packet processing func that does nothing.
 func (c *Conn) ignoreSTUNPackets() {
 	c.stunReceiveFunc.Store(func([]byte, *net.UDPAddr) {})
@@ -894,13 +884,15 @@ func (c *Conn) derpWriteChanOfAddr(addr *net.UDPAddr, peer key.Public) chan<- de
 	wg.Add(2)
 	c.prevDerp[nodeID] = wg
 
-	go func() {
-		dc.Connect(ctx)
-		if firstDerp {
+	if firstDerp {
+		startGate = c.derpStarted
+		go func() {
+			dc.Connect(ctx)
 			close(c.derpStarted)
-		}
-		c.runDerpReader(ctx, addr, dc, wg, startGate)
-	}()
+		}()
+	}
+
+	go c.runDerpReader(ctx, addr, dc, wg, startGate)
 	go c.runDerpWriter(ctx, addr, dc, ch, wg, startGate)
 
 	return ad.writeCh
diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go
index 52e8f4b9c..d365e1a03 100644
--- a/wgengine/magicsock/magicsock_test.go
+++ b/wgengine/magicsock/magicsock_test.go
@@ -34,6 +34,26 @@ import (
 	"tailscale.com/wgengine/tstun"
 )
 
+// WaitReady waits until the magicsock is entirely initialized and connected
+// to its home DERP server. This is normally not necessary, since magicsock
+// is intended to be entirely asynchronous, but it helps eliminate race
+// conditions in tests. In particular, you can't expect two test magicsocks
+// to be able to connect to each other through a test DERP unless they are
+// both fully initialized before you try.
+func (c *Conn) WaitReady(t *testing.T) {
+	t.Helper()
+	timer := time.NewTimer(10 * time.Second)
+	defer timer.Stop()
+	select {
+	case <-c.derpStarted:
+		return
+	case <-c.connCtx.Done():
+		t.Fatalf("magicsock.Conn closed while waiting for readiness")
+	case <-timer.C:
+		t.Fatalf("timeout waiting for readiness")
+	}
+}
+
 func TestListen(t *testing.T) {
 	tstest.PanicOnLog()
 	rc := tstest.NewResourceCheck()
@@ -406,8 +426,8 @@ func TestTwoDevicePing(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	conn1.WaitReady()
-	conn2.WaitReady()
+	conn1.WaitReady(t)
+	conn2.WaitReady(t)
 
 	ping1 := func(t *testing.T) {
 		msg2to1 := tuntest.Ping(net.ParseIP("1.0.0.1"), net.ParseIP("1.0.0.2"))