summaryrefslogtreecommitdiff
path: root/internal/ssh/client
diff options
context:
space:
mode:
Diffstat (limited to 'internal/ssh/client')
-rw-r--r--internal/ssh/client/authmethods.go178
-rw-r--r--internal/ssh/client/authmethods_test.go194
-rw-r--r--internal/ssh/client/customkeycallback.go6
-rw-r--r--internal/ssh/client/hostkeycallback.go7
-rw-r--r--internal/ssh/client/knownhostscallback.go205
-rw-r--r--internal/ssh/client/knownhostscallback_test.go305
-rw-r--r--internal/ssh/client/simplecallback.go5
7 files changed, 780 insertions, 120 deletions
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 :