feat(agent): Add SSH max timeout option for coder agent (#6596)

bensejas · web-flow · commit 7076dee52215 · 2023-03-15T09:08:50.000-05:00
* feat(agent): Add SSH max timeout option for coder agent

* Fix lint and update test golden snapshot
diff --git a/agent/agent.go b/agent/agent.go
@@ -78,6 +78,7 @@ type Options struct {
 	EnvironmentVariables   map[string]string
 	Logger                 slog.Logger
 	AgentPorts             map[int]string
+	SSHMaxTimeout          time.Duration
 }
 
 type Client interface {
@@ -126,6 +127,7 @@ func New(options Options) io.Closer {
 		lifecycleReported:      make(chan codersdk.WorkspaceAgentLifecycle, 1),
 		ignorePorts:            options.AgentPorts,
 		connStatsChan:          make(chan *agentsdk.Stats, 1),
+		sshMaxTimeout:          options.SSHMaxTimeout,
 	}
 	a.init(ctx)
 	return a
@@ -153,9 +155,10 @@ type agent struct {
 
 	envVars map[string]string
 	// metadata is atomic because values can change after reconnection.
-	metadata     atomic.Value
-	sessionToken atomic.Pointer[string]
-	sshServer    *ssh.Server
+	metadata      atomic.Value
+	sessionToken  atomic.Pointer[string]
+	sshServer     *ssh.Server
+	sshMaxTimeout time.Duration
 
 	lifecycleUpdate   chan struct{}
 	lifecycleReported chan codersdk.WorkspaceAgentLifecycle
@@ -780,6 +783,7 @@ func (a *agent) init(ctx context.Context) {
 				_ = session.Exit(1)
 			},
 		},
+		MaxTimeout: a.sshMaxTimeout,
 	}
 
 	go a.runLoop(ctx)
diff --git a/cli/agent.go b/cli/agent.go
@@ -31,10 +31,11 @@ import (
 
 func workspaceAgent() *cobra.Command {
 	var (
-		auth         string
-		logDir       string
-		pprofAddress string
-		noReap       bool
+		auth          string
+		logDir        string
+		pprofAddress  string
+		noReap        bool
+		sshMaxTimeout time.Duration
 	)
 	cmd := &cobra.Command{
 		Use: "agent",
@@ -208,7 +209,8 @@ func workspaceAgent() *cobra.Command {
 				EnvironmentVariables: map[string]string{
 					"GIT_ASKPASS": executablePath,
 				},
-				AgentPorts: agentPorts,
+				AgentPorts:    agentPorts,
+				SSHMaxTimeout: sshMaxTimeout,
 			})
 			<-ctx.Done()
 			return closer.Close()
@@ -219,6 +221,7 @@ func workspaceAgent() *cobra.Command {
 	cliflag.StringVarP(cmd.Flags(), &logDir, "log-dir", "", "CODER_AGENT_LOG_DIR", os.TempDir(), "Specify the location for the agent log files")
 	cliflag.StringVarP(cmd.Flags(), &pprofAddress, "pprof-address", "", "CODER_AGENT_PPROF_ADDRESS", "127.0.0.1:6060", "The address to serve pprof.")
 	cliflag.BoolVarP(cmd.Flags(), &noReap, "no-reap", "", "", false, "Do not start a process reaper.")
+	cliflag.DurationVarP(cmd.Flags(), &sshMaxTimeout, "ssh-max-timeout", "", "CODER_AGENT_SSH_MAX_TIMEOUT", time.Duration(0), "Specify the max timeout for a SSH connection")
 	return cmd
 }
 
diff --git a/cli/testdata/coder_agent_--help.golden b/cli/testdata/coder_agent_--help.golden
@@ -2,14 +2,16 @@ Usage:
   coder agent [flags]
 
 Flags:
-      --auth string            Specify the authentication type to use for the agent.
-                               Consumes $CODER_AGENT_AUTH (default "token")
-  -h, --help                   help for agent
-      --log-dir string         Specify the location for the agent log files.
-                               Consumes $CODER_AGENT_LOG_DIR (default "/tmp")
-      --no-reap                Do not start a process reaper.
-      --pprof-address string   The address to serve pprof.
-                               Consumes $CODER_AGENT_PPROF_ADDRESS (default "127.0.0.1:6060")
+      --auth string                Specify the authentication type to use for the agent.
+                                   Consumes $CODER_AGENT_AUTH (default "token")
+  -h, --help                       help for agent
+      --log-dir string             Specify the location for the agent log files.
+                                   Consumes $CODER_AGENT_LOG_DIR (default "/tmp")
+      --no-reap                    Do not start a process reaper.
+      --pprof-address string       The address to serve pprof.
+                                   Consumes $CODER_AGENT_PPROF_ADDRESS (default "127.0.0.1:6060")
+      --ssh-max-timeout duration   Specify the max timeout for a SSH connection.
+                                   Consumes $CODER_AGENT_SSH_MAX_TIMEOUT
 
 Global Flags:
       --global-config coder   Path to the global coder config directory.
