diff --git a/pkg/inject/inject.go b/pkg/inject/inject.go index a6875a5ea..f5517be11 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,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 } 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 {