Skip sendMaxBytes check for EndStream control frames (#907)

connect_ext_test.go

@@ -1814,6 +1814,70 @@ func TestHandlerWithSendMaxBytes(t *testing.T) {
 	})
 }
 
+func TestHandlerWithSendMaxBytesStreamEndStream(t *testing.T) {
+	// Regression test: when a streaming handler returns an error because a
+	// message exceeds sendMaxBytes, the protocol control frame carrying the
+	// error (Connect EndStream or gRPC-Web trailer frame) must not itself be
+	// rejected by the sendMaxBytes check.
+	t.Parallel()
+	const sendMaxBytes = 1 // Tiny limit so even a minimal message exceeds it.
+	testCases := []struct {
+		name             string
+		compressMinBytes int
+		clientOptions    []connect.ClientOption
+	}{
+		{
+			name:             "connect_uncompressed",
+			compressMinBytes: math.MaxInt,
+		},
+		{
+			name:             "connect_compressed",
+			compressMinBytes: 1,
+		},
+		{
+			name:             "grpcweb_uncompressed",
+			compressMinBytes: math.MaxInt,
+			clientOptions:    []connect.ClientOption{connect.WithGRPCWeb()},
+		},
+		{
+			name:             "grpcweb_compressed",
+			compressMinBytes: 1,
+			clientOptions:    []connect.ClientOption{connect.WithGRPCWeb()},
+		},
+	}
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			t.Parallel()
+			mux := http.NewServeMux()
+			mux.Handle(pingv1connect.NewPingServiceHandler(
+				&pluggablePingServer{
+					countUp: func(ctx context.Context, req *connect.Request[pingv1.CountUpRequest], stream *connect.ServerStream[pingv1.CountUpResponse]) error {
+						// Send a message that exceeds sendMaxBytes; this will fail
+						// with CodeResourceExhausted. The handler returns the error,
+						// which should be delivered to the client via the EndStream frame.
+						err := stream.Send(&pingv1.CountUpResponse{Number: 42})
+						if err != nil {
+							return err
+						}
+						return nil
+					},
+				},
+				connect.WithSendMaxBytes(sendMaxBytes),
+				connect.WithCompressMinBytes(testCase.compressMinBytes),
+			))
+			server := memhttptest.NewServer(t, mux)
+			client := pingv1connect.NewPingServiceClient(server.Client(), server.URL(), testCase.clientOptions...)
+			stream, err := client.CountUp(t.Context(), connect.NewRequest(&pingv1.CountUpRequest{Number: 1}))
+			assert.Nil(t, err)
+			// Drain the stream — we expect no messages and an error at the end.
+			assert.False(t, stream.Receive())
+			assert.NotNil(t, stream.Err(), assert.Sprintf("expected error from stream"))
+			assert.Equal(t, connect.CodeOf(stream.Err()), connect.CodeResourceExhausted)
+			assert.Nil(t, stream.Close())
+		})
+	}
+}
+
 func TestClientWithSendMaxBytes(t *testing.T) {
 	t.Parallel()
 	mux := http.NewServeMux()

envelope.go

@@ -28,11 +28,17 @@ import (
 // same meaning in the gRPC-Web, gRPC-HTTP2, and Connect protocols.
 const flagEnvelopeCompressed = 0b00000001
 
-var errSpecialEnvelope = errorf(
-	CodeUnknown,
-	"final message has protocol-specific flags: %w",
-	// User code checks for end of stream with errors.Is(err, io.EOF).
-	io.EOF,
+var (
+	errSpecialEnvelope = errorf(
+		CodeUnknown,
+		"final message has protocol-specific flags: %w",
+		// User code checks for end of stream with errors.Is(err, io.EOF).
+		io.EOF,
+	)
+	// errExceedsSendMax is a sentinel wrapped by sendMaxBytes errors so
+	// that callers (e.g. MarshalEndStream) can detect the violation and
+	// fall back to a smaller frame.
+	errExceedsSendMax = errors.New("")
 )
 
 // envelope is a block of arbitrary bytes wrapped in gRPC and Connect's framing
@@ -160,7 +166,7 @@ func (w *envelopeWriter) Write(env *envelope) *Error {
 		w.compressionPool == nil ||
 		env.Data.Len() < w.compressMinBytes {
 		if w.sendMaxBytes > 0 && env.Data.Len() > w.sendMaxBytes {
-			return errorf(CodeResourceExhausted, "message size %d exceeds sendMaxBytes %d", env.Data.Len(), w.sendMaxBytes)
+			return errorf(CodeResourceExhausted, "message size %d exceeds sendMaxBytes %d%w", env.Data.Len(), w.sendMaxBytes, errExceedsSendMax)
 		}
 		return w.write(env)
 	}
@@ -170,14 +176,52 @@ func (w *envelopeWriter) Write(env *envelope) *Error {
 		return err
 	}
 	if w.sendMaxBytes > 0 && data.Len() > w.sendMaxBytes {
-		return errorf(CodeResourceExhausted, "compressed message size %d exceeds sendMaxBytes %d", data.Len(), w.sendMaxBytes)
+		return errorf(CodeResourceExhausted, "compressed message size %d exceeds sendMaxBytes %d%w", data.Len(), w.sendMaxBytes, errExceedsSendMax)
 	}
 	return w.write(&envelope{
 		Data:  data,
 		Flags: env.Flags | flagEnvelopeCompressed,
 	})
 }
 
+// writeControlFrame implements a fallback strategy for writing protocol
+// control frames (Connect EndStream, gRPC-Web trailers) that may exceed
+// sendMaxBytes. It calls marshal to produce the frame bytes and attempts to
+// write the frame. If the frame exceeds sendMaxBytes:
+//  1. Calls reduce (if non-nil) to strip non-essential data (e.g. error
+//     details), re-marshals, and retries.
+//  2. If still too large, sends the frame anyway — the limit is unreasonably
+//     low and delivering the error is more important.
+func (w *envelopeWriter) writeControlFrame(flags uint8, marshal func() ([]byte, *Error), reduce func()) *Error {
+	write := func() *Error {
+		data, err := marshal()
+		if err != nil {
+			return err
+		}
+		raw := bytes.NewBuffer(data)
+		defer w.bufferPool.Put(raw)
+		return w.Write(&envelope{Data: raw, Flags: flags})
+	}
+	if writeErr := write(); writeErr == nil || !errors.Is(writeErr, errExceedsSendMax) {
+		return writeErr
+	}
+	if reduce != nil {
+		reduce()
+		if writeErr := write(); writeErr == nil || !errors.Is(writeErr, errExceedsSendMax) {
+			return writeErr
+		}
+	}
+	// Even the (possibly reduced) frame exceeds sendMaxBytes. The limit is
+	// unreasonably low, so send it anyway.
+	data, err := marshal()
+	if err != nil {
+		return err
+	}
+	raw := bytes.NewBuffer(data)
+	defer w.bufferPool.Put(raw)
+	return w.write(&envelope{Data: raw, Flags: flags})
+}
+
 func (w *envelopeWriter) marshalAppend(message any, codec marshalAppender) *Error {
 	// Codec supports MarshalAppend; try to re-use a []byte from the pool.
 	buffer := w.bufferPool.Get()

envelope_test.go

@@ -16,6 +16,7 @@ package connect
 
 import (
 	"bytes"
+	"errors"
 	"io"
 	"testing"
 
@@ -102,6 +103,67 @@ func TestEnvelope(t *testing.T) {
 	})
 }
 
+func TestEnvelopeWriteSendMaxBytes(t *testing.T) {
+	t.Parallel()
+	testCases := []struct {
+		name       string
+		flags      uint8
+		compressed bool
+	}{
+		{
+			name: "rejects_oversized_message_uncompressed",
+		},
+		{
+			name:       "rejects_oversized_message_compressed",
+			compressed: true,
+		},
+		{
+			name:  "rejects_oversized_end_stream_uncompressed",
+			flags: connectFlagEnvelopeEndStream,
+		},
+		{
+			name:       "rejects_oversized_end_stream_compressed",
+			flags:      connectFlagEnvelopeEndStream,
+			compressed: true,
+		},
+		{
+			name:  "rejects_oversized_grpc_web_trailer_uncompressed",
+			flags: grpcFlagEnvelopeTrailer,
+		},
+		{
+			name:       "rejects_oversized_grpc_web_trailer_compressed",
+			flags:      grpcFlagEnvelopeTrailer,
+			compressed: true,
+		},
+	}
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			t.Parallel()
+			out := &bytes.Buffer{}
+			writer := envelopeWriter{
+				sender:       writeSender{writer: out},
+				bufferPool:   newBufferPool(),
+				sendMaxBytes: 1,
+			}
+			if testCase.compressed {
+				gzipOption, ok := withGzip().(*compressionOption)
+				assert.True(t, ok)
+				writer.compressMinBytes = 1
+				writer.compressionPool = gzipOption.CompressionPool
+				writer.bufferPool = newBufferPool()
+			}
+			env := &envelope{
+				Data:  bytes.NewBuffer(make([]byte, 64)),
+				Flags: testCase.flags,
+			}
+			err := writer.Write(env)
+			assert.NotNil(t, err)
+			assert.Equal(t, CodeOf(err), CodeResourceExhausted)
+			assert.True(t, errors.Is(err, errExceedsSendMax))
+		})
+	}
+}
+
 // byteByByteReader is test reader that reads a single byte at a time.
 type byteByByteReader struct {
 	reader io.ByteReader

protocol_connect.go

@@ -855,16 +855,18 @@ func (m *connectStreamingMarshaler) MarshalEndStream(err error, trailer http.Hea
 			mergeNonProtocolHeaders(end.Trailer, connectErr.meta)
 		}
 	}
-	data, marshalErr := json.Marshal(end)
-	if marshalErr != nil {
-		return errorf(CodeInternal, "marshal end stream: %w", marshalErr)
-	}
-	raw := bytes.NewBuffer(data)
-	defer m.envelopeWriter.bufferPool.Put(raw)
-	return m.Write(&envelope{
-		Data:  raw,
-		Flags: connectFlagEnvelopeEndStream,
-	})
+	marshal := func() ([]byte, *Error) {
+		data, marshalErr := json.Marshal(end)
+		if marshalErr != nil {
+			return nil, errorf(CodeInternal, "marshal end stream: %w", marshalErr)
+		}
+		return data, nil
+	}
+	var reduce func()
+	if end.Error != nil && len(end.Error.Details) > 0 {
+		reduce = func() { end.Error.Details = nil }
+	}
+	return m.writeControlFrame(connectFlagEnvelopeEndStream, marshal, reduce)
 }
 
 type connectStreamingUnmarshaler struct {

protocol_grpc.go

@@ -16,6 +16,7 @@ package connect
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -577,8 +578,6 @@ type grpcMarshaler struct {
 }
 
 func (m *grpcMarshaler) MarshalWebTrailers(trailer http.Header) *Error {
-	raw := m.envelopeWriter.bufferPool.Get()
-	defer m.envelopeWriter.bufferPool.Put(raw)
 	for key, values := range trailer {
 		// Per the Go specification, keys inserted during iteration may be produced
 		// later in the iteration or may be skipped. For safety, avoid mutating the
@@ -590,13 +589,21 @@ func (m *grpcMarshaler) MarshalWebTrailers(trailer http.Header) *Error {
 		delete(trailer, key)
 		trailer[lower] = values
 	}
-	if err := trailer.Write(raw); err != nil {
-		return errorf(CodeInternal, "format trailers: %w", err)
+	marshal := func() ([]byte, *Error) {
+		var buf bytes.Buffer
+		if err := trailer.Write(&buf); err != nil {
+			return nil, errorf(CodeInternal, "format trailers: %w", err)
+		}
+		return buf.Bytes(), nil
 	}
-	return m.Write(&envelope{
-		Data:  raw,
-		Flags: grpcFlagEnvelopeTrailer,
-	})
+	// Use direct map access because the loop above lowercased all keys,
+	// so http.Header.Get/Del (which canonicalize) won't find them.
+	detailsKey := strings.ToLower(grpcHeaderDetails)
+	var reduce func()
+	if _, ok := trailer[detailsKey]; ok {
+		reduce = func() { delete(trailer, detailsKey) }
+	}
+	return m.writeControlFrame(grpcFlagEnvelopeTrailer, marshal, reduce)
 }
 
 type grpcUnmarshaler struct {

protocol_grpc_test.go

@@ -197,6 +197,49 @@ func TestGRPCWebTrailerMarshalling(t *testing.T) {
 	assert.Equal(t, marshalled, "grpc-message: Foo\r\ngrpc-status: 0\r\nuser-provided: bar\r\n")
 }
 
+func TestGRPCWebTrailerMarshallingSendMaxBytesControlFrame(t *testing.T) {
+	t.Parallel()
+	testCases := []struct {
+		name       string
+		compressed bool
+	}{
+		{
+			name: "uncompressed",
+		},
+		{
+			name:       "compressed",
+			compressed: true,
+		},
+	}
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			t.Parallel()
+			responseWriter := httptest.NewRecorder()
+			marshaler := grpcMarshaler{
+				envelopeWriter: envelopeWriter{
+					sender:       writeSender{writer: responseWriter},
+					bufferPool:   newBufferPool(),
+					sendMaxBytes: 1,
+				},
+			}
+			if testCase.compressed {
+				gzipOption, ok := withGzip().(*compressionOption)
+				assert.True(t, ok)
+				marshaler.compressMinBytes = 1
+				marshaler.compressionPool = gzipOption.CompressionPool
+			}
+			trailer := http.Header{}
+			trailer.Add("grpc-status", "13")
+			trailer.Add("grpc-message", strings.Repeat("x", 64))
+			err := marshaler.MarshalWebTrailers(trailer)
+			assert.Nil(t, err)
+			frame := responseWriter.Body.Bytes()
+			assert.True(t, len(frame) > 5, assert.Sprintf("expected trailer frame in response body"))
+			assert.True(t, frame[0]&grpcFlagEnvelopeTrailer == grpcFlagEnvelopeTrailer)
+		})
+	}
+}
+
 func BenchmarkGRPCPercentEncoding(b *testing.B) {
 	input := "Hello, 世界"
 	want := "Hello, %E4%B8%96%E7%95%8C"