diff --git a/cmd/agentssh/main.go b/cmd/agentssh/main.go index b138e6f..d82f8bf 100644 --- a/cmd/agentssh/main.go +++ b/cmd/agentssh/main.go @@ -49,10 +49,16 @@ var newExecutor = func(cfg *config.Config) executor.Executor { if cfg != nil { options.HostKeyPolicy = cfg.Inventory.HostKeyPolicy options.PasswordSource = passwordSourceForRun(cfg.Paths) + options.DisableConnectionReuse = !sshMultiplexingEnabled(cfg) + options.KeepAliveInterval = sshKeepAliveInterval(cfg) } return executor.NewNativeExecutor(options) default: - return executor.NewSSHExecutor(nil) + return executor.NewSSHExecutorWithOptions(nil, executor.SSHOptions{ + DisableMultiplexing: !sshMultiplexingEnabled(cfg), + ControlPersist: sshControlPersist(cfg), + KeepAliveInterval: sshKeepAliveInterval(cfg), + }) } } @@ -976,6 +982,46 @@ func selectedTransport(cfg *config.Config) string { } } +func sshMultiplexingEnabled(cfg *config.Config) bool { + if cfg == nil { + return true + } + switch strings.ToLower(strings.TrimSpace(cfg.Inventory.SSH.Multiplexing)) { + case "", "on", "true", "yes", "1", "auto": + return true + case "off", "false", "no", "0", "disabled": + return false + default: + return true + } +} + +func sshControlPersist(cfg *config.Config) time.Duration { + if cfg == nil { + return 0 + } + return sshDuration(cfg.Inventory.SSH.ControlPersist) +} + +func sshKeepAliveInterval(cfg *config.Config) time.Duration { + if cfg == nil { + return 0 + } + return sshDuration(cfg.Inventory.SSH.KeepAliveInterval) +} + +func sshDuration(value string) time.Duration { + value = strings.TrimSpace(value) + if value == "" { + return 0 + } + duration, err := time.ParseDuration(value) + if err != nil { + return 0 + } + return duration +} + type inventoryAddOptions struct { Name string Addr string @@ -1549,6 +1595,7 @@ func runInventoryDiscover(cmd *cobra.Command, opts inventoryDiscoverOptions) err // will run (env-only master; no prompting across many candidates). PasswordSource: passwordSourceForRun(cfg.Paths), }) + defer func() { _ = exec.Close() }() result.Candidates = discovery.Probe(context.Background(), result.Candidates, discovery.ProbeOptions{ Executor: exec, Timeout: executor.ProbeTimeout, @@ -1667,6 +1714,7 @@ func runInventoryTest(cmd *cobra.Command, name string) error { // (env-only master). PasswordSource: passwordSourceForRun(cfg.Paths), }) + defer func() { _ = exec.Close() }() result := exec.Probe(context.Background(), resolved.Targets[0]) if result.Err == nil && result.ExitCode == 0 { refreshInventoryHostOS(cfg.Paths, name, result.OS) @@ -1756,6 +1804,7 @@ func runDirect(cmd *cobra.Command, targetName string, remoteCommand string, flag sessionStore := approval.SessionStore{Dir: cfg.Paths.SessionsDir} pendingStore := approval.PendingStore{PendingDir: cfg.Paths.PendingDir, ResponsesDir: cfg.Paths.ResponsesDir} ssh := newExecutor(cfg) + defer func() { _ = ssh.Close() }() plans, err := buildRunPlans(cfg, resolved, remoteCommand, flags, runtime, sessionStore, runtime.Enabled) if err != nil { return err diff --git a/cmd/agentssh/main_test.go b/cmd/agentssh/main_test.go index 2ef7a9f..bbb688f 100644 --- a/cmd/agentssh/main_test.go +++ b/cmd/agentssh/main_test.go @@ -111,9 +111,13 @@ func TestTransportSelection(t *testing.T) { } cfg.Inventory.Transport = executor.TransportShell shellExec := newExecutor(cfg) - if _, ok := shellExec.(executor.SSHExecutor); !ok { + shellSSH, ok := shellExec.(executor.SSHExecutor) + if !ok { t.Fatalf("inventory shell transport = %T, want shell-out SSHExecutor", shellExec) } + if shellSSH.Options.DisableMultiplexing { + t.Fatal("default shell executor disabled multiplexing") + } t.Setenv("AGENTSSH_TRANSPORT", executor.TransportShell) cfg.Inventory.Transport = executor.TransportNative envExec := newExecutor(cfg) @@ -123,6 +127,38 @@ func TestTransportSelection(t *testing.T) { t.Logf("transport default=%T inventory_shell=%T env_shell_overrides=%T", defaultExec, shellExec, envExec) } +func TestNewExecutorSSHReuseConfig(t *testing.T) { + cfg := &config.Config{} + cfg.Inventory.SSH.Multiplexing = "off" + cfg.Inventory.SSH.ControlPersist = "2m" + cfg.Inventory.SSH.KeepAliveInterval = "15s" + + nativeExec := newExecutor(cfg) + nativeSSH, ok := nativeExec.(executor.NativeExecutor) + if !ok { + t.Fatalf("native executor = %T", nativeExec) + } + if !nativeSSH.Options.DisableConnectionReuse { + t.Fatal("native DisableConnectionReuse = false, want true") + } + if nativeSSH.Options.KeepAliveInterval != 15*time.Second { + t.Fatalf("native keepalive = %s, want 15s", nativeSSH.Options.KeepAliveInterval) + } + + cfg.Inventory.Transport = executor.TransportShell + shellExec := newExecutor(cfg) + shellSSH, ok := shellExec.(executor.SSHExecutor) + if !ok { + t.Fatalf("shell executor = %T", shellExec) + } + if !shellSSH.Options.DisableMultiplexing { + t.Fatal("shell DisableMultiplexing = false, want true") + } + if shellSSH.Options.ControlPersist != 2*time.Minute || shellSSH.Options.KeepAliveInterval != 15*time.Second { + t.Fatalf("shell options = %#v", shellSSH.Options) + } +} + func TestPrintSSHExit255MappingDemo(t *testing.T) { result := executor.Result{ ExitCode: 255, @@ -2135,6 +2171,8 @@ func (e fakeExecutor) Run(_ context.Context, request executor.Request) executor. } } +func (e fakeExecutor) Close() error { return nil } + var _ executor.Executor = fakeExecutor{} func (e fakeExecutor) RunStreaming(_ context.Context, request executor.Request, stdout io.Writer, stderr io.Writer) executor.Result { @@ -2170,6 +2208,8 @@ func (e osDetectingExecutor) Run(_ context.Context, request executor.Request) ex } } +func (e osDetectingExecutor) Close() error { return nil } + var _ executor.Executor = osDetectingExecutor{} func writeInChunks(w io.Writer, data []byte, size int) { @@ -2208,6 +2248,8 @@ func (e bufferedOnlyExecutor) Run(_ context.Context, request executor.Request) e } } +func (e bufferedOnlyExecutor) Close() error { return nil } + var _ executor.Executor = bufferedOnlyExecutor{} type cliTestSSHServer struct { diff --git a/internal/config/config.go b/internal/config/config.go index 0baa42a..5e63927 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -28,6 +28,10 @@ const seedInventoryYAML = `# AgentSSH inventory — manage via 'agentssh tui' (H version: 1 transport: native # built-in Go SSH client; set "ssh" to shell out to system ssh host_key_policy: strict # or "accept-new" for trust-on-first-use +ssh: + multiplexing: on # set "off" to disable connection reuse + control_persist: 60s # shell transport ControlPersist lifetime + keepalive_interval: 30s # pooled/native heartbeat and shell ServerAliveInterval hosts: {} ` diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 790e270..5795583 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -91,6 +91,9 @@ func TestEnsureHomeCreatesAndSeeds(t *testing.T) { if cfg.Inventory.Transport != "native" { t.Fatalf("seeded transport = %q, want native", cfg.Inventory.Transport) } + if cfg.Inventory.SSH.Multiplexing != "on" || cfg.Inventory.SSH.ControlPersist != "60s" || cfg.Inventory.SSH.KeepAliveInterval != "30s" { + t.Fatalf("seeded ssh config = %#v", cfg.Inventory.SSH) + } if len(cfg.Policy.Rules) != 0 || len(cfg.Policy.HostOverrides) != 0 { t.Fatalf("seeded policy should have zero active rules: %+v", cfg.Policy) } diff --git a/internal/executor/executor.go b/internal/executor/executor.go index dd6731d..78f5145 100644 --- a/internal/executor/executor.go +++ b/internal/executor/executor.go @@ -3,13 +3,17 @@ package executor import ( "bytes" "context" + "crypto/sha256" + "encoding/hex" "errors" "fmt" "io" "os" "os/exec" + "path/filepath" "strconv" "strings" + "sync" "time" "github.com/Praeviso/AgentSSH/internal/inventory" @@ -35,6 +39,7 @@ type Result struct { // Executor runs commands against remote hosts. type Executor interface { Run(ctx context.Context, request Request) Result + Close() error } // StreamingExecutor optionally streams stdout/stderr into supplied writers. @@ -67,20 +72,57 @@ type RunResult struct { // SSHExecutor shells out to the system ssh binary. type SSHExecutor struct { - Runner Runner + Runner Runner + Options SSHOptions + mux *sshMultiplexer } +// SSHOptions configures the shell-out SSH transport. +type SSHOptions struct { + // DisableMultiplexing preserves the old one-process-per-command argv shape. + DisableMultiplexing bool + ControlPersist time.Duration + ControlDir string + KeepAliveInterval time.Duration +} + +const maxControlPathLen = 90 + // NewSSHExecutor returns an executor backed by the system ssh binary. func NewSSHExecutor(runner Runner) SSHExecutor { + return NewSSHExecutorWithOptions(runner, SSHOptions{}) +} + +// NewSSHExecutorWithOptions returns a shell-out executor with explicit SSH +// transport options. +func NewSSHExecutorWithOptions(runner Runner, options SSHOptions) SSHExecutor { if runner == nil { runner = ExecRunner{} } - return SSHExecutor{Runner: runner} + if options.ControlPersist == 0 { + options.ControlPersist = 60 * time.Second + } + exec := SSHExecutor{Runner: runner, Options: options} + if !options.DisableMultiplexing { + exec.mux = newSSHMultiplexer(options.ControlDir) + } + return exec +} + +// Close drops this process's in-memory mux cache. It intentionally leaves +// ControlPath sockets on disk: OpenSSH masters remove them when their +// ControlPersist lifetime ends, and keeping the stable path available is what +// lets later AgentSSH invocations reuse the master. +func (e SSHExecutor) Close() error { + if e.mux == nil { + return nil + } + return e.mux.Close() } func (e SSHExecutor) Run(ctx context.Context, request Request) Result { start := time.Now() - argv := BuildSSHArgv(request.Target, request.Command) + argv := e.buildArgv(request.Target, request.Command) runResult := e.Runner.Run(ctx, argv) result := Result{ Stdout: runResult.Stdout, @@ -93,6 +135,7 @@ func (e SSHExecutor) Run(ctx context.Context, request Request) Result { if sshResultConnected(result) { if request.Command == OSProbeCommand { result.OS = NormalizeOS(result.Stdout) + e.cacheOS(request.Target, result.OS) } else { result.OS = e.detectOS(ctx, request.Target) } @@ -102,7 +145,7 @@ func (e SSHExecutor) Run(ctx context.Context, request Request) Result { func (e SSHExecutor) RunStreaming(ctx context.Context, request Request, stdout io.Writer, stderr io.Writer) Result { start := time.Now() - argv := BuildSSHArgv(request.Target, request.Command) + argv := e.buildArgv(request.Target, request.Command) runResult := runStreamingProcess(ctx, argv, stdout, stderr) result := Result{ ExitCode: runResult.ExitCode, @@ -123,13 +166,50 @@ func sshResultConnected(result Result) bool { } func (e SSHExecutor) detectOS(ctx context.Context, target inventory.Target) string { + if osName := e.cachedOS(target); osName != "" { + return osName + } probeCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() - result := e.Runner.Run(probeCtx, BuildSSHArgv(target, OSProbeCommand)) + result := e.Runner.Run(probeCtx, e.buildArgv(target, OSProbeCommand)) if result.Err != nil || result.ExitCode != 0 { return "" } - return NormalizeOS(result.Stdout) + osName := NormalizeOS(result.Stdout) + e.cacheOS(target, osName) + return osName +} + +func (e SSHExecutor) buildArgv(target inventory.Target, command string) []string { + if e.mux == nil || e.Options.DisableMultiplexing { + return BuildSSHArgv(target, command) + } + if target.Host.SSHConfigAlias != "" { + return BuildSSHArgv(target, command) + } + controlPath := e.mux.controlPath(target) + if controlPath == "" { + return BuildSSHArgv(target, command) + } + return BuildSSHArgvWithOptions(target, command, SSHArgvOptions{ + ControlPath: controlPath, + ControlPersist: e.Options.ControlPersist, + KeepAliveInterval: e.Options.KeepAliveInterval, + }) +} + +func (e SSHExecutor) cachedOS(target inventory.Target) string { + if e.mux == nil || e.Options.DisableMultiplexing { + return "" + } + return e.mux.cachedOS(sshControlKey(target)) +} + +func (e SSHExecutor) cacheOS(target inventory.Target, osName string) { + if e.mux == nil || e.Options.DisableMultiplexing || osName == "" { + return + } + e.mux.setOS(sshControlKey(target), osName) } // OSProbeCommand is the remote command AgentSSH uses for internal host metadata @@ -161,6 +241,20 @@ func NormalizeOS(value string) string { // The remote command is appended as a single argv element; AgentSSH does not // locally shell-join or reinterpret it. func BuildSSHArgv(target inventory.Target, command string) []string { + return BuildSSHArgvWithOptions(target, command, SSHArgvOptions{}) +} + +// SSHArgvOptions carries local ssh(1) process options that AgentSSH owns. +type SSHArgvOptions struct { + ControlPath string + ControlPersist time.Duration + KeepAliveInterval time.Duration +} + +// BuildSSHArgvWithOptions constructs the local ssh(1) argv with optional +// ControlMaster multiplexing. The remote command is still appended as one argv +// element and is never locally shell-joined. +func BuildSSHArgvWithOptions(target inventory.Target, command string, options SSHArgvOptions) []string { host := target.Host hostSpec := host.SSHConfigAlias if hostSpec == "" { @@ -174,9 +268,167 @@ func BuildSSHArgv(target inventory.Target, command string) []string { if host.Port != 0 && host.Port != 22 && host.SSHConfigAlias == "" { argv = append(argv, "-p", strconv.Itoa(host.Port)) } + if options.ControlPath != "" && host.SSHConfigAlias == "" { + persist := options.ControlPersist + if persist <= 0 { + persist = 60 * time.Second + } + argv = append(argv, + "-o", "ControlMaster=auto", + "-o", "ControlPersist="+formatControlPersist(persist), + "-o", "ControlPath="+options.ControlPath, + ) + if options.KeepAliveInterval > 0 { + argv = append(argv, "-o", "ServerAliveInterval="+formatSeconds(options.KeepAliveInterval)) + } + } return append(argv, hostSpec, command) } +func formatControlPersist(duration time.Duration) string { + return formatControlSeconds(duration) +} + +func formatControlSeconds(duration time.Duration) string { + return formatSeconds(duration) + "s" +} + +func formatSeconds(duration time.Duration) string { + seconds := int(duration.Round(time.Second) / time.Second) + if seconds < 1 { + seconds = 1 + } + return strconv.Itoa(seconds) +} + +type sshMultiplexer struct { + mu sync.Mutex + dir string + configured string + paths map[string]struct{} + disabled map[string]struct{} + osByKey map[string]string +} + +func newSSHMultiplexer(controlDir string) *sshMultiplexer { + return &sshMultiplexer{ + configured: controlDir, + paths: map[string]struct{}{}, + disabled: map[string]struct{}{}, + osByKey: map[string]string{}, + } +} + +func (m *sshMultiplexer) controlPath(target inventory.Target) string { + key := sshControlKey(target) + m.mu.Lock() + if _, ok := m.disabled[key]; ok { + m.mu.Unlock() + return "" + } + m.mu.Unlock() + + dir := m.dirPath() + if dir == "" { + return "" + } + sum := sha256.Sum256([]byte(key)) + path := filepath.Join(dir, "cm-"+hex.EncodeToString(sum[:])[:24]) + if len(path) > maxControlPathLen { + m.mu.Lock() + m.disabled[key] = struct{}{} + m.mu.Unlock() + return "" + } + + m.mu.Lock() + m.paths[path] = struct{}{} + m.mu.Unlock() + + return path +} + +func (m *sshMultiplexer) dirPath() string { + m.mu.Lock() + defer m.mu.Unlock() + if m.dir != "" { + return m.dir + } + if m.configured != "" { + dir := filepath.Clean(m.configured) + if err := os.MkdirAll(dir, 0o700); err != nil { + return "" + } + m.dir = dir + return m.dir + } + if dir, err := defaultControlDir(); err == nil { + m.dir = dir + return m.dir + } + dir, err := os.MkdirTemp("", "agentssh-ssh-*") + if err != nil { + return "" + } + m.dir = dir + return m.dir +} + +func (m *sshMultiplexer) Close() error { + m.mu.Lock() + m.dir = "" + m.paths = map[string]struct{}{} + m.disabled = map[string]struct{}{} + m.osByKey = map[string]string{} + m.mu.Unlock() + + return nil +} + +func (m *sshMultiplexer) cachedOS(key string) string { + m.mu.Lock() + defer m.mu.Unlock() + return m.osByKey[key] +} + +func (m *sshMultiplexer) setOS(key string, osName string) { + m.mu.Lock() + defer m.mu.Unlock() + m.osByKey[key] = osName +} + +func defaultControlDir() (string, error) { + cacheDir, err := os.UserCacheDir() + if err != nil { + return "", err + } + dir := filepath.Join(cacheDir, "agentssh", "ssh-mux") + if err := os.MkdirAll(dir, 0o700); err != nil { + return "", err + } + if err := os.Chmod(dir, 0o700); err != nil { + return "", err + } + return dir, nil +} + +func sshControlKey(target inventory.Target) string { + host := target.Host + if host.SSHConfigAlias != "" { + return "alias\x00" + host.SSHConfigAlias + } + port := host.Port + if port == 0 { + port = 22 + } + return strings.Join([]string{ + "direct", + host.User, + strings.ToLower(host.Addr), + strconv.Itoa(port), + }, "\x00") +} + // ExecRunner executes argv directly with os/exec. type ExecRunner struct{} diff --git a/internal/executor/executor_test.go b/internal/executor/executor_test.go index 5c57dfe..2cd5ef1 100644 --- a/internal/executor/executor_test.go +++ b/internal/executor/executor_test.go @@ -5,8 +5,12 @@ import ( "context" "errors" "fmt" + "os" + "path/filepath" "reflect" + "strings" "testing" + "time" "github.com/Praeviso/AgentSSH/internal/inventory" ) @@ -32,6 +36,39 @@ func TestBuildSSHArgvKeepsRemoteCommandAsSingleArg(t *testing.T) { } } +func TestBuildSSHArgvWithControlMasterOptions(t *testing.T) { + target := inventory.Target{ + Name: "web-1", + Host: inventory.Host{ + Addr: "10.0.0.11", + User: "deploy", + Port: 2222, + }, + } + command := `printf 'a|b;$(whoami)' | sed 's/a/x/'` + controlPath := filepath.Join(t.TempDir(), "cm-web-1") + + got := BuildSSHArgvWithOptions(target, command, SSHArgvOptions{ + ControlPath: controlPath, + ControlPersist: 90 * time.Second, + }) + want := []string{ + "ssh", + "-p", "2222", + "-o", "ControlMaster=auto", + "-o", "ControlPersist=90s", + "-o", "ControlPath=" + controlPath, + "deploy@10.0.0.11", + command, + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("argv = %#v, want %#v", got, want) + } + if got[len(got)-1] != command { + t.Fatalf("remote command was changed: %q", got[len(got)-1]) + } +} + func TestBuildSSHArgvUsesSSHConfigAlias(t *testing.T) { target := inventory.Target{ Name: "web-1", @@ -59,7 +96,7 @@ func TestSSHExecutorUsesInjectedRunner(t *testing.T) { } return RunResult{Stdout: "ok\n", ExitCode: 0} }) - exec := NewSSHExecutor(runner) + exec := NewSSHExecutorWithOptions(runner, SSHOptions{DisableMultiplexing: true}) result := exec.Run(context.Background(), Request{ Target: inventory.Target{ @@ -91,6 +128,254 @@ func TestSSHExecutorUsesInjectedRunner(t *testing.T) { } } +func TestSSHExecutorMultiplexesCommandAndProbe(t *testing.T) { + controlDir := t.TempDir() + var calls [][]string + runner := RunnerFunc(func(_ context.Context, argv []string) RunResult { + calls = append(calls, append([]string{}, argv...)) + if argv[len(argv)-1] == OSProbeCommand { + return RunResult{Stdout: "Linux\n", ExitCode: 0} + } + return RunResult{Stdout: "ok\n", ExitCode: 0} + }) + exec := NewSSHExecutorWithOptions(runner, SSHOptions{ControlDir: controlDir, ControlPersist: 42 * time.Second}) + defer func() { _ = exec.Close() }() + + result := exec.Run(context.Background(), Request{ + Target: inventory.Target{ + Name: "web-1", + Host: inventory.Host{Addr: "10.0.0.11", User: "deploy"}, + }, + Command: "echo ok", + }) + + if len(calls) != 2 { + t.Fatalf("runner calls = %#v, want command + OS probe", calls) + } + commandPath := sshOptionValue(t, calls[0], "ControlPath") + probePath := sshOptionValue(t, calls[1], "ControlPath") + if commandPath == "" || probePath == "" || commandPath != probePath { + t.Fatalf("control paths command=%q probe=%q", commandPath, probePath) + } + if !strings.HasPrefix(commandPath, controlDir+string(os.PathSeparator)) { + t.Fatalf("control path %q is not under %q", commandPath, controlDir) + } + for _, argv := range calls { + if got := sshOptionValue(t, argv, "ControlMaster"); got != "auto" { + t.Fatalf("ControlMaster = %q, want auto in %#v", got, argv) + } + if got := sshOptionValue(t, argv, "ControlPersist"); got != "42s" { + t.Fatalf("ControlPersist = %q, want 42s in %#v", got, argv) + } + } + if result.Stdout != "ok\n" || result.ExitCode != 0 || result.Err != nil || result.OS != "linux" { + t.Fatalf("result = %#v", result) + } + if !reflect.DeepEqual(result.Argv, calls[0]) { + t.Fatalf("result argv = %#v, want first runner argv %#v", result.Argv, calls[0]) + } +} + +func TestSSHExecutorCachesOSPerMuxKey(t *testing.T) { + controlDir := t.TempDir() + var probes int + runner := RunnerFunc(func(_ context.Context, argv []string) RunResult { + if argv[len(argv)-1] == OSProbeCommand { + probes++ + return RunResult{Stdout: "Linux\n", ExitCode: 0} + } + return RunResult{Stdout: "ok\n", ExitCode: 0} + }) + exec := NewSSHExecutorWithOptions(runner, SSHOptions{ControlDir: controlDir}) + defer func() { _ = exec.Close() }() + target := inventory.Target{Name: "web-1", Host: inventory.Host{Addr: "10.0.0.11", User: "deploy"}} + + for i := 0; i < 2; i++ { + result := exec.Run(context.Background(), Request{Target: target, Command: "echo ok"}) + if result.Err != nil || result.OS != "linux" { + t.Fatalf("run %d result = %#v", i, result) + } + } + if probes != 1 { + t.Fatalf("OS probes = %d, want 1 cached probe", probes) + } +} + +func TestSSHExecutorLeavesAliasMuxOptionsToSSHConfig(t *testing.T) { + controlDir := t.TempDir() + var call []string + runner := RunnerFunc(func(_ context.Context, argv []string) RunResult { + call = append([]string{}, argv...) + return RunResult{ExitCode: 255, Err: errors.New("connect failed")} + }) + exec := NewSSHExecutorWithOptions(runner, SSHOptions{ControlDir: controlDir}) + defer func() { _ = exec.Close() }() + + _ = exec.Run(context.Background(), Request{ + Target: inventory.Target{Name: "prod", Host: inventory.Host{SSHConfigAlias: "prod-web", Addr: "ignored", User: "ignored", Port: 2222}}, + Command: "uptime", + }) + + if got := sshOptionValue(t, call, "ControlPath"); got != "" { + t.Fatalf("alias ControlPath = %q, want no injected mux options in %#v", got, call) + } + want := []string{"ssh", "prod-web", "uptime"} + if !reflect.DeepEqual(call, want) { + t.Fatalf("alias argv = %#v, want %#v", call, want) + } +} + +func TestBuildSSHArgvWithOptionsDoesNotOverrideAliasMuxConfig(t *testing.T) { + target := inventory.Target{Name: "prod", Host: inventory.Host{SSHConfigAlias: "prod-web"}} + got := BuildSSHArgvWithOptions(target, "uptime", SSHArgvOptions{ + ControlPath: filepath.Join(t.TempDir(), "cm-prod"), + ControlPersist: time.Minute, + }) + want := []string{"ssh", "prod-web", "uptime"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("argv = %#v, want %#v", got, want) + } +} + +func TestSSHExecutorStableControlDirSurvivesClose(t *testing.T) { + cacheRoot, err := os.MkdirTemp("/tmp", "agentssh-cache-*") + if err != nil { + t.Fatalf("create short cache dir: %v", err) + } + t.Cleanup(func() { _ = os.RemoveAll(cacheRoot) }) + t.Setenv("XDG_CACHE_HOME", cacheRoot) + var controlPath string + runner := RunnerFunc(func(_ context.Context, argv []string) RunResult { + if controlPath == "" { + controlPath = sshOptionValue(t, argv, "ControlPath") + if err := os.WriteFile(controlPath, []byte("socket placeholder"), 0o600); err != nil { + t.Fatalf("write control path placeholder: %v", err) + } + } + return RunResult{ExitCode: 255, Err: errors.New("connect failed")} + }) + exec := NewSSHExecutor(runner) + + _ = exec.Run(context.Background(), Request{ + Target: inventory.Target{Name: "web-1", Host: inventory.Host{Addr: "10.0.0.11", User: "deploy"}}, + Command: "uptime", + }) + if controlPath == "" { + t.Fatal("runner did not receive a ControlPath") + } + controlDir := filepath.Dir(controlPath) + wantDir := filepath.Join(cacheRoot, "agentssh", "ssh-mux") + if controlDir != wantDir { + t.Fatalf("control dir = %q, want %q", controlDir, wantDir) + } + if err := exec.Close(); err != nil { + t.Fatalf("close executor: %v", err) + } + if _, err := os.Stat(controlPath); err != nil { + t.Fatalf("control path after close: %v", err) + } +} + +func TestSSHExecutorFallsBackWhenControlPathTooLong(t *testing.T) { + controlDir := filepath.Join(t.TempDir(), strings.Repeat("long-segment-", 12)) + var call []string + runner := RunnerFunc(func(_ context.Context, argv []string) RunResult { + call = append([]string{}, argv...) + return RunResult{ExitCode: 255, Err: errors.New("connect failed")} + }) + exec := NewSSHExecutorWithOptions(runner, SSHOptions{ControlDir: controlDir}) + defer func() { _ = exec.Close() }() + + _ = exec.Run(context.Background(), Request{ + Target: inventory.Target{Name: "web-1", Host: inventory.Host{Addr: "10.0.0.11", User: "deploy"}}, + Command: "uptime", + }) + + if got := sshOptionValue(t, call, "ControlPath"); got != "" { + t.Fatalf("long ControlPath still injected as %q in %#v", got, call) + } + want := []string{"ssh", "deploy@10.0.0.11", "uptime"} + if !reflect.DeepEqual(call, want) { + t.Fatalf("fallback argv = %#v, want %#v", call, want) + } +} + +func TestSSHExecutorInjectsServerAliveIntervalWhenConfigured(t *testing.T) { + target := inventory.Target{Name: "web-1", Host: inventory.Host{Addr: "10.0.0.11", User: "deploy"}} + got := BuildSSHArgvWithOptions(target, "uptime", SSHArgvOptions{ + ControlPath: filepath.Join(t.TempDir(), "cm-web-1"), + ControlPersist: time.Minute, + KeepAliveInterval: 7 * time.Second, + }) + if value := sshOptionValue(t, got, "ServerAliveInterval"); value != "7" { + t.Fatalf("ServerAliveInterval = %q, want 7 in %#v", value, got) + } +} + +func TestSSHExecutorMultiplexingCanBeDisabled(t *testing.T) { + var calls [][]string + runner := RunnerFunc(func(_ context.Context, argv []string) RunResult { + calls = append(calls, append([]string{}, argv...)) + return RunResult{ExitCode: 255, Err: errors.New("connect failed")} + }) + exec := NewSSHExecutorWithOptions(runner, SSHOptions{DisableMultiplexing: true}) + + result := exec.Run(context.Background(), Request{ + Target: inventory.Target{Name: "web-1", Host: inventory.Host{Addr: "10.0.0.11", User: "deploy"}}, + Command: "uptime", + }) + + want := []string{"ssh", "deploy@10.0.0.11", "uptime"} + if result.Err == nil { + t.Fatal("result err = nil") + } + if len(calls) != 1 || !reflect.DeepEqual(calls[0], want) || !reflect.DeepEqual(result.Argv, want) { + t.Fatalf("calls=%#v result argv=%#v, want %#v", calls, result.Argv, want) + } +} + +func TestSSHExecutorCloseLeavesConfiguredControlPath(t *testing.T) { + var controlPath string + controlDir, err := os.MkdirTemp("/tmp", "agentssh-mux-*") + if err != nil { + t.Fatalf("create short control dir: %v", err) + } + t.Cleanup(func() { _ = os.RemoveAll(controlDir) }) + runner := RunnerFunc(func(_ context.Context, argv []string) RunResult { + if controlPath == "" { + controlPath = sshOptionValue(t, argv, "ControlPath") + if err := os.WriteFile(controlPath, []byte("socket placeholder"), 0o600); err != nil { + t.Fatalf("write control path placeholder: %v", err) + } + } + return RunResult{ExitCode: 255, Err: errors.New("connect failed")} + }) + exec := NewSSHExecutorWithOptions(runner, SSHOptions{ControlDir: controlDir}) + + result := exec.Run(context.Background(), Request{ + Target: inventory.Target{Name: "web-1", Host: inventory.Host{Addr: "10.0.0.11", User: "deploy"}}, + Command: "uptime", + }) + if result.Err == nil { + t.Fatal("result err = nil") + } + if controlPath == "" { + t.Fatal("runner did not receive a ControlPath") + } + if _, err := os.Stat(controlPath); err != nil { + t.Fatalf("control path placeholder before close: %v", err) + } + if err := exec.Close(); err != nil { + t.Fatalf("close executor: %v", err) + } + if _, err := os.Stat(controlPath); err != nil { + t.Fatalf("control path after close: %v", err) + } + if _, err := os.Stat(controlDir); err != nil { + t.Fatalf("control dir after close: %v", err) + } +} + func TestSSHExecutorDoesNotParseUserStdoutAsOS(t *testing.T) { var probeCalled bool runner := RunnerFunc(func(_ context.Context, argv []string) RunResult { @@ -100,7 +385,7 @@ func TestSSHExecutorDoesNotParseUserStdoutAsOS(t *testing.T) { } return RunResult{Stdout: "Linux\n", ExitCode: 0} }) - exec := NewSSHExecutor(runner) + exec := NewSSHExecutorWithOptions(runner, SSHOptions{DisableMultiplexing: true}) result := exec.Run(context.Background(), Request{ Target: inventory.Target{Name: "web-1", Host: inventory.Host{Addr: "10.0.0.11", User: "deploy"}}, Command: "echo Linux", @@ -189,7 +474,7 @@ func TestPrintArgvDemo(t *testing.T) { fmt.Printf("argv=%#v\n", argv) return RunResult{ExitCode: 0} }) - exec := NewSSHExecutor(runner) + exec := NewSSHExecutorWithOptions(runner, SSHOptions{DisableMultiplexing: true}) result := exec.Run(context.Background(), Request{ Target: inventory.Target{ Name: "web-1", @@ -201,3 +486,17 @@ func TestPrintArgvDemo(t *testing.T) { t.Fatalf("demo result = %#v", result) } } + +func sshOptionValue(t *testing.T, argv []string, name string) string { + t.Helper() + prefix := name + "=" + for i := 0; i+1 < len(argv); i++ { + if argv[i] != "-o" { + continue + } + if strings.HasPrefix(argv[i+1], prefix) { + return strings.TrimPrefix(argv[i+1], prefix) + } + } + return "" +} diff --git a/internal/executor/native.go b/internal/executor/native.go index 83e4823..7c8c752 100644 --- a/internal/executor/native.go +++ b/internal/executor/native.go @@ -13,6 +13,7 @@ import ( "path/filepath" "strconv" "strings" + "sync" "syscall" "time" @@ -30,6 +31,10 @@ const ( nativeArgv0 = "native-ssh" ) +const defaultNativeKeepAliveInterval = 30 * time.Second + +var nativeKeepAliveTimeout = 5 * time.Second + // ProbeTimeout bounds connectivity probes (inventory test, discover --probe, the // TUI t/probe actions). It matches the default Run connect budget so a probe // fails only when a real run would: a shorter cap produced false "cannot reach @@ -42,13 +47,19 @@ type NativeOptions struct { ConfigPath string KnownHostsPath string ConnectTimeout time.Duration - HostKeyPolicy string // strict | accept-new - PasswordSource func(host string) (string, bool) + // DisableConnectionReuse preserves the old one-dial-per-command behavior. + DisableConnectionReuse bool + // KeepAliveInterval controls pooled-connection heartbeats. Zero uses the + // default; a negative value disables keepalives for tests and diagnostics. + KeepAliveInterval time.Duration + HostKeyPolicy string // strict | accept-new + PasswordSource func(host string) (string, bool) } // NativeExecutor executes commands using golang.org/x/crypto/ssh. type NativeExecutor struct { Options NativeOptions + pool *nativeClientPool } // NewNativeExecutor returns a native SSH executor. @@ -59,23 +70,39 @@ func NewNativeExecutor(options NativeOptions) NativeExecutor { if options.HostKeyPolicy == "" { options.HostKeyPolicy = "strict" } - return NativeExecutor{Options: options} + var pool *nativeClientPool + if !options.DisableConnectionReuse { + pool = newNativeClientPool(options.KeepAliveInterval) + } + return NativeExecutor{ + Options: options, + pool: pool, + } +} + +// Close releases every pooled SSH connection held by this executor, including +// ProxyJump tunnel and jump-client closers owned by each connection. +func (e NativeExecutor) Close() error { + if e.pool == nil { + return nil + } + return e.pool.Close() } func (e NativeExecutor) Run(ctx context.Context, request Request) Result { var stdout bytes.Buffer var stderr bytes.Buffer - result := e.runSession(ctx, request, &stdout, &stderr) + result := e.runSession(ctx, request, &stdout, &stderr, true) result.Stdout = stdout.String() result.Stderr = stderr.String() return result } func (e NativeExecutor) RunStreaming(ctx context.Context, request Request, stdout io.Writer, stderr io.Writer) Result { - return e.runSession(ctx, request, stdout, stderr) + return e.runSession(ctx, request, stdout, stderr, true) } -func (e NativeExecutor) runSession(ctx context.Context, request Request, stdout io.Writer, stderr io.Writer) Result { +func (e NativeExecutor) runSession(ctx context.Context, request Request, stdout io.Writer, stderr io.Writer, usePool bool) Result { start := time.Now() target, err := e.resolveTarget(request.Target) targetDisplay := target.display() @@ -87,15 +114,62 @@ func (e NativeExecutor) runSession(ctx context.Context, request Request, stdout return Result{ExitCode: -1, Duration: time.Since(start), Err: err, Argv: argv} } - client, err := e.dial(ctx, target) + client, poolKey, poolEntry, pooled, err := e.client(ctx, target, usePool) if err != nil { return Result{ExitCode: -1, Duration: time.Since(start), Err: err, Argv: argv} } + closeClient := !pooled + poolReleased := false + releasePool := func() { + if pooled && poolEntry != nil && !poolReleased { + e.pool.release(poolEntry) + poolReleased = true + } + } defer func() { - _ = client.Close() + releasePool() + if closeClient && client != nil { + _ = client.Close() + } }() - session, err := client.NewSession() + closeCurrentClient := func() { + if pooled && poolEntry != nil { + e.pool.evictNow(poolKey, poolEntry, client) + return + } + _ = client.Close() + } + session, err := e.newSession(ctx, client, e.connectTimeout(), closeCurrentClient) + if err != nil && pooled && isSessionCapacityError(err) { + releasePool() + client, _, _, pooled, err = e.client(ctx, target, false) + closeClient = true + poolEntry = nil + poolKey = "" + poolReleased = false + closeCurrentClient = func() { _ = client.Close() } + if err == nil { + session, err = e.newSession(ctx, client, e.connectTimeout(), closeCurrentClient) + } + } + if err != nil && pooled && e.shouldReconnect(poolKey, poolEntry, client, err) { + e.pool.evictNow(poolKey, poolEntry, client) + releasePool() + client, poolKey, poolEntry, pooled, err = e.client(ctx, target, usePool) + closeClient = !pooled + poolReleased = false + closeCurrentClient = func() { + if pooled && poolEntry != nil { + e.pool.evictNow(poolKey, poolEntry, client) + return + } + _ = client.Close() + } + if err == nil { + session, err = e.newSession(ctx, client, e.connectTimeout(), closeCurrentClient) + } + } if err != nil { return Result{ExitCode: -1, Duration: time.Since(start), Err: err, Argv: argv} } @@ -106,9 +180,12 @@ func (e NativeExecutor) runSession(ctx context.Context, request Request, stdout session.Stdout = stdout session.Stderr = stderr - runErr := session.Run(request.Command) + runErr := runSSHSession(ctx, session, request.Command, closeCurrentClient) exitCode, err := nativeExitCode(runErr) if err != nil { + if pooled && isBrokenSSHClientError(err) { + e.pool.depoolWhenIdle(poolKey, poolEntry, client) + } return Result{ExitCode: -1, Duration: time.Since(start), Err: err, Argv: argv} } osName := "" @@ -118,15 +195,42 @@ func (e NativeExecutor) runSession(ctx context.Context, request Request, stdout } } if osName == "" { - osName = e.detectOS(ctx, client) + if pooled { + osName = e.pool.cachedOS(poolEntry) + } + if osName == "" { + osName = e.detectOS(ctx, client, closeCurrentClient) + } + } + if pooled && osName != "" { + e.pool.setOS(poolEntry, osName) } return Result{ExitCode: exitCode, Duration: time.Since(start), Argv: argv, OS: osName} } -func (e NativeExecutor) detectOS(ctx context.Context, client *nativeClient) string { +func (e NativeExecutor) client(ctx context.Context, target nativeTarget, usePool bool) (*nativeClient, string, *nativePoolEntry, bool, error) { + if e.pool == nil || !usePool { + client, err := e.dial(ctx, target) + return client, "", nil, false, err + } + key := target.poolKey() + client, entry, err := e.pool.get(ctx, key, func(context.Context) (*nativeClient, error) { + return e.dial(ctx, target) + }) + return client, key, entry, true, err +} + +func (e NativeExecutor) shouldReconnect(key string, entry *nativePoolEntry, client *nativeClient, err error) bool { + if e.pool == nil { + return false + } + return e.pool.isDead(key, entry, client) || isBrokenSSHClientError(err) +} + +func (e NativeExecutor) detectOS(ctx context.Context, client *nativeClient, closeClient func()) string { probeCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() - session, err := client.NewSession() + session, err := e.newSession(probeCtx, client, 5*time.Second, closeClient) if err != nil { return "" } @@ -135,20 +239,76 @@ func (e NativeExecutor) detectOS(ctx context.Context, client *nativeClient) stri }() var stdout bytes.Buffer session.Stdout = &stdout + if err := runSSHSession(probeCtx, session, OSProbeCommand, func() { + if closeClient != nil { + closeClient() + } + _ = session.Close() + }); err != nil { + return "" + } + return NormalizeOS(stdout.String()) +} + +func (e NativeExecutor) connectTimeout() time.Duration { + if e.Options.ConnectTimeout != 0 { + return e.Options.ConnectTimeout + } + return 10 * time.Second +} + +func (e NativeExecutor) newSession(ctx context.Context, client *nativeClient, timeout time.Duration, closeClient func()) (*ssh.Session, error) { + if timeout <= 0 { + timeout = 10 * time.Second + } + sessionCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + type sessionResult struct { + session *ssh.Session + err error + } + done := make(chan sessionResult) + go func() { + session, err := client.NewSession() + select { + case done <- sessionResult{session: session, err: err}: + case <-sessionCtx.Done(): + if session != nil { + _ = session.Close() + } + } + }() + select { + case result := <-done: + return result.session, result.err + case <-sessionCtx.Done(): + if closeClient != nil { + closeClient() + } + return nil, sessionCtx.Err() + } +} + +func runSSHSession(ctx context.Context, session *ssh.Session, command string, closeClient func()) error { done := make(chan error, 1) go func() { - done <- session.Run(OSProbeCommand) + done <- session.Run(command) }() select { case err := <-done: - if err != nil { - return "" + return err + case <-ctx.Done(): + if closeClient != nil { + closeClient() + } else { + _ = session.Close() } - case <-probeCtx.Done(): - _ = session.Close() - return "" + // Closing the transport unblocks session.Run, and Run does not return + // until its stdout/stderr copies finish. Wait for it so the caller can + // read the output buffers without racing those copies. + <-done + return ctx.Err() } - return NormalizeOS(stdout.String()) } type nativeTarget struct { @@ -181,6 +341,17 @@ func (t nativeTarget) display() string { return t.User + "@" + t.address() } +func (t nativeTarget) poolKey() string { + parts := []string{ + strings.ToLower(t.HostName), + strconv.Itoa(t.Port), + t.User, + strings.TrimSpace(t.ProxyJump), + } + parts = append(parts, t.IdentityFiles...) + return strings.Join(parts, "\x00") +} + func (e NativeExecutor) resolveTarget(target inventory.Target) (nativeTarget, error) { host := target.Host if host.SSHConfigAlias != "" { @@ -342,6 +513,316 @@ func (c *nativeClient) Close() error { return err } +type nativeClientPool struct { + mu sync.Mutex + entries map[string]*nativePoolEntry + keepAliveInterval time.Duration + closed bool +} + +type nativePoolEntry struct { + client *nativeClient + dialing bool + ready chan struct{} + readyOnce sync.Once + stop chan struct{} + stopOnce sync.Once + dialErr error + inFlight int + removed bool + closeWhenIdle bool + forceClosed bool + osName string +} + +func (e *nativePoolEntry) closeReady() { + e.readyOnce.Do(func() { + if e.ready != nil { + close(e.ready) + } + }) +} + +func (e *nativePoolEntry) closeStop() { + e.stopOnce.Do(func() { + if e.stop != nil { + close(e.stop) + } + }) +} + +func newNativeClientPool(keepAliveInterval time.Duration) *nativeClientPool { + if keepAliveInterval == 0 { + keepAliveInterval = defaultNativeKeepAliveInterval + } + return &nativeClientPool{ + entries: map[string]*nativePoolEntry{}, + keepAliveInterval: keepAliveInterval, + } +} + +func (p *nativeClientPool) get(ctx context.Context, key string, dial func(context.Context) (*nativeClient, error)) (*nativeClient, *nativePoolEntry, error) { + for { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return nil, nil, errors.New("native SSH executor is closed") + } + entry := p.entries[key] + if entry == nil { + entry = &nativePoolEntry{dialing: true, ready: make(chan struct{})} + p.entries[key] = entry + p.mu.Unlock() + return p.dialEntry(ctx, key, entry, dial) + } + if entry.client != nil { + client := entry.client + entry.inFlight++ + p.mu.Unlock() + return client, entry, nil + } + if entry.dialing { + ready := entry.ready + p.mu.Unlock() + select { + case <-ready: + p.mu.Lock() + dialErr := entry.dialErr + p.mu.Unlock() + if dialErr != nil { + return nil, nil, dialErr + } + continue + case <-ctx.Done(): + return nil, nil, ctx.Err() + } + } + p.mu.Unlock() + } +} + +func (p *nativeClientPool) dialEntry(ctx context.Context, key string, entry *nativePoolEntry, dial func(context.Context) (*nativeClient, error)) (*nativeClient, *nativePoolEntry, error) { + client, err := dial(ctx) + + p.mu.Lock() + defer p.mu.Unlock() + current := p.entries[key] + if current != entry { + if client != nil { + _ = client.Close() + } + if err != nil { + return nil, nil, err + } + return nil, nil, errors.New("native SSH pooled connection was replaced while dialing") + } + entry.dialing = false + if err != nil { + entry.dialErr = err + delete(p.entries, key) + entry.closeReady() + return nil, nil, err + } + if p.closed { + delete(p.entries, key) + entry.closeReady() + _ = client.Close() + return nil, nil, errors.New("native SSH executor is closed") + } + entry.client = client + entry.inFlight = 1 + entry.stop = make(chan struct{}) + entry.closeReady() + p.startKeepAliveLocked(key, entry) + return client, entry, nil +} + +func (p *nativeClientPool) startKeepAliveLocked(key string, entry *nativePoolEntry) { + if p.keepAliveInterval < 0 || entry.client == nil || entry.stop == nil { + return + } + client := entry.client + stop := entry.stop + interval := p.keepAliveInterval + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + if err := sendKeepAlive(client, nativeKeepAliveTimeout); err != nil { + p.evictNow(key, entry, client) + return + } + case <-stop: + return + } + } + }() +} + +func (p *nativeClientPool) isDead(key string, entry *nativePoolEntry, client *nativeClient) bool { + p.mu.Lock() + defer p.mu.Unlock() + current := p.entries[key] + return current != entry || entry == nil || entry.client != client || entry.forceClosed +} + +func (p *nativeClientPool) release(entry *nativePoolEntry) { + var closeClient *nativeClient + var closeStop *nativePoolEntry + + p.mu.Lock() + if entry != nil && entry.inFlight > 0 { + entry.inFlight-- + if entry.inFlight == 0 && entry.closeWhenIdle && !entry.forceClosed { + closeClient = entry.client + closeStop = entry + entry.forceClosed = true + } + } + p.mu.Unlock() + + if closeStop != nil { + closeStop.closeStop() + } + if closeClient != nil { + _ = closeClient.Close() + } +} + +func (p *nativeClientPool) cachedOS(entry *nativePoolEntry) string { + p.mu.Lock() + defer p.mu.Unlock() + if entry == nil || entry.forceClosed { + return "" + } + return entry.osName +} + +func (p *nativeClientPool) setOS(entry *nativePoolEntry, osName string) { + p.mu.Lock() + defer p.mu.Unlock() + if entry != nil && !entry.forceClosed { + entry.osName = osName + } +} + +func (p *nativeClientPool) depoolWhenIdle(key string, entry *nativePoolEntry, client *nativeClient) { + var closeClient *nativeClient + var closeStop *nativePoolEntry + + p.mu.Lock() + if entry != nil && entry.client == client && !entry.forceClosed { + if p.entries[key] == entry { + delete(p.entries, key) + } + entry.removed = true + entry.closeWhenIdle = true + if entry.inFlight == 0 { + closeClient = entry.client + closeStop = entry + entry.forceClosed = true + } + } + p.mu.Unlock() + + if closeStop != nil { + closeStop.closeStop() + } + if closeClient != nil { + _ = closeClient.Close() + } +} + +func (p *nativeClientPool) evictNow(key string, entry *nativePoolEntry, client *nativeClient) { + var closeClient *nativeClient + var closeStop *nativePoolEntry + + p.mu.Lock() + if entry != nil && entry.client == client && !entry.forceClosed { + if p.entries[key] == entry { + delete(p.entries, key) + } + entry.removed = true + entry.closeWhenIdle = true + entry.forceClosed = true + closeClient = entry.client + closeStop = entry + } + p.mu.Unlock() + + if closeStop != nil { + closeStop.closeStop() + } + if closeClient != nil { + _ = closeClient.Close() + } +} + +func sendKeepAlive(client *nativeClient, timeout time.Duration) error { + if timeout <= 0 { + timeout = 5 * time.Second + } + done := make(chan error, 1) + go func() { + _, _, err := client.SendRequest("keepalive@openssh.com", true, nil) + done <- err + }() + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case err := <-done: + return err + case <-timer.C: + return context.DeadlineExceeded + } +} + +func (p *nativeClientPool) Close() error { + if p == nil { + return nil + } + + var clients []*nativeClient + var stops []*nativePoolEntry + var ready []*nativePoolEntry + + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return nil + } + p.closed = true + for key, entry := range p.entries { + delete(p.entries, key) + if entry.stop != nil { + stops = append(stops, entry) + } + if entry.dialing && entry.ready != nil { + ready = append(ready, entry) + } + if entry.client != nil { + clients = append(clients, entry.client) + } + } + p.mu.Unlock() + + for _, entry := range stops { + entry.closeStop() + } + for _, entry := range ready { + entry.closeReady() + } + var err error + for _, client := range clients { + if closeErr := client.Close(); closeErr != nil && err == nil { + err = closeErr + } + } + return err +} + func (e NativeExecutor) dialTCP(ctx context.Context, addr string) (net.Conn, error) { timeout := e.Options.ConnectTimeout if timeout == 0 { @@ -410,7 +891,12 @@ func (e NativeExecutor) clientConfig(target nativeTarget) (*ssh.ClientConfig, io // Probe opens a native SSH session and runs a no-op command. It is intended for // explicit operator diagnostics; callers must gate it behind a user action. func (e NativeExecutor) Probe(ctx context.Context, target inventory.Target) Result { - return e.Run(ctx, Request{Target: target, Command: OSProbeCommand}) + var stdout bytes.Buffer + var stderr bytes.Buffer + result := e.runSession(ctx, Request{Target: target, Command: OSProbeCommand}, &stdout, &stderr, false) + result.Stdout = stdout.String() + result.Stderr = stderr.String() + return result } // ProbeStatusForError maps a transport error into the probe status vocabulary. @@ -470,6 +956,30 @@ func isConnectionRefused(err error) bool { return errors.Is(err, syscall.ECONNREFUSED) || strings.Contains(strings.ToLower(err.Error()), "connection refused") } +func isBrokenSSHClientError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || errors.Is(err, syscall.EPIPE) || errors.Is(err, syscall.ECONNRESET) || errors.Is(err, syscall.ECONNABORTED) { + return true + } + message := strings.ToLower(err.Error()) + return strings.Contains(message, "use of closed network connection") || + strings.Contains(message, "broken pipe") || + strings.Contains(message, "connection reset") || + strings.Contains(message, "connection closed") || + strings.Contains(message, "unexpected packet in response to channel open: ") +} + +func isSessionCapacityError(err error) bool { + var openErr *ssh.OpenChannelError + if errors.As(err, &openErr) { + return openErr.Reason == ssh.Prohibited || openErr.Reason == ssh.ResourceShortage + } + message := strings.ToLower(err.Error()) + return strings.Contains(message, "administratively prohibited") || strings.Contains(message, "resource shortage") +} + func (e NativeExecutor) authMethods(target nativeTarget) ([]ssh.AuthMethod, io.Closer, error) { var methods []ssh.AuthMethod var agentConn io.Closer diff --git a/internal/executor/native_test.go b/internal/executor/native_test.go index b598c2e..2b99fcc 100644 --- a/internal/executor/native_test.go +++ b/internal/executor/native_test.go @@ -10,6 +10,7 @@ import ( "encoding/pem" "errors" "fmt" + "io" "net" "os" "path/filepath" @@ -18,6 +19,7 @@ import ( "sync" "sync/atomic" "testing" + "time" "github.com/Praeviso/AgentSSH/internal/inventory" "golang.org/x/crypto/ssh" @@ -36,6 +38,7 @@ func TestNativeExecutorExecSuccessAndRemoteNonZero(t *testing.T) { KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), ConfigPath: filepath.Join(home, ".ssh", "config"), }) + defer func() { _ = exec.Close() }() target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) @@ -69,6 +72,7 @@ func TestNativeExecutorRunStreaming(t *testing.T) { KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), ConfigPath: filepath.Join(home, ".ssh", "config"), }) + defer func() { _ = exec.Close() }() target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} var stdout bytes.Buffer @@ -86,6 +90,386 @@ func TestNativeExecutorRunStreaming(t *testing.T) { t.Logf("native streaming stdout=%q stderr=%q exit_code=%d", stdout.String(), stderr.String(), result.ExitCode) } +func TestNativeExecutorReusesPooledConnection(t *testing.T) { + home := t.TempDir() + clientSigner := writeClientKey(t, home) + server := newTestSSHServer(t, clientSigner.PublicKey()) + defer server.Close() + + writeKnownHosts(t, home, server.Addr(), server.HostSigner.PublicKey()) + + exec := NewNativeExecutor(NativeOptions{ + KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), + ConfigPath: filepath.Join(home, ".ssh", "config"), + KeepAliveInterval: -1, + }) + defer func() { _ = exec.Close() }() + target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} + + for _, command := range []string{"ok", "exit7"} { + result := exec.Run(context.Background(), Request{Target: target, Command: command}) + if result.Err != nil { + t.Fatalf("%s result err = %v", command, result.Err) + } + } + if got := atomic.LoadInt32(&server.accepted); got != 1 { + t.Fatalf("accepted connections = %d, want 1 pooled connection", got) + } +} + +func TestNativeExecutorCachesOSPerPooledConnection(t *testing.T) { + home := t.TempDir() + clientSigner := writeClientKey(t, home) + server := newTestSSHServer(t, clientSigner.PublicKey()) + defer server.Close() + + writeKnownHosts(t, home, server.Addr(), server.HostSigner.PublicKey()) + + exec := NewNativeExecutor(NativeOptions{ + KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), + ConfigPath: filepath.Join(home, ".ssh", "config"), + KeepAliveInterval: -1, + }) + defer func() { _ = exec.Close() }() + target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} + + for i := 0; i < 2; i++ { + result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) + if result.Err != nil || result.OS != "linux" { + t.Fatalf("run %d result = %#v err=%v", i, result, result.Err) + } + } + if got := atomic.LoadInt32(&server.osProbes); got != 1 { + t.Fatalf("OS probes = %d, want 1 cached probe", got) + } +} + +func TestNativeExecutorPoolSingleFlightsConcurrentDial(t *testing.T) { + home := t.TempDir() + clientSigner := writeClientKey(t, home) + server := newTestSSHServer(t, clientSigner.PublicKey()) + defer server.Close() + + writeKnownHosts(t, home, server.Addr(), server.HostSigner.PublicKey()) + + exec := NewNativeExecutor(NativeOptions{ + KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), + ConfigPath: filepath.Join(home, ".ssh", "config"), + KeepAliveInterval: -1, + }) + defer func() { _ = exec.Close() }() + target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} + + var wg sync.WaitGroup + errs := make(chan error, 8) + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) + if result.Err != nil || result.ExitCode != 0 || result.Stdout != "ok\n" { + errs <- fmt.Errorf("result = %#v err=%v", result, result.Err) + } + }() + } + wg.Wait() + close(errs) + for err := range errs { + t.Fatal(err) + } + if got := atomic.LoadInt32(&server.accepted); got != 1 { + t.Fatalf("accepted connections = %d, want 1 pooled connection", got) + } +} + +func TestNativeClientPoolBroadcastsDialErrorToWaiters(t *testing.T) { + pool := newNativeClientPool(-1) + wantErr := errors.New("dial failed") + dialEntered := make(chan struct{}) + releaseDial := make(chan struct{}) + var dials int32 + dial := func(context.Context) (*nativeClient, error) { + if atomic.AddInt32(&dials, 1) == 1 { + close(dialEntered) + } + <-releaseDial + return nil, wantErr + } + + errs := make(chan error, 6) + go func() { + _, _, err := pool.get(context.Background(), "web-1", dial) + errs <- err + }() + <-dialEntered + for i := 0; i < 5; i++ { + go func() { + _, _, err := pool.get(context.Background(), "web-1", dial) + errs <- err + }() + } + time.Sleep(20 * time.Millisecond) + close(releaseDial) + + for i := 0; i < 6; i++ { + if err := <-errs; !errors.Is(err, wantErr) { + t.Fatalf("waiter %d err = %v, want %v", i, err, wantErr) + } + } + if got := atomic.LoadInt32(&dials); got != 1 { + t.Fatalf("dials = %d, want one broadcast failed dial", got) + } +} + +func TestNativeClientPoolDepoolDefersCloseUntilInFlightDrains(t *testing.T) { + pool := newNativeClientPool(-1) + var closes int32 + entry := &nativePoolEntry{ + client: &nativeClient{closers: []io.Closer{ + closerFunc(func() error { + atomic.AddInt32(&closes, 1) + return nil + }), + }}, + stop: make(chan struct{}), + inFlight: 2, + } + pool.entries["web-1"] = entry + + pool.depoolWhenIdle("web-1", entry, entry.client) + if _, ok := pool.entries["web-1"]; ok { + t.Fatal("entry still pooled after suspicious run error") + } + if got := atomic.LoadInt32(&closes); got != 0 { + t.Fatalf("close count after depool = %d, want 0 while sessions are in flight", got) + } + pool.release(entry) + if got := atomic.LoadInt32(&closes); got != 0 { + t.Fatalf("close count after first release = %d, want 0", got) + } + pool.release(entry) + if got := atomic.LoadInt32(&closes); got != 1 { + t.Fatalf("close count after final release = %d, want 1", got) + } +} + +func TestNativeExecutorReconnectsWhenPooledClientDies(t *testing.T) { + home := t.TempDir() + clientSigner := writeClientKey(t, home) + server := newTestSSHServer(t, clientSigner.PublicKey()) + defer server.Close() + + writeKnownHosts(t, home, server.Addr(), server.HostSigner.PublicKey()) + + exec := NewNativeExecutor(NativeOptions{ + KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), + ConfigPath: filepath.Join(home, ".ssh", "config"), + KeepAliveInterval: -1, + }) + defer func() { _ = exec.Close() }() + target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} + + result := exec.Run(context.Background(), Request{Target: target, Command: "close-connection"}) + if result.Err != nil || result.ExitCode != 0 { + t.Fatalf("close-connection result = %#v err=%v", result, result.Err) + } + waitForAtomic(t, &server.closed, 1) + + result = exec.Run(context.Background(), Request{Target: target, Command: "ok"}) + if result.Err != nil || result.ExitCode != 0 || result.Stdout != "ok\n" { + t.Fatalf("reconnected result = %#v err=%v", result, result.Err) + } + if got := atomic.LoadInt32(&server.accepted); got != 2 { + t.Fatalf("accepted connections = %d, want reconnect to dial once", got) + } +} + +func TestNativeExecutorFallsBackToUnpooledWhenMaxSessionsRejects(t *testing.T) { + home := t.TempDir() + clientSigner := writeClientKey(t, home) + server := newTestSSHServer(t, clientSigner.PublicKey()) + server.maxSessions = 1 + server.blockCh = make(chan struct{}) + defer server.Close() + + writeKnownHosts(t, home, server.Addr(), server.HostSigner.PublicKey()) + + exec := NewNativeExecutor(NativeOptions{ + KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), + ConfigPath: filepath.Join(home, ".ssh", "config"), + KeepAliveInterval: -1, + }) + defer func() { _ = exec.Close() }() + target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} + + firstDone := make(chan Result, 1) + go func() { + firstDone <- exec.Run(context.Background(), Request{Target: target, Command: "block"}) + }() + waitForAtomic(t, &server.activeSessions, 1) + + second := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) + if second.Err != nil || second.ExitCode != 0 || second.Stdout != "ok\n" { + t.Fatalf("fallback result = %#v err=%v", second, second.Err) + } + if got := atomic.LoadInt32(&server.sessionRejects); got == 0 { + t.Fatal("pooled channel-open was not rejected by the session cap") + } + + close(server.blockCh) + first := <-firstDone + if first.Err != nil || first.ExitCode != 0 { + t.Fatalf("first pooled result = %#v err=%v", first, first.Err) + } + third := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) + if third.Err != nil || third.ExitCode != 0 { + t.Fatalf("third pooled result = %#v err=%v", third, third.Err) + } + if got := atomic.LoadInt32(&server.accepted); got != 2 { + t.Fatalf("accepted connections = %d, want pooled connection plus one unpooled fallback", got) + } +} + +func TestNativeExecutorNewSessionTimesOutAndClosesPooledClient(t *testing.T) { + home := t.TempDir() + clientSigner := writeClientKey(t, home) + server := newTestSSHServer(t, clientSigner.PublicKey()) + server.hangChannels = make(chan struct{}) + defer server.Close() + + writeKnownHosts(t, home, server.Addr(), server.HostSigner.PublicKey()) + + exec := NewNativeExecutor(NativeOptions{ + KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), + ConfigPath: filepath.Join(home, ".ssh", "config"), + ConnectTimeout: 30 * time.Millisecond, + KeepAliveInterval: -1, + }) + defer func() { _ = exec.Close() }() + target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} + + start := time.Now() + result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) + close(server.hangChannels) + if !errors.Is(result.Err, context.DeadlineExceeded) { + t.Fatalf("result err = %v, want context deadline exceeded", result.Err) + } + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("NewSession timeout took %s", elapsed) + } + waitForAtomic(t, &server.closed, 1) +} + +func TestNativeExecutorRunHonorsContextWithoutRunTimeout(t *testing.T) { + home := t.TempDir() + clientSigner := writeClientKey(t, home) + server := newTestSSHServer(t, clientSigner.PublicKey()) + server.blockCh = make(chan struct{}) + defer server.Close() + + writeKnownHosts(t, home, server.Addr(), server.HostSigner.PublicKey()) + + exec := NewNativeExecutor(NativeOptions{ + KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), + ConfigPath: filepath.Join(home, ".ssh", "config"), + KeepAliveInterval: -1, + }) + defer func() { _ = exec.Close() }() + target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + + result := exec.Run(ctx, Request{Target: target, Command: "block"}) + close(server.blockCh) + if !errors.Is(result.Err, context.DeadlineExceeded) { + t.Fatalf("result err = %v, want context deadline exceeded", result.Err) + } + waitForAtomic(t, &server.closed, 1) +} + +func TestNativeExecutorKeepAliveAndClose(t *testing.T) { + home := t.TempDir() + clientSigner := writeClientKey(t, home) + server := newTestSSHServer(t, clientSigner.PublicKey()) + defer server.Close() + + writeKnownHosts(t, home, server.Addr(), server.HostSigner.PublicKey()) + + exec := NewNativeExecutor(NativeOptions{ + KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), + ConfigPath: filepath.Join(home, ".ssh", "config"), + KeepAliveInterval: 5 * time.Millisecond, + }) + target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} + + result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) + if result.Err != nil || result.ExitCode != 0 { + t.Fatalf("run result = %#v err=%v", result, result.Err) + } + waitForAtomic(t, &server.globalRequests, 1) + + if err := exec.Close(); err != nil { + t.Fatalf("close executor: %v", err) + } + waitForAtomic(t, &server.closed, 1) +} + +func TestNativeExecutorKeepAliveTimeoutEvictsAndCloses(t *testing.T) { + home := t.TempDir() + clientSigner := writeClientKey(t, home) + server := newTestSSHServer(t, clientSigner.PublicKey()) + server.hangGlobalRequests = make(chan struct{}) + defer server.Close() + + writeKnownHosts(t, home, server.Addr(), server.HostSigner.PublicKey()) + + oldTimeout := nativeKeepAliveTimeout + nativeKeepAliveTimeout = 20 * time.Millisecond + defer func() { nativeKeepAliveTimeout = oldTimeout }() + + exec := NewNativeExecutor(NativeOptions{ + KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), + ConfigPath: filepath.Join(home, ".ssh", "config"), + KeepAliveInterval: 10 * time.Millisecond, + }) + defer func() { _ = exec.Close() }() + target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} + + result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) + if result.Err != nil || result.ExitCode != 0 { + t.Fatalf("run result = %#v err=%v", result, result.Err) + } + waitForAtomic(t, &server.globalRequests, 1) + waitForAtomic(t, &server.closed, 1) + close(server.hangGlobalRequests) +} + +func TestNativeClientPoolCloseClosesCloserChain(t *testing.T) { + pool := newNativeClientPool(-1) + var first int32 + var second int32 + pool.entries["web-1"] = &nativePoolEntry{ + client: &nativeClient{closers: []io.Closer{ + closerFunc(func() error { + atomic.AddInt32(&first, 1) + return nil + }), + closerFunc(func() error { + atomic.AddInt32(&second, 1) + return nil + }), + }}, + stop: make(chan struct{}), + } + + if err := pool.Close(); err != nil { + t.Fatalf("pool close: %v", err) + } + if atomic.LoadInt32(&first) != 1 || atomic.LoadInt32(&second) != 1 { + t.Fatalf("closer counts first=%d second=%d", first, second) + } +} + func TestNativeExecutorHostKeyRejected(t *testing.T) { home := t.TempDir() clientSigner := writeClientKey(t, home) @@ -106,6 +490,7 @@ func TestNativeExecutorHostKeyRejected(t *testing.T) { KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), ConfigPath: filepath.Join(home, ".ssh", "config"), }) + defer func() { _ = exec.Close() }() target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) if result.Err == nil || result.ExitCode != -1 { @@ -126,6 +511,7 @@ func TestNativeExecutorAcceptNewTrustsUnknownHost(t *testing.T) { ConfigPath: filepath.Join(home, ".ssh", "config"), HostKeyPolicy: "accept-new", }) + defer func() { _ = exec.Close() }() target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) if result.Err != nil || result.ExitCode != 0 { @@ -144,6 +530,7 @@ func TestNativeExecutorAcceptNewTrustsUnknownHost(t *testing.T) { KnownHostsPath: knownHostsPath, ConfigPath: filepath.Join(home, ".ssh", "config"), }) + defer func() { _ = strictExec.Close() }() result = strictExec.Run(context.Background(), Request{Target: target, Command: "ok"}) if result.Err != nil || result.ExitCode != 0 { t.Fatalf("strict result after accept-new = %#v err=%v", result, result.Err) @@ -208,6 +595,7 @@ func TestNativeExecutorUsesPerHostIdentityBeforeDefault(t *testing.T) { KnownHostsPath: filepath.Join(home, ".ssh", "known_hosts"), ConfigPath: filepath.Join(home, ".ssh", "config"), }) + defer func() { _ = exec.Close() }() target := inventory.Target{Name: "test", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test", IdentityFile: "~/keys/web-1"}} result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) if result.Err != nil || result.ExitCode != 0 { @@ -282,6 +670,7 @@ func TestNativeExecutorPasswordAuthEndToEnd(t *testing.T) { return "secretpw", true }, }) + defer func() { _ = exec.Close() }() target := inventory.Target{Name: "web-1", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) if result.Err != nil || result.ExitCode != 0 || result.Stdout != "ok\n" { @@ -304,6 +693,7 @@ func TestNativeExecutorPrefersKeyAuthBeforePassword(t *testing.T) { ConfigPath: filepath.Join(home, ".ssh", "config"), PasswordSource: func(string) (string, bool) { return "secretpw", true }, }) + defer func() { _ = exec.Close() }() target := inventory.Target{Name: "web-1", Host: inventory.Host{Addr: server.Host(), Port: server.Port(), User: "test"}} result := exec.Run(context.Background(), Request{Target: target, Command: "ok"}) if result.Err != nil || result.ExitCode != 0 { @@ -339,10 +729,26 @@ func TestConnectHintMapping(t *testing.T) { } type testSSHServer struct { - Listener net.Listener - HostSigner ssh.Signer - wg sync.WaitGroup - allowedKey ssh.PublicKey + Listener net.Listener + HostSigner ssh.Signer + wg sync.WaitGroup + allowedKey ssh.PublicKey + accepted int32 + closed int32 + globalRequests int32 + osProbes int32 + activeSessions int32 + sessionRejects int32 + maxSessions int32 + blockCh chan struct{} + hangChannels chan struct{} + hangGlobalRequests chan struct{} +} + +type closerFunc func() error + +func (fn closerFunc) Close() error { + return fn() } type passwordTestSSHServer struct { @@ -441,7 +847,7 @@ func (s *passwordTestSSHServer) handle(conn net.Conn) { if err != nil { continue } - go handleSession(channel, requests) + go handleSession(sshConn, channel, requests, nil) } } @@ -498,6 +904,7 @@ func (s *testSSHServer) accept() { func (s *testSSHServer) handle(conn net.Conn) { defer s.wg.Done() + atomic.AddInt32(&s.accepted, 1) cfg := &ssh.ServerConfig{ PublicKeyCallback: func(_ ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) { if string(key.Marshal()) == string(s.allowedKey.Marshal()) { @@ -512,22 +919,66 @@ func (s *testSSHServer) handle(conn net.Conn) { _ = conn.Close() return } - defer func() { _ = sshConn.Close() }() - go ssh.DiscardRequests(reqs) + defer func() { + atomic.AddInt32(&s.closed, 1) + _ = sshConn.Close() + }() + go s.handleGlobalRequests(reqs) + var connectionSessions int32 for ch := range chans { + if s.hangChannels != nil { + <-s.hangChannels + } if ch.ChannelType() != "session" { _ = ch.Reject(ssh.UnknownChannelType, "session only") continue } + if maxSessions := atomic.LoadInt32(&s.maxSessions); maxSessions > 0 && atomic.LoadInt32(&connectionSessions) >= maxSessions { + atomic.AddInt32(&s.sessionRejects, 1) + _ = ch.Reject(ssh.Prohibited, "session limit reached") + continue + } channel, requests, err := ch.Accept() if err != nil { continue } - go handleSession(channel, requests) + atomic.AddInt32(&connectionSessions, 1) + atomic.AddInt32(&s.activeSessions, 1) + go func() { + defer func() { + atomic.AddInt32(&connectionSessions, -1) + atomic.AddInt32(&s.activeSessions, -1) + }() + handleSession(sshConn, channel, requests, s) + }() + } +} + +func (s *testSSHServer) handleGlobalRequests(requests <-chan *ssh.Request) { + for req := range requests { + atomic.AddInt32(&s.globalRequests, 1) + if s.hangGlobalRequests != nil { + <-s.hangGlobalRequests + } + if req.WantReply { + _ = req.Reply(true, nil) + } } } -func handleSession(channel ssh.Channel, requests <-chan *ssh.Request) { +func waitForAtomic(t *testing.T, value *int32, want int32) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if got := atomic.LoadInt32(value); got >= want { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("atomic value = %d, want at least %d", atomic.LoadInt32(value), want) +} + +func handleSession(conn *ssh.ServerConn, channel ssh.Channel, requests <-chan *ssh.Request, server *testSSHServer) { defer func() { _ = channel.Close() }() for req := range requests { if req.Type != "exec" { @@ -539,12 +990,27 @@ func handleSession(channel ssh.Channel, requests <-chan *ssh.Request) { _ = req.Reply(true, nil) code := uint32(0) if payload.Command == OSProbeCommand { + if server != nil { + atomic.AddInt32(&server.osProbes, 1) + } _, _ = channel.Write([]byte("Linux\n")) + } else if strings.Contains(payload.Command, "block") { + if server != nil && server.blockCh != nil { + <-server.blockCh + } + _, _ = channel.Write([]byte("ok\n")) } else if strings.Contains(payload.Command, "stream-secret") { _, _ = channel.Write([]byte("line1\n")) _, _ = channel.Write([]byte("password=")) _, _ = channel.Write([]byte("secret123\nline3\n")) _, _ = channel.Stderr().Write([]byte("warn=secret\n")) + } else if strings.Contains(payload.Command, "close-connection") { + _, _ = channel.Write([]byte("ok\n")) + _, _ = channel.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{Status: code})) + if conn != nil { + go func() { _ = conn.Close() }() + } + return } else if strings.Contains(payload.Command, "exit7") { code = 7 _, _ = channel.Stderr().Write([]byte("failed\n")) diff --git a/internal/inventory/inventory.go b/internal/inventory/inventory.go index aa94497..08b6693 100644 --- a/internal/inventory/inventory.go +++ b/internal/inventory/inventory.go @@ -12,10 +12,29 @@ type Inventory struct { Version int `yaml:"version,omitempty" json:"version"` Transport string `yaml:"transport,omitempty" json:"transport"` HostKeyPolicy string `yaml:"host_key_policy,omitempty" json:"host_key_policy"` + SSH SSHConfig `yaml:"ssh,omitempty" json:"ssh,omitempty"` Hosts map[string]Host `yaml:"hosts,omitempty" json:"hosts"` Groups map[string]Group `yaml:"groups,omitempty" json:"groups"` } +// SSHConfig contains local SSH transport tuning. It never stores credentials; +// the fields only decide how AgentSSH reuses already-authenticated connections. +type SSHConfig struct { + // Multiplexing is "on" (default) or "off". For shell transport it controls + // OpenSSH ControlMaster injection; for native transport it controls pooled + // client reuse. + Multiplexing string `yaml:"multiplexing,omitempty" json:"multiplexing,omitempty"` + ControlPersist string `yaml:"control_persist,omitempty" json:"control_persist,omitempty"` + KeepAliveInterval string `yaml:"keepalive_interval,omitempty" json:"keepalive_interval,omitempty"` +} + +// IsZero lets yaml/json omit the ssh section when every value is defaulted. +func (c SSHConfig) IsZero() bool { + return strings.TrimSpace(c.Multiplexing) == "" && + strings.TrimSpace(c.ControlPersist) == "" && + strings.TrimSpace(c.KeepAliveInterval) == "" +} + // Host describes a named SSH target without storing private keys or passwords. type Host struct { Addr string `yaml:"addr,omitempty" json:"addr"`