Track connections for SSH sessions to prevent leaks

kylecarbs · kylecarbs · commit 80f9cdaa1224 · 2023-02-03T15:28:03.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -645,15 +645,24 @@ func (a *agent) init(ctx context.Context) {
 			sshLogger.Info(ctx, "ssh connection ended", slog.Error(err))
 		},
 		Handler: func(session ssh.Session) {
-			err := a.handleSSHSession(session)
-			var exitError *exec.ExitError
-			if xerrors.As(err, &exitError) {
-				a.logger.Debug(ctx, "ssh session returned", slog.Error(exitError))
-				_ = session.Exit(exitError.ExitCode())
-				return
-			}
+			err := a.trackConnGoroutine(func() {
+				err := a.handleSSHSession(session)
+				var exitError *exec.ExitError
+				if xerrors.As(err, &exitError) {
+					a.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))
+					// This exit code is designed to be unlikely to be confused for a legit exit code
+					// from the process.
+					_ = session.Exit(MagicSessionErrorCode)
+					return
+				}
+			})
 			if err != nil {
-				a.logger.Warn(ctx, "ssh session failed", slog.Error(err))
+				a.logger.Warn(ctx, "track 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)
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -305,7 +305,7 @@ func TestAgent_TCPLocalForwarding(t *testing.T) {
 		}
 	}()
 
-	cmd := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%d:127.0.0.1:%d", randomPort, remotePort)}, []string{"sleep", "10"})
+	cmd := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%d:127.0.0.1:%d", randomPort, remotePort)}, []string{"sleep", "5"})
 	err = cmd.Start()
 	require.NoError(t, err)
 
@@ -372,7 +372,7 @@ func TestAgent_TCPRemoteForwarding(t *testing.T) {
 		}
 	}()
 
-	cmd := setupSSHCommand(t, []string{"-R", fmt.Sprintf("127.0.0.1:%d:127.0.0.1:%d", randomPort, localPort)}, []string{"sleep", "10"})
+	cmd := setupSSHCommand(t, []string{"-R", fmt.Sprintf("127.0.0.1:%d:127.0.0.1:%d", randomPort, localPort)}, []string{"sleep", "5"})
 	err = cmd.Start()
 	require.NoError(t, err)
 
@@ -437,7 +437,7 @@ func TestAgent_UnixLocalForwarding(t *testing.T) {
 		}
 	}()
 
-	cmd := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%s:%s", localSocketPath, remoteSocketPath)}, []string{"sleep", "10"})
+	cmd := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%s:%s", localSocketPath, remoteSocketPath)}, []string{"sleep", "5"})
 	err = cmd.Start()
 	require.NoError(t, err)
 
@@ -495,7 +495,7 @@ func TestAgent_UnixRemoteForwarding(t *testing.T) {
 		}
 	}()
 
-	cmd := setupSSHCommand(t, []string{"-R", fmt.Sprintf("%s:%s", remoteSocketPath, localSocketPath)}, []string{"sleep", "10"})
+	cmd := setupSSHCommand(t, []string{"-R", fmt.Sprintf("%s:%s", remoteSocketPath, localSocketPath)}, []string{"sleep", "5"})
 	err = cmd.Start()
 	require.NoError(t, err)
 
diff --git a/cli/root_test.go b/cli/root_test.go
@@ -24,10 +24,6 @@ import (
 	"github.com/coder/coder/testutil"
 )
 
-func init() {
-	// os.Setenv("NO_COLOR", "1")
-}
-
 // To update the golden files:
 // make update-golden-files
 var updateGoldenFiles = flag.Bool("update", false, "update .golden files")
