Merge pull request #3790 from crazy-max/dial-stop

commands: stop dial-stdio when the builder connection closes

commands/dial_stdio.go

@@ -1,6 +1,7 @@
 package commands
 
 import (
+	"context"
 	"io"
 	"net"
 	"os"
@@ -15,7 +16,6 @@ import (
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
-	"golang.org/x/sync/errgroup"
 )
 
 type stdioOptions struct {
@@ -79,27 +79,43 @@ func runDialStdio(dockerCli command.Cli, opts stdioOptions) error {
 			return err
 		}
 
+		return proxyConn(ctx, conn, os.Stdin, os.Stdout)
+	})
+}
+
+func proxyConn(ctx context.Context, conn net.Conn, stdin io.Reader, stdout io.Writer) error {
 	defer conn.Close()
 
+	stdinDone := make(chan error, 1)
+	stdoutDone := make(chan error, 1)
+
 	go func() {
-			<-ctx.Done()
+		_, err := io.Copy(conn, stdin)
 		closeWrite(conn)
+		stdinDone <- err
+	}()
+	go func() {
+		_, err := io.Copy(stdout, conn)
+		closeRead(conn)
+		stdoutDone <- err
 	}()
 
-		var eg errgroup.Group
-
-		eg.Go(func() error {
-			_, err := io.Copy(conn, os.Stdin)
-			closeWrite(conn)
+	for {
+		select {
+		case <-ctx.Done():
+			return context.Cause(ctx)
+		case err := <-stdinDone:
+			if err != nil && !errors.Is(err, net.ErrClosed) && !errors.Is(err, io.ErrClosedPipe) {
 				return err
-		})
-		eg.Go(func() error {
-			_, err := io.Copy(os.Stdout, conn)
-			closeRead(conn)
+			}
+			stdinDone = nil
+		case err := <-stdoutDone:
+			if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) && !errors.Is(err, io.ErrClosedPipe) {
 				return err
-		})
-		return eg.Wait()
-	})
+			}
+			return nil
+		}
+	}
 }
 
 func closeRead(conn net.Conn) error {

commands/dial_stdio_test.go

@@ -0,0 +1,56 @@
+package commands
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestProxyConnRemoteClose(t *testing.T) {
+	clientConn, serverConn := net.Pipe()
+	defer serverConn.Close()
+
+	stdin := &blockingReader{waitCh: make(chan struct{})}
+	defer stdin.Close()
+
+	var stdout bytes.Buffer
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- proxyConn(context.Background(), clientConn, stdin, &stdout)
+	}()
+
+	go func() {
+		_, _ = serverConn.Write([]byte("hello"))
+		_ = serverConn.Close()
+	}()
+
+	select {
+	case err := <-errCh:
+		require.NoError(t, err)
+		require.Equal(t, "hello", stdout.String())
+	case <-time.After(2 * time.Second):
+		t.Fatal("proxyConn did not return after the remote side closed")
+	}
+}
+
+type blockingReader struct {
+	waitCh    chan struct{}
+	closeOnce sync.Once
+}
+
+func (r *blockingReader) Read([]byte) (int, error) {
+	<-r.waitCh
+	return 0, io.EOF
+}
+
+func (r *blockingReader) Close() {
+	r.closeOnce.Do(func() {
+		close(r.waitCh)
+	})
+}