properly handle nil input/output

Author: znull

pipe/function.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"strings"
 )
 
 // StageFunc is a function that can be used to power a `goStage`. It
@@ -66,11 +67,19 @@ func (s *goStage) Start(
 	if stdin, ok := stdin.(readerNopCloser); ok {
 		r = stdin.Reader
 	}
+	if r == nil {
+		// treat nil as empty input.
+		r = strings.NewReader("")
+	}
 
 	var w io.Writer = stdout
 	if stdout, ok := stdout.(writerNopCloser); ok {
 		w = stdout.Writer
 	}
+	if w == nil {
+		// treat nil output as /dev/null
+		w = io.Discard
+	}
 
 	go func() {
 		defer close(s.done)

pipe/memorylimit_test.go

@@ -210,7 +210,8 @@ func testMemoryLimit(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (str
 				for i := 0; i < mbs; i++ {
 					_, err := stdout.Write(bytes[:])
 					if err != nil {
-						require.ErrorIs(t, err, syscall.EPIPE)
+						assert.ErrorIs(t, err, syscall.EPIPE)
+						return nil
 					}
 				}
 
@@ -244,7 +245,8 @@ func testMemoryLimitWithObserver(t *testing.T, mbs int, limit uint64, stage pipe
 				for i := 0; i < mbs; i++ {
 					_, err := stdout.Write(bytes[:])
 					if err != nil {
-						require.ErrorIs(t, err, syscall.EPIPE)
+						assert.ErrorIs(t, err, syscall.EPIPE)
+						return nil
 					}
 				}
 				return nil

pipe/pipeline_test.go

@@ -798,6 +798,26 @@ func TestPrintf(t *testing.T) {
 	}
 }
 
+func TestPrintlnNoOutput(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	p := pipe.New()
+	p.Add(pipe.Println("Look Ma, no output!"))
+	assert.NoError(t, p.Run(ctx))
+}
+
+func TestFunctionNoInput(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	p := pipe.New()
+	p.Add(pipe.Function("read-all", func(_ context.Context, _ pipe.Env, stdin io.Reader, _ io.Writer) error {
+		n, err := io.Copy(io.Discard, stdin)
+		assert.Equal(t, int64(0), n)
+		return err
+	}))
+	assert.NoError(t, p.Run(ctx))
+}
+
 func TestErrors(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()