fix: stop tearing down non-TTY processes on SSH session end (cherry-pick #18673) (#18676)

gcp-cherry-pick-bot[bot] · spikecurtis · web-flow · commit 774792476c0c · 2025-06-30T23:09:37.000+04:00
Cherry-picked fix: stop tearing down non-TTY processes on SSH session end (#18673) (possibly temporary) fix for #18519 Matches OpenSSH for non-tty sessions, where we don't actively terminate the process. Adds explicit tracking to the SSH server for these processes so that if we are shutting down we terminate them: this ensures that we can shut down quickly to allow shutdown scripts to run. It also ensures our tests don't leak system resources. Co-authored-by: Spike Curtis <spike@coder.com>
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -125,6 +125,7 @@ type Server struct {
 	listeners map[net.Listener]struct{}
 	conns     map[net.Conn]struct{}
 	sessions  map[ssh.Session]struct{}
+	processes map[*os.Process]struct{}
 	closing   chan struct{}
 	// Wait for goroutines to exit, waited without
 	// a lock on mu but protected by closing.
@@ -183,6 +184,7 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
 		fs:        fs,
 		conns:     make(map[net.Conn]struct{}),
 		sessions:  make(map[ssh.Session]struct{}),
+		processes: make(map[*os.Process]struct{}),
 		logger:    logger,
 
 		config: config,
@@ -587,7 +589,10 @@ func (s *Server) startNonPTYSession(logger slog.Logger, session ssh.Session, mag
 	// otherwise context cancellation will not propagate properly
 	// and SSH server close may be delayed.
 	cmd.SysProcAttr = cmdSysProcAttr()
-	cmd.Cancel = cmdCancel(session.Context(), logger, cmd)
+
+	// to match OpenSSH, we don't actually tear a non-TTY command down, even if the session ends.
+	// c.f. https://github.com/coder/coder/issues/18519#issuecomment-3019118271
+	cmd.Cancel = nil
 
 	cmd.Stdout = session
 	cmd.Stderr = session.Stderr()
@@ -610,6 +615,16 @@ func (s *Server) startNonPTYSession(logger slog.Logger, session ssh.Session, mag
 		s.metrics.sessionErrors.WithLabelValues(magicTypeLabel, "no", "start_command").Add(1)
 		return xerrors.Errorf("start: %w", err)
 	}
+
+	// Since we don't cancel the process when the session stops, we still need to tear it down if we are closing. So
+	// track it here.
+	if !s.trackProcess(cmd.Process, true) {
+		// must be closing
+		err = cmdCancel(logger, cmd.Process)
+		return xerrors.Errorf("failed to track process: %w", err)
+	}
+	defer s.trackProcess(cmd.Process, false)
+
 	sigs := make(chan ssh.Signal, 1)
 	session.Signals(sigs)
 	defer func() {
@@ -1070,6 +1085,27 @@ func (s *Server) trackSession(ss ssh.Session, add bool) (ok bool) {
 	return true
 }
 
+// trackCommand registers the process with the server. If the server is
+// closing, the process is not registered and should be closed.
+//
+//nolint:revive
+func (s *Server) trackProcess(p *os.Process, add bool) (ok bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if add {
+		if s.closing != nil {
+			// Server closed.
+			return false
+		}
+		s.wg.Add(1)
+		s.processes[p] = struct{}{}
+		return true
+	}
+	s.wg.Done()
+	delete(s.processes, p)
+	return true
+}
+
 // Close the server and all active connections. Server can be re-used
 // after Close is done.
 func (s *Server) Close() error {
@@ -1109,6 +1145,10 @@ func (s *Server) Close() error {
 		_ = c.Close()
 	}
 
+	for p := range s.processes {
+		_ = cmdCancel(s.logger, p)
+	}
+
 	s.logger.Debug(ctx, "closing SSH server")
 	err := s.srv.Close()
 
diff --git a/agent/agentssh/exec_other.go b/agent/agentssh/exec_other.go
@@ -4,7 +4,7 @@ package agentssh
 
 import (
 	"context"
-	"os/exec"
+	"os"
 	"syscall"
 
 	"cdr.dev/slog"
@@ -16,9 +16,7 @@ func cmdSysProcAttr() *syscall.SysProcAttr {
 	}
 }
 
-func cmdCancel(ctx context.Context, logger slog.Logger, cmd *exec.Cmd) func() error {
-	return func() error {
-		logger.Debug(ctx, "cmdCancel: sending SIGHUP to process and children", slog.F("pid", cmd.Process.Pid))
-		return syscall.Kill(-cmd.Process.Pid, syscall.SIGHUP)
-	}
+func cmdCancel(logger slog.Logger, p *os.Process) error {
+	logger.Debug(context.Background(), "cmdCancel: sending SIGHUP to process and children", slog.F("pid", p.Pid))
+	return syscall.Kill(-p.Pid, syscall.SIGHUP)
 }
diff --git a/agent/agentssh/exec_windows.go b/agent/agentssh/exec_windows.go
@@ -2,7 +2,7 @@ package agentssh
 
 import (
 	"context"
-	"os/exec"
+	"os"
 	"syscall"
 
 	"cdr.dev/slog"
@@ -12,14 +12,12 @@ func cmdSysProcAttr() *syscall.SysProcAttr {
 	return &syscall.SysProcAttr{}
 }
 
-func cmdCancel(ctx context.Context, logger slog.Logger, cmd *exec.Cmd) func() error {
-	return func() error {
-		logger.Debug(ctx, "cmdCancel: killing process", slog.F("pid", cmd.Process.Pid))
-		// Windows doesn't support sending signals to process groups, so we
-		// have to kill the process directly. In the future, we may want to
-		// implement a more sophisticated solution for process groups on
-		// Windows, but for now, this is a simple way to ensure that the
-		// process is terminated when the context is cancelled.
-		return cmd.Process.Kill()
-	}
+func cmdCancel(logger slog.Logger, p *os.Process) error {
+	logger.Debug(context.Background(), "cmdCancel: killing process", slog.F("pid", p.Pid))
+	// Windows doesn't support sending signals to process groups, so we
+	// have to kill the process directly. In the future, we may want to
+	// implement a more sophisticated solution for process groups on
+	// Windows, but for now, this is a simple way to ensure that the
+	// process is terminated when the context is cancelled.
+	return p.Kill()
 }

Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@ package agentssh`
`4`	`4`
`5`	`5`	`import (`
`6`	`6`	`"context"`
`7`		`- "os/exec"`
	`7`	`+ "os"`
`8`	`8`	`"syscall"`
`9`	`9`
`10`	`10`	`"cdr.dev/slog"`
`@@ -16,9 +16,7 @@ func cmdSysProcAttr() *syscall.SysProcAttr {`
`16`	`16`	`}`
`17`	`17`	`}`
`18`	`18`
`19`		`-func cmdCancel(ctx context.Context, logger slog.Logger, cmd *exec.Cmd) func() error {`
`20`		`- return func() error {`
`21`		`- logger.Debug(ctx, "cmdCancel: sending SIGHUP to process and children", slog.F("pid", cmd.Process.Pid))`
`22`		`- return syscall.Kill(-cmd.Process.Pid, syscall.SIGHUP)`
`23`		`- }`
	`19`	`+func cmdCancel(logger slog.Logger, p *os.Process) error {`
	`20`	`+ logger.Debug(context.Background(), "cmdCancel: sending SIGHUP to process and children", slog.F("pid", p.Pid))`
	`21`	`+ return syscall.Kill(-p.Pid, syscall.SIGHUP)`
`24`	`22`	`}`