/*
 * 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 (
	"encoding/binary"
	"fmt"

	"github.com/cloudwego/netpoll"
	"golang.org/x/net/http2"
	"golang.org/x/net/http2/hpack"
)

/* struct:
 * DataFrame
 * HeadersFrame
 * SettingsFrame
 * PushPromiseFrame
 * GoAwayFrame
 * ContinuationFrame
 * MetaHeadersFrame
 */

// a frameParser parses a frame given its FrameHeader and payload
// bytes. The length of payload will always equal fh.Length (which
// might be 0).
type frameParser func(fc *frameCache, fh http2.FrameHeader, payload []byte) (http2.Frame, error)

var frameParsers = []frameParser{
	// FrameData:         parseDataFrame,
	http2.FrameHeaders:      parseHeadersFrame,
	http2.FramePriority:     parsePriorityFrame,
	http2.FrameRSTStream:    parseRSTStreamFrame,
	http2.FrameSettings:     parseSettingsFrame,
	http2.FramePushPromise:  parsePushPromise,
	http2.FramePing:         parsePingFrame,
	http2.FrameGoAway:       parseGoAwayFrame,
	http2.FrameWindowUpdate: parseWindowUpdateFrame,
	http2.FrameContinuation: parseContinuationFrame,
}

func typeFrameParser(t http2.FrameType) frameParser {
	if 1 <= t && t <= 9 {
		return frameParsers[t]
	}
	return parseUnknownFrame
}

// A DataFrame conveys arbitrary, variable-length sequences of octets
// associated with a stream.
// See http://http2.github.io/http2-spec/#rfc.section.6.1
type DataFrame struct {
	http2.FrameHeader
	data []byte // TODO: data = netpoll.Reader here
}

func (f *DataFrame) StreamEnded() bool {
	return f.FrameHeader.Flags.Has(http2.FlagDataEndStream)
}

// Data returns the frame's data octets, not including any padding
// size byte or padding suffix bytes.
// The caller must not retain the returned memory past the next
// call to ReadFrame.
func (f *DataFrame) Data() []byte {
	return f.data
}

func parseDataFrame(fc *frameCache, fh http2.FrameHeader, payload netpoll.Reader) (http2.Frame, error) {
	if fh.StreamID == 0 {
		// DATA frames MUST be associated with a stream. If a
		// DATA frame is received whose stream identifier
		// field is 0x0, the recipient MUST respond with a
		// connection error (Section 5.4.1) of type
		// PROTOCOL_ERROR.
		return nil, connError{http2.ErrCodeProtocol, "DATA frame with stream ID 0"}
	}
	f := fc.getDataFrame()
	f.FrameHeader = fh

	var padSize byte
	payloadLen := int(fh.Length)
	if fh.Flags.Has(http2.FlagDataPadded) {
		var err error
		padSize, err = payload.ReadByte()
		payloadLen--
		if err != nil {
			return nil, err
		}
	}
	if int(padSize) > payloadLen {
		// If the length of the padding is greater than the
		// length of the frame payload, the recipient MUST
		// treat this as a connection error.
		// Filed: https://github.com/http2/http2-spec/issues/610
		return nil, connError{http2.ErrCodeProtocol, "pad size larger than data payload"}
	}
	data, err := payload.Next(payloadLen)
	if err != nil {
		return nil, err
	}
	f.data = data[:payloadLen-int(padSize)]
	return f, nil
}

// A HeadersFrame is used to open a stream and additionally carries a
// header block fragment.
type HeadersFrame struct {
	http2.FrameHeader

	// Priority is set if FlagHeadersPriority is set in the FrameHeader.
	Priority http2.PriorityParam

	headerFragBuf []byte // not owned
}

func (f *HeadersFrame) HeaderBlockFragment() []byte {
	return f.headerFragBuf
}

func (f *HeadersFrame) HeadersEnded() bool {
	return f.FrameHeader.Flags.Has(http2.FlagHeadersEndHeaders)
}

func (f *HeadersFrame) StreamEnded() bool {
	return f.FrameHeader.Flags.Has(http2.FlagHeadersEndStream)
}

func (f *HeadersFrame) HasPriority() bool {
	return f.FrameHeader.Flags.Has(http2.FlagHeadersPriority)
}

func parseHeadersFrame(_ *frameCache, fh http2.FrameHeader, p []byte) (_ http2.Frame, err error) {
	hf := &HeadersFrame{
		FrameHeader: fh,
	}
	if fh.StreamID == 0 {
		// HEADERS frames MUST be associated with a stream. If a HEADERS frame
		// is received whose stream identifier field is 0x0, the recipient MUST
		// respond with a connection error (Section 5.4.1) of type
		// PROTOCOL_ERROR.
		return nil, connError{http2.ErrCodeProtocol, "HEADERS frame with stream ID 0"}
	}
	var padLength uint8
	if fh.Flags.Has(http2.FlagHeadersPadded) {
		if p, padLength, err = readByte(p); err != nil {
			return
		}
	}
	if fh.Flags.Has(http2.FlagHeadersPriority) {
		var v uint32
		p, v, err = readUint32(p)
		if err != nil {
			return nil, err
		}
		hf.Priority.StreamDep = v & 0x7fffffff
		hf.Priority.Exclusive = (v != hf.Priority.StreamDep) // high bit was set
		p, hf.Priority.Weight, err = readByte(p)
		if err != nil {
			return nil, err
		}
	}
	// note: len(p)-int(padLength) == 0 is valid
	if len(p)-int(padLength) < 0 {
		return nil, http2.StreamError{StreamID: fh.StreamID, Code: http2.ErrCodeProtocol}
	}
	hf.headerFragBuf = p[:len(p)-int(padLength)]
	return hf, nil
}

func parsePriorityFrame(_ *frameCache, fh http2.FrameHeader, payload []byte) (http2.Frame, error) {
	if fh.StreamID == 0 {
		return nil, connError{http2.ErrCodeProtocol, "PRIORITY frame with stream ID 0"}
	}
	if len(payload) != 5 {
		return nil, connError{http2.ErrCodeFrameSize, fmt.Sprintf("PRIORITY frame payload size was %d; want 5", len(payload))}
	}
	v := binary.BigEndian.Uint32(payload[:4])
	streamID := v & 0x7fffffff // mask off high bit
	return &http2.PriorityFrame{
		FrameHeader: fh,
		PriorityParam: http2.PriorityParam{
			Weight:    payload[4],
			StreamDep: streamID,
			Exclusive: streamID != v, // was high bit set?
		},
	}, nil
}

func parseRSTStreamFrame(_ *frameCache, fh http2.FrameHeader, p []byte) (http2.Frame, error) {
	if len(p) != 4 {
		return nil, http2.ConnectionError(http2.ErrCodeFrameSize)
	}
	if fh.StreamID == 0 {
		return nil, http2.ConnectionError(http2.ErrCodeProtocol)
	}
	return &http2.RSTStreamFrame{FrameHeader: fh, ErrCode: http2.ErrCode(binary.BigEndian.Uint32(p[:4]))}, nil
}

// A SettingsFrame conveys configuration parameters that affect how
// endpoints communicate, such as preferences and constraints on peer
// behavior.
//
// See http://http2.github.io/http2-spec/#SETTINGS
type SettingsFrame struct {
	http2.FrameHeader
	p []byte
}

func parseSettingsFrame(_ *frameCache, fh http2.FrameHeader, p []byte) (http2.Frame, error) {
	if fh.Flags.Has(http2.FlagSettingsAck) && fh.Length > 0 {
		// When this (ACK 0x1) bit is set, the payload of the
		// SETTINGS frame MUST be empty. Receipt of a
		// SETTINGS frame with the ACK flag set and a length
		// field value other than 0 MUST be treated as a
		// connection error (Section 5.4.1) of type
		// FRAME_SIZE_ERROR.
		return nil, http2.ConnectionError(http2.ErrCodeFrameSize)
	}
	if fh.StreamID != 0 {
		// SETTINGS frames always apply to a connection,
		// never a single stream. The stream identifier for a
		// SETTINGS frame MUST be zero (0x0).  If an endpoint
		// receives a SETTINGS frame whose stream identifier
		// field is anything other than 0x0, the endpoint MUST
		// respond with a connection error (Section 5.4.1) of
		// type PROTOCOL_ERROR.
		return nil, http2.ConnectionError(http2.ErrCodeProtocol)
	}
	if len(p)%6 != 0 {
		// Expecting even number of 6 byte settings.
		return nil, http2.ConnectionError(http2.ErrCodeFrameSize)
	}
	f := &SettingsFrame{FrameHeader: fh, p: p}
	if v, ok := f.Value(http2.SettingInitialWindowSize); ok && v > (1<<31)-1 {
		// Values above the maximum flow control window size of 2^31 - 1 MUST
		// be treated as a connection error (Section 5.4.1) of type
		// FLOW_CONTROL_ERROR.
		return nil, http2.ConnectionError(http2.ErrCodeFlowControl)
	}
	return f, nil
}

func (f *SettingsFrame) IsAck() bool {
	return f.FrameHeader.Flags.Has(http2.FlagSettingsAck)
}

func (f *SettingsFrame) Value(id http2.SettingID) (v uint32, ok bool) {
	for i := 0; i < f.NumSettings(); i++ {
		if s := f.Setting(i); s.ID == id {
			return s.Val, true
		}
	}
	return 0, false
}

// Setting returns the setting from the frame at the given 0-based index.
// The index must be >= 0 and less than f.NumSettings().
func (f *SettingsFrame) Setting(i int) http2.Setting {
	buf := f.p
	return http2.Setting{
		ID:  http2.SettingID(binary.BigEndian.Uint16(buf[i*6 : i*6+2])),
		Val: binary.BigEndian.Uint32(buf[i*6+2 : i*6+6]),
	}
}

func (f *SettingsFrame) NumSettings() int { return len(f.p) / 6 }

// HasDuplicates reports whether f contains any duplicate setting IDs.
func (f *SettingsFrame) HasDuplicates() bool {
	num := f.NumSettings()
	if num == 0 {
		return false
	}
	// If it's small enough (the common case), just do the n^2
	// thing and avoid a map allocation.
	if num < 10 {
		for i := 0; i < num; i++ {
			idi := f.Setting(i).ID
			for j := i + 1; j < num; j++ {
				idj := f.Setting(j).ID
				if idi == idj {
					return true
				}
			}
		}
		return false
	}
	seen := map[http2.SettingID]bool{}
	for i := 0; i < num; i++ {
		id := f.Setting(i).ID
		if seen[id] {
			return true
		}
		seen[id] = true
	}
	return false
}

// ForeachSetting runs fn for each setting.
// It stops and returns the first error.
func (f *SettingsFrame) ForeachSetting(fn func(http2.Setting) error) error {
	for i := 0; i < f.NumSettings(); i++ {
		if err := fn(f.Setting(i)); err != nil {
			return err
		}
	}
	return nil
}

// A PushPromiseFrame is used to initiate a server stream.
// See http://http2.github.io/http2-spec/#rfc.section.6.6

// A PushPromiseFrame is used to initiate a server stream.
// See http://http2.github.io/http2-spec/#rfc.section.6.6
type PushPromiseFrame struct {
	http2.FrameHeader
	PromiseID     uint32
	headerFragBuf []byte // not owned
}

func (f *PushPromiseFrame) HeaderBlockFragment() []byte {
	return f.headerFragBuf
}

func (f *PushPromiseFrame) HeadersEnded() bool {
	return f.FrameHeader.Flags.Has(http2.FlagPushPromiseEndHeaders)
}

func parsePushPromise(_ *frameCache, fh http2.FrameHeader, p []byte) (_ http2.Frame, err error) {
	pp := &PushPromiseFrame{
		FrameHeader: fh,
	}
	if pp.StreamID == 0 {
		// PUSH_PROMISE frames MUST be associated with an existing,
		// peer-initiated stream. The stream identifier of a
		// PUSH_PROMISE frame indicates the stream it is associated
		// with. If the stream identifier field specifies the value
		// 0x0, a recipient MUST respond with a connection error
		// (Section 5.4.1) of type PROTOCOL_ERROR.
		return nil, http2.ConnectionError(http2.ErrCodeProtocol)
	}
	// The PUSH_PROMISE frame includes optional padding.
	// Padding fields and flags are identical to those defined for DATA frames
	var padLength uint8
	if fh.Flags.Has(http2.FlagPushPromisePadded) {
		if p, padLength, err = readByte(p); err != nil {
			return
		}
	}

	p, pp.PromiseID, err = readUint32(p)
	if err != nil {
		return
	}
	pp.PromiseID = pp.PromiseID & (1<<31 - 1)

	if int(padLength) > len(p) {
		// like the DATA frame, error out if padding is longer than the body.
		return nil, http2.ConnectionError(http2.ErrCodeProtocol)
	}
	pp.headerFragBuf = p[:len(p)-int(padLength)]
	return pp, nil
}

func parsePingFrame(_ *frameCache, fh http2.FrameHeader, payload []byte) (http2.Frame, error) {
	if len(payload) != 8 {
		return nil, http2.ConnectionError(http2.ErrCodeFrameSize)
	}
	if fh.StreamID != 0 {
		return nil, http2.ConnectionError(http2.ErrCodeProtocol)
	}
	f := &http2.PingFrame{FrameHeader: fh}
	copy(f.Data[:], payload)
	return f, nil
}

// A GoAwayFrame informs the remote peer to stop creating streams on this connection.
// See http://http2.github.io/http2-spec/#rfc.section.6.8
type GoAwayFrame struct {
	http2.FrameHeader
	LastStreamID uint32
	ErrCode      http2.ErrCode
	debugData    []byte
}

// DebugData returns any debug data in the GOAWAY frame. Its contents
// are not defined.
// The caller must not retain the returned memory past the next
// call to ReadFrame.
func (f *GoAwayFrame) DebugData() []byte {
	return f.debugData
}

func parseGoAwayFrame(_ *frameCache, fh http2.FrameHeader, p []byte) (http2.Frame, error) {
	if fh.StreamID != 0 {
		return nil, http2.ConnectionError(http2.ErrCodeProtocol)
	}
	if len(p) < 8 {
		return nil, http2.ConnectionError(http2.ErrCodeFrameSize)
	}
	return &GoAwayFrame{
		FrameHeader:  fh,
		LastStreamID: binary.BigEndian.Uint32(p[:4]) & (1<<31 - 1),
		ErrCode:      http2.ErrCode(binary.BigEndian.Uint32(p[4:8])),
		debugData:    p[8:],
	}, nil
}

func parseWindowUpdateFrame(_ *frameCache, fh http2.FrameHeader, p []byte) (http2.Frame, error) {
	if len(p) != 4 {
		return nil, http2.ConnectionError(http2.ErrCodeFrameSize)
	}
	inc := binary.BigEndian.Uint32(p[:4]) & 0x7fffffff // mask off high reserved bit
	if inc == 0 {
		// A receiver MUST treat the receipt of a
		// WINDOW_UPDATE frame with an flow control window
		// increment of 0 as a stream error (Section 5.4.2) of
		// type PROTOCOL_ERROR; errors on the connection flow
		// control window MUST be treated as a connection
		// error (Section 5.4.1).
		if fh.StreamID == 0 {
			return nil, http2.ConnectionError(http2.ErrCodeProtocol)
		}
		return nil, http2.StreamError{StreamID: fh.StreamID, Code: http2.ErrCodeProtocol}
	}
	return &http2.WindowUpdateFrame{
		FrameHeader: fh,
		Increment:   inc,
	}, nil
}

// A ContinuationFrame is used to continue a sequence of header block fragments.
// See http://http2.github.io/http2-spec/#rfc.section.6.10
type ContinuationFrame struct {
	http2.FrameHeader
	headerFragBuf []byte
}

func parseContinuationFrame(_ *frameCache, fh http2.FrameHeader, p []byte) (http2.Frame, error) {
	if fh.StreamID == 0 {
		return nil, connError{http2.ErrCodeProtocol, "CONTINUATION frame with stream ID 0"}
	}
	return &ContinuationFrame{FrameHeader: fh, headerFragBuf: p}, nil
}

func (f *ContinuationFrame) HeaderBlockFragment() []byte {
	return f.headerFragBuf
}

func (f *ContinuationFrame) HeadersEnded() bool {
	return f.FrameHeader.Flags.Has(http2.FlagContinuationEndHeaders)
}

// An UnknownFrame is the frame type returned when the frame type is unknown
// or no specific frame type parser exists.
type UnknownFrame struct {
	http2.FrameHeader
	p []byte
}

// Payload returns the frame's payload (after the header).  It is not
// valid to call this method after a subsequent call to
// Framer.ReadFrame, nor is it valid to retain the returned slice.
// The memory is owned by the Framer and is invalidated when the next
// frame is read.
func (f *UnknownFrame) Payload() []byte {
	return f.p
}

func parseUnknownFrame(_ *frameCache, fh http2.FrameHeader, p []byte) (http2.Frame, error) {
	return &UnknownFrame{fh, p}, nil
}

// A MetaHeadersFrame is the representation of one HEADERS frame and
// zero or more contiguous CONTINUATION frames and the decoding of
// their HPACK-encoded contents.
//
// This type of frame does not appear on the wire and is only returned
// by the Framer when Framer.ReadMetaHeaders is set.
type MetaHeadersFrame struct {
	*HeadersFrame

	// Fields are the fields contained in the HEADERS and
	// CONTINUATION frames. The underlying slice is owned by the
	// Framer and must not be retained after the next call to
	// ReadFrame.
	//
	// Fields are guaranteed to be in the correct http2 order and
	// not have unknown pseudo header fields or invalid header
	// field names or values. Required pseudo header fields may be
	// missing, however. Use the MetaHeadersFrame.Pseudo accessor
	// method access pseudo headers.
	Fields []hpack.HeaderField

	// Truncated is whether the max header list size limit was hit
	// and Fields is incomplete. The hpack decoder state is still
	// valid, however.
	Truncated bool
}

// PseudoValue returns the given pseudo header field's value.
// The provided pseudo field should not contain the leading colon.
func (mh *MetaHeadersFrame) PseudoValue(pseudo string) string {
	for _, hf := range mh.Fields {
		if !hf.IsPseudo() {
			return ""
		}
		if hf.Name[1:] == pseudo {
			return hf.Value
		}
	}
	return ""
}

// RegularFields returns the regular (non-pseudo) header fields of mh.
// The caller does not own the returned slice.
func (mh *MetaHeadersFrame) RegularFields() []hpack.HeaderField {
	for i, hf := range mh.Fields {
		if !hf.IsPseudo() {
			return mh.Fields[i:]
		}
	}
	return nil
}

// PseudoFields returns the pseudo header fields of mh.
// The caller does not own the returned slice.
func (mh *MetaHeadersFrame) PseudoFields() []hpack.HeaderField {
	for i, hf := range mh.Fields {
		if !hf.IsPseudo() {
			return mh.Fields[:i]
		}
	}
	return mh.Fields
}

func (mh *MetaHeadersFrame) checkPseudos() error {
	var isRequest, isResponse bool
	pf := mh.PseudoFields()
	for i, hf := range pf {
		switch hf.Name {
		case ":method", ":path", ":scheme", ":authority":
			isRequest = true
		case ":status":
			isResponse = true
		default:
			return pseudoHeaderError(hf.Name)
		}
		// Check for duplicates.
		// This would be a bad algorithm, but N is 4.
		// And this doesn't allocate.
		for _, hf2 := range pf[:i] {
			if hf.Name == hf2.Name {
				return duplicatePseudoHeaderError(hf.Name)
			}
		}
	}
	if isRequest && isResponse {
		return errMixPseudoHeaderTypes
	}
	return nil
}
