fix(conn): plug goroutine leaks and WaitGroup race in ssh/agent packages

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,13 @@ 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()
+
+	<-errChan
+
+	return first
 }

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))
+}

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()
 }

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 {