RemoteDesk/transport/hysteria2-agent/media_ring.go

package main

import (
	"context"
	"errors"
	"fmt"
)

// EncodedRingRecord is the cross-process contract produced by the Rust
// Windows agent. PTS is the Media Foundation 100-ns media timeline.
type EncodedRingRecord struct {
	Sequence uint64
	PTS      uint64
	Kind     uint32
	KeyFrame bool
	Payload  []byte
}

// EncodedRing is intentionally a single-consumer interface. The producer may
// drop a video record when this queue is full; the consumer must never block
// input or audio while polling it.
type EncodedRing interface {
	Pop() (*EncodedRingRecord, error)
	Close() error
}

// AudioRing is a separate consumer contract even though its mmap slot layout
// matches the video ring. Keeping the interfaces distinct prevents audio from
// accidentally sharing video backpressure or drop policy.
type AudioRing interface {
	Pop() (*EncodedRingRecord, error)
	Close() error
}

func validateRingRecord(record *EncodedRingRecord) error {
	if record == nil || record.Sequence == 0 || len(record.Payload) == 0 || len(record.Payload) > videoFrameMaxBytes {
		return errors.New("encoded ring record is invalid")
	}
	return nil
}

// pumpEncodedRing drains the ring into RDV1 datagrams. The callback is the
// already-authenticated Hysteria2 UDP stream; this function never retries a
// failed video packet.
func pumpEncodedRing(ctx context.Context, ring EncodedRing, streamID uint32, generation uint64, send func([]byte) error) error {
	if streamID == 0 || generation == 0 {
		return errors.New("encoded ring stream identity is invalid")
	}
	if ring == nil || send == nil {
		return errors.New("encoded ring pump arguments are invalid")
	}
	ticker := newMediaPollTicker()
	defer ticker.Stop()
	for {
		drained := false
		for {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
			}
			record, err := ring.Pop()
			if err != nil {
				return fmt.Errorf("read encoded media ring: %w", err)
			}
			if record == nil {
				break
			}
			drained = true
			if err := validateRingRecord(record); err != nil {
				return err
			}
			flags := uint16(0)
			if record.KeyFrame {
				flags |= videoFlagKeyFrame
			}
			datagrams, err := fragmentVideoAccessUnit(VideoAccessUnit{
				StreamID: streamID, Generation: generation, Sequence: record.Sequence,
				FrameID: record.Sequence, PTS: record.PTS, DTS: record.PTS,
				Flags: flags, Payload: record.Payload,
			})
			if err != nil {
				return fmt.Errorf("fragment encoded media: %w", err)
			}
			for _, datagram := range datagrams {
				if err := send(datagram); err != nil {
					return fmt.Errorf("send encoded video datagram: %w", err)
				}
			}
		}
		if drained {
			continue
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
		}
	}
}

// pumpAudioRing sends independent Opus packets as RDA1 datagrams. Audio does
// not wait for video and does not use video frame reassembly.
func pumpAudioRing(ctx context.Context, ring AudioRing, streamID uint32, generation uint64, send func([]byte) error) error {
	if streamID == 0 || generation == 0 {
		return errors.New("audio ring stream identity is invalid")
	}
	if ring == nil || send == nil {
		return errors.New("audio ring pump arguments are invalid")
	}
	ticker := newMediaPollTicker()
	defer ticker.Stop()
	for {
		drained := false
		for {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
			}
			record, err := ring.Pop()
			if err != nil {
				return fmt.Errorf("read audio media ring: %w", err)
			}
			if record == nil {
				break
			}
			drained = true
			if record.Sequence == 0 || len(record.Payload) == 0 || len(record.Payload) > audioPacketMaxBytes {
				return errors.New("audio ring record is invalid")
			}
			flags := uint16(0)
			if record.Kind == 2 {
				flags |= audioFlagDiscontinuity
			}
			datagram, err := encodeAudioDatagram(AudioPacket{
				StreamID: streamID, Generation: generation, Sequence: record.Sequence,
				PTS: record.PTS, DurationMS: 20, Flags: flags, Payload: record.Payload,
			})
			if err != nil {
				return fmt.Errorf("encode audio datagram: %w", err)
			}
			if err := send(datagram); err != nil {
				return fmt.Errorf("send audio datagram: %w", err)
			}
		}
		if drained {
			continue
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
		}
	}
}