From a78a02f125b5a69f909277df1fe566d41f98546d Mon Sep 17 00:00:00 2001 From: Kritoooo Date: Wed, 1 Jul 2026 19:51:53 -0700 Subject: [PATCH 1/2] feat(executor): reuse SSH connections instead of dialing per command Both transports previously opened a fresh connection (dial + handshake + auth) for every command and closed it immediately. Add connection reuse. Native transport (golang.org/x/crypto/ssh): - Concurrency-safe per-host connection pool keyed by resolved target shape (host, port, user, ProxyJump, identity files); concurrent dials for the same key are single-flighted, different keys proceed independently. - Pooled clients keep the existing nativeClient closer chain, so direct clients, ProxyJump tunnels, and jump clients still close in order. - keepalive@openssh.com heartbeat per pooled connection (default 30s); a failed heartbeat evicts and closes the client. - Transparent one-shot reconnect when a stale pooled client fails NewSession. Never retries after session.Run has started, preserving remote side-effect and exit-code semantics. - NativeExecutor.Close() drains the pool; Probe stays one-shot so diagnostics never leave a long-lived socket. Shell transport (system ssh binary): - Enable OpenSSH multiplexing by default: ControlMaster=auto, ControlPersist=60s, stable per-target ControlPath under a 0700 per-run temp dir; command and hidden OS probe share the socket. - SSHOptions allows disabling or customizing persist/dir; BuildSSHArgv stays backward-compatible, multiplexing rides on BuildSSHArgvWithOptions. - SSHExecutor.Close() removes the owned temp dir / tracked sockets. Tests cover pool reuse, concurrent single-flight dialing, reconnect on a dead pooled client, keepalive+close, closer-chain shutdown, ControlMaster argv construction, shared command/probe socket, disable path, and temp-dir cleanup. gofmt/go build/go vet/go test all pass. Implemented via Codex. Co-Authored-By: Claude Opus 4.8 (1M context) --- cmd/agentssh/main.go | 5 + internal/executor/executor.go | 188 ++++++++++++++++++- internal/executor/executor_test.go | 162 +++++++++++++++- internal/executor/native.go | 290 +++++++++++++++++++++++++++-- internal/executor/native_test.go | 218 +++++++++++++++++++++- 5 files changed, 835 insertions(+), 28 deletions(-) diff --git a/cmd/agentssh/main.go b/cmd/agentssh/main.go index b138e6f..c1d357e 100644 --- a/cmd/agentssh/main.go +++ b/cmd/agentssh/main.go @@ -1549,6 +1549,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 +1668,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 +1758,9 @@ 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) + if closer, ok := ssh.(interface{ Close() error }); ok { + defer func() { _ = closer.Close() }() + } plans, err := buildRunPlans(cfg, resolved, remoteCommand, flags, runtime, sessionStore, runtime.Enabled) if err != nil { return err diff --git a/internal/executor/executor.go b/internal/executor/executor.go index dd6731d..22bb6a6 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" @@ -67,20 +71,53 @@ 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 } // 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 removes the per-run ControlPath directory used by shell transport +// multiplexing. Any OpenSSH master process left by ControlPersist exits on its +// own bounded timer, but the socket path is no longer exposed to future runs. +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, @@ -102,7 +139,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, @@ -125,13 +162,27 @@ func sshResultConnected(result Result) bool { func (e SSHExecutor) detectOS(ctx context.Context, target inventory.Target) string { 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) } +func (e SSHExecutor) buildArgv(target inventory.Target, command string) []string { + if e.mux == nil || e.Options.DisableMultiplexing { + 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, + }) +} + // OSProbeCommand is the remote command AgentSSH uses for internal host metadata // refreshes. Callers must never mix its stdout into agent-visible command output. const OSProbeCommand = "uname -s" @@ -161,6 +212,19 @@ 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 +} + +// 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 +238,123 @@ 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 != "" { + persist := options.ControlPersist + if persist <= 0 { + persist = 60 * time.Second + } + argv = append(argv, + "-o", "ControlMaster=auto", + "-o", "ControlPersist="+formatControlPersist(persist), + "-o", "ControlPath="+options.ControlPath, + ) + } return append(argv, hostSpec, command) } +func formatControlPersist(duration time.Duration) string { + seconds := int(duration.Round(time.Second) / time.Second) + if seconds < 1 { + seconds = 1 + } + return strconv.Itoa(seconds) + "s" +} + +type sshMultiplexer struct { + mu sync.Mutex + dir string + configured string + owned bool + paths map[string]struct{} +} + +func newSSHMultiplexer(controlDir string) *sshMultiplexer { + return &sshMultiplexer{configured: controlDir, paths: map[string]struct{}{}} +} + +func (m *sshMultiplexer) controlPath(target inventory.Target) string { + dir := m.dirPath() + if dir == "" { + return "" + } + sum := sha256.Sum256([]byte(sshControlKey(target))) + path := filepath.Join(dir, "cm-"+hex.EncodeToString(sum[:])[:24]) + + 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 + } + dir, err := os.MkdirTemp("", "agentssh-ssh-*") + if err != nil { + return "" + } + m.dir = dir + m.owned = true + return m.dir +} + +func (m *sshMultiplexer) Close() error { + m.mu.Lock() + dir := m.dir + owned := m.owned + paths := make([]string, 0, len(m.paths)) + for path := range m.paths { + paths = append(paths, path) + } + m.dir = "" + m.owned = false + m.paths = map[string]struct{}{} + m.mu.Unlock() + + if dir == "" { + return nil + } + if owned { + return os.RemoveAll(dir) + } + var err error + for _, path := range paths { + if removeErr := os.Remove(path); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) && err == nil { + err = removeErr + } + } + return err +} + +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..eda9b44 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,111 @@ 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 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 TestSSHExecutorCloseCleansOwnedControlDir(t *testing.T) { + 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) + + 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") + } + controlDir := filepath.Dir(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(controlDir); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("control dir after close err=%v, want not exist", err) + } +} + func TestSSHExecutorDoesNotParseUserStdoutAsOS(t *testing.T) { var probeCalled bool runner := RunnerFunc(func(_ context.Context, argv []string) RunResult { @@ -100,7 +242,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 +331,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 +343,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..74cf733 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,8 @@ const ( nativeArgv0 = "native-ssh" ) +const defaultNativeKeepAliveInterval = 30 * 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 +45,17 @@ type NativeOptions struct { ConfigPath string KnownHostsPath string ConnectTimeout time.Duration - HostKeyPolicy string // strict | accept-new - PasswordSource func(host string) (string, 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 +66,35 @@ func NewNativeExecutor(options NativeOptions) NativeExecutor { if options.HostKeyPolicy == "" { options.HostKeyPolicy = "strict" } - return NativeExecutor{Options: options} + return NativeExecutor{ + Options: options, + pool: newNativeClientPool(options.KeepAliveInterval), + } +} + +// 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 +106,24 @@ 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, pooled, err := e.client(ctx, target, usePool) if err != nil { return Result{ExitCode: -1, Duration: time.Since(start), Err: err, Argv: argv} } - defer func() { - _ = client.Close() - }() + if !pooled { + defer func() { + _ = client.Close() + }() + } session, err := client.NewSession() + if err != nil && pooled && e.shouldReconnect(poolKey, client, err) { + e.pool.evict(poolKey, client) + client, _, _, err = e.client(ctx, target, usePool) + if err == nil { + session, err = client.NewSession() + } + } if err != nil { return Result{ExitCode: -1, Duration: time.Since(start), Err: err, Argv: argv} } @@ -109,6 +137,9 @@ func (e NativeExecutor) runSession(ctx context.Context, request Request, stdout runErr := session.Run(request.Command) exitCode, err := nativeExitCode(runErr) if err != nil { + if pooled && isBrokenSSHClientError(err) { + e.pool.evict(poolKey, client) + } return Result{ExitCode: -1, Duration: time.Since(start), Err: err, Argv: argv} } osName := "" @@ -123,6 +154,25 @@ func (e NativeExecutor) runSession(ctx context.Context, request Request, stdout return Result{ExitCode: exitCode, Duration: time.Since(start), Argv: argv, OS: osName} } +func (e NativeExecutor) client(ctx context.Context, target nativeTarget, usePool bool) (*nativeClient, string, bool, error) { + if e.pool == nil || !usePool { + client, err := e.dial(ctx, target) + return client, "", false, err + } + key := target.poolKey() + client, err := e.pool.get(ctx, key, func(context.Context) (*nativeClient, error) { + return e.dial(ctx, target) + }) + return client, key, true, err +} + +func (e NativeExecutor) shouldReconnect(key string, client *nativeClient, err error) bool { + if e.pool == nil { + return false + } + return e.pool.isDead(key, client) || isBrokenSSHClientError(err) +} + func (e NativeExecutor) detectOS(ctx context.Context, client *nativeClient) string { probeCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() @@ -181,6 +231,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 +403,193 @@ 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{} + stop chan struct{} +} + +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, error) { + for { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return 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 + p.mu.Unlock() + return client, nil + } + if entry.dialing { + ready := entry.ready + p.mu.Unlock() + select { + case <-ready: + continue + case <-ctx.Done(): + return nil, ctx.Err() + } + } + p.mu.Unlock() + } +} + +func (p *nativeClientPool) dialEntry(ctx context.Context, key string, entry *nativePoolEntry, dial func(context.Context) (*nativeClient, error)) (*nativeClient, 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, err + } + return nil, errors.New("native SSH pooled connection was replaced while dialing") + } + entry.dialing = false + if err != nil { + delete(p.entries, key) + close(entry.ready) + return nil, err + } + if p.closed { + delete(p.entries, key) + close(entry.ready) + _ = client.Close() + return nil, errors.New("native SSH executor is closed") + } + entry.client = client + entry.stop = make(chan struct{}) + close(entry.ready) + p.startKeepAliveLocked(key, entry) + return client, 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 := client.SendRequest("keepalive@openssh.com", true, nil); err != nil { + p.evict(key, client) + return + } + case <-stop: + return + } + } + }() +} + +func (p *nativeClientPool) isDead(key string, client *nativeClient) bool { + p.mu.Lock() + defer p.mu.Unlock() + entry := p.entries[key] + return entry == nil || entry.client != client +} + +func (p *nativeClientPool) evict(key string, client *nativeClient) { + var closeClient *nativeClient + var stop chan struct{} + + p.mu.Lock() + entry := p.entries[key] + if entry != nil && entry.client == client { + delete(p.entries, key) + closeClient = entry.client + stop = entry.stop + } + p.mu.Unlock() + + if stop != nil { + close(stop) + } + if closeClient != nil { + _ = closeClient.Close() + } +} + +func (p *nativeClientPool) Close() error { + if p == nil { + return nil + } + + var clients []*nativeClient + var stops []chan struct{} + var ready []chan struct{} + + 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.stop) + } + if entry.dialing && entry.ready != nil { + ready = append(ready, entry.ready) + } + if entry.client != nil { + clients = append(clients, entry.client) + } + } + p.mu.Unlock() + + for _, stop := range stops { + close(stop) + } + for _, ch := range ready { + close(ch) + } + 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 +658,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 +723,21 @@ 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 (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..946b805 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,155 @@ 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 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 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 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 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 +259,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 +280,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 +299,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 +364,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 +439,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 +462,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 +498,19 @@ 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 +} + +type closerFunc func() error + +func (fn closerFunc) Close() error { + return fn() } type passwordTestSSHServer struct { @@ -441,7 +609,7 @@ func (s *passwordTestSSHServer) handle(conn net.Conn) { if err != nil { continue } - go handleSession(channel, requests) + go handleSession(sshConn, channel, requests) } } @@ -498,6 +666,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,8 +681,11 @@ 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) for ch := range chans { if ch.ChannelType() != "session" { _ = ch.Reject(ssh.UnknownChannelType, "session only") @@ -523,11 +695,32 @@ func (s *testSSHServer) handle(conn net.Conn) { if err != nil { continue } - go handleSession(channel, requests) + go handleSession(sshConn, channel, requests) } } -func handleSession(channel ssh.Channel, requests <-chan *ssh.Request) { +func (s *testSSHServer) handleGlobalRequests(requests <-chan *ssh.Request) { + for req := range requests { + atomic.AddInt32(&s.globalRequests, 1) + if req.WantReply { + _ = req.Reply(true, nil) + } + } +} + +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) { defer func() { _ = channel.Close() }() for req := range requests { if req.Type != "exec" { @@ -545,6 +738,13 @@ func handleSession(channel ssh.Channel, requests <-chan *ssh.Request) { _, _ = 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")) From fc081e4c75af4137ef638b7e513c516efe93bf4c Mon Sep 17 00:00:00 2001 From: Kritoooo Date: Thu, 2 Jul 2026 00:09:46 -0700 Subject: [PATCH 2/2] fix(executor): resolve 10 code-review findings in connection reuse Correctness (native pool): - Bound pooled NewSession with ConnectTimeout and make session.Run ctx-aware (close transport on cancel, no command timeout), so a half-open pooled connection fails fast instead of hanging until kernel TCP timeout. Wait for the Run goroutine after cancel so output buffers are never read while its copies still write. - Run keepalive SendRequest under a short timeout; timeout/error evicts and closes, restoring dead-connection detection. - Track in-flight sessions per pool entry: suspicious post-Run errors de-pool immediately but defer Close until in-flight sessions drain, so one command's channel-level error no longer kills concurrent sessions. Provably-dead paths still close immediately. - Broadcast dial errors to all waiters via entry.dialErr, so N concurrent commands to an unreachable host fail in parallel. - Fall back to a one-off unpooled dial when NewSession is rejected with Prohibited/ResourceShortage (sshd MaxSessions), without evicting the healthy shared client. Shell transport (multiplexing): - Use a stable per-user control dir (UserCacheDir/agentssh/ssh-mux, 0700); Close no longer removes live sockets, so masters are actually reused across invocations instead of leaking orphaned authenticated masters with unlinked sockets. - Do not inject ControlMaster/ControlPath/ControlPersist for ssh_config alias targets - the user's ssh_config governs. - Fall back to the non-multiplexed argv when the ControlPath would exceed ~90 bytes (sun_path limit), caching the decision. Config & lifecycle: - Expose ssh.multiplexing / ssh.control_persist / ssh.keepalive_interval in config and wire through newExecutor. - Cache the normalized OS per host (pool entry / mux key) instead of re-probing uname -s on every command. - Add Close() error to the Executor interface; drop the ad-hoc io.Closer type assertion in runDirect. All fixes verified with gofmt, go vet, and go test -count=1 -race ./... (full suite green). Implemented via Codex; the ctx-cancel buffer race was found by the race detector and fixed on top. Co-Authored-By: Claude Opus 4.8 (1M context) --- cmd/agentssh/main.go | 52 +++- cmd/agentssh/main_test.go | 44 +++- internal/config/config.go | 4 + internal/config/config_test.go | 3 + internal/executor/executor.go | 136 ++++++++--- internal/executor/executor_test.go | 153 +++++++++++- internal/executor/native.go | 376 ++++++++++++++++++++++++----- internal/executor/native_test.go | 286 +++++++++++++++++++++- internal/inventory/inventory.go | 19 ++ 9 files changed, 955 insertions(+), 118 deletions(-) diff --git a/cmd/agentssh/main.go b/cmd/agentssh/main.go index c1d357e..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 @@ -1758,9 +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) - if closer, ok := ssh.(interface{ Close() error }); ok { - defer func() { _ = closer.Close() }() - } + 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 22bb6a6..78f5145 100644 --- a/internal/executor/executor.go +++ b/internal/executor/executor.go @@ -39,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. @@ -82,8 +83,11 @@ type SSHOptions struct { 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{}) @@ -105,9 +109,10 @@ func NewSSHExecutorWithOptions(runner Runner, options SSHOptions) SSHExecutor { return exec } -// Close removes the per-run ControlPath directory used by shell transport -// multiplexing. Any OpenSSH master process left by ControlPersist exits on its -// own bounded timer, but the socket path is no longer exposed to future runs. +// 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 @@ -130,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) } @@ -160,29 +166,52 @@ 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, 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, + 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 // refreshes. Callers must never mix its stdout into agent-visible command output. const OSProbeCommand = "uname -s" @@ -217,8 +246,9 @@ func BuildSSHArgv(target inventory.Target, command string) []string { // SSHArgvOptions carries local ssh(1) process options that AgentSSH owns. type SSHArgvOptions struct { - ControlPath string - ControlPersist time.Duration + ControlPath string + ControlPersist time.Duration + KeepAliveInterval time.Duration } // BuildSSHArgvWithOptions constructs the local ssh(1) argv with optional @@ -238,7 +268,7 @@ func BuildSSHArgvWithOptions(target inventory.Target, command string, options SS if host.Port != 0 && host.Port != 22 && host.SSHConfigAlias == "" { argv = append(argv, "-p", strconv.Itoa(host.Port)) } - if options.ControlPath != "" { + if options.ControlPath != "" && host.SSHConfigAlias == "" { persist := options.ControlPersist if persist <= 0 { persist = 60 * time.Second @@ -248,37 +278,68 @@ func BuildSSHArgvWithOptions(target inventory.Target, command string, options SS "-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) + "s" + return strconv.Itoa(seconds) } type sshMultiplexer struct { mu sync.Mutex dir string configured string - owned bool 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{}{}} + 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(sshControlKey(target))) + 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{}{} @@ -301,41 +362,54 @@ func (m *sshMultiplexer) dirPath() string { 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 - m.owned = true return m.dir } func (m *sshMultiplexer) Close() error { m.mu.Lock() - dir := m.dir - owned := m.owned - paths := make([]string, 0, len(m.paths)) - for path := range m.paths { - paths = append(paths, path) - } m.dir = "" - m.owned = false m.paths = map[string]struct{}{} + m.disabled = map[string]struct{}{} + m.osByKey = map[string]string{} m.mu.Unlock() - if dir == "" { - return nil + 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 } - if owned { - return os.RemoveAll(dir) + dir := filepath.Join(cacheDir, "agentssh", "ssh-mux") + if err := os.MkdirAll(dir, 0o700); err != nil { + return "", err } - var err error - for _, path := range paths { - if removeErr := os.Remove(path); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) && err == nil { - err = removeErr - } + if err := os.Chmod(dir, 0o700); err != nil { + return "", err } - return err + return dir, nil } func sshControlKey(target inventory.Target) string { diff --git a/internal/executor/executor_test.go b/internal/executor/executor_test.go index eda9b44..2cd5ef1 100644 --- a/internal/executor/executor_test.go +++ b/internal/executor/executor_test.go @@ -176,6 +176,142 @@ func TestSSHExecutorMultiplexesCommandAndProbe(t *testing.T) { } } +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 { @@ -198,8 +334,13 @@ func TestSSHExecutorMultiplexingCanBeDisabled(t *testing.T) { } } -func TestSSHExecutorCloseCleansOwnedControlDir(t *testing.T) { +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") @@ -209,7 +350,7 @@ func TestSSHExecutorCloseCleansOwnedControlDir(t *testing.T) { } return RunResult{ExitCode: 255, Err: errors.New("connect failed")} }) - exec := NewSSHExecutor(runner) + 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"}}, @@ -221,15 +362,17 @@ func TestSSHExecutorCloseCleansOwnedControlDir(t *testing.T) { if controlPath == "" { t.Fatal("runner did not receive a ControlPath") } - controlDir := filepath.Dir(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(controlDir); !errors.Is(err, os.ErrNotExist) { - t.Fatalf("control dir after close err=%v, want not exist", 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) } } diff --git a/internal/executor/native.go b/internal/executor/native.go index 74cf733..7c8c752 100644 --- a/internal/executor/native.go +++ b/internal/executor/native.go @@ -33,6 +33,8 @@ const ( 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 @@ -45,6 +47,8 @@ type NativeOptions struct { ConfigPath string KnownHostsPath string ConnectTimeout time.Duration + // 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 @@ -66,9 +70,13 @@ func NewNativeExecutor(options NativeOptions) NativeExecutor { if options.HostKeyPolicy == "" { options.HostKeyPolicy = "strict" } + var pool *nativeClientPool + if !options.DisableConnectionReuse { + pool = newNativeClientPool(options.KeepAliveInterval) + } return NativeExecutor{ Options: options, - pool: newNativeClientPool(options.KeepAliveInterval), + pool: pool, } } @@ -106,22 +114,60 @@ func (e NativeExecutor) runSession(ctx context.Context, request Request, stdout return Result{ExitCode: -1, Duration: time.Since(start), Err: err, Argv: argv} } - client, poolKey, pooled, err := e.client(ctx, target, usePool) + 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} } - if !pooled { - defer func() { - _ = client.Close() - }() + closeClient := !pooled + poolReleased := false + releasePool := func() { + if pooled && poolEntry != nil && !poolReleased { + e.pool.release(poolEntry) + poolReleased = true + } } + defer func() { + releasePool() + if closeClient && client != nil { + _ = client.Close() + } + }() - session, err := client.NewSession() - if err != nil && pooled && e.shouldReconnect(poolKey, client, err) { - e.pool.evict(poolKey, client) - client, _, _, err = e.client(ctx, target, usePool) + 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 = client.NewSession() + session, err = e.newSession(ctx, client, e.connectTimeout(), closeCurrentClient) } } if err != nil { @@ -134,11 +180,11 @@ 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.evict(poolKey, client) + e.pool.depoolWhenIdle(poolKey, poolEntry, client) } return Result{ExitCode: -1, Duration: time.Since(start), Err: err, Argv: argv} } @@ -149,34 +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) client(ctx context.Context, target nativeTarget, usePool bool) (*nativeClient, string, bool, error) { +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, "", false, err + return client, "", nil, false, err } key := target.poolKey() - client, err := e.pool.get(ctx, key, func(context.Context) (*nativeClient, error) { + client, entry, err := e.pool.get(ctx, key, func(context.Context) (*nativeClient, error) { return e.dial(ctx, target) }) - return client, key, true, err + return client, key, entry, true, err } -func (e NativeExecutor) shouldReconnect(key string, client *nativeClient, err error) bool { +func (e NativeExecutor) shouldReconnect(key string, entry *nativePoolEntry, client *nativeClient, err error) bool { if e.pool == nil { return false } - return e.pool.isDead(key, client) || isBrokenSSHClientError(err) + return e.pool.isDead(key, entry, client) || isBrokenSSHClientError(err) } -func (e NativeExecutor) detectOS(ctx context.Context, client *nativeClient) string { +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 "" } @@ -185,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 { @@ -411,10 +521,34 @@ type nativeClientPool struct { } type nativePoolEntry struct { - client *nativeClient - dialing bool - ready chan struct{} - stop chan 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 { @@ -427,12 +561,12 @@ func newNativeClientPool(keepAliveInterval time.Duration) *nativeClientPool { } } -func (p *nativeClientPool) get(ctx context.Context, key string, dial func(context.Context) (*nativeClient, error)) (*nativeClient, error) { +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, errors.New("native SSH executor is closed") + return nil, nil, errors.New("native SSH executor is closed") } entry := p.entries[key] if entry == nil { @@ -443,24 +577,31 @@ func (p *nativeClientPool) get(ctx context.Context, key string, dial func(contex } if entry.client != nil { client := entry.client + entry.inFlight++ p.mu.Unlock() - return client, nil + 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, ctx.Err() + 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, error) { +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() @@ -471,27 +612,29 @@ func (p *nativeClientPool) dialEntry(ctx context.Context, key string, entry *nat _ = client.Close() } if err != nil { - return nil, err + return nil, nil, err } - return nil, errors.New("native SSH pooled connection was replaced while dialing") + 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) - close(entry.ready) - return nil, err + entry.closeReady() + return nil, nil, err } if p.closed { delete(p.entries, key) - close(entry.ready) + entry.closeReady() _ = client.Close() - return nil, errors.New("native SSH executor is closed") + return nil, nil, errors.New("native SSH executor is closed") } entry.client = client + entry.inFlight = 1 entry.stop = make(chan struct{}) - close(entry.ready) + entry.closeReady() p.startKeepAliveLocked(key, entry) - return client, nil + return client, entry, nil } func (p *nativeClientPool) startKeepAliveLocked(key string, entry *nativePoolEntry) { @@ -507,8 +650,8 @@ func (p *nativeClientPool) startKeepAliveLocked(key string, entry *nativePoolEnt for { select { case <-ticker.C: - if _, _, err := client.SendRequest("keepalive@openssh.com", true, nil); err != nil { - p.evict(key, client) + if err := sendKeepAlive(client, nativeKeepAliveTimeout); err != nil { + p.evictNow(key, entry, client) return } case <-stop: @@ -518,42 +661,132 @@ func (p *nativeClientPool) startKeepAliveLocked(key string, entry *nativePoolEnt }() } -func (p *nativeClientPool) isDead(key string, client *nativeClient) bool { +func (p *nativeClientPool) isDead(key string, entry *nativePoolEntry, client *nativeClient) bool { p.mu.Lock() defer p.mu.Unlock() - entry := p.entries[key] - return entry == nil || entry.client != client + current := p.entries[key] + return current != entry || entry == nil || entry.client != client || entry.forceClosed } -func (p *nativeClientPool) evict(key string, client *nativeClient) { +func (p *nativeClientPool) release(entry *nativePoolEntry) { var closeClient *nativeClient - var stop chan struct{} + var closeStop *nativePoolEntry p.mu.Lock() - entry := p.entries[key] - if entry != nil && entry.client == client { - delete(p.entries, key) + 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 - stop = entry.stop + closeStop = entry } p.mu.Unlock() - if stop != nil { - close(stop) + 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 []chan struct{} - var ready []chan struct{} + var stops []*nativePoolEntry + var ready []*nativePoolEntry p.mu.Lock() if p.closed { @@ -564,10 +797,10 @@ func (p *nativeClientPool) Close() error { for key, entry := range p.entries { delete(p.entries, key) if entry.stop != nil { - stops = append(stops, entry.stop) + stops = append(stops, entry) } if entry.dialing && entry.ready != nil { - ready = append(ready, entry.ready) + ready = append(ready, entry) } if entry.client != nil { clients = append(clients, entry.client) @@ -575,11 +808,11 @@ func (p *nativeClientPool) Close() error { } p.mu.Unlock() - for _, stop := range stops { - close(stop) + for _, entry := range stops { + entry.closeStop() } - for _, ch := range ready { - close(ch) + for _, entry := range ready { + entry.closeReady() } var err error for _, client := range clients { @@ -738,6 +971,15 @@ func isBrokenSSHClientError(err error) bool { 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 946b805..2b99fcc 100644 --- a/internal/executor/native_test.go +++ b/internal/executor/native_test.go @@ -117,6 +117,33 @@ func TestNativeExecutorReusesPooledConnection(t *testing.T) { } } +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) @@ -155,6 +182,77 @@ func TestNativeExecutorPoolSingleFlightsConcurrentDial(t *testing.T) { } } +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) @@ -186,6 +284,109 @@ func TestNativeExecutorReconnectsWhenPooledClientDies(t *testing.T) { } } +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) @@ -213,6 +414,36 @@ func TestNativeExecutorKeepAliveAndClose(t *testing.T) { 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 @@ -498,13 +729,20 @@ func TestConnectHintMapping(t *testing.T) { } type testSSHServer struct { - Listener net.Listener - HostSigner ssh.Signer - wg sync.WaitGroup - allowedKey ssh.PublicKey - accepted int32 - closed int32 - globalRequests int32 + 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 @@ -609,7 +847,7 @@ func (s *passwordTestSSHServer) handle(conn net.Conn) { if err != nil { continue } - go handleSession(sshConn, channel, requests) + go handleSession(sshConn, channel, requests, nil) } } @@ -686,22 +924,42 @@ func (s *testSSHServer) handle(conn net.Conn) { _ = 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(sshConn, 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) } @@ -720,7 +978,7 @@ func waitForAtomic(t *testing.T, value *int32, want int32) { 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) { +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" { @@ -732,7 +990,15 @@ func handleSession(conn *ssh.ServerConn, channel ssh.Channel, requests <-chan *s _ = 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=")) 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"`