From 81ad0bc7fdecb0e2fe5237a01e4153dec7983890 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Fri, 18 Oct 2024 14:20:41 -0400 Subject: [PATCH 1/2] Use concurrency-safe bytes buffer to avoid race Signed-off-by: Danny Kopping --- provisionersdk/agent_test.go | 43 ++++++++++++++++++++++++++++-------- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/provisionersdk/agent_test.go b/provisionersdk/agent_test.go index 5dbdc41a89cd2..147d8cc86dbd7 100644 --- a/provisionersdk/agent_test.go +++ b/provisionersdk/agent_test.go @@ -8,7 +8,6 @@ package provisionersdk_test import ( "bytes" - "context" "errors" "fmt" "net/http" @@ -47,12 +46,10 @@ func TestAgentScript(t *testing.T) { t.Run("Valid", func(t *testing.T) { t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) script := serveScript(t, bashEcho) - ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort) - t.Cleanup(cancel) - - var output bytes.Buffer + var output safeBuffer // This is intentionally ran in single quotes to mimic how a customer may // embed our script. Our scripts should not include any single quotes. // nolint:gosec @@ -84,12 +81,10 @@ func TestAgentScript(t *testing.T) { t.Run("Invalid", func(t *testing.T) { t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) script := serveScript(t, unexpectedEcho) - ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort) - t.Cleanup(cancel) - - var output bytes.Buffer + var output safeBuffer // This is intentionally ran in single quotes to mimic how a customer may // embed our script. Our scripts should not include any single quotes. // nolint:gosec @@ -159,3 +154,33 @@ func serveScript(t *testing.T, in string) string { script = strings.ReplaceAll(script, "${AUTH_TYPE}", "token") return script } + +// safeBuffer is a concurrency-safe bytes.Buffer +type safeBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (sb *safeBuffer) Write(p []byte) (n int, err error) { + sb.mu.Lock() + defer sb.mu.Unlock() + return sb.buf.Write(p) +} + +func (sb *safeBuffer) Read(p []byte) (n int, err error) { + sb.mu.Lock() + defer sb.mu.Unlock() + return sb.buf.Read(p) +} + +func (sb *safeBuffer) Bytes() []byte { + sb.mu.Lock() + defer sb.mu.Unlock() + return sb.buf.Bytes() +} + +func (sb *safeBuffer) String() string { + sb.mu.Lock() + defer sb.mu.Unlock() + return sb.buf.String() +} \ No newline at end of file From b90dbca1406abc8e810ba1ed5590fcb44071b977 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Fri, 18 Oct 2024 14:31:10 -0400 Subject: [PATCH 2/2] make fmt Signed-off-by: Danny Kopping --- provisionersdk/agent_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/provisionersdk/agent_test.go b/provisionersdk/agent_test.go index 147d8cc86dbd7..3be01e20dce6f 100644 --- a/provisionersdk/agent_test.go +++ b/provisionersdk/agent_test.go @@ -183,4 +183,4 @@ func (sb *safeBuffer) String() string { sb.mu.Lock() defer sb.mu.Unlock() return sb.buf.String() -} \ No newline at end of file +}