From 849951be1d1a7ee9f9302006ccb187bf5b4e36f3 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Wed, 22 Jul 2026 23:51:18 +0300 Subject: =?UTF-8?q?feat:=20DTail=20fork=20=E2=80=94=20server/client=20feat?= =?UTF-8?q?ure=20development?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Squashed development of the snonux/dtail fork's product code (internal/, cmd/) since diverging from mimecast/dtail. Major areas: - Read/output path: the former "turbo" channel-less path is now the single, default server-side read/output path for cat/grep/tail and MapReduce; the old channel-based path and its config/env toggles were removed. - MapReduce: single aggregate implementation (server + serverless) fed directly by a processor pipeline, with input-exhausted finalization via the shutdown coordinator; high-concurrency and data-race fixes. - Journal source reads (journal:unit.service) via journalctl, Linux-gated behind a journal-v1 capability. - Auth-key fast reconnect: in-memory per-user public-key cache with TTL/max-keys, registered over an authenticated session (AUTHKEY), checked before authorized_keys. - Interactive query reload (--interactive-query) with SESSION START/UPDATE generation boundaries and capability negotiation. - Client-side deadlines: --timeout / --shutdownAfter as context deadlines; follow shutdown handling. - Client logging: diagnostics-only daily log by default, opt-in payload tee via --log-payload. - Numerous correctness fixes (buffer-pool double-recycle races, EOF-sentinel leaks, glob-expansion cap, TOCTOU in CSV parsing) with accompanying unit tests. Co-Authored-By: Claude Opus 4.8 --- internal/ssh/client/authmethods.go | 178 +++++++++------ internal/ssh/client/authmethods_test.go | 194 ++++++++++++++++ internal/ssh/client/customkeycallback.go | 6 +- internal/ssh/client/hostkeycallback.go | 7 +- internal/ssh/client/knownhostscallback.go | 205 +++++++++++++---- internal/ssh/client/knownhostscallback_test.go | 305 +++++++++++++++++++++++++ internal/ssh/client/simplecallback.go | 5 +- 7 files changed, 780 insertions(+), 120 deletions(-) create mode 100644 internal/ssh/client/authmethods_test.go create mode 100644 internal/ssh/client/knownhostscallback_test.go (limited to 'internal/ssh/client') diff --git a/internal/ssh/client/authmethods.go b/internal/ssh/client/authmethods.go index 6128018..67985a4 100644 --- a/internal/ssh/client/authmethods.go +++ b/internal/ssh/client/authmethods.go @@ -2,6 +2,7 @@ package client import ( "fmt" + "io" "os" "github.com/mimecast/dtail/internal/config" @@ -11,106 +12,151 @@ import ( gossh "golang.org/x/crypto/ssh" ) -const addedPathStr string = "Added path to list of auth methods, not adding further methods" +// noopCloser lets callers unconditionally defer closer.Close() when a code +// path does not own a real resource (e.g. no agent available). +type noopCloserFunc struct{} + +func (noopCloserFunc) Close() error { return nil } + +var noAuthCloser io.Closer = noopCloserFunc{} + +var ( + privateKeySigner = ssh.PrivateKeySigner + agentSigners = ssh.AgentSignersWithKeyIndex +) // InitSSHAuthMethods initialises all known SSH auth methods on the client side. +// The returned io.Closer owns any ssh-agent connection acquired while building +// the auth methods and must be closed by the caller once all SSH handshakes +// that consume the returned auth methods have completed. The closer is always +// non-nil so callers can unconditionally `defer closer.Close()`. func InitSSHAuthMethods(sshAuthMethods []gossh.AuthMethod, - hostKeyCallback gossh.HostKeyCallback, trustAllHosts bool, throttleCh chan struct{}, - privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) { + hostKeyCallback gossh.HostKeyCallback, trustAllHosts bool, + privateKeyPath string, agentKeyIndex int) ([]gossh.AuthMethod, HostKeyCallback, io.Closer) { if len(sshAuthMethods) > 0 { simpleCallback, err := NewSimpleCallback() if err != nil { dlog.Client.FatalPanic(err) } - return sshAuthMethods, simpleCallback - } - return initKnownHostsAuthMethods(trustAllHosts, throttleCh, privateKeyPath) -} - -func initIntegrationTestKnownHostsAuthMethods() []gossh.AuthMethod { - var sshAuthMethods []gossh.AuthMethod - privateKeyPath := "./id_rsa" - - GeneratePrivatePublicKeyPairIfNotExists(privateKeyPath, 4096) - authMethod, err := ssh.PrivateKey(privateKeyPath) - if err != nil { - dlog.Client.FatalPanic("Unable to use private SSH key", privateKeyPath, err) + return sshAuthMethods, simpleCallback, noAuthCloser } - - sshAuthMethods = append(sshAuthMethods, authMethod) - dlog.Client.Debug("initKnownHostsAuthMethods", addedPathStr, privateKeyPath) - return sshAuthMethods + return initKnownHostsAuthMethods(trustAllHosts, privateKeyPath, agentKeyIndex) } -func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{}, - privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) { +func initKnownHostsAuthMethods(trustAllHosts bool, + privateKeyPath string, agentKeyIndex int) ([]gossh.AuthMethod, HostKeyCallback, io.Closer) { - var sshAuthMethods []gossh.AuthMethod knownHostsFile := fmt.Sprintf("%s/.ssh/known_hosts", os.Getenv("HOME")) if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") { // In case of integration test, override known hosts file path. knownHostsFile = "./known_hosts" } - knownHostsCallback, err := NewKnownHostsCallback(knownHostsFile, trustAllHosts, throttleCh) + knownHostsCallback, err := NewKnownHostsCallback(knownHostsFile, trustAllHosts) if err != nil { dlog.Client.FatalPanic(knownHostsFile, err) } dlog.Client.Debug("initKnownHostsAuthMethods", "Added known hosts file path", knownHostsFile) if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") { - return initIntegrationTestKnownHostsAuthMethods(), knownHostsCallback + if privateKeyPath == "" { + privateKeyPath = "./id_rsa" + } + GeneratePrivatePublicKeyPairIfNotExists(privateKeyPath, 4096) + } + + sshAuthMethods, agentCloser := collectKnownHostsAuthMethods(privateKeyPath, agentKeyIndex) + if len(sshAuthMethods) == 0 { + _ = agentCloser.Close() + dlog.Client.FatalPanic("Unable to find private SSH key information") } - // Try to read custom private key path. - if privateKeyPath != "" { - authMethod, err := ssh.PrivateKey(privateKeyPath) - if err == nil { - sshAuthMethods = append(sshAuthMethods, authMethod) - dlog.Client.Debug("initKnownHostsAuthMethods", addedPathStr, privateKeyPath) - return sshAuthMethods, knownHostsCallback + return sshAuthMethods, knownHostsCallback, agentCloser +} + +func collectKnownHostsAuthMethods(privateKeyPath string, agentKeyIndex int) ([]gossh.AuthMethod, io.Closer) { + signers, agentCloser := collectKnownHostsSigners(privateKeyPath, agentKeyIndex) + if len(signers) == 0 { + return nil, agentCloser + } + return []gossh.AuthMethod{gossh.PublicKeys(signers...)}, agentCloser +} + +func collectKnownHostsSigners(privateKeyPath string, agentKeyIndex int) ([]gossh.Signer, io.Closer) { + var signers []gossh.Signer + + home := os.Getenv("HOME") + defaultPrivateKeyPaths := []string{ + home + "/.ssh/id_rsa", + home + "/.ssh/id_dsa", + home + "/.ssh/id_ecdsa", + home + "/.ssh/id_ed25519", + } + if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") { + defaultPrivateKeyPaths = append([]string{"./id_rsa"}, defaultPrivateKeyPaths...) + } + + if privateKeyPath == "" { + privateKeyPath = defaultPrivateKeyPaths[0] + } + + addedPrivateKeyPaths := make(map[string]bool, len(defaultPrivateKeyPaths)+1) + addedPublicKeys := make(map[string]bool, len(defaultPrivateKeyPaths)+1) + addSigner := func(source string, signer gossh.Signer) { + if signer == nil { + return + } + + pubKey := string(signer.PublicKey().Marshal()) + if addedPublicKeys[pubKey] { + dlog.Client.Debug("initKnownHostsAuthMethods", "Skipping duplicate signer", source) + return + } + + addedPublicKeys[pubKey] = true + signers = append(signers, signer) + dlog.Client.Debug("initKnownHostsAuthMethods", "Added signer", source) + } + addPrivateKeySigner := func(path string) { + if path == "" { + return + } + if addedPrivateKeyPaths[path] { + return } - dlog.Client.FatalPanic("Unable to use private SSH key", privateKeyPath, err) + + signer, err := privateKeySigner(path) + if err != nil { + dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to load private key signer", path, err) + return + } + + addedPrivateKeyPaths[path] = true + addSigner(path, signer) } - // Second, try SSH Agent - authMethod, err := ssh.Agent() - if err == nil { - sshAuthMethods = append(sshAuthMethods, authMethod) - dlog.Client.Debug("initKnownHostsAuthMethods", "Added SSH Agent (SSH_AUTH_SOCK)"+ - "to list of auth methods, not adding further methods") - return sshAuthMethods, knownHostsCallback + // First, the explicit auth key path (or default ~/.ssh/id_rsa). + addPrivateKeySigner(privateKeyPath) + + // Second, SSH agent (YubiKey-backed keys are typically exposed here). + // The agent signers sign lazily over the agent connection, so its + // io.Closer must live until the caller is done with the signers. + loadedAgentSigners, agentCloser, err := agentSigners(agentKeyIndex) + if err != nil { + dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to load SSH agent signers", err) } - dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to init SSH Agent auth method", err) - - // Third, try Linux/UNIX default key paths - privateKeyPath = os.Getenv("HOME") + "/.ssh/id_rsa" - authMethod, err = ssh.PrivateKey(privateKeyPath) - if err == nil { - sshAuthMethods = append(sshAuthMethods, authMethod) - dlog.Client.Debug("initKnownHostsAuthmethods", addedPathStr, privateKeyPath) - return sshAuthMethods, knownHostsCallback + if agentCloser == nil { + agentCloser = noAuthCloser } - dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to use private key", privateKeyPath, err) - - privateKeyPath = os.Getenv("HOME") + "/.ssh/id_dsa" - authMethod, err = ssh.PrivateKey(privateKeyPath) - if err == nil { - sshAuthMethods = append(sshAuthMethods, authMethod) - dlog.Client.Debug("initKnownHostsAuthmethods", addedPathStr, privateKeyPath) - return sshAuthMethods, knownHostsCallback + for i, signer := range loadedAgentSigners { + addSigner(fmt.Sprintf("agent:%d:%d", agentKeyIndex, i), signer) } - privateKeyPath = os.Getenv("HOME") + "/.ssh/id_ecdsa" - authMethod, err = ssh.PrivateKey(privateKeyPath) - if err == nil { - sshAuthMethods = append(sshAuthMethods, authMethod) - dlog.Client.Debug("initKnownHostsAuthmethods", addedPathStr, privateKeyPath) - return sshAuthMethods, knownHostsCallback + // Third, additional default private key paths. + for _, path := range defaultPrivateKeyPaths { + addPrivateKeySigner(path) } - dlog.Client.FatalPanic("Unable to find private SSH key information", privateKeyPath, err) - // Never reach this point. - return sshAuthMethods, knownHostsCallback + return signers, agentCloser } diff --git a/internal/ssh/client/authmethods_test.go b/internal/ssh/client/authmethods_test.go new file mode 100644 index 0000000..3811c91 --- /dev/null +++ b/internal/ssh/client/authmethods_test.go @@ -0,0 +1,194 @@ +package client + +import ( + "fmt" + "io" + "reflect" + "testing" + + "github.com/mimecast/dtail/internal/io/dlog" + + gossh "golang.org/x/crypto/ssh" +) + +// testCloser is a sentinel io.Closer used by tests to assert that callers +// release ssh-agent connections returned by the mocked agentSigners hook. +type testCloser struct { + closed int +} + +func (c *testCloser) Close() error { + c.closed++ + return nil +} + +type mockPublicKey struct { + id string +} + +func (k *mockPublicKey) Type() string { + return "ssh-rsa" +} + +func (k *mockPublicKey) Marshal() []byte { + return []byte(k.id) +} + +func (k *mockPublicKey) Verify(_ []byte, _ *gossh.Signature) error { + return nil +} + +type mockSigner struct { + key gossh.PublicKey +} + +func newMockSigner(id string) gossh.Signer { + return &mockSigner{key: &mockPublicKey{id: id}} +} + +func (s *mockSigner) PublicKey() gossh.PublicKey { + return s.key +} + +func (s *mockSigner) Sign(_ io.Reader, _ []byte) (*gossh.Signature, error) { + return &gossh.Signature{ + Format: "ssh-rsa", + Blob: []byte("sig"), + }, nil +} + +func TestCollectKnownHostsAuthMethodsOrder(t *testing.T) { + homeDir := "/tmp/dtail-auth-order" + t.Setenv("HOME", homeDir) + // Keep this unit test deterministic regardless of integration-mode env. + t.Setenv("DTAIL_INTEGRATION_TEST_RUN_MODE", "") + + originalPrivateKeySigner := privateKeySigner + originalAgentSigners := agentSigners + originalLogger := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + privateKeySigner = originalPrivateKeySigner + agentSigners = originalAgentSigners + dlog.Client = originalLogger + }) + + var callOrder []string + successfulPrivateKeys := map[string]gossh.Signer{ + "/custom/id_fast": newMockSigner("custom"), + homeDir + "/.ssh/id_rsa": newMockSigner("default-rsa"), + homeDir + "/.ssh/id_dsa": newMockSigner("default-dsa"), + } + + privateKeySigner = func(path string) (gossh.Signer, error) { + callOrder = append(callOrder, "private:"+path) + signer, found := successfulPrivateKeys[path] + if !found { + return nil, fmt.Errorf("missing private key: %s", path) + } + return signer, nil + } + agentCloser := &testCloser{} + agentSigners = func(keyIndex int) ([]gossh.Signer, io.Closer, error) { + callOrder = append(callOrder, fmt.Sprintf("agent:%d", keyIndex)) + return []gossh.Signer{newMockSigner("agent")}, agentCloser, nil + } + + methods, closer := collectKnownHostsAuthMethods("/custom/id_fast", 7) + if len(methods) != 1 { + t.Fatalf("Expected 1 auth method, got %d", len(methods)) + } + if closer == nil { + t.Fatalf("Expected non-nil agent closer from collectKnownHostsAuthMethods") + } + if err := closer.Close(); err != nil { + t.Fatalf("agent closer returned error: %v", err) + } + if agentCloser.closed < 1 { + t.Fatalf("Expected caller to be able to close agent conn; closed=%d", agentCloser.closed) + } + + callOrder = nil + signers, sCloser := collectKnownHostsSigners("/custom/id_fast", 7) + if len(signers) != 4 { + t.Fatalf("Expected 4 signers, got %d", len(signers)) + } + if sCloser == nil { + t.Fatalf("Expected non-nil agent closer from collectKnownHostsSigners") + } + _ = sCloser.Close() + + expectedOrder := []string{ + "private:/custom/id_fast", + "agent:7", + "private:/tmp/dtail-auth-order/.ssh/id_rsa", + "private:/tmp/dtail-auth-order/.ssh/id_dsa", + "private:/tmp/dtail-auth-order/.ssh/id_ecdsa", + "private:/tmp/dtail-auth-order/.ssh/id_ed25519", + } + if !reflect.DeepEqual(callOrder, expectedOrder) { + t.Fatalf("Unexpected auth method call order.\nexpected: %v\ngot: %v", expectedOrder, callOrder) + } +} + +func TestCollectKnownHostsAuthMethodsSkipsDuplicateDefaultPath(t *testing.T) { + homeDir := "/tmp/dtail-auth-dedupe" + t.Setenv("HOME", homeDir) + // Keep this unit test deterministic regardless of integration-mode env. + t.Setenv("DTAIL_INTEGRATION_TEST_RUN_MODE", "") + + originalPrivateKeySigner := privateKeySigner + originalAgentSigners := agentSigners + originalLogger := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + privateKeySigner = originalPrivateKeySigner + agentSigners = originalAgentSigners + dlog.Client = originalLogger + }) + + sharedSigner := newMockSigner("shared") + var callOrder []string + privateKeySigner = func(path string) (gossh.Signer, error) { + callOrder = append(callOrder, "private:"+path) + if path == homeDir+"/.ssh/id_rsa" { + return sharedSigner, nil + } + return nil, fmt.Errorf("missing private key: %s", path) + } + agentCloser := &testCloser{} + agentSigners = func(keyIndex int) ([]gossh.Signer, io.Closer, error) { + callOrder = append(callOrder, fmt.Sprintf("agent:%d", keyIndex)) + return []gossh.Signer{sharedSigner}, agentCloser, nil + } + + methods, closer := collectKnownHostsAuthMethods(homeDir+"/.ssh/id_rsa", 2) + if len(methods) != 1 { + t.Fatalf("Expected 1 auth method, got %d", len(methods)) + } + if closer == nil { + t.Fatalf("Expected non-nil agent closer from collectKnownHostsAuthMethods") + } + _ = closer.Close() + + callOrder = nil + signers, sCloser := collectKnownHostsSigners(homeDir+"/.ssh/id_rsa", 2) + if len(signers) != 1 { + t.Fatalf("Expected duplicate keys to collapse to 1 signer, got %d", len(signers)) + } + if sCloser == nil { + t.Fatalf("Expected non-nil agent closer from collectKnownHostsSigners") + } + _ = sCloser.Close() + + expectedOrder := []string{ + "private:/tmp/dtail-auth-dedupe/.ssh/id_rsa", + "agent:2", + "private:/tmp/dtail-auth-dedupe/.ssh/id_dsa", + "private:/tmp/dtail-auth-dedupe/.ssh/id_ecdsa", + "private:/tmp/dtail-auth-dedupe/.ssh/id_ed25519", + } + if !reflect.DeepEqual(callOrder, expectedOrder) { + t.Fatalf("Unexpected auth method call order.\nexpected: %v\ngot: %v", expectedOrder, callOrder) + } +} diff --git a/internal/ssh/client/customkeycallback.go b/internal/ssh/client/customkeycallback.go index 53b8e3c..0107895 100644 --- a/internal/ssh/client/customkeycallback.go +++ b/internal/ssh/client/customkeycallback.go @@ -1,6 +1,7 @@ package client import ( + "context" "net" "golang.org/x/crypto/ssh" @@ -15,8 +16,9 @@ func NewCustomCallback() (*CustomCallback, error) { return &h, nil } -// Wrap the host key callback. -func (h *CustomCallback) Wrap() ssh.HostKeyCallback { +// Wrap the host key callback. ctx is accepted for interface compatibility +// but the custom callback never blocks so it has nothing to abort. +func (h *CustomCallback) Wrap(_ context.Context) ssh.HostKeyCallback { return func(server string, remote net.Addr, key ssh.PublicKey) error { return nil } diff --git a/internal/ssh/client/hostkeycallback.go b/internal/ssh/client/hostkeycallback.go index 95543f2..1fcd9c6 100644 --- a/internal/ssh/client/hostkeycallback.go +++ b/internal/ssh/client/hostkeycallback.go @@ -8,8 +8,13 @@ import ( // HostKeyCallback is a wrapper around ssh.KnownHosts so that we can add all // unknown hosts in a single batch to the known_hosts file. +// +// Wrap returns an ssh.HostKeyCallback that is bound to the lifetime of the +// supplied SSH handshake context. Implementations MUST abort any blocking +// operations (e.g. prompting the user for unknown hosts) once ctx is +// cancelled so the SSH handshake does not hang and no goroutines leak. type HostKeyCallback interface { - Wrap() ssh.HostKeyCallback + Wrap(ctx context.Context) ssh.HostKeyCallback Untrusted(server string) bool PromptAddHosts(ctx context.Context) } diff --git a/internal/ssh/client/knownhostscallback.go b/internal/ssh/client/knownhostscallback.go index fe3543c..aee675d 100644 --- a/internal/ssh/client/knownhostscallback.go +++ b/internal/ssh/client/knownhostscallback.go @@ -11,6 +11,7 @@ import ( "time" "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/io/fs" "github.com/mimecast/dtail/internal/io/prompt" "golang.org/x/crypto/ssh" @@ -38,36 +39,78 @@ type unknownHost struct { // unknown hosts in a single batch to the known_hosts file. type KnownHostsCallback struct { knownHostsPath string + knownHostsFile fs.RootedPath unknownCh chan unknownHost - throttleCh chan struct{} trustAllHostsCh chan struct{} - untrustedHosts map[string]bool - mutex *sync.Mutex + // trustAllOnce guards the single close of trustAllHostsCh. The old + // select/default/close pattern was not atomic: two concurrent callers + // could both observe the channel open, both fall through to the default + // branch, and both call close() — causing a panic. sync.Once makes the + // close idempotent and race-free without any additional locking. + trustAllOnce sync.Once + untrustedHosts map[string]bool + mutex *sync.Mutex } +var _ HostKeyCallback = (*KnownHostsCallback)(nil) + // NewKnownHostsCallback returns a new wrapper. -func NewKnownHostsCallback(knownHostsPath string, trustAllHosts bool, - throttleCh chan struct{}) (HostKeyCallback, error) { +func NewKnownHostsCallback(knownHostsPath string, trustAllHosts bool) (HostKeyCallback, error) { - os.OpenFile(knownHostsPath, os.O_RDONLY|os.O_CREATE, 0666) + knownHostsFile, err := fs.NewRootedPath(knownHostsPath) + if err != nil { + return nil, err + } + ensureKnownHostsFile(knownHostsFile) untrustedHosts := make(map[string]bool) c := KnownHostsCallback{ knownHostsPath: knownHostsPath, + knownHostsFile: knownHostsFile, unknownCh: make(chan unknownHost), trustAllHostsCh: make(chan struct{}), - throttleCh: throttleCh, untrustedHosts: untrustedHosts, mutex: &sync.Mutex{}, } if trustAllHosts { - close(c.trustAllHostsCh) + // Use the same sync.Once path so both the constructor and the + // interactive "all" prompt are idempotent and race-free. + c.closeTrustAllHostsCh() } - return c, nil + return &c, nil } -// Wrap the host key callback. -func (c KnownHostsCallback) Wrap() ssh.HostKeyCallback { +// closeTrustAllHostsCh closes trustAllHostsCh exactly once via sync.Once, +// regardless of how many goroutines call it concurrently. This replaces the +// former select/default/close pattern which was not atomic: two concurrent +// callers could both observe the channel open, both take the default branch, +// and both call close() — causing a panic. +func (c *KnownHostsCallback) closeTrustAllHostsCh() { + c.trustAllOnce.Do(func() { close(c.trustAllHostsCh) }) +} + +func ensureKnownHostsFile(knownHostsFile fs.RootedPath) { + root, err := knownHostsFile.OpenRoot() + if err != nil { + return + } + defer root.Close() + + fd, err := root.OpenFile(knownHostsFile.Name(), os.O_RDONLY|os.O_CREATE, 0o666) + if err != nil { + return + } + fd.Close() +} + +// Wrap the host key callback. The returned ssh.HostKeyCallback is bound to +// ctx: if ctx is cancelled while we are waiting for the PromptAddHosts +// goroutine to consume an unknown host or to return a user decision, the +// callback aborts with ctx.Err() instead of blocking forever. This prevents +// a stuck SSH handshake (and a leaked goroutine per unknown host) when the +// client shuts down before the user responds, or when PromptAddHosts has +// already returned because its ctx was cancelled. +func (c *KnownHostsCallback) Wrap(ctx context.Context) ssh.HostKeyCallback { return func(server string, remote net.Addr, key ssh.PublicKey) error { // Parse known_hosts file knownHostsCb, err := knownhosts.New(c.knownHostsPath) @@ -80,10 +123,6 @@ func (c KnownHostsCallback) Wrap() ssh.HostKeyCallback { // OK return nil } - // Make sure that interactive user callback does not interfere with - // SSH connection throttler. - <-c.throttleCh - defer func() { c.throttleCh <- struct{}{} }() unknown := unknownHost{ server: server, @@ -91,13 +130,26 @@ func (c KnownHostsCallback) Wrap() ssh.HostKeyCallback { key: key, hostLine: knownhosts.Line([]string{server}, key), ipLine: knownhosts.Line([]string{remote.String()}, key), - responseCh: make(chan response), + responseCh: make(chan response, 1), + } + // Keep host trust discovery diagnostics out of normal command output. + // In trust-all and plain modes this warning can corrupt tool output. + dlog.Client.Debug("Encountered unknown host", unknown.server, unknown.remote.String()) + // Notify user that there is an unknown host. Honour ctx cancellation + // so we do not block forever when PromptAddHosts has already exited. + select { + case c.unknownCh <- unknown: + case <-ctx.Done(): + return fmt.Errorf("host key callback cancelled for %s: %w", server, ctx.Err()) } - dlog.Client.Warn("Encountered unknown host", unknown) - // Notify user that there is an unknown host - c.unknownCh <- unknown - // Wait for user input. - switch <-unknown.responseCh { + // Wait for user input. Same contract as above: abort on ctx cancel. + var resp response + select { + case resp = <-unknown.responseCh: + case <-ctx.Done(): + return fmt.Errorf("host key callback cancelled for %s: %w", server, ctx.Err()) + } + switch resp { case trustHost: // End user acknowledged host key return nil @@ -113,7 +165,7 @@ func (c KnownHostsCallback) Wrap() ssh.HostKeyCallback { // PromptAddHosts prompts a question to the user whether unknown hosts should // be added to the known hosts or not. -func (c KnownHostsCallback) PromptAddHosts(ctx context.Context) { +func (c *KnownHostsCallback) PromptAddHosts(ctx context.Context) { var hosts []unknownHost for { // Check whether there is a unknown host @@ -138,7 +190,7 @@ func (c KnownHostsCallback) PromptAddHosts(ctx context.Context) { } } -func (c KnownHostsCallback) promptAddHosts(hosts []unknownHost) { +func (c *KnownHostsCallback) promptAddHosts(hosts []unknownHost) { var servers []string for _, host := range hosts { servers = append(servers, host.server) @@ -146,8 +198,12 @@ func (c KnownHostsCallback) promptAddHosts(hosts []unknownHost) { select { case <-c.trustAllHostsCh: - dlog.Client.Warn("Trusting host keys of servers", servers) - c.trustHosts(hosts) + // Trust-all mode is non-interactive; avoid warning-level noise on stdout. + dlog.Client.Debug("Trusting host keys of servers", servers) + if err := c.trustHosts(hosts); err != nil { + dlog.Client.Error("Unable to update known hosts file", c.knownHostsPath, err) + c.dontTrustHosts(hosts) + } return default: } @@ -163,9 +219,11 @@ func (c KnownHostsCallback) promptAddHosts(hosts []unknownHost) { Long: "yes", Short: "y", Callback: func() { - c.trustHosts(hosts) - }, - EndCallback: func() { + if err := c.trustHosts(hosts); err != nil { + dlog.Client.Error("Unable to update known hosts file", c.knownHostsPath, err) + c.dontTrustHosts(hosts) + return + } dlog.Client.Info("Added hosts to known hosts file", c.knownHostsPath) }, } @@ -175,10 +233,14 @@ func (c KnownHostsCallback) promptAddHosts(hosts []unknownHost) { Long: "all", Short: "a", Callback: func() { - close(c.trustAllHostsCh) - c.trustHosts(hosts) - }, - EndCallback: func() { + if err := c.trustHosts(hosts); err != nil { + dlog.Client.Error("Unable to update known hosts file", c.knownHostsPath, err) + c.dontTrustHosts(hosts) + return + } + // Mark trust-all atomically so that concurrent "all" callbacks + // from other batches do not double-close the channel. + c.closeTrustAllHostsCh() dlog.Client.Info("Added hosts to known hosts file", c.knownHostsPath) }, } @@ -212,41 +274,59 @@ func (c KnownHostsCallback) promptAddHosts(hosts []unknownHost) { p.Ask() } -func (c KnownHostsCallback) trustHosts(hosts []unknownHost) { +func (c *KnownHostsCallback) trustHosts(hosts []unknownHost) error { + root, err := c.knownHostsFile.OpenRoot() + if err != nil { + return err + } + defer root.Close() + + tmpKnownHostsName := fmt.Sprintf("%s.tmp", c.knownHostsFile.Name()) tmpKnownHostsPath := fmt.Sprintf("%s.tmp", c.knownHostsPath) + cleanupTmp := func() { + if err := root.Remove(tmpKnownHostsName); err != nil && !os.IsNotExist(err) { + dlog.Client.Debug("Unable to remove temporary known hosts file", tmpKnownHostsPath, err) + } + } - newFd, err := os.OpenFile(tmpKnownHostsPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0600) + newFd, err := root.OpenFile(tmpKnownHostsName, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600) if err != nil { - panic(fmt.Sprintf("%s: %s", tmpKnownHostsPath, err.Error())) + return fmt.Errorf("open temp known hosts file %s: %w", tmpKnownHostsPath, err) + } + if err := newFd.Chmod(0o600); err != nil { + newFd.Close() + cleanupTmp() + return fmt.Errorf("chmod temp known hosts file %s: %w", tmpKnownHostsPath, err) } - defer newFd.Close() // Newly trusted hosts in normalized form addresses := make(map[string]struct{}) // First write to new known hosts file, and keep track of addresses for _, unknown := range hosts { - unknown.responseCh <- trustHost - // Add once as [HOSTNAME]:PORT addresses[knownhosts.Normalize(unknown.server)] = struct{}{} // And once as [IP]:PORT addresses[knownhosts.Normalize(unknown.remote.String())] = struct{}{} if _, err := newFd.WriteString(fmt.Sprintf("%s\n", unknown.hostLine)); err != nil { - panic(err) + newFd.Close() + cleanupTmp() + return fmt.Errorf("write host known_hosts entry: %w", err) } if _, err := newFd.WriteString(fmt.Sprintf("%s\n", unknown.ipLine)); err != nil { - panic(err) + newFd.Close() + cleanupTmp() + return fmt.Errorf("write ip known_hosts entry: %w", err) } } // Read old known hosts file, to see which are old and new entries - os.OpenFile(c.knownHostsPath, os.O_RDONLY|os.O_CREATE, 0666) - oldFd, err := os.Open(c.knownHostsPath) + oldFd, err := root.OpenFile(c.knownHostsFile.Name(), os.O_RDONLY|os.O_CREATE, 0o600) if err != nil { - panic(err) + newFd.Close() + cleanupTmp() + return fmt.Errorf("open known hosts file %s: %w", c.knownHostsPath, err) } - defer oldFd.Close() scanner := bufio.NewScanner(oldFd) // Now, append all still valid old entries to the new host file @@ -255,24 +335,51 @@ func (c KnownHostsCallback) trustHosts(hosts []unknownHost) { address := strings.SplitN(line, " ", 2)[0] if _, ok := addresses[address]; !ok { - newFd.WriteString(fmt.Sprintf("%s\n", line)) + if _, err := newFd.WriteString(fmt.Sprintf("%s\n", line)); err != nil { + oldFd.Close() + newFd.Close() + cleanupTmp() + return fmt.Errorf("append existing known_hosts entry: %w", err) + } } } + if err := scanner.Err(); err != nil { + oldFd.Close() + newFd.Close() + cleanupTmp() + return fmt.Errorf("scan existing known_hosts entries: %w", err) + } + + if err := oldFd.Close(); err != nil { + newFd.Close() + cleanupTmp() + return fmt.Errorf("close known hosts file %s: %w", c.knownHostsPath, err) + } + if err := newFd.Close(); err != nil { + cleanupTmp() + return fmt.Errorf("close temp known hosts file %s: %w", tmpKnownHostsPath, err) + } // Now, replace old known hosts file - if err := os.Rename(tmpKnownHostsPath, c.knownHostsPath); err != nil { - panic(err) + if err := root.Rename(tmpKnownHostsName, c.knownHostsFile.Name()); err != nil { + cleanupTmp() + return fmt.Errorf("replace known_hosts file %s: %w", c.knownHostsPath, err) + } + + for _, unknown := range hosts { + unknown.responseCh <- trustHost } + return nil } -func (c KnownHostsCallback) dontTrustHosts(hosts []unknownHost) { +func (c *KnownHostsCallback) dontTrustHosts(hosts []unknownHost) { for _, unknown := range hosts { unknown.responseCh <- dontTrustHost } } // Untrusted returns true if the host is not trusted. False otherwise. -func (c KnownHostsCallback) Untrusted(server string) bool { +func (c *KnownHostsCallback) Untrusted(server string) bool { c.mutex.Lock() defer c.mutex.Unlock() _, ok := c.untrustedHosts[server] diff --git a/internal/ssh/client/knownhostscallback_test.go b/internal/ssh/client/knownhostscallback_test.go new file mode 100644 index 0000000..ac9d438 --- /dev/null +++ b/internal/ssh/client/knownhostscallback_test.go @@ -0,0 +1,305 @@ +package client + +import ( + "context" + "errors" + "net" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/mimecast/dtail/internal/io/dlog" + + "golang.org/x/crypto/ssh/knownhosts" +) + +func TestTrustHostsAppendsDistinctExistingEntries(t *testing.T) { + knownHostsPath := filepath.Join(t.TempDir(), "known_hosts") + existingLine := knownhosts.Line([]string{"old.example:2222"}, &mockPublicKey{id: "old"}) + if err := os.WriteFile(knownHostsPath, []byte(existingLine+"\n"), 0o600); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + callback := testKnownHostsCallback(t, knownHostsPath) + unknown := testUnknownHost("new.example:2222", "127.0.0.1:2222", "new") + + if err := callback.trustHosts([]unknownHost{unknown}); err != nil { + t.Fatalf("trustHosts failed: %v", err) + } + + got, err := os.ReadFile(knownHostsPath) + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + + want := strings.Join([]string{ + unknown.hostLine, + unknown.ipLine, + existingLine, + "", + }, "\n") + if string(got) != want { + t.Fatalf("trustHosts wrote:\n%s\nwant:\n%s", got, want) + } + + if response := <-unknown.responseCh; response != trustHost { + t.Fatalf("unexpected trust response: %v", response) + } +} + +func TestTrustHostsReplacesExistingEntriesForSameHostAndIP(t *testing.T) { + knownHostsPath := filepath.Join(t.TempDir(), "known_hosts") + oldUnknown := testUnknownHost("replace.example:2222", "127.0.0.1:2222", "old") + keepLine := knownhosts.Line([]string{"keep.example:2222"}, &mockPublicKey{id: "keep"}) + initialContents := strings.Join([]string{ + oldUnknown.hostLine, + oldUnknown.ipLine, + keepLine, + "", + }, "\n") + if err := os.WriteFile(knownHostsPath, []byte(initialContents), 0o600); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + callback := testKnownHostsCallback(t, knownHostsPath) + newUnknown := testUnknownHost("replace.example:2222", "127.0.0.1:2222", "new") + + if err := callback.trustHosts([]unknownHost{newUnknown}); err != nil { + t.Fatalf("trustHosts failed: %v", err) + } + + got, err := os.ReadFile(knownHostsPath) + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + + want := strings.Join([]string{ + newUnknown.hostLine, + newUnknown.ipLine, + keepLine, + "", + }, "\n") + if string(got) != want { + t.Fatalf("trustHosts wrote:\n%s\nwant:\n%s", got, want) + } + + if response := <-newUnknown.responseCh; response != trustHost { + t.Fatalf("unexpected trust response: %v", response) + } +} + +func TestTrustHostsRejectsEscapingKnownHostsSymlink(t *testing.T) { + rootDir := filepath.Join(t.TempDir(), "ssh") + if err := os.MkdirAll(rootDir, 0o755); err != nil { + t.Fatalf("MkdirAll failed: %v", err) + } + + outsidePath := filepath.Join(filepath.Dir(rootDir), "outside_known_hosts") + if err := os.WriteFile(outsidePath, nil, 0o600); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + knownHostsPath := filepath.Join(rootDir, "known_hosts") + if err := os.Symlink(filepath.Join("..", "outside_known_hosts"), knownHostsPath); err != nil { + t.Fatalf("Symlink failed: %v", err) + } + + callback := testKnownHostsCallback(t, knownHostsPath) + unknown := testUnknownHost("escape.example:2222", "127.0.0.1:2222", "new") + + if err := callback.trustHosts([]unknownHost{unknown}); err == nil { + t.Fatalf("trustHosts succeeded for escaping known_hosts symlink") + } +} + +// stubClientLogger installs a no-op dlog.Client for tests that exercise the +// host-key callback (which emits a Debug log on the unknown-host path). +func stubClientLogger(t *testing.T) { + t.Helper() + original := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { dlog.Client = original }) +} + +// TestWrapReturnsWhenCtxCancelledBeforeUnknownChSend verifies that when +// PromptAddHosts has already exited (no consumer on unknownCh), the wrapped +// host-key callback unblocks on ctx cancel instead of hanging the SSH +// handshake and leaking a goroutine. Pre-fix this test times out because +// `c.unknownCh <- unknown` blocks forever. +func TestWrapReturnsWhenCtxCancelledBeforeUnknownChSend(t *testing.T) { + stubClientLogger(t) + knownHostsPath := filepath.Join(t.TempDir(), "known_hosts") + if err := os.WriteFile(knownHostsPath, nil, 0o600); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + callback := testKnownHostsCallback(t, knownHostsPath) + ctx, cancel := context.WithCancel(context.Background()) + + wrapped := callback.Wrap(ctx) + errCh := make(chan error, 1) + go func() { + errCh <- wrapped("host.example:2222", testTCPAddr("127.0.0.1:2222"), + &mockPublicKey{id: "new"}) + }() + + // Give the goroutine a moment to park on the unknownCh send, then cancel. + time.Sleep(10 * time.Millisecond) + cancel() + + select { + case err := <-errCh: + if err == nil { + t.Fatalf("expected non-nil error after ctx cancel, got nil") + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected error to wrap context.Canceled, got %v", err) + } + case <-time.After(100 * time.Millisecond): + t.Fatalf("host key callback did not return within 100ms after ctx cancel") + } +} + +// TestWrapReturnsWhenCtxCancelledBeforeResponse verifies that when a consumer +// has picked the unknown host off unknownCh but never writes a response +// (e.g. PromptAddHosts was cancelled mid-batch), the callback still unblocks +// on ctx cancel rather than blocking on responseCh forever. +func TestWrapReturnsWhenCtxCancelledBeforeResponse(t *testing.T) { + stubClientLogger(t) + knownHostsPath := filepath.Join(t.TempDir(), "known_hosts") + if err := os.WriteFile(knownHostsPath, nil, 0o600); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + callback := testKnownHostsCallback(t, knownHostsPath) + ctx, cancel := context.WithCancel(context.Background()) + + // Simulate a consumer that drains unknownCh but never writes to + // responseCh, mimicking PromptAddHosts buffering a batch and then exiting. + consumed := make(chan struct{}) + go func() { + <-callback.unknownCh + close(consumed) + }() + + wrapped := callback.Wrap(ctx) + errCh := make(chan error, 1) + go func() { + errCh <- wrapped("host.example:2222", testTCPAddr("127.0.0.1:2222"), + &mockPublicKey{id: "new"}) + }() + + select { + case <-consumed: + case <-time.After(100 * time.Millisecond): + t.Fatalf("consumer never received unknown host") + } + + cancel() + + select { + case err := <-errCh: + if err == nil { + t.Fatalf("expected non-nil error after ctx cancel, got nil") + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected error to wrap context.Canceled, got %v", err) + } + case <-time.After(100 * time.Millisecond): + t.Fatalf("host key callback did not return within 100ms after ctx cancel") + } +} + +// TestCloseTrustAllHostsChConcurrentNoPanic verifies that concurrent calls to +// closeTrustAllHostsCh — the method that replaces the racy select/default/close +// pattern — never panic and leave the channel durably closed. With -race this +// also detects any data race on the underlying sync.Once / trustAllHostsCh. +// +// Pre-fix, the equivalent inline select/default/close code in the "all" answer +// callback was not atomic: two goroutines could both observe the channel open, +// both take the default branch, and both call close() — causing a panic. +func TestCloseTrustAllHostsChConcurrentNoPanic(t *testing.T) { + knownHostsPath := filepath.Join(t.TempDir(), "known_hosts") + if err := os.WriteFile(knownHostsPath, nil, 0o600); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + callback := testKnownHostsCallback(t, knownHostsPath) + + // Run many goroutines concurrently to maximise the chance of hitting the + // race window that existed before the sync.Once fix. + const concurrency = 50 + ready := make(chan struct{}) + done := make(chan struct{}, concurrency) + + for i := 0; i < concurrency; i++ { + go func() { + <-ready + // closeTrustAllHostsCh is the idempotent replacement for the racy + // select/default/close sequence. All calls must be safe. + callback.closeTrustAllHostsCh() + done <- struct{}{} + }() + } + + // Release all goroutines simultaneously to maximise contention. + close(ready) + for i := 0; i < concurrency; i++ { + <-done + } + + // The channel must be closed exactly once: a receive on a closed channel + // returns immediately with the zero value. + select { + case <-callback.trustAllHostsCh: + // OK – closed by exactly one goroutine via sync.Once. + default: + t.Fatalf("trustAllHostsCh is still open after concurrent closeTrustAllHostsCh calls") + } +} + +func testKnownHostsCallback(t *testing.T, knownHostsPath string) *KnownHostsCallback { + t.Helper() + + callback, err := NewKnownHostsCallback(knownHostsPath, false) + if err != nil { + t.Fatalf("NewKnownHostsCallback failed: %v", err) + } + + knownHostsCallback, ok := callback.(*KnownHostsCallback) + if !ok { + t.Fatalf("unexpected callback type %T", callback) + } + + return knownHostsCallback +} + +func testUnknownHost(server, remoteAddr, keyID string) unknownHost { + key := &mockPublicKey{id: keyID} + remote := testTCPAddr(remoteAddr) + + return unknownHost{ + server: server, + remote: remote, + key: key, + hostLine: knownhosts.Line([]string{server}, key), + ipLine: knownhosts.Line([]string{remote.String()}, key), + responseCh: make(chan response, 1), + } +} + +func testTCPAddr(address string) *net.TCPAddr { + host, portStr, err := net.SplitHostPort(address) + if err != nil { + panic(err) + } + + port, err := net.LookupPort("tcp", portStr) + if err != nil { + panic(err) + } + + return &net.TCPAddr{IP: net.ParseIP(host), Port: port} +} diff --git a/internal/ssh/client/simplecallback.go b/internal/ssh/client/simplecallback.go index 580fa36..fd0e95b 100644 --- a/internal/ssh/client/simplecallback.go +++ b/internal/ssh/client/simplecallback.go @@ -17,8 +17,9 @@ func NewSimpleCallback() (SimpleCallback, error) { return SimpleCallback{}, nil } -// Wrap the host key callback. -func (SimpleCallback) Wrap() ssh.HostKeyCallback { +// Wrap the host key callback. ctx is accepted for interface compatibility +// but the simple callback never blocks so it has nothing to abort. +func (SimpleCallback) Wrap(_ context.Context) ssh.HostKeyCallback { return func(server string, remote net.Addr, key ssh.PublicKey) error { return nil } -- cgit v1.2.3