fix(proxy): don't shutdown 'uc proxy' when a client connection aborts

cmd/uc/proxy.go

@@ -132,35 +132,34 @@ func runProxy(ctx context.Context, uncli *cli.CLI, opts proxyOptions) error {
 	// endpoint and shuffles the data, *it* will actually experience errors.
 	remoteAddr := net.JoinHostPort(ip.String(), strconv.Itoa(opts.remotePort))
 
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
 	p := &proxy.Proxy{
 		Listener:    listener,
 		RemoteAddr:  remoteAddr,
 		DialContext: dialer.DialContext,
 		OnError: func(err error) {
-			fmt.Printf("Failed to proxy to '%s': %v\n", remoteAddr, err)
-			cancel()
+			if proxy.IsConnectionClosedError(err) {
+				return
+			}
+			// A more actionable error instead of the cryptic [ssh -W] command error.
+			if strings.Contains(err.Error(), "Session open refused by peer") {
+				fmt.Printf("Could not connect to '%s': connection refused. "+
+					"Check that the service is running and listening on port %d inside the container.\n",
+					remoteAddr, opts.remotePort)
+				return
+			}
+
+			fmt.Printf("Failed to proxy a connection to '%s': %v\n", remoteAddr, err)
 		},
 	}
 
-	// Run the proxy in the background and signal when it has fully shut down.
-	done := make(chan struct{})
-	go func() {
-		p.Run(ctx)
-		close(done)
-	}()
-
 	// Prefix the local address with the scheme for common HTTP ports so it becomes control-clickable in most
 	// terminals. We assume plain HTTP since TLS is typically terminated by Caddy in front of the service.
 	fmt.Printf("%s%s → %s (%s%s%s)\n", schemeForPort(opts.remotePort), p.Listener.Addr().String(),
 		remoteAddr, opts.service, tui.Faint.Render("/"), containerID)
 
-	<-ctx.Done()
-	// Wait for the proxy to drain in-flight connections and shut down gracefully.
-	<-done
-
+	if err = p.Run(ctx); err != nil {
+		return fmt.Errorf("run proxy to '%s': %w", remoteAddr, err)
+	}
 	return nil
 }
 

internal/proxy/proxy.go

@@ -2,11 +2,12 @@ package proxy
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net"
-	"os"
 	"sync"
+	"syscall"
 	"time"
 )
 
@@ -15,58 +16,53 @@ type Proxy struct {
 	Listener    net.Listener
 	RemoteAddr  string
 	DialContext func(ctx context.Context, network, address string) (net.Conn, error)
+	// OnError is called for errors that occur during proxying individual connections. It may be called concurrently
+	// for different connections.
 	OnError     func(error)
 	activeConns sync.WaitGroup
 }
 
-// deadliner is an interface for listeners that support setting deadlines.
-type deadliner interface {
-	SetDeadline(t time.Time) error
-}
-
 // halfCloser is an interface for connections that support half-close.
 type halfCloser interface {
 	CloseWrite() error
 }
 
-// Run starts the proxy and runs until the context is canceled.
-func (p *Proxy) Run(ctx context.Context) {
+// IsConnectionClosedError reports whether err indicates that a connection was closed or aborted by either peer.
+// Callers can use it to ignore routine connection shutdown or broken pipe errors reported to Proxy.OnError.
+func IsConnectionClosedError(err error) bool {
+	return errors.Is(err, net.ErrClosed) || errors.Is(err, io.ErrClosedPipe) ||
+		errors.Is(err, syscall.EPIPE) || errors.Is(err, syscall.ECONNRESET)
+}
+
+// Run starts the proxy and runs until the context is canceled or the listener fails. It returns nil when the context
+// is canceled. Errors handling individual connections are reported to OnError and do not stop the proxy.
+func (p *Proxy) Run(ctx context.Context) error {
 	if p.DialContext == nil {
 		p.DialContext = (&net.Dialer{}).DialContext
 	}
 
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 	defer p.Listener.Close()
 
-	// Handle incoming connections until context is canceled.
-Loop:
+	// Closing the listener unblocks Accept when the context is canceled. This works for both TCP and Unix listeners
+	// and avoids polling with listener deadlines.
+	stopClose := context.AfterFunc(ctx, func() {
+		p.Listener.Close()
+	})
+	defer stopClose()
+
+	var runErr error
 	for {
-		select {
-		case <-ctx.Done():
-			break Loop
-		default:
-		}
-
-		// Set a deadline on the listener if supported to check context periodically.
-		if dl, ok := p.Listener.(deadliner); ok {
-			dl.SetDeadline(time.Now().Add(1 * time.Second))
-		}
-
 		conn, err := p.Listener.Accept()
 		if err != nil {
-			if os.IsTimeout(err) {
-				// Just a timeout, continue to check context and accept again.
-				continue
+			if ctx.Err() != nil {
+				break
 			}
 
-			select {
-			case <-ctx.Done():
-				break Loop
-			default:
-				if p.OnError != nil {
-					p.OnError(fmt.Errorf("accept local connection: %w", err))
-				}
-				continue
-			}
+			runErr = fmt.Errorf("accept local connection: %w", err)
+			cancel()
+			break
 		}
 
 		p.activeConns.Add(1)
@@ -75,6 +71,7 @@ Loop:
 
 	// Wait for all connections to finish.
 	p.activeConns.Wait()
+	return runErr
 }
 
 func (p *Proxy) handleConnection(ctx context.Context, localConn net.Conn) {
@@ -87,46 +84,62 @@ func (p *Proxy) handleConnection(ctx context.Context, localConn net.Conn) {
 
 	remoteConn, err := p.DialContext(dialCtx, "tcp", p.RemoteAddr)
 	if err != nil {
-		if p.OnError != nil {
+		if ctx.Err() == nil && p.OnError != nil {
 			p.OnError(fmt.Errorf("connect remote address '%s': %w", p.RemoteAddr, err))
 		}
 		return
 	}
 	defer remoteConn.Close()
 
-	// Bidirectional copy with proper half-close handling.
+	// Closing both connections aborts both copies after cancellation or a copy error. A clean EOF still uses
+	// half-close so the other direction can finish sending any remaining data.
+	closeConnections := func() {
+		localConn.Close()
+		remoteConn.Close()
+	}
+	stopClose := context.AfterFunc(ctx, closeConnections)
+	defer stopClose()
+
 	done := make(chan error, 2)
 
 	go func() {
 		_, err := io.Copy(remoteConn, localConn)
+		if err != nil {
+			done <- err
+			closeConnections()
+			return
+		}
 		// Close write half of remote connection if supported.
 		if hc, ok := remoteConn.(halfCloser); ok {
 			hc.CloseWrite()
 		}
-		done <- err
+		done <- nil
 	}()
 
 	go func() {
 		_, err := io.Copy(localConn, remoteConn)
+		if err != nil {
+			done <- err
+			closeConnections()
+			return
+		}
 		// Close write half of local connection if supported.
 		if hc, ok := localConn.(halfCloser); ok {
 			hc.CloseWrite()
 		}
-		done <- err
+		done <- nil
 	}()
 
-	// Wait for both copies to complete or context cancel.
+	// Wait for both copies to complete. The first error is the original failure because a copy reports it before
+	// closing the connections to unblock the other copy.
+	var copyErr error
 	for range 2 {
-		select {
-		case <-ctx.Done():
-			// Close connections to abort ongoing copies.
-			localConn.Close()
-			remoteConn.Close()
-			return
-		case err = <-done:
-			if err != nil && p.OnError != nil {
-				p.OnError(fmt.Errorf("data copy: %w", err))
-			}
+		if err = <-done; err != nil && copyErr == nil {
+			copyErr = err
 		}
 	}
+
+	if copyErr != nil && ctx.Err() == nil && p.OnError != nil {
+		p.OnError(fmt.Errorf("data copy: %w", copyErr))
+	}
 }

internal/proxy/proxy_test.go

@@ -0,0 +1,188 @@
+package proxy
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"sync/atomic"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsConnectionClosedError(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{name: "closed network connection", err: net.ErrClosed, want: true},
+		{name: "closed pipe", err: io.ErrClosedPipe, want: true},
+		{name: "broken pipe", err: syscall.EPIPE, want: true},
+		{name: "connection reset", err: syscall.ECONNRESET, want: true},
+		{name: "wrapped connection error", err: fmt.Errorf("copy data: %w", syscall.EPIPE), want: true},
+		{name: "other error", err: errors.New("copy failed"), want: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(t, tt.want, IsConnectionClosedError(tt.err))
+		})
+	}
+}
+
+func TestRunContinuesAfterClosedConnectionError(t *testing.T) {
+	t.Parallel()
+
+	listener := newTestListener()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	closedErrCh := make(chan error, 1)
+	unexpectedErrCh := make(chan error, 1)
+	var dialCount atomic.Int32
+
+	p := &Proxy{
+		Listener:   listener,
+		RemoteAddr: "remote:80",
+		DialContext: func(context.Context, string, string) (net.Conn, error) {
+			if dialCount.Add(1) == 1 {
+				return readErrorConn{err: syscall.EPIPE}, nil
+			}
+
+			proxyConn, upstreamConn := net.Pipe()
+			go func() {
+				defer upstreamConn.Close()
+				_, _ = upstreamConn.Write([]byte("ok"))
+			}()
+			return proxyConn, nil
+		},
+		OnError: func(err error) {
+			if IsConnectionClosedError(err) {
+				closedErrCh <- err
+				return
+			}
+			unexpectedErrCh <- err
+		},
+	}
+
+	runErrCh := make(chan error, 1)
+	go func() {
+		runErrCh <- p.Run(ctx)
+	}()
+
+	firstConn := listener.connect()
+	defer firstConn.Close()
+
+	select {
+	case connErr := <-closedErrCh:
+		require.ErrorIs(t, connErr, syscall.EPIPE)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for closed connection error")
+	}
+
+	secondConn := listener.connect()
+	defer secondConn.Close()
+	require.NoError(t, secondConn.SetReadDeadline(time.Now().Add(time.Second)))
+
+	got := make([]byte, 2)
+	_, err := io.ReadFull(secondConn, got)
+	require.NoError(t, err)
+	require.Equal(t, "ok", string(got))
+
+	select {
+	case unexpectedErr := <-unexpectedErrCh:
+		t.Fatalf("unexpected connection error: %v", unexpectedErr)
+	default:
+	}
+
+	cancel()
+	select {
+	case runErr := <-runErrCh:
+		require.NoError(t, runErr)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for proxy to stop")
+	}
+}
+
+func TestRunReturnsListenerError(t *testing.T) {
+	t.Parallel()
+
+	listenerErr := errors.New("listener failed")
+	listener := errorListener{err: listenerErr}
+
+	p := &Proxy{Listener: listener}
+	err := p.Run(context.Background())
+	require.Error(t, err)
+	require.ErrorContains(t, err, "accept local connection")
+	require.ErrorIs(t, err, listenerErr)
+}
+
+type testListener struct {
+	conns     chan net.Conn
+	closed    chan struct{}
+	closeOnce sync.Once
+}
+
+func newTestListener() *testListener {
+	return &testListener{
+		conns:  make(chan net.Conn),
+		closed: make(chan struct{}),
+	}
+}
+
+func (l *testListener) connect() net.Conn {
+	clientConn, proxyConn := net.Pipe()
+	l.conns <- proxyConn
+	return clientConn
+}
+
+func (l *testListener) Accept() (net.Conn, error) {
+	select {
+	case conn := <-l.conns:
+		return conn, nil
+	case <-l.closed:
+		return nil, net.ErrClosed
+	}
+}
+
+func (l *testListener) Close() error {
+	l.closeOnce.Do(func() {
+		close(l.closed)
+	})
+	return nil
+}
+
+func (l *testListener) Addr() net.Addr {
+	return &net.TCPAddr{}
+}
+
+type errorListener struct {
+	err error
+}
+
+func (l errorListener) Accept() (net.Conn, error) { return nil, l.err }
+func (errorListener) Close() error                { return nil }
+func (errorListener) Addr() net.Addr              { return &net.TCPAddr{} }
+
+// readErrorConn fails reads immediately so tests can deterministically exercise a proxy copy failure.
+type readErrorConn struct {
+	err error
+}
+
+func (c readErrorConn) Read([]byte) (int, error)       { return 0, c.err }
+func (readErrorConn) Write(p []byte) (int, error)      { return len(p), nil }
+func (readErrorConn) Close() error                     { return nil }
+func (readErrorConn) LocalAddr() net.Addr              { return &net.TCPAddr{} }
+func (readErrorConn) RemoteAddr() net.Addr             { return &net.TCPAddr{} }
+func (readErrorConn) SetDeadline(time.Time) error      { return nil }
+func (readErrorConn) SetReadDeadline(time.Time) error  { return nil }
+func (readErrorConn) SetWriteDeadline(time.Time) error { return nil }

pkg/client/image.go

@@ -252,13 +252,19 @@ func (cli *Client) pushImageToMachine(
 	// The proxy runs in a goroutine. Capture the first error in a channel
 	// so we can surface it alongside the push error if push fails.
 	proxyErrCh := make(chan error, 1)
-	onProxyError := func(err error) {
+	recordProxyError := func(err error) {
 		select {
 		case proxyErrCh <- fmt.Errorf("proxy to unregistry: %w", err):
 		default:
 		}
 		pw.Event(progress.NewEvent(proxyEventID, progress.Error, err.Error()))
 	}
+	onProxyError := func(err error) {
+		if proxy.IsConnectionClosedError(err) {
+			return
+		}
+		recordProxyError(err)
+	}
 
 	// socketPath is set for plain rootless Docker (not running inside a VM): the Go proxy listens on a unix
 	// socket that is bind-mounted into the socat container, bypassing slirp4netns network routing entirely.
@@ -315,7 +321,11 @@ func (cli *Client) pushImageToMachine(
 	}
 	defer cleanup()
 
-	go unregProxy.Run(proxyCtx)
+	go func() {
+		if err := unregProxy.Run(proxyCtx); err != nil {
+			recordProxyError(err)
+		}
+	}()
 
 	if dockerEnv.Virtualised {
 		// VM-based Docker (Docker Desktop, Rancher Desktop, etc.): run a socat container inside the VM