From f94467e9dfff1b494fae0cd1b4faeb7cca7f5509 Mon Sep 17 00:00:00 2001 From: Samuel K Date: Sat, 25 Apr 2026 14:13:22 -0500 Subject: [PATCH 1/2] fix(conn): plug goroutine leaks and WaitGroup race in ssh/agent packages - helper.go: setupContextCancellation returns cleanup func to unblock goroutine via done channel instead of self-referencing exit channel - forward.go: convert both forward() and reverseForward() first goroutine to waitGroup.Go() eliminating Add/Done race with wg.Wait() - inject.go: pipe() accepts io.WriteCloser/io.ReadCloser and closes both endpoints after first copy completes, preventing orphaned goroutine - Add comprehensive tests for pipe(), readLine(), waitForMessage(), and performMutualHandshake() --- pkg/inject/inject.go | 13 +- pkg/inject/inject_test.go | 251 ++++++++++++++++++++++++++++++++++++++ pkg/ssh/forward.go | 12 +- pkg/ssh/helper.go | 17 +-- 4 files changed, 275 insertions(+), 18 deletions(-) create mode 100644 pkg/inject/inject_test.go diff --git a/pkg/inject/inject.go b/pkg/inject/inject.go index a6875a5ea..76d9f9a54 100644 --- a/pkg/inject/inject.go +++ b/pkg/inject/inject.go @@ -305,7 +305,10 @@ func readLine(reader io.Reader) (string, error) { } } -func pipe(toStdin io.Writer, fromStdin io.Reader, toStdout io.Writer, fromStdout io.Reader) error { +func pipe( + toStdin io.WriteCloser, fromStdin io.Reader, + toStdout io.Writer, fromStdout io.ReadCloser, +) error { errChan := make(chan error, 2) go func() { _, err := io.Copy(toStdout, fromStdout) @@ -315,5 +318,11 @@ func pipe(toStdin io.Writer, fromStdin io.Reader, toStdout io.Writer, fromStdout _, err := io.Copy(toStdin, fromStdin) errChan <- err }() - return <-errChan + + first := <-errChan + + _ = toStdin.Close() + _ = fromStdout.Close() + + return first } diff --git a/pkg/inject/inject_test.go b/pkg/inject/inject_test.go new file mode 100644 index 000000000..845627d0f --- /dev/null +++ b/pkg/inject/inject_test.go @@ -0,0 +1,251 @@ +package inject + +import ( + "bytes" + "context" + "errors" + "io" + "os" + "runtime" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/suite" +) + +// errWriter is an io.Writer that always returns a configured error. +type errWriter struct { + err error +} + +func (w *errWriter) Write(_ []byte) (int, error) { + return 0, w.err +} + +// nopWriteCloser wraps an io.Writer with a no-op Close. +type nopWriteCloser struct { + io.Writer +} + +func (nopWriteCloser) Close() error { return nil } + +// --- PipeTestSuite --- + +type PipeTestSuite struct { + suite.Suite +} + +func (s *PipeTestSuite) TestPipe_NormalBidirectionalCopy() { + fromStdinReader, fromStdinWriter := io.Pipe() + toStdoutBuf := &bytes.Buffer{} + toStdinBuf := &bytes.Buffer{} + + fromStdoutReader, fromStdoutWriter := io.Pipe() + toStdinPipeReader, toStdinPipeWriter := io.Pipe() + + errCh := make(chan error, 1) + go func() { + errCh <- pipe(toStdinPipeWriter, fromStdinReader, toStdoutBuf, fromStdoutReader) + }() + + _, err := fromStdinWriter.Write([]byte("hello from stdin")) + s.Require().NoError(err) + _ = fromStdinWriter.Close() + + _, err = fromStdoutWriter.Write([]byte("hello from stdout")) + s.Require().NoError(err) + _ = fromStdoutWriter.Close() + + received, err := io.ReadAll(toStdinPipeReader) + s.Require().NoError(err) + toStdinBuf.Write(received) + + s.NoError(<-errCh) + s.Equal("hello from stdin", toStdinBuf.String()) + s.Equal("hello from stdout", toStdoutBuf.String()) +} + +func (s *PipeTestSuite) TestPipe_WriterSideClosesFirst() { + stdinReader, stdinWriter, err := os.Pipe() + s.Require().NoError(err) + + stdoutReader, stdoutWriter, err := os.Pipe() + s.Require().NoError(err) + _ = stdoutWriter.Close() + + errCh := make(chan error, 1) + go func() { + errCh <- pipe(stdinWriter, strings.NewReader("data"), io.Discard, stdoutReader) + }() + + s.NoError(<-errCh) + _ = stdinReader.Close() +} + +func (s *PipeTestSuite) TestPipe_ReaderSideClosesFirst() { + _, stdinWriter, err := os.Pipe() + s.Require().NoError(err) + + pr, pw := io.Pipe() + + errCh := make(chan error, 1) + go func() { + errCh <- pipe(stdinWriter, &bytes.Buffer{}, io.Discard, pr) + }() + + _ = pw.Close() + s.NoError(<-errCh) +} + +func (s *PipeTestSuite) TestPipe_ErrorPropagation() { + expectedErr := errors.New("write boom") + toStdin := nopWriteCloser{&errWriter{err: expectedErr}} + + fromStdoutReader, fromStdoutWriter := io.Pipe() + go func() { + _, _ = fromStdoutWriter.Write([]byte("data")) + _ = fromStdoutWriter.Close() + }() + + err := pipe(toStdin, strings.NewReader("trigger write"), io.Discard, fromStdoutReader) + s.ErrorIs(err, expectedErr) +} + +func (s *PipeTestSuite) TestPipe_BothEndpointsClosedAfterReturn() { + stdinReader, stdinWriter, err := os.Pipe() + s.Require().NoError(err) + + stdoutReader, stdoutWriter, err := os.Pipe() + s.Require().NoError(err) + _ = stdoutWriter.Close() + + err = pipe(stdinWriter, &bytes.Buffer{}, io.Discard, stdoutReader) + s.NoError(err) + + _, writeErr := stdinWriter.Write([]byte("test")) + s.Error(writeErr, "stdinWriter should be closed after pipe returns") + + buf := make([]byte, 1) + _, readErr := stdoutReader.Read(buf) + s.Error(readErr, "stdoutReader should be closed after pipe returns") + + _ = stdinReader.Close() +} + +func (s *PipeTestSuite) TestPipe_NoGoroutineLeak() { + before := runtime.NumGoroutine() + + for range 10 { + stdinReader, stdinWriter, err := os.Pipe() + s.Require().NoError(err) + + stdoutReader, stdoutWriter, err := os.Pipe() + s.Require().NoError(err) + _ = stdoutWriter.Close() + + _ = pipe(stdinWriter, &bytes.Buffer{}, io.Discard, stdoutReader) + _ = stdinReader.Close() + } + + time.Sleep(50 * time.Millisecond) + after := runtime.NumGoroutine() + s.LessOrEqual(after, before+5, "goroutine leak detected: before=%d after=%d", before, after) +} + +func TestPipeSuite(t *testing.T) { + suite.Run(t, new(PipeTestSuite)) +} + +// --- ReadLineTestSuite --- + +type ReadLineTestSuite struct { + suite.Suite +} + +func (s *ReadLineTestSuite) TestReadLine_NormalLine() { + r := strings.NewReader("hello\n") + line, err := readLine(r) + s.NoError(err) + s.Equal("hello", line) +} + +func (s *ReadLineTestSuite) TestReadLine_MultipleLines() { + r := strings.NewReader("first\nsecond\n") + line, err := readLine(r) + s.NoError(err) + s.Equal("first", line) +} + +func (s *ReadLineTestSuite) TestReadLine_EOFBeforeNewline() { + r := strings.NewReader("no newline") + _, err := readLine(r) + s.ErrorIs(err, io.EOF) +} + +func (s *ReadLineTestSuite) TestReadLine_EmptyReader() { + r := strings.NewReader("") + _, err := readLine(r) + s.ErrorIs(err, io.EOF) +} + +func TestReadLineSuite(t *testing.T) { + suite.Run(t, new(ReadLineTestSuite)) +} + +// --- WaitForMessageTestSuite --- + +type WaitForMessageTestSuite struct { + suite.Suite +} + +func (s *WaitForMessageTestSuite) TestWaitForMessage_Success() { + ch := make(chan error, 1) + ch <- nil + err := waitForMessage(ch, time.Second) + s.NoError(err) +} + +func (s *WaitForMessageTestSuite) TestWaitForMessage_SuccessWithError() { + ch := make(chan error, 1) + expected := errors.New("something failed") + ch <- expected + err := waitForMessage(ch, time.Second) + s.ErrorIs(err, expected) +} + +func (s *WaitForMessageTestSuite) TestWaitForMessage_Timeout() { + ch := make(chan error) + err := waitForMessage(ch, 10*time.Millisecond) + s.ErrorIs(err, context.DeadlineExceeded) +} + +func TestWaitForMessageSuite(t *testing.T) { + suite.Run(t, new(WaitForMessageTestSuite)) +} + +// --- PerformMutualHandshakeTestSuite --- + +type PerformMutualHandshakeTestSuite struct { + suite.Suite +} + +func (s *PerformMutualHandshakeTestSuite) TestPerformMutualHandshake_ValidPing() { + buf := &bytes.Buffer{} + wc := nopWriteCloser{buf} + err := performMutualHandshake("ping\n", wc) + s.NoError(err) + s.Equal("pong\n", buf.String()) +} + +func (s *PerformMutualHandshakeTestSuite) TestPerformMutualHandshake_InvalidInput() { + buf := &bytes.Buffer{} + wc := nopWriteCloser{buf} + err := performMutualHandshake("hello\n", wc) + s.Error(err) + s.Contains(err.Error(), "unexpected start line") +} + +func TestPerformMutualHandshakeSuite(t *testing.T) { + suite.Run(t, new(PerformMutualHandshakeTestSuite)) +} diff --git a/pkg/ssh/forward.go b/pkg/ssh/forward.go index abaa4de86..c01be4fd0 100644 --- a/pkg/ssh/forward.go +++ b/pkg/ssh/forward.go @@ -115,15 +115,14 @@ func forward( // Copy localConn.Reader to sshConn.Writer waitGroup := sync.WaitGroup{} - go func() { - defer waitGroup.Done() + waitGroup.Go(func() { defer func() { _ = sshConn.Close() }() _, err = io.Copy(sshConn, localConn) if err != nil { log.Debugf("error copying to remote: %v", err) } - }() + }) // Copy sshConn.Reader to localConn.Writer waitGroup.Go(func() { @@ -134,7 +133,6 @@ func forward( log.Debugf("error copying to local: %v", err) } }) - waitGroup.Add(1) waitGroup.Wait() } @@ -153,15 +151,14 @@ func reverseForward( // Copy localConn.Reader to sshConn.Writer waitGroup := sync.WaitGroup{} - go func() { - defer waitGroup.Done() + waitGroup.Go(func() { defer func() { _ = localConn.Close() }() _, err = io.Copy(localConn, remoteConn) if err != nil { log.Debugf("error copying to local: %v", err) } - }() + }) // Copy sshConn.Reader to localConn.Writer waitGroup.Go(func() { @@ -172,6 +169,5 @@ func reverseForward( log.Debugf("error copying to remote: %v", err) } }) - waitGroup.Add(1) waitGroup.Wait() } diff --git a/pkg/ssh/helper.go b/pkg/ssh/helper.go index bd399bf47..7c0b36dc9 100644 --- a/pkg/ssh/helper.go +++ b/pkg/ssh/helper.go @@ -153,9 +153,11 @@ func Run(ctx context.Context, opts RunOptions) error { _ = sess.Setenv(k, v) // Ignore errors - command should work without env vars } - if err := setupContextCancellation(ctx, sess); err != nil { + cleanup, err := setupContextCancellation(ctx, sess) + if err != nil { return err } + defer cleanup() sess.Stdin = opts.Stdin sess.Stdout = opts.Stdout @@ -169,20 +171,19 @@ func Run(ctx context.Context, opts RunOptions) error { return nil } -func setupContextCancellation(ctx context.Context, sess *ssh.Session) error { +func setupContextCancellation(ctx context.Context, sess *ssh.Session) (func(), error) { if err := ctx.Err(); err != nil { - return fmt.Errorf("context already cancelled: %w", err) + return nil, fmt.Errorf("context already cancelled: %w", err) } - exit := make(chan struct{}) + done := make(chan struct{}) go func() { - defer close(exit) select { case <-ctx.Done(): - _ = sess.Signal(ssh.SIGINT) // Send interrupt, let defer handle close - case <-exit: + _ = sess.Signal(ssh.SIGINT) + case <-done: } }() - return nil + return func() { close(done) }, nil } func handleRunError(ctx context.Context, err error, command string) error { From 9858cb2d781b435c49398390c7ef5d457206ca8c Mon Sep 17 00:00:00 2001 From: Samuel K Date: Sat, 25 Apr 2026 14:29:10 -0500 Subject: [PATCH 2/2] fix(inject): wait for second goroutine in pipe() to prevent data race Drain errChan after closing endpoints so both io.Copy goroutines have fully completed before pipe() returns, making output buffers safe to read by callers. --- pkg/inject/inject.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pkg/inject/inject.go b/pkg/inject/inject.go index 76d9f9a54..f5517be11 100644 --- a/pkg/inject/inject.go +++ b/pkg/inject/inject.go @@ -324,5 +324,7 @@ func pipe( _ = toStdin.Close() _ = fromStdout.Close() + <-errChan + return first }