Wrap errors with context cancellation codes (#659)

This PR wraps errors with the appropriate connect code of Cancelled or
DeadlineExceeded if the context error is not nil.

Improves error handling for some well known error cases that do not
surface context.Cancelled errors. For example HTTP2 "client disconnect"
string errors are now raised with a Cancelled code not an Unknown. This
lets handlers check the error code for better handling and reporting of
errors.

Fix for #645

connect_ext_test.go

@@ -19,12 +19,14 @@ import (
 	"compress/flate"
 	"compress/gzip"
 	"context"
+	"crypto/tls"
 	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
 	"math"
 	"math/rand"
+	"net"
 	"net/http"
 	"runtime"
 	"strings"
@@ -2307,6 +2309,139 @@ func TestStreamUnexpectedEOF(t *testing.T) {
 	}
 }
 
+// TestClientDisconnect tests that the handler receives a CodeCanceled error when
+// the client abruptly disconnects.
+func TestClientDisconnect(t *testing.T) {
+	t.Parallel()
+	type httpRoundTripFunc func(server *memhttp.Server, clientConn *net.Conn, onError chan struct{}) http.RoundTripper
+	http1RoundTripper := func(server *memhttp.Server, clientConn *net.Conn, onError chan struct{}) http.RoundTripper {
+		transport := server.TransportHTTP1()
+		dialContext := transport.DialContext
+		transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
+			conn, err := dialContext(ctx, network, addr)
+			if err != nil {
+				close(onError)
+				return nil, err
+			}
+			*clientConn = conn // Capture the client connection.
+			return conn, nil
+		}
+		return transport
+	}
+	http2RoundTripper := func(server *memhttp.Server, clientConn *net.Conn, onError chan struct{}) http.RoundTripper {
+		transport := server.Transport()
+		dialContext := transport.DialTLSContext
+		transport.DialTLSContext = func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) {
+			conn, err := dialContext(ctx, network, addr, cfg)
+			if err != nil {
+				close(onError)
+				return nil, err
+			}
+			*clientConn = conn // Capture the client connection.
+			return conn, nil
+		}
+		return transport
+	}
+	testTransportClosure := func(t *testing.T, captureTransport httpRoundTripFunc) { //nolint:thelper
+		t.Run("handler_reads", func(t *testing.T) {
+			var (
+				handlerReceiveErr error
+				handlerContextErr error
+				gotRequest        = make(chan struct{})
+				gotResponse       = make(chan struct{})
+			)
+			pingServer := &pluggablePingServer{
+				sum: func(ctx context.Context, stream *connect.ClientStream[pingv1.SumRequest]) (*connect.Response[pingv1.SumResponse], error) {
+					close(gotRequest)
+					for stream.Receive() {
+						// Do nothing
+					}
+					handlerReceiveErr = stream.Err()
+					handlerContextErr = ctx.Err()
+					close(gotResponse)
+					return connect.NewResponse(&pingv1.SumResponse{}), nil
+				},
+			}
+			mux := http.NewServeMux()
+			mux.Handle(pingv1connect.NewPingServiceHandler(pingServer))
+			server := memhttptest.NewServer(t, mux)
+			var clientConn net.Conn
+			transport := captureTransport(server, &clientConn, gotRequest)
+			serverClient := &http.Client{Transport: transport}
+			client := pingv1connect.NewPingServiceClient(serverClient, server.URL())
+			stream := client.Sum(context.Background())
+			// Send header.
+			assert.Nil(t, stream.Send(nil))
+			<-gotRequest
+			// Client abruptly disconnects.
+			if !assert.NotNil(t, clientConn) {
+				return
+			}
+			assert.Nil(t, clientConn.Close())
+			_, err := stream.CloseAndReceive()
+			assert.NotNil(t, err)
+			<-gotResponse
+			assert.NotNil(t, handlerReceiveErr)
+			assert.Equal(t, connect.CodeOf(handlerReceiveErr), connect.CodeCanceled)
+			assert.ErrorIs(t, handlerContextErr, context.Canceled)
+		})
+		t.Run("handler_writes", func(t *testing.T) {
+			var (
+				handlerReceiveErr error
+				handlerContextErr error
+				gotRequest        = make(chan struct{})
+				gotResponse       = make(chan struct{})
+			)
+			pingServer := &pluggablePingServer{
+				countUp: func(ctx context.Context, req *connect.Request[pingv1.CountUpRequest], stream *connect.ServerStream[pingv1.CountUpResponse]) error {
+					close(gotRequest)
+					var err error
+					for err == nil {
+						err = stream.Send(&pingv1.CountUpResponse{})
+					}
+					handlerReceiveErr = err
+					handlerContextErr = ctx.Err()
+					close(gotResponse)
+					return nil
+				},
+			}
+			mux := http.NewServeMux()
+			mux.Handle(pingv1connect.NewPingServiceHandler(pingServer))
+			server := memhttptest.NewServer(t, mux)
+			var clientConn net.Conn
+			transport := captureTransport(server, &clientConn, gotRequest)
+			serverClient := &http.Client{Transport: transport}
+			client := pingv1connect.NewPingServiceClient(serverClient, server.URL())
+			stream, err := client.CountUp(context.Background(), connect.NewRequest(&pingv1.CountUpRequest{}))
+			if !assert.Nil(t, err) {
+				return
+			}
+			<-gotRequest
+			// Client abruptly disconnects.
+			if !assert.NotNil(t, clientConn) {
+				return
+			}
+			assert.Nil(t, clientConn.Close())
+			for stream.Receive() {
+				// Do nothing
+			}
+			assert.NotNil(t, stream.Err())
+			<-gotResponse
+			assert.NotNil(t, handlerReceiveErr)
+			assert.Equal(t, connect.CodeOf(handlerReceiveErr), connect.CodeCanceled)
+			assert.ErrorIs(t, handlerContextErr, context.Canceled)
+		})
+	}
+	t.Run("http1", func(t *testing.T) {
+		t.Parallel()
+		testTransportClosure(t, http1RoundTripper)
+	})
+	t.Run("http2", func(t *testing.T) {
+		t.Parallel()
+		testTransportClosure(t, http2RoundTripper)
+	})
+}
+
 func TestTrailersOnlyErrors(t *testing.T) {
 	t.Parallel()
 

envelope.go

@@ -16,6 +16,7 @@ package connect
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"errors"
 	"io"
@@ -117,6 +118,7 @@ func (e *envelope) Len() int {
 }
 
 type envelopeWriter struct {
+	ctx              context.Context //nolint:containedctx
 	sender           messageSender
 	codec            Codec
 	compressMinBytes int
@@ -208,7 +210,7 @@ func (w *envelopeWriter) marshal(message any) *Error {
 
 func (w *envelopeWriter) write(env *envelope) *Error {
 	if _, err := w.sender.Send(env); err != nil {
-		err = wrapIfContextError(err)
+		err = wrapIfContextDone(w.ctx, err)
 		if connectErr, ok := asError(err); ok {
 			return connectErr
 		}
@@ -218,6 +220,7 @@ func (w *envelopeWriter) write(env *envelope) *Error {
 }
 
 type envelopeReader struct {
+	ctx             context.Context //nolint:containedctx
 	reader          io.Reader
 	codec           Codec
 	last            envelope
@@ -312,15 +315,12 @@ func (r *envelopeReader) Read(env *envelope) *Error {
 			// add any alarming text about protocol errors, though.
 			return NewError(CodeUnknown, err)
 		}
-		err = wrapIfContextError(err)
-		// Something else has gone wrong - the stream didn't end cleanly.
+		err = wrapIfMaxBytesError(err, "read 5 byte message prefix")
+		err = wrapIfContextDone(r.ctx, err)
 		if connectErr, ok := asError(err); ok {
 			return connectErr
 		}
-		if maxBytesErr := asMaxBytesError(err, "read 5 byte message prefix"); maxBytesErr != nil {
-			// We're reading from an http.MaxBytesHandler, and we've exceeded the read limit.
-			return maxBytesErr
-		}
+		// Something else has gone wrong - the stream didn't end cleanly.
 		return errorf(
 			CodeInvalidArgument,
 			"protocol error: incomplete envelope: %w", err,
@@ -338,10 +338,6 @@ func (r *envelopeReader) Read(env *envelope) *Error {
 	// CopyN will return an error if it doesn't read the requested
 	// number of bytes.
 	if readN, err := io.CopyN(env.Data, r.reader, size); err != nil {
-		if maxBytesErr := asMaxBytesError(err, "read %d byte message", size); maxBytesErr != nil {
-			// We're reading from an http.MaxBytesHandler, and we've exceeded the read limit.
-			return maxBytesErr
-		}
 		if errors.Is(err, io.EOF) {
 			// We've gotten fewer bytes than we expected, so the stream has ended
 			// unexpectedly.
@@ -352,7 +348,8 @@ func (r *envelopeReader) Read(env *envelope) *Error {
 				readN,
 			)
 		}
-		err = wrapIfContextError(err)
+		err = wrapIfMaxBytesError(err, "read %d byte message", size)
+		err = wrapIfContextDone(r.ctx, err)
 		if connectErr, ok := asError(err); ok {
 			return connectErr
 		}

envelope_test.go

@@ -16,6 +16,7 @@ package connect
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"testing"
 
@@ -44,6 +45,7 @@ func TestEnvelope(t *testing.T) {
 			t.Parallel()
 			env := &envelope{Data: &bytes.Buffer{}}
 			rdr := envelopeReader{
+				ctx: context.Background(),
 				reader: byteByByteReader{
 					reader: bytes.NewReader(buf.Bytes()),
 				},

error.go

@@ -308,6 +308,25 @@ func wrapIfContextError(err error) error {
 	return err
 }
 
+// wrapIfContextDone wraps errors with CodeCanceled or CodeDeadlineExceeded
+// if the context is done. It leaves already-wrapped errors unchanged.
+func wrapIfContextDone(ctx context.Context, err error) error {
+	if err == nil {
+		return nil
+	}
+	err = wrapIfContextError(err)
+	if _, ok := asError(err); ok {
+		return err
+	}
+	ctxErr := ctx.Err()
+	if errors.Is(ctxErr, context.Canceled) {
+		return NewError(CodeCanceled, err)
+	} else if errors.Is(ctxErr, context.DeadlineExceeded) {
+		return NewError(CodeDeadlineExceeded, err)
+	}
+	return err
+}
+
 // wrapIfLikelyH2CNotConfiguredError adds a wrapping error that has a message
 // telling the caller that they likely need to use h2c but are using a raw http.Client{}.
 //
@@ -414,10 +433,18 @@ func wrapIfRSTError(err error) error {
 	}
 }
 
-func asMaxBytesError(err error, tmpl string, args ...any) *Error {
+// wrapIfMaxBytesError wraps errors returned reading from a http.MaxBytesHandler
+// whose limit has been exceeded.
+func wrapIfMaxBytesError(err error, tmpl string, args ...any) error {
+	if err == nil {
+		return nil
+	}
+	if _, ok := asError(err); ok {
+		return err
+	}
 	var maxBytesErr *http.MaxBytesError
 	if ok := errors.As(err, &maxBytesErr); !ok {
-		return nil
+		return err
 	}
 	prefix := fmt.Sprintf(tmpl, args...)
 	return errorf(CodeResourceExhausted, "%s: exceeded %d byte http.MaxBytesReader limit", prefix, maxBytesErr.Limit)

protocol_connect.go

@@ -148,6 +148,7 @@ func (h *connectHandler) NewConn(
 	responseWriter http.ResponseWriter,
 	request *http.Request,
 ) (handlerConnCloser, bool) {
+	ctx := request.Context()
 	query := request.URL.Query()
 	// We need to parse metadata before entering the interceptor stack; we'll
 	// send the error to the client later on.
@@ -255,6 +256,7 @@ func (h *connectHandler) NewConn(
 			request:        request,
 			responseWriter: responseWriter,
 			marshaler: connectUnaryMarshaler{
+				ctx:              ctx,
 				sender:           writeSender{writer: responseWriter},
 				codec:            codec,
 				compressMinBytes: h.CompressMinBytes,
@@ -265,6 +267,7 @@ func (h *connectHandler) NewConn(
 				sendMaxBytes:     h.SendMaxBytes,
 			},
 			unmarshaler: connectUnaryUnmarshaler{
+				ctx:             ctx,
 				reader:          requestBody,
 				codec:           codec,
 				compressionPool: h.CompressionPools.Get(requestCompression),
@@ -281,6 +284,7 @@ func (h *connectHandler) NewConn(
 			responseWriter: responseWriter,
 			marshaler: connectStreamingMarshaler{
 				envelopeWriter: envelopeWriter{
+					ctx:              ctx,
 					sender:           writeSender{responseWriter},
 					codec:            codec,
 					compressMinBytes: h.CompressMinBytes,
@@ -291,6 +295,7 @@ func (h *connectHandler) NewConn(
 			},
 			unmarshaler: connectStreamingUnmarshaler{
 				envelopeReader: envelopeReader{
+					ctx:             ctx,
 					reader:          requestBody,
 					codec:           codec,
 					compressionPool: h.CompressionPools.Get(requestCompression),
@@ -376,6 +381,7 @@ func (c *connectClient) NewConn(
 			bufferPool:       c.BufferPool,
 			marshaler: connectUnaryRequestMarshaler{
 				connectUnaryMarshaler: connectUnaryMarshaler{
+					ctx:              ctx,
 					sender:           duplexCall,
 					codec:            c.Codec,
 					compressMinBytes: c.CompressMinBytes,
@@ -387,6 +393,7 @@ func (c *connectClient) NewConn(
 				},
 			},
 			unmarshaler: connectUnaryUnmarshaler{
+				ctx:          ctx,
 				reader:       duplexCall,
 				codec:        c.Codec,
 				bufferPool:   c.BufferPool,
@@ -416,6 +423,7 @@ func (c *connectClient) NewConn(
 			codec:            c.Codec,
 			marshaler: connectStreamingMarshaler{
 				envelopeWriter: envelopeWriter{
+					ctx:              ctx,
 					sender:           duplexCall,
 					codec:            c.Codec,
 					compressMinBytes: c.CompressMinBytes,
@@ -426,6 +434,7 @@ func (c *connectClient) NewConn(
 			},
 			unmarshaler: connectStreamingUnmarshaler{
 				envelopeReader: envelopeReader{
+					ctx:          ctx,
 					reader:       duplexCall,
 					codec:        c.Codec,
 					bufferPool:   c.BufferPool,
@@ -912,6 +921,7 @@ func (u *connectStreamingUnmarshaler) EndStreamError() *Error {
 }
 
 type connectUnaryMarshaler struct {
+	ctx              context.Context //nolint:containedctx
 	sender           messageSender
 	codec            Codec
 	compressMinBytes int
@@ -1077,6 +1087,7 @@ func (m *connectUnaryRequestMarshaler) writeWithGet(url *url.URL) *Error {
 }
 
 type connectUnaryUnmarshaler struct {
+	ctx             context.Context //nolint:containedctx
 	reader          io.Reader
 	codec           Codec
 	compressionPool *compressionPool
@@ -1103,13 +1114,11 @@ func (u *connectUnaryUnmarshaler) UnmarshalFunc(message any, unmarshal func([]by
 	// ReadFrom ignores io.EOF, so any error here is real.
 	bytesRead, err := data.ReadFrom(reader)
 	if err != nil {
-		err = wrapIfContextError(err)
+		err = wrapIfMaxBytesError(err, "read first %d bytes of message", bytesRead)
+		err = wrapIfContextDone(u.ctx, err)
 		if connectErr, ok := asError(err); ok {
 			return connectErr
 		}
-		if readMaxBytesErr := asMaxBytesError(err, "read first %d bytes of message", bytesRead); readMaxBytesErr != nil {
-			return readMaxBytesErr
-		}
 		return errorf(CodeUnknown, "read message: %w", err)
 	}
 	if u.readMaxBytes > 0 && bytesRead > int64(u.readMaxBytes) {