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/baseclient.go | 213 ++++- internal/clients/baseclient_retry_test.go | 280 +++++++ internal/clients/catclient.go | 24 +- internal/clients/client_benchmark_test.go | 136 ++++ internal/clients/connectors/connector.go | 18 + internal/clients/connectors/serverconnection.go | 248 +++++- .../clients/connectors/serverconnection_test.go | 880 +++++++++++++++++++++ internal/clients/connectors/serverless.go | 213 ++++- internal/clients/connectors/sessiontransport.go | 213 +++++ internal/clients/grepclient.go | 24 +- internal/clients/handlers/basehandler.go | 319 +++++++- internal/clients/handlers/basehandler_test.go | 356 +++++++++ internal/clients/handlers/clienthandler.go | 15 +- internal/clients/handlers/handler.go | 6 + internal/clients/handlers/healthhandler.go | 13 +- internal/clients/handlers/maprhandler.go | 55 +- internal/clients/handlers/maprhandler_test.go | 278 +++++++ internal/clients/healthclient.go | 22 +- internal/clients/interactive_control.go | 425 ++++++++++ internal/clients/interactive_control_test.go | 641 +++++++++++++++ internal/clients/maker.go | 8 + internal/clients/maprclient.go | 269 ++++--- internal/clients/maprclient_test.go | 153 ++++ internal/clients/query_regex.go | 30 + internal/clients/runtime_boundary.go | 219 +++++ internal/clients/runtime_boundary_test.go | 89 +++ internal/clients/session_spec.go | 14 + internal/clients/session_spec_test.go | 147 ++++ internal/clients/stats.go | 28 +- internal/clients/tailclient.go | 25 +- 30 files changed, 5111 insertions(+), 250 deletions(-) create mode 100644 internal/clients/baseclient_retry_test.go create mode 100644 internal/clients/client_benchmark_test.go create mode 100644 internal/clients/connectors/serverconnection_test.go create mode 100644 internal/clients/connectors/sessiontransport.go create mode 100644 internal/clients/handlers/basehandler_test.go create mode 100644 internal/clients/handlers/maprhandler_test.go create mode 100644 internal/clients/interactive_control.go create mode 100644 internal/clients/interactive_control_test.go create mode 100644 internal/clients/maprclient_test.go create mode 100644 internal/clients/query_regex.go create mode 100644 internal/clients/runtime_boundary.go create mode 100644 internal/clients/runtime_boundary_test.go create mode 100644 internal/clients/session_spec.go create mode 100644 internal/clients/session_spec_test.go (limited to 'internal/clients') diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go index 013f2f2..165fbd5 100644 --- a/internal/clients/baseclient.go +++ b/internal/clients/baseclient.go @@ -2,6 +2,8 @@ package clients import ( "context" + "io" + "math/rand" "sync" "time" @@ -15,29 +17,52 @@ import ( gossh "golang.org/x/crypto/ssh" ) +const ( + initialRetryDelay = 2 * time.Second + maxRetryDelay = 60 * time.Second + retryJitterFactor = 0.2 // +/-20% jitter to avoid synchronized reconnect storms. +) + // This is the main client data structure. type baseClient struct { + mu *sync.RWMutex config.Args + runtime *clientRuntimeBoundary // To display client side stats stats *stats // We have one connection per remote server. connections []connectors.Connector // SSH auth methods to use to connect to the remote servers. sshAuthMethods []gossh.AuthMethod + // authCloser owns any ssh-agent connection acquired while building the + // auth methods; it must be closed once all SSH handshakes that consume + // sshAuthMethods have completed. + authCloser io.Closer // To deal with SSH host keys hostKeyCallback client.HostKeyCallback // Throttle how fast we initiate SSH connections concurrently throttleCh chan struct{} // Retry connection upon failure? retry bool + // The current connection-wide session specification. + sessionSpec SessionSpec // Connection maker helper. maker maker + // Optional factory override for retry/reconnect tests. + connectionFactory func(server string, authMethods []gossh.AuthMethod, + hostKeyCallback client.HostKeyCallback, sessionSpec SessionSpec, + interactive bool) connectors.Connector + // Optional sleep override for retry tests. + sleepFn func(context.Context, time.Duration) bool // Regex is the regular expresion object for line filtering Regex regex.Regex } func (c *baseClient) init() { dlog.Client.Debug("Initiating base client", c.Args.String()) + if c.runtime == nil { + c.runtime = newClientRuntimeBoundary(config.CurrentRuntime()) + } flag := regex.Default if c.Args.RegexInvert { @@ -52,25 +77,52 @@ func (c *baseClient) init() { if c.Args.Serverless { return } - c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods( + c.sshAuthMethods, c.hostKeyCallback, c.authCloser = client.InitSSHAuthMethods( c.Args.SSHAuthMethods, c.Args.SSHHostKeyCallback, c.Args.TrustAllHosts, - c.throttleCh, c.Args.SSHPrivateKeyFilePath) + c.Args.SSHPrivateKeyFilePath, c.Args.SSHAgentKeyIndex) } -func (c *baseClient) makeConnections(maker maker) { +func (c *baseClient) makeConnections(maker maker) error { c.maker = maker + if builder, ok := maker.(sessionSpecMaker); ok { + sessionSpec, err := builder.makeSessionSpec() + if err != nil { + dlog.Client.FatalPanic("unable to build session specification", err) + } + c.sessionSpec = sessionSpec + } - discoveryService := discovery.New(c.Discovery, c.ServersStr, discovery.Shuffle) + discoveryService, err := discovery.New(c.Discovery, c.ServersStr, discovery.Shuffle) + if err != nil { + return err + } for _, server := range discoveryService.ServerList() { c.connections = append(c.connections, c.makeConnection(server, c.sshAuthMethods, c.hostKeyCallback)) } - c.stats = newTailStats(len(c.connections)) + c.stats = newTailStats(len(c.connections), c.runtime.output, c.runtime.InterruptPause()) + return nil } func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status int) { + if c.Args.InteractiveQuery { + return c.startInteractiveControl(ctx, statsCh) + } + return c.runConnections(ctx, statsCh) +} + +func (c *baseClient) runConnections(ctx context.Context, statsCh <-chan string) (status int) { dlog.Client.Trace("Starting base client") + // Release the ssh-agent connection (if any) once all handshakes and + // reconnect attempts that consume c.sshAuthMethods have finished. + if c.authCloser != nil { + defer func() { + if err := c.authCloser.Close(); err != nil { + dlog.Client.Debug("baseClient", "failed to close ssh-agent connection", err) + } + }() + } // Can be nil when serverless. if c.hostKeyCallback != nil { // Periodically check for unknown hosts, and ask the user whether to trust them or not. @@ -80,10 +132,11 @@ func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status i go c.stats.Start(ctx, c.throttleCh, statsCh, c.Args.Quiet) var wg sync.WaitGroup - wg.Add(len(c.connections)) + connections := c.snapshotConnections() + wg.Add(len(connections)) var mutex sync.Mutex - for i, conn := range c.connections { + for i, conn := range connections { go func(i int, conn connectors.Connector) { defer wg.Done() connStatus := c.startConnection(ctx, i, conn) @@ -102,11 +155,14 @@ func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status i func (c *baseClient) startConnection(ctx context.Context, i int, conn connectors.Connector) (status int) { + retryDelay := initialRetryDelay + retryRandom := newRetryRandom(i) + for { connCtx, cancel := context.WithCancel(ctx) - defer cancel() conn.Start(connCtx, cancel, c.throttleCh, c.stats.connectionsEstCh) + cancel() // Retrieve status code from handler (dtail client will exit with that status) status = conn.Handler().Status() @@ -122,20 +178,147 @@ func (c *baseClient) startConnection(ctx context.Context, i int, default: } - // Yes, we want to retry. - time.Sleep(time.Second * 2) - dlog.Client.Debug(conn.Server(), "Reconnecting") + // Yes, we want to retry with exponential backoff and jitter. + sleepDuration := jitterRetryDelay(retryDelay, retryRandom) + dlog.Client.Debug(conn.Server(), "Reconnecting", "backoff", sleepDuration) + if !c.sleepRetry(ctx, sleepDuration) { + return + } + + retryDelay = nextRetryDelay(retryDelay) conn = c.makeConnection(conn.Server(), c.sshAuthMethods, c.hostKeyCallback) - c.connections[i] = conn + c.replaceConnection(i, conn) + } +} + +func nextRetryDelay(current time.Duration) time.Duration { + if current <= 0 { + return initialRetryDelay + } + + next := current * 2 + if next > maxRetryDelay || next < current { + return maxRetryDelay + } + return next +} + +func jitterRetryDelay(base time.Duration, random *rand.Rand) time.Duration { + if base <= 0 || random == nil { + return base + } + + jitter := time.Duration(float64(base) * retryJitterFactor) + if jitter <= 0 { + return base + } + + minDelay := base - jitter + maxDelay := base + jitter + if maxDelay < minDelay { + return base + } + + return minDelay + time.Duration(random.Int63n(int64(maxDelay-minDelay+1))) +} + +func sleepWithContext(ctx context.Context, delay time.Duration) bool { + if delay <= 0 { + return true } + + timer := time.NewTimer(delay) + defer timer.Stop() + + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} + +func newRetryRandom(seedOffset int) *rand.Rand { + return rand.New(rand.NewSource(time.Now().UnixNano() + int64(seedOffset))) } func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback client.HostKeyCallback) connectors.Connector { - if c.Args.Serverless { + args, sessionSpec := c.snapshotConnectionState() + return c.makeConnectionWithState(server, sshAuthMethods, hostKeyCallback, args, sessionSpec) +} + +func (c *baseClient) makeConnectionWithState(server string, sshAuthMethods []gossh.AuthMethod, + hostKeyCallback client.HostKeyCallback, args config.Args, sessionSpec SessionSpec) connectors.Connector { + if c.connectionFactory != nil { + return c.connectionFactory(server, sshAuthMethods, hostKeyCallback, + sessionSpec, args.InteractiveQuery) + } + if args.Serverless { return connectors.NewServerless(c.UserName, c.maker.makeHandler(server), - c.maker.makeCommands()) + c.maker.makeCommands(), sessionSpec, args.InteractiveQuery, c.runtime) } return connectors.NewServerConnection(server, c.UserName, sshAuthMethods, - hostKeyCallback, c.maker.makeHandler(server), c.maker.makeCommands()) + hostKeyCallback, c.maker.makeHandler(server), c.maker.makeCommands(), + sessionSpec, args.InteractiveQuery, args.SSHPrivateKeyFilePath, + args.NoAuthKey, c.runtime) +} + +func (c *baseClient) sleepRetry(ctx context.Context, delay time.Duration) bool { + if c.sleepFn != nil { + return c.sleepFn(ctx, delay) + } + return sleepWithContext(ctx, delay) +} + +func (c *baseClient) snapshotConnectionState() (config.Args, SessionSpec) { + mu := c.stateMu() + mu.RLock() + defer mu.RUnlock() + + return c.Args, c.sessionSpec +} + +func (c *baseClient) snapshotMutableState() (config.Args, SessionSpec, []connectors.Connector) { + mu := c.stateMu() + mu.RLock() + defer mu.RUnlock() + + return c.Args, c.sessionSpec, append([]connectors.Connector(nil), c.connections...) +} + +func (c *baseClient) snapshotConnections() []connectors.Connector { + mu := c.stateMu() + mu.RLock() + defer mu.RUnlock() + + return append([]connectors.Connector(nil), c.connections...) +} + +func (c *baseClient) storeReloadState(args config.Args, spec SessionSpec) { + mu := c.stateMu() + mu.Lock() + defer mu.Unlock() + + c.Args = args + c.sessionSpec = spec +} + +func (c *baseClient) replaceConnection(i int, conn connectors.Connector) { + mu := c.stateMu() + mu.Lock() + defer mu.Unlock() + + c.connections[i] = conn +} + +func (c *baseClient) stateMu() *sync.RWMutex { + if c.mu == nil { + c.mu = newBaseClientMu() + } + return c.mu +} + +func newBaseClientMu() *sync.RWMutex { + return &sync.RWMutex{} } diff --git a/internal/clients/baseclient_retry_test.go b/internal/clients/baseclient_retry_test.go new file mode 100644 index 0000000..100f171 --- /dev/null +++ b/internal/clients/baseclient_retry_test.go @@ -0,0 +1,280 @@ +package clients + +import ( + "context" + "math/rand" + "sync/atomic" + "testing" + "time" + + "github.com/mimecast/dtail/internal/clients/connectors" + "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/omode" + sshclient "github.com/mimecast/dtail/internal/ssh/client" + + gossh "golang.org/x/crypto/ssh" +) + +func TestNextRetryDelay(t *testing.T) { + tests := []struct { + name string + current time.Duration + want time.Duration + }{ + {name: "zero uses initial", current: 0, want: initialRetryDelay}, + {name: "doubles normally", current: 4 * time.Second, want: 8 * time.Second}, + {name: "caps at max", current: 40 * time.Second, want: maxRetryDelay}, + {name: "stays max at max", current: maxRetryDelay, want: maxRetryDelay}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := nextRetryDelay(tt.current); got != tt.want { + t.Fatalf("nextRetryDelay(%v) = %v, want %v", tt.current, got, tt.want) + } + }) + } +} + +func TestJitterRetryDelayWithinBounds(t *testing.T) { + base := 10 * time.Second + random := rand.New(rand.NewSource(1)) + + min := 8 * time.Second + max := 12 * time.Second + + for i := 0; i < 100; i++ { + got := jitterRetryDelay(base, random) + if got < min || got > max { + t.Fatalf("jitterRetryDelay() = %v, expected between %v and %v", got, min, max) + } + } +} + +func TestSleepWithContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + start := time.Now() + if sleepWithContext(ctx, time.Second) { + t.Fatalf("sleepWithContext should stop when context is canceled") + } + + if time.Since(start) > 100*time.Millisecond { + t.Fatalf("sleepWithContext took too long to exit on canceled context") + } +} + +func TestStartConnectionReconnectsWithLatestSessionSpec(t *testing.T) { + originalLogger := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + dlog.Client = originalLogger + }) + + first := &retryTestConnector{ + server: "srv1", + handler: &retryTestHandler{}, + } + second := &retryTestConnector{ + server: "srv1", + handler: &retryTestHandler{}, + } + + originalSpec := SessionSpec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + updatedSpec := SessionSpec{ + Mode: omode.TailClient, + Files: []string{"/var/log/next.log"}, + Regex: "WARN", + } + + sleepCalls := 0 + var capturedSpec SessionSpec + client := &baseClient{ + mu: newBaseClientMu(), + retry: true, + sessionSpec: originalSpec, + stats: &stats{ + connectionsEstCh: make(chan struct{}, 1), + }, + connections: []connectors.Connector{first}, + connectionFactory: func(server string, _ []gossh.AuthMethod, + _ sshclient.HostKeyCallback, sessionSpec SessionSpec, _ bool) connectors.Connector { + if server != "srv1" { + t.Fatalf("unexpected reconnect server %q", server) + } + capturedSpec = sessionSpec + return second + }, + } + client.sleepFn = func(context.Context, time.Duration) bool { + if sleepCalls == 0 { + sleepCalls++ + client.sessionSpec = updatedSpec + return true + } + return false + } + + status := client.startConnection(context.Background(), 0, first) + if status != 0 { + t.Fatalf("startConnection() status = %d, want 0", status) + } + if capturedSpec.Regex != updatedSpec.Regex || len(capturedSpec.Files) != 1 || capturedSpec.Files[0] != updatedSpec.Files[0] { + t.Fatalf("reconnect used stale session spec: got %#v want %#v", capturedSpec, updatedSpec) + } + if client.connections[0] != second { + t.Fatalf("expected retried connector to replace the original connection") + } +} + +func TestApplyInteractiveReloadConcurrentWithReconnect(t *testing.T) { + resetClientLogger(t) + + originalSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + nextSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/next.log"}, + Regex: "WARN", + } + originalArgs := config.Args{ + Mode: omode.GrepClient, + What: "/var/log/app.log", + RegexStr: "ERROR", + } + nextArgs := config.Args{ + Mode: omode.GrepClient, + What: "/var/log/next.log", + RegexStr: "WARN", + } + + var reconnects atomic.Int32 + client := &baseClient{ + mu: newBaseClientMu(), + Args: originalArgs, + retry: true, + sessionSpec: originalSpec, + stats: &stats{ + connectionsEstCh: make(chan struct{}, 1), + }, + connections: []connectors.Connector{ + newReloadRetryConnector("srv1", originalSpec), + }, + maker: &interactiveReloadMaker{}, + connectionFactory: func(server string, _ []gossh.AuthMethod, + _ sshclient.HostKeyCallback, sessionSpec SessionSpec, _ bool) connectors.Connector { + reconnects.Add(1) + return newReloadRetryConnector(server, sessionSpec) + }, + } + + client.sleepFn = func(context.Context, time.Duration) bool { + return reconnects.Load() < 200 + } + + done := make(chan int, 1) + go func() { + done <- client.startConnection(context.Background(), 0, client.connections[0]) + }() + + for i := 0; i < 200; i++ { + if err := client.applyInteractiveReload(nextArgs, nextSpec); err != nil { + t.Fatalf("applyInteractiveReload() error = %v", err) + } + if err := client.applyInteractiveReload(originalArgs, originalSpec); err != nil { + t.Fatalf("applyInteractiveReload() error = %v", err) + } + } + + if status := <-done; status != 0 { + t.Fatalf("startConnection() status = %d, want 0", status) + } +} + +func newReloadRetryConnector(server string, spec SessionSpec) *reloadRetryConnector { + return &reloadRetryConnector{ + interactiveReloadConnector: interactiveReloadConnector{ + server: server, + supported: true, + committedSpec: spec, + liveSpec: spec, + generation: 1, + }, + handler: &retryTestHandler{}, + } +} + +type reloadRetryConnector struct { + interactiveReloadConnector + handler handlers.Handler +} + +func (c *reloadRetryConnector) Handler() handlers.Handler { return c.handler } + +type retryTestConnector struct { + handler handlers.Handler + server string +} + +func (c *retryTestConnector) Start(context.Context, context.CancelFunc, chan struct{}, chan struct{}) { +} + +func (c *retryTestConnector) Server() string { return c.server } + +func (c *retryTestConnector) Handler() handlers.Handler { return c.handler } + +func (*retryTestConnector) SupportsQueryUpdates(time.Duration) bool { return false } + +func (*retryTestConnector) ApplySessionSpec(SessionSpec, time.Duration) error { return nil } + +func (*retryTestConnector) ApplySessionSpecWithGeneration(SessionSpec, uint64, time.Duration) error { + return nil +} + +func (*retryTestConnector) CommittedSession() (SessionSpec, uint64, bool) { + return SessionSpec{}, 0, false +} + +func (*retryTestConnector) RestoreCommittedSession(SessionSpec, uint64, bool) {} + +type retryTestHandler struct{} + +func (*retryTestHandler) Read([]byte) (int, error) { return 0, nil } + +func (*retryTestHandler) Write(p []byte) (int, error) { return len(p), nil } + +func (*retryTestHandler) Capabilities() []string { return nil } + +func (*retryTestHandler) HasCapability(string) bool { return false } + +func (*retryTestHandler) ReportServerError(string) {} + +func (*retryTestHandler) SendMessage(string) error { return nil } + +func (*retryTestHandler) Server() string { return "srv1" } + +func (*retryTestHandler) Status() int { return 0 } + +func (*retryTestHandler) Shutdown() {} + +func (*retryTestHandler) Done() <-chan struct{} { + done := make(chan struct{}) + close(done) + return done +} + +func (*retryTestHandler) WaitForCapabilities(time.Duration) bool { return false } + +func (*retryTestHandler) WaitForSessionAck(time.Duration) (handlers.SessionAck, bool) { + return handlers.SessionAck{}, false +} diff --git a/internal/clients/catclient.go b/internal/clients/catclient.go index bd65560..e2e247e 100644 --- a/internal/clients/catclient.go +++ b/internal/clients/catclient.go @@ -2,9 +2,7 @@ package clients import ( "errors" - "fmt" "runtime" - "strings" "github.com/mimecast/dtail/internal/clients/handlers" "github.com/mimecast/dtail/internal/config" @@ -26,14 +24,18 @@ func NewCatClient(args config.Args) (*CatClient, error) { c := CatClient{ baseClient: baseClient{ + mu: newBaseClientMu(), Args: args, throttleCh: make(chan struct{}, args.ConnectionsPerCPU*runtime.NumCPU()), retry: false, + runtime: newClientRuntimeBoundary(config.CurrentRuntime()), }, } c.init() - c.makeConnections(c) + if err := c.makeConnections(c); err != nil { + return nil, err + } return &c, nil } @@ -41,14 +43,18 @@ func (c CatClient) makeHandler(server string) handlers.Handler { return handlers.NewClientHandler(server) } +func (c CatClient) makeSessionSpec() (SessionSpec, error) { + return NewSessionSpec(c.Args), nil +} + func (c CatClient) makeCommands() (commands []string) { - regex, err := c.Regex.Serialize() + sessionSpec, err := c.makeSessionSpec() if err != nil { - dlog.Client.FatalPanic(err) + dlog.Client.FatalPanic("unable to build cat session spec", err) } - for _, file := range strings.Split(c.What, ",") { - commands = append(commands, fmt.Sprintf("%s:%s %s %s", - c.Mode.String(), c.Args.SerializeOptions(), file, regex)) + commands, err = sessionSpec.Commands() + if err != nil { + dlog.Client.FatalPanic("unable to build cat commands from session spec", err) } - return + return commands } diff --git a/internal/clients/client_benchmark_test.go b/internal/clients/client_benchmark_test.go new file mode 100644 index 0000000..9423274 --- /dev/null +++ b/internal/clients/client_benchmark_test.go @@ -0,0 +1,136 @@ +package clients + +import ( + "bytes" + "context" + "fmt" + "os" + "path/filepath" + "testing" + "time" + + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/source" + "sync" +) + +func setupBenchmarkData(b *testing.B, lines int) string { + b.Helper() + + tmpDir := b.TempDir() + testFile := filepath.Join(tmpDir, "benchmark_data.log") + + f, err := os.Create(testFile) + if err != nil { + b.Fatalf("Failed to create test file: %v", err) + } + defer f.Close() + + // Create test data + for i := 0; i < lines; i++ { + line := fmt.Sprintf("INFO|1002-071143|1|test.go:%d|8|%d|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=%d|lifetimeConnections=%d|pattern=test-%d|data=%s\n", + i%100, i%50, i%10, i, i%5, "some-test-data-that-makes-the-line-longer") + f.WriteString(line) + } + + return testFile +} + +func BenchmarkDGrep(b *testing.B) { + benchmarkDGrepWithSize(b, 100000) // 100k lines +} + +// Benchmark with different file sizes +func BenchmarkDGrepSmallFile(b *testing.B) { + benchmarkDGrepWithSize(b, 1000) // 1k lines +} + +func BenchmarkDGrepMediumFile(b *testing.B) { + benchmarkDGrepWithSize(b, 50000) // 50k lines +} + +func BenchmarkDGrepLargeFile(b *testing.B) { + benchmarkDGrepWithSize(b, 500000) // 500k lines +} + +func benchmarkDGrepWithSize(b *testing.B, lines int) { + // Setup config. The direct-output read path is the only runtime path. + config.Server = &config.ServerConfig{ + MaxConcurrentCats: 10, + MaxConcurrentTails: 50, + MaxLineLength: 1024 * 1024, + } + + config.Common = &config.CommonConfig{ + Logger: "none", + LogLevel: "error", + } + + config.Client = &config.ClientConfig{ + TermColorsEnable: false, + } + + // Initialize logging + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + wg := &sync.WaitGroup{} + wg.Add(1) + dlog.Start(ctx, wg, source.Client) + + // Create test data + testFile := setupBenchmarkData(b, lines) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + // Create grep client + args := config.Args{ + ServersStr: "serverless", + QueryStr: "", + What: testFile, + RegexStr: "pattern=test-1", + Serverless: true, + Plain: true, + } + + client, err := NewGrepClient(args) + if err != nil { + b.Fatalf("Failed to create grep client: %v", err) + } + + // Capture output + oldStdout := os.Stdout + r, w, _ := os.Pipe() + os.Stdout = w + + // Run grep + statusCh := make(chan int, 1) + go func() { + status := client.Start(ctx, nil) // nil for statsCh + statusCh <- status + }() + + // Wait for completion or timeout + select { + case status := <-statusCh: + if status != 0 { + b.Errorf("Grep failed with status: %d", status) + } + case <-time.After(30 * time.Second): + b.Error("Grep timed out") + } + + // Restore stdout + w.Close() + os.Stdout = oldStdout + + // Read captured output + var buf bytes.Buffer + buf.ReadFrom(r) + } + + // Report custom metrics + b.ReportMetric(float64(lines), "lines/op") + b.ReportMetric(float64(lines)/b.Elapsed().Seconds(), "lines/sec") +} diff --git a/internal/clients/connectors/connector.go b/internal/clients/connectors/connector.go index 3ab6a08..00de32e 100644 --- a/internal/clients/connectors/connector.go +++ b/internal/clients/connectors/connector.go @@ -2,8 +2,10 @@ package connectors import ( "context" + "time" "github.com/mimecast/dtail/internal/clients/handlers" + sessionspec "github.com/mimecast/dtail/internal/session" ) // Connector interface. @@ -14,4 +16,20 @@ type Connector interface { Server() string // Handler for the connection. Handler() handlers.Handler + // SupportsQueryUpdates reports whether the connected server advertised + // runtime query replacement support within the given timeout. + SupportsQueryUpdates(timeout time.Duration) bool + // ApplySessionSpec starts or updates the interactive session workload on an + // already connected server when query updates are supported. + ApplySessionSpec(spec sessionspec.Spec, timeout time.Duration) error + // ApplySessionSpecWithGeneration starts or updates the interactive session + // workload using the provided committed generation as the base for the + // session update command. + ApplySessionSpecWithGeneration(spec sessionspec.Spec, generation uint64, timeout time.Duration) error + // CommittedSession returns the last session spec and generation that the + // server acknowledged for this connection. + CommittedSession() (sessionspec.Spec, uint64, bool) + // RestoreCommittedSession resets the local committed session snapshot without + // advancing the generation. + RestoreCommittedSession(spec sessionspec.Spec, generation uint64, committed bool) } 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) +} diff --git a/internal/clients/connectors/serverconnection_test.go b/internal/clients/connectors/serverconnection_test.go new file mode 100644 index 0000000..4e22dc7 --- /dev/null +++ b/internal/clients/connectors/serverconnection_test.go @@ -0,0 +1,880 @@ +package connectors + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/mimecast/dtail/internal/clients/handlers" + "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/omode" + "github.com/mimecast/dtail/internal/protocol" + sessionspec "github.com/mimecast/dtail/internal/session" + + "golang.org/x/crypto/ssh" +) + +// TestResolveAuthKeyPathNoLiteralPath is a regression test for the bug +// described in task l6: when authKeyPath is empty and HOME is also unset, +// resolveAuthKeyPath must return "" instead of a mangled path like +// "/.ssh/id_rsa" or the literal "~/.ssh/id_rsa" that the SSH stack cannot use. +func TestResolveAuthKeyPathNoLiteralPath(t *testing.T) { + // Unset HOME so the environment fallback is also empty. + t.Setenv("HOME", "") + + got := resolveAuthKeyPath("") + if got != "" { + t.Fatalf("resolveAuthKeyPath(\"\") with empty HOME = %q; want \"\"", got) + } +} + +// TestResolveAuthKeyPathExplicitPathPassedThrough verifies that a non-empty +// explicit auth key path is returned unchanged. +func TestResolveAuthKeyPathExplicitPathPassedThrough(t *testing.T) { + got := resolveAuthKeyPath("/custom/key") + if got != "/custom/key" { + t.Fatalf("resolveAuthKeyPath(\"/custom/key\") = %q; want \"/custom/key\"", got) + } +} + +// TestResolveAuthKeyPathFallsBackToHome verifies that when authKeyPath is empty +// but HOME is set, the function returns the expected default path. +func TestResolveAuthKeyPathFallsBackToHome(t *testing.T) { + t.Setenv("HOME", "/home/testuser") + + got := resolveAuthKeyPath("") + want := "/home/testuser/.ssh/id_rsa" + if got != want { + t.Fatalf("resolveAuthKeyPath(\"\") = %q; want %q", got, want) + } +} + +func TestExtractAuthKeyBase64(t *testing.T) { + originalLogger := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + dlog.Client = originalLogger + }) + + t.Run("valid authorized key line", func(t *testing.T) { + pubKey := []byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA user@host\n") + + got, err := extractAuthKeyBase64(pubKey) + if err != nil { + t.Fatalf("Expected valid key, got error: %v", err) + } + if got != "AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" { + t.Fatalf("Unexpected base64 payload: %s", got) + } + }) + + t.Run("invalid key format", func(t *testing.T) { + _, err := extractAuthKeyBase64([]byte("not-a-valid-authorized-key-line")) + if err == nil { + t.Fatalf("Expected parse error for invalid key format") + } + }) + + t.Run("invalid base64 payload", func(t *testing.T) { + _, err := extractAuthKeyBase64([]byte("ssh-ed25519 !!! not-valid\n")) + if err == nil { + t.Fatalf("Expected error for invalid base64 payload") + } + }) +} + +func TestSendAuthKeyRegistrationCommand(t *testing.T) { + originalLogger := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + dlog.Client = originalLogger + }) + + tempDir := t.TempDir() + privateKeyPath := filepath.Join(tempDir, "id_rsa") + publicKeyPath := privateKeyPath + ".pub" + if err := os.WriteFile(publicKeyPath, + []byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA user@host\n"), 0600); err != nil { + t.Fatalf("Unable to write public key test file: %v", err) + } + + handler := &mockHandler{} + conn := &ServerConnection{ + server: "srv1", + handler: handler, + authKeyPath: privateKeyPath, + } + + conn.sendAuthKeyRegistrationCommand() + + if len(handler.commands) != 1 { + t.Fatalf("Expected one AUTHKEY command, got %d", len(handler.commands)) + } + expected := "AUTHKEY AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" + if handler.commands[0] != expected { + t.Fatalf("Unexpected AUTHKEY command.\nexpected: %s\ngot: %s", expected, handler.commands[0]) + } +} + +func TestNewServerConnectionUsesInjectedSettings(t *testing.T) { + resetClientLogger(t) + + conn := NewServerConnection( + "srv1", + "user", + nil, + testHostKeyCallback{}, + &mockHandler{}, + nil, + sessionspec.Spec{}, + false, + "", + false, + testSSHSettings{port: 3022, timeout: 5 * time.Second}, + ) + + if conn.hostname != "srv1" { + t.Fatalf("Expected hostname srv1, got %q", conn.hostname) + } + if conn.port != 3022 { + t.Fatalf("Expected injected port 3022, got %d", conn.port) + } + if conn.config.Timeout != 5*time.Second { + t.Fatalf("Expected injected timeout 5s, got %v", conn.config.Timeout) + } +} + +func TestNewServerConnectionFallsBackToDefaults(t *testing.T) { + resetClientLogger(t) + + conn := NewServerConnection( + "srv1", + "user", + nil, + testHostKeyCallback{}, + &mockHandler{}, + nil, + sessionspec.Spec{}, + false, + "", + false, + testSSHSettings{}, + ) + + if conn.port != defaultSSHPort { + t.Fatalf("Expected default port %d, got %d", defaultSSHPort, conn.port) + } + if conn.config.Timeout != defaultSSHConnectTimeout { + t.Fatalf("Expected default timeout %v, got %v", defaultSSHConnectTimeout, conn.config.Timeout) + } +} + +func TestServerConnectionSupportsQueryUpdates(t *testing.T) { + resetClientLogger(t) + + conn := &ServerConnection{ + handler: &mockHandler{ + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + }, + } + + if !conn.SupportsQueryUpdates(10 * time.Millisecond) { + t.Fatalf("expected query-update capability to be detected") + } +} + +func TestServerConnectionSupportsQueryUpdatesFallsBackForOlderServers(t *testing.T) { + resetClientLogger(t) + + conn := &ServerConnection{ + handler: &mockHandler{}, + } + + if conn.SupportsQueryUpdates(5 * time.Millisecond) { + t.Fatalf("expected old-server fallback when no capability is advertised") + } +} + +func TestServerConnectionSupportsQueryUpdatesRequiresCapabilityFlag(t *testing.T) { + resetClientLogger(t) + + conn := &ServerConnection{ + handler: &mockHandler{ + waitForCapabilities: true, + }, + } + + if conn.SupportsQueryUpdates(10 * time.Millisecond) { + t.Fatalf("expected capability wait success alone to be insufficient") + } +} + +func TestServerConnectionApplySessionSpecStart(t *testing.T) { + resetClientLogger(t) + + conn := &ServerConnection{ + server: "srv1", + handler: &mockHandler{ + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + sessionAcks: []handlers.SessionAck{{ + Action: "start", + Generation: 1, + }}, + }, + } + + spec := sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + if err := conn.ApplySessionSpec(spec, 10*time.Millisecond); err != nil { + t.Fatalf("ApplySessionSpec() error = %v", err) + } + + mock := conn.handler.(*mockHandler) + if len(mock.commands) != 1 { + t.Fatalf("expected one session command, got %d", len(mock.commands)) + } + if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 1 || committedSpec.Regex != "ERROR" { + t.Fatalf("unexpected committed session: spec=%#v generation=%d ok=%v", committedSpec, generation, ok) + } +} + +func TestServerConnectionApplySessionSpecUpdateUsesNextGeneration(t *testing.T) { + resetClientLogger(t) + + mock := &mockHandler{ + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + sessionAcks: []handlers.SessionAck{ + {Action: "start", Generation: 4}, + {Action: "update", Generation: 5}, + }, + } + conn := &ServerConnection{ + server: "srv1", + handler: mock, + } + + startSpec := sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + updateSpec := sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "WARN", + } + + if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil { + t.Fatalf("start ApplySessionSpec() error = %v", err) + } + if err := conn.ApplySessionSpec(updateSpec, 10*time.Millisecond); err != nil { + t.Fatalf("update ApplySessionSpec() error = %v", err) + } + if len(mock.commands) != 2 { + t.Fatalf("expected two session commands, got %d", len(mock.commands)) + } + if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 5 || committedSpec.Regex != "WARN" { + t.Fatalf("unexpected committed session after update: spec=%#v generation=%d ok=%v", committedSpec, generation, ok) + } +} + +func TestServerConnectionApplySessionSpecReappliesPreviousSpecForRollback(t *testing.T) { + resetClientLogger(t) + + mock := &mockHandler{ + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + sessionAcks: []handlers.SessionAck{ + {Action: "start", Generation: 4}, + {Action: "update", Generation: 5}, + {Action: "update", Generation: 6}, + }, + } + conn := &ServerConnection{ + server: "srv1", + handler: mock, + } + + startSpec := sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + updateSpec := sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "WARN", + } + + if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil { + t.Fatalf("start ApplySessionSpec() error = %v", err) + } + if err := conn.ApplySessionSpec(updateSpec, 10*time.Millisecond); err != nil { + t.Fatalf("update ApplySessionSpec() error = %v", err) + } + if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil { + t.Fatalf("rollback ApplySessionSpec() error = %v", err) + } + if len(mock.commands) != 3 { + t.Fatalf("expected three session commands, got %d", len(mock.commands)) + } + if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 6 || committedSpec.Regex != "ERROR" { + t.Fatalf("unexpected committed session after rollback: spec=%#v generation=%d ok=%v", committedSpec, generation, ok) + } +} + +func TestServerConnectionApplySessionSpecFallsBackForUnsupportedServer(t *testing.T) { + resetClientLogger(t) + + conn := &ServerConnection{ + handler: &mockHandler{}, + } + + err := conn.ApplySessionSpec(sessionspec.Spec{Mode: omode.TailClient, Regex: "ERROR"}, 5*time.Millisecond) + if !errors.Is(err, ErrSessionUnsupported) { + t.Fatalf("expected ErrSessionUnsupported, got %v", err) + } +} + +func TestRequireJournalCapability(t *testing.T) { + tests := []struct { + name string + spec sessionspec.Spec + waitForCapabilities bool + capabilities map[string]bool + wantErr error + wantServerError bool + }{ + { + name: "journal file with journal capability", + spec: sessionspec.Spec{ + Mode: omode.CatClient, + Files: []string{"journal:ssh.service"}, + }, + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityJournalV1: true, + }, + }, + { + name: "journal file without journal capability", + spec: sessionspec.Spec{ + Mode: omode.CatClient, + Files: []string{"journal:ssh.service"}, + }, + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + wantErr: ErrJournalUnsupported, + wantServerError: true, + }, + { + name: "journal file without capabilities advertisement", + spec: sessionspec.Spec{ + Mode: omode.CatClient, + Files: []string{"journal:ssh.service"}, + }, + wantErr: ErrJournalUnsupported, + wantServerError: true, + }, + { + name: "regular file without journal capability", + spec: sessionspec.Spec{ + Mode: omode.CatClient, + Files: []string{"/var/log/app.log"}, + }, + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + handler := &mockHandler{ + waitForCapabilities: tc.waitForCapabilities, + capabilities: tc.capabilities, + } + + err := requireJournalCapability("srv1", handler, tc.spec, 10*time.Millisecond) + if !errors.Is(err, tc.wantErr) { + t.Fatalf("requireJournalCapability() error = %v, want %v", err, tc.wantErr) + } + if got := handler.serverError != ""; got != tc.wantServerError { + t.Fatalf("server error recorded = %v, want %v", got, tc.wantServerError) + } + if tc.wantServerError && !strings.Contains(handler.serverError, protocol.CapabilityJournalV1) { + t.Fatalf("server error %q does not mention %s", handler.serverError, protocol.CapabilityJournalV1) + } + }) + } +} + +func TestDispatchInitialCommandsRejectsJournalWithoutCapability(t *testing.T) { + resetClientLogger(t) + + handler := &mockHandler{ + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + } + spec := sessionspec.Spec{ + Mode: omode.CatClient, + Files: []string{"journal:ssh.service"}, + } + + err := dispatchInitialCommands("srv1", handler, []string{"cat: journal:ssh.service ."}, false, spec, &committedSessionState{}) + if !errors.Is(err, ErrJournalUnsupported) { + t.Fatalf("expected ErrJournalUnsupported, got %v", err) + } + if len(handler.commands) != 0 { + t.Fatalf("expected no commands to be sent, got %#v", handler.commands) + } + if handler.Status() != 1 { + t.Fatalf("handler status = %d, want 1", handler.Status()) + } +} + +func TestDispatchInitialCommandsRejectsInteractiveJournalWithoutCapability(t *testing.T) { + resetClientLogger(t) + + handler := &mockHandler{ + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + } + spec := sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"journal:ssh.service"}, + } + + err := dispatchInitialCommands("srv1", handler, []string{"tail: journal:ssh.service ."}, true, spec, &committedSessionState{}) + if !errors.Is(err, ErrJournalUnsupported) { + t.Fatalf("expected ErrJournalUnsupported, got %v", err) + } + if len(handler.commands) != 0 { + t.Fatalf("expected no commands to be sent, got %#v", handler.commands) + } + if handler.Status() != 1 { + t.Fatalf("handler status = %d, want 1", handler.Status()) + } +} + +func TestServerConnectionApplySessionSpecPreservesCommittedStateOnRejectedUpdate(t *testing.T) { + resetClientLogger(t) + + mock := &mockHandler{ + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + sessionAcks: []handlers.SessionAck{ + {Action: "start", Generation: 2}, + {Action: "error", Error: "bad reload"}, + }, + } + conn := &ServerConnection{ + server: "srv1", + handler: mock, + } + + startSpec := sessionspec.Spec{Mode: omode.TailClient, Regex: "ERROR"} + if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil { + t.Fatalf("start ApplySessionSpec() error = %v", err) + } + + err := conn.ApplySessionSpec(sessionspec.Spec{Mode: omode.TailClient, Regex: "WARN"}, 10*time.Millisecond) + if !errors.Is(err, ErrSessionRejected) { + t.Fatalf("expected ErrSessionRejected, got %v", err) + } + if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 2 || committedSpec.Regex != "ERROR" { + t.Fatalf("unexpected committed session after rejected update: spec=%#v generation=%d ok=%v", committedSpec, generation, ok) + } +} + +func TestServerConnectionApplySessionSpecRejectsUnexpectedAck(t *testing.T) { + resetClientLogger(t) + + mock := &mockHandler{ + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + sessionAcks: []handlers.SessionAck{ + {Action: "update", Generation: 1}, + }, + } + conn := &ServerConnection{ + server: "srv1", + handler: mock, + } + + err := conn.ApplySessionSpec(sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + }, 10*time.Millisecond) + if !errors.Is(err, ErrUnexpectedSessionAck) { + t.Fatalf("expected ErrUnexpectedSessionAck, got %v", err) + } + if _, _, ok := conn.CommittedSession(); ok { + t.Fatalf("unexpected committed session after mismatched ack") + } +} + +func TestServerConnectionApplySessionSpecTimesOutWaitingForAck(t *testing.T) { + resetClientLogger(t) + + mock := &mockHandler{ + waitForCapabilities: true, + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + } + conn := &ServerConnection{ + server: "srv1", + handler: mock, + } + + err := conn.ApplySessionSpec(sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + }, 10*time.Millisecond) + if !errors.Is(err, ErrSessionAckTimeout) { + t.Fatalf("expected ErrSessionAckTimeout, got %v", err) + } + if len(mock.commands) != 1 { + t.Fatalf("expected session command to be sent before timeout, got %d", len(mock.commands)) + } + if _, _, ok := conn.CommittedSession(); ok { + t.Fatalf("unexpected committed session after missing ack") + } +} + +func TestApplySessionSpecSerializesConcurrentBootstrapAndReload(t *testing.T) { + resetClientLogger(t) + + handler := newBlockingSessionHandler() + state := &committedSessionState{} + + initialSpec := sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + reloadSpec := sessionspec.Spec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "WARN", + } + + initialErrCh := make(chan error, 1) + go func() { + initialErrCh <- dispatchInitialCommands("srv1", handler, nil, true, initialSpec, state) + }() + + firstCommand := <-handler.commandsCh + if !strings.HasPrefix(firstCommand, "SESSION START ") { + t.Fatalf("expected initial SESSION START command, got %q", firstCommand) + } + + reloadErrCh := make(chan error, 1) + go func() { + reloadErrCh <- applySessionSpec("srv1", handler, state, reloadSpec, 50*time.Millisecond) + }() + + select { + case command := <-handler.commandsCh: + t.Fatalf("unexpected concurrent session command before bootstrap ack: %q", command) + case <-time.After(10 * time.Millisecond): + } + + handler.ackCh <- handlers.SessionAck{Action: "start", Generation: 1} + if err := <-initialErrCh; err != nil { + t.Fatalf("dispatchInitialCommands() error = %v", err) + } + + secondCommand := <-handler.commandsCh + if !strings.HasPrefix(secondCommand, "SESSION UPDATE 2 ") { + t.Fatalf("expected reload to send SESSION UPDATE after bootstrap, got %q", secondCommand) + } + + handler.ackCh <- handlers.SessionAck{Action: "update", Generation: 2} + if err := <-reloadErrCh; err != nil { + t.Fatalf("applySessionSpec() error = %v", err) + } + + committedSpec, generation, ok := state.snapshot() + if !ok || generation != 2 || committedSpec.Regex != "WARN" { + t.Fatalf("unexpected committed session after reload: spec=%#v generation=%d ok=%v", committedSpec, generation, ok) + } +} + +// TestThrottleReleasedIsIdempotent verifies that calling the throttle-release +// logic from two concurrent goroutines drains throttleCh exactly once and does +// not panic or block. This is a regression test for the data race that existed +// when the old bool guard (throttlingDone) was read and written without +// synchronization: under -race two goroutines could both observe the bool as +// false and both attempt to drain the channel, stealing an extra slot. +func TestThrottleReleasedIsIdempotent(t *testing.T) { + t.Parallel() + + // throttleCh is buffered with 1 slot, as in the real Start() path. + throttleCh := make(chan struct{}, 1) + throttleCh <- struct{}{} // occupy the one slot + + conn := &ServerConnection{} + + const workers = 64 + var wg sync.WaitGroup + wg.Add(workers) + + // Simulate workers racing to release the throttle slot (e.g. handle() + // early-release and the defer cleanup in Start() firing around the same + // time). Only one drain must succeed; the rest must be no-ops. + for range workers { + go func() { + defer wg.Done() + conn.throttleReleased.Do(func() { + <-throttleCh + }) + }() + } + + wg.Wait() + + // throttleCh must be empty: exactly one goroutine drained it. + if len(throttleCh) != 0 { + t.Fatalf("throttleCh length = %d, want 0 (slot was not released)", len(throttleCh)) + } + + // Confirm a second occupant can now be added, proving the slot is free. + select { + case throttleCh <- struct{}{}: + // expected: slot was freed exactly once + default: + t.Fatal("throttleCh full after release, expected one free slot") + } +} + +type testSSHSettings struct { + port int + timeout time.Duration +} + +func (s testSSHSettings) SSHPort() int { + return s.port +} + +func (s testSSHSettings) SSHConnectTimeout() time.Duration { + return s.timeout +} + +type testHostKeyCallback struct{} + +func (testHostKeyCallback) Wrap(context.Context) ssh.HostKeyCallback { + return ssh.InsecureIgnoreHostKey() +} + +func (testHostKeyCallback) Untrusted(string) bool { + return false +} + +func (testHostKeyCallback) PromptAddHosts(context.Context) {} + +func resetClientLogger(t *testing.T) { + t.Helper() + + originalLogger := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + dlog.Client = originalLogger + }) +} + +type mockHandler struct { + commands []string + capabilities map[string]bool + waitForCapabilities bool + sessionAcks []handlers.SessionAck + serverError string + status int +} + +var _ handlers.Handler = (*mockHandler)(nil) + +func (m *mockHandler) SendMessage(command string) error { + m.commands = append(m.commands, command) + return nil +} + +func (m *mockHandler) Capabilities() []string { + var capabilities []string + for capability := range m.capabilities { + capabilities = append(capabilities, capability) + } + return capabilities +} + +func (m *mockHandler) HasCapability(name string) bool { + return m.capabilities[name] +} + +func (m *mockHandler) ReportServerError(message string) { + m.serverError = message + m.status = 1 +} + +func (m *mockHandler) Server() string { + return "mock" +} + +func (m *mockHandler) Status() int { + return m.status +} + +func (m *mockHandler) Shutdown() {} + +func (m *mockHandler) Done() <-chan struct{} { + ch := make(chan struct{}) + close(ch) + return ch +} + +func (m *mockHandler) WaitForCapabilities(timeout time.Duration) bool { + return m.waitForCapabilities +} + +func (m *mockHandler) WaitForSessionAck(timeout time.Duration) (handlers.SessionAck, bool) { + if timeout <= 0 { + return handlers.SessionAck{}, false + } + if len(m.sessionAcks) == 0 { + return handlers.SessionAck{}, false + } + + ack := m.sessionAcks[0] + m.sessionAcks = m.sessionAcks[1:] + return ack, true +} + +func (m *mockHandler) Read(_ []byte) (int, error) { + return 0, nil +} + +func (m *mockHandler) Write(p []byte) (int, error) { + return len(p), nil +} + +type blockingSessionHandler struct { + mu sync.Mutex + comman