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 +- internal/ssh/server/authkeystore.go | 186 +++++++++++++++ internal/ssh/server/authkeystore_test.go | 178 +++++++++++++++ internal/ssh/server/hostkey.go | 70 ++++-- internal/ssh/server/hostkey_test.go | 37 +++ internal/ssh/server/publickeycallback.go | 159 ++++++++++--- internal/ssh/server/publickeycallback_test.go | 297 ++++++++++++++++++++++++ internal/ssh/ssh.go | 145 ++++++++++-- internal/ssh/ssh_agent_test.go | 152 ++++++++++++ internal/ssh/ssh_test.go | 113 +++++++++ 16 files changed, 2046 insertions(+), 191 deletions(-) create mode 100644 internal/ssh/client/authmethods_test.go create mode 100644 internal/ssh/client/knownhostscallback_test.go create mode 100644 internal/ssh/server/authkeystore.go create mode 100644 internal/ssh/server/authkeystore_test.go create mode 100644 internal/ssh/server/hostkey_test.go create mode 100644 internal/ssh/server/publickeycallback_test.go create mode 100644 internal/ssh/ssh_agent_test.go create mode 100644 internal/ssh/ssh_test.go (limited to 'internal/ssh') 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 } diff --git a/internal/ssh/server/authkeystore.go b/internal/ssh/server/authkeystore.go new file mode 100644 index 0000000..4de71ee --- /dev/null +++ b/internal/ssh/server/authkeystore.go @@ -0,0 +1,186 @@ +package server + +import ( + "sync" + "time" + + gossh "golang.org/x/crypto/ssh" +) + +const ( + defaultAuthKeyTTL = 24 * time.Hour + defaultAuthKeyMaxPerUser = 5 +) + +type authKeyEntry struct { + pubKey gossh.PublicKey + registeredAt time.Time +} + +// AuthKeyStore is an in-memory, per-user cache of SSH public keys. +// Each Server instance owns exactly one AuthKeyStore, constructed via +// NewAuthKeyStore. There is no shared package-level instance; callers +// must always supply a non-nil store. +type AuthKeyStore struct { + mu sync.RWMutex + keysByUser map[string][]authKeyEntry + ttl time.Duration + maxKeysPerUser int + now func() time.Time +} + +// NewAuthKeyStore builds a thread-safe auth key store. +func NewAuthKeyStore(ttl time.Duration, maxKeysPerUser int) *AuthKeyStore { + return newAuthKeyStoreWithClock(ttl, maxKeysPerUser, time.Now) +} + +func newAuthKeyStoreWithClock(ttl time.Duration, maxKeysPerUser int, + nowFn func() time.Time) *AuthKeyStore { + + if ttl <= 0 { + ttl = defaultAuthKeyTTL + } + if maxKeysPerUser <= 0 { + maxKeysPerUser = defaultAuthKeyMaxPerUser + } + if nowFn == nil { + nowFn = time.Now + } + + return &AuthKeyStore{ + keysByUser: make(map[string][]authKeyEntry), + ttl: ttl, + maxKeysPerUser: maxKeysPerUser, + now: nowFn, + } +} + +// Add stores or refreshes a key for a user. +func (s *AuthKeyStore) Add(user string, pubKey gossh.PublicKey) { + if user == "" || pubKey == nil { + return + } + + now := s.now() + offeredKey := marshalKey(pubKey) + + s.mu.Lock() + defer s.mu.Unlock() + + userEntries := s.pruneExpiredLocked(user, now) + + newEntries := make([]authKeyEntry, 0, len(userEntries)+1) + for _, entry := range userEntries { + if marshalKey(entry.pubKey) == offeredKey { + continue + } + newEntries = append(newEntries, entry) + } + + newEntries = append(newEntries, authKeyEntry{ + pubKey: pubKey, + registeredAt: now, + }) + if len(newEntries) > s.maxKeysPerUser { + newEntries = newEntries[len(newEntries)-s.maxKeysPerUser:] + } + + s.keysByUser[user] = newEntries +} + +// Has returns true if a non-expired key exists for a user. +func (s *AuthKeyStore) Has(user string, pubKey gossh.PublicKey) bool { + if user == "" || pubKey == nil { + return false + } + + now := s.now() + offeredKey := marshalKey(pubKey) + + s.mu.Lock() + defer s.mu.Unlock() + + userEntries := s.pruneExpiredLocked(user, now) + for _, entry := range userEntries { + if marshalKey(entry.pubKey) == offeredKey { + return true + } + } + + return false +} + +// Remove deletes a key for a user if it exists. +func (s *AuthKeyStore) Remove(user string, pubKey gossh.PublicKey) { + if user == "" || pubKey == nil { + return + } + + offeredKey := marshalKey(pubKey) + + s.mu.Lock() + defer s.mu.Unlock() + + userEntries := s.pruneExpiredLocked(user, s.now()) + if len(userEntries) == 0 { + return + } + + remaining := make([]authKeyEntry, 0, len(userEntries)) + for _, entry := range userEntries { + if marshalKey(entry.pubKey) == offeredKey { + continue + } + remaining = append(remaining, entry) + } + + if len(remaining) == 0 { + delete(s.keysByUser, user) + return + } + + s.keysByUser[user] = remaining +} + +func (s *AuthKeyStore) pruneExpiredLocked(user string, now time.Time) []authKeyEntry { + userEntries, ok := s.keysByUser[user] + if !ok || len(userEntries) == 0 { + delete(s.keysByUser, user) + return nil + } + + hasExpiredEntries := false + for _, entry := range userEntries { + if s.expired(entry, now) { + hasExpiredEntries = true + break + } + } + if !hasExpiredEntries { + return userEntries + } + + activeEntries := make([]authKeyEntry, 0, len(userEntries)) + for _, entry := range userEntries { + if s.expired(entry, now) { + continue + } + activeEntries = append(activeEntries, entry) + } + + if len(activeEntries) == 0 { + delete(s.keysByUser, user) + return nil + } + + s.keysByUser[user] = activeEntries + return activeEntries +} + +func (s *AuthKeyStore) expired(entry authKeyEntry, now time.Time) bool { + return !entry.registeredAt.Add(s.ttl).After(now) +} + +func marshalKey(pubKey gossh.PublicKey) string { + return string(pubKey.Marshal()) +} diff --git a/internal/ssh/server/authkeystore_test.go b/internal/ssh/server/authkeystore_test.go new file mode 100644 index 0000000..056db7b --- /dev/null +++ b/internal/ssh/server/authkeystore_test.go @@ -0,0 +1,178 @@ +package server + +import ( + "crypto/ed25519" + "sync" + "testing" + "time" + + gossh "golang.org/x/crypto/ssh" +) + +func TestAuthKeyStoreAddHasRemove(t *testing.T) { + store := NewAuthKeyStore(time.Hour, 5) + key := testPublicKey(t, 1) + + if store.Has("alice", key) { + t.Fatalf("Store should not contain key before add") + } + + store.Add("alice", key) + if !store.Has("alice", key) { + t.Fatalf("Store should contain key after add") + } + + store.Remove("alice", key) + if store.Has("alice", key) { + t.Fatalf("Store should not contain key after remove") + } +} + +func TestAuthKeyStoreHasExpiresKeysLazily(t *testing.T) { + now := time.Date(2026, 3, 3, 10, 0, 0, 0, time.UTC) + store := newAuthKeyStoreWithClock(10*time.Second, 5, func() time.Time { return now }) + key := testPublicKey(t, 2) + + store.Add("alice", key) + if !store.Has("alice", key) { + t.Fatalf("Store should contain fresh key") + } + + now = now.Add(11 * time.Second) + if store.Has("alice", key) { + t.Fatalf("Store should expire key when ttl is exceeded") + } + + store.mu.RLock() + defer store.mu.RUnlock() + if len(store.keysByUser["alice"]) != 0 { + t.Fatalf("Expired entries should be removed on Has call") + } +} + +func TestAuthKeyStoreEnforcesPerUserKeyLimit(t *testing.T) { + now := time.Date(2026, 3, 3, 10, 0, 0, 0, time.UTC) + store := newAuthKeyStoreWithClock(time.Hour, 2, func() time.Time { return now }) + + keyOne := testPublicKey(t, 3) + keyTwo := testPublicKey(t, 4) + keyThree := testPublicKey(t, 5) + + store.Add("alice", keyOne) + now = now.Add(1 * time.Second) + store.Add("alice", keyTwo) + now = now.Add(1 * time.Second) + store.Add("alice", keyThree) + + if store.Has("alice", keyOne) { + t.Fatalf("Oldest key should be evicted once max key limit is reached") + } + if !store.Has("alice", keyTwo) { + t.Fatalf("Second key should remain in store") + } + if !store.Has("alice", keyThree) { + t.Fatalf("Newest key should remain in store") + } +} + +func TestAuthKeyStoreAddRefreshesExistingKey(t *testing.T) { + now := time.Date(2026, 3, 3, 10, 0, 0, 0, time.UTC) + store := newAuthKeyStoreWithClock(10*time.Second, 5, func() time.Time { return now }) + key := testPublicKey(t, 6) + + store.Add("alice", key) + now = now.Add(9 * time.Second) + store.Add("alice", key) + + now = now.Add(5 * time.Second) + if !store.Has("alice", key) { + t.Fatalf("Key should stay valid after it is refreshed") + } + + now = now.Add(6 * time.Second) + if store.Has("alice", key) { + t.Fatalf("Refreshed key should expire once ttl is exceeded from latest add") + } +} + +func TestAuthKeyStoreUserIsolation(t *testing.T) { + store := NewAuthKeyStore(time.Hour, 5) + key := testPublicKey(t, 7) + + store.Add("alice", key) + if store.Has("bob", key) { + t.Fatalf("Key lookup must be isolated by user") + } +} + +func TestAuthKeyStoreIgnoresInvalidInput(t *testing.T) { + store := NewAuthKeyStore(time.Hour, 5) + key := testPublicKey(t, 8) + + store.Add("", key) + store.Add("alice", nil) + store.Remove("", key) + store.Remove("alice", nil) + + if store.Has("", key) { + t.Fatalf("Empty user should not match") + } + if store.Has("alice", nil) { + t.Fatalf("Nil key should not match") + } +} + +func TestAuthKeyStoreConcurrentAccess(t *testing.T) { + store := NewAuthKeyStore(time.Hour, 5) + users := []string{"alice", "bob", "carol"} + keys := []gossh.PublicKey{ + testPublicKey(t, 11), + testPublicKey(t, 12), + testPublicKey(t, 13), + testPublicKey(t, 14), + } + + var wg sync.WaitGroup + for worker := 0; worker < 32; worker++ { + wg.Add(1) + go func(workerID int) { + defer wg.Done() + + user := users[workerID%len(users)] + for i := 0; i < 200; i++ { + key := keys[(workerID+i)%len(keys)] + store.Add(user, key) + _ = store.Has(user, key) + if i%3 == 0 { + store.Remove(user, key) + } + } + }(worker) + } + wg.Wait() + + store.mu.RLock() + defer store.mu.RUnlock() + for user, userEntries := range store.keysByUser { + if len(userEntries) > store.maxKeysPerUser { + t.Fatalf("User %s exceeded max key limit: %d", user, len(userEntries)) + } + } +} + +func testPublicKey(t *testing.T, seedByte byte) gossh.PublicKey { + t.Helper() + + seed := make([]byte, ed25519.SeedSize) + for i := range seed { + seed[i] = seedByte + } + + privateKey := ed25519.NewKeyFromSeed(seed) + publicKey, err := gossh.NewPublicKey(privateKey.Public()) + if err != nil { + t.Fatalf("Unable to build ssh public key: %s", err.Error()) + } + + return publicKey +} diff --git a/internal/ssh/server/hostkey.go b/internal/ssh/server/hostkey.go index b2d4569..1315351 100644 --- a/internal/ssh/server/hostkey.go +++ b/internal/ssh/server/hostkey.go @@ -1,41 +1,75 @@ package server import ( - "os" + "errors" + iofs "io/fs" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/io/fs" "github.com/mimecast/dtail/internal/ssh" ) +const ( + defaultHostKeyBits = 4096 + defaultHostKeyFile = "./cache/ssh_host_key" +) + // PrivateHostKey retrieves the private server RSA host key. -func PrivateHostKey() []byte { - hostKeyFile := config.Server.HostKeyFile +func PrivateHostKey(hostKeyFile string, hostKeyBits int) []byte { + if hostKeyFile == "" { + hostKeyFile = defaultHostKeyFile + } + if hostKeyBits <= 0 { + hostKeyBits = defaultHostKeyBits + } if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") { hostKeyFile = "./ssh_host_key" } - _, err := os.Stat(hostKeyFile) - - if os.IsNotExist(err) { - dlog.Server.Info("Generating private server RSA host key") - privateKey, err := ssh.GeneratePrivateRSAKey(config.Server.HostKeyBits) - - if err != nil { - dlog.Server.FatalPanic("Failed to generate private server RSA host key", err) - } + hostKeyPath, err := fs.NewRootedPath(hostKeyFile) + if err != nil { + dlog.Server.FatalPanic("Invalid private server RSA host key path", hostKeyFile, err) + } - pem := ssh.EncodePrivateKeyToPEM(privateKey) - if err := os.WriteFile(hostKeyFile, pem, 0600); err != nil { - dlog.Server.Error("Unable to write private server RSA host key to file", - hostKeyFile, err) + _, err = hostKeyPath.Stat() + if err != nil { + // os.IsNotExist does not unwrap fmt.Errorf chains from RootedPath.Stat; use errors.Is. + if errors.Is(err, iofs.ErrNotExist) { + dlog.Server.Info("Generating private server RSA host key") + pem, genErr := generatePrivateHostKey(hostKeyBits) + if genErr != nil { + dlog.Server.FatalPanic("Failed to generate private server RSA host key", genErr) + } + if storeErr := storePrivateHostKey(hostKeyPath, pem); storeErr != nil { + dlog.Server.Error("Unable to write private server RSA host key to file", + hostKeyFile, storeErr) + } + return pem } - return pem + dlog.Server.FatalPanic("Cannot stat private server RSA host key path", hostKeyFile, err) } dlog.Server.Info("Reading private server RSA host key from file", hostKeyFile) - pem, err := os.ReadFile(hostKeyFile) + pem, err := readPrivateHostKey(hostKeyPath) if err != nil { dlog.Server.FatalPanic("Failed to load private server RSA host key", err) } return pem } + +func generatePrivateHostKey(hostKeyBits int) ([]byte, error) { + privateKey, err := ssh.GeneratePrivateRSAKey(hostKeyBits) + if err != nil { + return nil, err + } + + return ssh.EncodePrivateKeyToPEM(privateKey), nil +} + +func storePrivateHostKey(hostKeyPath fs.RootedPath, pem []byte) error { + return hostKeyPath.WriteFile(pem, 0o600) +} + +func readPrivateHostKey(hostKeyPath fs.RootedPath) ([]byte, error) { + return hostKeyPath.ReadFile() +} diff --git a/internal/ssh/server/hostkey_test.go b/internal/ssh/server/hostkey_test.go new file mode 100644 index 0000000..e318fb4 --- /dev/null +++ b/internal/ssh/server/hostkey_test.go @@ -0,0 +1,37 @@ +package server + +import ( + "bytes" + "github.com/mimecast/dtail/internal/io/fs" + "os" + "path/filepath" + "testing" +) + +func TestPrivateHostKeyGeneratesAndReloadsExistingKey(t *testing.T) { + hostKeyFile := filepath.Join(t.TempDir(), "cache", "ssh_host_key") + if err := os.MkdirAll(filepath.Dir(hostKeyFile), 0o755); err != nil { + t.Fatalf("MkdirAll failed: %v", err) + } + + hostKeyPath, err := fs.NewRootedPath(hostKeyFile) + if err != nil { + t.Fatalf("NewRootedPath failed: %v", err) + } + + firstPEM, err := generatePrivateHostKey(1024) + if err != nil { + t.Fatalf("generatePrivateHostKey failed: %v", err) + } + if err := storePrivateHostKey(hostKeyPath, firstPEM); err != nil { + t.Fatalf("storePrivateHostKey failed: %v", err) + } + + secondPEM, err := readPrivateHostKey(hostKeyPath) + if err != nil { + t.Fatalf("readPrivateHostKey failed: %v", err) + } + if !bytes.Equal(secondPEM, firstPEM) { + t.Fatalf("readPrivateHostKey returned different key data") + } +} diff --git a/internal/ssh/server/publickeycallback.go b/internal/ssh/server/publickeycallback.go index bcc9004..8a23384 100644 --- a/internal/ssh/server/publickeycallback.go +++ b/internal/ssh/server/publickeycallback.go @@ -1,38 +1,64 @@ package server import ( + "bytes" + "errors" "fmt" + iofs "io/fs" "os" goUser "os/user" + "path/filepath" "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/io/fs" user "github.com/mimecast/dtail/internal/user/server" gossh "golang.org/x/crypto/ssh" ) -// PublicKeyCallback is for the server to check whether a public SSH key is -// authorized ot not. -func PublicKeyCallback(c gossh.ConnMetadata, - offeredPubKey gossh.PublicKey) (*gossh.Permissions, error) { +type authorizedKeyParser func([]byte) (gossh.PublicKey, string, []string, []byte, error) + +// NewPublicKeyCallback creates an instance-scoped SSH public key callback. +// keyStore must be non-nil; callers are responsible for constructing and +// wiring the store. There is no shared package-level fallback. +func NewPublicKeyCallback(authKeyEnabled bool, cacheDir string, + keyStore *AuthKeyStore) func(gossh.ConnMetadata, gossh.PublicKey) (*gossh.Permissions, error) { - user, err := user.New(c.User(), c.RemoteAddr().String()) + if keyStore == nil { + panic("NewPublicKeyCallback: keyStore must not be nil") + } + return func(c gossh.ConnMetadata, offeredPubKey gossh.PublicKey) (*gossh.Permissions, error) { + return publicKeyCallback(c, offeredPubKey, authKeyEnabled, cacheDir, keyStore) + } +} + +func publicKeyCallback(c gossh.ConnMetadata, offeredPubKey gossh.PublicKey, + authKeyEnabled bool, cacheDir string, keyStore *AuthKeyStore) (*gossh.Permissions, error) { + + user, err := user.New(c.User(), c.RemoteAddr().String(), nil) if err != nil { return nil, err } dlog.Server.Info(user, "Incoming authorization") - authorizedKeysFile, err := authorizedKeysFile(user) + if authKeyEnabled { + if permissions := authKeyStorePermissions(keyStore, user.Name, offeredPubKey); permissions != nil { + dlog.Server.Info(user, "Authorized by in-memory auth key store") + return permissions, nil + } + } + + authorizedKeysPath, err := authorizedKeysPathForUser(user, cacheDir) if err != nil { return nil, err } - dlog.Server.Info(user, "Reading", authorizedKeysFile) - authorizedKeysBytes, err := os.ReadFile(authorizedKeysFile) + dlog.Server.Info(user, "Reading", authorizedKeysPath.Path()) + authorizedKeysBytes, err := authorizedKeysPath.ReadFile() if err != nil { return nil, fmt.Errorf("Unable to read authorized keys file|%s|%s|%s", - authorizedKeysFile, user, err.Error()) + authorizedKeysPath.Path(), user, err.Error()) } return verifyAuthorizedKeys(user, authorizedKeysBytes, offeredPubKey) @@ -40,57 +66,126 @@ func PublicKeyCallback(c gossh.ConnMetadata, func verifyAuthorizedKeys(user *user.User, authorizedKeysBytes []byte, offeredPubKey gossh.PublicKey) (*gossh.Permissions, error) { + return verifyAuthorizedKeysWithParser(user, authorizedKeysBytes, offeredPubKey, gossh.ParseAuthorizedKey) +} + +func verifyAuthorizedKeysWithParser(user *user.User, authorizedKeysBytes []byte, + offeredPubKey gossh.PublicKey, parseAuthorizedKey authorizedKeyParser) (*gossh.Permissions, error) { authorizedKeysMap := map[string]bool{} for len(authorizedKeysBytes) > 0 { - authorizedPubKey, _, _, restBytes, err := gossh.ParseAuthorizedKey(authorizedKeysBytes) + authorizedPubKey, _, _, restBytes, err := parseAuthorizedKey(authorizedKeysBytes) if err != nil { - return nil, fmt.Errorf("unable to parse authorized keys bytes|%s|%s", - user, err.Error()) + if dlog.Server != nil { + dlog.Server.Warn(user, "Skipping unparseable authorized_keys line", err) + } + nextAuthorizedKeysBytes, ok := advanceToNextAuthorizedKeysLine(authorizedKeysBytes) + if !ok { + break + } + authorizedKeysBytes = nextAuthorizedKeysBytes + continue } authorizedKeysMap[string(authorizedPubKey.Marshal())] = true authorizedKeysBytes = restBytes - dlog.Server.Debug(user, "Authorized public key fingerprint", - gossh.FingerprintSHA256(authorizedPubKey)) + if dlog.Server != nil { + dlog.Server.Debug(user, "Authorized public key fingerprint", + gossh.FingerprintSHA256(authorizedPubKey)) + } } - dlog.Server.Debug(user, "Offered public key fingerprint", gossh.FingerprintSHA256(offeredPubKey)) + if dlog.Server != nil { + dlog.Server.Debug(user, "Offered public key fingerprint", gossh.FingerprintSHA256(offeredPubKey)) + } if authorizedKeysMap[string(offeredPubKey.Marshal())] { - return &gossh.Permissions{ - Extensions: map[string]string{"pubkey-fp": gossh.FingerprintSHA256(offeredPubKey)}, - }, nil + return permissionsFromPublicKey(offeredPubKey), nil } return nil, fmt.Errorf("%s|public key of user not authorized", user) } -func authorizedKeysFile(user *user.User) (string, error) { +func advanceToNextAuthorizedKeysLine(authorizedKeysBytes []byte) ([]byte, bool) { + lineEnd := bytes.IndexByte(authorizedKeysBytes, '\n') + if lineEnd == -1 { + return nil, false + } + + nextBytes := authorizedKeysBytes[lineEnd+1:] + return nextBytes, true +} + +func authKeyStorePermissions(keyStore *AuthKeyStore, userName string, + offeredPubKey gossh.PublicKey) *gossh.Permissions { + + if keyStore == nil || !keyStore.Has(userName, offeredPubKey) { + return nil + } + + return permissionsFromPublicKey(offeredPubKey) +} + +func permissionsFromPublicKey(offeredPubKey gossh.PublicKey) *gossh.Permissions { + return &gossh.Permissions{ + Extensions: map[string]string{"pubkey-fp": gossh.FingerprintSHA256(offeredPubKey)}, + } +} + +type userLookupFunc func(string) (*goUser.User, error) + +func authorizedKeysPathForUser(user *user.User, cacheDir string) (fs.RootedPath, error) { if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") { // In this case, we expect a pub key in the current directory. - return "./id_rsa.pub", nil + return fs.NewRootedPath("./id_rsa.pub") } cwd, err := os.Getwd() if err != nil { - return "", err + return fs.RootedPath{}, err } - // Check for cached version in the dserver directory. - authorizedKeysFile := fmt.Sprintf("%s/%s/%s.authorized_keys", cwd, - config.Common.CacheDir, user.Name) - if _, err = os.Stat(authorizedKeysFile); err == nil { - return authorizedKeysFile, nil + return findAuthorizedKeysPath(user, cacheDir, cwd, goUser.Lookup) +} + +func findAuthorizedKeysPath(user *user.User, cacheDir, cwd string, + lookupUser userLookupFunc) (fs.RootedPath, error) { + + // Check for cached version in the dserver directory. An absolute + // CacheDir (as used by the BSD packages, e.g. /var/run/dserver/cache) + // must be used as-is: joining it with the CWD would break auth + // whenever dserver is started from a directory other than / (e.g. a + // manual rc.d restart from a home directory). Relative CacheDirs stay + // relative to the CWD as before. + if cacheDir != "" { + cacheBase := cacheDir + if !filepath.IsAbs(cacheDir) { + cacheBase = filepath.Join(cwd, cacheDir) + } + cachePath := filepath.Join(cacheBase, fmt.Sprintf("%s.authorized_keys", user.Name)) + rootedCachePath, err := fs.NewRootedPath(cachePath) + if err != nil { + return fs.RootedPath{}, err + } + if _, err := rootedCachePath.Stat(); err == nil { + return rootedCachePath, nil + } } // As the last option, check the regular SSH path. - osUser, err := goUser.Lookup(user.Name) + osUser, err := lookupUser(user.Name) + if err != nil { + return fs.RootedPath{}, err + } + authorizedKeysPath := filepath.Join(osUser.HomeDir, ".ssh", "authorized_keys") + rootedAuthorizedKeysPath, err := fs.NewRootedPath(authorizedKeysPath) if err != nil { - return "", err + return fs.RootedPath{}, err + } + if _, err = rootedAuthorizedKeysPath.Stat(); err == nil { + return rootedAuthorizedKeysPath, nil } - authorizedKeysFile = fmt.Sprintf("%s/.ssh/authorized_keys", osUser.HomeDir) - if _, err = os.Stat(authorizedKeysFile); err == nil { - return authorizedKeysFile, nil + if !errors.Is(err, iofs.ErrNotExist) { + return fs.RootedPath{}, err } - return "", fmt.Errorf("unable to find a any authorized keys file") + return fs.RootedPath{}, fmt.Errorf("unable to find any authorized keys file") } diff --git a/internal/ssh/server/publickeycallback_test.go b/internal/ssh/server/publickeycallback_test.go new file mode 100644 index 0000000..2f597cb --- /dev/null +++ b/internal/ssh/server/publickeycallback_test.go @@ -0,0 +1,297 @@ +package server + +import ( + "bytes" + "errors" + "os" + goUser "os/user" + "path/filepath" + "testing" + "time" + + serveruser "github.com/mimecast/dtail/internal/user/server" + + gossh "golang.org/x/crypto/ssh" +) + +func TestAuthKeyStorePermissions(t *testing.T) { + // Create an isolated store for this test — there is no package-level global. + store := NewAuthKeyStore(time.Hour, 5) + + key := testPublicKey(t, 21) + + if