VerseVDI-Data-Plane/gateway/frame.go

package gateway

import (
	"encoding/binary"
	"errors"
	"fmt"
)

const (
	frameV1HeaderSize     = 21
	frameV2HeaderSize     = 23
	frameV1PayloadSize    = 1179
	frameV2PayloadSize    = 1177
	maxV1FragmentCount    = 16
	maxV2FragmentCount    = 891
	maxCompleteFrameBytes = 1 << 20
	maxFrameSize          = 1 << 16
	frameHeaderSize       = frameV2HeaderSize
)

const (
	ChannelControl = byte(1)
	ChannelAck     = byte(2)
	ChannelText    = byte(3)
	ChannelVideo   = byte(10)
	ChannelAudio   = byte(11)
	ChannelInput   = byte(12)
)

var (
	ErrFrameTruncated       = errors.New("gateway frame truncated")
	ErrFrameMagic           = errors.New("gateway frame magic mismatch")
	ErrFrameVersion         = errors.New("gateway frame version unsupported")
	ErrFrameChannel         = errors.New("gateway frame channel unsupported")
	ErrFrameFlags           = errors.New("gateway frame flags unsupported")
	ErrFrameFragment        = errors.New("gateway frame fragment invalid")
	ErrFrameLength          = errors.New("gateway frame length mismatch")
	ErrFramePayloadLimit    = errors.New("gateway frame payload exceeds channel limit")
	ErrFrameSize            = errors.New("gateway frame exceeds size limit")
	ErrFrameFragmentedLimit = errors.New("gateway payload requires too many fragments")
)

type Frame struct {
	Version       byte
	Channel       byte
	Flags         byte
	Sequence      uint32
	TimestampMS   uint64
	FragmentIndex uint16
	FragmentCount uint16
	Payload       []byte
}

func channelLimit(version, channel byte) (int, bool) {
	if version == 2 {
		switch channel {
		case ChannelVideo, ChannelAudio:
			return frameV2PayloadSize, true
		default:
			return 0, false
		}
	}
	switch channel {
	case ChannelControl:
		return 1024, true
	case ChannelAck:
		return 2048, true
	case ChannelText:
		return 65515, true
	case ChannelVideo, ChannelAudio, ChannelInput:
		return frameV1PayloadSize, true
	default:
		return 0, false
	}
}

func EncodeFrame(frame Frame) ([]byte, error) {
	version := frame.Version
	if version == 0 {
		version = 1
	}
	if version != 1 && version != 2 {
		return nil, ErrFrameVersion
	}
	limit, ok := channelLimit(version, frame.Channel)
	if !ok {
		return nil, ErrFrameChannel
	}
	if frame.Flags != 0 {
		return nil, ErrFrameFlags
	}
	maxFragments := uint16(maxV1FragmentCount)
	headerSize := frameV1HeaderSize
	if version == 2 {
		maxFragments = maxV2FragmentCount
		headerSize = frameV2HeaderSize
	}
	if frame.FragmentCount == 0 || frame.FragmentCount > maxFragments || frame.FragmentIndex >= frame.FragmentCount {
		return nil, ErrFrameFragment
	}
	if len(frame.Payload) > limit {
		return nil, ErrFramePayloadLimit
	}
	if len(frame.Payload) > 1<<16-headerSize {
		return nil, ErrFrameSize
	}
	encoded := make([]byte, headerSize+len(frame.Payload))
	encoded[0], encoded[1], encoded[2], encoded[3], encoded[4] = 'V', 'D', version, frame.Channel, frame.Flags
	binary.BigEndian.PutUint32(encoded[5:9], frame.Sequence)
	binary.BigEndian.PutUint64(encoded[9:17], frame.TimestampMS)
	if version == 1 {
		encoded[17], encoded[18] = byte(frame.FragmentIndex), byte(frame.FragmentCount)
		binary.BigEndian.PutUint16(encoded[19:21], uint16(len(frame.Payload)))
	} else {
		binary.BigEndian.PutUint16(encoded[17:19], frame.FragmentIndex)
		binary.BigEndian.PutUint16(encoded[19:21], frame.FragmentCount)
		binary.BigEndian.PutUint16(encoded[21:23], uint16(len(frame.Payload)))
	}
	copy(encoded[headerSize:], frame.Payload)
	return encoded, nil
}

func DecodeFrame(raw []byte) (Frame, error) {
	if len(raw) < 3 {
		return Frame{}, ErrFrameTruncated
	}
	if raw[0] != 'V' || raw[1] != 'D' {
		return Frame{}, ErrFrameMagic
	}
	version := raw[2]
	if version != 1 && version != 2 {
		return Frame{}, ErrFrameVersion
	}
	headerSize := frameV1HeaderSize
	maxFragments := uint16(maxV1FragmentCount)
	if version == 2 {
		headerSize = frameV2HeaderSize
		maxFragments = maxV2FragmentCount
	}
	if len(raw) < headerSize {
		return Frame{}, ErrFrameTruncated
	}
	if version == 1 && len(raw) > 1<<16 || version == 2 && len(raw) > 1200 {
		return Frame{}, ErrFrameSize
	}
	limit, ok := channelLimit(version, raw[3])
	if !ok {
		return Frame{}, ErrFrameChannel
	}
	if raw[4] != 0 {
		return Frame{}, ErrFrameFlags
	}
	var fragmentIndex, fragmentCount uint16
	payloadOffset := 19
	if version == 1 {
		fragmentIndex, fragmentCount = uint16(raw[17]), uint16(raw[18])
	} else {
		fragmentIndex = binary.BigEndian.Uint16(raw[17:19])
		fragmentCount = binary.BigEndian.Uint16(raw[19:21])
		payloadOffset = 21
	}
	if fragmentCount == 0 || fragmentCount > maxFragments || fragmentIndex >= fragmentCount {
		return Frame{}, ErrFrameFragment
	}
	payloadLength := int(binary.BigEndian.Uint16(raw[payloadOffset : payloadOffset+2]))
	if payloadLength > limit {
		return Frame{}, ErrFramePayloadLimit
	}
	if len(raw) != headerSize+payloadLength {
		return Frame{}, ErrFrameLength
	}
	return Frame{
		Version:       version,
		Channel:       raw[3],
		Flags:         raw[4],
		Sequence:      binary.BigEndian.Uint32(raw[5:9]),
		TimestampMS:   binary.BigEndian.Uint64(raw[9:17]),
		FragmentIndex: fragmentIndex,
		FragmentCount: fragmentCount,
		Payload:       append([]byte(nil), raw[headerSize:]...),
	}, nil
}

func FragmentPayload(channel byte, sequence uint32, timestampMS uint64, payload []byte) ([]Frame, error) {
	version := byte(1)
	limit := frameV1PayloadSize
	maxFragments := maxV1FragmentCount
	if channel == ChannelVideo || channel == ChannelAudio {
		version = 2
		limit = frameV2PayloadSize
		maxFragments = maxV2FragmentCount
	}
	if _, ok := channelLimit(version, channel); !ok {
		return nil, ErrFrameChannel
	}
	if len(payload) > maxCompleteFrameBytes {
		return nil, ErrFrameFragmentedLimit
	}
	count := (len(payload) + limit - 1) / limit
	if count == 0 {
		count = 1
	}
	if count > maxFragments {
		return nil, ErrFrameFragmentedLimit
	}
	frames := make([]Frame, 0, count)
	for index := 0; index < count; index++ {
		start := index * limit
		end := start + limit
		if end > len(payload) {
			end = len(payload)
		}
		frames = append(frames, Frame{
			Version:       version,
			Channel:       channel,
			Sequence:      sequence,
			TimestampMS:   timestampMS,
			FragmentIndex: uint16(index),
			FragmentCount: uint16(count),
			Payload:       append([]byte(nil), payload[start:end]...),
		})
	}
	return frames, nil
}

func ValidateFrame(raw []byte) error {
	_, err := DecodeFrame(raw)
	return err
}

func FrameError(channel byte, err error) error {
	return fmt.Errorf("channel %d: %w", channel, err)
}