Repository: grpc/grpc-go
File: internal/transport/http2_server.go (Go)

Code before:
func checkForHeaderListSize(hf []hpack.HeaderField, maxSendHeaderListSize *uint32) bool {
	if maxSendHeaderListSize == nil {
		return true
	}
	var sz int64
	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,

Review comment (refactor):
nit: Since there's only one caller for this function, can we merge it into `writeEarlyAbort`? We can split the method in the future if it becomes too long.

Code after:
	if t.maxSendHeaderListSize == nil {
		return true
	}
	var sz int64
	for _, f := range 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)
			}
			return false
		}
	}
	return true
}

// 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 := []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 {
			t.logger.Errorf("Failed to marshal rpc status: %s, error: %v", pretty.ToJSON(p), err)
		}
		if err == nil {
			hf = append(hf, hpack.HeaderField{Name: grpcStatusDetailsBinHeader, Value: encodeBinHeader(stBytes)})
		}
	}
	success, _ := t.controlBuf.executeAndPut(func() bool {
		return t.checkForHeaderListSize(hf)
	}, &earlyAbortStream{
		streamID: streamID,
		rst:      rst,
		hf:       hf,
	})
	if !success {
		t.controlBuf.put(&cleanupStream{
			streamID: streamID,
			rst:      true,
			rstCode:  http2.ErrCodeInternal,
			onWrite:  func() {},
		})
	}
}