From 71c5360f98ffb135d4e35d83d10467c2fe179569 Mon Sep 17 00:00:00 2001 From: Vitalii Levitskii Date: Wed, 24 Aug 2022 14:44:41 +0300 Subject: [PATCH] Fixed infinite loop in ExpectProcess.ExpectFunc Signed-off-by: Vitalii Levitskii --- pkg/expect/expect.go | 19 ++++++++++++++++--- pkg/expect/expect_test.go | 32 +++++++++++++++++++++++++++++++- tests/e2e/ctl_v3_elect_test.go | 3 ++- tests/e2e/ctl_v3_lock_test.go | 3 ++- tests/e2e/v3_curl_test.go | 3 ++- tests/framework/e2e/util.go | 3 ++- 6 files changed, 55 insertions(+), 8 deletions(-) diff --git a/pkg/expect/expect.go b/pkg/expect/expect.go index 648dea34f..b6cbe0114 100644 --- a/pkg/expect/expect.go +++ b/pkg/expect/expect.go @@ -18,6 +18,7 @@ package expect import ( "bufio" + "context" "fmt" "io" "os" @@ -96,7 +97,7 @@ func (ep *ExpectProcess) read() { } // ExpectFunc returns the first line satisfying the function f. -func (ep *ExpectProcess) ExpectFunc(f func(string) bool) (string, error) { +func (ep *ExpectProcess) ExpectFunc(ctx context.Context, f func(string) bool) (string, error) { i := 0 for { @@ -114,7 +115,13 @@ func (ep *ExpectProcess) ExpectFunc(f func(string) bool) (string, error) { break } ep.mu.Unlock() - time.Sleep(time.Millisecond * 10) + + select { + case <-ctx.Done(): + return "", fmt.Errorf("failed to find match string: %w", ctx.Err()) + case <-time.After(time.Millisecond * 10): + // continue loop + } } ep.mu.Lock() lastLinesIndex := len(ep.lines) - DEBUG_LINES_TAIL @@ -128,9 +135,15 @@ func (ep *ExpectProcess) ExpectFunc(f func(string) bool) (string, error) { ep.err, lastLines) } +// ExpectWithContext returns the first line containing the given string. +func (ep *ExpectProcess) ExpectWithContext(ctx context.Context, s string) (string, error) { + return ep.ExpectFunc(ctx, func(txt string) bool { return strings.Contains(txt, s) }) +} + // Expect returns the first line containing the given string. +// Deprecated: please use ExpectWithContext instead. func (ep *ExpectProcess) Expect(s string) (string, error) { - return ep.ExpectFunc(func(txt string) bool { return strings.Contains(txt, s) }) + return ep.ExpectWithContext(context.Background(), s) } // LineCount returns the number of recorded lines since diff --git a/pkg/expect/expect_test.go b/pkg/expect/expect_test.go index f091f1377..d5a22848d 100644 --- a/pkg/expect/expect_test.go +++ b/pkg/expect/expect_test.go @@ -17,9 +17,12 @@ package expect import ( + "context" "os" "testing" "time" + + "github.com/stretchr/testify/require" ) func TestExpectFunc(t *testing.T) { @@ -28,7 +31,7 @@ func TestExpectFunc(t *testing.T) { t.Fatal(err) } wstr := "hello world\r\n" - l, eerr := ep.ExpectFunc(func(a string) bool { return len(a) > 10 }) + l, eerr := ep.ExpectFunc(context.Background(), func(a string) bool { return len(a) > 10 }) if eerr != nil { t.Fatal(eerr) } @@ -40,6 +43,33 @@ func TestExpectFunc(t *testing.T) { } } +func TestExpectFuncTimeout(t *testing.T) { + ep, err := NewExpect("tail", "-f", "/dev/null") + if err != nil { + t.Fatal(err) + } + go func() { + // It's enough to have "talkative" process to stuck in the infinite loop of reading + for { + err := ep.Send("new line\n") + if err != nil { + return + } + } + }() + + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + + _, err = ep.ExpectFunc(ctx, func(a string) bool { return false }) + + require.ErrorAs(t, err, &context.DeadlineExceeded) + + if err = ep.Stop(); err != nil { + t.Fatal(err) + } +} + func TestEcho(t *testing.T) { ep, err := NewExpect("echo", "hello world") if err != nil { diff --git a/tests/e2e/ctl_v3_elect_test.go b/tests/e2e/ctl_v3_elect_test.go index 40b15f69b..4aecae07d 100644 --- a/tests/e2e/ctl_v3_elect_test.go +++ b/tests/e2e/ctl_v3_elect_test.go @@ -15,6 +15,7 @@ package e2e import ( + "context" "os" "strings" "testing" @@ -106,7 +107,7 @@ func ctlV3Elect(cx ctlCtx, name, proposal string) (*expect.ExpectProcess, <-chan return proc, outc, err } go func() { - s, xerr := proc.ExpectFunc(func(string) bool { return true }) + s, xerr := proc.ExpectFunc(context.TODO(), func(string) bool { return true }) if xerr != nil { cx.t.Errorf("expect failed (%v)", xerr) } diff --git a/tests/e2e/ctl_v3_lock_test.go b/tests/e2e/ctl_v3_lock_test.go index 26ea22501..ef9d8a73f 100644 --- a/tests/e2e/ctl_v3_lock_test.go +++ b/tests/e2e/ctl_v3_lock_test.go @@ -15,6 +15,7 @@ package e2e import ( + "context" "fmt" "os" "strings" @@ -127,7 +128,7 @@ func ctlV3Lock(cx ctlCtx, name string) (*expect.ExpectProcess, <-chan string, er return proc, outc, err } go func() { - s, xerr := proc.ExpectFunc(func(string) bool { return true }) + s, xerr := proc.ExpectFunc(context.TODO(), func(string) bool { return true }) if xerr != nil { cx.t.Errorf("expect failed (%v)", xerr) } diff --git a/tests/e2e/v3_curl_test.go b/tests/e2e/v3_curl_test.go index 35450bbe4..e01e8dd44 100644 --- a/tests/e2e/v3_curl_test.go +++ b/tests/e2e/v3_curl_test.go @@ -15,6 +15,7 @@ package e2e import ( + "context" "encoding/base64" "encoding/json" "fmt" @@ -249,7 +250,7 @@ func testV3CurlAuth(cx ctlCtx) { testutil.AssertNil(cx.t, err) defer proc.Close() - cURLRes, err := proc.ExpectFunc(lineFunc) + cURLRes, err := proc.ExpectFunc(context.Background(), lineFunc) testutil.AssertNil(cx.t, err) authRes := make(map[string]interface{}) diff --git a/tests/framework/e2e/util.go b/tests/framework/e2e/util.go index 6db4e404d..dc454715d 100644 --- a/tests/framework/e2e/util.go +++ b/tests/framework/e2e/util.go @@ -15,6 +15,7 @@ package e2e import ( + "context" "encoding/json" "fmt" "math/rand" @@ -36,7 +37,7 @@ func WaitReadyExpectProc(exproc *expect.ExpectProcess, readyStrs []string) error } return false } - _, err := exproc.ExpectFunc(matchSet) + _, err := exproc.ExpectFunc(context.Background(), matchSet) return err }