diff options
Diffstat (limited to 'internal/clients/interactive_control_test.go')
| -rw-r--r-- | internal/clients/interactive_control_test.go | 641 |
1 files changed, 641 insertions, 0 deletions
diff --git a/internal/clients/interactive_control_test.go b/internal/clients/interactive_control_test.go new file mode 100644 index 0000000..cb779df --- /dev/null +++ b/internal/clients/interactive_control_test.go @@ -0,0 +1,641 @@ +package clients + +import ( + "context" + "errors" + "reflect" + "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" + "github.com/mimecast/dtail/internal/protocol" + sessionspec "github.com/mimecast/dtail/internal/session" + + "golang.org/x/crypto/ssh" +) + +func TestParseInteractiveCommandForGrepReload(t *testing.T) { + current := config.Args{ + Mode: omode.GrepClient, + What: "/var/log/app.log", + RegexStr: "ERROR", + } + + command, err := parseInteractiveCommand(current, + `:reload --grep WARN --before 2 --after 3 --max 4 --invert --files "/tmp/other.log" --plain --quiet --timeout 7`) + if err != nil { + t.Fatalf("parseInteractiveCommand() error = %v", err) + } + + if command.kind != "reload" { + t.Fatalf("command kind = %q, want reload", command.kind) + } + if command.next.What != "/tmp/other.log" { + t.Fatalf("files = %q, want /tmp/other.log", command.next.What) + } + if command.next.RegexStr != "WARN" { + t.Fatalf("regex = %q, want WARN", command.next.RegexStr) + } + if !command.next.RegexInvert || !command.next.Plain || !command.next.Quiet { + t.Fatalf("expected invert/plain/quiet flags to be set: %#v", command.next) + } + if command.next.LContext.BeforeContext != 2 || command.next.LContext.AfterContext != 3 || command.next.LContext.MaxCount != 4 { + t.Fatalf("unexpected context values: %#v", command.next.LContext) + } + if command.spec.Regex != "WARN" { + t.Fatalf("session spec regex = %q, want WARN", command.spec.Regex) + } +} + +func TestParseInteractiveCommandForMapReloadDerivesRegex(t *testing.T) { + current := config.Args{ + Mode: omode.MapClient, + What: "/var/log/app.log", + QueryStr: "select count(status) from stats group by status", + } + + command, err := parseInteractiveCommand(current, + `:reload --query "select count(status) from warnings group by status" --files /tmp/new.log --plain --timeout 9`) + if err != nil { + t.Fatalf("parseInteractiveCommand() error = %v", err) + } + + if command.kind != "reload" { + t.Fatalf("command kind = %q, want reload", command.kind) + } + if command.next.QueryStr != "select count(status) from warnings group by status" { + t.Fatalf("query = %q", command.next.QueryStr) + } + if command.spec.Regex != "\\|MAPREDUCE:WARNINGS\\|" { + t.Fatalf("session spec regex = %q, want WARNINGS table regex", command.spec.Regex) + } +} + +func TestParseInteractiveCommandRejectsUnterminatedQuotes(t *testing.T) { + current := config.Args{ + Mode: omode.MapClient, + QueryStr: "select count(status) from stats group by status", + } + + if _, err := parseInteractiveCommand(current, `:reload --query "select count(status) from stats`); err == nil { + t.Fatalf("expected parseInteractiveCommand() to reject unterminated quoted input") + } +} + +func TestApplyInteractiveReloadRejectsUnsupportedConnections(t *testing.T) { + client := &baseClient{ + mu: newBaseClientMu(), + Args: config.Args{ + Mode: omode.GrepClient, + What: "/var/log/app.log", + RegexStr: "ERROR", + }, + sessionSpec: SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + }, + connections: []connectors.Connector{ + &interactiveReloadConnector{server: "srv1", supported: false}, + }, + } + + err := client.applyInteractiveReload(config.Args{ + Mode: omode.GrepClient, + What: "/tmp/next.log", + RegexStr: "WARN", + }, SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/tmp/next.log"}, + Regex: "WARN", + }) + if !errors.Is(err, connectors.ErrSessionUnsupported) { + t.Fatalf("expected ErrSessionUnsupported, got %v", err) + } + if client.Args.What != "/var/log/app.log" || client.sessionSpec.Regex != "ERROR" { + t.Fatalf("client state changed on unsupported reload: args=%#v spec=%#v", client.Args, client.sessionSpec) + } +} + +func TestApplyInteractiveReloadRollsBackEarlyFailure(t *testing.T) { + oldArgs := config.Args{ + Mode: omode.GrepClient, + What: "/var/log/app.log", + RegexStr: "ERROR", + } + oldSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + rollbackErr := errors.New("boom") + + connA := &interactiveReloadConnector{server: "srv1", supported: true, committedSpec: oldSpec, liveSpec: oldSpec, generation: 4} + connB := &interactiveReloadConnector{server: "srv2", supported: true, committedSpec: oldSpec, liveSpec: oldSpec, generation: 4} + connC := &interactiveReloadConnector{server: "srv3", supported: true, committedSpec: oldSpec, liveSpec: oldSpec, applyErr: rollbackErr, generation: 4} + maker := &interactiveReloadMaker{} + + client := &baseClient{ + mu: newBaseClientMu(), + Args: oldArgs, + sessionSpec: oldSpec, + connections: []connectors.Connector{connA, connB, connC}, + maker: maker, + } + + nextArgs := config.Args{ + Mode: omode.GrepClient, + What: "/tmp/new.log", + RegexStr: "WARN", + } + nextSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/tmp/new.log"}, + Regex: "WARN", + } + + err := client.applyInteractiveReload(nextArgs, nextSpec) + if err == nil || !errors.Is(err, rollbackErr) { + t.Fatalf("expected reload error from failing connection, got %v", err) + } + + if !reflect.DeepEqual(client.Args, oldArgs) || !reflect.DeepEqual(client.sessionSpec, oldSpec) { + t.Fatalf("client state changed on partial failure: args=%#v spec=%#v", client.Args, client.sessionSpec) + } + if len(maker.commits) != 0 { + t.Fatalf("expected no committed shared state, got %#v", maker.commits) + } + if !reflect.DeepEqual(connA.committedSpec, oldSpec) || !reflect.DeepEqual(connB.committedSpec, oldSpec) { + t.Fatalf("expected successful connections to roll back to old spec: %#v %#v", connA.committedSpec, connB.committedSpec) + } + if connA.generation != 6 || connB.generation != 6 { + t.Fatalf("expected rollback to advance generations through the remote update path, got %d and %d", connA.generation, connB.generation) + } + if connA.applyCount != 2 || connB.applyCount != 2 || connC.applyCount != 1 { + t.Fatalf("unexpected apply counts during rollback: %#v %#v %#v", connA.applyCount, connB.applyCount, connC.applyCount) + } + if !reflect.DeepEqual(connA.liveSpec, oldSpec) || !reflect.DeepEqual(connB.liveSpec, oldSpec) { + t.Fatalf("expected rollback to restore live specs on successful connections: %#v %#v", connA.liveSpec, connB.liveSpec) + } + if !reflect.DeepEqual(connC.committedSpec, oldSpec) || connC.generation != 4 { + t.Fatalf("failed connection should remain on original committed session: %#v generation=%d", connC.committedSpec, connC.generation) + } +} + +func TestApplyInteractiveReloadRollsBackLateAckFailure(t *testing.T) { + oldArgs := config.Args{ + Mode: omode.GrepClient, + What: "/var/log/app.log", + RegexStr: "ERROR", + } + oldSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + nextSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/tmp/new.log"}, + Regex: "WARN", + } + + connA := &interactiveReloadConnector{server: "srv1", supported: true, committedSpec: oldSpec, liveSpec: oldSpec, generation: 4} + connB := &interactiveReloadConnector{server: "srv2", supported: true, committedSpec: oldSpec, liveSpec: oldSpec, generation: 4} + connC := &interactiveReloadConnector{ + server: "srv3", + supported: true, + committedSpec: oldSpec, + liveSpec: oldSpec, + generation: 4, + applyErrAfterAdvance: connectors.ErrSessionAckTimeout, + applyErrAfterAdvanceSpec: nextSpec, + } + + client := &baseClient{ + mu: newBaseClientMu(), + Args: oldArgs, + sessionSpec: oldSpec, + connections: []connectors.Connector{connA, connB, connC}, + } + + nextArgs := config.Args{ + Mode: omode.GrepClient, + What: "/tmp/new.log", + RegexStr: "WARN", + } + err := client.applyInteractiveReload(nextArgs, nextSpec) + if err == nil || !errors.Is(err, connectors.ErrSessionAckTimeout) { + t.Fatalf("expected late ack timeout error, got %v", err) + } + + if !reflect.DeepEqual(client.Args, oldArgs) || !reflect.DeepEqual(client.sessionSpec, oldSpec) { + t.Fatalf("client state changed on late failure: args=%#v spec=%#v", client.Args, client.sessionSpec) + } + if !reflect.DeepEqual(connA.committedSpec, oldSpec) || !reflect.DeepEqual(connB.committedSpec, oldSpec) { + t.Fatalf("expected successful connections to roll back to old spec: %#v %#v", connA.committedSpec, connB.committedSpec) + } + if connA.generation != 6 || connB.generation != 6 { + t.Fatalf("expected rollback to advance generations through the remote update path, got %d and %d", connA.generation, connB.generation) + } + if connA.applyCount != 2 || connB.applyCount != 2 || connC.applyCount != 2 { + t.Fatalf("unexpected apply counts during late rollback: %#v %#v %#v", connA.applyCount, connB.applyCount, connC.applyCount) + } + if !reflect.DeepEqual(connA.liveSpec, oldSpec) || !reflect.DeepEqual(connB.liveSpec, oldSpec) { + t.Fatalf("expected rollback to restore live specs on successful connections: %#v %#v", connA.liveSpec, connB.liveSpec) + } + if !reflect.DeepEqual(connC.committedSpec, oldSpec) || !reflect.DeepEqual(connC.liveSpec, oldSpec) || connC.generation != 6 { + t.Fatalf("late-failing connection should also roll back to original committed session: committed=%#v live=%#v generation=%d", connC.committedSpec, connC.liveSpec, connC.generation) + } +} + +func TestApplyInteractiveReloadRollsBackLateAckFailureWithRealServerConnection(t *testing.T) { + resetClientLogger(t) + + oldArgs := config.Args{ + Mode: omode.GrepClient, + What: "/var/log/app.log", + RegexStr: "ERROR", + } + oldSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + nextSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/tmp/new.log"}, + Regex: "WARN", + } + + handlerA := newInteractiveReloadSessionHandler( + handlers.SessionAck{Action: "update", Generation: 5}, + handlers.SessionAck{Action: "update", Generation: 6}, + ) + handlerB := newInteractiveReloadSessionHandler( + handlers.SessionAck{Action: "update", Generation: 5}, + handlers.SessionAck{Action: "update", Generation: 6}, + ) + handlerC := newInteractiveReloadSessionHandler() + connA := newInteractiveReloadServerConnection(t, "srv1", handlerA, oldSpec) + connB := newInteractiveReloadServerConnection(t, "srv2", handlerB, oldSpec) + connC := newInteractiveReloadServerConnection(t, "srv3", handlerC, oldSpec) + + for _, conn := range []*connectors.ServerConnection{connA, connB, connC} { + conn.RestoreCommittedSession(oldSpec, 4, true) + } + + client := &baseClient{ + mu: newBaseClientMu(), + Args: oldArgs, + sessionSpec: oldSpec, + connections: []connectors.Connector{connA, connB, connC}, + } + + nextArgs := config.Args{ + Mode: omode.GrepClient, + What: "/tmp/new.log", + RegexStr: "WARN", + } + err := client.applyInteractiveReload(nextArgs, nextSpec) + if err == nil || !errors.Is(err, connectors.ErrSessionAckTimeout) { + t.Fatalf("expected late ack timeout error, got %v", err) + } + + if !reflect.DeepEqual(client.Args, oldArgs) || !reflect.DeepEqual(client.sessionSpec, oldSpec) { + t.Fatalf("client state changed on late failure: args=%#v spec=%#v", client.Args, client.sessionSpec) + } + + for _, conn := range []*connectors.ServerConnection{connA, connB} { + committedSpec, generation, ok := conn.CommittedSession() + if !ok || generation != 6 || !reflect.DeepEqual(committedSpec, oldSpec) { + t.Fatalf("successful connection did not roll back cleanly: spec=%#v generation=%d ok=%v", committedSpec, generation, ok) + } + } + + if committedSpec, generation, ok := connC.CommittedSession(); !ok || generation != 4 || !reflect.DeepEqual(committedSpec, oldSpec) { + t.Fatalf("timed-out rollback should leave failed connection at its original committed session: spec=%#v generation=%d ok=%v", committedSpec, generation, ok) + } + if len(handlerC.commands) != 2 { + t.Fatalf("expected the failed connection to receive both the reload and rollback commands, got %#v", handlerC.commands) + } +} + +func TestApplyInteractiveReloadCommitsSharedState(t *testing.T) { + oldSpec := SessionSpec{ + Mode: omode.MapClient, + Files: []string{"/var/log/app.log"}, + Query: "select count(status) from stats group by status", + Regex: "\\|MAPREDUCE:STATS\\|", + } + connA := &interactiveReloadConnector{server: "srv1", supported: true, committedSpec: oldSpec, liveSpec: oldSpec, generation: 4} + connB := &interactiveReloadConnector{server: "srv2", supported: true, committedSpec: oldSpec, liveSpec: oldSpec, generation: 4} + maker := &interactiveReloadMaker{} + + client := &baseClient{ + mu: newBaseClientMu(), + Args: config.Args{ + Mode: omode.MapClient, + What: "/var/log/app.log", + QueryStr: "select count(status) from stats group by status", + }, + sessionSpec: SessionSpec{ + Mode: omode.MapClient, + Files: []string{"/var/log/app.log"}, + Query: "select count(status) from stats group by status", + Regex: "\\|MAPREDUCE:STATS\\|", + }, + connections: []connectors.Connector{connA, connB}, + maker: maker, + } + + nextArgs := config.Args{ + Mode: omode.MapClient, + What: "/tmp/new.log", + QueryStr: "select count(status) from warnings group by status", + Plain: true, + Timeout: 5, + } + nextSpec := SessionSpec{ + Mode: omode.MapClient, + Files: []string{"/tmp/new.log"}, + Query: nextArgs.QueryStr, + Regex: "\\|MAPREDUCE:WARNINGS\\|", + Timeout: 5, + } + + if err := client.applyInteractiveReload(nextArgs, nextSpec); err != nil { + t.Fatalf("applyInteractiveReload() error = %v", err) + } + + if client.Args.What != "/tmp/new.log" || client.sessionSpec.Query != nextArgs.QueryStr { + t.Fatalf("client state not committed: args=%#v spec=%#v", client.Args, client.sessionSpec) + } + if len(maker.commits) != 1 { + t.Fatalf("expected one sessionCommitter call, got %d", len(maker.commits)) + } + if maker.commits[0].generation != 5 || maker.commits[0].spec.Query != nextArgs.QueryStr { + t.Fatalf("unexpected commit payload: %#v", maker.commits[0]) + } + if connA.committedSpec.Query != nextArgs.QueryStr || connB.committedSpec.Query != nextArgs.QueryStr { + t.Fatalf("connectors did not receive new session spec: %#v %#v", connA.committedSpec, connB.committedSpec) + } +} + +func TestApplyInteractiveReloadRejectsMismatchedCommittedGenerations(t *testing.T) { + oldSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + connA := &interactiveReloadConnector{server: "srv1", supported: true, committedSpec: oldSpec, liveSpec: oldSpec, generation: 4} + connB := &interactiveReloadConnector{server: "srv2", supported: true, committedSpec: oldSpec, liveSpec: oldSpec, generation: 5} + maker := &interactiveReloadMaker{} + + client := &baseClient{ + mu: newBaseClientMu(), + Args: config.Args{ + Mode: omode.GrepClient, + What: "/var/log/app.log", + RegexStr: "ERROR", + }, + sessionSpec: SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + }, + connections: []connectors.Connector{connA, connB}, + maker: maker, + } + + nextArgs := config.Args{ + Mode: omode.GrepClient, + What: "/tmp/new.log", + RegexStr: "WARN", + } + nextSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/tmp/new.log"}, + Regex: "WARN", + } + + err := client.applyInteractiveReload(nextArgs, nextSpec) + if err == nil || err.Error() != "mismatched committed generations: got 5 and 6" { + t.Fatalf("expected mismatched generation error, got %v", err) + } + if client.Args.What != "/var/log/app.log" || client.sessionSpec.Regex != "ERROR" { + t.Fatalf("client state changed on mismatched generations: args=%#v spec=%#v", client.Args, client.sessionSpec) + } + if len(maker.commits) != 0 { + t.Fatalf("expected no committed shared state, got %#v", maker.commits) + } + if connA.generation != 6 || connB.generation != 7 { + t.Fatalf("expected rollback to replay the previous spec remotely, got %d and %d", connA.generation, connB.generation) + } + if connA.applyCount != 2 || connB.applyCount != 2 { + t.Fatalf("expected rollback to reapply the previous spec on both connections, got %d and %d", connA.applyCount, connB.applyCount) + } + if !reflect.DeepEqual(connA.committedSpec, oldSpec) || !reflect.DeepEqual(connB.committedSpec, oldSpec) { + t.Fatalf("expected rollback to restore committed specs, got %#v %#v", connA.committedSpec, connB.committedSpec) + } +} + +type interactiveReloadConnector struct { + committedSpec sessionspec.Spec + liveSpec sessionspec.Spec + applyErr error + applyErrAfterAdvance error + applyErrAfterAdvanceSpec sessionspec.Spec + applyCount int + generation uint64 + server string + supported bool +} + +type interactiveReloadSessionHandler struct { + commands []string + capabilities map[string]bool + sessionAcks []handlers.SessionAck + waitForCapabilities bool +} + +var _ handlers.Handler = (*interactiveReloadSessionHandler)(nil) + +func newInteractiveReloadSessionHandler(acks ...handlers.SessionAck) *interactiveReloadSessionHandler { + return &interactiveReloadSessionHandler{ + capabilities: map[string]bool{ + protocol.CapabilityQueryUpdateV1: true, + }, + sessionAcks: append([]handlers.SessionAck(nil), acks...), + waitForCapabilities: true, + } +} + +func (h *interactiveReloadSessionHandler) SendMessage(command string) error { + h.commands = append(h.commands, command) + return nil +} + +func (h *interactiveReloadSessionHandler) Capabilities() []string { + capabilities := make([]string, 0, len(h.capabilities)) + for capability := range h.capabilities { + capabilities = append(capabilities, capability) + } + return capabilities +} + +func (h *interactiveReloadSessionHandler) HasCapability(name string) bool { + return h.capabilities[name] +} + +func (*interactiveReloadSessionHandler) ReportServerError(string) {} + +func (*interactiveReloadSessionHandler) Server() string { return "mock" } + +func (*interactiveReloadSessionHandler) Status() int { return 0 } + +func (*interactiveReloadSessionHandler) Shutdown() {} + +func (*interactiveReloadSessionHandler) Done() <-chan struct{} { + ch := make(chan struct{}) + close(ch) + return ch +} + +func (h *interactiveReloadSessionHandler) WaitForCapabilities(time.Duration) bool { + return h.waitForCapabilities +} + +func (h *interactiveReloadSessionHandler) WaitForSessionAck(timeout time.Duration) (handlers.SessionAck, bool) { + if timeout <= 0 { + return handlers.SessionAck{}, false + } + if len(h.sessionAcks) == 0 { + return handlers.SessionAck{}, false + } + + ack := h.sessionAcks[0] + h.sessionAcks = h.sessionAcks[1:] + return ack, true +} + +func (*interactiveReloadSessionHandler) Read(_ []byte) (int, error) { + return 0, nil +} + +func (*interactiveReloadSessionHandler) Write(p []byte) (int, error) { + return len(p), nil +} + +type interactiveReloadHostKeyCallback struct{} + +func (interactiveReloadHostKeyCallback) Wrap(context.Context) ssh.HostKeyCallback { + return ssh.InsecureIgnoreHostKey() +} + +func (interactiveReloadHostKeyCallback) Untrusted(string) bool { return false } + +func (interactiveReloadHostKeyCallback) PromptAddHosts(context.Context) {} + +func newInteractiveReloadServerConnection(t *testing.T, server string, handler handlers.Handler, spec SessionSpec) *connectors.ServerConnection { + t.Helper() + + conn := connectors.NewServerConnection( + server, + "user", + nil, + interactiveReloadHostKeyCallback{}, + handler, + nil, + spec, + true, + "", + false, + nil, + ) + conn.RestoreCommittedSession(spec, 4, true) + return conn +} + +func resetClientLogger(t *testing.T) { + t.Helper() + + originalLogger := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + dlog.Client = originalLogger + }) +} + +func (*interactiveReloadConnector) Start(context.Context, context.CancelFunc, chan struct{}, chan struct{}) { +} + +func (c *interactiveReloadConnector) Server() string { return c.server } + +func (*interactiveReloadConnector) Handler() handlers.Handler { return nil } + +func (c *interactiveReloadConnector) SupportsQueryUpdates(time.Duration) bool { return c.supported } + +func (c *interactiveReloadConnector) ApplySessionSpec(spec sessionspec.Spec, _ time.Duration) error { + _, generation, ok := c.CommittedSession() + if !ok { + generation = 0 + } + return c.ApplySessionSpecWithGeneration(spec, generation, 0) +} + +func (c *interactiveReloadConnector) ApplySessionSpecWithGeneration(spec sessionspec.Spec, generation uint64, _ time.Duration) error { + c.applyCount++ + if c.applyErrAfterAdvance != nil && reflect.DeepEqual(spec, c.applyErrAfterAdvanceSpec) { + c.liveSpec = spec + return c.applyErrAfterAdvance + } + if c.applyErr != nil { + return c.applyErr + } + if generation == 0 { + c.generation = 1 + } else { + c.generation = generation + 1 + } + c.liveSpec = spec + c.committedSpec = spec + return nil +} + +func (c *interactiveReloadConnector) CommittedSession() (sessionspec.Spec, uint64, bool) { + if c.generation == 0 { + return sessionspec.Spec{}, 0, false + } + return c.committedSpec, c.generation, true +} + +func (c *interactiveReloadConnector) RestoreCommittedSession(spec sessionspec.Spec, generation uint64, committed bool) { + c.liveSpec = spec + c.committedSpec = spec + c.generation = generation +} + +type interactiveReloadMaker struct { + commits []interactiveReloadCommit +} + +type interactiveReloadCommit struct { + generation uint64 + spec SessionSpec +} + +func (*interactiveReloadMaker) makeHandler(string) handlers.Handler { return nil } + +func (*interactiveReloadMaker) makeCommands() []string { return nil } + +func (m *interactiveReloadMaker) commitSessionSpec(spec SessionSpec, generation uint64) error { + m.commits = append(m.commits, interactiveReloadCommit{ + generation: generation, + spec: spec, + }) + return nil +} |