From b743b85dad8c96f453834fbbaf2d87aa094da46e Mon Sep 17 00:00:00 2001
From: Andrew Lytvynov <awly@tailscale.com>
Date: Mon, 22 Apr 2024 09:27:12 -0700
Subject: [PATCH] ipn/ipnlocal,ssh/tailssh: reject c2n /update if SSH conns are
 active (#11820)

Since we already track active SSH connections, it's not hard to
proactively reject updates until those finish. We attempt to do the same
on the control side, but the detection latency for new connections is in
the minutes, which is not fast enough for common short sessions.

Handle a `force=true` query parameter to override this behavior, so that
control can still trigger an update on a server where some long-running
abandoned SSH session is open.

Updates https://github.com/tailscale/corp/issues/18556

Signed-off-by: Andrew Lytvynov <awly@tailscale.com>
---
 ipn/ipnlocal/c2n.go    | 7 +++++++
 ipn/ipnlocal/local.go  | 4 ++++
 ssh/tailssh/tailssh.go | 7 +++++++
 3 files changed, 18 insertions(+)

diff --git a/ipn/ipnlocal/c2n.go b/ipn/ipnlocal/c2n.go
index 26435570e..ec7dc8bc4 100644
--- a/ipn/ipnlocal/c2n.go
+++ b/ipn/ipnlocal/c2n.go
@@ -300,6 +300,13 @@ func handleC2NUpdatePost(b *LocalBackend, w http.ResponseWriter, r *http.Request
 		return
 	}
 
+	// Do not update if we have active inbound SSH connections. Control can set
+	// force=true query parameter to override this.
+	if r.FormValue("force") != "true" && b.sshServer != nil && b.sshServer.NumActiveConns() > 0 {
+		res.Err = "not updating due to active SSH connections"
+		return
+	}
+
 	// Check if update was already started, and mark as started.
 	if !b.trySetC2NUpdateStarted() {
 		res.Err = "update already started"
diff --git a/ipn/ipnlocal/local.go b/ipn/ipnlocal/local.go
index 4c57502aa..62e8c7905 100644
--- a/ipn/ipnlocal/local.go
+++ b/ipn/ipnlocal/local.go
@@ -123,6 +123,10 @@ func getControlDebugFlags() []string {
 type SSHServer interface {
 	HandleSSHConn(net.Conn) error
 
+	// NumActiveConns returns the number of connections passed to HandleSSHConn
+	// that are still active.
+	NumActiveConns() int
+
 	// OnPolicyChange is called when the SSH access policy changes,
 	// so that existing sessions can be re-evaluated for validity
 	// and closed if they'd no longer be accepted.
diff --git a/ssh/tailssh/tailssh.go b/ssh/tailssh/tailssh.go
index 69df3add2..2bfb645f3 100644
--- a/ssh/tailssh/tailssh.go
+++ b/ssh/tailssh/tailssh.go
@@ -143,6 +143,13 @@ func (srv *server) trackActiveConn(c *conn, add bool) {
 	delete(srv.activeConns, c)
 }
 
+// NumActiveConns returns the number of active SSH connections.
+func (srv *server) NumActiveConns() int {
+	srv.mu.Lock()
+	defer srv.mu.Unlock()
+	return len(srv.activeConns)
+}
+
 // HandleSSHConn handles a Tailscale SSH connection from c.
 // This is the entry point for all SSH connections.
 // When this returns, the connection is closed.