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

2026-08-26 19:13:34 +00:00 · 2026-07-21 14:16:46 +10:00
parent b99abb7c39
commit fa77edf53e
4 changed files with 276 additions and 66 deletions
@@ -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))
+	}
 }
@@ -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 }