diff options
Diffstat (limited to 'internal/clients')
30 files changed, 5111 insertions, 250 deletions
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() + } |