address arjun's comment

Author: joybestourous

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,12 +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
-	maxSendHeaderListSize *uint32
+	streamID uint32
+	rst      bool
+	hf       []hpack.HeaderField // Pre-built header fields
 }
 
 func (*earlyAbortStream) isTransportResponseFrame() bool { return false }
@@ -847,34 +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 !checkForHeaderListSize(headerFields, eas.maxSendHeaderListSize) {
-		if l.logger.V(logLevel) {
-			l.logger.Infof("Header list size to send violates the maximum size (%d bytes) set by client", *eas.maxSendHeaderListSize)
-		}
-		return l.framer.fr.WriteRSTStream(eas.streamID, http2.ErrCodeInternal)
-	}
-
-	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,17 +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)
 		}
-		eas := &earlyAbortStream{
-			httpStatus:     http.StatusBadRequest,
-			streamID:       streamID,
-			contentSubtype: s.contentSubtype,
-			status:         status.New(codes.Internal, errMsg),
-			rst:            !frame.StreamEnded(),
-		}
-		t.controlBuf.executeAndPut(func() bool {
-			eas.maxSendHeaderListSize = t.maxSendHeaderListSize
-			return true
-		}, eas)
+		t.writeEarlyAbort(streamID, s.contentSubtype, status.New(codes.Internal, errMsg), http.StatusBadRequest, !frame.StreamEnded())
 		return nil
 	}
 
@@ -503,31 +493,11 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
 		return nil
 	}
 	if !isGRPC {
-		eas := &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.controlBuf.executeAndPut(func() bool {
-			eas.maxSendHeaderListSize = t.maxSendHeaderListSize
-			return true
-		}, eas)
+		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 {
-		eas := &earlyAbortStream{
-			httpStatus:     http.StatusBadRequest,
-			streamID:       streamID,
-			contentSubtype: s.contentSubtype,
-			status:         headerError,
-			rst:            !frame.StreamEnded(),
-		}
-		t.controlBuf.executeAndPut(func() bool {
-			eas.maxSendHeaderListSize = t.maxSendHeaderListSize
-			return true
-		}, eas)
+		t.writeEarlyAbort(streamID, s.contentSubtype, headerError, http.StatusBadRequest, !frame.StreamEnded())
 		return nil
 	}
 
@@ -581,17 +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)
 		}
-		eas := &earlyAbortStream{
-			httpStatus:     http.StatusMethodNotAllowed,
-			streamID:       streamID,
-			contentSubtype: s.contentSubtype,
-			status:         status.New(codes.Internal, errMsg),
-			rst:            !frame.StreamEnded(),
-		}
-		t.controlBuf.executeAndPut(func() bool {
-			eas.maxSendHeaderListSize = t.maxSendHeaderListSize
-			return true
-		}, eas)
+		t.writeEarlyAbort(streamID, s.contentSubtype, status.New(codes.Internal, errMsg), http.StatusMethodNotAllowed, !frame.StreamEnded())
 		s.cancel()
 		return nil
 	}
@@ -606,35 +566,15 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
 			if !ok {
 				stat = status.New(codes.PermissionDenied, err.Error())
 			}
-			eas := &earlyAbortStream{
-				httpStatus:     http.StatusOK,
-				streamID:       s.id,
-				contentSubtype: s.contentSubtype,
-				status:         stat,
-				rst:            !frame.StreamEnded(),
-			}
-			t.controlBuf.executeAndPut(func() bool {
-				eas.maxSendHeaderListSize = t.maxSendHeaderListSize
-				return true
-			}, eas)
+			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.
-		eas := &earlyAbortStream{
-			httpStatus:     http.StatusOK,
-			streamID:       s.id,
-			contentSubtype: s.contentSubtype,
-			status:         status.New(codes.DeadlineExceeded, context.DeadlineExceeded.Error()),
-			rst:            !frame.StreamEnded(),
-		}
-		t.controlBuf.executeAndPut(func() bool {
-			eas.maxSendHeaderListSize = t.maxSendHeaderListSize
-			return true
-		}, eas)
+		t.writeEarlyAbort(s.id, s.contentSubtype, status.New(codes.DeadlineExceeded, context.DeadlineExceeded.Error()), http.StatusOK, !frame.StreamEnded())
 		return nil
 	}
 
@@ -1019,6 +959,44 @@ func (t *http2Server) checkForHeaderListSize(it any) bool {
 	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: