session.go

package gem

import (
	"bufio"
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

const controlOpenTimeout = 5 * time.Second

type CoreSession struct {
	transport Transport
	cfg       Config
	sessionID uint64

	nextStreamID atomic.Uint32

	mu      sync.Mutex
	videoTx map[uint32]*videoSendState
	videoRx map[uint32]*videoReceiveState

	videoPipeline *videoPipeline
	filePipeline  *filePipeline
}

func NewSession(transport Transport, cfg Config) (Session, error) {
	return NewCoreSession(transport, cfg)
}

func NewCoreSession(transport Transport, cfg Config) (*CoreSession, error) {
	videoPipe, err := newVideoPipeline(cfg.Video)
	if err != nil {
		return nil, err
	}
	filePipe, err := newFilePipeline(cfg.File)
	if err != nil {
		videoPipe.Close()
		return nil, err
	}

	sessionID := cfg.SessionID
	if sessionID == 0 {
		sessionID, err = randomUint64()
		if err != nil {
			videoPipe.Close()
			filePipe.Close()
			return nil, err
		}
	}

	s := &CoreSession{
		transport:     transport,
		cfg:           cfg,
		sessionID:     sessionID,
		videoTx:       make(map[uint32]*videoSendState),
		videoRx:       make(map[uint32]*videoReceiveState),
		videoPipeline: videoPipe,
		filePipeline:  filePipe,
	}
	s.nextStreamID.Store(1)
	return s, nil
}

func (s *CoreSession) ID() uint64 {
	return s.sessionID
}

func (s *CoreSession) Transport() Transport {
	return s.transport
}

func (s *CoreSession) UpdateCompressionGraph(graph []byte) error {
	return s.videoPipeline.UpdateGraph(graph)
}

func (s *CoreSession) Close() error {
	if err := s.videoPipeline.Close(); err != nil {
		_ = s.filePipeline.Close()
		_ = s.transport.Close()
		return err
	}
	if err := s.filePipeline.Close(); err != nil {
		_ = s.transport.Close()
		return err
	}
	return s.transport.Close()
}

func (s *CoreSession) Run(ctx context.Context, hooks ReceiveHooks) error {
	errCh := make(chan error, 2)
	var started int

	started++
	go func() {
		errCh <- s.runUniStreamLoop(ctx, hooks)
	}()

	if s.transport.SupportsDatagrams() {
		started++
		go func() {
			errCh <- s.runDatagramLoop(ctx, hooks)
		}()
	}

	for started > 0 {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case err := <-errCh:
			started--
			if err != nil && !errors.Is(err, context.Canceled) {
				return err
			}
		}
	}

	return nil
}

func (s *CoreSession) runUniStreamLoop(ctx context.Context, hooks ReceiveHooks) error {
	for {
		stream, err := s.transport.AcceptStream(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			return err
		}

		go func(reader Stream) {
			_ = s.handleUniStream(ctx, reader, hooks)
		}(stream)
	}
}

func (s *CoreSession) handleUniStream(ctx context.Context, stream Stream, hooks ReceiveHooks) error {
	buffered := bufio.NewReader(stream)
	header, err := buffered.Peek(17)
	if err != nil {
		return err
	}
	if string(header[:4]) != streamMagic {
		return fmt.Errorf("unknown stream magic %q", string(header[:4]))
	}

	switch header[4] {
	case streamClassControl:
		return s.handleControlStream(buffered)
	case streamClassFile:
		return s.receiveFileStream(ctx, buffered, hooks)
	default:
		return fmt.Errorf("unsupported stream class %d", header[4])
	}
}

func (s *CoreSession) handleControlStream(reader io.Reader) error {
	packet, err := readControlStream(reader)
	if err != nil {
		return err
	}
	if packet.SessionID != s.sessionID {
		return nil
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	state, ok := s.videoTx[packet.StreamID]
	if !ok {
		return nil
	}

	switch packet.Type {
	case controlTypeLossRatio:
		lossRatio, err := decodeLossRatioBody(packet.Body)
		if err != nil {
			return err
		}
		state.PacketLossRatio = lossRatio
	case controlTypeFECFeedback:
		body, err := decodeFECFeedbackBody(packet.Body)
		if err != nil {
			return err
		}
		if body.FrameIndex <= state.LastFeedbackFrame {
			return nil
		}
		state.Redundancy = adjustRedundancy(state.Redundancy, body, s.cfg.Video.FEC)
		state.LastFeedbackFrame = body.FrameIndex
	case controlTypeTelemetry:
		body, err := decodeTelemetryBody(packet.Body)
		if err != nil {
			return err
		}
		if body.FrameIndex <= state.LastTelemetryFrame {
			return nil
		}
		state.Redundancy = adjustRedundancyWithTelemetry(state.Redundancy, body, s.cfg.Video.FEC)
		state.Jitter = time.Duration(body.JitterMicros) * time.Microsecond
		state.LastTelemetryFrame = body.FrameIndex
	case controlTypeStreamFin:
		delete(s.videoRx, packet.StreamID)
	}

	return nil
}

func (s *CoreSession) nextLogicalStreamID() uint32 {
	return s.nextStreamID.Add(1) - 1
}

func (s *CoreSession) sendControl(ctx context.Context, packet controlPacket) error {
	controlCtx, cancel := context.WithTimeout(ctx, controlOpenTimeout)
	defer cancel()

	stream, err := s.transport.OpenStream(controlCtx, StreamMetadata{
		SessionID:   packet.SessionID,
		StreamID:    packet.StreamID,
		Type:        StreamTypeControl,
		Reliability: ReliabilityReliable,
		ContentType: "application/x-gem-control",
		Connection:  s.transport.ConnectionContext(),
	})
	if err != nil {
		return err
	}
	defer stream.Close()

	return writeControlStream(stream, packet)
}

func randomUint64() (uint64, error) {
	var raw [8]byte
	if _, err := rand.Read(raw[:]); err != nil {
		return 0, err
	}
	value := binary.LittleEndian.Uint64(raw[:])
	if value == 0 {
		value = 1
	}
	return value, nil
}