packet.go

package gem

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math"
	"time"
)

const (
	streamMagic        = "GEM1"
	streamClassFile    = 0x01
	streamClassControl = 0x02

	controlTypeFECFeedback = 'F'
	controlTypeTelemetry   = 'T'
	controlTypeStreamFin   = 'E'
	controlTypeLossRatio   = 'B'

	fileFrameTypeBlock = 0x01
	fileFrameTypeFin   = 0x02

	videoGeometryHeaderSize = 24
	videoDatagramPrefixSize = 1 + 8 + videoGeometryHeaderSize
)

type videoDatagramHeader struct {
	SessionID     uint64
	StreamID      uint32
	FrameIndex    uint64
	SequenceID    uint32
	ProtectedSize uint32
	PacketSize    uint16
	Kind          VideoPayloadKind
	Codec         uint8
}

func encodeVideoEnvelope(frame VideoFrame) ([]byte, error) {
	var buf bytes.Buffer
	if err := binary.Write(&buf, binary.LittleEndian, frame.TimestampMicros); err != nil {
		return nil, err
	}
	var keyframe uint8
	if frame.Keyframe {
		keyframe = 1
	}
	if err := buf.WriteByte(keyframe); err != nil {
		return nil, err
	}
	if err := binary.Write(&buf, binary.LittleEndian, uint32(len(frame.Payload))); err != nil {
		return nil, err
	}
	if _, err := buf.Write(frame.Payload); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decodeVideoEnvelope(raw []byte) (VideoFrame, error) {
	if len(raw) < 13 {
		return VideoFrame{}, errors.New("video envelope too short")
	}
	frame := VideoFrame{}
	frame.TimestampMicros = binary.LittleEndian.Uint64(raw[:8])
	frame.Keyframe = raw[8] == 1
	size := binary.LittleEndian.Uint32(raw[9:13])
	if int(13+size) != len(raw) {
		return VideoFrame{}, errors.New("video envelope payload length mismatch")
	}
	frame.Payload = append([]byte(nil), raw[13:]...)
	return frame, nil
}

func encodeVideoDatagram(h videoDatagramHeader, payload []byte) ([]byte, error) {
	buf := make([]byte, videoDatagramPrefixSize+len(payload))
	buf[0] = AppDatagramTypeVideo
	binary.LittleEndian.PutUint64(buf[1:9], h.SessionID)
	binary.LittleEndian.PutUint32(buf[9:13], h.StreamID)
	binary.LittleEndian.PutUint64(buf[13:21], h.FrameIndex)
	binary.LittleEndian.PutUint32(buf[21:25], h.SequenceID)
	binary.LittleEndian.PutUint32(buf[25:29], h.ProtectedSize)
	binary.LittleEndian.PutUint16(buf[29:31], h.PacketSize)
	buf[31] = byte(h.Kind)
	buf[32] = h.Codec
	copy(buf[33:], payload)
	return buf, nil
}

func decodeVideoDatagram(raw []byte) (videoDatagramHeader, []byte, error) {
	if len(raw) < videoDatagramPrefixSize {
		return videoDatagramHeader{}, nil, errors.New("video datagram too short")
	}
	if raw[0] != AppDatagramTypeVideo {
		return videoDatagramHeader{}, nil, errors.New("unexpected datagram type")
	}
	h := videoDatagramHeader{
		SessionID:     binary.LittleEndian.Uint64(raw[1:9]),
		StreamID:      binary.LittleEndian.Uint32(raw[9:13]),
		FrameIndex:    binary.LittleEndian.Uint64(raw[13:21]),
		SequenceID:    binary.LittleEndian.Uint32(raw[21:25]),
		ProtectedSize: binary.LittleEndian.Uint32(raw[25:29]),
		PacketSize:    binary.LittleEndian.Uint16(raw[29:31]),
		Kind:          VideoPayloadKind(raw[31]),
		Codec:         raw[32],
	}
	if h.PacketSize == 0 {
		return videoDatagramHeader{}, nil, errors.New("video datagram packet size is zero")
	}
	payloadLen := minInt(len(raw)-videoDatagramPrefixSize, int(h.PacketSize))
	if payloadLen <= 0 {
		return videoDatagramHeader{}, nil, errors.New("video datagram missing payload")
	}
	return h, append([]byte(nil), raw[33:33+payloadLen]...), nil
}

type controlPacket struct {
	Type      byte
	SessionID uint64
	StreamID  uint32
	Body      []byte
}

func writeStreamPreamble(w io.Writer, class byte, sessionID uint64, streamID uint32) error {
	var fixed [17]byte
	copy(fixed[:4], []byte(streamMagic))
	fixed[4] = class
	binary.LittleEndian.PutUint64(fixed[5:13], sessionID)
	binary.LittleEndian.PutUint32(fixed[13:17], streamID)
	_, err := w.Write(fixed[:])
	return err
}

func readStreamPreamble(r io.Reader) (class byte, sessionID uint64, streamID uint32, err error) {
	var fixed [17]byte
	if _, err = io.ReadFull(r, fixed[:]); err != nil {
		return 0, 0, 0, err
	}
	if string(fixed[:4]) != streamMagic {
		return 0, 0, 0, errors.New("gem stream magic mismatch")
	}
	return fixed[4], binary.LittleEndian.Uint64(fixed[5:13]), binary.LittleEndian.Uint32(fixed[13:17]), nil
}

func encodeControlPacket(packet controlPacket) ([]byte, error) {
	var buf bytes.Buffer
	if err := buf.WriteByte(packet.Type); err != nil {
		return nil, err
	}
	if _, err := buf.Write(packet.Body); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

type fecFeedback struct {
	FrameIndex   uint64
	IdealPackets uint32
	UsedPackets  uint32
}

func encodeFECFeedbackBody(body fecFeedback) []byte {
	raw := make([]byte, 16)
	binary.LittleEndian.PutUint64(raw[:8], body.FrameIndex)
	binary.LittleEndian.PutUint32(raw[8:12], body.IdealPackets)
	binary.LittleEndian.PutUint32(raw[12:16], body.UsedPackets)
	return raw
}

func decodeFECFeedbackBody(raw []byte) (fecFeedback, error) {
	if len(raw) != 16 {
		return fecFeedback{}, errors.New("fec feedback size mismatch")
	}
	return fecFeedback{
		FrameIndex:   binary.LittleEndian.Uint64(raw[:8]),
		IdealPackets: binary.LittleEndian.Uint32(raw[8:12]),
		UsedPackets:  binary.LittleEndian.Uint32(raw[12:16]),
	}, nil
}

type telemetryPacket struct {
	FrameIndex    uint64
	PacketLossPPM uint32
	JitterMicros  uint32
}

func encodeTelemetryBody(body telemetryPacket) []byte {
	raw := make([]byte, 16)
	binary.LittleEndian.PutUint64(raw[:8], body.FrameIndex)
	binary.LittleEndian.PutUint32(raw[8:12], body.PacketLossPPM)
	binary.LittleEndian.PutUint32(raw[12:16], body.JitterMicros)
	return raw
}

func decodeTelemetryBody(raw []byte) (telemetryPacket, error) {
	if len(raw) != 16 {
		return telemetryPacket{}, errors.New("telemetry size mismatch")
	}
	return telemetryPacket{
		FrameIndex:    binary.LittleEndian.Uint64(raw[:8]),
		PacketLossPPM: binary.LittleEndian.Uint32(raw[8:12]),
		JitterMicros:  binary.LittleEndian.Uint32(raw[12:16]),
	}, nil
}

type streamFinPacket struct {
	FinalBytes  uint64
	FinalFrames uint64
}

func encodeStreamFinBody(body streamFinPacket) []byte {
	raw := make([]byte, 16)
	binary.LittleEndian.PutUint64(raw[:8], body.FinalBytes)
	binary.LittleEndian.PutUint64(raw[8:16], body.FinalFrames)
	return raw
}

func decodeStreamFinBody(raw []byte) (streamFinPacket, error) {
	if len(raw) != 16 {
		return streamFinPacket{}, errors.New("stream fin size mismatch")
	}
	return streamFinPacket{
		FinalBytes:  binary.LittleEndian.Uint64(raw[:8]),
		FinalFrames: binary.LittleEndian.Uint64(raw[8:16]),
	}, nil
}

func encodeLossRatioBody(ratio float32) []byte {
	raw := make([]byte, 4)
	binary.LittleEndian.PutUint32(raw, math.Float32bits(ratio))
	return raw
}

func decodeLossRatioBody(raw []byte) (float32, error) {
	if len(raw) != 4 {
		return 0, errors.New("loss ratio size mismatch")
	}
	return math.Float32frombits(binary.LittleEndian.Uint32(raw)), nil
}

func writeControlStream(w io.Writer, packet controlPacket) error {
	if err := writeStreamPreamble(w, streamClassControl, packet.SessionID, packet.StreamID); err != nil {
		return err
	}
	raw, err := encodeControlPacket(packet)
	if err != nil {
		return err
	}
	_, err = w.Write(raw)
	return err
}

func readControlStream(r io.Reader) (controlPacket, error) {
	class, sessionID, streamID, err := readStreamPreamble(r)
	if err != nil {
		return controlPacket{}, err
	}
	if class != streamClassControl {
		return controlPacket{}, errors.New("unexpected stream class")
	}
	raw, err := io.ReadAll(r)
	if err != nil {
		return controlPacket{}, err
	}
	if len(raw) == 0 {
		return controlPacket{}, errors.New("empty control payload")
	}
	return controlPacket{
		Type:      raw[0],
		SessionID: sessionID,
		StreamID:  streamID,
		Body:      append([]byte(nil), raw[1:]...),
	}, nil
}

type fileStreamDescriptor struct {
	Manifest          FileManifest
	BlockSize         uint32
	PacketPayloadSize uint16
	Redundancy        float32
	Compression       FileCompressionMode
}

func writeFileStreamHeader(w io.Writer, sessionID uint64, streamID uint32, desc fileStreamDescriptor) error {
	if err := writeStreamPreamble(w, streamClassFile, sessionID, streamID); err != nil {
		return err
	}
	nameBytes := []byte(desc.Manifest.Name)
	if len(nameBytes) > math.MaxUint16 {
		return fmt.Errorf("file name too long: %d bytes", len(nameBytes))
	}
	var fixed [35]byte
	binary.LittleEndian.PutUint64(fixed[:8], desc.Manifest.TraceID)
	binary.LittleEndian.PutUint64(fixed[8:16], uint64(desc.Manifest.DeclaredSize))
	binary.LittleEndian.PutUint64(fixed[16:24], uint64(desc.Manifest.Timestamp.Unix()))
	fixed[24] = byte(desc.Manifest.Semantics)
	binary.LittleEndian.PutUint32(fixed[25:29], desc.BlockSize)
	binary.LittleEndian.PutUint16(fixed[29:31], desc.PacketPayloadSize)
	binary.LittleEndian.PutUint32(fixed[31:35], math.Float32bits(desc.Redundancy))
	if _, err := w.Write(fixed[:]); err != nil {
		return err
	}
	if _, err := w.Write([]byte{byte(desc.Compression)}); err != nil {
		return err
	}
	if err := binary.Write(w, binary.LittleEndian, uint16(len(nameBytes))); err != nil {
		return err
	}
	_, err := w.Write(nameBytes)
	return err
}

func readFileStreamHeader(r io.Reader) (uint64, uint32, fileStreamDescriptor, error) {
	class, sessionID, streamID, err := readStreamPreamble(r)
	if err != nil {
		return 0, 0, fileStreamDescriptor{}, err
	}
	if class != streamClassFile {
		return 0, 0, fileStreamDescriptor{}, errors.New("unexpected stream class")
	}
	var fixed [35]byte
	if _, err := io.ReadFull(r, fixed[:]); err != nil {
		return 0, 0, fileStreamDescriptor{}, err
	}
	compression := make([]byte, 1)
	if _, err := io.ReadFull(r, compression); err != nil {
		return 0, 0, fileStreamDescriptor{}, err
	}
	var nameLen uint16
	if err := binary.Read(r, binary.LittleEndian, &nameLen); err != nil {
		return 0, 0, fileStreamDescriptor{}, err
	}
	nameBytes := make([]byte, nameLen)
	if _, err := io.ReadFull(r, nameBytes); err != nil {
		return 0, 0, fileStreamDescriptor{}, err
	}
	declaredSize := int64(binary.LittleEndian.Uint64(fixed[8:16]))
	desc := fileStreamDescriptor{
		Manifest: FileManifest{
			Name:         string(nameBytes),
			TraceID:      binary.LittleEndian.Uint64(fixed[:8]),
			DeclaredSize: declaredSize,
			Timestamp:    time.Unix(int64(binary.LittleEndian.Uint64(fixed[16:24])), 0),
			Semantics:    StreamSemantics(fixed[24]),
		},
		BlockSize:         binary.LittleEndian.Uint32(fixed[25:29]),
		PacketPayloadSize: binary.LittleEndian.Uint16(fixed[29:31]),
		Redundancy:        math.Float32frombits(binary.LittleEndian.Uint32(fixed[31:35])),
		Compression:       FileCompressionMode(compression[0]),
	}
	if declaredSize < 0 {
		desc.Manifest.DeclaredSize = -1
	}
	return sessionID, streamID, desc, nil
}

type fileBlockFrameHeader struct {
	BlockID        uint64
	ProtectedSize  uint32
	PacketSize     uint16
	PacketCount    uint32
	UncompressedSz uint32
}

func writeFileBlockFrameHeader(w io.Writer, h fileBlockFrameHeader) error {
	var fixed [1 + 8 + 4 + 2 + 4 + 4]byte
	fixed[0] = fileFrameTypeBlock
	binary.LittleEndian.PutUint64(fixed[1:9], h.BlockID)
	binary.LittleEndian.PutUint32(fixed[9:13], h.ProtectedSize)
	binary.LittleEndian.PutUint16(fixed[13:15], h.PacketSize)
	binary.LittleEndian.PutUint32(fixed[15:19], h.PacketCount)
	binary.LittleEndian.PutUint32(fixed[19:23], h.UncompressedSz)
	_, err := w.Write(fixed[:])
	return err
}

func readFileBlockFrameHeader(r io.Reader) (fileBlockFrameHeader, error) {
	var fixed [23]byte
	if _, err := io.ReadFull(r, fixed[:]); err != nil {
		return fileBlockFrameHeader{}, err
	}
	if fixed[0] != fileFrameTypeBlock {
		return fileBlockFrameHeader{}, fmt.Errorf("unexpected file block frame type: %d", fixed[0])
	}
	return fileBlockFrameHeader{
		BlockID:        binary.LittleEndian.Uint64(fixed[1:9]),
		ProtectedSize:  binary.LittleEndian.Uint32(fixed[9:13]),
		PacketSize:     binary.LittleEndian.Uint16(fixed[13:15]),
		PacketCount:    binary.LittleEndian.Uint32(fixed[15:19]),
		UncompressedSz: binary.LittleEndian.Uint32(fixed[19:23]),
	}, nil
}

func writeDroplet(w io.Writer, seqID uint32, payload []byte) error {
	var fixed [6]byte
	binary.LittleEndian.PutUint32(fixed[:4], seqID)
	binary.LittleEndian.PutUint16(fixed[4:6], uint16(len(payload)))
	if _, err := w.Write(fixed[:]); err != nil {
		return err
	}
	_, err := w.Write(payload)
	return err
}

func readDroplet(r io.Reader) (uint32, []byte, error) {
	var fixed [6]byte
	if _, err := io.ReadFull(r, fixed[:]); err != nil {
		return 0, nil, err
	}
	seqID := binary.LittleEndian.Uint32(fixed[:4])
	size := binary.LittleEndian.Uint16(fixed[4:6])
	payload := make([]byte, size)
	if _, err := io.ReadFull(r, payload); err != nil {
		return 0, nil, err
	}
	return seqID, payload, nil
}

func writeFileFinFrame(w io.Writer, finalBytes, finalBlocks uint64) error {
	var fixed [17]byte
	fixed[0] = fileFrameTypeFin
	binary.LittleEndian.PutUint64(fixed[1:9], finalBytes)
	binary.LittleEndian.PutUint64(fixed[9:17], finalBlocks)
	_, err := w.Write(fixed[:])
	return err
}

func readFileFinFrame(r io.Reader) (streamFinPacket, error) {
	var fixed [17]byte
	if _, err := io.ReadFull(r, fixed[:]); err != nil {
		return streamFinPacket{}, err
	}
	if fixed[0] != fileFrameTypeFin {
		return streamFinPacket{}, fmt.Errorf("unexpected file fin frame type: %d", fixed[0])
	}
	return streamFinPacket{
		FinalBytes:  binary.LittleEndian.Uint64(fixed[1:9]),
		FinalFrames: binary.LittleEndian.Uint64(fixed[9:17]),
	}, nil
}

func minInt(a, b int) int {
	if a < b {
		return a
	}
	return b
}