transport: http2 server must validates header list size when early aborting stream

internal/transport/controlbuf.go

@@ -24,19 +24,13 @@ import (
 	"fmt"
 	"net"
 	"runtime"
-	"strconv"
 	"sync"
 	"sync/atomic"
 
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/hpack"
 	"google.golang.org/grpc/internal/grpclog"
-	"google.golang.org/grpc/internal/grpcutil"
-	"google.golang.org/grpc/internal/pretty"
-	istatus "google.golang.org/grpc/internal/status"
 	"google.golang.org/grpc/mem"
-	"google.golang.org/grpc/status"
-	"google.golang.org/protobuf/proto"
 )
 
 var updateHeaderTblSize = func(e *hpack.Encoder, v uint32) {
@@ -150,11 +144,9 @@ type cleanupStream struct {
 func (c *cleanupStream) isTransportResponseFrame() bool { return c.rst } // Results in a RST_STREAM
 
 type earlyAbortStream struct {
-	httpStatus     uint32
-	streamID       uint32
-	contentSubtype string
-	status         *status.Status
-	rst            bool
+	streamID uint32
+	rst      bool
+	hf       []hpack.HeaderField // Pre-built header fields
 }
 
 func (*earlyAbortStream) isTransportResponseFrame() bool { return false }
@@ -846,27 +838,7 @@ func (l *loopyWriter) earlyAbortStreamHandler(eas *earlyAbortStream) error {
 	if l.side == clientSide {
 		return errors.New("earlyAbortStream not handled on client")
 	}
-	// In case the caller forgets to set the http status, default to 200.
-	if eas.httpStatus == 0 {
-		eas.httpStatus = 200
-	}
-	headerFields := []hpack.HeaderField{
-		{Name: ":status", Value: strconv.Itoa(int(eas.httpStatus))},
-		{Name: "content-type", Value: grpcutil.ContentType(eas.contentSubtype)},
-		{Name: "grpc-status", Value: strconv.Itoa(int(eas.status.Code()))},
-		{Name: "grpc-message", Value: encodeGrpcMessage(eas.status.Message())},
-	}
-
-	if p := istatus.RawStatusProto(eas.status); len(p.GetDetails()) > 0 {
-		stBytes, err := proto.Marshal(p)
-		if err != nil {
-			l.logger.Errorf("Failed to marshal rpc status: %s, error: %v", pretty.ToJSON(p), err)
-		} else {
-			headerFields = append(headerFields, hpack.HeaderField{Name: grpcStatusDetailsBinHeader, Value: encodeBinHeader(stBytes)})
-		}
-	}
-
-	if err := l.writeHeader(eas.streamID, true, headerFields, nil); err != nil {
+	if err := l.writeHeader(eas.streamID, true, eas.hf, nil); err != nil {
 		return err
 	}
 	if eas.rst {

internal/transport/http2_server.go

@@ -479,13 +479,7 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
 		if t.logger.V(logLevel) {
 			t.logger.Infof("Aborting the stream early: %v", errMsg)
 		}
-		t.controlBuf.put(&earlyAbortStream{
-			httpStatus:     http.StatusBadRequest,
-			streamID:       streamID,
-			contentSubtype: s.contentSubtype,
-			status:         status.New(codes.Internal, errMsg),
-			rst:            !frame.StreamEnded(),
-		})
+		t.writeEarlyAbort(streamID, s.contentSubtype, status.New(codes.Internal, errMsg), http.StatusBadRequest, !frame.StreamEnded())
 		return nil
 	}
 
@@ -499,23 +493,11 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
 		return nil
 	}
 	if !isGRPC {
-		t.controlBuf.put(&earlyAbortStream{
-			httpStatus:     http.StatusUnsupportedMediaType,
-			streamID:       streamID,
-			contentSubtype: s.contentSubtype,
-			status:         status.Newf(codes.InvalidArgument, "invalid gRPC request content-type %q", contentType),
-			rst:            !frame.StreamEnded(),
-		})
+		t.writeEarlyAbort(streamID, s.contentSubtype, status.Newf(codes.InvalidArgument, "invalid gRPC request content-type %q", contentType), http.StatusUnsupportedMediaType, !frame.StreamEnded())
 		return nil
 	}
 	if headerError != nil {
-		t.controlBuf.put(&earlyAbortStream{
-			httpStatus:     http.StatusBadRequest,
-			streamID:       streamID,
-			contentSubtype: s.contentSubtype,
-			status:         headerError,
-			rst:            !frame.StreamEnded(),
-		})
+		t.writeEarlyAbort(streamID, s.contentSubtype, headerError, http.StatusBadRequest, !frame.StreamEnded())
 		return nil
 	}
 
@@ -569,13 +551,7 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
 		if t.logger.V(logLevel) {
 			t.logger.Infof("Aborting the stream early: %v", errMsg)
 		}
-		t.controlBuf.put(&earlyAbortStream{
-			httpStatus:     http.StatusMethodNotAllowed,
-			streamID:       streamID,
-			contentSubtype: s.contentSubtype,
-			status:         status.New(codes.Internal, errMsg),
-			rst:            !frame.StreamEnded(),
-		})
+		t.writeEarlyAbort(streamID, s.contentSubtype, status.New(codes.Internal, errMsg), http.StatusMethodNotAllowed, !frame.StreamEnded())
 		s.cancel()
 		return nil
 	}
@@ -590,27 +566,15 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
 			if !ok {
 				stat = status.New(codes.PermissionDenied, err.Error())
 			}
-			t.controlBuf.put(&earlyAbortStream{
-				httpStatus:     http.StatusOK,
-				streamID:       s.id,
-				contentSubtype: s.contentSubtype,
-				status:         stat,
-				rst:            !frame.StreamEnded(),
-			})
+			t.writeEarlyAbort(s.id, s.contentSubtype, stat, http.StatusOK, !frame.StreamEnded())
 			return nil
 		}
 	}
 
 	if s.ctx.Err() != nil {
 		t.mu.Unlock()
 		// Early abort in case the timeout was zero or so low it already fired.
-		t.controlBuf.put(&earlyAbortStream{
-			httpStatus:     http.StatusOK,
-			streamID:       s.id,
-			contentSubtype: s.contentSubtype,
-			status:         status.New(codes.DeadlineExceeded, context.DeadlineExceeded.Error()),
-			rst:            !frame.StreamEnded(),
-		})
+		t.writeEarlyAbort(s.id, s.contentSubtype, status.New(codes.DeadlineExceeded, context.DeadlineExceeded.Error()), http.StatusOK, !frame.StreamEnded())
 		return nil
 	}
 
@@ -969,23 +933,70 @@ func appendHeaderFieldsFromMD(headerFields []hpack.HeaderField, md metadata.MD)
 	return headerFields
 }
 
-func (t *http2Server) checkForHeaderListSize(it any) bool {
-	if t.maxSendHeaderListSize == nil {
+// checkForHeaderListSize checks if the header list size exceeds the limit set
+// by the peer. It returns false if the limit is exceeded.
+func checkForHeaderListSize(hf []hpack.HeaderField, maxSendHeaderListSize *uint32) bool {
+	if maxSendHeaderListSize == nil {
 		return true
 	}
-	hdrFrame := it.(*headerFrame)
 	var sz int64
-	for _, f := range hdrFrame.hf {
-		if sz += int64(f.Size()); sz > int64(*t.maxSendHeaderListSize) {
-			if t.logger.V(logLevel) {
-				t.logger.Infof("Header list size to send violates the maximum size (%d bytes) set by client", *t.maxSendHeaderListSize)
-			}
+	for _, f := range hf {
+		if sz += int64(f.Size()); sz > int64(*maxSendHeaderListSize) {
 			return false
 		}
 	}
 	return true
 }
 
+func (t *http2Server) checkForHeaderListSize(it any) bool {
+	hdrFrame := it.(*headerFrame)
+	if !checkForHeaderListSize(hdrFrame.hf, t.maxSendHeaderListSize) {
+		if t.logger.V(logLevel) {
+			t.logger.Infof("Header list size to send violates the maximum size (%d bytes) set by client", *t.maxSendHeaderListSize)
+		}
+		return false
+	}
+	return true
+}
+
+// buildEarlyAbortHF builds the header fields for an early abort response.
+func buildEarlyAbortHF(httpStatus uint32, contentSubtype string, stat *status.Status) []hpack.HeaderField {
+	hf := []hpack.HeaderField{
+		{Name: ":status", Value: strconv.Itoa(int(httpStatus))},
+		{Name: "content-type", Value: grpcutil.ContentType(contentSubtype)},
+		{Name: "grpc-status", Value: strconv.Itoa(int(stat.Code()))},
+		{Name: "grpc-message", Value: encodeGrpcMessage(stat.Message())},
+	}
+	if p := istatus.RawStatusProto(stat); len(p.GetDetails()) > 0 {
+		stBytes, err := proto.Marshal(p)
+		if err == nil {
+			hf = append(hf, hpack.HeaderField{Name: grpcStatusDetailsBinHeader, Value: encodeBinHeader(stBytes)})
+		}
+	}
+	return hf
+}
+
+// writeEarlyAbort sends an early abort response with the given HTTP status and gRPC status.
+// If the header list size exceeds the peer's limit, it sends a RST_STREAM instead.
+func (t *http2Server) writeEarlyAbort(streamID uint32, contentSubtype string, stat *status.Status, httpStatus uint32, rst bool) {
+	hf := buildEarlyAbortHF(httpStatus, contentSubtype, stat)
+	success, _ := t.controlBuf.executeAndPut(func() bool {
+		return checkForHeaderListSize(hf, t.maxSendHeaderListSize)
+	}, &earlyAbortStream{
+		streamID: streamID,
+		rst:      rst,
+		hf:       hf,
+	})
+	if !success {
+		t.controlBuf.put(&cleanupStream{
+			streamID: streamID,
+			rst:      true,
+			rstCode:  http2.ErrCodeInternal,
+			onWrite:  func() {},
+		})
+	}
+}
+
 func (t *http2Server) streamContextErr(s *ServerStream) error {
 	select {
 	case <-t.done:

test/end2end_test.go

@@ -6134,6 +6134,43 @@ func testClientMaxHeaderListSizeServerIntentionalViolation(t *testing.T, e env)
 	}
 }
 
+func (s) TestEarlyAbortStreamHeaderListSizeCheck(t *testing.T) {
+	lis, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatalf("Failed to listen: %v", err)
+	}
+	s := grpc.NewServer()
+	defer s.Stop()
+	go s.Serve(lis)
+
+	conn, err := net.DialTimeout("tcp", lis.Addr().String(), defaultTestTimeout)
+	if err != nil {
+		t.Fatalf("Failed to dial: %v", err)
+	}
+	defer conn.Close()
+	st := newServerTesterFromConn(t, conn)
+
+	// Set a very small MaxHeaderListSize that any response headers would violate.
+	st.greetWithSettings(http2.Setting{ID: http2.SettingMaxHeaderListSize, Val: 1})
+
+	// Send a request with an invalid content-type to trigger early abort.
+	st.writeHeaders(http2.HeadersFrameParam{
+		StreamID: 1,
+		BlockFragment: st.encodeHeader(
+			":method", "POST",
+			":path", "/grpc.testing.TestService/UnaryCall",
+			"content-type", "text/plain", // Invalid content-type to trigger early abort
+			"te", "trailers",
+		),
+		EndStream:  true,
+		EndHeaders: true,
+	})
+
+	// We should receive a RST_STREAM with ErrCodeInternal because the response
+	// headers exceed the MaxHeaderListSize limit.
+	st.wantRSTStream(http2.ErrCodeInternal)
+}
+
 func (s) TestNetPipeConn(t *testing.T) {
 	// This test will block indefinitely if grpc writes both client and server
 	// prefaces without either reading from the Conn.

test/servertester.go

@@ -91,8 +91,19 @@ func (st *serverTester) readFrame() (http2.Frame, error) {
 // greet initiates the client's HTTP/2 connection into a state where
 // frames may be sent.
 func (st *serverTester) greet() {
+	st.greetWithSettings()
+}
+
+// greetWithSettings initiates the client's HTTP/2 connection with custom settings.
+func (st *serverTester) greetWithSettings(settings ...http2.Setting) {
 	st.writePreface()
-	st.writeInitialSettings()
+	if len(settings) > 0 {
+		if err := st.fr.WriteSettings(settings...); err != nil {
+			st.t.Fatalf("Error writing initial SETTINGS frame from client to server: %v", err)
+		}
+	} else {
+		st.writeInitialSettings()
+	}
 	st.wantSettings()
 	st.writeSettingsAck()
 	for {