RemoteDesk/transport/hysteria2-agent/media_ring_windows.go

//go:build windows

package main

import (
	"encoding/binary"
	"errors"
	"fmt"
	"os"
	"sync/atomic"
	"unsafe"

	"golang.org/x/sys/windows"
)

const (
	ringMagic        = "RDMRING1"
	ringVersion      = uint32(1)
	ringHeaderBytes  = 40
	slotHeaderBytes  = 24
	maxRingSlots     = uint32(256)
	maxRingSlotBytes = uint32(16 * 1024 * 1024)
)

type mappedEncodedRing struct {
	file      *os.File
	path      string
	mapping   windows.Handle
	address   uintptr
	data      []byte
	slotCount uint32
	slotBytes uint32
}

func openEncodedRing(path string) (EncodedRing, error) {
	file, err := os.OpenFile(path, os.O_RDWR, 0)
	if err != nil {
		return nil, fmt.Errorf("open encoded ring: %w", err)
	}
	closeOnError := true
	defer func() {
		if closeOnError {
			_ = file.Close()
		}
	}()
	stat, err := file.Stat()
	if err != nil {
		return nil, fmt.Errorf("stat encoded ring: %w", err)
	}
	length := stat.Size()
	if length < ringHeaderBytes {
		return nil, errors.New("encoded ring is shorter than its header")
	}
	mapping, err := windows.CreateFileMapping(windows.Handle(file.Fd()), nil, windows.PAGE_READWRITE, 0, 0, nil)
	if err != nil {
		return nil, fmt.Errorf("map encoded ring file: %w", err)
	}
	address, err := windows.MapViewOfFile(mapping, windows.FILE_MAP_READ|windows.FILE_MAP_WRITE, 0, 0, uintptr(length))
	if err != nil {
		_ = windows.CloseHandle(mapping)
		return nil, fmt.Errorf("map encoded ring view: %w", err)
	}
	if uintptr(length) > uintptr(^uint(0)>>1) {
		_ = windows.UnmapViewOfFile(address)
		_ = windows.CloseHandle(mapping)
		return nil, errors.New("encoded ring is too large for this process")
	}
	data := unsafe.Slice((*byte)(unsafe.Pointer(address)), int(length))
	if string(data[:8]) != ringMagic || binary.LittleEndian.Uint32(data[8:12]) != ringVersion {
		_ = windows.UnmapViewOfFile(address)
		_ = windows.CloseHandle(mapping)
		return nil, errors.New("encoded ring header magic or version is invalid")
	}
	slots := binary.LittleEndian.Uint32(data[12:16])
	slotBytes := binary.LittleEndian.Uint32(data[16:20])
	if slots < 2 || slots > maxRingSlots || slots&(slots-1) != 0 || slotBytes == 0 || slotBytes > maxRingSlotBytes {
		_ = windows.UnmapViewOfFile(address)
		_ = windows.CloseHandle(mapping)
		return nil, errors.New("encoded ring geometry is invalid")
	}
	expected := uint64(ringHeaderBytes) + uint64(slots)*uint64(slotHeaderBytes+slotBytes)
	if expected != uint64(length) {
		_ = windows.UnmapViewOfFile(address)
		_ = windows.CloseHandle(mapping)
		return nil, errors.New("encoded ring file size does not match its header")
	}
	closeOnError = false
	return &mappedEncodedRing{file: file, path: path, mapping: mapping, address: address, data: data, slotCount: slots, slotBytes: slotBytes}, nil
}

func openAudioRing(path string) (AudioRing, error) {
	ring, err := openEncodedRing(path)
	if err != nil {
		return nil, err
	}
	return ring, nil
}

func (r *mappedEncodedRing) Pop() (*EncodedRingRecord, error) {
	read := atomic.LoadUint64((*uint64)(unsafe.Pointer(&r.data[32])))
	write := atomic.LoadUint64((*uint64)(unsafe.Pointer(&r.data[24])))
	if read == write {
		return nil, nil
	}
	offset := ringHeaderBytes + (read%uint64(r.slotCount))*uint64(slotHeaderBytes+r.slotBytes)
	start := int(offset)
	if start+slotHeaderBytes > len(r.data) {
		return nil, errors.New("encoded ring slot header is out of bounds")
	}
	length := binary.LittleEndian.Uint32(r.data[start+20 : start+24])
	if length == 0 || length > r.slotBytes || uint64(start+slotHeaderBytes)+uint64(length) > uint64(len(r.data)) {
		return nil, errors.New("encoded ring record length is invalid")
	}
	payload := append([]byte(nil), r.data[start+slotHeaderBytes:start+slotHeaderBytes+int(length)]...)
	record := &EncodedRingRecord{
		Sequence: binary.LittleEndian.Uint64(r.data[start : start+8]),
		PTS:      binary.LittleEndian.Uint64(r.data[start+8 : start+16]),
		Kind:     binary.LittleEndian.Uint32(r.data[start+16 : start+20]),
		Payload:  payload,
	}
	record.KeyFrame = record.Kind == 2
	if record.Sequence == 0 {
		return nil, errors.New("encoded ring record sequence is invalid")
	}
	atomic.StoreUint64((*uint64)(unsafe.Pointer(&r.data[32])), read+1)
	return record, nil
}

func (r *mappedEncodedRing) Close() error {
	if r.address != 0 {
		_ = windows.UnmapViewOfFile(r.address)
		r.address = 0
	}
	if r.mapping != 0 {
		_ = windows.CloseHandle(r.mapping)
		r.mapping = 0
	}
	if r.file != nil {
		err := r.file.Close()
		if removeErr := os.Remove(r.path); err == nil && removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) {
			err = removeErr
		}
		return err
	}
	return nil
}