summaryrefslogtreecommitdiff
path: root/internal/clients/connectors
diff options
context:
space:
mode:
Diffstat (limited to 'internal/clients/connectors')
-rw-r--r--internal/clients/connectors/connector.go18
-rw-r--r--internal/clients/connectors/serverconnection.go248
-rw-r--r--internal/clients/connectors/serverconnection_test.go880
-rw-r--r--internal/clients/connectors/serverless.go213
-rw-r--r--internal/clients/connectors/sessiontransport.go213
5 files changed, 1502 insertions, 70 deletions
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 '<type> <base64-key> [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