diff --git a/docs/gop/gop0.md b/docs/gop/gop0.md index 695845aa..5d0be052 100644 --- a/docs/gop/gop0.md +++ b/docs/gop/gop0.md @@ -59,3 +59,4 @@ The numerical index contains a table of all GOPs, ordered by number. | 1 | Core CLI Primitives | Active | Informational | Core | | 2 | Login Expression Detection in Execution Loop for Console Streamers | Accepted | Standards Track | Drivers/Transports | | 3 | Repeated Question Detection in Execution Loop | Accepted | Standards Track | Core | +| 4 | Recovering Command Echo After a Transient Prompt | Accepted | Standards Track | Core | diff --git a/docs/gop/gop4.md b/docs/gop/gop4.md new file mode 100644 index 00000000..a8499087 --- /dev/null +++ b/docs/gop/gop4.md @@ -0,0 +1,95 @@ +# GOP 4 – Recovering Command Echo After a Transient Prompt + +- **Author:** Reydant +- **Status:** Accepted +- **Type:** Standards Track +- **Topic:** Core +- **Created:** 2026-08-12 +- **Last-Modified:** 2026-08-12 + +## Abstract + +Some interactive CLIs redraw their prompt while they receive a command. If a transport read ends at that redraw, `GenericExecute` can match the prompt before it has observed the complete command echo and return an `EchoReadException`. + +This GOP adds recovery for that condition. When a prompt is matched before the command echo, gnetcli preserves the consumed bytes by prepending them to the connector's unread buffer, then continues reading. This lets the normal echo matcher evaluate the preserved bytes together with later terminal output. + +## Motivation + +RouterOS redraws the command line while each character is being entered. A stream can contain data such as: + +``` +\r[robot@router] > /user exp +``` + +At a network-read boundary, the `\r[robot@router] > ` part is indistinguishable from a final prompt. The command echo is incomplete, however, and more redraw data or the submitted command echo may arrive later. + +Previously, matching this prompt before the expected echo caused command execution to fail immediately with `echo read error`. + +## Failure dialog from RouterOS logs + +The following is a shortened and sanitized dialog from the original failure. Runs of terminal-padding spaces are replaced with `…`; `\x1b[K` clears the rest of the line. + +```text +client -> /user export\n + +device -> \r[robot@router] > /user exp\r\r[robot@router] > + ^ prompt matcher stops here + +device -> \x1b[K\r[robot@router] > /user ex\r\r[robot@router] > /user ex… +device -> \x1b[K\r[robot@router] > /user exp\r\r[robot@router] > /user exp… +``` + +The first read ends immediately after a redraw prompt. At that point `/user exp` is only a partial line, so it cannot match the expected `/user export\r\n` echo. Before GOP 4, `GenericExecute` returned `EchoReadException` at the marked prompt and never read the later redraw data. + +## Proposal + +`streamer.Connector` provides: + +```go +PrependBuffer(data []byte) error +``` + +Each connector prepends `data` to its internal unread buffer. The next `Read` or `ReadTo` consumes these bytes before reading from the network. + +During `GenericExecute`: + +1. gnetcli writes the command and waits for echo, prompt, pager, question, login, and callbacks as usual. +2. If the complete echo is matched, execution proceeds normally. +3. If a prompt is matched before echo, gnetcli attempts to reconstruct the echo from the bytes before that prompt. +4. If reconstruction fails, gnetcli records the associated `EchoReadException`, prepends those bytes through `PrependBuffer`, and continues reading with the same expression set. +5. If the complete echo later matches, normal execution resumes. If reading ends or times out first, gnetcli returns the saved echo error; a device error detected in a read timeout still takes precedence. + +The behavior applies to all `streamer.Connector` implementations. SSH returns an error if `PrependBuffer` is called before its shell session is initialized; `GenericExecute` propagates that failure. + +## Compatibility + +- Normal command execution is unchanged when echo is received before prompt. +- `Read(n)` retains its existing meaning: it returns immediately when at least `n` bytes are already buffered, including prepended bytes. This is required for protocol readers such as NETCONF chunk framing. +- `PrependBuffer` is a new method on the public `streamer.Connector` interface. Third-party connector implementations must add it. + +## Rationale + +The connector already owns unread transport bytes. Returning consumed pre-prompt data to that buffer preserves the existing command-execution matcher order and avoids adding a separate pre-echo state machine to `GenericExecute`. + +The alternative of disabling prompt matching until echo is found avoids the premature match, but changes matching behavior and requires managing separate matcher states. Buffer restoration instead retries with the original expressions and preserves the normal execution flow. + +## Alternatives + +### Explicit pre-echo state machine + +`GenericExecute` could track separate phases such as `awaitEchoWithPrompt`, `awaitEchoWithoutPrompt`, and `afterEcho`. After a prompt arrives before echo, it would rebuild its expression list without the prompt matcher; after echo is found, it would rebuild the normal prompt-first matcher. + +This makes the execution phases explicit, but it adds state transitions, duplicates matcher construction, and requires care to restore the correct matcher after every transition, including questions and callbacks. It also changes the expressions that are active during execution. + +The accepted approach keeps the existing matcher set and returns consumed bytes to the transport's unread buffer. The next read evaluates the original expressions against the complete reconstructed stream, so the recovery logic stays localized to the prompt-before-echo case. + +## Testing Plan + +- RouterOS mock tests split a redraw sequence at a prompt boundary and verify that the complete command echo is later recognized. +- Streamer tests verify that buffered data is consumed before network data and that `Read(n)` preserves its total-buffer size semantics. + +## References + +- `pkg/device/genericcli/genericcli.go` — `GenericExecute` echo recovery. +- `pkg/streamer/streamer.go` — `Connector` and unread-buffer handling. +- `pkg/device/ros/device_mock_test.go` — RouterOS redraw regression test. diff --git a/internal/gvendor/gvendor.go b/internal/gvendor/gvendor.go index 8d37a1c3..8a8a7922 100644 --- a/internal/gvendor/gvendor.go +++ b/internal/gvendor/gvendor.go @@ -459,7 +459,7 @@ func (m *connWrapper) Close() error { func (m *connWrapper) ReadTo(ctx context.Context, expr expr.Expr) (streamer.ReadRes, error) { m.log.Debug("read to", zap.String("expr", expr.Repr())) - res, extra, read, err := streamer.GenericReadX(ctx, m.stdoutBufferExtra, m.stdoutBuffer, defaultReadSize, readTimeout, expr, 0, 0) + res, extra, read, err := streamer.GenericReadX(ctx, m.stdoutBufferExtra, m.stdoutBuffer, defaultReadSize, readTimeout, streamer.WithRegExpr(expr)) m.stdoutBufferExtra = extra if err != nil { return nil, err diff --git a/pkg/device/errors.go b/pkg/device/errors.go index d506d48b..54cd485c 100644 --- a/pkg/device/errors.go +++ b/pkg/device/errors.go @@ -26,8 +26,9 @@ func ThrowExecException(data string) error { } type EchoReadException struct { - lastRead []byte - promptFound bool // indicates if we found prompt after echo read error + lastRead []byte + promptFound bool // indicates if we found prompt after echo read error + questionFound bool // indicates if we found question after echo error } func (e *EchoReadException) Error() string { @@ -39,10 +40,16 @@ func (e *EchoReadException) PromptFound() bool { return e.promptFound } -func ThrowEchoReadException(lastRead []byte, promptFound bool) error { +// QuestionFound indicates if gnetcli succeeded in reading question after echo read failure +func (e *EchoReadException) QuestionFound() bool { + return e.questionFound +} + +func ThrowEchoReadException(lastRead []byte, promptFound bool, questionFound bool) error { return &EchoReadException{ - lastRead: lastRead, - promptFound: promptFound, + lastRead: lastRead, + promptFound: promptFound, + questionFound: questionFound, } } diff --git a/pkg/device/genericcli/genericcli.go b/pkg/device/genericcli/genericcli.go index a48f9009..1114bd33 100644 --- a/pkg/device/genericcli/genericcli.go +++ b/pkg/device/genericcli/genericcli.go @@ -593,6 +593,8 @@ func GenericExecute(command cmd.Cmd, connector streamer.Connector, cli GenericCL if len(cmdQuestions) > 0 { questions = append(cmdQuestions, questions...) } + + exprsAdd, exprsAddMap := command.GetExprCallback() checkExprs := []expr.NamedExpr{ {Name: echoExprName, Exprs: []expr.Expr{expCmdEcho}}, {Name: promptExprName, Exprs: []expr.Expr{cli.prompt}}, @@ -603,14 +605,18 @@ func GenericExecute(command cmd.Cmd, connector streamer.Connector, cli GenericCL checkExprs = append(checkExprs, expr.NamedExpr{Name: loginExprName, Exprs: []expr.Expr{cli.login}}) } exprs := expr.NewSimpleExprListNamedOrdered(checkExprs) - - exprsAdd, exprsAddMap := command.GetExprCallback() + callbackPatternStart := 0 + for _, namedExpr := range checkExprs { + callbackPatternStart += len(namedExpr.Exprs) + } for _, exprCB := range exprsAdd { exprs.Add("cb", expr.NewSimpleExpr().FromPattern(exprCB)) } + cbLimit := 100 seenEcho := false var lastQuestion []byte // GOP3 + var lastPromptBeforeEchoError error repeatedQuestionCount := 0 for { // pager loop match, err := connector.ReadTo(ctx, exprs) @@ -623,6 +629,15 @@ func GenericExecute(command cmd.Cmd, connector streamer.Connector, cli GenericCL return nil, outputErr } } + // This case means we got prompt without echo previously. + // This could mean 2 separate problems, which are hard to distinguish: + // 1) Device redraws terminal, partially echoing; we got chunk ending on prompt before device wrote full echo + // 2) Prompt/echo is configured incorrect for this device. + // For the 1) case we prepend existing buffer and retry read until we read echo. If we receive read error - it was actually 2) case - so we return original error. + // gop4 for more details. + if lastPromptBeforeEchoError != nil { + return nil, lastPromptBeforeEchoError + } return nil, err } matchId := match.GetPatternNo() @@ -631,65 +646,71 @@ func GenericExecute(command cmd.Cmd, connector streamer.Connector, cli GenericCL if matchName == echoExprName { seenEcho = true exprs.Delete(echoExprName) + callbackPatternStart-- + lastPromptBeforeEchoError = nil continue } mbefore := match.GetBefore() - if !seenEcho { - if matchName == questionExprName { // caught question before echo - // check for echo, drop it and proceed with question - termParsedEcho, err := terminal.ParseDropLastReturn(mbefore) - if err != nil { - return nil, fmt.Errorf("echo terminal parse error %w", err) - } - mres, ok := exprs.Match(termParsedEcho) - if !ok { - return nil, device.ThrowEchoReadException(mbefore, true) - } - if exprs.GetName(mres.PatternNo) == echoExprName { - seenEcho = true - } - mbefore = termParsedEcho[mres.End:] - } - } - - if !seenEcho { - promptFound := matchName == promptExprName - // case where we caught prompt before echo because of term codes in echo - if len(mbefore) < 2 || !promptFound { // don't bother to do complex logic - return nil, device.ThrowEchoReadException(mbefore, promptFound) - } - - termParsedEcho, err := terminal.ParseDropLastReturn(mbefore) + seenPrompt := matchName == promptExprName + seenQuestion := matchName == questionExprName + checkEcho := func(mBefore []byte) ([]byte, error) { + // Check for echo, drop it, and continue handling the matched expression. + termParsedEcho, err := terminal.ParseDropLastReturn(mBefore) if err != nil { - return nil, fmt.Errorf("echo terminal parse error %w", err) + return nil, fmt.Errorf("echo terminal before question parse error %w", err) } mres, ok := exprs.Match(termParsedEcho) if !ok { - // prompt expression may consume newline from echo, but it must be presented in echo - if mbefore[len(mbefore)-1] != '\n' { - mbefore = append(mbefore, '\n') - } - termParsedEcho, err = terminal.ParseDropLastReturn(mbefore) - if err != nil { - return nil, fmt.Errorf("echo terminal parse error %w", err) - } - mres, ok = exprs.Match(termParsedEcho) - if !ok { - return nil, device.ThrowEchoReadException(mbefore, promptFound) - } + return nil, device.ThrowEchoReadException(mbefore, seenPrompt, seenQuestion) } - // assuring that it is echo if exprs.GetName(mres.PatternNo) != echoExprName { - return nil, device.ThrowEchoReadException(mbefore, promptFound) - } - if mres.End > len(termParsedEcho) { - return nil, errors.New("termParsedEcho len less than mres.End") + return nil, device.ThrowEchoReadException(mbefore, seenPrompt, seenQuestion) } seenEcho = true exprs.Delete(echoExprName) - // delete echo - mbefore = termParsedEcho[mres.End:] + callbackPatternStart-- + lastPromptBeforeEchoError = nil + return termParsedEcho[mres.End:], nil } + if !seenEcho { + if len(mbefore) < 2 { + return nil, device.ThrowEchoReadException(mbefore, seenPrompt, seenQuestion) + } + switch { + case matchName == questionExprName: + // check for echo, drop it and proceed with question + mBefore, err := checkEcho(mbefore) + if err != nil { + return nil, err + } + mbefore = mBefore + case matchName == promptExprName: + mBefore, err := checkEcho(mbefore) + // prompt expression may consume newline from echo, but it must be presented in echo + if err != nil && mbefore[len(mbefore)-1] != '\n' { + mBefore, err = checkEcho(append(mbefore, '\n')) + } + // GOP 4: a CLI may redraw its prompt before the complete echo arrives. + // Preserve consumed bytes and retry the normal matcher with them prepended. + if err != nil { + lastPromptBeforeEchoError = err + err := connector.PrependBuffer(mbefore) + if err != nil { + return nil, fmt.Errorf( + "prepend consumed output after prompt-before-echo: %w; %w", + err, + device.ThrowEchoReadException(mbefore, true, false), + ) + } + continue + } + + mbefore = mBefore + default: + return nil, device.ThrowEchoReadException(mbefore, false, false) + } + } + if matchName == promptExprName { buffer.Write(mbefore) if store, ok := match.GetMatchedGroups()["store"]; ok { @@ -732,7 +753,11 @@ func GenericExecute(command cmd.Cmd, connector streamer.Connector, cli GenericCL return nil, fmt.Errorf("callback limit") } cbLimit-- - wr := exprsAddMap[exprsAdd[matchId-3]] + callbackIndex := matchId - callbackPatternStart + if callbackIndex < 0 || callbackIndex >= len(exprsAdd) { + return nil, fmt.Errorf("invalid callback pattern index %d", matchId) + } + wr := exprsAddMap[exprsAdd[callbackIndex]] logger.Debug("write callback result") err := connector.Write([]byte(wr)) if err != nil { diff --git a/pkg/device/genericcli/genericcli_test.go b/pkg/device/genericcli/genericcli_test.go index a7e38555..465a52d7 100644 --- a/pkg/device/genericcli/genericcli_test.go +++ b/pkg/device/genericcli/genericcli_test.go @@ -314,3 +314,33 @@ func TestQuestionWithAnswerNotSendNL(t *testing.T) { require.NoError(t, resErr) require.Equal(t, cmdRes, []cmd.CmdRes{cmd.NewCmdRes(nil)}) } + +func TestCommandCallbackWithAdditionalQuestionExpression(t *testing.T) { + logger := zap.Must(zap.NewDevelopmentConfig().Build()) + dialog := [][]gmock.Action{ + { + gmock.Send(""), + gmock.Expect("test\n"), + gmock.SendEcho("test\r\n"), + gmock.Send("callback"), + gmock.Expect("answer"), + gmock.Send("done\r\n"), + gmock.Close(), + }, + } + + command := cmd.NewCmd( + "test", + cmd.WithAddAnswers(cmd.NewAnswer("unused-question", "unused", true)), + cmd.WithExprCallback(cmd.NewExprCallback("callback", "answer")), + ) + cmdRes, resErr, serverErr, err := gmock.RunCmd(func(connector streamer.Connector) device.Device { + dev := newDevice(fullQuestion, connector, logger) + return &dev + }, gmock.ConcatMultipleSlices(dialog), []cmd.Cmd{command}, logger) + + require.NoError(t, err) + require.NoError(t, serverErr) + require.NoError(t, resErr) + require.Equal(t, []cmd.CmdRes{cmd.NewCmdRes([]byte("done"))}, cmdRes) +} diff --git a/pkg/device/juniper/device_mock_test.go b/pkg/device/juniper/device_mock_test.go index 069226ff..611a23bc 100644 --- a/pkg/device/juniper/device_mock_test.go +++ b/pkg/device/juniper/device_mock_test.go @@ -208,7 +208,7 @@ func TestInvalidShowCommandsWithException(t *testing.T) { }, everyDayByeBye, }, - err: device.ThrowEchoReadException([]byte("dis \r\n ^\r\nunknown command.\r\n"), true), + err: device.ThrowEchoReadException([]byte("dis \r\n ^\r\nunknown command.\r\n"), true, false), }, } diff --git a/pkg/device/ros/device_mock_test.go b/pkg/device/ros/device_mock_test.go index c3989a08..a26519aa 100644 --- a/pkg/device/ros/device_mock_test.go +++ b/pkg/device/ros/device_mock_test.go @@ -124,168 +124,172 @@ func TestRos(t *testing.T) { " " + " " + " \u001b[K\r[username12345" + - "@mk-rb3011-test] > /ip ser\r\r[username12345@mk-rb3011-test] > /ip ser " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b" + - "[K\r[username12345@mk-rb3011-test] > /ip serv\r\r[username12345@mk-rb3011-test] > /ip serv " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[username12345@mk-rb3011-test] > /ip servi\r\r[username12345@mk-rb3011-" + - "test] > /ip servi " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[username12345@mk-rb3011-test] > /ip servic\r\r[use" + - "rname12345@mk-rb3011-test] > /ip servic " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[username12345@mk-rb3011-test]" + - " > /ip service\r\r[username12345@mk-rb3011-test] > /ip service " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[usernam" + - "e12345@mk-rb3011-test] > /ip service \r\r[username12345@mk-rb3011-test] > /ip service " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[username12345@mk-rb3011-test] > /ip service e\r\r[username12345@mk-rb3011-test" + - "] > /ip service e " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[username12345@mk-rb3011-test] > /ip service ex\r\r[use" + - "rname12345@mk-rb3011-test] > /ip service ex " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[username12345@mk-rb3011-test]" + - " > /ip service exp\r\r[username12345@mk-rb3011-test] > /ip service exp " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[use" + - "rname12345@mk-rb3011-test] > /ip service expo\r\r[username12345@mk-rb3011-test] > /ip service " + - "expo " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[username12345@mk-rb3011-test] > /ip service expor\r\r[username12345@mk" + - "-rb3011-test] > /ip service expor " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[username12345@mk-rb3011-test] > /ip servic" + - "e export\r\r[username12345@mk-rb3011-test] > /ip service export " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r[username12345" + - "@mk-rb3011-test] > /ip service export\r[username12345@mk-rb3011-test] > /ip service export\r\n" + - "\r# oct/17/2024 10:17:11 by RouterOS 6.49.17\r\n# software id = 1111-1111\r\n#\r\n# model = RB" + - "3011UiAS\r\n# serial number = 111111111111\r\n/ip service\r\nset telnet disabled=yes\r\nset ft" + - "p disabled=yes\r\nset www disabled=yes\r\nset api disabled=yes\r\nset api-ssl disabled=yes\r\n" + - "\r\r\r\r[username12345@mk-rb3011-test] > " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \r[username12345@mk-rb3011-t" + - "est] > \r\r[username12345@mk-rb3011-test3] > " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " " + - " \u001b[K\r" + - "", + "@mk-rb3011-test] > ", + ), + m.Sleep(1), + m.Send( + "/ip ser\r\r[username12345@mk-rb3011-test] > /ip ser " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b" + + "[K\r[username12345@mk-rb3011-test] > /ip serv\r\r[username12345@mk-rb3011-test] > /ip serv " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[username12345@mk-rb3011-test] > /ip servi\r\r[username12345@mk-rb3011-" + + "test] > /ip servi " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[username12345@mk-rb3011-test] > /ip servic\r\r[use" + + "rname12345@mk-rb3011-test] > /ip servic " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[username12345@mk-rb3011-test]" + + " > /ip service\r\r[username12345@mk-rb3011-test] > /ip service " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[usernam" + + "e12345@mk-rb3011-test] > /ip service \r\r[username12345@mk-rb3011-test] > /ip service " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[username12345@mk-rb3011-test] > /ip service e\r\r[username12345@mk-rb3011-test" + + "] > /ip service e " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[username12345@mk-rb3011-test] > /ip service ex\r\r[use" + + "rname12345@mk-rb3011-test] > /ip service ex " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[username12345@mk-rb3011-test]" + + " > /ip service exp\r\r[username12345@mk-rb3011-test] > /ip service exp " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[use" + + "rname12345@mk-rb3011-test] > /ip service expo\r\r[username12345@mk-rb3011-test] > /ip service " + + "expo " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[username12345@mk-rb3011-test] > /ip service expor\r\r[username12345@mk" + + "-rb3011-test] > /ip service expor " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[username12345@mk-rb3011-test] > /ip servic" + + "e export\r\r[username12345@mk-rb3011-test] > /ip service export " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r[username12345" + + "@mk-rb3011-test] > /ip service export\r[username12345@mk-rb3011-test] > /ip service export\r\n" + + "\r# oct/17/2024 10:17:11 by RouterOS 6.49.17\r\n# software id = 1111-1111\r\n#\r\n# model = RB" + + "3011UiAS\r\n# serial number = 111111111111\r\n/ip service\r\nset telnet disabled=yes\r\nset ft" + + "p disabled=yes\r\nset www disabled=yes\r\nset api disabled=yes\r\nset api-ssl disabled=yes\r\n" + + "\r\r\r\r[username12345@mk-rb3011-test] > " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \r[username12345@mk-rb3011-t" + + "est] > \r\r[username12345@mk-rb3011-test3] > " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " " + + " \u001b[K\r" + + "", ), m.Send("[username12345@mk-rb3011-test2] > "), }, diff --git a/pkg/streamer/console/console.go b/pkg/streamer/console/console.go index 405657e2..9a2754bf 100644 --- a/pkg/streamer/console/console.go +++ b/pkg/streamer/console/console.go @@ -117,10 +117,11 @@ func (m *Streamer) GetBuffer() []byte { // PrependBuffer makes already consumed data available to the next read. // It is intended for protocol detectors which peek at the stream before the // device implementation starts its own login sequence. -func (m *Streamer) PrependBuffer(data []byte) { +func (m *Streamer) PrependBuffer(data []byte) error { buffer := make([]byte, 0, len(data)+len(m.bufferExtra)) buffer = append(buffer, data...) m.bufferExtra = append(buffer, m.bufferExtra...) + return nil } func (m *Streamer) FlushBuffer() { @@ -897,7 +898,7 @@ func (m *Streamer) closeForChangePort() error { func (m *Streamer) ReadTo(ctx context.Context, exp expr.Expr) (streamer.ReadRes, error) { m.logger.Debug("read to", zap.String("expr", exp.Repr())) exprs := expr.NewSimpleExprList(exp, expr.NewSimpleExpr().FromPattern(regExErrors)) - res, extra, read, err := streamer.GenericReadX(ctx, m.bufferExtra, m.buffer, readBufferSize, m.readTimeout, exprs, 0, 0) + res, extra, read, err := streamer.GenericReadX(ctx, m.bufferExtra, m.buffer, readBufferSize, m.readTimeout, streamer.WithRegExpr(exprs)) if m.trace != nil { m.trace(trace.Read, read) } @@ -940,7 +941,7 @@ func (m *Streamer) CheckConsoleError(readRes streamer.ReadRes) error { func (m *Streamer) Read(ctx context.Context, size int) ([]byte, error) { m.logger.Debug("read", zap.Int("size", size)) - res, extra, read, err := streamer.GenericReadX(ctx, m.bufferExtra, m.buffer, readBufferSize, m.readTimeout, nil, size, 0) + res, extra, read, err := streamer.GenericReadX(ctx, m.bufferExtra, m.buffer, readBufferSize, m.readTimeout, streamer.WithMaxReadSize(size)) if err == nil && res.RetType != streamer.Size { return nil, fmt.Errorf("unexpected res type %d", res.RetType) } @@ -953,7 +954,7 @@ func (m *Streamer) Read(ctx context.Context, size int) ([]byte, error) { func (m *Streamer) XRead(ctx context.Context, size int, duration time.Duration, expr expr.Expr) (*streamer.ReadXRes, error) { m.logger.Debug("read to", zap.Int("size", size), zap.Any("expr", expr), zap.Duration("duration", duration)) - res, extra, read, err := streamer.GenericReadX(ctx, m.bufferExtra, m.buffer, size, duration, expr, size, duration) + res, extra, read, err := streamer.GenericReadX(ctx, m.bufferExtra, m.buffer, size, duration, streamer.WithRegExpr(expr), streamer.WithMaxReadSize(size), streamer.WithMaxDuration(duration)) m.bufferExtra = extra if m.trace != nil { m.trace(trace.Read, read) diff --git a/pkg/streamer/rfc2217/rfc2217.go b/pkg/streamer/rfc2217/rfc2217.go index 42694c45..3ac49d42 100644 --- a/pkg/streamer/rfc2217/rfc2217.go +++ b/pkg/streamer/rfc2217/rfc2217.go @@ -95,6 +95,13 @@ func (m *Streamer) SetReadTimeout(duration time.Duration) time.Duration { return prev } +func (m *Streamer) PrependBuffer(data []byte) error { + buffer := make([]byte, 0, len(data)+len(m.stdoutBufferExtra)) + buffer = append(buffer, data...) + m.stdoutBufferExtra = append(buffer, m.stdoutBufferExtra...) + return nil +} + func (m *Streamer) SetTrace(cb trace.CB) { m.trace = cb } @@ -253,7 +260,7 @@ func (m *Streamer) Read(context.Context, int) ([]byte, error) { func (m *Streamer) ReadTo(ctx context.Context, expr expr.Expr) (streamer.ReadRes, error) { m.logger.Debug("read to", zap.String("expr", expr.Repr())) - res, extra, read, err := streamer.GenericReadX(ctx, m.stdoutBufferExtra, m.stdoutBuffer, defaultReadSize, m.readTimeout, expr, 0, 0) + res, extra, read, err := streamer.GenericReadX(ctx, m.stdoutBufferExtra, m.stdoutBuffer, defaultReadSize, m.readTimeout, streamer.WithRegExpr(expr)) if m.trace != nil { m.trace(trace.Read, read) } diff --git a/pkg/streamer/ssh/ssh.go b/pkg/streamer/ssh/ssh.go index e6a877af..845a6b08 100644 --- a/pkg/streamer/ssh/ssh.go +++ b/pkg/streamer/ssh/ssh.go @@ -324,7 +324,7 @@ func (m *Streamer) Read(ctx context.Context, size int) ([]byte, error) { return nil, err } } - res, extra, read, err := streamer.GenericReadX(ctx, m.session.stdoutBufferExtra, m.session.stdoutBuffer, defaultReadSize, m.readTimeout, nil, size, 0) + res, extra, read, err := streamer.GenericReadX(ctx, m.session.stdoutBufferExtra, m.session.stdoutBuffer, defaultReadSize, m.readTimeout, streamer.WithMaxReadSize(size)) if m.trace != nil { m.trace(trace.Read, read) } @@ -347,7 +347,7 @@ func (m *Streamer) ReadTo(ctx context.Context, expr expr.Expr) (streamer.ReadRes return nil, err } } - res, extra, read, err := streamer.GenericReadX(ctx, m.session.stdoutBufferExtra, m.session.stdoutBuffer, defaultReadSize, m.readTimeout, expr, 0, 0) + res, extra, read, err := streamer.GenericReadX(ctx, m.session.stdoutBufferExtra, m.session.stdoutBuffer, defaultReadSize, m.readTimeout, streamer.WithRegExpr(expr)) if m.trace != nil { m.trace(trace.Read, read) } @@ -365,6 +365,17 @@ func (m *Streamer) ReadTo(ctx context.Context, expr expr.Expr) (streamer.ReadRes return res.ExprRes, nil } +func (m *Streamer) PrependBuffer(data []byte) error { + if m.session == nil { + return errors.New("ssh session is not initialized") + } + + buffer := make([]byte, 0, len(data)+len(m.session.stdoutBufferExtra)) + buffer = append(buffer, data...) + m.session.stdoutBufferExtra = append(buffer, m.session.stdoutBufferExtra...) + return nil +} + func (m *Streamer) HasFeature(feature streamer.Const) bool { if feature == streamer.AutoLogin || feature == streamer.Cmd { return true diff --git a/pkg/streamer/streamer.go b/pkg/streamer/streamer.go index 1f36896b..d67245ca 100644 --- a/pkg/streamer/streamer.go +++ b/pkg/streamer/streamer.go @@ -7,7 +7,6 @@ package streamer import ( "context" "errors" - "fmt" "net" "os" "time" @@ -27,6 +26,7 @@ type Connector interface { SetCredentialsInterceptor(func(credentials.Credentials) credentials.Credentials) SetTrace(trace.CB) SetReadTimeout(time.Duration) time.Duration + PrependBuffer([]byte) error Close() ReadTo(context.Context, expr.Expr) (ReadRes, error) Read(ctx context.Context, n int) ([]byte, error) @@ -249,19 +249,49 @@ func flushCh(ch <-chan []byte) []byte { } } +type GenericReadConfig struct { + maxDuration time.Duration + maxReadSize int + regExpr expr.Expr +} + +type GenericReadOption func(*GenericReadConfig) + +// WithRegExpr stops read on regExpr match +func WithRegExpr(regExpr expr.Expr) GenericReadOption { + return func(grc *GenericReadConfig) { + grc.regExpr = regExpr + } +} + +// WithMaxDuration specifies maximum time for reading. Results in timeout result without error +func WithMaxDuration(maxDuration time.Duration) GenericReadOption { + return func(grc *GenericReadConfig) { + grc.maxDuration = maxDuration + } +} + +// WithMaxReadSize specifies maximum size of bytes returned. Any leftover bytes from read will be returned as left bytes +func WithMaxReadSize(maxReadSize int) GenericReadOption { + return func(grc *GenericReadConfig) { + grc.maxReadSize = maxReadSize + } +} + // GenericReadX reads from readCh till expr matched, exceeded time or read more than size. // Returns error if nothing was read during readTimeout or ctx was Done -// readSize - maximum read size -// maxDuration - maximum time for reading -// regExpr - read till regex match // Returns read res, left bytes, read bytes, error func GenericReadX(ctx context.Context, inBuffer []byte, readCh chan []byte, readSize int, readTimeout time.Duration, - regExpr expr.Expr, maxReadSize int, maxDuration time.Duration) (*ReadXRes, []byte, []byte, error) { - if maxDuration == 0 && maxReadSize == 0 && regExpr == nil { - return nil, nil, nil, fmt.Errorf("specify maxDuration, maxReadSize or regExpr") + requiredOpt GenericReadOption, opts ...GenericReadOption) (*ReadXRes, []byte, []byte, error) { + cfg := GenericReadConfig{} + for _, v := range append(opts, requiredOpt) { + v(&cfg) + } + if cfg.maxDuration == 0 && cfg.maxReadSize == 0 && cfg.regExpr == nil { + return nil, nil, nil, errors.New("specify maxDuration, maxReadSize or regExpr via options") } buffer := inBuffer - maxDurationTimeout := NewTimerWithDefault(maxDuration) + maxDurationTimeout := NewTimerWithDefault(cfg.maxDuration) for { select { case <-ctx.Done(): @@ -272,16 +302,16 @@ func GenericReadX(ctx context.Context, inBuffer []byte, readCh chan []byte, read } readIterTimeout := NewTimerWithDefault(readTimeout) // check size - if maxReadSize > 0 && len(buffer) >= maxReadSize { - data, extra := splitBytes(buffer, maxReadSize) + if cfg.maxReadSize > 0 && len(buffer) >= cfg.maxReadSize { + data, extra := splitBytes(buffer, cfg.maxReadSize) StopTimer(readIterTimeout) StopTimer(maxDurationTimeout) return NewReadXRes(Size, data, nil, []byte{}), extra, buffer[len(inBuffer):], nil } - if regExpr != nil { + if cfg.regExpr != nil { // check expr - mRes, ok := regExpr.Match(buffer) + mRes, ok := cfg.regExpr.Match(buffer) if ok { var underlyingRes ReadRes if mRes.Underlying != nil { diff --git a/pkg/streamer/streamer_test.go b/pkg/streamer/streamer_test.go index ed3aadd6..13d4228f 100644 --- a/pkg/streamer/streamer_test.go +++ b/pkg/streamer/streamer_test.go @@ -19,7 +19,7 @@ func TestGenericReadNSimple(t *testing.T) { buffer := []byte{} readSize := 2 readTimeout := 2 * time.Second - res, extra, read, err := GenericReadX(ctx, buffer, ch, readSize, readTimeout, nil, 5, 0) + res, extra, read, err := GenericReadX(ctx, buffer, ch, readSize, readTimeout, WithMaxReadSize(5)) left := readAll(ch) assert.NoError(t, err) @@ -37,7 +37,7 @@ func TestGenericReadNBuff(t *testing.T) { for i := 0; i < len(data); i++ { ch <- []byte{data[i]} } - res, extra, read, err := GenericReadX(ctx, buffer, ch, 1, 2*time.Second, nil, 3, 0) + res, extra, read, err := GenericReadX(ctx, buffer, ch, 1, 2*time.Second, WithMaxReadSize(3)) left := readAll(ch) assert.NoError(t, err) @@ -67,7 +67,7 @@ func TestGenericReadToSimple(t *testing.T) { readTimeout := 2 * time.Second ch := setupChan([]byte("aest")) pat := expr.NewSimpleExpr().FromPattern("es") - res, extra, read, err := GenericReadX(ctx, buffer, ch, readSize, readTimeout, pat, 0, 0) + res, extra, read, err := GenericReadX(ctx, buffer, ch, readSize, readTimeout, WithRegExpr(pat)) left := readAll(ch) assert.NoError(t, err) @@ -93,7 +93,7 @@ func TestGenericReadXCtxDoneFlushChannel(t *testing.T) { cancel() pat := expr.NewSimpleExpr().FromPattern("never-matches") - _, _, _, err := GenericReadX(ctx, nil, ch, 4096, time.Second, pat, 0, 0) + _, _, _, err := GenericReadX(ctx, nil, ch, 4096, time.Second, WithRegExpr(pat)) require.Error(t, err) require.ErrorIs(t, err, context.Canceled) @@ -113,7 +113,7 @@ func TestGenericReadXTimeoutNoExtraLeak(t *testing.T) { ctx := context.Background() pat := expr.NewSimpleExpr().FromPattern("never-matches") // readTimeout > maxDuration so that maxDurationTimeout fires first (Timeout path). - res, extra, _, err := GenericReadX(ctx, nil, ch, 4096, time.Second, pat, 0, 50*time.Millisecond) + res, extra, _, err := GenericReadX(ctx, nil, ch, 4096, time.Second, WithRegExpr(pat), WithMaxDuration(50*time.Millisecond)) require.NoError(t, err) assert.Equal(t, Timeout, res.RetType) assert.Equal(t, []byte("hello"), res.BytesRes) diff --git a/pkg/streamer/telnet/telnet.go b/pkg/streamer/telnet/telnet.go index 0f586600..b6ac88a0 100644 --- a/pkg/streamer/telnet/telnet.go +++ b/pkg/streamer/telnet/telnet.go @@ -80,6 +80,13 @@ func (m *Streamer) SetReadTimeout(duration time.Duration) time.Duration { return prev } +func (m *Streamer) PrependBuffer(data []byte) error { + buffer := make([]byte, 0, len(data)+len(m.stdoutBufferExtra)) + buffer = append(buffer, data...) + m.stdoutBufferExtra = append(buffer, m.stdoutBufferExtra...) + return nil +} + func (m *Streamer) SetTrace(cb trace.CB) { m.trace = cb } @@ -174,7 +181,7 @@ func (m *Streamer) Read(context.Context, int) ([]byte, error) { func (m *Streamer) ReadTo(ctx context.Context, expr expr.Expr) (streamer.ReadRes, error) { m.logger.Debug("read to", zap.String("expr", expr.Repr())) - res, extra, read, err := streamer.GenericReadX(ctx, m.stdoutBufferExtra, m.stdoutBuffer, defaultReadSize, m.readTimeout, expr, 0, 0) + res, extra, read, err := streamer.GenericReadX(ctx, m.stdoutBufferExtra, m.stdoutBuffer, defaultReadSize, m.readTimeout, streamer.WithRegExpr(expr)) if m.trace != nil { m.trace(trace.Read, read) }