From b50d32059f1b33311dbba96a57c82d33a28f0e1f Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Tue, 21 Jan 2025 09:50:45 -0800
Subject: [PATCH] tsnet: block in Server.Dial until backend is Running

Updates #14715

Change-Id: I8c91e94fd1c6278c7f94a6b890274ed8a01e6f25
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 tsnet/tsnet.go      | 32 ++++++++++++++++++++++++++++++++
 tsnet/tsnet_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/tsnet/tsnet.go b/tsnet/tsnet.go
index fd894c38a..b769e719c 100644
--- a/tsnet/tsnet.go
+++ b/tsnet/tsnet.go
@@ -169,9 +169,41 @@ func (s *Server) Dial(ctx context.Context, network, address string) (net.Conn, e
 	if err := s.Start(); err != nil {
 		return nil, err
 	}
+	if err := s.awaitRunning(ctx); err != nil {
+		return nil, err
+	}
 	return s.dialer.UserDial(ctx, network, address)
 }
 
+// awaitRunning waits until the backend is in state Running.
+// If the backend is in state Starting, it blocks until it reaches
+// a terminal state (such as Stopped, NeedsMachineAuth)
+// or the context expires.
+func (s *Server) awaitRunning(ctx context.Context) error {
+	st := s.lb.State()
+	for {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		switch st {
+		case ipn.Running:
+			return nil
+		case ipn.NeedsLogin, ipn.Starting:
+			// Even after LocalBackend.Start, the state machine is still briefly
+			// in the "NeedsLogin" state. So treat that as also "Starting" and
+			// wait for us to get out of that state.
+			s.lb.WatchNotifications(ctx, ipn.NotifyInitialState, nil, func(n *ipn.Notify) (keepGoing bool) {
+				if n.State != nil {
+					st = *n.State
+				}
+				return st == ipn.NeedsLogin || st == ipn.Starting
+			})
+		default:
+			return fmt.Errorf("tsnet: backend in state %v", st)
+		}
+	}
+}
+
 // HTTPClient returns an HTTP client that is configured to connect over Tailscale.
 //
 // This is useful if you need to have your tsnet services connect to other devices on
diff --git a/tsnet/tsnet_test.go b/tsnet/tsnet_test.go
index c2f27d0f3..552e8dbee 100644
--- a/tsnet/tsnet_test.go
+++ b/tsnet/tsnet_test.go
@@ -232,6 +232,46 @@ func startServer(t *testing.T, ctx context.Context, controlURL, hostname string)
 	return s, status.TailscaleIPs[0], status.Self.PublicKey
 }
 
+func TestDialBlocks(t *testing.T) {
+	tstest.ResourceCheck(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	controlURL, _ := startControl(t)
+
+	// Make one tsnet that blocks until it's up.
+	s1, _, _ := startServer(t, ctx, controlURL, "s1")
+
+	ln, err := s1.Listen("tcp", ":8080")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	// Then make another tsnet node that will only be woken up
+	// upon the first dial.
+	tmp := filepath.Join(t.TempDir(), "s2")
+	os.MkdirAll(tmp, 0755)
+	s2 := &Server{
+		Dir:               tmp,
+		ControlURL:        controlURL,
+		Hostname:          "s2",
+		Store:             new(mem.Store),
+		Ephemeral:         true,
+		getCertForTesting: testCertRoot.getCert,
+	}
+	if *verboseNodes {
+		s2.Logf = log.Printf
+	}
+	t.Cleanup(func() { s2.Close() })
+
+	c, err := s2.Dial(ctx, "tcp", "s1:8080")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+}
+
 func TestConn(t *testing.T) {
 	tstest.ResourceCheck(t)
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)