RemoteDesk/transport/hysteria2-agent/media.go

package main

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

const (
	videoDatagramMagic       = "RDV1"
	videoDatagramVersion     = 1
	videoDatagramHeaderBytes = 58
	videoDatagramPayloadMax  = 1200
	videoDatagramMaxBytes    = videoDatagramHeaderBytes + videoDatagramPayloadMax
	videoFrameMaxBytes       = 16 * 1024 * 1024
	videoFrameMaxFragments   = 4096
	videoReassemblyMaxFrames = 8
	videoReassemblyTimeout   = 100 * time.Millisecond
)

const (
	audioDatagramMagic      = "RDA1"
	audioDatagramVersion    = 1
	audioDatagramHeaderSize = 40
	audioPacketMaxBytes     = 4 * 1024
	audioDatagramMaxBytes   = audioDatagramHeaderSize + audioPacketMaxBytes
)

const (
	videoFlagConfig        uint16 = 1 << 0
	videoFlagKeyFrame      uint16 = 1 << 1
	videoFlagFirstFragment uint16 = 1 << 2
	videoFlagLastFragment  uint16 = 1 << 3
	videoFlagDiscontinuity uint16 = 1 << 4
)

const audioFlagDiscontinuity uint16 = 1 << 0

// VideoAccessUnit is a complete encoded frame before transport fragmentation.
// The Hysteria2 relay must treat the resulting datagrams as opaque bytes.
type VideoAccessUnit struct {
	StreamID   uint32
	Generation uint64
	Sequence   uint64
	FrameID    uint64
	PTS        uint64
	DTS        uint64
	Flags      uint16
	Payload    []byte
}

// AudioPacket is one independently decodable Opus packet. Audio is never
// reassembled with video and is forwarded on its own priority path.
type AudioPacket struct {
	StreamID   uint32
	Generation uint64
	Sequence   uint64
	PTS        uint64
	DurationMS uint16
	Flags      uint16
	Payload    []byte
}

func encodeAudioDatagram(packet AudioPacket) ([]byte, error) {
	if err := validateAudioPacket(packet); err != nil {
		return nil, err
	}
	output := make([]byte, audioDatagramHeaderSize+len(packet.Payload))
	copy(output, audioDatagramMagic)
	output[4] = audioDatagramVersion
	binary.LittleEndian.PutUint16(output[6:8], audioDatagramHeaderSize)
	binary.LittleEndian.PutUint32(output[8:12], packet.StreamID)
	binary.LittleEndian.PutUint64(output[12:20], packet.Generation)
	binary.LittleEndian.PutUint64(output[20:28], packet.Sequence)
	binary.LittleEndian.PutUint64(output[28:36], packet.PTS)
	binary.LittleEndian.PutUint16(output[36:38], packet.DurationMS)
	binary.LittleEndian.PutUint16(output[38:40], packet.Flags)
	copy(output[audioDatagramHeaderSize:], packet.Payload)
	return output, nil
}

func decodeAudioDatagram(input []byte) (AudioPacket, error) {
	if len(input) < audioDatagramHeaderSize || len(input) > audioDatagramMaxBytes {
		return AudioPacket{}, errors.New("audio datagram size is invalid")
	}
	if string(input[:4]) != audioDatagramMagic || input[4] != audioDatagramVersion || input[5] != 0 {
		return AudioPacket{}, errors.New("audio datagram version is invalid")
	}
	headerBytes := int(binary.LittleEndian.Uint16(input[6:8]))
	if headerBytes != audioDatagramHeaderSize || headerBytes > len(input) {
		return AudioPacket{}, errors.New("audio datagram header is invalid")
	}
	packet := AudioPacket{
		StreamID:   binary.LittleEndian.Uint32(input[8:12]),
		Generation: binary.LittleEndian.Uint64(input[12:20]),
		Sequence:   binary.LittleEndian.Uint64(input[20:28]),
		PTS:        binary.LittleEndian.Uint64(input[28:36]),
		DurationMS: binary.LittleEndian.Uint16(input[36:38]),
		Flags:      binary.LittleEndian.Uint16(input[38:40]),
		Payload:    input[headerBytes:],
	}
	if err := validateAudioPacket(packet); err != nil {
		return AudioPacket{}, err
	}
	return packet, nil
}

func validateAudioPacket(packet AudioPacket) error {
	if packet.StreamID == 0 || packet.Generation == 0 || packet.Sequence == 0 {
		return errors.New("audio packet identity is invalid")
	}
	if packet.DurationMS != 10 && packet.DurationMS != 20 && packet.DurationMS != 40 && packet.DurationMS != 60 {
		return errors.New("audio packet duration is invalid")
	}
	if len(packet.Payload) == 0 || len(packet.Payload) > audioPacketMaxBytes {
		return errors.New("audio packet payload size is invalid")
	}
	if packet.Flags & ^audioFlagDiscontinuity != 0 {
		return errors.New("audio packet flags are invalid")
	}
	return nil
}

type videoDatagram struct {
	StreamID      uint32
	Generation    uint64
	Sequence      uint64
	FrameID       uint64
	PTS           uint64
	DTS           uint64
	FragmentID    uint16
	FragmentCount uint16
	Flags         uint16
	Payload       []byte
}

func fragmentVideoAccessUnit(unit VideoAccessUnit) ([][]byte, error) {
	if unit.StreamID == 0 || unit.Generation == 0 {
		return nil, errors.New("video stream identity is invalid")
	}
	if len(unit.Payload) == 0 || len(unit.Payload) > videoFrameMaxBytes {
		return nil, errors.New("video access unit size is invalid")
	}
	fragmentCount := (len(unit.Payload) + videoDatagramPayloadMax - 1) / videoDatagramPayloadMax
	if fragmentCount > videoFrameMaxFragments {
		return nil, errors.New("video access unit has too many fragments")
	}

	datagrams := make([][]byte, 0, fragmentCount)
	for fragmentID, offset := 0, 0; offset < len(unit.Payload); fragmentID++ {
		end := min(offset+videoDatagramPayloadMax, len(unit.Payload))
		flags := unit.Flags
		if fragmentID == 0 {
			flags |= videoFlagFirstFragment
		}
		if end == len(unit.Payload) {
			flags |= videoFlagLastFragment
		}
		packet := videoDatagram{
			StreamID:      unit.StreamID,
			Generation:    unit.Generation,
			Sequence:      unit.Sequence,
			FrameID:       unit.FrameID,
			PTS:           unit.PTS,
			DTS:           unit.DTS,
			FragmentID:    uint16(fragmentID),
			FragmentCount: uint16(fragmentCount),
			Flags:         flags,
			Payload:       unit.Payload[offset:end],
		}
		datagrams = append(datagrams, encodeVideoDatagram(packet))
		offset = end
	}
	return datagrams, nil
}

func encodeVideoDatagram(packet videoDatagram) []byte {
	output := make([]byte, videoDatagramHeaderBytes+len(packet.Payload))
	copy(output, videoDatagramMagic)
	output[4] = videoDatagramVersion
	binary.LittleEndian.PutUint16(output[6:8], videoDatagramHeaderBytes)
	binary.LittleEndian.PutUint32(output[8:12], packet.StreamID)
	binary.LittleEndian.PutUint64(output[12:20], packet.Generation)
	binary.LittleEndian.PutUint64(output[20:28], packet.Sequence)
	binary.LittleEndian.PutUint64(output[28:36], packet.FrameID)
	binary.LittleEndian.PutUint64(output[36:44], packet.PTS)
	binary.LittleEndian.PutUint64(output[44:52], packet.DTS)
	binary.LittleEndian.PutUint16(output[52:54], packet.FragmentID)
	binary.LittleEndian.PutUint16(output[54:56], packet.FragmentCount)
	binary.LittleEndian.PutUint16(output[56:58], packet.Flags)
	copy(output[videoDatagramHeaderBytes:], packet.Payload)
	return output
}

func decodeVideoDatagram(input []byte) (videoDatagram, error) {
	if len(input) < videoDatagramHeaderBytes || len(input) > videoDatagramHeaderBytes+videoDatagramPayloadMax {
		return videoDatagram{}, errors.New("video datagram size is invalid")
	}
	if string(input[:4]) != videoDatagramMagic || input[4] != videoDatagramVersion {
		return videoDatagram{}, errors.New("video datagram version is invalid")
	}
	headerBytes := int(binary.LittleEndian.Uint16(input[6:8]))
	if headerBytes != videoDatagramHeaderBytes || headerBytes > len(input) {
		return videoDatagram{}, errors.New("video datagram header is invalid")
	}
	packet := videoDatagram{
		StreamID:      binary.LittleEndian.Uint32(input[8:12]),
		Generation:    binary.LittleEndian.Uint64(input[12:20]),
		Sequence:      binary.LittleEndian.Uint64(input[20:28]),
		FrameID:       binary.LittleEndian.Uint64(input[28:36]),
		PTS:           binary.LittleEndian.Uint64(input[36:44]),
		DTS:           binary.LittleEndian.Uint64(input[44:52]),
		FragmentID:    binary.LittleEndian.Uint16(input[52:54]),
		FragmentCount: binary.LittleEndian.Uint16(input[54:56]),
		Flags:         binary.LittleEndian.Uint16(input[56:58]),
		Payload:       input[headerBytes:],
	}
	if packet.StreamID == 0 || packet.Generation == 0 || packet.FragmentCount == 0 ||
		packet.FragmentCount > videoFrameMaxFragments || packet.FragmentID >= packet.FragmentCount ||
		len(packet.Payload) == 0 || len(packet.Payload) > videoDatagramPayloadMax {
		return videoDatagram{}, errors.New("video datagram fields are invalid")
	}
	if packet.Flags&videoFlagFirstFragment != 0 != (packet.FragmentID == 0) ||
		packet.Flags&videoFlagLastFragment != 0 != (packet.FragmentID+1 == packet.FragmentCount) {
		return videoDatagram{}, errors.New("video datagram fragment flags are invalid")
	}
	return packet, nil
}

type videoFrameKey struct {
	streamID   uint32
	generation uint64
	frameID    uint64
}

type videoFrameAssembly struct {
	created       time.Time
	sequence      uint64
	pts           uint64
	dts           uint64
	flags         uint16
	fragmentCount uint16
	fragments     map[uint16][]byte
	bytes         int
}

// videoReassembler is intentionally bounded. It never waits for a missing
// datagram past the deadline and never retransmits an expired video frame.
type videoReassembler struct {
	frames map[videoFrameKey]*videoFrameAssembly
}

func newVideoReassembler() *videoReassembler {
	return &videoReassembler{frames: make(map[videoFrameKey]*videoFrameAssembly)}
}

func (r *videoReassembler) push(input []byte, now time.Time) (*VideoAccessUnit, error) {
	packet, err := decodeVideoDatagram(input)
	if err != nil {
		return nil, err
	}
	r.expire(now)
	key := videoFrameKey{streamID: packet.StreamID, generation: packet.Generation, frameID: packet.FrameID}
	assembly := r.frames[key]
	if assembly == nil {
		if len(r.frames) >= videoReassemblyMaxFrames {
			r.dropOldest()
		}
		assembly = &videoFrameAssembly{
			created:       now,
			sequence:      packet.Sequence,
			pts:           packet.PTS,
			dts:           packet.DTS,
			flags:         packet.Flags &^ (videoFlagFirstFragment | videoFlagLastFragment),
			fragmentCount: packet.FragmentCount,
			fragments:     make(map[uint16][]byte, packet.FragmentCount),
		}
		r.frames[key] = assembly
	} else if assembly.fragmentCount != packet.FragmentCount || assembly.sequence != packet.Sequence {
		delete(r.frames, key)
		return nil, errors.New("video frame metadata changed during reassembly")
	}
	if _, exists := assembly.fragments[packet.FragmentID]; !exists {
		assembly.fragments[packet.FragmentID] = append([]byte(nil), packet.Payload...)
		assembly.bytes += len(packet.Payload)
	}
	if assembly.bytes > videoFrameMaxBytes || len(assembly.fragments) != int(assembly.fragmentCount) {
		return nil, nil
	}
	payload := make([]byte, 0, assembly.bytes)
	for index := uint16(0); index < assembly.fragmentCount; index++ {
		fragment, ok := assembly.fragments[index]
		if !ok {
			return nil, nil
		}
		payload = append(payload, fragment...)
	}
	delete(r.frames, key)
	return &VideoAccessUnit{
		StreamID:   packet.StreamID,
		Generation: packet.Generation,
		Sequence:   assembly.sequence,
		FrameID:    packet.FrameID,
		PTS:        assembly.pts,
		DTS:        assembly.dts,
		Flags:      assembly.flags | (packet.Flags & (videoFlagConfig | videoFlagKeyFrame | videoFlagDiscontinuity)),
		Payload:    payload,
	}, nil
}

func (r *videoReassembler) expire(now time.Time) {
	for key, assembly := range r.frames {
		if now.Sub(assembly.created) >= videoReassemblyTimeout {
			delete(r.frames, key)
		}
	}
}

func (r *videoReassembler) dropOldest() {
	var oldest videoFrameKey
	var oldestTime time.Time
	for key, assembly := range r.frames {
		if oldestTime.IsZero() || assembly.created.Before(oldestTime) {
			oldest, oldestTime = key, assembly.created
		}
	}
	if !oldestTime.IsZero() {
		delete(r.frames, oldest)
	}
}

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

func validateVideoDatagramSize(input []byte) error {
	if len(input) > videoDatagramMaxBytes {
		return fmt.Errorf("video datagram exceeds %d bytes", videoDatagramMaxBytes)
	}
	return nil
}

func validateAudioDatagramSize(input []byte) error {
	if len(input) > audioDatagramMaxBytes {
		return fmt.Errorf("audio datagram exceeds %d bytes", audioDatagramMaxBytes)
	}
	return nil
}