Wrap errors with context cancellation codes

.golangci.yml

@@ -132,3 +132,9 @@ issues:
     # We want to show examples with http.Get
     - linters: [noctx]
       path: internal/memhttp/memhttp_test.go
+    # We need envelope readers and writers to have access to the context for error handling.
+    - linters: [containedctx]
+      path: envelope.go
+    # We need marshallers and unmarshallers to have access to the context for error handling.
+    - linters: [containedctx]
+      path: protocol_connect.go

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"
@@ -2276,6 +2278,132 @@ 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()
+	captureTransportConn := func(server *memhttp.Server, clientConn *net.Conn, onError chan struct{}, http2 bool) http.RoundTripper {
+		if http2 {
+			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
+		}
+		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
+	}
+	testTransportClosure := func(t *testing.T, http2 bool) { //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 := captureTransportConn(server, &clientConn, gotRequest, http2)
+			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 := captureTransportConn(server, &clientConn, gotRequest, http2)
+			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)
+		})
+	}
+	testTransportClosure(t, true)
+	testTransportClosure(t, false)
+}
+
 // TestBlankImportCodeGeneration tests that services.connect.go is generated with
 // blank import statements to services.pb.go so that the service's Descriptor is
 // available in the global proto registry.

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
 	sender           messageSender
 	codec            Codec
 	compressMinBytes int
@@ -209,6 +211,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 = wrapWithContextError(w.ctx, err)
 		if connectErr, ok := asError(err); ok {
 			return connectErr
 		}
@@ -218,6 +221,7 @@ func (w *envelopeWriter) write(env *envelope) *Error {
 }
 
 type envelopeReader struct {
+	ctx             context.Context
 	reader          io.Reader
 	codec           Codec
 	last            envelope
@@ -305,17 +309,12 @@ func (r *envelopeReader) Read(env *envelope) *Error {
 			return NewError(CodeUnknown, err)
 		}
 		err = wrapIfContextError(err)
+		err = wrapWithContextError(r.ctx, err)
+		err = wrapIfMaxBytesError(err, "read 5 byte message prefix")
 		if connectErr, ok := asError(err); ok {
 			return connectErr
 		}
 		// Something else has gone wrong - the stream didn't end cleanly.
-		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
-		}
 		return errorf(
 			CodeInvalidArgument,
 			"protocol error: incomplete envelope: %w", err,
@@ -333,10 +332,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.
@@ -348,6 +343,8 @@ func (r *envelopeReader) Read(env *envelope) *Error {
 			)
 		}
 		err = wrapIfContextError(err)
+		err = wrapWithContextError(r.ctx, err)
+		err = wrapIfMaxBytesError(err, "read %d byte message", size)
 		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

@@ -302,6 +302,26 @@ func wrapIfContextError(err error) error {
 	return err
 }
 
+// wrapWithContextError wraps errors with CodeCanceled or CodeDeadlineExceeded
+// if the context is done. It leaves already-wrapped errors unchanged.
+func wrapWithContextError(ctx context.Context, err error) error {
+	if err == nil {
+		return nil
+	}
+	if _, ok := asError(err); ok {
+		return err
+	}
+	ctxErr := ctx.Err()
+	switch {
+	case errors.Is(ctxErr, context.Canceled):
+		return NewError(CodeCanceled, err)
+	case errors.Is(ctxErr, context.DeadlineExceeded):
+		return NewError(CodeDeadlineExceeded, err)
+	default:
+		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{}.
 //
@@ -408,10 +428,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

@@ -147,6 +147,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.
@@ -254,6 +255,7 @@ func (h *connectHandler) NewConn(
 			request:        request,
 			responseWriter: responseWriter,
 			marshaler: connectUnaryMarshaler{
+				ctx:              ctx,
 				sender:           writeSender{writer: responseWriter},
 				codec:            codec,
 				compressMinBytes: h.CompressMinBytes,
@@ -264,6 +266,7 @@ func (h *connectHandler) NewConn(
 				sendMaxBytes:     h.SendMaxBytes,
 			},
 			unmarshaler: connectUnaryUnmarshaler{
+				ctx:             ctx,
 				reader:          requestBody,
 				codec:           codec,
 				compressionPool: h.CompressionPools.Get(requestCompression),
@@ -280,6 +283,7 @@ func (h *connectHandler) NewConn(
 			responseWriter: responseWriter,
 			marshaler: connectStreamingMarshaler{
 				envelopeWriter: envelopeWriter{
+					ctx:              ctx,
 					sender:           writeSender{responseWriter},
 					codec:            codec,
 					compressMinBytes: h.CompressMinBytes,
@@ -290,6 +294,7 @@ func (h *connectHandler) NewConn(
 			},
 			unmarshaler: connectStreamingUnmarshaler{
 				envelopeReader: envelopeReader{
+					ctx:             ctx,
 					reader:          requestBody,
 					codec:           codec,
 					compressionPool: h.CompressionPools.Get(requestCompression),
@@ -375,6 +380,7 @@ func (c *connectClient) NewConn(
 			bufferPool:       c.BufferPool,
 			marshaler: connectUnaryRequestMarshaler{
 				connectUnaryMarshaler: connectUnaryMarshaler{
+					ctx:              ctx,
 					sender:           duplexCall,
 					codec:            c.Codec,
 					compressMinBytes: c.CompressMinBytes,
@@ -386,6 +392,7 @@ func (c *connectClient) NewConn(
 				},
 			},
 			unmarshaler: connectUnaryUnmarshaler{
+				ctx:          ctx,
 				reader:       duplexCall,
 				codec:        c.Codec,
 				bufferPool:   c.BufferPool,
@@ -415,6 +422,7 @@ func (c *connectClient) NewConn(
 			codec:            c.Codec,
 			marshaler: connectStreamingMarshaler{
 				envelopeWriter: envelopeWriter{
+					ctx:              ctx,
 					sender:           duplexCall,
 					codec:            c.Codec,
 					compressMinBytes: c.CompressMinBytes,
@@ -425,6 +433,7 @@ func (c *connectClient) NewConn(
 			},
 			unmarshaler: connectStreamingUnmarshaler{
 				envelopeReader: envelopeReader{
+					ctx:          ctx,
 					reader:       duplexCall,
 					codec:        c.Codec,
 					bufferPool:   c.BufferPool,
@@ -892,6 +901,7 @@ func (u *connectStreamingUnmarshaler) EndStreamError() *Error {
 }
 
 type connectUnaryMarshaler struct {
+	ctx              context.Context
 	sender           messageSender
 	codec            Codec
 	compressMinBytes int
@@ -1057,6 +1067,7 @@ func (m *connectUnaryRequestMarshaler) writeWithGet(url *url.URL) *Error {
 }
 
 type connectUnaryUnmarshaler struct {
+	ctx             context.Context
 	reader          io.Reader
 	codec           Codec
 	compressionPool *compressionPool
@@ -1084,12 +1095,11 @@ func (u *connectUnaryUnmarshaler) UnmarshalFunc(message any, unmarshal func([]by
 	bytesRead, err := data.ReadFrom(reader)
 	if err != nil {
 		err = wrapIfContextError(err)
+		err = wrapWithContextError(u.ctx, err)
+		err = wrapIfMaxBytesError(err, "read first %d bytes of message", bytesRead)
 		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) {