From aa660e0631cfd862775ef7f1dd31c1e7a5b69250 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Thu, 6 Apr 2023 19:57:30 +0300
Subject: [PATCH] feat(agentssh): Gracefully close SSH sessions on Close
 (#7027)

By tracking and closing sessions manually before closing the underlying
connections, we ensure that the termination is propagated to SSH/SFTP
clients and they're not left waiting for a connection timeout.

Refs: #6177
---
 agent/agentssh/agentssh.go | 51 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 50 insertions(+), 1 deletion(-)

diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
index c882380bac..86e1eb9e36 100644
--- a/agent/agentssh/agentssh.go
+++ b/agent/agentssh/agentssh.go
@@ -50,6 +50,7 @@ type Server struct {
 	mu        sync.RWMutex // Protects following.
 	listeners map[net.Listener]struct{}
 	conns     map[net.Conn]struct{}
+	sessions  map[ssh.Session]struct{}
 	closing   chan struct{}
 	// Wait for goroutines to exit, waited without
 	// a lock on mu but protected by closing.
@@ -86,6 +87,7 @@ func NewServer(ctx context.Context, logger slog.Logger, maxTimeout time.Duration
 	s := &Server{
 		listeners: make(map[net.Listener]struct{}),
 		conns:     make(map[net.Conn]struct{}),
+		sessions:  make(map[ssh.Session]struct{}),
 		logger:    logger,
 	}
 
@@ -129,7 +131,7 @@ func NewServer(ctx context.Context, logger slog.Logger, maxTimeout time.Duration
 			}
 		},
 		SubsystemHandlers: map[string]ssh.SubsystemHandler{
-			"sftp": s.sftpHandler,
+			"sftp": s.sessionHandler,
 		},
 		MaxTimeout: maxTimeout,
 	}
@@ -152,7 +154,26 @@ func (s *Server) ConnStats() ConnStats {
 }
 
 func (s *Server) sessionHandler(session ssh.Session) {
+	if !s.trackSession(session, true) {
+		// See (*Server).Close() for why we call Close instead of Exit.
+		_ = session.Close()
+		return
+	}
+	defer s.trackSession(session, false)
+
 	ctx := session.Context()
+
+	switch ss := session.Subsystem(); ss {
+	case "":
+	case "sftp":
+		s.sftpHandler(session)
+		return
+	default:
+		s.logger.Debug(ctx, "unsupported subsystem", slog.F("subsystem", ss))
+		_ = session.Exit(1)
+		return
+	}
+
 	err := s.sessionStart(session)
 	var exitError *exec.ExitError
 	if xerrors.As(err, &exitError) {
@@ -560,6 +581,25 @@ func (s *Server) trackConn(l net.Listener, c net.Conn, add bool) (ok bool) {
 	return true
 }
 
+// trackSession registers the session with the server. If the server is
+// closing, the session is not registered and should be closed.
+//
+//nolint:revive
+func (s *Server) trackSession(ss ssh.Session, add bool) (ok bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if add {
+		if s.closing != nil {
+			// Server closed.
+			return false
+		}
+		s.sessions[ss] = struct{}{}
+		return true
+	}
+	delete(s.sessions, ss)
+	return true
+}
+
 // Close the server and all active connections. Server can be re-used
 // after Close is done.
 func (s *Server) Close() error {
@@ -573,6 +613,15 @@ func (s *Server) Close() error {
 	}
 	s.closing = make(chan struct{})
 
+	// Close all active sessions to gracefully
+	// terminate client connections.
+	for ss := range s.sessions {
+		// We call Close on the underlying channel here because we don't
+		// want to send an exit status to the client (via Exit()).
+		// Typically OpenSSH clients will return 255 as the exit status.
+		_ = ss.Close()
+	}
+
 	// Close all active listeners and connections.
 	for l := range s.listeners {
 		_ = l.Close()