sshd: listen on the tsnet interface too

docs/rescue.md

@@ -191,8 +191,7 @@ pivot; pass `--keep-firewall` to keep it.
 
 ## Auto-reset on hang
 
-If the entrypoint hangs (stuck `apk add`, DNS timeout, kernel deadlock)
-the pivoted rootfs may be unreachable. `--watchdog-timeout` arms
+If the entrypoint hangs, the pivoted rootfs may be unreachable. `--watchdog-timeout` arms
 `/dev/watchdog` before the entrypoint runs and pets it from a goroutine;
 if anything stops that goroutine the kernel resets and the box comes
 back on the original OS.

internal/postpivot/run.go

@@ -44,24 +44,24 @@ func Run(argv []string) int {
 		}
 	}
 
-	if cfg != nil && cfg.SSH != nil {
-		go func() {
-			if err := StartSSHServer(context.Background(), cfg.SSH); err != nil {
-				slog.Error("sshd", "err", err)
-			}
-		}()
-	}
-
-	// Tailscale: re-open the tsnet state persisted by the pre-pivot
-	// PreAuth and serve SSH on the tailnet. Runs in a goroutine so the
-	// entrypoint supervisor still takes the foreground.
+	var tailnet TailnetListener
 	if cfg != nil && cfg.Tailscale != nil {
 		hostname := hostnameFromArgs(cfg.Tailscale.Args)
+		srv, err := tsnetauth.NewPostPivotServer(context.Background(), tsnetauth.PostPivotOptions{
+			Hostname: hostname,
+		})
+		if err != nil {
+			slog.Error("tsnet post-pivot", "err", err)
+		} else {
+			defer srv.Close()
+			tailnet = srv
+		}
+	}
+
+	if cfg != nil && cfg.SSH != nil {
 		go func() {
-			if err := tsnetauth.PostPivot(context.Background(), tsnetauth.PostPivotOptions{
-				Hostname: hostname,
-			}); err != nil {
-				slog.Error("tsnet post-pivot", "err", err)
+			if err := StartSSHServer(context.Background(), cfg.SSH, tailnet); err != nil {
+				slog.Error("sshd", "err", err)
 			}
 		}()
 	}

internal/postpivot/sshd_linux.go

@@ -22,44 +22,85 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
-// StartSSHServer runs an SSH server for the post-pivot rescue rootfs.
-// Blocks until ctx is cancelled or the listener errors. Public-key
-// (from cfg.AuthorizedKeys, OpenSSH `authorized_keys` format) and
-// password (cfg.Password) auth are both accepted when configured; the
-// server rejects clients when neither is set. Sessions run `/bin/sh`
-// as root — the pivoted rootfs has no user database.
-func StartSSHServer(ctx context.Context, cfg *SSHConfig) error {
+// TailnetListener is anything that can hand out net.Listeners on a
+// user-space tailnet interface (e.g. *tsnet.Server).
+type TailnetListener interface {
+	Listen(network, addr string) (net.Listener, error)
+}
+
+// StartSSHServer serves an SSH server on kernel :port and, when
+// tailnet is non-nil, additionally on the tailnet interface. Pubkey
+// (cfg.AuthorizedKeys, one per line) and password (cfg.Password) auth
+// are both accepted when configured; server refuses to start with
+// neither. Sessions run /bin/sh as root.
+func StartSSHServer(ctx context.Context, cfg *SSHConfig, tailnet TailnetListener) error {
 	if cfg == nil {
 		return errors.New("sshd: nil SSH config")
 	}
 	port := cfg.Port
 	if port == 0 {
 		port = 22
 	}
-	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
-	if err != nil {
-		return fmt.Errorf("sshd: listen :%d: %w", port, err)
+	addr := fmt.Sprintf(":%d", port)
+
+	var listeners []net.Listener
+	if kln, err := net.Listen("tcp", addr); err == nil {
+		listeners = append(listeners, kln)
+	} else {
+		slog.Warn("sshd: kernel listen failed", "err", err)
 	}
-	return serveSSH(ctx, ln, cfg)
+	if tailnet != nil {
+		if tln, err := tailnet.Listen("tcp", addr); err == nil {
+			listeners = append(listeners, tln)
+		} else {
+			slog.Warn("sshd: tailnet listen failed", "err", err)
+		}
+	}
+	if len(listeners) == 0 {
+		return fmt.Errorf("sshd: no listeners could bind :%d", port)
+	}
+	return serveSSH(ctx, listeners, cfg)
 }
 
-// serveSSH is the injectable core so tests can pass a random-port
-// listener. Closes ln on exit.
-func serveSSH(ctx context.Context, ln net.Listener, cfg *SSHConfig) error {
-	defer ln.Close()
-
+// serveSSH is the injectable core so tests can pass random-port
+// listeners. Closes every listener on exit.
+func serveSSH(ctx context.Context, listeners []net.Listener, cfg *SSHConfig) error {
 	sc, err := buildServerConfig(cfg)
 	if err != nil {
+		for _, ln := range listeners {
+			ln.Close()
+		}
 		return err
 	}
 
 	go func() {
 		<-ctx.Done()
-		ln.Close()
+		for _, ln := range listeners {
+			ln.Close()
+		}
 	}()
 
-	slog.Info("sshd: listening", "addr", ln.Addr())
+	var wg sync.WaitGroup
+	errCh := make(chan error, len(listeners))
+	for _, ln := range listeners {
+		slog.Info("sshd: listening", "addr", ln.Addr())
+		wg.Add(1)
+		go func(ln net.Listener) {
+			defer wg.Done()
+			errCh <- acceptLoop(ctx, ln, sc)
+		}(ln)
+	}
+	wg.Wait()
+	close(errCh)
+	for err := range errCh {
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
 
+func acceptLoop(ctx context.Context, ln net.Listener, sc *ssh.ServerConfig) error {
 	var wg sync.WaitGroup
 	defer wg.Wait()
 	for {
@@ -68,7 +109,7 @@ func serveSSH(ctx context.Context, ln net.Listener, cfg *SSHConfig) error {
 			if ctx.Err() != nil {
 				return nil
 			}
-			return fmt.Errorf("sshd: accept: %w", err)
+			return fmt.Errorf("sshd: accept on %s: %w", ln.Addr(), err)
 		}
 		wg.Add(1)
 		go func() {

internal/postpivot/sshd_linux_test.go

@@ -26,7 +26,7 @@ func serveSSHTest(t *testing.T, cfg *SSHConfig) string {
 	}
 	ctx, cancel := context.WithCancel(t.Context())
 	done := make(chan error, 1)
-	go func() { done <- serveSSH(ctx, ln, cfg) }()
+	go func() { done <- serveSSH(ctx, []net.Listener{ln}, cfg) }()
 	t.Cleanup(func() {
 		cancel()
 		select {
@@ -148,7 +148,7 @@ func TestSSHNoAuthConfiguredRefuses(t *testing.T) {
 	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 
-	err = serveSSH(ctx, ln, &SSHConfig{}) // no keys, no password
+	err = serveSSH(ctx, []net.Listener{ln}, &SSHConfig{}) // no keys, no password
 	if err == nil || !strings.Contains(err.Error(), "no auth method") {
 		t.Errorf("serveSSH err = %v, want 'no auth method'", err)
 	}
@@ -168,6 +168,30 @@ func TestParseAuthorizedKeys(t *testing.T) {
 	}
 }
 
+type fakeTailnet struct{}
+
+func (fakeTailnet) Listen(network, addr string) (net.Listener, error) {
+	return net.Listen(network, "127.0.0.1:0")
+}
+
+func TestStartSSHServerAddsTailnetListener(t *testing.T) {
+	_, authorized := genClientKey(t)
+	ctx, cancel := context.WithCancel(t.Context())
+	done := make(chan error, 1)
+	go func() { done <- StartSSHServer(ctx, &SSHConfig{AuthorizedKeys: authorized}, fakeTailnet{}) }()
+	// Give both listeners a moment to bind, then shut down.
+	time.Sleep(100 * time.Millisecond)
+	cancel()
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Errorf("StartSSHServer: %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Error("StartSSHServer did not exit")
+	}
+}
+
 func TestParseAuthorizedKeysBadLineReturnsError(t *testing.T) {
 	_, err := parseAuthorizedKeys("this-is-not-a-key")
 	if err == nil {

internal/postpivot/sshd_other.go

@@ -5,11 +5,17 @@ package postpivot
 import (
 	"context"
 	"errors"
+	"net"
 )
 
+// TailnetListener parity with the Linux build.
+type TailnetListener interface {
+	Listen(network, addr string) (net.Listener, error)
+}
+
 // StartSSHServer is a stub on non-Linux platforms. The pty + os/exec
 // shell handling requires Linux semantics, so we don't build the
 // server elsewhere.
-func StartSSHServer(_ context.Context, _ *SSHConfig) error {
+func StartSSHServer(_ context.Context, _ *SSHConfig, _ TailnetListener) error {
 	return errors.New("sshd: not supported on this platform")
 }

internal/tsnetauth/postpivot.go

@@ -8,52 +8,49 @@ import (
 	"tailscale.com/tsnet"
 )
 
-// PostPivotOptions configures PostPivot.
+// PostPivotOptions configures NewPostPivotServer.
 type PostPivotOptions struct {
-	// Hostname must match what PreAuth used; reusing the persisted
-	// state requires the same hostname.
+	// Hostname must match what PreAuth used.
 	Hostname string
-	// SSHAddr is the bind address for Tailscale SSH (default ":22").
-	SSHAddr string
 }
 
-// PostPivot is invoked from `xmorph --init` after the pivot. It opens
-// tsnet at /var/lib/tailscale (the path the persisted PreAuth state
-// landed at) and listens for Tailscale SSH on SSHAddr. AuthKey is empty
-// — the state on disk has us already authenticated.
-//
-// Blocks until ctx is cancelled or the listener errors. The caller
-// runs this in a goroutine alongside Supervise.
-func PostPivot(ctx context.Context, opts PostPivotOptions) error {
-	addr := opts.SSHAddr
-	if addr == "" {
-		addr = ":22"
-	}
-
+// NewPostPivotServer opens tsnet from the state persisted by PreAuth
+// and brings the node up on the tailnet. Caller owns Close.
+func NewPostPivotServer(ctx context.Context, opts PostPivotOptions) (*tsnet.Server, error) {
 	srv := &tsnet.Server{
 		Dir:      "/var/lib/tailscale",
 		Hostname: opts.Hostname,
 	}
-	defer srv.Close()
-
 	if _, err := srv.Up(ctx); err != nil {
-		return fmt.Errorf("post-pivot tsnet up: %w", err)
+		srv.Close()
+		return nil, fmt.Errorf("post-pivot tsnet up: %w", err)
 	}
+	return srv, nil
+}
 
-	// ListenSSH wires the in-process Tailscale SSH server to the
-	// tailnet listener (tsnet.Server.ListenSSH in v1.100.0). The
-	// returned net.Listener serves SSH connections — accepting and
-	// closing them is handled internally by tsnet.
+// ServeSSH runs Tailscale-native SSH (tsnet ListenSSH). Blocks until
+// ctx is cancelled. addr empty defaults to ":22".
+func ServeSSH(ctx context.Context, srv *tsnet.Server, addr string) error {
+	if addr == "" {
+		addr = ":22"
+	}
 	ln, err := srv.ListenSSH(addr)
 	if err != nil {
 		return fmt.Errorf("tsnet ListenSSH %s: %w", addr, err)
 	}
 	defer ln.Close()
-
-	slog.Info("tsnet SSH up", "hostname", opts.Hostname, "addr", addr)
-
-	// Block until ctx is done; the SSH listener runs in the background
-	// inside tsnet. Closing srv on defer tears down the listener.
+	slog.Info("tsnet SSH up", "hostname", srv.Hostname, "addr", addr)
 	<-ctx.Done()
 	return ctx.Err()
 }
+
+// PostPivot preserves the old one-call API: create + serve. Kept for
+// callers that don't need the *tsnet.Server for their own listeners.
+func PostPivot(ctx context.Context, opts PostPivotOptions) error {
+	srv, err := NewPostPivotServer(ctx, opts)
+	if err != nil {
+		return err
+	}
+	defer srv.Close()
+	return ServeSSH(ctx, srv, ":22")
+}