derp/derphttp: add watch reconnection tests from #9719

Co-authored-by: Val <valerie@tailscale.com>
Signed-off-by: Anton Tolchanov <anton@tailscale.com>

cmd/derper/mesh.go

@@ -41,6 +41,7 @@ func startMeshWithHost(s *derp.Server, host string) error {
 		return err
 	}
 	c.MeshKey = s.MeshKey()
+	c.WatchConnectionChanges = true
 
 	// For meshed peers within a region, connect via VPC addresses.
 	c.SetURLDialer(func(ctx context.Context, network, addr string) (net.Conn, error) {

derp/derphttp/derphttp_client.go

@@ -56,6 +56,12 @@ type Client struct {
 	MeshKey   string             // optional; for trusted clients
 	IsProber  bool               // optional; for probers to optional declare themselves as such
 
+	// WatchConnectionChanges is whether the client wishes to subscribe to
+	// notifications about clients connecting & disconnecting.
+	//
+	// Only trusted connections (using MeshKey) are allowed to use this.
+	WatchConnectionChanges bool
+
 	// BaseContext, if non-nil, returns the base context to use for dialing a
 	// new derp server. If nil, context.Background is used.
 	// In either case, additional timeouts may be added to the base context.
@@ -80,6 +86,7 @@ type Client struct {
 	addrFamSelAtomic syncs.AtomicValue[AddressFamilySelector]
 
 	mu           sync.Mutex
+	started      bool // true upon first connect, never transitions to false
 	preferred    bool
 	canAckPings  bool
 	closed       bool
@@ -142,6 +149,15 @@ func NewClient(privateKey key.NodePrivate, serverURL string, logf logger.Logf) (
 	return c, nil
 }
 
+// isStarted reports whether this client has been used yet.
+//
+// If if reports false, it may still have its exported fields configured.
+func (c *Client) isStarted() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.started
+}
+
 // Connect connects or reconnects to the server, unless already connected.
 // It returns nil if there was already a good connection, or if one was made.
 func (c *Client) Connect(ctx context.Context) error {
@@ -284,6 +300,7 @@ func useWebsockets() bool {
 func (c *Client) connect(ctx context.Context, caller string) (client *derp.Client, connGen int, err error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
+	c.started = true
 	if c.closed {
 		return nil, 0, ErrClientClosed
 	}
@@ -495,6 +512,13 @@ func (c *Client) connect(ctx context.Context, caller string) (client *derp.Clien
 		}
 	}
 
+	if c.WatchConnectionChanges {
+		if err := derpClient.WatchConnectionChanges(); err != nil {
+			go httpConn.Close()
+			return nil, 0, err
+		}
+	}
+
 	c.serverPubKey = derpClient.ServerPublicKey()
 	c.client = derpClient
 	c.netConn = tcpConn
@@ -956,22 +980,6 @@ func (c *Client) NotePreferred(v bool) {
 	}
 }
 
-// WatchConnectionChanges sends a request to subscribe to
-// notifications about clients connecting & disconnecting.
-//
-// Only trusted connections (using MeshKey) are allowed to use this.
-func (c *Client) WatchConnectionChanges() error {
-	client, _, err := c.connect(c.newContext(), "derphttp.Client.WatchConnectionChanges")
-	if err != nil {
-		return err
-	}
-	err = client.WatchConnectionChanges()
-	if err != nil {
-		c.closeForReconnect(client)
-	}
-	return err
-}
-
 // ClosePeer asks the server to close target's TCP connection.
 //
 // Only trusted connections (using MeshKey) are allowed to use this.

derp/derphttp/derphttp_test.go

@@ -9,6 +9,7 @@ import (
 	"crypto/tls"
 	"net"
 	"net/http"
+	"net/netip"
 	"sync"
 	"testing"
 	"time"
@@ -206,3 +207,193 @@ func TestPing(t *testing.T) {
 		t.Fatalf("Ping: %v", err)
 	}
 }
+
+func newTestServer(t *testing.T, k key.NodePrivate) (serverURL string, s *derp.Server) {
+	s = derp.NewServer(k, t.Logf)
+	httpsrv := &http.Server{
+		TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler)),
+		Handler:      Handler(s),
+	}
+
+	ln, err := net.Listen("tcp4", "localhost:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	serverURL = "http://" + ln.Addr().String()
+	s.SetMeshKey("1234")
+
+	go func() {
+		if err := httpsrv.Serve(ln); err != nil {
+			if err == http.ErrServerClosed {
+				t.Logf("server closed")
+				return
+			}
+			panic(err)
+		}
+	}()
+	return
+}
+
+func newWatcherClient(t *testing.T, watcherPrivateKey key.NodePrivate, serverToWatchURL string) (c *Client) {
+	c, err := NewClient(watcherPrivateKey, serverToWatchURL, t.Logf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	c.MeshKey = "1234"
+	return
+}
+
+// breakConnection breaks the connection, which should trigger a reconnect.
+func (c *Client) breakConnection(brokenClient *derp.Client) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.client != brokenClient {
+		return
+	}
+	if c.netConn != nil {
+		c.netConn.Close()
+		c.netConn = nil
+	}
+	c.client = nil
+}
+
+// Test that a watcher connection successfully reconnects and processes peer
+// updates after a different thread breaks and reconnects the connection, while
+// the watcher is waiting on recv().
+func TestBreakWatcherConnRecv(t *testing.T) {
+	var wg sync.WaitGroup
+	defer wg.Wait()
+	// Make the watcher server
+	serverPrivateKey1 := key.NewNode()
+	_, s1 := newTestServer(t, serverPrivateKey1)
+	defer s1.Close()
+
+	// Make the watched server
+	serverPrivateKey2 := key.NewNode()
+	serverURL2, s2 := newTestServer(t, serverPrivateKey2)
+	defer s2.Close()
+
+	// Make the watcher (but it is not connected yet)
+	watcher1 := newWatcherClient(t, serverPrivateKey1, serverURL2)
+	defer watcher1.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	watcherChan := make(chan int, 1)
+
+	// Set the wait time after a connection fails to much lower
+	origRetryInterval := retryInterval
+	retryInterval = 50 * time.Millisecond
+	defer func() { retryInterval = origRetryInterval }()
+
+	// Start the watcher thread (which connects to the watched server)
+	wg.Add(1) // To avoid using t.Logf after the test ends. See https://golang.org/issue/40343
+	go func() {
+		defer wg.Done()
+		var peers int
+		add := func(k key.NodePublic, _ netip.AddrPort) {
+			t.Logf("add: %v", k.ShortString())
+			peers++
+			// Signal that the watcher has run
+			watcherChan <- peers
+		}
+		remove := func(k key.NodePublic) { t.Logf("remove: %v", k.ShortString()); peers-- }
+
+		watcher1.RunWatchConnectionLoop(ctx, serverPrivateKey1.Public(), t.Logf, add, remove)
+	}()
+
+	timer := time.NewTimer(5 * time.Second)
+	defer timer.Stop()
+
+	// Wait for the watcher to run, then break the connection and check if it
+	// reconnected and received peer updates.
+	for i := 0; i < 10; i++ {
+		select {
+		case peers := <-watcherChan:
+			if peers != 1 {
+				t.Fatal("wrong number of peers added during watcher connection")
+			}
+		case <-timer.C:
+			t.Fatalf("watcher did not process the peer update")
+		}
+		watcher1.breakConnection(watcher1.client)
+		// re-establish connection by sending a packet
+		watcher1.ForwardPacket(key.NodePublic{}, key.NodePublic{}, []byte("bogus"))
+
+		timer.Reset(5 * time.Second)
+	}
+}
+
+// Test that a watcher connection successfully reconnects and processes peer
+// updates after a different thread breaks and reconnects the connection, while
+// the watcher is not waiting on recv().
+func TestBreakWatcherConn(t *testing.T) {
+	var wg sync.WaitGroup
+	defer wg.Wait()
+	// Make the watcher server
+	serverPrivateKey1 := key.NewNode()
+	_, s1 := newTestServer(t, serverPrivateKey1)
+	defer s1.Close()
+
+	// Make the watched server
+	serverPrivateKey2 := key.NewNode()
+	serverURL2, s2 := newTestServer(t, serverPrivateKey2)
+	defer s2.Close()
+
+	// Make the watcher (but it is not connected yet)
+	watcher1 := newWatcherClient(t, serverPrivateKey1, serverURL2)
+	defer watcher1.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	watcherChan := make(chan int, 1)
+	breakerChan := make(chan bool, 1)
+
+	// Set the wait time after a connection fails to much lower
+	origRetryInterval := retryInterval
+	retryInterval = 50 * time.Millisecond
+	defer func() { retryInterval = origRetryInterval }()
+
+	// Start the watcher thread (which connects to the watched server)
+	wg.Add(1) // To avoid using t.Logf after the test ends. See https://golang.org/issue/40343
+	go func() {
+		defer wg.Done()
+		var peers int
+		add := func(k key.NodePublic, _ netip.AddrPort) {
+			t.Logf("add: %v", k.ShortString())
+			peers++
+			// Signal that the watcher has run
+			watcherChan <- peers
+			// Wait for breaker to run
+			<-breakerChan
+		}
+		remove := func(k key.NodePublic) { t.Logf("remove: %v", k.ShortString()); peers-- }
+
+		watcher1.RunWatchConnectionLoop(ctx, serverPrivateKey1.Public(), t.Logf, add, remove)
+	}()
+
+	timer := time.NewTimer(5 * time.Second)
+	defer timer.Stop()
+
+	// Wait for the watcher to run, then break the connection and check if it
+	// reconnected and received peer updates.
+	for i := 0; i < 10; i++ {
+		select {
+		case peers := <-watcherChan:
+			if peers != 1 {
+				t.Fatal("wrong number of peers added during watcher connection")
+			}
+		case <-timer.C:
+			t.Fatalf("watcher did not process the peer update")
+		}
+		watcher1.breakConnection(watcher1.client)
+		// re-establish connection by sending a packet
+		watcher1.ForwardPacket(key.NodePublic{}, key.NodePublic{}, []byte("bogus"))
+		// signal that the breaker is done
+		breakerChan <- true
+
+		timer.Reset(5 * time.Second)
+	}
+}

derp/derphttp/mesh_client.go

@@ -14,25 +14,36 @@ import (
 	"tailscale.com/types/logger"
 )
 
-// RunWatchConnectionLoop loops until ctx is done, sending WatchConnectionChanges and subscribing to
-// connection changes.
+var retryInterval = 5 * time.Second
+
+// RunWatchConnectionLoop loops until ctx is done, sending
+// WatchConnectionChanges and subscribing to connection changes.
 //
-// If the server's public key is ignoreServerKey, RunWatchConnectionLoop returns.
+// If the server's public key is ignoreServerKey, RunWatchConnectionLoop
+// returns.
 //
 // Otherwise, the add and remove funcs are called as clients come & go.
 //
-// infoLogf, if non-nil, is the logger to write periodic status
-// updates about how many peers are on the server. Error log output is
-// set to the c's logger, regardless of infoLogf's value.
+// infoLogf, if non-nil, is the logger to write periodic status updates about
+// how many peers are on the server. Error log output is set to the c's logger,
+// regardless of infoLogf's value.
 //
-// To force RunWatchConnectionLoop to return quickly, its ctx needs to
-// be closed, and c itself needs to be closed.
+// To force RunWatchConnectionLoop to return quickly, its ctx needs to be
+// closed, and c itself needs to be closed.
+//
+// It is a fatal error to call this on an already-started Client withoutq having
+// initialized Client.WatchConnectionChanges to true.
 func (c *Client) RunWatchConnectionLoop(ctx context.Context, ignoreServerKey key.NodePublic, infoLogf logger.Logf, add func(key.NodePublic, netip.AddrPort), remove func(key.NodePublic)) {
+	if !c.WatchConnectionChanges {
+		if c.isStarted() {
+			panic("invalid use of RunWatchConnectionLoop on already-started Client without setting Client.RunWatchConnectionLoop")
+		}
+		c.WatchConnectionChanges = true
+	}
 	if infoLogf == nil {
 		infoLogf = logger.Discard
 	}
 	logf := c.logf
-	const retryInterval = 5 * time.Second
 	const statusInterval = 10 * time.Second
 	var (
 		mu              sync.Mutex
@@ -101,14 +112,6 @@ func (c *Client) RunWatchConnectionLoop(ctx context.Context, ignoreServerKey key
 	}
 
 	for ctx.Err() == nil {
-		err := c.WatchConnectionChanges()
-		if err != nil {
-			clear()
-			logf("WatchConnectionChanges: %v", err)
-			sleep(retryInterval)
-			continue
-		}
-
 		if c.ServerPublicKey() == ignoreServerKey {
 			logf("detected self-connect; ignoring host")
 			return