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/clients/connectors/serverconnection.go | 248 +++++++++++++++++++++--- 1 file changed, 218 insertions(+), 30 deletions(-) (limited to 'internal/clients/connectors/serverconnection.go') diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go index 5c3d455..b98f815 100644 --- a/internal/clients/connectors/serverconnection.go +++ b/internal/clients/connectors/serverconnection.go @@ -2,20 +2,38 @@ package connectors import ( "context" + "encoding/base64" "fmt" "io" + "net" + "os" + "path/filepath" "strconv" "strings" + "sync" "time" "github.com/mimecast/dtail/internal/clients/handlers" - "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/protocol" + sessionspec "github.com/mimecast/dtail/internal/session" "github.com/mimecast/dtail/internal/ssh/client" "golang.org/x/crypto/ssh" ) +// SSHSettings provides the connection settings needed by ServerConnection. +type SSHSettings interface { + SSHPort() int + SSHConnectTimeout() time.Duration +} + +const ( + defaultSSHConnectTimeout = 2 * time.Second + defaultSSHPort = 2222 + defaultCapabilityWait = 250 * time.Millisecond +) + // ServerConnection represents a connection to a single remote dtail server via // SSH protocol. type ServerConnection struct { @@ -28,30 +46,61 @@ type ServerConnection struct { config *ssh.ClientConfig handler handlers.Handler commands []string + sessionSpec sessionspec.Spec + sessionState committedSessionState + interactive bool + authKeyPath string + authKeyDisabled bool hostKeyCallback client.HostKeyCallback - throttlingDone bool + // throttleReleased ensures the throttle slot is returned to throttleCh + // exactly once, even if both the early-release path in handle() and the + // deferred cleanup in Start() execute concurrently or the same goroutine + // hits both paths. sync.Once is safe for concurrent callers whereas the + // previous bool guard was not synchronized. + throttleReleased sync.Once } +var _ Connector = (*ServerConnection)(nil) + // NewServerConnection returns a new DTail SSH server connection. func NewServerConnection(server string, userName string, authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback, - handler handlers.Handler, commands []string) *ServerConnection { + handler handlers.Handler, commands []string, sessionSpec sessionspec.Spec, + interactive bool, authKeyPath string, authKeyDisabled bool, settings SSHSettings) *ServerConnection { dlog.Client.Debug(server, "Creating new connection", server, handler, commands) + sshConnectTimeout := defaultSSHConnectTimeout + defaultPort := defaultSSHPort + if settings != nil { + sshConnectTimeout = settings.SSHConnectTimeout() + defaultPort = settings.SSHPort() + } + if sshConnectTimeout <= 0 { + sshConnectTimeout = defaultSSHConnectTimeout + } + if defaultPort <= 0 { + defaultPort = defaultSSHPort + } + c := ServerConnection{ hostKeyCallback: hostKeyCallback, server: server, handler: handler, commands: commands, + sessionSpec: sessionSpec, + interactive: interactive, + authKeyPath: resolveAuthKeyPath(authKeyPath), + authKeyDisabled: authKeyDisabled, config: &ssh.ClientConfig{ - User: userName, - Auth: authMethods, - HostKeyCallback: hostKeyCallback.Wrap(), - Timeout: time.Second * 2, + User: userName, + Auth: authMethods, + Timeout: sshConnectTimeout, + // HostKeyCallback is assigned per-handshake in dial() so the + // callback can honour the handshake's context (see dial()). }, } - c.initServerPort() + c.initServerPort(defaultPort) return &c } @@ -61,12 +110,42 @@ func (c *ServerConnection) Server() string { return c.server } // Handler returns the handler used for the connection. func (c *ServerConnection) Handler() handlers.Handler { return c.handler } +// SupportsQueryUpdates reports whether the remote server advertised the +// runtime query replacement capability. Older servers simply time out and +// return false here without affecting the legacy command path. +func (c *ServerConnection) SupportsQueryUpdates(timeout time.Duration) bool { + return supportsQueryUpdates(c.handler, timeout) +} + +// ApplySessionSpec starts or updates the interactive session state on the +// existing SSH connection when runtime query updates are supported. +func (c *ServerConnection) ApplySessionSpec(spec sessionspec.Spec, timeout time.Duration) error { + return applySessionSpec(c.server, c.handler, &c.sessionState, spec, timeout) +} + +// ApplySessionSpecWithGeneration starts or updates the interactive session +// state using an explicit committed generation as the base for the update. +func (c *ServerConnection) ApplySessionSpecWithGeneration(spec sessionspec.Spec, generation uint64, timeout time.Duration) error { + return applySessionSpecWithGeneration(c.server, c.handler, &c.sessionState, spec, generation, false, timeout) +} + +// CommittedSession returns the last server-acknowledged session state. +func (c *ServerConnection) CommittedSession() (sessionspec.Spec, uint64, bool) { + return c.sessionState.snapshot() +} + +// RestoreCommittedSession resets the local session snapshot without advancing +// the generation. +func (c *ServerConnection) RestoreCommittedSession(spec sessionspec.Spec, generation uint64, committed bool) { + c.sessionState.restore(spec, generation, committed) +} + // Attempt to parse the server port address from the provided server FQDN. -func (c *ServerConnection) initServerPort() { +func (c *ServerConnection) initServerPort(defaultPort int) { parts := strings.Split(c.server, ":") if len(parts) == 1 { c.hostname = c.server - c.port = config.Common.SSHPort + c.port = defaultPort return } @@ -99,12 +178,15 @@ func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc, go func() { defer func() { - if !c.throttlingDone { - dlog.Client.Debug(c.server, "Unthrottling connection (1)", + // Release the throttle slot on the way out regardless of which + // code path already attempted it. throttleReleased.Do guarantees + // the drain happens exactly once even if handle() already released + // the slot early (the fast path when the session is fully up). + c.throttleReleased.Do(func() { + dlog.Client.Debug(c.server, "Unthrottling connection (cleanup)", len(throttleCh), cap(throttleCh)) - c.throttlingDone = true <-throttleCh - } + }) cancel() }() @@ -133,10 +215,33 @@ func (c *ServerConnection) dial(ctx context.Context, cancel context.CancelFunc, address := fmt.Sprintf("%s:%d", c.hostname, c.port) dlog.Client.Debug(c.server, "Dialing into the connection", address) - client, err := ssh.Dial("tcp", address, c.config) + // Use context-aware dialing to enable proper cancellation during connection establishment. + // TCP KeepAlive (30s) prevents silent connection failures on long-lived connections. + dialer := &net.Dialer{ + Timeout: c.config.Timeout, // Use the SSH config timeout (2 seconds) + KeepAlive: 30 * time.Second, // Standard Go default for connection health monitoring + } + + // Establish TCP connection with context support for cancellation + conn, err := dialer.DialContext(ctx, "tcp", address) if err != nil { - return err + return fmt.Errorf("failed to dial TCP connection to %s: %w", address, err) } + + // Perform SSH handshake over the established TCP connection. Build a + // per-handshake ssh.ClientConfig so the host-key callback is bound to + // ctx and unblocks cleanly if the handshake is cancelled (e.g. when the + // user aborts before responding to the unknown-host prompt). + handshakeConfig := *c.config + handshakeConfig.HostKeyCallback = c.hostKeyCallback.Wrap(ctx) + sshConn, chans, reqs, err := ssh.NewClientConn(conn, address, &handshakeConfig) + if err != nil { + conn.Close() + return fmt.Errorf("SSH handshake failed for %s: %w", address, err) + } + + // Create SSH client from the connection components + client := ssh.NewClient(sshConn, chans, reqs) defer client.Close() return c.session(ctx, cancel, client, throttleCh) @@ -149,7 +254,7 @@ func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFun dlog.Client.Debug(c.server, "Creating SSH session") session, err := client.NewSession() if err != nil { - return err + return fmt.Errorf("failed to create SSH session for %s: %w", c.server, err) } defer session.Close() return c.handle(ctx, cancel, session, throttleCh) @@ -161,14 +266,14 @@ func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc dlog.Client.Debug(c.server, "Creating handler for SSH session") stdinPipe, err := session.StdinPipe() if err != nil { - return err + return fmt.Errorf("failed to get SSH session stdin pipe for %s: %w", c.server, err) } stdoutPipe, err := session.StdoutPipe() if err != nil { - return err + return fmt.Errorf("failed to get SSH session stdout pipe for %s: %w", c.server, err) } if err := session.Shell(); err != nil { - return err + return fmt.Errorf("failed to start SSH shell for %s: %w", c.server, err) } go func() { @@ -191,22 +296,105 @@ func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc } }() - // Send all commands to client. - for _, command := range c.commands { - dlog.Client.Debug(command) - if err := c.handler.SendMessage(command); err != nil { - dlog.Client.Debug(err) - } + if c.authKeyDisabled { + dlog.Client.Debug(c.server, "Skipping AUTHKEY registration because auth-key is disabled") + } else { + c.sendAuthKeyRegistrationCommand() + } + + if err := dispatchInitialCommands(c.server, c.handler, c.commands, c.interactive, c.sessionSpec, &c.sessionState); err != nil { + c.handler.Shutdown() + return err } - if !c.throttlingDone { - dlog.Client.Debug(c.server, "Unthrottling connection (2)", + // Release the throttle slot as soon as the session is fully established so + // the next pending connection can proceed without waiting for this session + // to finish. throttleReleased.Do is idempotent: if the deferred cleanup + // in Start() fires first (e.g. on a dial error path that never reaches + // here), the slot is still returned exactly once. + c.throttleReleased.Do(func() { + dlog.Client.Debug(c.server, "Unthrottling connection (session up)", len(throttleCh), cap(throttleCh)) - c.throttlingDone = true <-throttleCh - } + }) <-ctx.Done() c.handler.Shutdown() return nil } + +// resolveAuthKeyPath returns the effective auth-key path. When the provided +// path is non-empty it is used as-is. Otherwise the function falls back to +// $HOME/.ssh/id_rsa. If HOME is also empty it returns "" so that the AUTHKEY +// registration step (sendAuthKeyRegistrationCommand) will skip gracefully +// instead of trying to open a path that the SSH library cannot expand. +func resolveAuthKeyPath(authKeyPath string) string { + if strings.TrimSpace(authKeyPath) != "" { + return authKeyPath + } + homeDir := os.Getenv("HOME") + if homeDir == "" { + return "" + } + return filepath.Join(homeDir, ".ssh", "id_rsa") +} + +func (c *ServerConnection) sendAuthKeyRegistrationCommand() { + authKeyPubPath := c.authKeyPath + ".pub" + authKeyPubBytes, err := os.ReadFile(authKeyPubPath) + if err != nil { + dlog.Client.Debug(c.server, "Skipping AUTHKEY registration, unable to read public key", authKeyPubPath, err) + return + } + + authKeyBase64, err := extractAuthKeyBase64(authKeyPubBytes) + if err != nil { + dlog.Client.Debug(c.server, "Skipping AUTHKEY registration, invalid public key file", authKeyPubPath, err) + return + } + + if err := c.handler.SendMessage("AUTHKEY " + authKeyBase64); err != nil { + dlog.Client.Debug(c.server, "Unable to send AUTHKEY registration command", err) + return + } + dlog.Client.Debug(c.server, "Sent AUTHKEY registration command", authKeyPubPath) +} + +func extractAuthKeyBase64(authKeyPubBytes []byte) (string, error) { + authKeyPubContent := string(authKeyPubBytes) + for _, line := range strings.Split(authKeyPubContent, "\n") { + trimmedLine := strings.TrimSpace(line) + if trimmedLine == "" || strings.HasPrefix(trimmedLine, "#") { + continue + } + + fields := strings.Fields(trimmedLine) + if len(fields) < 2 { + return "", fmt.Errorf("expected authorized key format ' [comment]'") + } + + authKeyBase64 := strings.TrimSpace(fields[1]) + if _, err := base64.StdEncoding.DecodeString(authKeyBase64); err != nil { + return "", fmt.Errorf("invalid base64 public key: %w", err) + } + + return authKeyBase64, nil + } + + return "", fmt.Errorf("no public key found") +} + +func supportsQueryUpdates(handler handlers.Handler, timeout time.Duration) bool { + if handler == nil { + return false + } + + if timeout <= 0 { + timeout = defaultCapabilityWait + } + if !handler.WaitForCapabilities(timeout) { + return false + } + + return handler.HasCapability(protocol.CapabilityQueryUpdateV1) +} -- cgit v1.2.3