From 7457031ae601ac4a3c09b0e95fc710c52ab1e885 Mon Sep 17 00:00:00 2001 From: jingyu Date: Sat, 7 Feb 2026 13:17:09 +0800 Subject: [PATCH] Use two-stage fallback for oversized EndStream control frames Instead of unconditionally bypassing the sendMaxBytes check for control frames, implement a graduated fallback strategy in envelopeWriter: 1. Try to send the full EndStream frame (error code + message + details). 2. If it exceeds sendMaxBytes, call a reduce closure to strip error details in-place, re-marshal via a marshal closure, and retry. 3. If even the minimal frame is too large (unreasonably low limit), force-send it anyway. The fallback logic lives entirely in envelope.go's writeControlFrame, which takes two closures from the protocol layer: - marshal: serializes the current frame state into bytes. Because it closes over the frame struct, it automatically reflects any mutations made by reduce. - reduce: mutates the captured frame struct to strip non-essential data (error details for Connect, grpc-status-details-bin for gRPC-Web). This keeps sendMaxBytes enforcement intact for control frames while guaranteeing that the client always receives at least a minimal error. Fixes #907 Signed-off-by: jingyu --- connect_ext_test.go | 64 +++++++++++++++++++++++++++++++++++++++++++ envelope.go | 58 ++++++++++++++++++++++++++++++++++----- envelope_test.go | 62 +++++++++++++++++++++++++++++++++++++++++ protocol_connect.go | 22 ++++++++------- protocol_grpc.go | 23 ++++++++++------ protocol_grpc_test.go | 43 +++++++++++++++++++++++++++++ 6 files changed, 247 insertions(+), 25 deletions(-) diff --git a/connect_ext_test.go b/connect_ext_test.go index 8612865d..f8597819 100644 --- a/connect_ext_test.go +++ b/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() diff --git a/envelope.go b/envelope.go index 21a897de..e853ab3a 100644 --- a/envelope.go +++ b/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,7 +176,7 @@ 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, @@ -178,6 +184,44 @@ func (w *envelopeWriter) Write(env *envelope) *Error { }) } +// 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() diff --git a/envelope_test.go b/envelope_test.go index dfecc797..2d64763e 100644 --- a/envelope_test.go +++ b/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 diff --git a/protocol_connect.go b/protocol_connect.go index 371f382e..a09a7464 100644 --- a/protocol_connect.go +++ b/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 { diff --git a/protocol_grpc.go b/protocol_grpc.go index 4c0e8125..4f4c41ec 100644 --- a/protocol_grpc.go +++ b/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 { diff --git a/protocol_grpc_test.go b/protocol_grpc_test.go index 1e16b73f..09d73f1f 100644 --- a/protocol_grpc_test.go +++ b/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"