refactor(agent): Move SSH server into agentssh package - #7004

--- agent/agent.go	2023-04-04 14:57:21.344043135 +0000
+++ agent/agentssh/agentssh.go	2023-04-04 14:55:18.747492195 +0000
@@ -1,4 +1,4 @@
-package agent
+package agentssh
 
 import (
 	"bufio"
@@ -16,6 +16,8 @@
 	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
+	"time"
 
 	"github.com/gliderlabs/ssh"
 	"github.com/pkg/sftp"
@@ -23,7 +25,9 @@
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+
 	"github.com/coder/coder/agent/usershell"
+	"github.com/coder/coder/codersdk/agentsdk"
 	"github.com/coder/coder/pty"
 )
 
@@ -33,46 +37,65 @@
 	// unlikely to shadow other exit codes, which are typically 1, 2, 3, etc.
 	MagicSessionErrorCode = 229
 
-	// MagicSSHSessionTypeEnvironmentVariable is used to track the purpose behind an SSH connection.
+	// MagicSessionTypeEnvironmentVariable is used to track the purpose behind an SSH connection.
 	// This is stripped from any commands being executed, and is counted towards connection stats.
-	MagicSSHSessionTypeEnvironmentVariable = "CODER_SSH_SESSION_TYPE"
-	// MagicSSHSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
-	MagicSSHSessionTypeVSCode = "vscode"
-	// MagicSSHSessionTypeJetBrains is set in the SSH config by the JetBrains extension to identify itself.
-	MagicSSHSessionTypeJetBrains = "jetbrains"
+	MagicSessionTypeEnvironmentVariable = "CODER_SSH_SESSION_TYPE"
+	// MagicSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
+	MagicSessionTypeVSCode = "vscode"
+	// MagicSessionTypeJetBrains is set in the SSH config by the JetBrains extension to identify itself.
+	MagicSessionTypeJetBrains = "jetbrains"
 )
 
-func (a *agent) init(ctx context.Context) {
+type Server struct {
+	serveWg sync.WaitGroup
+	logger  slog.Logger
+
+	srv *ssh.Server
+
+	Env        map[string]string
+	AgentToken func() string
+
+	manifest atomic.Pointer[agentsdk.Manifest]
+
+	connCountVSCode     atomic.Int64
+	connCountJetBrains  atomic.Int64
+	connCountSSHSession atomic.Int64
+}
+
+func NewServer(ctx context.Context, logger slog.Logger, maxTimeout time.Duration) (*Server, error) {
 	// Clients' should ignore the host key when connecting.
 	// The agent needs to authenticate with coderd to SSH,
 	// so SSH authentication doesn't improve security.
 	randomHostKey, err := rsa.GenerateKey(rand.Reader, 2048)
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 	randomSigner, err := gossh.NewSignerFromKey(randomHostKey)
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 
-	sshLogger := a.logger.Named("ssh-server")
 	forwardHandler := &ssh.ForwardedTCPHandler{}
-	unixForwardHandler := &forwardedUnixHandler{log: a.logger}
+	unixForwardHandler := &forwardedUnixHandler{log: logger}
+
+	s := &Server{
+		logger: logger,
+	}
 
-	a.sshServer = &ssh.Server{
+	s.srv = &ssh.Server{
 		ChannelHandlers: map[string]ssh.ChannelHandler{
 			"direct-tcpip":                   ssh.DirectTCPIPHandler,
 			"[email protected]": directStreamLocalHandler,
 			"session":                        ssh.DefaultSessionHandler,
 		},
-		ConnectionFailedCallback: func(conn net.Conn, err error) {
-			sshLogger.Info(ctx, "ssh connection ended", slog.Error(err))
+		ConnectionFailedCallback: func(_ net.Conn, err error) {
+			s.logger.Info(ctx, "ssh connection ended", slog.Error(err))
 		},
-		Handler:     a.sessionHandler,
+		Handler:     s.sessionHandler,
 		HostSigners: []ssh.Signer{randomSigner},
 		LocalPortForwardingCallback: func(ctx ssh.Context, destinationHost string, destinationPort uint32) bool {
 			// Allow local port forwarding all!
-			sshLogger.Debug(ctx, "local port forward",
+			s.logger.Debug(ctx, "local port forward",
 				slog.F("destination-host", destinationHost),
 				slog.F("destination-port", destinationPort))
 			return true
@@ -82,7 +105,7 @@
 		},
 		ReversePortForwardingCallback: func(ctx ssh.Context, bindHost string, bindPort uint32) bool {
 			// Allow reverse port forwarding all!
-			sshLogger.Debug(ctx, "local port forward",
+			s.logger.Debug(ctx, "local port forward",
 				slog.F("bind-host", bindHost),
 				slog.F("bind-port", bindPort))
 			return true
@@ -99,24 +122,44 @@
 			}
 		},
 		SubsystemHandlers: map[string]ssh.SubsystemHandler{
-			"sftp": a.sftpHandler,
+			"sftp": s.sftpHandler,
 		},
-		MaxTimeout: a.sshMaxTimeout,
+		MaxTimeout: maxTimeout,
+	}
+
+	return s, nil
+}
+
+// SetManifest sets the manifest used for starting commands.
+func (s *Server) SetManifest(m *agentsdk.Manifest) {
+	s.manifest.Store(m)
+}
+
+type ConnStats struct {
+	Sessions  int64
+	VSCode    int64
+	JetBrains int64
 	}
 
-	go a.runLoop(ctx)
+func (s *Server) ConnStats() ConnStats {
+	return ConnStats{
+		Sessions:  s.connCountSSHSession.Load(),
+		VSCode:    s.connCountVSCode.Load(),
+		JetBrains: s.connCountJetBrains.Load(),
+	}
 }
 
-func (a *agent) sessionHandler(session ssh.Session) {
-	err := a.handleSSHSession(session)
+func (s *Server) sessionHandler(session ssh.Session) {
+	ctx := session.Context()
+	err := s.sessionStart(session)
 	var exitError *exec.ExitError
 	if xerrors.As(err, &exitError) {
-		a.logger.Debug(ctx, "ssh session returned", slog.Error(exitError))
+		s.logger.Debug(ctx, "ssh session returned", slog.Error(exitError))
 		_ = session.Exit(exitError.ExitCode())
 		return
 	}
 	if err != nil {
-		a.logger.Warn(ctx, "ssh session failed", slog.Error(err))
+		s.logger.Warn(ctx, "ssh session failed", slog.Error(err))
 		// This exit code is designed to be unlikely to be confused for a legit exit code
 		// from the process.
 		_ = session.Exit(MagicSessionErrorCode)
@@ -125,32 +168,32 @@
 	_ = session.Exit(0)
 }
 
-func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
+func (s *Server) sessionStart(session ssh.Session) (retErr error) {
 	ctx := session.Context()
 	env := session.Environ()
 	var magicType string
 	for index, kv := range env {
-		if !strings.HasPrefix(kv, MagicSSHSessionTypeEnvironmentVariable) {
+		if !strings.HasPrefix(kv, MagicSessionTypeEnvironmentVariable) {
 			continue
 		}
-		magicType = strings.TrimPrefix(kv, MagicSSHSessionTypeEnvironmentVariable+"=")
+		magicType = strings.TrimPrefix(kv, MagicSessionTypeEnvironmentVariable+"=")
 		env = append(env[:index], env[index+1:]...)
 	}
 	switch magicType {
-	case MagicSSHSessionTypeVSCode:
-		a.connCountVSCode.Add(1)
-		defer a.connCountVSCode.Add(-1)
-	case MagicSSHSessionTypeJetBrains:
-		a.connCountJetBrains.Add(1)
-		defer a.connCountJetBrains.Add(-1)
+	case MagicSessionTypeVSCode:
+		s.connCountVSCode.Add(1)
+		defer s.connCountVSCode.Add(-1)
+	case MagicSessionTypeJetBrains:
+		s.connCountJetBrains.Add(1)
+		defer s.connCountJetBrains.Add(-1)
 	case "":
-		a.connCountSSHSession.Add(1)
-		defer a.connCountSSHSession.Add(-1)
+		s.connCountSSHSession.Add(1)
+		defer s.connCountSSHSession.Add(-1)
 	default:
-		a.logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("type", magicType))
+		s.logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("type", magicType))
 	}
 
-	cmd, err := a.createCommand(ctx, session.RawCommand(), env)
+	cmd, err := s.CreateCommand(ctx, session.RawCommand(), env)
 	if err != nil {
 		return err
 	}
@@ -172,14 +215,14 @@
 		session.DisablePTYEmulation()
 
 		if !isQuietLogin(session.RawCommand()) {
-			manifest := a.manifest.Load()
+			manifest := s.manifest.Load()
 			if manifest != nil {
 				err = showMOTD(session, manifest.MOTDFile)
 				if err != nil {
-					a.logger.Error(ctx, "show MOTD", slog.Error(err))
+					s.logger.Error(ctx, "show MOTD", slog.Error(err))
 				}
 			} else {
-				a.logger.Warn(ctx, "metadata lookup failed, unable to show MOTD")
+				s.logger.Warn(ctx, "metadata lookup failed, unable to show MOTD")
 			}
 		}
 
@@ -188,7 +231,7 @@
 		// The pty package sets `SSH_TTY` on supported platforms.
 		ptty, process, err := pty.Start(cmd, pty.WithPTYOption(
 			pty.WithSSHRequest(sshPty),
-			pty.WithLogger(slog.Stdlib(ctx, a.logger, slog.LevelInfo)),
+			pty.WithLogger(slog.Stdlib(ctx, s.logger, slog.LevelInfo)),
 		))
 		if err != nil {
 			return xerrors.Errorf("start command: %w", err)
@@ -198,7 +241,7 @@
 			defer wg.Wait()
 			closeErr := ptty.Close()
 			if closeErr != nil {
-				a.logger.Warn(ctx, "failed to close tty", slog.Error(closeErr))
+				s.logger.Warn(ctx, "failed to close tty", slog.Error(closeErr))
 				if retErr == nil {
 					retErr = closeErr
 				}
@@ -209,7 +252,7 @@
 				resizeErr := ptty.Resize(uint16(win.Height), uint16(win.Width))
 				// If the pty is closed, then command has exited, no need to log.
 				if resizeErr != nil && !errors.Is(resizeErr, pty.ErrClosed) {
-					a.logger.Warn(ctx, "failed to resize tty", slog.Error(resizeErr))
+					s.logger.Warn(ctx, "failed to resize tty", slog.Error(resizeErr))
 				}
 			}
 		}()
@@ -258,7 +301,7 @@
 		// ExitErrors just mean the command we run returned a non-zero exit code, which is normal
 		// and not something to be concerned about.  But, if it's something else, we should log it.
 		if err != nil && !xerrors.As(err, &exitErr) {
-			a.logger.Warn(ctx, "wait error", slog.Error(err))
+			s.logger.Warn(ctx, "wait error", slog.Error(err))
 		}
 		return err
 	}
@@ -287,7 +330,7 @@
 // Close implements io.Closer.
 func (readNopCloser) Close() error { return nil }
 
-func (a *agent) sftpHandler(session ssh.Session) {
+func (s *Server) sftpHandler(session ssh.Session) {
 	ctx := session.Context()
 
 	// Typically sftp sessions don't request a TTY, but if they do,
@@ -301,14 +344,14 @@
 	// directory so that SFTP connections land there.
 	homedir, err := userHomeDir()
 	if err != nil {
-		sshLogger.Warn(ctx, "get sftp working directory failed, unable to get home dir", slog.Error(err))
+		s.logger.Warn(ctx, "get sftp working directory failed, unable to get home dir", slog.Error(err))
 	} else {
 		opts = append(opts, sftp.WithServerWorkingDirectory(homedir))
 	}
 
 	server, err := sftp.NewServer(session, opts...)
 	if err != nil {
-		sshLogger.Debug(ctx, "initialize sftp server", slog.Error(err))
+		s.logger.Debug(ctx, "initialize sftp server", slog.Error(err))
 		return
 	}
 	defer server.Close()
@@ -326,14 +369,14 @@
 		_ = session.Exit(0)
 		return
 	}
-	sshLogger.Warn(ctx, "sftp server closed with error", slog.Error(err))
+	s.logger.Warn(ctx, "sftp server closed with error", slog.Error(err))
 	_ = session.Exit(1)
 }
 
-// createCommand processes raw command input with OpenSSH-like behavior.
+// CreateCommand processes raw command input with OpenSSH-like behavior.
 // If the script provided is empty, it will default to the users shell.
 // This injects environment variables specified by the user at launch too.
-func (a *agent) createCommand(ctx context.Context, script string, env []string) (*exec.Cmd, error) {
+func (s *Server) CreateCommand(ctx context.Context, script string, env []string) (*exec.Cmd, error) {
 	currentUser, err := user.Current()
 	if err != nil {
 		return nil, xerrors.Errorf("get current user: %w", err)
@@ -345,7 +388,7 @@
 		return nil, xerrors.Errorf("get user shell: %w", err)
 	}
 
-	manifest := a.manifest.Load()
+	manifest := s.manifest.Load()
 	if manifest == nil {
 		return nil, xerrors.Errorf("no metadata was provided")
 	}
@@ -398,7 +441,7 @@
 	cmd.Env = append(cmd.Env, fmt.Sprintf(`GIT_SSH_COMMAND=%s gitssh --`, unixExecutablePath))
 
 	// Specific Coder subcommands require the agent token exposed!
-	cmd.Env = append(cmd.Env, fmt.Sprintf("CODER_AGENT_TOKEN=%s", *a.sessionToken.Load()))
+	cmd.Env = append(cmd.Env, fmt.Sprintf("CODER_AGENT_TOKEN=%s", s.AgentToken()))
 
 	// Set SSH connection environment variables (these are also set by OpenSSH
 	// and thus expected to be present by SSH clients). Since the agent does
@@ -427,13 +470,34 @@
 
 	// Agent-level environment variables should take over all!
 	// This is used for setting agent-specific variables like "CODER_AGENT_TOKEN".
-	for envKey, value := range a.envVars {
+	for envKey, value := range s.Env {
 		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", envKey, value))
 	}
 
 	return cmd, nil
 }
 
+func (s *Server) Serve(l net.Listener) error {
+	s.serveWg.Add(1)
+	defer s.serveWg.Done()
+	return s.srv.Serve(l)
+}
+
+func (s *Server) Close() error {
+	err := s.srv.Close()
+	s.serveWg.Wait()
+	return err
+}
+
+// Shutdown gracefully closes all active SSH connections and stops
+// accepting new connections.
+//
+// Shutdown is not implemented.
+func (*Server) Shutdown(ctx context.Context) error {
+	// TODO(mafredri): Implement shutdown, SIGHUP running commands, etc.
+	return nil
+}
+
 // isQuietLogin checks if the SSH server should perform a quiet login or not.
 //
 // https://github.com/openssh/openssh-portable/blob/25bd659cc72268f2858c5415740c442ee950049f/session.c#L816