VerseVDI-Data-Plane/gateway/transport.go

package gateway

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"slices"
	"sync"
	"sync/atomic"
	"time"

	"github.com/quic-go/quic-go"

	protocol "git.sechmachine.io.vn/sechmachine/VerseVDI-Protocol/gen/go/protocol"
)

const (
	defaultHelloLimit    = 16 * 1024
	defaultHelloTimeout  = 10 * time.Second
	defaultControlLimit  = 128 * 1024
	clientControlBacklog = 64
	terminalAckTimeout   = 2 * time.Second
	applicationError     = quic.ApplicationErrorCode(0x100)
	controlFlowID        = "control.ack.v1"
	inputFlowID          = "input.sequenced.v1"
	clipboardFlowID      = "clipboard.text.v1"
)

var (
	ErrAdmissionRejected = errors.New("gateway admission rejected")
	ErrGatewayDraining   = errors.New("gateway draining")
	ErrGatewayTLS        = errors.New("gateway requires TLS 1.3 client authentication")
	ErrAuthorityExpired  = errors.New("gateway authority expired")
)

type Admission interface {
	Admit(context.Context, protocol.TunnelAdmissionRequest) (protocol.SessionAuthority, error)
	ProviderWork(context.Context, protocol.SessionAuthority) (protocol.ProviderSessionWork, error)
	Release(context.Context, protocol.SessionAuthority) error
}

type AdmissionFunc func(context.Context, protocol.TunnelAdmissionRequest) (protocol.SessionAuthority, error)

func (f AdmissionFunc) Admit(ctx context.Context, request protocol.TunnelAdmissionRequest) (protocol.SessionAuthority, error) {
	return f(ctx, request)
}

func (AdmissionFunc) Release(context.Context, protocol.SessionAuthority) error { return nil }

func (AdmissionFunc) ProviderWork(context.Context, protocol.SessionAuthority) (protocol.ProviderSessionWork, error) {
	return protocol.ProviderSessionWork{}, ErrAdmissionRejected
}

type ProviderStateReporter interface {
	ReportProviderState(context.Context, protocol.ProviderState) error
}

type ClipboardAuditReporter interface {
	ReportClipboardAudit(context.Context, protocol.GatewayClipboardAudit) error
}

type ServerConfig struct {
	ListenAddress          string
	TLSConfig              *tls.Config
	QUICConfig             *quic.Config
	GatewayID              string
	Features               []string
	Capabilities           protocol.CapabilityProfile
	ProviderCapabilities   protocol.CapabilityProfile
	Admission              Admission
	ProviderStateReporter  ProviderStateReporter
	ClipboardAuditReporter ClipboardAuditReporter
	Provider               Provider
	ProviderProfile        string
	ProviderIdentity       string
	PacerKbps              int64
	mediaObserver          func(mediaTimingObservation)
}

type mediaTimingObservation struct {
	QueueDelay      time.Duration
	ProcessingDelay time.Duration
	PacingDelay     time.Duration
}

type Server struct {
	listener     *quic.Listener
	config       ServerConfig
	metrics      *Metrics
	pacer        *fairPacer
	mu           sync.Mutex
	sessions     map[*gatewaySession]struct{}
	connections  map[*quic.Conn]struct{}
	helloTimeout time.Duration
	draining     atomic.Bool
	closed       atomic.Bool
	closeOnce    sync.Once
	workers      sync.WaitGroup
}

func NewServer(config ServerConfig) (*Server, error) {
	if config.ListenAddress == "" {
		config.ListenAddress = "127.0.0.1:0"
	}
	if config.GatewayID == "" || config.Admission == nil || config.Provider == nil {
		return nil, errors.New("gateway id, admission, and provider are required")
	}
	if err := validateServerTLS(config.TLSConfig); err != nil {
		return nil, err
	}
	if capabilityProfileUnset(config.Capabilities) {
		config.Capabilities = DefaultCapabilities()
	}
	if capabilityProfileUnset(config.ProviderCapabilities) {
		config.ProviderCapabilities = DefaultCapabilities()
	}
	if config.Capabilities.Validate() != nil || config.ProviderCapabilities.Validate() != nil {
		return nil, ErrNoCapabilityOverlap
	}
	if config.ProviderProfile == "" {
		config.ProviderProfile = ProviderProfileApollo
	}
	if config.PacerKbps < 0 {
		return nil, errors.New("negative pacing limit")
	}
	tlsConfig := config.TLSConfig.Clone()
	if len(tlsConfig.NextProtos) == 0 {
		tlsConfig.NextProtos = []string{"versevdi-gateway-v1"}
	}
	quicConfig := &quic.Config{EnableDatagrams: true, MaxIdleTimeout: 30 * time.Second, MaxIncomingStreams: 2, MaxIncomingUniStreams: 2}
	if config.QUICConfig != nil {
		quicConfig = config.QUICConfig.Clone()
		quicConfig.EnableDatagrams = true
	}
	listener, err := quic.ListenAddr(config.ListenAddress, tlsConfig, quicConfig)
	if err != nil {
		return nil, err
	}
	return &Server{listener: listener, config: config, metrics: &Metrics{}, pacer: newFairPacer(config.PacerKbps), sessions: make(map[*gatewaySession]struct{}), connections: make(map[*quic.Conn]struct{}), helloTimeout: defaultHelloTimeout}, nil
}

func validateServerTLS(config *tls.Config) error {
	if config == nil || config.MinVersion < tls.VersionTLS13 || config.ClientAuth != tls.RequireAndVerifyClientCert || config.ClientCAs == nil || len(config.Certificates) == 0 {
		return ErrGatewayTLS
	}
	return nil
}

func (s *Server) Addr() net.Addr           { return s.listener.Addr() }
func (s *Server) Metrics() MetricsSnapshot { return s.metrics.Snapshot() }
func (s *Server) Draining() bool           { return s.draining.Load() }

func (s *Server) BeginDrain() {
	if s.draining.CompareAndSwap(false, true) {
		s.metrics.DrainTransitions.Add(1)
	}
}

func (s *Server) Serve(ctx context.Context) error {
	if s.closed.Load() {
		return net.ErrClosed
	}
	go func() {
		<-ctx.Done()
		_ = s.Close()
	}()
	for {
		connection, err := s.listener.Accept(ctx)
		if err != nil {
			if s.closed.Load() || errors.Is(err, context.Canceled) || errors.Is(err, net.ErrClosed) {
				return nil
			}
			return err
		}
		s.mu.Lock()
		if s.closed.Load() {
			s.mu.Unlock()
			_ = connection.CloseWithError(applicationError, "server closed")
			continue
		}
		s.connections[connection] = struct{}{}
		s.workers.Add(1)
		s.mu.Unlock()
		go func() {
			defer s.workers.Done()
			defer func() {
				s.mu.Lock()
				delete(s.connections, connection)
				s.mu.Unlock()
			}()
			s.handleConnection(ctx, connection)
		}()
	}
}

func (s *Server) Close() error {
	var err error
	s.closeOnce.Do(func() {
		s.BeginDrain()
		s.mu.Lock()
		s.closed.Store(true)
		sessions := make([]*gatewaySession, 0, len(s.sessions))
		for session := range s.sessions {
			sessions = append(sessions, session)
		}
		connections := make([]*quic.Conn, 0, len(s.connections))
		for connection := range s.connections {
			connections = append(connections, connection)
		}
		s.mu.Unlock()
		err = s.listener.Close()
		for _, session := range sessions {
			session.cancel()
		}
		for _, connection := range connections {
			_ = connection.CloseWithError(applicationError, "server closed")
		}
	})
	s.workers.Wait()
	return err
}

func (s *Server) handleConnection(parent context.Context, connection *quic.Conn) {
	defer connection.CloseWithError(applicationError, "connection closed")
	helloDeadline := time.Now().Add(s.helloTimeout)
	ctx, cancel := context.WithDeadline(parent, helloDeadline)
	defer cancel()
	stream, err := connection.AcceptStream(ctx)
	if err != nil {
		return
	}
	if err := stream.SetDeadline(helloDeadline); err != nil {
		return
	}
	writeError := func(code string, err error, retryable bool) {
		responseDeadline := time.Now().Add(time.Second)
		if stream.SetDeadline(responseDeadline) != nil {
			return
		}
		if writeStableError(stream, code, err, retryable) == nil && stream.Close() == nil {
			responseCtx, responseCancel := context.WithDeadline(context.Background(), responseDeadline)
			defer responseCancel()
			select {
			case <-connection.Context().Done():
			case <-responseCtx.Done():
			}
		}
	}
	requestBytes, err := readWire(stream, defaultHelloLimit)
	if err != nil {
		writeError("invalid_hello", err, false)
		return
	}
	request, err := protocol.DecodeTunnelAdmissionRequest(requestBytes)
	if err != nil {
		writeError("invalid_hello", err, false)
		return
	}
	if s.Draining() {
		writeError("gateway_draining", ErrGatewayDraining, true)
		return
	}
	if request.GatewayID != s.config.GatewayID {
		writeError("wrong_gateway", ErrAdmissionRejected, false)
		return
	}
	authority, err := s.config.Admission.Admit(ctx, request)
	if err != nil {
		s.metrics.AdmissionRejects.Add(1)
		writeError(stableAdmissionCode(err), err, false)
		return
	}
	if s.Draining() {
		_ = s.config.Admission.Release(context.Background(), authority)
		writeError("gateway_draining", ErrGatewayDraining, false)
		return
	}
	if err := s.validateAuthority(authority, request); err != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		writeError("invalid_authority", err, false)
		return
	}
	work, err := s.config.Admission.ProviderWork(ctx, authority)
	if err != nil || s.validateProviderWork(work, authority) != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		writeError("provider_work_unavailable", ErrAdmissionRejected, false)
		return
	}
	selected, err := IntersectCapabilities(s.config.Capabilities, s.config.ProviderCapabilities, request.Capabilities, authority.Capabilities)
	if err != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		s.metrics.AdmissionRejects.Add(1)
		writeError("no_capability_overlap", err, false)
		return
	}
	selected, err = selectApolloPolicyCapabilities(work.StreamPolicy, selected)
	if err != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		s.metrics.AdmissionRejects.Add(1)
		writeError("no_capability_overlap", err, false)
		return
	}
	clipboard, err := newClipboardGate(work.ClipboardPolicy, time.Now)
	if err != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		writeError("provider_work_unavailable", ErrAdmissionRejected, false)
		return
	}
	if (work.ClipboardPolicy.ClientToProviderEnabled || work.ClipboardPolicy.ProviderToClientEnabled) && s.config.ClipboardAuditReporter == nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		writeError("clipboard_audit_unavailable", ErrAdmissionRejected, false)
		return
	}
	if err := s.reportProviderState(ctx, protocol.ProviderState{Version: "1", SessionID: request.SessionID, State: ProviderStateStarting, CleanupPending: false, Channels: []string{"video", "audio", "input", "feedback"}}); err != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		writeError("provider_state_unavailable", err, false)
		return
	}
	providerSession, err := s.config.Provider.Start(ctx, LaunchRequest{SessionID: request.SessionID, Capabilities: selected, ProviderProfile: authority.ProviderProfile, ProviderIdentity: work.ProviderIdentity, ProviderWork: work})
	if err != nil {
		s.metrics.ProviderErrors.Add(1)
		_ = s.reportProviderState(context.Background(), protocol.ProviderState{Version: "1", SessionID: request.SessionID, State: ProviderStateFailed, CleanupPending: false, Channels: []string{"video", "audio", "input", "feedback"}})
		_ = s.config.Admission.Release(context.Background(), authority)
		writeError(stableProviderCode(err), err, false)
		return
	}
	if err := s.reportProviderState(ctx, providerSession.State()); err != nil {
		_ = providerSession.ReleaseAll(context.Background())
		_ = providerSession.Terminate(context.Background())
		_ = s.config.Admission.Release(context.Background(), authority)
		writeError("provider_state_unavailable", err, false)
		return
	}
	clientAuthority := protocol.ClientSessionAuthority{
		Version: authority.Version, SessionID: authority.SessionID, GatewayID: authority.GatewayID, Audience: authority.Audience,
		ReconnectSequence: authority.ReconnectSequence, ExpiresAt: authority.ExpiresAt, Capabilities: selected,
	}
	authorityBytes, err := protocol.EncodeClientSessionAuthority(clientAuthority)
	if err != nil || writeWire(stream, authorityBytes, defaultHelloLimit) != nil || stream.SetDeadline(time.Time{}) != nil {
		_ = providerSession.ReleaseAll(context.Background())
		_ = providerSession.Terminate(context.Background())
		_ = s.config.Admission.Release(context.Background(), authority)
		return
	}
	session := newGatewaySession(s, connection, stream, request, authority, providerSession, clipboard)
	s.addSession(session)
	s.metrics.ActiveSessions.Add(1)
	s.metrics.AdmittedSessions.Add(1)
	if authority.ReconnectSequence > 0 {
		s.metrics.Reconnects.Add(1)
	}
	defer func() {
		s.removeSession(session)
		s.metrics.ActiveSessions.Add(-1)
	}()
	session.run()
}

func (s *Server) reportProviderState(ctx context.Context, state protocol.ProviderState) error {
	if err := state.Validate(); err != nil {
		return err
	}
	s.metrics.observeProviderState(state.State)
	if s.config.ProviderStateReporter == nil {
		return nil
	}
	return s.config.ProviderStateReporter.ReportProviderState(ctx, state)
}

func (s *Server) validateAuthority(authority protocol.SessionAuthority, request protocol.TunnelAdmissionRequest) error {
	if err := authority.Validate(); err != nil {
		return err
	}
	if authority.SessionID != request.SessionID || authority.GatewayID != request.GatewayID || authority.Audience != request.Audience || authority.ProviderProfile != s.config.ProviderProfile {
		return ErrAdmissionRejected
	}
	expires, err := time.Parse(time.RFC3339Nano, authority.ExpiresAt)
	if err != nil || !time.Now().Before(expires) {
		return ErrAuthorityExpired
	}
	return nil
}

func (s *Server) validateProviderWork(work protocol.ProviderSessionWork, authority protocol.SessionAuthority) error {
	if err := work.Validate(); err != nil {
		return err
	}
	if work.SessionID != authority.SessionID || work.GatewayID != authority.GatewayID ||
		work.ReconnectSequence != authority.ReconnectSequence || work.ExpiresAt != authority.ExpiresAt ||
		work.ProviderProfile != authority.ProviderProfile || !apolloPolicyMatchesCapabilities(work.StreamPolicy, authority.Capabilities) {
		return ErrAdmissionRejected
	}
	return nil
}

func apolloPolicyMatchesCapabilities(policy protocol.ProviderStreamPolicy, capabilities protocol.CapabilityProfile) bool {
	if validateApolloStreamPolicy(policy) != nil || capabilities.Audio != "encoded" {
		return false
	}
	required := apolloPolicyProfile(policy)
	return required != "" && slices.Contains(capabilities.ClientDecode, required)
}

func selectApolloPolicyCapabilities(policy protocol.ProviderStreamPolicy, capabilities protocol.CapabilityProfile) (protocol.CapabilityProfile, error) {
	if !apolloPolicyMatchesCapabilities(policy, capabilities) {
		return protocol.CapabilityProfile{}, ErrNoCapabilityOverlap
	}
	capabilities.ClientDecode = []string{apolloPolicyProfile(policy)}
	return capabilities, nil
}

func apolloPolicyProfile(policy protocol.ProviderStreamPolicy) string {
	switch policy.Codec {
	case "H264":
		return "h264-opus"
	case "HEVC":
		return "hevc-opus"
	default:
		return ""
	}
}

func (s *Server) addSession(session *gatewaySession) {
	s.mu.Lock()
	s.sessions[session] = struct{}{}
	s.mu.Unlock()
}

func (s *Server) removeSession(session *gatewaySession) {
	s.mu.Lock()
	delete(s.sessions, session)
	s.mu.Unlock()
}

type gatewaySession struct {
	server         *Server
	connection     *quic.Conn
	control        *quic.Stream
	request        protocol.TunnelAdmissionRequest
	authority      protocol.SessionAuthority
	provider       ProviderSession
	clipboard      *clipboardGate
	ctx            context.Context
	cancel         context.CancelFunc
	cleanupOnce    sync.Once
	inputMu        sync.Mutex
	controlWriteMu sync.Mutex
	outputMu       sync.Mutex
	terminalMu     sync.Mutex
	pressed        map[string]struct{}
	sequence       atomic.Uint32
	mediaDrops     uint64
	mediaQuiesced  bool
	terminalSent   atomic.Bool
	terminalAwait  bool
	terminalAck    chan struct{}
	endReason      error
	result         chan error
}

func newGatewaySession(server *Server, connection *quic.Conn, control *quic.Stream, request protocol.TunnelAdmissionRequest, authority protocol.SessionAuthority, provider ProviderSession, clipboard *clipboardGate) *gatewaySession {
	ctx, cancel := context.WithCancel(context.Background())
	return &gatewaySession{server: server, connection: connection, control: control, request: request, authority: authority, provider: provider, clipboard: clipboard, ctx: ctx, cancel: cancel, pressed: make(map[string]struct{}), terminalAck: make(chan struct{}, 1), result: make(chan error, 3)}
}

func (s *gatewaySession) run() {
	defer s.cleanup()
	deadline, err := time.Parse(time.RFC3339Nano, s.authority.ExpiresAt)
	if err != nil {
		return
	}
	if deadline.Before(time.Now()) {
		return
	}
	timer := time.NewTimer(time.Until(deadline))
	defer timer.Stop()
	go s.controlLoop()
	go s.datagramLoop()
	go s.mediaLoop()
	go s.providerEventLoop()
	go s.providerTelemetryLoop()
	if s.clipboard != nil && s.clipboard.policy.ProviderToClientEnabled {
		go s.clipboardLoop()
	}
	select {
	case <-timer.C:
		s.server.metrics.InputRejected.Add(1)
	case <-s.ctx.Done():
	case s.endReason = <-s.result:
	}
	s.cancel()
	if s.terminalSent.Load() {
		s.cleanup()
	}
}

func (s *gatewaySession) providerEventLoop() {
	events := s.provider.Events()
	for events != nil {
		select {
		case <-s.ctx.Done():
			return
		case event, ok := <-events:
			if !ok {
				return
			}
			terminal := event.Kind == ProviderEventTerminated || event.Kind == ProviderEventDisconnected
			if terminal {
				s.outputMu.Lock()
				s.mediaQuiesced = true
				s.terminalMu.Lock()
			}
			payload, err := EncodeProviderEvent(event)
			if err == nil {
				err = s.sendControl(s.sequence.Add(1), payload)
			}
			if terminal {
				s.terminalAwait = err == nil
				s.terminalMu.Unlock()
				s.outputMu.Unlock()
			}
			if err != nil {
				s.result <- err
				return
			}
			if terminal {
				s.terminalSent.Store(true)
				timer := time.NewTimer(terminalAckTimeout)
				select {
				case <-s.terminalAck:
				case <-timer.C:
					s.terminalMu.Lock()
					s.terminalAwait = false
					s.terminalMu.Unlock()
					s.result <- context.DeadlineExceeded
					return
				case <-s.ctx.Done():
					timer.Stop()
					s.terminalMu.Lock()
					s.terminalAwait = false
					s.terminalMu.Unlock()
					return
				}
				timer.Stop()
			}
			switch event.Kind {
			case ProviderEventTerminated:
				s.result <- ErrProviderTerminated
				return
			case ProviderEventDisconnected:
				s.result <- ErrProviderDisconnected
				return
			}
		}
	}
}

func (s *gatewaySession) clipboardLoop() {
	ticker := time.NewTicker(500 * time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case <-s.ctx.Done():
			return
		case <-ticker.C:
			text, err := s.provider.ReadClipboard(s.ctx)
			if err != nil {
				if auditErr := s.reportClipboardAudit("provider_to_client", "rejected", clipboardAuditTextBytes(text), clipboardAuditReason(err)); auditErr != nil {
					s.result <- auditErr
					return
				}
				s.result <- err
				return
			}
			value, suppress, err := s.clipboard.fromProvider(text)
			if err != nil {
				if auditErr := s.reportClipboardAudit("provider_to_client", "rejected", clipboardAuditTextBytes(text), clipboardAuditReason(err)); auditErr != nil {
					s.result <- auditErr
					return
				}
				s.result <- err
				return
			}
			if suppress {
				if err := s.reportClipboardAudit("provider_to_client", "suppressed", clipboardAuditTextBytes(text), "loop"); err != nil {
					s.result <- err
					return
				}
				continue
			}
			if err := s.sendClipboard(value); err != nil {
				s.result <- err
				return
			}
			if err := s.reportClipboardAudit(value.Direction, "forwarded", clipboardAuditTextBytes(value.Text), "forwarded"); err != nil {
				s.result <- err
				return
			}
		}
	}
}

func (s *gatewaySession) providerTelemetryLoop() {
	s.observeProviderTelemetry()
	ticker := time.NewTicker(500 * time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case <-s.ctx.Done():
			return
		case <-ticker.C:
			s.observeProviderTelemetry()
		}
	}
}

func (s *gatewaySession) observeProviderTelemetry() {
	telemetry := s.provider.Telemetry()
	if telemetry.MediaDrops >= s.mediaDrops {
		s.server.metrics.MediaDrops.Add(telemetry.MediaDrops - s.mediaDrops)
		s.mediaDrops = telemetry.MediaDrops
	}
	s.server.metrics.observeProviderTelemetry(telemetry)
}

func (s *gatewaySession) controlLoop() {
	for {
		data, err := readWire(s.control, defaultControlLimit)
		if err != nil {
			s.result <- err
			return
		}
		frame, err := protocol.DecodeChannelFrame(data)
		if err != nil {
			s.result <- err
			return
		}
		payload, err := base64.StdEncoding.DecodeString(frame.Payload)
		if err != nil || len(payload) > maxFrameSize {
			s.result <- ErrFramePayloadLimit
			return
		}
		switch frame.FlowID {
		case controlFlowID:
			sequence, sequenceErr := channelSequence(frame.Sequence)
			if sequenceErr != nil {
				s.result <- sequenceErr
				return
			}
			if err := s.handleControl(payload, sequence); err != nil {
				s.result <- err
				return
			}
		case inputFlowID:
			sequence, sequenceErr := channelSequence(frame.Sequence)
			if sequenceErr != nil {
				s.result <- sequenceErr
				return
			}
			if err := s.handleInput(payload, sequence); err != nil {
				s.result <- err
				return
			}
		case clipboardFlowID:
			value, decodeErr := protocol.DecodeGatewayClipboardText(payload)
			if decodeErr != nil {
				s.result <- ErrProviderMalformed
				return
			}
			if err := s.handleClipboard(value); err != nil {
				s.result <- err
				return
			}
		default:
			s.result <- ErrFrameChannel
			return
		}
	}
}

func (s *gatewaySession) datagramLoop() {
	for {
		data, err := s.connection.ReceiveDatagram(s.ctx)
		if err != nil {
			s.result <- err
			return
		}
		frame, err := DecodeFrame(data)
		if err != nil {
			s.result <- err
			return
		}
		switch frame.Channel {
		case ChannelInput:
			if err := s.handleInput(frame.Payload, frame.Sequence); err != nil {
				s.result <- err
				return
			}
		case ChannelText:
			s.result <- ErrFrameChannel
			return
		default:
			s.result <- ErrFrameChannel
			return
		}
	}
}

func (s *gatewaySession) mediaLoop() {
	video, audio := s.provider.Video(), s.provider.Audio()
	for video != nil || audio != nil {
		select {
		case <-s.ctx.Done():
			return
		case media, ok := <-video:
			if !ok {
				video = nil
				continue
			}
			if media.expiry != nil {
				media.expiry.Stop()
			}
			media.releaseQueue()
			if !media.EnqueuedAt.IsZero() && time.Since(media.EnqueuedAt) > nativeApolloVideoQueueLatency {
				s.server.metrics.MediaDrops.Add(1)
				continue
			}
			if err := s.forwardMedia(ChannelVideo, media); err != nil {
				s.result <- err
				return
			}
		case media, ok := <-audio:
			if !ok {
				audio = nil
				continue
			}
			if err := s.forwardMedia(ChannelAudio, media); err != nil {
				s.result <- err
				return
			}
		}
	}
	s.result <- ErrProviderDisconnected
}

func (s *gatewaySession) forwardMedia(channel byte, media ProviderMedia) error {
	s.outputMu.Lock()
	defer s.outputMu.Unlock()
	state := s.provider.State().State
	if s.mediaQuiesced || state == ProviderStateTerminated || state == ProviderStateDisconnected {
		s.mediaQuiesced = true
		return nil
	}
	return s.sendMedia(channel, media)
}

func (s *gatewaySession) sendMedia(channel byte, media ProviderMedia) error {
	dequeuedAt := time.Now()
	if media.EnqueuedAt.IsZero() || media.EnqueuedAt.After(dequeuedAt) {
		media.EnqueuedAt = dequeuedAt
	}
	if media.ReceivedAt.IsZero() || media.ReceivedAt.After(media.EnqueuedAt) {
		media.ReceivedAt = media.EnqueuedAt
	}
	processingStarted := time.Now()
	frames, err := FragmentPayload(channel, s.sequence.Add(1), uint64(time.Now().UnixMilli()), media.Payload)
	if err != nil {
		return err
	}
	var pacingDelay time.Duration
	for _, frame := range frames {
		encoded, err := EncodeFrame(frame)
		if err != nil {
			return err
		}
		pacingStarted := time.Now()
		if err := s.server.pacer.wait(s.ctx, s.authority.SessionID, len(encoded)); err != nil {
			return err
		}
		pacingDelay += time.Since(pacingStarted)
		if err := s.connection.SendDatagram(encoded); err != nil {
			return err
		}
		s.server.metrics.MediaPackets.Add(1)
		s.server.metrics.MediaBytes.Add(uint64(len(encoded)))
	}
	queueDelay := dequeuedAt.Sub(media.EnqueuedAt)
	processingDelay := max(media.EnqueuedAt.Sub(media.ReceivedAt)+time.Since(processingStarted)-pacingDelay, 0)
	s.server.metrics.QueueDelayNanos.Add(uint64(queueDelay))
	s.server.metrics.ProcessingDelayNanos.Add(uint64(processingDelay))
	s.server.metrics.PacingDelayNanos.Add(uint64(pacingDelay))
	s.server.metrics.ProcessingSamples.Add(1)
	if s.server.config.mediaObserver != nil {
		s.server.config.mediaObserver(mediaTimingObservation{
			QueueDelay: queueDelay, ProcessingDelay: processingDelay, PacingDelay: pacingDelay,
		})
	}
	return nil
}

func (s *gatewaySession) sendControl(sequence uint32, payload []byte) error {
	if len(payload) > 1024 {
		return ErrFramePayloadLimit
	}
	frame := protocol.ChannelFrame{Version: "1", FlowID: controlFlowID, Sequence: int64(sequence), Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
	encoded, err := protocol.EncodeChannelFrame(frame)
	if err != nil {
		return err
	}
	s.controlWriteMu.Lock()
	defer s.controlWriteMu.Unlock()
	return writeWire(s.control, encoded, defaultControlLimit)
}

func (s *gatewaySession) sendClipboard(value protocol.GatewayClipboardText) error {
	payload, err := protocol.EncodeGatewayClipboardText(value)
	if err != nil {
		return err
	}
	frame := protocol.ChannelFrame{Version: "1", FlowID: clipboardFlowID, Sequence: int64(s.sequence.Add(1)), Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
	encoded, err := protocol.EncodeChannelFrame(frame)
	if err != nil {
		return err
	}
	s.controlWriteMu.Lock()
	defer s.controlWriteMu.Unlock()
	return writeWire(s.control, encoded, defaultControlLimit)
}

func (s *gatewaySession) handleControl(payload []byte, sequence uint32) error {
	feedback, err := DecodeClientFeedback(payload)
	if err != nil {
		return err
	}
	if feedback.Kind == FeedbackTerminalReceipt {
		s.terminalMu.Lock()
		defer s.terminalMu.Unlock()
		if !s.terminalAwait {
			return ErrProviderMalformed
		}
		s.terminalAwait = false
		select {
		case s.terminalAck <- struct{}{}:
			return nil
		default:
			return ErrProviderMalformed
		}
	}
	feedback.Sequence = sequence
	return s.provider.Feedback(s.ctx, feedback)
}

func (s *gatewaySession) handleClipboard(value protocol.GatewayClipboardText) error {
	if s.clipboard == nil {
		return ErrClipboardDenied
	}
	suppress, err := s.clipboard.fromClient(value)
	if err != nil {
		if auditErr := s.reportClipboardAudit(value.Direction, "rejected", clipboardAuditTextBytes(value.Text), clipboardAuditReason(err)); auditErr != nil {
			return auditErr
		}
		return err
	}
	if suppress {
		return s.reportClipboardAudit(value.Direction, "suppressed", clipboardAuditTextBytes(value.Text), "loop")
	}
	if err := s.provider.WriteClipboard(s.ctx, value.Text); err != nil {
		s.clipboard.retractClient(value)
		if auditErr := s.reportClipboardAudit(value.Direction, "rejected", clipboardAuditTextBytes(value.Text), "provider"); auditErr != nil {
			return auditErr
		}
		return err
	}
	return s.reportClipboardAudit(value.Direction, "forwarded", clipboardAuditTextBytes(value.Text), "forwarded")
}

func (s *gatewaySession) reportClipboardAudit(direction, outcome string, textBytes int, reason string) error {
	if s.server.config.ClipboardAuditReporter == nil {
		return ErrClipboardDenied
	}
	ctx, cancel := context.WithTimeout(s.ctx, 5*time.Second)
	defer cancel()
	return s.server.config.ClipboardAuditReporter.ReportClipboardAudit(ctx, protocol.GatewayClipboardAudit{
		Version: "1", SessionID: s.authority.SessionID, Direction: direction, Outcome: outcome, TextBytes: int64(textBytes), Reason: reason,
	})
}

func clipboardAuditTextBytes(text string) int {
	if len(text) > 65536 {
		return 65536
	}
	return len(text)
}

func clipboardAuditReason(err error) string {
	switch {
	case errors.Is(err, ErrClipboardRate):
		return "rate"
	case errors.Is(err, ErrProviderMalformed):
		return "malformed"
	case errors.Is(err, ErrClipboardDenied):
		return "policy"
	default:
		return "provider"
	}
}

func (s *gatewaySession) handleInput(payload []byte, sequence uint32) error {
	event, err := DecodeInputEvent(payload)
	if err != nil {
		return err
	}
	requiredFeature := ""
	switch event.Device {
	case "mouse-absolute":
		requiredFeature = "input.absolute.v1"
	case "mouse-scroll":
		requiredFeature = "input.scroll.v1"
	}
	if requiredFeature != "" && !slices.Contains(s.server.config.Features, requiredFeature) {
		return ErrInputMalformed
	}
	event.Sequence = sequence
	if err := s.provider.Input(s.ctx, event); err != nil {
		s.server.metrics.InputRejected.Add(1)
		return err
	}
	key := fmt.Sprintf("%s:%d", event.Device, event.Code)
	s.inputMu.Lock()
	if event.Pressed {
		s.pressed[key] = struct{}{}
	} else {
		delete(s.pressed, key)
	}
	s.inputMu.Unlock()
	return nil
}

func channelSequence(sequence int64) (uint32, error) {
	if sequence < 0 || sequence > int64(^uint32(0)) {
		return 0, ErrProviderMalformed
	}
	return uint32(sequence), nil
}

func (s *gatewaySession) cleanup() {
	s.cleanupOnce.Do(func() {
		s.cancel()
		s.server.pacer.remove(s.authority.SessionID)
		cleanupCtx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		releaseInputsErr := s.provider.ReleaseAll(cleanupCtx)
		if releaseInputsErr != nil {
			s.server.metrics.ProviderErrors.Add(1)
		}
		terminateErr := s.provider.Terminate(cleanupCtx)
		if terminateErr != nil {
			s.server.metrics.ProviderErrors.Add(1)
		}
		state := s.provider.State()
		if releaseInputsErr != nil || terminateErr != nil {
			state.State = ProviderStateCleanup
			state.CleanupPending = true
			if err := s.server.reportProviderState(cleanupCtx, state); err != nil {
				s.server.metrics.ProviderErrors.Add(1)
			} else if err := s.server.config.Admission.Release(cleanupCtx, s.authority); err != nil {
				s.server.metrics.ProviderErrors.Add(1)
			}
			_ = s.connection.CloseWithError(applicationError, "session closed")
			return
		}
		if errors.Is(s.endReason, ErrProviderDisconnected) {
			state.State = ProviderStateDisconnected
			state.CleanupPending = false
		}
		if err := s.server.config.Admission.Release(cleanupCtx, s.authority); err != nil {
			s.server.metrics.ProviderErrors.Add(1)
		}
		if err := s.server.reportProviderState(cleanupCtx, state); err != nil {
			s.server.metrics.ProviderErrors.Add(1)
		}
		_ = s.connection.CloseWithError(applicationError, "session closed")
	})
}

func writeStableError(writer io.Writer, code string, _ error, retryable bool) error {
	message := stableErrorMessage(code)
	payload, encodeErr := protocol.EncodeStableError(protocol.StableError{Version: "1", Code: code, Message: message, Retryable: retryable})
	if encodeErr != nil {
		return encodeErr
	}
	return writeWire(writer, payload, defaultHelloLimit)
}

func stableErrorMessage(code string) string {
	switch code {
	case "invalid_hello":
		return "invalid client hello"
	case "gateway_draining":
		return "gateway is draining"
	case "wrong_gateway":
		return "gateway does not match admission request"
	case "admission_rejected", "expired_grant":
		return "admission rejected"
	case "invalid_authority":
		return "invalid session authority"
	case "no_capability_overlap":
		return "no compatible capability"
	case "clipboard_audit_unavailable":
		return "clipboard audit unavailable"
	case "provider_work_unavailable", "provider_identity_rejected", "provider_malformed", "provider_timeout", "provider_unavailable", "provider_state_unavailable":
		return "provider unavailable"
	default:
		return "request failed"
	}
}

func stableAdmissionCode(err error) string {
	if errors.Is(err, ErrGatewayDraining) {
		return "gateway_draining"
	}
	if errors.Is(err, ErrAuthorityExpired) {
		return "expired_grant"
	}
	return "admission_rejected"
}

func stableProviderCode(err error) string {
	switch {
	case errors.Is(err, ErrProviderIdentity):
		return "provider_identity_rejected"
	case errors.Is(err, ErrProviderMalformed):
		return "provider_malformed"
	case errors.Is(err, ErrProviderTimeout), errors.Is(err, context.DeadlineExceeded):
		return "provider_timeout"
	default:
		return "provider_unavailable"
	}
}

func writeWire(writer io.Writer, payload []byte, max int) error {
	if len(payload) > max {
		return ErrFrameSize
	}
	header := [4]byte{}
	binary.BigEndian.PutUint32(header[:], uint32(len(payload)))
	if _, err := writer.Write(header[:]); err != nil {
		return err
	}
	_, err := writer.Write(payload)
	return err
}

func readWire(reader io.Reader, max int) ([]byte, error) {
	var header [4]byte
	if _, err := io.ReadFull(reader, header[:]); err != nil {
		return nil, err
	}
	length := binary.BigEndian.Uint32(header[:])
	if length > uint32(max) {
		return nil, ErrFrameSize
	}
	payload := make([]byte, int(length))
	if _, err := io.ReadFull(reader, payload); err != nil {
		return nil, err
	}
	return payload, nil
}

type Client struct {
	connection     *quic.Conn
	control        *quic.Stream
	controlReadMu  sync.Mutex
	controlWriteMu sync.Mutex
	pendingControl map[string][][]byte
	Authority      protocol.ClientSessionAuthority
}

func Dial(ctx context.Context, address string, tlsConfig *tls.Config, request protocol.TunnelAdmissionRequest) (*Client, error) {
	if tlsConfig == nil || tlsConfig.MinVersion < tls.VersionTLS13 || tlsConfig.RootCAs == nil || len(tlsConfig.RootCAs.Subjects()) == 0 || len(tlsConfig.Certificates) == 0 {
		return nil, ErrGatewayTLS
	}
	tlsConfig = tlsConfig.Clone()
	if len(tlsConfig.NextProtos) == 0 {
		tlsConfig.NextProtos = []string{"versevdi-gateway-v1"}
	}
	quicConfig := &quic.Config{EnableDatagrams: true, MaxIdleTimeout: 30 * time.Second}
	connection, err := quic.DialAddr(ctx, address, tlsConfig, quicConfig)
	if err != nil {
		return nil, err
	}
	stream, err := connection.OpenStreamSync(ctx)
	if err != nil {
		_ = connection.CloseWithError(applicationError, "stream unavailable")
		return nil, err
	}
	payload, err := protocol.EncodeTunnelAdmissionRequest(request)
	if err != nil {
		_ = connection.CloseWithError(applicationError, "invalid hello")
		return nil, err
	}
	if err := writeWire(stream, payload, defaultHelloLimit); err != nil {
		_ = connection.CloseWithError(applicationError, "invalid hello")
		return nil, err
	}
	response, err := readWire(stream, defaultHelloLimit)
	if err != nil {
		_ = connection.CloseWithError(applicationError, "no authority")
		return nil, err
	}
	authority, authorityErr := protocol.DecodeClientSessionAuthority(response)
	if authorityErr != nil {
		stable, stableErr := protocol.DecodeStableError(response)
		if stableErr == nil {
			return nil, fmt.Errorf("%s: %s", stable.Code, stable.Message)
		}
		return nil, authorityErr
	}
	return &Client{connection: connection, control: stream, pendingControl: make(map[string][][]byte), Authority: authority}, nil
}

func (c *Client) SendInput(event InputEvent) error {
	payload, err := EncodeInputEvent(event)
	if err != nil {
		return err
	}
	frame := protocol.ChannelFrame{Version: "1", FlowID: inputFlowID, Sequence: int64(event.Sequence), Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
	encoded, err := protocol.EncodeChannelFrame(frame)
	if err != nil {
		return err
	}
	return c.writeControl(encoded)
}

func (c *Client) SendControl(payload []byte) error {
	return c.sendControl(0, payload)
}

func (c *Client) SendFeedback(feedback Feedback) error {
	payload, err := EncodeClientFeedback(feedback)
	if err != nil {
		return err
	}
	return c.sendControl(feedback.Sequence, payload)
}

func (c *Client) SendClipboard(value protocol.GatewayClipboardText) error {
	payload, err := protocol.EncodeGatewayClipboardText(value)
	if err != nil {
		return err
	}
	frame := protocol.ChannelFrame{Version: "1", FlowID: clipboardFlowID, Sequence: 0, Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
	encoded, err := protocol.EncodeChannelFrame(frame)
	if err != nil {
		return err
	}
	return c.writeControl(encoded)
}

func (c *Client) sendControl(sequence uint32, payload []byte) error {
	frame := protocol.ChannelFrame{Version: "1", FlowID: controlFlowID, Sequence: int64(sequence), Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
	encoded, err := protocol.EncodeChannelFrame(frame)
	if err != nil {
		return err
	}
	return c.writeControl(encoded)
}

func (c *Client) writeControl(encoded []byte) error {
	c.controlWriteMu.Lock()
	defer c.controlWriteMu.Unlock()
	return writeWire(c.control, encoded, defaultControlLimit)
}

func (c *Client) ReceiveFrame(ctx context.Context) (Frame, error) {
	data, err := c.connection.ReceiveDatagram(ctx)
	if err != nil {
		return Frame{}, err
	}
	return DecodeFrame(data)
}

func (c *Client) ReceiveProviderEvent(ctx context.Context) (ProviderEvent, error) {
	payload, err := c.receiveControlPayload(ctx, controlFlowID)
	if err != nil {
		return ProviderEvent{}, err
	}
	if len(payload) > 1024 {
		return ProviderEvent{}, ErrProviderMalformed
	}
	event, err := DecodeProviderEvent(payload)
	if err == nil && (event.Kind == ProviderEventTerminated || event.Kind == ProviderEventDisconnected) {
		err = c.SendFeedback(Feedback{Kind: FeedbackTerminalReceipt})
	}
	return event, err
}

func (c *Client) ReceiveClipboard(ctx context.Context) (protocol.GatewayClipboardText, error) {
	payload, err := c.receiveControlPayload(ctx, clipboardFlowID)
	if err != nil {
		return protocol.GatewayClipboardText{}, err
	}
	return protocol.DecodeGatewayClipboardText(payload)
}

func (c *Client) receiveControlPayload(ctx context.Context, flowID string) ([]byte, error) {
	if c == nil || c.control == nil || (flowID != controlFlowID && flowID != clipboardFlowID) {
		return nil, ErrProviderMalformed
	}
	c.controlReadMu.Lock()
	defer c.controlReadMu.Unlock()
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if queued := c.pendingControl[flowID]; len(queued) > 0 {
		payload := queued[0]
		c.pendingControl[flowID] = queued[1:]
		return payload, nil
	}
	if deadline, ok := ctx.Deadline(); ok {
		if err := c.control.SetReadDeadline(deadline); err != nil {
			return nil, err
		}
		defer c.control.SetReadDeadline(time.Time{})
	}
	for {
		data, err := readWire(c.control, defaultControlLimit)
		if err != nil {
			return nil, err
		}
		frame, err := protocol.DecodeChannelFrame(data)
		if err != nil || (frame.FlowID != controlFlowID && frame.FlowID != clipboardFlowID) {
			return nil, ErrProviderMalformed
		}
		payload, err := base64.StdEncoding.DecodeString(frame.Payload)
		if err != nil || len(payload) > maxFrameSize {
			return nil, ErrProviderMalformed
		}
		if frame.FlowID == flowID {
			return payload, nil
		}
		if len(c.pendingControl[frame.FlowID]) >= clientControlBacklog {
			return nil, ErrFramePayloadLimit
		}
		c.pendingControl[frame.FlowID] = append(c.pendingControl[frame.FlowID], payload)
	}
}

func (c *Client) Close() error {
	return c.connection.CloseWithError(applicationError, "client closed")
}