From 21cb0b361fbbb7580fae8f8a01193b3ec361e56a Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Tue, 31 Aug 2021 14:36:10 -0700
Subject: [PATCH] safesocket: add connect retry loop to wait for tailscaled

Updates #2708

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 cmd/tailscale/depaware.txt  |  2 +-
 cmd/tailscaled/depaware.txt |  1 +
 safesocket/safesocket.go    | 33 ++++++++++++++++++++++++++++++++-
 safesocket/safesocket_ps.go | 37 +++++++++++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 2 deletions(-)
 create mode 100644 safesocket/safesocket_ps.go

diff --git a/cmd/tailscale/depaware.txt b/cmd/tailscale/depaware.txt
index c9720cd75..180f07656 100644
--- a/cmd/tailscale/depaware.txt
+++ b/cmd/tailscale/depaware.txt
@@ -4,7 +4,7 @@ tailscale.com/cmd/tailscale dependencies: (generated by github.com/tailscale/dep
    W    github.com/alexbrainman/sspi/internal/common                 from github.com/alexbrainman/sspi/negotiate
    W 💣 github.com/alexbrainman/sspi/negotiate                       from tailscale.com/net/tshttpproxy
         github.com/kballard/go-shellquote                            from tailscale.com/cmd/tailscale/cli
-     💣 github.com/mitchellh/go-ps                                   from tailscale.com/cmd/tailscale/cli
+     💣 github.com/mitchellh/go-ps                                   from tailscale.com/cmd/tailscale/cli+
         github.com/peterbourgon/ff/v2                                from github.com/peterbourgon/ff/v2/ffcli
         github.com/peterbourgon/ff/v2/ffcli                          from tailscale.com/cmd/tailscale/cli
         github.com/tailscale/goupnp                                  from github.com/tailscale/goupnp/dcps/internetgateway2+
diff --git a/cmd/tailscaled/depaware.txt b/cmd/tailscaled/depaware.txt
index 8e1416352..8a8cae982 100644
--- a/cmd/tailscaled/depaware.txt
+++ b/cmd/tailscaled/depaware.txt
@@ -25,6 +25,7 @@ tailscale.com/cmd/tailscaled dependencies: (generated by github.com/tailscale/de
    L 💣 github.com/mdlayher/netlink/nlenc                            from github.com/mdlayher/netlink+
    L    github.com/mdlayher/sdnotify                                 from tailscale.com/util/systemd
    L 💣 github.com/mdlayher/socket                                   from github.com/mdlayher/netlink
+     💣 github.com/mitchellh/go-ps                                   from tailscale.com/safesocket
    W    github.com/pkg/errors                                        from github.com/tailscale/certstore
    W 💣 github.com/tailscale/certstore                               from tailscale.com/control/controlclient
         github.com/tailscale/goupnp                                  from github.com/tailscale/goupnp/dcps/internetgateway2+
diff --git a/safesocket/safesocket.go b/safesocket/safesocket.go
index f415bc11e..87fcc9d12 100644
--- a/safesocket/safesocket.go
+++ b/safesocket/safesocket.go
@@ -10,6 +10,7 @@ import (
 	"errors"
 	"net"
 	"runtime"
+	"time"
 )
 
 type closeable interface {
@@ -29,9 +30,39 @@ func ConnCloseWrite(c net.Conn) error {
 	return c.(closeable).CloseWrite()
 }
 
+var processStartTime = time.Now()
+var tailscaledProcExists = func() bool { return false } // set by safesocket_ps.go
+
+// tailscaledStillStarting reports whether tailscaled is probably
+// still starting up. That is, it reports whether the caller should
+// keep retrying to connect.
+func tailscaledStillStarting() bool {
+	d := time.Since(processStartTime)
+	if d < 2*time.Second {
+		// Without even checking the process table, assume
+		// that for the first two seconds that tailscaled is
+		// probably still starting.  That is, assume they're
+		// running "tailscaled & tailscale up ...." and make
+		// the tailscale client block for a bit for tailscaled
+		// to start accepting on the socket.
+		return true
+	}
+	if d > 5*time.Second {
+		return false
+	}
+	return tailscaledProcExists()
+}
+
 // Connect connects to either path (on Unix) or the provided localhost port (on Windows).
 func Connect(path string, port uint16) (net.Conn, error) {
-	return connect(path, port)
+	for {
+		c, err := connect(path, port)
+		if err != nil && tailscaledStillStarting() {
+			time.Sleep(250 * time.Millisecond)
+			continue
+		}
+		return c, err
+	}
 }
 
 // Listen returns a listener either on Unix socket path (on Unix), or
diff --git a/safesocket/safesocket_ps.go b/safesocket/safesocket_ps.go
new file mode 100644
index 000000000..b693501d8
--- /dev/null
+++ b/safesocket/safesocket_ps.go
@@ -0,0 +1,37 @@
+// Copyright (c) 2021 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build linux || windows || darwin || freebsd
+// +build linux windows darwin freebsd
+
+package safesocket
+
+import (
+	"strings"
+
+	ps "github.com/mitchellh/go-ps"
+)
+
+func init() {
+	tailscaledProcExists = func() bool {
+		procs, err := ps.Processes()
+		if err != nil {
+			return false
+		}
+		for _, proc := range procs {
+			name := proc.Executable()
+			const tailscaled = "tailscaled"
+			if len(name) < len(tailscaled) {
+				continue
+			}
+			// Do case insensitive comparison for Windows,
+			// notably, and ignore any ".exe" suffix.
+			if strings.EqualFold(name[:len(tailscaled)], tailscaled) {
+				return true
+			}
+		}
+		return false
+
+	}
+}