<commit_before>package quiclatest

import (
	"encoding/binary"

	"github.com/ami-GS/gQUIC/latest/qtype"
)

type Packet interface {
	GetWire() ([]byte, error)
	GetHeader() PacketHeader
	SetHeader(ph PacketHeader)
	GetFrames() []Frame
	SetFrames(fs []Frame)
}

func ParsePacket(data []byte) (packet Packet, idx int, err error) {
	version := binary.BigEndian.Uint32(data[1:5])
	if version == 0 {
		// VersionNegotiationPacket use all UDP datagram (doesn't have frame)
		return ParseVersionNegotiationPacket(data)
	}
	header, idx, err := PacketHeaderParserMap[PacketHeaderType(data[0]&0x80)](data) // ParseHeader
	if err != nil {
		return nil, 0, err
	}

	fs, idxTmp, err := ParseFrames(data[idx:])
	if err != nil {
		return nil, 0, err
	}

	packet, err = newPacket(data[0], header, fs)
	if err != nil {
		return nil, 0, err
	}

	return packet, idx + idxTmp, err
}

type BasePacket struct {
	Header     PacketHeader
	Frames     []Frame
	PaddingNum int
}

func (bp *BasePacket) GetHeader() PacketHeader {
	return bp.Header
}
func (bp *BasePacket) SetHeader(h PacketHeader) {
	bp.Header = h
}

func (bp *BasePacket) SetFrames(fs []Frame) {
	bp.Frames = fs
}

func (bp *BasePacket) GetFrames() []Frame {
	return bp.Frames
}

func (bp *BasePacket) GetWire() (wire []byte, err error) {
	hWire := bp.Header.GetWire()
	if err != nil {
		return nil, err
	}
	fsWire, err := GetFrameWires(bp.Frames)
	if err != nil {
		return nil, err
	}
	if bp.PaddingNum != 0 {
		fsWire = append(fsWire, make([]byte, bp.PaddingNum)...)
	}
	// TODO: protect for short header?
	return append(hWire, fsWire...), nil
}

func newPacket(firstByte byte, ph PacketHeader, fs []Frame) (Packet, error) {
	// TODO: needs frame type validation for each packet type
	if firstByte&byte(LongHeaderType) == byte(LongHeaderType) {
		payloadLen := 0
		for _, frame := range fs {
			payloadLen += frame.GetWireSize()
		}

		switch LongHeaderPacketType(firstByte & 0x7f) {
		case InitialPacketType:
			// TODO: check whether it is stream frame or not
			return &InitialPacket{
				BasePacket: &BasePacket{
					Header:     ph,
					Frames:     fs,
					PaddingNum: InitialPacketMinimumPayloadSize - payloadLen,
				},
			}, nil
		case RetryPacketType:
			return &RetryPacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case HandshakePacketType:
			return &HandshakePacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case ZeroRTTProtectedPacketType:
			return &ProtectedPacket{
				BasePacket: &BasePacket{
					Header: ph,
					Frames: fs,
				},
				RTT: 0,
			}, nil
		default:
			// error
			return nil, nil
		}
	} else if firstByte&byte(ShortHeaderType) == byte(ShortHeaderType) {
		return &ProtectedPacket{
			BasePacket: &BasePacket{
				Header: ph,
				Frames: fs,
			},
			RTT: 1,
		}, nil
	}
	// error
	return nil, nil
}

// long header with type of 0x7F
type InitialPacket struct {
	*BasePacket
}

func NewInitialPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, payloadLen uint64) *InitialPacket {
	return &InitialPacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, payloadLen),
		},
	}

}

// long header with type of 0x7E
type RetryPacket struct {
	*BasePacket
}

func NewRetryPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, payloadLen uint64) *RetryPacket {
	return &RetryPacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(RetryPacketType, version, destConnID, srcConnID, packetNumber, payloadLen),
		},
	}
}

// long header with type of 0x7D
type HandshakePacket struct {
	*BasePacket
}

func NewHandshakePacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, payloadLen uint64) *HandshakePacket {
	return &HandshakePacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(HandshakePacketType, version, destConnID, srcConnID, packetNumber, payloadLen),
		},
	}
}

// long header with 0-RTT (type:0x7C)
// short header with 1-RTT
type ProtectedPacket struct {
	*BasePacket
	RTT byte
}

func NewProtectedPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, payloadLen uint64, rtt byte) *ProtectedPacket {
	var header PacketHeader
	if rtt == 0 {
		header = NewShortHeader(destConnID, packetNumber)
	} else if rtt == 1 {
		header = NewLongHeader(ZeroRTTProtectedPacketType, version, destConnID, srcConnID, packetNumber, payloadLen)
	} else {
		// error
	}

	return &ProtectedPacket{
		BasePacket: &BasePacket{
			Header: header,
		},
		RTT: rtt,
	}
}

/*
    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
   +-+-+-+-+-+-+-+-+
   |1|  Unused (7) |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                          Version (32)                         |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |DCIL(4)|SCIL(4)|
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |               Destination Connection ID (0/32..144)         ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                 Source Connection ID (0/32..144)            ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                    Supported Version 1 (32)                 ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version 2 (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
                                  ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version N (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
*/

// Version negotiation doesn't use long header, but have similar form
type VersionNegotiationPacket struct {<commit_msg>change packet creation interface to more realistic<commit_after>package quiclatest

import (
	"encoding/binary"

	"github.com/ami-GS/gQUIC/latest/qtype"
)

type Packet interface {
	GetWire() ([]byte, error)
	GetHeader() PacketHeader
	SetHeader(ph PacketHeader)
	GetFrames() []Frame
	SetFrames(fs []Frame)
}

func ParsePacket(data []byte) (packet Packet, idx int, err error) {
	version := binary.BigEndian.Uint32(data[1:5])
	if version == 0 {
		// VersionNegotiationPacket use all UDP datagram (doesn't have frame)
		return ParseVersionNegotiationPacket(data)
	}
	header, idx, err := PacketHeaderParserMap[PacketHeaderType(data[0]&0x80)](data) // ParseHeader
	if err != nil {
		return nil, 0, err
	}

	fs, idxTmp, err := ParseFrames(data[idx:])
	if err != nil {
		return nil, 0, err
	}

	packet, err = newPacket(data[0], header, fs)
	if err != nil {
		return nil, 0, err
	}

	return packet, idx + idxTmp, err
}

type BasePacket struct {
	Header     PacketHeader
	Frames     []Frame
	PaddingNum int
}

func (bp *BasePacket) GetHeader() PacketHeader {
	return bp.Header
}
func (bp *BasePacket) SetHeader(h PacketHeader) {
	bp.Header = h
}

func (bp *BasePacket) SetFrames(fs []Frame) {
	bp.Frames = fs
}

func (bp *BasePacket) GetFrames() []Frame {
	return bp.Frames
}

func (bp *BasePacket) GetWire() (wire []byte, err error) {
	hWire := bp.Header.GetWire()
	if err != nil {
		return nil, err
	}
	fsWire, err := GetFrameWires(bp.Frames)
	if err != nil {
		return nil, err
	}
	if bp.PaddingNum != 0 {
		fsWire = append(fsWire, make([]byte, bp.PaddingNum)...)
	}
	// TODO: protect for short header?
	return append(hWire, fsWire...), nil
}

func newPacket(firstByte byte, ph PacketHeader, fs []Frame) (Packet, error) {
	// TODO: needs frame type validation for each packet type
	if firstByte&byte(LongHeaderType) == byte(LongHeaderType) {
		payloadLen := 0
		for _, frame := range fs {
			payloadLen += frame.GetWireSize()
		}

		switch LongHeaderPacketType(firstByte & 0x7f) {
		case InitialPacketType:
			// TODO: check whether it is stream frame or not
			return &InitialPacket{
				BasePacket: &BasePacket{
					Header:     ph,
					Frames:     fs,
					PaddingNum: InitialPacketMinimumPayloadSize - payloadLen,
				},
			}, nil
		case RetryPacketType:
			return &RetryPacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case HandshakePacketType:
			return &HandshakePacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case ZeroRTTProtectedPacketType:
			return &ProtectedPacket{
				BasePacket: &BasePacket{
					Header: ph,
					Frames: fs,
				},
				RTT: 0,
			}, nil
		default:
			// error
			return nil, nil
		}
	} else if firstByte&byte(ShortHeaderType) == byte(ShortHeaderType) {
		return &ProtectedPacket{
			BasePacket: &BasePacket{
				Header: ph,
				Frames: fs,
			},
			RTT: 1,
		}, nil
	}
	// error
	return nil, nil
}

// long header with type of 0x7F
type InitialPacket struct {
	*BasePacket
}

const InitialPacketMinimumPayloadSize = 1200

func NewInitialPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, sFrame *StreamFrame) *InitialPacket {
	sFrameLen := sFrame.GetWireSize()
	var frames []Frame
	var lh *LongHeader
	paddingNum := 0
	if InitialPacketMinimumPayloadSize <= sFrameLen {
		// TODO: need to check when sFrameLen is over MTUIPv4
		lh = NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, uint64(sFrameLen))
		frames = []Frame{sFrame}
	} else {
		lh = NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, InitialPacketMinimumPayloadSize)
		frames = []Frame{sFrame}
		paddingNum = InitialPacketMinimumPayloadSize - sFrameLen
	}
	p := &InitialPacket{
		BasePacket: &BasePacket{
			Header:     lh,
			Frames:     frames,
			PaddingNum: paddingNum,
		},
	}
	return p
}

// long header with type of 0x7E
type RetryPacket struct {
	*BasePacket
}

func NewRetryPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, frames []Frame) *RetryPacket {
	if len(frames) < 2 {
		return nil
	}
	if _, ok := frames[0].(*StreamFrame); !ok {
		return nil
	}
	if _, ok := frames[1].(*AckFrame); !ok {
		return nil
	}
	for i := 2; i < len(frames); i++ {
		if _, ok := frames[i].(*PaddingFrame); !ok {
			return nil
		}
	}

	payloadLen := 0
	for i := 0; i < len(frames); i++ {
		payloadLen += frames[i].GetWireSize()
	}

	return &RetryPacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(RetryPacketType, version, destConnID, srcConnID, packetNumber, uint64(payloadLen)),
			Frames: frames,
		},
	}
}

// long header with type of 0x7D
type HandshakePacket struct {
	*BasePacket
}

func NewHandshakePacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, frames []Frame) *HandshakePacket {
	minimumReq := false
	payloadLen := 0
	for i := 0; i < len(frames); i++ {
		switch frames[i].(type) {
		case *StreamFrame:
			minimumReq = true
			payloadLen += frames[i].GetWireSize()
		case *AckFrame, *PathChallengeFrame /*or*/, *PathResponseFrame, *ConnectionCloseFrame:
			payloadLen += frames[i].GetWireSize()
		default:
			return nil
		}
	}
	if !minimumReq {
		return nil
	}

	return &HandshakePacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(HandshakePacketType, version, destConnID, srcConnID, packetNumber, uint64(payloadLen)),
			Frames: frames,
		},
	}
}

// long header with 0-RTT (type:0x7C)
// short header with 1-RTT
type ProtectedPacket struct {
	*BasePacket
	RTT byte
}

func NewProtectedPacket(version qtype.Version, key bool, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, rtt byte, frames []Frame) *ProtectedPacket {
	var header PacketHeader
	if rtt == 0 {
		payloadLen := 0
		for _, frame := range frames {
			payloadLen += frame.GetWireSize()
		}
		header = NewLongHeader(ZeroRTTProtectedPacketType, version, destConnID, srcConnID, packetNumber, uint64(payloadLen))
	} else if rtt == 1 {
		header = NewShortHeader(key, destConnID, packetNumber)
	} else {
		// error
	}

	return &ProtectedPacket{
		BasePacket: &BasePacket{
			Header: header,
			Frames: frames,
		},
		RTT: rtt,
	}

}

/*
    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
   +-+-+-+-+-+-+-+-+
   |1|  Unused (7) |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                          Version (32)                         |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |DCIL(4)|SCIL(4)|
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |               Destination Connection ID (0/32..144)         ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                 Source Connection ID (0/32..144)            ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                    Supported Version 1 (32)                 ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version 2 (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
                                  ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version N (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
*/

// Version negotiation doesn't use long header, but have similar form
type VersionNegotiationPacket struct {
{'fraction_non_alphanumeric': 0.19180715764566697, 'fraction_numerical': 0.015454007690032535, 'mean_word_length': 2.738253178551686, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5279008', 'n_tokens_mistral': 5279, 'n_tokens_neox': 4367, 'n_words': 1139}