github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/nphttp2/grpc/grpcframe/frame_writer.go

/*
 * Copyright 2021 The Go Authors. All rights reserved.
 *
 * Use of this source code is governed by a BSD-style
 * license that can be found in the LICENSE file.
 *
 * This file may have been modified by CloudWeGo authors. All CloudWeGo
 * Modifications are Copyright 2022 CloudWeGo Authors.
 *
 * Code forked and modified from golang v1.17.4
 */

package grpcframe

import (
	"errors"

	"golang.org/x/net/http2"
)

const frameHeaderLen = 9

var padZeros = make([]byte, 255) // zeros for padding

var (
	errStreamID    = errors.New("invalid stream ID")
	errDepStreamID = errors.New("invalid dependent stream ID")
)

func (fr *Framer) startWrite(ftype http2.FrameType, flags http2.Flags, streamID uint32, payloadLen int) (err error) {
	if payloadLen >= (1 << 24) {
		return http2.ErrFrameTooLarge
	}
	fr.wbuf = append(fr.wbuf[:0],
		byte(payloadLen>>16),
		byte(payloadLen>>8),
		byte(payloadLen),
		byte(ftype),
		byte(flags),
		byte(streamID>>24),
		byte(streamID>>16),
		byte(streamID>>8),
		byte(streamID))
	return nil
}

func (fr *Framer) endWrite() (err error) {
	_, err = fr.writer.Write(fr.wbuf)
	return err
}

func (fr *Framer) writeByte(v byte)     { fr.wbuf = append(fr.wbuf, v) }
func (fr *Framer) writeBytes(v []byte)  { fr.wbuf = append(fr.wbuf, v...) }
func (fr *Framer) writeUint16(v uint16) { fr.wbuf = append(fr.wbuf, byte(v>>8), byte(v)) }
func (fr *Framer) writeUint32(v uint32) {
	fr.wbuf = append(fr.wbuf, byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
}

// WriteData writes a DATA frame.
//
// It will perform exactly one Write to the underlying Writer.
// It is the caller's responsibility not to violate the maximum frame size
// and to not call other Write methods concurrently.
func (fr *Framer) WriteData(streamID uint32, endStream bool, data []byte) error {
	if !validStreamID(streamID) && !fr.AllowIllegalWrites {
		return errStreamID
	}

	var flags http2.Flags
	if endStream {
		flags |= http2.FlagDataEndStream
	}

	err := fr.startWrite(http2.FrameData, flags, streamID, len(data))
	if err != nil {
		return err
	}
	fr.writeBytes(data)
	return fr.endWrite()
}

// WriteHeaders writes a single HEADERS frame.
//
// This is a low-level header writing method. Encoding headers and
// splitting them into any necessary CONTINUATION frames is handled
// elsewhere.
//
// It will perform exactly one Write to the underlying Writer.
// It is the caller's responsibility to not call other Write methods concurrently.
func (fr *Framer) WriteHeaders(p http2.HeadersFrameParam) error {
	if !validStreamID(p.StreamID) && !fr.AllowIllegalWrites {
		return errStreamID
	}
	var payloadLen int
	var flags http2.Flags
	if p.PadLength != 0 {
		flags |= http2.FlagHeadersPadded
		payloadLen += 1 + int(p.PadLength)
	}
	if p.EndStream {
		flags |= http2.FlagHeadersEndStream
	}
	if p.EndHeaders {
		flags |= http2.FlagHeadersEndHeaders
	}
	if !p.Priority.IsZero() {
		v := p.Priority.StreamDep
		if !validStreamIDOrZero(v) && !fr.AllowIllegalWrites {
			return errDepStreamID
		}
		flags |= http2.FlagHeadersPriority
		payloadLen += 5
	}
	payloadLen += len(p.BlockFragment)
	err := fr.startWrite(http2.FrameHeaders, flags, p.StreamID, payloadLen)
	if err != nil {
		return err
	}
	if p.PadLength != 0 {
		fr.writeByte(p.PadLength)
	}
	if !p.Priority.IsZero() {
		v := p.Priority.StreamDep
		if p.Priority.Exclusive {
			v |= 1 << 31
		}
		fr.writeUint32(v)
		fr.writeByte(p.Priority.Weight)
	}
	fr.writeBytes(p.BlockFragment)
	fr.writeBytes(padZeros[:p.PadLength])
	return fr.endWrite()
}

// WritePriority writes a PRIORITY frame.
//
// It will perform exactly one Write to the underlying Writer.
// It is the caller's responsibility to not call other Write methods concurrently.
func (fr *Framer) WritePriority(streamID uint32, p http2.PriorityParam) error {
	if !validStreamID(streamID) && !fr.AllowIllegalWrites {
		return errStreamID
	}
	if !validStreamIDOrZero(p.StreamDep) {
		return errDepStreamID
	}

	err := fr.startWrite(http2.FramePriority, 0, streamID, 5)
	if err != nil {
		return err
	}

	v := p.StreamDep
	if p.Exclusive {
		v |= 1 << 31
	}
	fr.writeUint32(v)
	fr.writeByte(p.Weight)
	return fr.endWrite()
}

// WriteRSTStream writes a RST_STREAM frame.
//
// It will perform exactly one Write to the underlying Writer.
// It is the caller's responsibility to not call other Write methods concurrently.
func (fr *Framer) WriteRSTStream(streamID uint32, code http2.ErrCode) error {
	if !validStreamID(streamID) && !fr.AllowIllegalWrites {
		return errStreamID
	}
	err := fr.startWrite(http2.FrameRSTStream, 0, streamID, 4)
	if err != nil {
		return err
	}
	fr.writeUint32(uint32(code))
	return fr.endWrite()
}

// WriteSettings writes a SETTINGS frame with zero or more settings
// specified and the ACK bit not set.
//
// It will perform exactly one Write to the underlying Writer.
// It is the caller's responsibility to not call other Write methods concurrently.
func (fr *Framer) WriteSettings(settings ...http2.Setting) error {
	payloadLen := 6 * len(settings)
	err := fr.startWrite(http2.FrameSettings, 0, 0, payloadLen)
	if err != nil {
		return err
	}
	for _, s := range settings {
		fr.writeUint16(uint16(s.ID))
		fr.writeUint32(s.Val)
	}
	return fr.endWrite()
}

// WriteSettingsAck writes an empty SETTINGS frame with the ACK bit set.
//
// It will perform exactly one Write to the underlying Writer.
// It is the caller's responsibility to not call other Write methods concurrently.
func (fr *Framer) WriteSettingsAck() error {
	err := fr.startWrite(http2.FrameSettings, http2.FlagSettingsAck, 0, 0)
	if err != nil {
		return err
	}
	return fr.endWrite()
}

// WritePushPromise writes a single PushPromise Frame.
//
// As with Header Frames, This is the low level call for writing
// individual frames. Continuation frames are handled elsewhere.
//
// It will perform exactly one Write to the underlying Writer.
// It is the caller's responsibility to not call other Write methods concurrently.
func (fr *Framer) WritePushPromise(p http2.PushPromiseParam) error {
	if !fr.AllowIllegalWrites && (!validStreamID(p.StreamID) || !validStreamID(p.PromiseID)) {
		return errStreamID
	}
	if !validStreamID(p.PromiseID) && !fr.AllowIllegalWrites {
		return errStreamID
	}
	var payloadLen int
	var flags http2.Flags
	if p.PadLength != 0 {
		flags |= http2.FlagPushPromisePadded
		payloadLen += 1 + int(p.PadLength)
	}
	if p.EndHeaders {
		flags |= http2.FlagPushPromiseEndHeaders
	}
	payloadLen += 4 + len(p.BlockFragment)
	err := fr.startWrite(http2.FramePushPromise, flags, p.StreamID, payloadLen)
	if err != nil {
		return err
	}
	if p.PadLength != 0 {
		fr.writeByte(p.PadLength)
	}
	fr.writeUint32(p.PromiseID)
	fr.writeBytes(p.BlockFragment)
	fr.writeBytes(padZeros[:p.PadLength])
	return fr.endWrite()
}

func (fr *Framer) WritePing(ack bool, data [8]byte) error {
	var flags http2.Flags
	if ack {
		flags = http2.FlagPingAck
	}
	err := fr.startWrite(http2.FramePing, flags, 0, len(data))
	if err != nil {
		return err
	}
	fr.writeBytes(data[:])
	return fr.endWrite()
}

func (fr *Framer) WriteGoAway(maxStreamID uint32, code http2.ErrCode, debugData []byte) error {
	payloadLen := 8 + len(debugData)
	err := fr.startWrite(http2.FrameGoAway, 0, 0, payloadLen)
	if err != nil {
		return err
	}
	fr.writeUint32(maxStreamID & (1<<31 - 1))
	fr.writeUint32(uint32(code))
	fr.writeBytes(debugData)
	return fr.endWrite()
}

// WriteWindowUpdate writes a WINDOW_UPDATE frame.
// The increment value must be between 1 and 2,147,483,647, inclusive.
// If the Stream ID is zero, the window update applies to the
// connection as a whole.
func (fr *Framer) WriteWindowUpdate(streamID, incr uint32) error {
	// "The legal range for the increment to the flow control window is 1 to 2^31-1 (2,147,483,647) octets."
	if (incr < 1 || incr > 2147483647) && !fr.AllowIllegalWrites {
		return errors.New("illegal window increment value")
	}

	err := fr.startWrite(http2.FrameWindowUpdate, 0, streamID, 4)
	if err != nil {
		return err
	}
	fr.writeUint32(incr)
	return fr.endWrite()
}

// WriteContinuation writes a CONTINUATION frame.
//
// It will perform exactly one Write to the underlying Writer.
// It is the caller's responsibility to not call other Write methods concurrently.
func (fr *Framer) WriteContinuation(streamID uint32, endHeaders bool, headerBlockFragment []byte) error {
	if !validStreamID(streamID) && !fr.AllowIllegalWrites {
		return errStreamID
	}
	var flags http2.Flags
	if endHeaders {
		flags |= http2.FlagContinuationEndHeaders
	}
	err := fr.startWrite(http2.FrameContinuation, flags, streamID, len(headerBlockFragment))
	if err != nil {
		return err
	}
	fr.writeBytes(headerBlockFragment)
	return fr.endWrite()
}

// WriteRawFrame writes a raw frame. This can be used to write
// extension frames unknown to this package.
func (fr *Framer) WriteRawFrame(t http2.FrameType, flags http2.Flags, streamID uint32, payload []byte) error {
	err := fr.startWrite(t, flags, streamID, len(payload))
	if err != nil {
		return err
	}
	fr.writeBytes(payload)
	return fr.endWrite()
}

func validStreamIDOrZero(streamID uint32) bool {
	return streamID&(1<<31) == 0
}

func validStreamID(streamID uint32) bool {
	return streamID != 0 && streamID&(1<<31) == 0
}