package loc

// EncodeAudioLevel packs an RFC 6464 audio level and voice-activity bit
// into the byte stored in LOC's AudioLevel property (§2.3.3.2).
//
// level is the magnitude in -dBov in the range [0, 127] (0 = loudest,
// 127 = silence). voiceActivity is the V flag from RFC 6464 §3. Bits
// above the 7-bit level range are clipped.
//
// Wire layout (MSB to LSB):
//
//        V | L L L L L L L
//          bit 7   bits 0-6
func EncodeAudioLevel(level uint8, voiceActivity bool) uint8 {
        b := level & 0x7F
        if voiceActivity {
                b |= 0x80
        }
        return b
}

// DecodeAudioLevel splits the LOC AudioLevel byte into the RFC 6464
// level magnitude (bits 0-6) and voice-activity flag (bit 7).
func DecodeAudioLevel(b uint8) (level uint8, voiceActivity bool) {
        return b & 0x7F, b&0x80 != 0
}

package loc

// Object is one LOC-packaged media chunk: the LOC Public Properties
// that travel in the MOQ Object Properties block, plus the codec
// elementary stream bytes that travel in the MOQ Object Payload.
//
// Object does not own a MOQ message — it produces the bytes a caller
// drops into a [github.com/floatdrop/moq-go/pkg/moqt/message.SubgroupObject]
// (or any other MOQ message that carries an object payload). The
// caller controls Group ID, Object ID, subgroup framing, and stream
// scheduling.
//
// LOC Private Properties (SecureObjects, §3.1.3) are not modelled
// here. When that support lands, Object grows a Private field and
// Encode/Decode learn the length-prefixed-prepended-to-payload layout
// the SecureObjects spec defines.
type Object struct {
        Properties Properties
        // Payload is the codec elementary stream — the "internal data" of
        // an EncodedAudioChunk / EncodedVideoChunk in the WebCodecs Codec
        // Registry. Nil and empty are equivalent.
        Payload []byte
}

// Encode returns the bytes ready to plug into a MOQ Object:
//
//   - props goes into the surrounding message's Properties slot
//     (e.g. [message.SubgroupObject.Properties]).
//   - payload is the MOQ Object Payload.
//
// The returned props slice is the inner KV-pair blob without an outer
// length prefix; the surrounding MOQ message frames it. The returned
// payload aliases [Object.Payload]; the caller must not mutate the
// originating slice after calling Encode if it is still being read.
func (o *Object) Encode() (props, payload []byte) {
        return o.Properties.Encode(), o.Payload
}

// Decode reconstructs an Object from the Properties bytes (the inner
// KV-pair blob, no length prefix) and the Object Payload bytes. Both
// slices may be nil.
//
// The returned Object's Payload aliases the input slice. Callers that
// need to retain Payload past the lifetime of the input buffer must
// copy.
func Decode(props, payload []byte) (Object, error) {
        p, err := ParseProperties(props)
        if err != nil {
                return Object{}, err
        }
        return Object{Properties: p, Payload: payload}, nil
}

package loc

import (
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// Properties carries the LOC metadata that travels in the MOQ Object
// Properties block. The wire encoding is a sequence of [wire.KVPair]
// values, identical to [message.ObjectProperties] but without the
// outer length prefix — that prefix is owned by the containing
// MOQ message (e.g. [message.SubgroupObject] applies it on write).
//
// The well-known LOC fields are exposed as typed accessors. Pairs not
// recognised by the typed accessors land in Extras and round-trip
// unchanged.
//
// Zero values for Timestamp / Timescale / AudioLevel are valid wire
// values; absence is tracked by the matching Has-bit so callers can
// distinguish "field not present" from "field present and zero". The
// byte-slice fields (VideoConfig / VideoFrameMarking / AudioConfig)
// instead use nil to mean absent.
type Properties struct {
        Timestamp uint64
        Timescale uint64

        // VideoConfig is the codec extradata (§2.3.2.1). Absent when nil.
        VideoConfig []byte

        // AudioConfig is the codec configuration (§2.3.3.1). Absent when nil.
        AudioConfig []byte

        // VideoFrameMarking carries the RFC 9626 frame-marking flags and
        // layer identifiers (§2.3.2.2) as a 1-4 byte string. Absent when nil.
        VideoFrameMarking []byte

        AudioLevel uint8

        HasTimestamp  bool
        HasTimescale  bool
        HasAudioLevel bool

        // Extras carries KV pairs whose Type is not one of the well-known
        // LOC properties. They round-trip verbatim. Extras MUST NOT contain
        // any of the well-known LOC property IDs — use the typed fields
        // instead. Append does not validate this; mixing the two leads to
        // undefined ordering.
        Extras []wire.KVPair
}

// Append serialises Properties as a flat sequence of KV pairs (no length
// prefix). The byte slice it produces (via [wire.Writer.Bytes]) goes
// directly into [message.SubgroupObject.Properties]; the
// SubgroupObject's own writer adds the outer length prefix.
func (p *Properties) Append(w *wire.Writer) {
        pairs := p.toPairs()
        w.KVPairs(pairs)
}

// Parse consumes KV pairs from r until r is empty and populates the
// fields of p. Any unknown property IDs land in Extras. Returns an
// error only if the wire data is malformed.
//
// Callers obtain a *wire.Reader bounded to the Properties bytes from
// the surrounding message — e.g. [message.SubgroupObject.Properties]
// is already the inner KV-pair blob with no length prefix.
func (p *Properties) Parse(r *wire.Reader) error {
        pairs, err := r.KVPairsRemaining()
        if err != nil {
                return fmt.Errorf("moqt/loc: parsing properties: %w", err)
        }
        *p = Properties{}
        for _, kv := range pairs {
                switch kv.Type {
                case PropTimestamp:
                        p.Timestamp = kv.IntVal
                        p.HasTimestamp = true
                case PropTimescale:
                        p.Timescale = kv.IntVal
                        p.HasTimescale = true
                case PropVideoFrameMarking:
                        p.VideoFrameMarking = kv.ByteVal
                case PropAudioLevel:
                        if kv.IntVal > 0xFF {
                                return fmt.Errorf("moqt/loc: audio level %d exceeds 0xFF", kv.IntVal)
                        }
                        p.AudioLevel = uint8(kv.IntVal)
                        p.HasAudioLevel = true
                case PropVideoConfig:
                        p.VideoConfig = kv.ByteVal
                case PropAudioConfig:
                        p.AudioConfig = kv.ByteVal
                default:
                        p.Extras = append(p.Extras, kv)
                }
        }
        return nil
}

// Encode is a convenience that returns the serialised bytes ready to
// drop into [message.SubgroupObject.Properties].
func (p *Properties) Encode() []byte {
        var w wire.Writer
        p.Append(&w)
        return w.Bytes()
}

// ParseProperties decodes a Properties value from raw KV-pair bytes
// (e.g. [message.SubgroupObject.Properties]). It is the inverse of
// [Properties.Encode].
func ParseProperties(raw []byte) (Properties, error) {
        var p Properties
        if len(raw) == 0 {
                return p, nil
        }
        r := wire.NewReader(raw)
        if err := p.Parse(r); err != nil {
                return Properties{}, err
        }
        return p, nil
}

// toPairs collects the typed fields and Extras into a single KV slice.
// [wire.Writer.KVPairs] sorts by Type before encoding, so the order
// here does not matter.
func (p *Properties) toPairs() []wire.KVPair {
        n := len(p.Extras)
        if p.HasTimestamp {
                n++
        }
        if p.HasTimescale {
                n++
        }
        if p.VideoFrameMarking != nil {
                n++
        }
        if p.HasAudioLevel {
                n++
        }
        if p.VideoConfig != nil {
                n++
        }
        if p.AudioConfig != nil {
                n++
        }
        if n == 0 {
                return nil
        }
        pairs := make([]wire.KVPair, 0, n)
        if p.HasTimestamp {
                pairs = append(pairs, wire.KVPair{Type: PropTimestamp, IntVal: p.Timestamp})
        }
        if p.HasTimescale {
                pairs = append(pairs, wire.KVPair{Type: PropTimescale, IntVal: p.Timescale})
        }
        if p.VideoFrameMarking != nil {
                pairs = append(pairs, wire.KVPair{Type: PropVideoFrameMarking, ByteVal: p.VideoFrameMarking})
        }
        if p.HasAudioLevel {
                pairs = append(pairs, wire.KVPair{Type: PropAudioLevel, IntVal: uint64(p.AudioLevel)})
        }
        if p.VideoConfig != nil {
                pairs = append(pairs, wire.KVPair{Type: PropVideoConfig, ByteVal: p.VideoConfig})
        }
        if p.AudioConfig != nil {
                pairs = append(pairs, wire.KVPair{Type: PropAudioConfig, ByteVal: p.AudioConfig})
        }
        pairs = append(pairs, p.Extras...)
        return pairs
}

package loc

import "encoding/binary"

// NALFraming describes how NAL units are delimited inside an
// AVC/HEVC LOC payload. See LOC §2.1.3 and §2.1.4.
type NALFraming int

const (
        // NALFramingUnknown means the payload does not begin with a
        // recognisable NAL framing (e.g. it is a non-NAL codec like AV1,
        // or the buffer is too short to tell).
        NALFramingUnknown NALFraming = iota

        // NALFramingStartCode4 means the payload begins with the 4-byte
        // AnnexB start code 0x00 0x00 0x00 0x01.
        NALFramingStartCode4

        // NALFramingStartCode3 means the payload begins with the 3-byte
        // AnnexB start code 0x00 0x00 0x01. §2.1.4 permits this only when
        // the track never uses length prefixes or Video Config.
        NALFramingStartCode3

        // NALFramingLengthPrefix means the payload begins with a 4-byte
        // big-endian length followed by that many bytes of NAL unit data.
        // §2.1.3: a length value of 1 SHOULD be interpreted as a start
        // code rather than a length, so the length-prefix detector rejects
        // that ambiguous case.
        NALFramingLengthPrefix
)

// DetectNALFraming inspects the first bytes of a video payload and
// guesses how its NAL units are delimited. Detection is heuristic:
// the result is reliable only for AVC/HEVC payloads that start with a
// NAL unit. Returns [NALFramingUnknown] when the buffer does not
// match any of the three patterns or is shorter than 4 bytes.
//
// Detection order (matches §2.1.3's tie-breaker for length == 1):
//  1. The 4-byte AnnexB start code 0x00 0x00 0x00 0x01.
//  2. The 3-byte AnnexB start code 0x00 0x00 0x01.
//  3. A 4-byte length prefix whose value is > 1 and does not exceed
//     the remaining payload length.
func DetectNALFraming(payload []byte) NALFraming {
        if len(payload) < 3 {
                return NALFramingUnknown
        }
        if len(payload) >= 4 && payload[0] == 0x00 && payload[1] == 0x00 &&
                payload[2] == 0x00 && payload[3] == 0x01 {
                return NALFramingStartCode4
        }
        if payload[0] == 0x00 && payload[1] == 0x00 && payload[2] == 0x01 {
                return NALFramingStartCode3
        }
        if len(payload) < 4 {
                return NALFramingUnknown
        }
        length := binary.BigEndian.Uint32(payload[:4])
        if length <= 1 {
                return NALFramingUnknown
        }
        if uint64(length)+4 > uint64(len(payload)) {
                return NALFramingUnknown
        }
        return NALFramingLengthPrefix
}

package message

import (
        "errors"
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// Datagram type field bit constants (§11.3).
const (
        DatagramPropertiesBit      = 0x01 // Properties field present
        DatagramEndOfGroupBit      = 0x02 // End of group marker
        DatagramZeroObjectIDBit    = 0x04 // Object ID omitted (treated as 0)
        DatagramDefaultPriorityBit = 0x08 // Priority omitted (use subscription default)
        DatagramStatusBit          = 0x20 // Object Status present instead of payload
)

// Valid datagram type ranges.
const (
        DatagramTypeMin       = 0x00
        DatagramTypeMax       = 0x0F
        DatagramTypeStatusMin = 0x20
        DatagramTypeStatusMax = 0x2F
)

// ObjectDatagram represents a MoQT object sent via QUIC datagram (§11.3).
type ObjectDatagram struct {
        Type              uint64 // Complex bit field
        TrackAlias        uint64
        GroupID           uint64
        ObjectID          uint64 // Optional based on ZERO_OBJECT_ID bit
        PublisherPriority uint8  // Optional based on DEFAULT_PRIORITY bit
        Properties        []byte // Optional based on PROPERTIES bit
        ObjectStatus      uint64 // Optional based on STATUS bit
        ObjectPayload     []byte // Present when STATUS bit is 0
}

// IsValidDatagramType checks if a datagram Type Flags value is valid per
// §11.3.1. The only bits with a specified meaning are the five Datagram*Bit
// flags, so the valid values are 0x00..0x0F / 0x20..0x21 / 0x24..0x25 /
// 0x28..0x29 / 0x2C..0x2D.
//
// §11.3.1 lists the invalid values, which MUST close the session with a
// PROTOCOL_VIOLATION:
//
//   - bit 4 (0x10) set, or any bit set whose meaning is not specified (i.e.
//     not 0x00..0x0F / 0x20..0x2F);
//   - both STATUS and END_OF_GROUP set (0x22,0x23,0x26,0x27,0x2A,0x2B,0x2E,
//     0x2F).
//
// Note STATUS+PROPERTIES (0x21,0x25,0x29,0x2D) IS a valid type: it only
// becomes an error when the Object Status is not Normal (0x0) — a per-value
// rule enforced by [ObjectDatagram.Validate], not a type-level one.
func IsValidDatagramType(typ uint64) bool {
        if typ > DatagramTypeStatusMax || (typ > DatagramTypeMax && typ < DatagramTypeStatusMin) {
                return false
        }
        if typ&DatagramStatusBit != 0 && typ&DatagramEndOfGroupBit != 0 {
                return false
        }
        return true
}

// HasProperties returns true if the PROPERTIES bit is set.
func (d *ObjectDatagram) HasProperties() bool {
        return d.Type&DatagramPropertiesBit != 0
}

// HasEndOfGroup returns true if the END_OF_GROUP bit is set.
func (d *ObjectDatagram) HasEndOfGroup() bool {
        return d.Type&DatagramEndOfGroupBit != 0
}

// HasZeroObjectID returns true if the ZERO_OBJECT_ID bit is set.
func (d *ObjectDatagram) HasZeroObjectID() bool {
        return d.Type&DatagramZeroObjectIDBit != 0
}

// HasDefaultPriority returns true if the DEFAULT_PRIORITY bit is set.
func (d *ObjectDatagram) HasDefaultPriority() bool {
        return d.Type&DatagramDefaultPriorityBit != 0
}

// HasStatus returns true if the STATUS bit is set.
func (d *ObjectDatagram) HasStatus() bool {
        return d.Type&DatagramStatusBit != 0
}

// Validate checks if the datagram is valid according to MoQT spec §11.3.1.
// Every violation below is a session-level PROTOCOL_VIOLATION at the
// receiver.
func (d *ObjectDatagram) Validate() error {
        if !IsValidDatagramType(d.Type) {
                return fmt.Errorf("invalid datagram type: 0x%02X", d.Type)
        }

        // Per §11.3.1: PROPERTIES bit set with a Properties Length of 0 MUST
        // close the session with a PROTOCOL_VIOLATION.
        if d.HasProperties() && len(d.Properties) == 0 {
                return errors.New("invalid datagram: PROPERTIES bit set with zero-length Properties")
        }

        if d.HasStatus() {
                // §11.2.1.1: the defined Object Status values are Normal (0x0),
                // End of Group (0x3), and End of Track (0x4); any other value
                // SHOULD be treated as a protocol error. Matches the subgroup
                // object codec's enforcement.
                switch d.ObjectStatus {
                case ObjectStatusNormal, ObjectStatusEndOfGroup, ObjectStatusEndOfTrack:
                default:
                        return fmt.Errorf("invalid datagram: unknown object status 0x%X", d.ObjectStatus)
                }

                // §11.3.1: "If an Object Datagram includes both the STATUS bit and
                // PROPERTIES bit, and the Object Status is not Normal (0x0), the
                // endpoint MUST close the session with a PROTOCOL_VIOLATION,
                // because only Normal Objects can have Properties."
                if d.HasProperties() && d.ObjectStatus != ObjectStatusNormal {
                        return fmt.Errorf("invalid datagram: non-Normal status 0x%X with Properties", d.ObjectStatus)
                }
        }

        return nil
}

// Append serializes the datagram to a wire.Writer.
func (d *ObjectDatagram) Append(w *wire.Writer) {
        w.Varint(d.Type)
        w.Varint(d.TrackAlias)
        w.Varint(d.GroupID)

        if !d.HasZeroObjectID() {
                w.Varint(d.ObjectID)
        }
        if !d.HasDefaultPriority() {
                w.UInt8(d.PublisherPriority)
        }
        if d.HasProperties() {
                w.VarintBytes(d.Properties)
        }

        if d.HasStatus() {
                w.Varint(d.ObjectStatus)
        } else {
                w.FixedBytes(d.ObjectPayload)
        }
}

// Parse deserializes a datagram from a wire.Reader into d.
func (d *ObjectDatagram) Parse(r *wire.Reader) error {
        typ, err := r.Varint()
        if err != nil {
                return fmt.Errorf("failed to read datagram type: %w", err)
        }
        d.Type = typ

        // Validate the type before parsing fields — the layout depends on its
        // bits. Rejects STATUS+END_OF_GROUP and out-of-form values (§11.3.1);
        // per-value rules (e.g. non-Normal status with Properties) run in
        // Validate once the fields are read.
        if !IsValidDatagramType(d.Type) {
                return fmt.Errorf("invalid datagram type: 0x%02X", d.Type)
        }

        d.TrackAlias, err = r.Varint()
        if err != nil {
                return fmt.Errorf("failed to read track alias: %w", err)
        }

        d.GroupID, err = r.Varint()
        if err != nil {
                return fmt.Errorf("failed to read group ID: %w", err)
        }

        if !d.HasZeroObjectID() {
                d.ObjectID, err = r.Varint()
                if err != nil {
                        return fmt.Errorf("failed to read object ID: %w", err)
                }
        } else {
                d.ObjectID = 0
        }

        if !d.HasDefaultPriority() {
                d.PublisherPriority, err = r.UInt8()
                if err != nil {
                        return fmt.Errorf("failed to read publisher priority: %w", err)
                }
        }

        if d.HasProperties() {
                d.Properties, err = r.VarintBytes()
                if err != nil {
                        return fmt.Errorf("failed to read properties: %w", err)
                }
        }

        // Read ObjectStatus or ObjectPayload based on STATUS bit
        if d.HasStatus() {
                d.ObjectStatus, err = r.Varint()
                if err != nil {
                        return fmt.Errorf("failed to read object status: %w", err)
                }
                // The status varint is the last field (§11.3.1 Figure 24); trailing
                // bytes mean the sender and receiver disagree on the layout.
                if !r.Empty() {
                        return fmt.Errorf("invalid datagram: %d trailing byte(s) after Object Status", r.Remaining())
                }
        } else {
                d.ObjectPayload = r.RemainingBytes()
        }

        // Semantic checks (zero-length Properties, non-Normal status with
        // Properties) live in Validate so parsing and standalone validation
        // cannot drift.
        return d.Validate()
}

package message

import (
        "fmt"
        "io"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// UnknownDataStreamTypeError is returned when the leading Type of an inbound
// data uni-stream is not one of the recognized data-stream types. It is
// session-fatal (§3.4): session.AcceptDataStream closes the session with
// PROTOCOL_VIOLATION before returning it.
type UnknownDataStreamTypeError struct {
        Type uint64
}

func (e *UnknownDataStreamTypeError) Error() string {
        return fmt.Sprintf("moqt/message: unknown data stream type %#x", e.Type)
}

// ReservedSubgroupIDModeError is returned when the leading Type of an inbound
// data uni-stream matches the SUBGROUP_HEADER pattern (bit 4 set, bit 7 clear)
// but carries the reserved SUBGROUP_ID_MODE value 0b11 in bits 1-2. Per
// §11.4.2, this MUST be treated as a session-level PROTOCOL_VIOLATION.
type ReservedSubgroupIDModeError struct {
        Type uint64
}

func (e *ReservedSubgroupIDModeError) Error() string {
        return fmt.Sprintf(
                "moqt/message: SUBGROUP_HEADER type %#x has reserved SUBGROUP_ID_MODE 0b11 — PROTOCOL_VIOLATION",
                e.Type,
        )
}

// ReadDataStreamType reads the leading Type varint that prefixes every MoQT
// uni-stream data header (SUBGROUP_HEADER §11.4.2, FETCH_HEADER §11.4.4,
// padding §11.5.1, ...). A dispatcher uses this together with type predicates
// such as IsSubgroupHeaderType to decide how to consume the remainder of the
// stream.
func ReadDataStreamType(r io.Reader) (uint64, error) {
        typ, err := wire.ReadVarint(wire.NewByteReader(r))
        if err != nil {
                return 0, fmt.Errorf("moqt/message: read uni-stream type: %w", err)
        }
        return typ, nil
}

// PaddingStreamType is the leading Type varint of a padding uni-stream
// (§11.5.1). Receivers MUST silently discard padding streams.
const PaddingStreamType uint64 = 0x132B3E28

package message

import (
        "fmt"
        "io"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// Fetch is a FETCH message per §10.13.
//
// draft-20 removed the Fetch Type discriminant and with it the Joining
// variant: a FETCH now just names a track, and its range travels in the
// LOCATION_FILTER parameter (§5.1.2) like every other filter. The backfill
// that a Joining FETCH used to provide is now a fill fetch stream, requested
// with FILL_PARAMETERS on the SUBSCRIBE itself (§5.1.3).
//
//        FETCH Message {
//          Type (vi64) = 0x16,
//          Length (16),
//          Request ID (vi64),
//          Track Namespace (..),
//          Track Name Length (vi64),
//          Track Name (..),
//          Number of Parameters (vi64),
//          Parameters (..) ...
//        }
type Fetch struct {
        RequestID  uint64
        Namespace  wire.TrackNamespace
        Name       []byte
        Parameters Parameters
}

// Append serializes the FETCH message to w.
func (m *Fetch) Append(w *wire.Writer) {
        w.Varint(m.RequestID)
        w.TrackNamespace(m.Namespace)
        w.VarintBytes(m.Name)
        m.Parameters.append(w)
}

// Parse deserializes the FETCH message from r.
func (m *Fetch) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.Varint(&m.RequestID)
        s.TrackNamespace(&m.Namespace)
        s.VarintBytes(&m.Name)
        if err := s.Err(); err != nil {
                return err
        }
        return m.Parameters.parse(r)
}

// Type returns the wire type ID for FETCH.
func (m *Fetch) Type() Type             { return TypeFetch }
func (m *Fetch) GetRequestID() uint64   { return m.RequestID }
func (m *Fetch) SetRequestID(id uint64) { m.RequestID = id }

// validateFullTrackName enforces §2.4.1: "If an endpoint receives a Track
// Namespace or a Full Track Name exceeding 4,096 bytes, it MUST close the
// session with a PROTOCOL_VIOLATION." The namespace-only half is already
// enforced at parse time by wire.Reader.TrackNamespace; this adds the Track
// Name's length for messages that carry a full name.
func validateFullTrackName(ns wire.TrackNamespace, name []byte) error {
        if total := ns.ByteLen() + len(name); total > wire.MaxFullTrackNameBytes {
                return fmt.Errorf("moqt/message: full track name is %d bytes, max %d (§2.4.1)",
                        total, wire.MaxFullTrackNameBytes)
        }
        return nil
}

// Validate enforces the FETCH invariant the wire decoder cannot: §2.4.1's
// 4,096-byte cap applies to the full track name, not just the namespace.
// ParsePayload invokes this automatically after decoding a FETCH frame.
//
// The range is no longer a FETCH field in draft-20, so its validation lives
// on the LOCATION_FILTER parameter ([LocationFilter.Validate]).
func (m *Fetch) Validate() error {
        return validateFullTrackName(m.Namespace, m.Name)
}

// FetchOK is a FETCH_OK message per §10.14.
type FetchOK struct {
        EndOfTrack      bool
        EndLocation     Location
        Parameters      Parameters
        TrackProperties []byte
}

// Append serializes the FETCH_OK message to w.
func (m *FetchOK) Append(w *wire.Writer) {
        if m.EndOfTrack {
                w.UInt8(1)
        } else {
                w.UInt8(0)
        }
        w.Varint(m.EndLocation.Group)
        w.Varint(m.EndLocation.Object)
        m.Parameters.append(w)
        w.FixedBytes(m.TrackProperties)
}

// Parse deserializes the FETCH_OK message from r.
func (m *FetchOK) Parse(r *wire.Reader) error {
        s := r.Scanner()
        var eot uint8
        s.UInt8(&eot)
        s.Varint(&m.EndLocation.Group)
        s.Varint(&m.EndLocation.Object)
        if err := s.Err(); err != nil {
                return err
        }
        m.EndOfTrack = eot == 1
        if err := m.Parameters.parse(r); err != nil {
                return err
        }
        m.TrackProperties = r.RemainingBytes()
        return nil
}

// Type returns the wire type ID for FETCH_OK.
func (m *FetchOK) Type() Type {
        return TypeFetchOK
}

// FetchHeader is the header of a FETCH_HEADER stream (§11.4.4). It identifies
// which FETCH request this stream responds to.
type FetchHeader struct {
        RequestID uint64
}

// RawType returns the leading Type varint as it appeared on the wire.
func (h FetchHeader) RawType() uint64 {
        return 0x05
}

// WriteFetchHeader writes the FETCH_HEADER wire Type and Request ID.
func WriteFetchHeader(w io.Writer, h FetchHeader) error {
        buf := wire.AppendVarint(nil, h.RawType())
        buf = wire.AppendVarint(buf, h.RequestID)
        _, err := w.Write(buf)
        return err
}

// ReadFetchHeader reads a FETCH_HEADER from r. The caller must have already
// read the stream type (0x05) via ReadDataStreamType.
func ReadFetchHeader(r io.Reader) (FetchHeader, error) {
        requestID, err := wire.ReadVarint(wire.NewByteReader(r))
        if err != nil {
                return FetchHeader{}, fmt.Errorf("moqt/message: read FETCH_HEADER Request ID: %w", err)
        }
        return FetchHeader{RequestID: requestID}, nil
}

// IsFetchHeaderType reports whether typ is a FETCH_HEADER stream type (0x05).
func IsFetchHeaderType(typ uint64) bool {
        return typ == 0x05
}

package message

import (
        "errors"
        "fmt"
        "io"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// FetchObject represents a single object in a FETCH response stream per §11.4.4.
type FetchObject struct {
        // SerializationFlags control which fields are present and how they're encoded.
        SerializationFlags uint64

        // GroupIDDelta is the delta from the previous Group ID. Present when
        // FetchFlagGroupIDDelta bit (0x08) is set.
        GroupIDDelta uint64

        // SubgroupID is encoded based on the two LSBs of SerializationFlags (mask 0x03).
        // Only present on the wire when the mode is FetchSubgroupIDExplicit (0x03).
        SubgroupID uint64

        // ObjectIDDelta is the delta from the previous Object ID (no +1), or the
        // absolute Object ID on the first object and whenever GroupIDDelta is
        // present (§11.4.4.1). Present when the
        // FetchFlagObjectIDDelta bit (0x04) is set.
        ObjectIDDelta uint64

        // PublisherPriority is present when FetchFlagPriority (0x10) is set.
        PublisherPriority uint8

        // Properties are present when FetchFlagProperties (0x20) is set.
        Properties []byte

        // ObjectPayload is always present, encoded on the wire with a varint
        // length prefix (§11.4.4 Figure 28). FETCH objects carry no Object
        // Status field (§11.2.1.1); absent ranges are expressed with the
        // end-of-range markers instead.
        ObjectPayload []byte
}

// Serialization flag bits per §11.4.4.1 (Table 8 & 9).
//
// Bits 0–1 (mask 0x03): Subgroup ID mode — see FetchSubgroupIDMode.
// Bit 2  (0x04): Object ID Delta present.
// Bit 3  (0x08): Group ID Delta present.
// Bit 4  (0x10): Priority field present.
// Bit 5  (0x20): Properties field present.
// Bit 6  (0x40): Datagram — no Subgroup ID; the subgroup-mode LSBs are ignored.
// Bit 7+ : reserved / end-of-range special values.
const (
        FetchFlagSubgroupIDMode uint64 = 0x03 // bits 0–1: subgroup encoding mode
        FetchFlagObjectIDDelta  uint64 = 0x04 // bit 2: Object ID Delta present
        FetchFlagGroupIDDelta   uint64 = 0x08 // bit 3: Group ID Delta present
        FetchFlagPriority       uint64 = 0x10 // bit 4: Priority present
        FetchFlagProperties     uint64 = 0x20 // bit 5: Properties present
        FetchFlagDatagram       uint64 = 0x40 // bit 6: Datagram — ignore subgroup bits
)

// FetchSubgroupIDMode encodes how the Subgroup ID is determined (bits 0–1).
type FetchSubgroupIDMode uint8

const (
        FetchSubgroupIDZero         FetchSubgroupIDMode = 0x00 // Subgroup ID is zero
        FetchSubgroupIDPrior        FetchSubgroupIDMode = 0x01 // Subgroup ID = prior object's Subgroup ID
        FetchSubgroupIDPriorPlusOne FetchSubgroupIDMode = 0x02 // Subgroup ID = prior + 1
        FetchSubgroupIDExplicit     FetchSubgroupIDMode = 0x03 // Subgroup ID field is present
)

// End of range markers per §11.4.4.2. Each stands for every Object between
// the previously serialized one and the Location this marker carries.
const (
        FetchEndOfNonExistentRange = 0x8C  // End of Non-Existent Range
        FetchEndOfUnknownRange     = 0x10C // End of Unknown Range
        FetchEndOfTimedOutRange    = 0x20C // End of Timed-Out Range (draft-20)
)

// isEndOfRange reports whether flags is any of the three §11.4.4.2 end-of-range
// markers, which share a wire shape: Group ID and Object ID follow the flags
// varint, and Subgroup ID, Priority and Properties are all absent.
func isEndOfRange(flags uint64) bool {
        switch flags {
        case FetchEndOfNonExistentRange, FetchEndOfUnknownRange, FetchEndOfTimedOutRange:
                return true
        }
        return false
}

// Append serializes a FetchObject to w.
//
// For end-of-range markers (SerializationFlags == 0x8C, 0x10C or 0x20C), the spec
// requires Group ID and Object ID fields to follow the flags varint.
// For normal objects, the payload is length-prefixed (varint + bytes).
func (o *FetchObject) Append(w *wire.Writer) {
        flags := o.SerializationFlags
        // §11.4.4.1: when the Datagram bit is set the publisher "SHOULD set the
        // two least significant bits to zero"; mask them so hand-built flag
        // combinations stay conformant on the wire.
        if flags&FetchFlagDatagram != 0 && flags < 128 {
                flags &^= FetchFlagSubgroupIDMode
        }
        w.Varint(flags)

        // End-of-range markers: Group ID and Object ID are always present (§11.4.4.2).
        if isEndOfRange(o.SerializationFlags) {
                w.Varint(o.GroupIDDelta)  // used as absolute Group ID for end-of-range
                w.Varint(o.ObjectIDDelta) // used as absolute Object ID for end-of-range
                return
        }

        if o.SerializationFlags&FetchFlagGroupIDDelta != 0 {
                w.Varint(o.GroupIDDelta)
        }

        if o.hasSubgroupIDField() {
                w.Varint(o.SubgroupID)
        }

        if o.SerializationFlags&FetchFlagObjectIDDelta != 0 {
                w.Varint(o.ObjectIDDelta)
        }

        if o.SerializationFlags&FetchFlagPriority != 0 {
                w.UInt8(o.PublisherPriority)
        }

        if o.SerializationFlags&FetchFlagProperties != 0 {
                w.VarintBytes(o.Properties)
        }

        // Object Payload Length (vi64) + Object Payload (..) per §11.4.4 Figure 28.
        w.VarintBytes(o.ObjectPayload)
}

// truncated converts a bare io.EOF into io.ErrUnexpectedEOF. Parse reads
// after the leading flags varint use it: at that point part of an object has
// been consumed, so a FIN is a truncated object, not a clean end-of-stream.
// Callers (e.g. the relay's upstream stitcher) rely on that distinction to
// tell "the sender FIN'd between objects, vouching for the rest of the range
// (§11.4.4)" from "the response broke off mid-object".
func truncated(err error) error {
        if errors.Is(err, io.EOF) {
                return io.ErrUnexpectedEOF
        }
        return err
}

// Parse deserializes a FetchObject from r.
// r may be a *wire.Reader (in-memory) or a *wire.StreamReader (streaming).
func (o *FetchObject) Parse(r wire.Decoder) error {
        flags, err := r.Varint()
        if err != nil {
                return err
        }
        o.SerializationFlags = flags
        // Decided on the flags alone: their bits must not be read as fields.
        if err := checkFetchFlags(flags); err != nil {
                return err
        }

        // End-of-range markers: Group ID and Object ID follow (§11.4.4.2).
        if isEndOfRange(flags) {
                groupID, err := r.Varint()
                if err != nil {
                        return fmt.Errorf("moqt/message: end-of-range group ID: %w", truncated(err))
                }
                o.GroupIDDelta = groupID // stored in GroupIDDelta as absolute Group ID

                objectID, err := r.Varint()
                if err != nil {
                        return fmt.Errorf("moqt/message: end-of-range object ID: %w", truncated(err))
                }
                o.ObjectIDDelta = objectID // stored in ObjectIDDelta as absolute Object ID
                return nil
        }

        if flags&FetchFlagGroupIDDelta != 0 {
                delta, err := r.Varint()
                if err != nil {
                        return truncated(err)
                }
                o.GroupIDDelta = delta
        }

        if o.hasSubgroupIDField() {
                subgroupID, err := r.Varint()
                if err != nil {
                        return truncated(err)
                }
                o.SubgroupID = subgroupID
        }

        if flags&FetchFlagObjectIDDelta != 0 {
                delta, err := r.Varint()
                if err != nil {
                        return truncated(err)
                }
                o.ObjectIDDelta = delta
        }

        if flags&FetchFlagPriority != 0 {
                priority, err := r.UInt8()
                if err != nil {
                        return truncated(err)
                }
                o.PublisherPriority = priority
        }

        if flags&FetchFlagProperties != 0 {
                props, err := r.VarintBytes()
                if err != nil {
                        return truncated(err)
                }
                o.Properties = props
        }

        // Object Payload Length (vi64) + Object Payload (..) per §11.4.4 Figure 28.
        payload, err := r.VarintBytes()
        if err != nil {
                return truncated(err)
        }
        o.ObjectPayload = payload

        return nil
}

// IsEndOfNonExistentRange reports whether this is an End of Non-Existent Range
// marker (0x8C): the Objects it covers are known not to exist.
func (o *FetchObject) IsEndOfNonExistentRange() bool {
        return o.SerializationFlags == FetchEndOfNonExistentRange
}

// IsEndOfUnknownRange reports whether this is an End of Unknown Range marker
// (0x10C): no source could vouch for the Objects it covers either way.
func (o *FetchObject) IsEndOfUnknownRange() bool {
        return o.SerializationFlags == FetchEndOfUnknownRange
}

// IsEndOfTimedOutRange reports whether this is an End of Timed-Out Range marker
// (0x20C): the Objects it covers were abandoned when FILL_TIMEOUT expired
// (§10.2.5), as opposed to being known absent (0x8C) or of unknown status
// (0x10C).
func (o *FetchObject) IsEndOfTimedOutRange() bool {
        return o.SerializationFlags == FetchEndOfTimedOutRange
}

// IsEndOfRange reports whether this is any §11.4.4.2 end-of-range marker rather
// than a serialized Object.
func (o *FetchObject) IsEndOfRange() bool { return isEndOfRange(o.SerializationFlags) }

// IsDatagram reports whether the Datagram bit (0x40) is set: the object was
// published with Forwarding Preference "Datagram" and carries no Subgroup ID.
func (o *FetchObject) IsDatagram() bool {
        return o.SerializationFlags&FetchFlagDatagram != 0
}

// hasSubgroupIDField reports whether a Subgroup ID field is present on the
// wire: the subgroup mode must be Explicit AND the Datagram bit must be
// clear — §11.4.4.1 Table 9 says 0x40 means "ignore the two least
// significant bits". Shared by Append and Parse so the encoder and decoder
// cannot disagree on field presence.
func (o *FetchObject) hasSubgroupIDField() bool {
        return o.SubgroupMode() == FetchSubgroupIDExplicit && !o.IsDatagram()
}

// SubgroupMode returns the subgroup ID encoding mode from the two LSBs.
func (o *FetchObject) SubgroupMode() FetchSubgroupIDMode {
        return FetchSubgroupIDMode(o.SerializationFlags & FetchFlagSubgroupIDMode)
}

// Validate checks the fetch object for protocol violations.
func (o *FetchObject) Validate() error {
        // Note: 0x40 with non-zero subgroup-mode LSBs stays valid — the publisher
        // only SHOULD zero them and the subscriber MUST ignore them (§11.4.4.1),
        // so rejecting the combination would itself be non-conformant.
        return checkFetchFlags(o.SerializationFlags)
}

// ErrInvalidFetchFlags is a Serialization Flags value of 128 or more that is
// not an End of Range marker: "Any other value is a PROTOCOL_VIOLATION"
// (§11.4.4). [FetchObject.Parse] returns it right after the flags.
var ErrInvalidFetchFlags = errors.New("moqt/message: invalid fetch object serialization flags")

func checkFetchFlags(flags uint64) error {
        if flags >= 128 && !isEndOfRange(flags) {
                return fmt.Errorf("%w 0x%X", ErrInvalidFetchFlags, flags)
        }
        return nil
}

package message

import (
        "errors"
        "fmt"
        "slices"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// fillParamsAllowed is Table 6 of §10.2.15: the only parameters that may
// appear inside FILL_PARAMETERS. Note TRACK_PROPERTY_FILTER (0x29) is absent —
// a fill is scoped to Objects, so only the Object-scoped filters carry over.
var fillParamsAllowed = []ParamID{
        ParamFillTimeout,
        ParamSubscriberPriority,
        ParamLocationFilter,
        ParamGroupOrder,
        ParamSubgroupFilter,
        ParamObjectIDFilter,
        ParamPriorityFilter,
        ParamObjectPropertyFilter,
}

// FillParametersParam builds FILL_PARAMETERS (§10.2.15) from the parameters
// that apply to the fill fetch stream. Its presence on a SUBSCRIBE or
// REQUEST_UPDATE is what asks the publisher to open a fill fetch stream
// (§5.1.3) — an empty inner list still requests one, filling the whole track
// up to Largest Object.
//
// The value is a nested parameter sequence: it is a separate parameter scope,
// so a type may appear both here and in the enclosing message (§10.2.15).
func FillParametersParam(inner Parameters) Parameter {
        var w wire.Writer
        inner.append(&w)
        return BytesParam(ParamFillParameters, w.Bytes())
}

// FillParametersFromParam extracts and parses FILL_PARAMETERS from a parameter
// list. ok is false when the parameter is absent, which per §5.1.3 means no
// fill fetch stream is requested — distinct from a present-but-empty list.
//
// An inner parameter outside Table 6 is an error the caller MUST map to a
// session-level PROTOCOL_VIOLATION (§10.2.15).
func FillParametersFromParam(ps Parameters) (inner Parameters, ok bool, err error) {
        p, found := ps.Find(ParamFillParameters)
        if !found {
                return nil, false, nil
        }
        if err := inner.parse(wire.NewReader(p.Bytes)); err != nil {
                // §10.2: an unknown parameter, or a Type past 2^64-1, is a
                // PROTOCOL_VIOLATION here as in any message; anything else is a
                // value that does not parse (§1.4.3).
                if errors.Is(err, ErrUnknownParameter) || errors.Is(err, errParamTypeOverflow) {
                        return nil, true, fmt.Errorf("moqt/message: FILL_PARAMETERS: %w", err)
                }
                return nil, true, fmt.Errorf("%w: FILL_PARAMETERS: %w", ErrValueFormatting, err)
        }
        for _, ip := range inner {
                if !slices.Contains(fillParamsAllowed, ip.Type) {
                        return nil, true, fmt.Errorf(
                                "moqt/message: %s not allowed inside FILL_PARAMETERS (PROTOCOL_VIOLATION §10.2.15)", ip.Type)
                }
        }
        // A separate message's parameters (§10.2.15): §10.2's duplicate rule applies.
        if t, dup := inner.firstDuplicate(); dup {
                return nil, true, fmt.Errorf("moqt/message: duplicate %s inside FILL_PARAMETERS (PROTOCOL_VIOLATION §10.2)", t)
        }
        return inner, true, nil
}

// IncludePropertiesParam builds INCLUDE_PROPERTIES (§10.2.21): whether the
// response should carry Track Properties. The default is 1, so this is only
// worth sending to suppress them.
func IncludePropertiesParam(include bool) Parameter {
        var v uint8
        if include {
                v = 1
        }
        return ByteParam(ParamIncludeProperties, v)
}

// IncludePropertiesFromParam reads INCLUDE_PROPERTIES (§10.2.21) from a
// parameter list, defaulting to true when absent. A value outside {0, 1} is an
// error the caller MUST map to a session-level PROTOCOL_VIOLATION.
func IncludePropertiesFromParam(ps Parameters) (bool, error) {
        p, ok := ps.Find(ParamIncludeProperties)
        if !ok {
                return true, nil
        }
        switch p.Byte {
        case 0:
                return false, nil
        case 1:
                return true, nil
        default:
                return false, fmt.Errorf(
                        "moqt/message: INCLUDE_PROPERTIES value %d outside {0,1} (PROTOCOL_VIOLATION §10.2.21)", p.Byte)
        }
}

package message

import (
        "errors"
        "fmt"
        "math"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// LocationFilter is the LOCATION_FILTER parameter value from §5.1.2.
//
// Wire format (the enclosing parameter is length-prefixed, and that Length is
// what selects how many of the four optional fields are present):
//
//        LOCATION_FILTER Parameter {
//          Parameter Type (vi64) = 0x21,
//          Length (vi64),
//          [StartGroup (vi64),]
//          [StartObject (vi64),]
//          [EndGroupDelta (vi64),]
//          [EndObject (vi64),]
//        }
//
// draft-20 replaced draft-19's Filter Type enum (NextGroupStart /
// LargestObject / AbsoluteStart / AbsoluteRange) with this positional
// encoding, so the field count *is* the discriminant:
//
//        0 fields  unfiltered — and, in REQUEST_UPDATE, removes the filter
//        1 field   StartGroup is RELATIVE: start = {Largest.Group + 1 - StartGroup, 0}
//        2 fields  {0,0} means the Next Object; otherwise an absolute start
//        3 fields  absolute start, end group = StartGroup + EndGroupDelta
//        4 fields  ...plus an explicit last Object in the end group
//
// The range is inclusive at both ends. An omitted end is open-ended on a
// subscription and means Largest Object on a Fetch (§5.1.2).
type LocationFilter struct {
        // Fields is how many of the four optional vi64s were on the wire (0-4).
        // It selects the interpretation of the rest, so it is part of the value
        // rather than a decoding artifact.
        Fields int

        StartGroup    uint64
        StartObject   uint64
        EndGroupDelta uint64
        EndObject     uint64
}

// Unfiltered reports whether the filter selects the whole track (no fields).
func (f *LocationFilter) Unfiltered() bool { return f.Fields == 0 }

// RelativeStart reports whether StartGroup counts back from the Next Group
// rather than naming an absolute Group (the one-field form).
func (f *LocationFilter) RelativeStart() bool { return f.Fields == 1 }

// NextObject reports whether the filter starts at the Object after Largest
// Object — the two-field all-zero form, draft-19's LargestObject filter.
func (f *LocationFilter) NextObject() bool {
        return f.Fields == 2 && f.StartGroup == 0 && f.StartObject == 0
}

// HasEnd reports whether the filter bounds the end of the range.
func (f *LocationFilter) HasEnd() bool { return f.Fields >= 3 }

// HasEndObject reports whether the filter names a last Object in the end
// Group. When false but HasEnd is true, every Object in the end Group passes.
func (f *LocationFilter) HasEndObject() bool { return f.Fields == 4 }

// Validate enforces the §5.1.2 rules the decoder cannot: a field count in
// range, and the end-group sum staying inside the 64-bit Group space ("If
// StartGroup + EndGroupDelta exceeds 2^64 - 1, the endpoint MUST close the
// session with a PROTOCOL_VIOLATION"). Callers map a non-nil error to
// PROTOCOL_VIOLATION.
//
// Note the asymmetry with a relative start, which §5.1.2 clamps rather than
// rejects — see [LocationFilter.Start].
func (f *LocationFilter) Validate() error {
        if f.Fields < 0 || f.Fields > 4 {
                return fmt.Errorf("moqt/message: LOCATION_FILTER has %d fields, want 0-4 (§5.1.2)", f.Fields)
        }
        if f.HasEnd() && f.StartGroup > math.MaxUint64-f.EndGroupDelta {
                return fmt.Errorf("%w (start=%d delta=%d)", errEndGroupOverflow, f.StartGroup, f.EndGroupDelta)
        }
        return nil
}

// errEndGroupOverflow is a LOCATION_FILTER whose StartGroup + EndGroupDelta
// exceeds 2^64 - 1: a PROTOCOL_VIOLATION (§5.1.2), unlike a value that does
// not parse ([ErrValueFormatting]).
var errEndGroupOverflow = errors.New("moqt/message: LOCATION_FILTER end group overflow (PROTOCOL_VIOLATION §5.1.2)")

// Start resolves the first Location that passes the filter, given the
// publisher's current Largest Object. hasLargest is false before anything has
// been published on the track, which §5.1.2 pins to {0, 0}.
//
// A relative StartGroup is clamped, not rejected (§5.1.2): a computed absolute
// group below 0 is set to 0, and one above 2^64 - 1 is set to 2^64 - 1.
func (f *LocationFilter) Start(largest Location, hasLargest bool) Location {
        switch {
        case f.Unfiltered():
                return Location{}

        case f.RelativeStart():
                // {Largest.Group + 1 - StartGroup, 0}, clamped at both ends.
                if !hasLargest {
                        return Location{}
                }
                if f.StartGroup > largest.Group {
                        // Largest.Group + 1 - StartGroup would go below 0.
                        return Location{}
                }
                if largest.Group == math.MaxUint64 && f.StartGroup == 0 {
                        // Largest.Group + 1 would exceed 2^64 - 1.
                        return Location{Group: math.MaxUint64}
                }
                return Location{Group: largest.Group + 1 - f.StartGroup}

        case f.NextObject():
                if !hasLargest {
                        return Location{}
                }
                if largest.Object == math.MaxUint64 {
                        // No Object can follow it within this Group.
                        return Location{Group: largest.Group, Object: math.MaxUint64}
                }
                return Location{Group: largest.Group, Object: largest.Object + 1}

        default:
                return Location{Group: f.StartGroup, Object: f.StartObject}
        }
}

// End resolves the last Location that passes the filter. ok is false when the
// filter is open-ended, which on a subscription means "no end" and on a Fetch
// means Largest Object (§5.1.2) — a distinction the caller owns.
//
// Call Validate first: an unvalidated end-group sum can wrap.
func (f *LocationFilter) End() (loc Location, ok bool) {
        if !f.HasEnd() {
                return Location{}, false
        }
        end := Location{Group: f.StartGroup + f.EndGroupDelta, Object: math.MaxUint64}
        if f.HasEndObject() {
                end.Object = f.EndObject
        }
        return end, true
}

// Matches reports whether the Object at loc passes this filter on a
// subscription, given the publisher's Largest Object. Both ends are inclusive
// and an absent end is open-ended (§5.1.2).
func (f *LocationFilter) Matches(loc Location, largest Location, hasLargest bool) bool {
        if loc.Less(f.Start(largest, hasLargest)) {
                return false
        }
        if end, ok := f.End(); ok && end.Less(loc) {
                return false
        }
        return true
}

// Append serialises the filter's fields to w. The caller writes the enclosing
// parameter's Length, which is what tells the peer how many fields follow.
func (f *LocationFilter) Append(w *wire.Writer) {
        if f.Fields >= 1 {
                w.Varint(f.StartGroup)
        }
        if f.Fields >= 2 {
                w.Varint(f.StartObject)
        }
        if f.Fields >= 3 {
                w.Varint(f.EndGroupDelta)
        }
        if f.Fields >= 4 {
                w.Varint(f.EndObject)
        }
}

// Parse deserialises a filter from r, consuming every remaining byte: r must
// be bounded to the LOCATION_FILTER parameter's value, since the byte count is
// the only thing that distinguishes the five forms (§5.1.2).
func (f *LocationFilter) Parse(r *wire.Reader) error {
        *f = LocationFilter{}
        dst := [4]*uint64{&f.StartGroup, &f.StartObject, &f.EndGroupDelta, &f.EndObject}
        for !r.Empty() {
                if f.Fields == len(dst) {
                        return fmt.Errorf("moqt/message: LOCATION_FILTER has %d trailing bytes after 4 fields (§5.1.2)",
                                r.Remaining())
                }
                v, err := r.Varint()
                if err != nil {
                        return fmt.Errorf("moqt/message: LOCATION_FILTER field %d: %w", f.Fields, err)
                }
                *dst[f.Fields] = v
                f.Fields++
        }
        return f.Validate()
}

// Bytes serialises the filter's fields to a fresh slice, for use as the
// LOCATION_FILTER parameter value.
func (f *LocationFilter) Bytes() []byte {
        var w wire.Writer
        f.Append(&w)
        return w.Bytes()
}

// ParseLocationFilter deserialises a LocationFilter from a LOCATION_FILTER
// parameter value.
func ParseLocationFilter(raw []byte) (*LocationFilter, error) {
        f := &LocationFilter{}
        if err := f.Parse(wire.NewReader(raw)); err != nil {
                return nil, err
        }
        return f, nil
}

// UnfilteredFilter returns a zero-length LOCATION_FILTER parameter (§5.1.2):
// the whole track. On REQUEST_UPDATE it removes an existing filter.
func UnfilteredFilter() Parameter {
        return LocationFilterParam(&LocationFilter{})
}

// NextObjectFilter returns a LOCATION_FILTER parameter (§5.1.2) starting at
// the Object after the publisher's Largest Object — the live edge, and the
// filter to pair with a fill so each Object arrives exactly once (§5.1.3).
//
// This is draft-19's LargestObject filter.
func NextObjectFilter() Parameter {
        return LocationFilterParam(&LocationFilter{Fields: 2})
}

// RelativeStartFilter returns an open-ended LOCATION_FILTER parameter (§5.1.2)
// starting groupsBack groups before the Next Group: 0 is the Next Group (which
// is draft-19's NextGroupStart filter), 1 the current group, N the group N-1
// before the current one.
func RelativeStartFilter(groupsBack uint64) Parameter {
        return LocationFilterParam(&LocationFilter{Fields: 1, StartGroup: groupsBack})
}

// AbsoluteStartFilter returns an open-ended LOCATION_FILTER parameter (§5.1.2)
// starting at an explicit Location. A start of {0, 0} is equivalent to
// unfiltered, and is encoded that way — the two-field all-zero form is the
// Next Object filter, not an absolute {0, 0}.
func AbsoluteStartFilter(start Location) Parameter {
        if start == (Location{}) {
                return UnfilteredFilter()
        }
        return LocationFilterParam(&LocationFilter{
                Fields:      2,
                StartGroup:  start.Group,
                StartObject: start.Object,
        })
}

// AbsoluteRangeFilter returns a LOCATION_FILTER parameter (§5.1.2) covering
// start through the end of group (start.Group + endGroupDelta), inclusive.
func AbsoluteRangeFilter(start Location, endGroupDelta uint64) Parameter {
        return LocationFilterParam(&LocationFilter{
                Fields:        3,
                StartGroup:    start.Group,
                StartObject:   start.Object,
                EndGroupDelta: endGroupDelta,
        })
}

// AbsoluteRangeObjectFilter returns a LOCATION_FILTER parameter (§5.1.2)
// covering the inclusive range start..{start.Group + endGroupDelta, endObject}.
func AbsoluteRangeObjectFilter(start Location, endGroupDelta, endObject uint64) Parameter {
        return LocationFilterParam(&LocationFilter{
                Fields:        4,
                StartGroup:    start.Group,
                StartObject:   start.Object,
                EndGroupDelta: endGroupDelta,
                EndObject:     endObject,
        })
}

package message

import (
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// MaxGoawayURIBytes is the maximum New Session URI length per §10.4.
const MaxGoawayURIBytes = 8192

// Goaway is the GOAWAY message (§10.4).
type Goaway struct {
        NewSessionURI []byte
        Timeout       uint64
}

func (m *Goaway) Type() Type { return TypeGoaway }

func (m *Goaway) Append(w *wire.Writer) {
        w.VarintBytes(m.NewSessionURI)
        w.Varint(m.Timeout)
}

func (m *Goaway) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.VarintBytes(&m.NewSessionURI)
        if err := s.Err(); err != nil {
                return err
        }
        if len(m.NewSessionURI) > MaxGoawayURIBytes {
                return fmt.Errorf("moqt/message: GOAWAY URI length %d exceeds %d", len(m.NewSessionURI), MaxGoawayURIBytes)
        }
        s.Varint(&m.Timeout)
        if err := s.Err(); err != nil {
                return err
        }
        return nil
}

package message

import (
        "math/rand/v2"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// GREASE (Generate Random Extensions And Sustain Extensibility) support per
// §14 of draft-ietf-moq-transport-20 and RFC 9170 §3.3.
//
// GREASE values follow the pattern 0x7F * N + 0x9D for non-negative integer
// values of N (that is, 0x9D, 0x11C, 0x19B, ..., 0x3FFFFFFFFFFFFFDE).
//
// §14 reserves GREASE values in the Setup Options, Properties, error-code, and
// Auth Token Type registries: implementations "MUST handle unknown values
// gracefully", and endpoints "MUST NOT close the session solely because they
// received an unknown value".

// greaseBase and greaseStep define the GREASE value pattern: base + step*N.
const (
        greaseBase uint64 = 0x9D
        greaseStep uint64 = 0x7F
)

// maxGreaseN is the largest N in §14's reserved range, whose last value is
// 0x3fffffffffffffde.
// 0x7F * N + 0x9D ≤ 0x3FFFFFFFFFFFFFFF  →  N ≤ (0x3FFFFFFFFFFFFFFF - 0x9D) / 0x7F.
const maxGreaseN uint64 = (0x3FFFFFFFFFFFFFFF - greaseBase) / greaseStep

// GreaseValue returns a random GREASE value from the reserved range. The
// returned value is suitable for use as a Setup Option type, Property type,
// or error code. Each call returns a fresh random value.
func GreaseValue() uint64 {
        //nolint:gosec // G404: GREASE values are deliberately non-cryptographic; randomness only spreads coverage.
        n := rand.Uint64N(maxGreaseN + 1)
        return greaseBase + greaseStep*n
}

// GreaseSetupOption returns a KVPair with a random GREASE type suitable for
// inclusion in a SETUP message's option list. Per §1.4.3, even types carry a
// varint value and odd types carry bytes; the GREASE pattern produces both
// parities, so the helper picks a value and fills the appropriate field with
// a small random payload.
func GreaseSetupOption() wire.KVPair {
        v := GreaseValue()
        kv := wire.KVPair{Type: v}
        if kv.IsBytes() {
                // Odd type → length-prefixed bytes. Send a small random payload.
                kv.ByteVal = []byte{byte(rand.UintN(256))} //nolint:gosec // G404: non-cryptographic GREASE payload by design.
        } else {
                // Even type → varint. Send a small random value.
                kv.IntVal = rand.Uint64N(256) //nolint:gosec // G404: non-cryptographic GREASE payload by design.
        }
        return kv
}

package message

import "cmp"

// Location represents a track location per §1.4.2.
type Location struct {
        Group  uint64
        Object uint64
}

// Compare returns -1, 0, or +1 according to whether l sorts before, equal
// to, or after other in the (Group, Object) lexicographic order. This is
// the total order MoQT uses for §10.2.17 (LARGEST_OBJECT monotonicity),
// §11.2 (intra-track Object ordering), and Fetch/Cache range scans
// (§5.1.2).
//
// The signature matches [cmp.Compare] so callers can pass
// Location.Compare directly to [slices.SortFunc] and
// [slices.BinarySearchFunc].
func (l Location) Compare(other Location) int {
        return cmp.Or(
                cmp.Compare(l.Group, other.Group),
                cmp.Compare(l.Object, other.Object),
        )
}

// Less reports whether l comes strictly before other in the (Group, Object)
// order described on [Location.Compare].
func (l Location) Less(other Location) bool { return l.Compare(other) < 0 }

package message

import "github.com/floatdrop/moq-go/pkg/moqt/wire"

// PublishNamespace is the PUBLISH_NAMESPACE message (§10.16). It announces
// that the publisher will publish tracks within a namespace.
type PublishNamespace struct {
        RequestID  uint64
        Namespace  wire.TrackNamespace
        Parameters Parameters
}

// Type returns the wire type ID for PUBLISH_NAMESPACE.
func (m *PublishNamespace) Type() Type             { return TypePublishNamespace }
func (m *PublishNamespace) GetRequestID() uint64   { return m.RequestID }
func (m *PublishNamespace) SetRequestID(id uint64) { m.RequestID = id }

// Append serializes the PUBLISH_NAMESPACE message to w.
func (m *PublishNamespace) Append(w *wire.Writer) {
        w.Varint(m.RequestID)
        w.TrackNamespace(m.Namespace)
        m.Parameters.append(w)
}

// Parse deserializes the PUBLISH_NAMESPACE message from r.
func (m *PublishNamespace) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.Varint(&m.RequestID)
        s.TrackNamespace(&m.Namespace)
        if err := s.Err(); err != nil {
                return err
        }
        return m.Parameters.parse(r)
}

// Namespace is the NAMESPACE message (§10.17). It announces a track
// namespace suffix on a PUBLISH_NAMESPACE or SUBSCRIBE_NAMESPACE request stream.
type Namespace struct {
        TrackNamespaceSuffix wire.TrackNamespace
}

// Type returns the wire type ID for NAMESPACE.
func (m *Namespace) Type() Type {
        return TypeNamespace
}

// Append serializes the NAMESPACE message to w.
func (m *Namespace) Append(w *wire.Writer) {
        w.TrackNamespace(m.TrackNamespaceSuffix)
}

// Parse deserializes the NAMESPACE message from r.
func (m *Namespace) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.TrackNamespace(&m.TrackNamespaceSuffix)
        return s.Err()
}

// NamespaceDone is the NAMESPACE_DONE message (§10.18). It signals that
// no more tracks will be published within a namespace.
type NamespaceDone struct {
        TrackNamespaceSuffix wire.TrackNamespace
}

// Type returns the wire type ID for NAMESPACE_DONE.
func (m *NamespaceDone) Type() Type {
        return TypeNamespaceDone
}

// Append serializes the NAMESPACE_DONE message to w.
func (m *NamespaceDone) Append(w *wire.Writer) {
        w.TrackNamespace(m.TrackNamespaceSuffix)
}

// Parse deserializes the NAMESPACE_DONE message from r.
func (m *NamespaceDone) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.TrackNamespace(&m.TrackNamespaceSuffix)
        return s.Err()
}

// SubscribeNamespace is the SUBSCRIBE_NAMESPACE message (§10.19). It
// subscribes to all tracks within a namespace prefix.
type SubscribeNamespace struct {
        RequestID            uint64
        TrackNamespacePrefix wire.TrackNamespace
        Parameters           Parameters
}

// Type returns the wire type ID for SUBSCRIBE_NAMESPACE.
func (m *SubscribeNamespace) Type() Type             { return TypeSubscribeNamespace }
func (m *SubscribeNamespace) GetRequestID() uint64   { return m.RequestID }
func (m *SubscribeNamespace) SetRequestID(id uint64) { m.RequestID = id }

// Append serializes the SUBSCRIBE_NAMESPACE message to w.
func (m *SubscribeNamespace) Append(w *wire.Writer) {
        w.Varint(m.RequestID)
        w.TrackNamespace(m.TrackNamespacePrefix)
        m.Parameters.append(w)
}

// Parse deserializes the SUBSCRIBE_NAMESPACE message from r.
func (m *SubscribeNamespace) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.Varint(&m.RequestID)
        s.TrackNamespace(&m.TrackNamespacePrefix)
        if err := s.Err(); err != nil {
                return err
        }
        return m.Parameters.parse(r)
}

// SubscribeTracks is the SUBSCRIBE_TRACKS message (§10.20). It subscribes
// to all tracks within a namespace prefix.
type SubscribeTracks struct {
        RequestID            uint64
        TrackNamespacePrefix wire.TrackNamespace
        Parameters           Parameters
}

// Type returns the wire type ID for SUBSCRIBE_TRACKS.
func (m *SubscribeTracks) Type() Type             { return TypeSubscribeTracks }
func (m *SubscribeTracks) GetRequestID() uint64   { return m.RequestID }
func (m *SubscribeTracks) SetRequestID(id uint64) { m.RequestID = id }

// Append serializes the SUBSCRIBE_TRACKS message to w.
func (m *SubscribeTracks) Append(w *wire.Writer) {
        w.Varint(m.RequestID)
        w.TrackNamespace(m.TrackNamespacePrefix)
        m.Parameters.append(w)
}

// Parse deserializes the SUBSCRIBE_TRACKS message from r.
func (m *SubscribeTracks) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.Varint(&m.RequestID)
        s.TrackNamespace(&m.TrackNamespacePrefix)
        if err := s.Err(); err != nil {
                return err
        }
        return m.Parameters.parse(r)
}

// PublishSkipped is the PUBLISH_SKIPPED message (§10.21). It signals that a
// specific track's Subscription was not created for this SUBSCRIBE_TRACKS.
type PublishSkipped struct {
        TrackNamespaceSuffix wire.TrackNamespace
        TrackName            []byte
}

// Type returns the wire type ID for PUBLISH_SKIPPED.
func (m *PublishSkipped) Type() Type {
        return TypePublishSkipped
}

// Append serializes the PUBLISH_SKIPPED message to w.
func (m *PublishSkipped) Append(w *wire.Writer) {
        w.TrackNamespace(m.TrackNamespaceSuffix)
        w.VarintBytes(m.TrackName)
}

// Parse deserializes the PUBLISH_SKIPPED message from r.
func (m *PublishSkipped) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.TrackNamespace(&m.TrackNamespaceSuffix)
        s.VarintBytes(&m.TrackName)
        return s.Err()
}

package message

import (
        "errors"
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// CheckObjectProperties reports whether an Object's raw Properties make its
// track malformed (§2.4.2) by a rule decidable from the Object alone: an
// unparsable pair or nested Immutable Properties (§12.7), a repeated Prior
// Group/Object ID Gap or one exceeding the Object's ID (§12.8, §12.9), or a
// Mandatory Track Property (§2.5.1). A repeated Immutable Properties is
// treated as malformed too, an interpretation: §12.7 forbids it but does not
// list it as malformed. Rules that need earlier Objects are not checked.
//
// Must not allocate: per-Object path.
func CheckObjectProperties(raw []byte, groupID, objectID uint64) error {
        var c objectPropertiesCheck
        if err := c.walk(raw, false); err != nil {
                return err
        }
        if c.groupGap > groupID {
                return fmt.Errorf("moqt/message: Prior Group ID Gap %d exceeds Group ID %d (§12.8)", c.groupGap, groupID)
        }
        if c.objectGap > objectID {
                return fmt.Errorf("moqt/message: Prior Object ID Gap %d exceeds Object ID %d (§12.9)", c.objectGap, objectID)
        }
        return nil
}

// objectPropertiesCheck carries the counts across the mutable list and the
// contents of Immutable Properties.
type objectPropertiesCheck struct {
        immutables, groupGaps, objectGaps int
        groupGap, objectGap               uint64
}

var errTooManyInstances = errors.New("more than one instance")

func (c *objectPropertiesCheck) walk(raw []byte, nested bool) error {
        r := wire.NewReader(raw)
        var prev uint64
        for !r.Empty() {
                kv, next, err := r.KVPairView(prev)
                if err != nil {
                        return fmt.Errorf("moqt/message: object properties: %w", err)
                }
                prev = next
                switch {
                case kv.Type == PropertyImmutableProperties:
                        if nested {
                                return errors.New("moqt/message: Immutable Properties inside Immutable Properties (§12.7)")
                        }
                        if c.immutables++; c.immutables > 1 {
                                // An interpretation, see CheckObjectProperties.
                                return fmt.Errorf("moqt/message: Immutable Properties: %w (§12.7, §2.4.2)", errTooManyInstances)
                        }
                        if err := c.walk(kv.ByteVal, true); err != nil {
                                return err
                        }
                case kv.Type == PropertyPriorGroupIDGap:
                        if c.groupGaps++; c.groupGaps > 1 {
                                return fmt.Errorf("moqt/message: Prior Group ID Gap: %w (§12.8)", errTooManyInstances)
                        }
                        c.groupGap = kv.IntVal
                case kv.Type == PropertyPriorObjectIDGap:
                        if c.objectGaps++; c.objectGaps > 1 {
                                return fmt.Errorf("moqt/message: Prior Object ID Gap: %w (§12.9)", errTooManyInstances)
                        }
                        c.objectGap = kv.IntVal
                case IsMandatoryTrackProperty(kv.Type):
                        return fmt.Errorf("moqt/message: Mandatory Track Property %#x as an Object Property (§2.5.1)", kv.Type)
                }
        }
        return nil
}

// PriorGaps are an Object's Prior Group ID Gap (§12.8) and Prior Object ID
// Gap (§12.9), each if present.
type PriorGaps struct {
        Group, Object       uint64
        HasGroup, HasObject bool
}

// ObjectPriorGaps returns the Prior Group and Object ID Gaps in an Object's raw
// Properties, searching Immutable Properties too (§12.7). Properties that
// [CheckObjectProperties] rejects for a reason it can see without the Object's
// IDs carry none.
//
// Must not allocate: per-Object path.
func ObjectPriorGaps(raw []byte) PriorGaps {
        var c objectPropertiesCheck
        if c.walk(raw, false) != nil {
                return PriorGaps{}
        }
        return PriorGaps{
                Group: c.groupGap, Object: c.objectGap,
                HasGroup: c.groupGaps == 1, HasObject: c.objectGaps == 1,
        }
}

// PriorObjectIDGap is [ObjectPriorGaps] reduced to the Prior Object ID Gap.
func PriorObjectIDGap(raw []byte) (uint64, bool) {
        g := ObjectPriorGaps(raw)
        return g.Object, g.HasObject
}

// ImmutableProperties returns the value of the Immutable Properties (§12.7) in
// an Object's raw Properties, serialized as received, and whether there is one.
// Properties that do not parse carry none.
//
// Must not allocate.
func ImmutableProperties(raw []byte) ([]byte, bool) {
        r := wire.NewReader(raw)
        var prev uint64
        for !r.Empty() {
                kv, next, err := r.KVPairView(prev)
                if err != nil {
                        return nil, false
                }
                if kv.Type == PropertyImmutableProperties {
                        return kv.ByteVal, true
                }
                prev = next
        }
        return nil, false
}

package message

import (
        "errors"
        "fmt"
        "math/bits"
        "slices"
        "strings"
)

// ErrUnknownParameter is wrapped by the parse error for an undefined Message
// Parameter type, a PROTOCOL_VIOLATION (§10.2).
var ErrUnknownParameter = errors.New("moqt/message: unknown parameter type")

// ParamScope is the message form a parameter block arrived in (§10.2.1). A
// REQUEST_OK takes its form from the request it answers (§10.5), and a
// REQUEST_UPDATE from the request it updates. Values are bit flags.
type ParamScope uint32

const (
        ScopeSubscribe ParamScope = 1 << iota
        ScopeSubscribeOK
        ScopePublish
        ScopePublishOK
        ScopeFetch
        ScopeFetchOK
        ScopeTrackStatus
        ScopeTrackStatusOK
        ScopeSubscribeNamespace
        ScopeSubscribeNamespaceOK
        ScopeSubscribeTracks
        ScopeSubscribeTracksOK
        ScopePublishNamespace
        ScopePublishNamespaceOK
        ScopePublishStateNotify
        ScopeRequestUpdateOK
        // ScopeUpdateFromSubscriber is a REQUEST_UPDATE on a subscription sent by
        // its subscriber, ScopeUpdateFromPublisher one sent by the publisher of a
        // PUBLISH. §5.1.4 allows the Range Filters only "from the subscriber".
        ScopeUpdateFromSubscriber
        ScopeUpdateFromPublisher
        ScopeUpdateFetch
        ScopeUpdateTrackStatus
        ScopeUpdateSubscribeNamespace
        ScopeUpdateSubscribeTracks
        ScopeUpdatePublishNamespace
)

const (
        scopeUpdateSubscription = ScopeUpdateFromSubscriber | ScopeUpdateFromPublisher
        scopeAnyUpdate          = scopeUpdateSubscription | ScopeUpdateFetch | ScopeUpdateTrackStatus |
                ScopeUpdateSubscribeNamespace | ScopeUpdateSubscribeTracks | ScopeUpdatePublishNamespace
        // §5.1.4 is ambiguous on where each filter may appear; all five share the
        // widest reading, which closes fewer sessions.
        scopeRangeFilter = ScopeSubscribe | ScopeFetch | ScopeSubscribeTracks |
                ScopeUpdateFromSubscriber | ScopeUpdateSubscribeTracks
)

// paramScopes is each parameter's "MAY appear in" list, from its definition.
// [Parameters.CheckScope] adds SUBSCRIBE_TRACKS to every SUBSCRIBE parameter
// (§10.20.1).
var paramScopes = map[ParamID]ParamScope{
        // §10.2.2
        ParamAuthorizationToken: ScopePublish | ScopeSubscribe | scopeAnyUpdate | ScopeSubscribeNamespace |
                ScopeSubscribeTracks | ScopePublishNamespace | ScopeTrackStatus | ScopeFetch,
        // §10.2.3, §10.2.4
        ParamSubgroupDeliveryTimeout: ScopeSubscribe | ScopePublish | scopeAnyUpdate,
        ParamObjectDeliveryTimeout:   ScopeSubscribe | ScopePublish | scopeAnyUpdate,
        // §10.2.5 (also inside FILL_PARAMETERS, a separate scope)
        ParamFillTimeout: ScopeFetch,
        // §10.2.6
        ParamRendezvousTimeout: ScopeSubscribe,
        // §10.2.7
        ParamSubscriberPriority: ScopeSubscribe | ScopePublish | ScopeFetch | scopeUpdateSubscription | ScopeUpdateFetch,
        // §10.2.8
        ParamGroupOrder: ScopeSubscribe | ScopePublish | ScopeSubscribeTracks | ScopeFetch,
        // §10.2.9
        ParamLocationFilter: ScopeFetch | ScopeSubscribe | ScopePublish | scopeUpdateSubscription |
                ScopePublishStateNotify,
        // §10.2.10–§10.2.14, §5.1.4
        ParamSubgroupFilter:       scopeRangeFilter,
        ParamObjectIDFilter:       scopeRangeFilter,
        ParamPriorityFilter:       scopeRangeFilter,
        ParamObjectPropertyFilter: scopeRangeFilter,
        ParamTrackPropertyFilter:  scopeRangeFilter,
        // §10.2.15
        ParamFillParameters: ScopeSubscribe | scopeUpdateSubscription,
        // §10.2.16
        ParamExpires: ScopeSubscribeOK | ScopePublish | ScopePublishOK | ScopeSubscribeNamespaceOK |
                ScopeSubscribeTracksOK | ScopePublishNamespaceOK | ScopeRequestUpdateOK,
        // §10.2.17
        ParamLargestObject: ScopeSubscribeOK | ScopePublish | ScopeRequestUpdateOK | ScopeTrackStatusOK |
                ScopePublishStateNotify,
        // §10.2.18
        ParamForward: ScopeSubscribe | scopeUpdateSubscription | ScopeUpdateSubscribeTracks | ScopePublish |
                ScopeSubscribeTracks | ScopePublishStateNotify,
        // §10.2.19
        ParamNewGroupRequest: ScopeSubscribe | scopeUpdateSubscription,
        // §10.2.20
        ParamTrackNamespacePrefix: ScopeUpdateSubscribeNamespace | ScopeUpdateSubscribeTracks,
        // §10.2.21
        ParamIncludeProperties: ScopeSubscribe | ScopeTrackStatus | ScopeFetch | ScopeSubscribeTracks,
}

var scopeNames = [...]string{
        "SUBSCRIBE", "SUBSCRIBE_OK", "PUBLISH", "PUBLISH_OK", "FETCH", "FETCH_OK",
        "TRACK_STATUS", "TRACK_STATUS_OK", "SUBSCRIBE_NAMESPACE", "SUBSCRIBE_NAMESPACE_OK",
        "SUBSCRIBE_TRACKS", "SUBSCRIBE_TRACKS_OK", "PUBLISH_NAMESPACE", "PUBLISH_NAMESPACE_OK",
        "PUBLISH_STATE_NOTIFY", "REQUEST_UPDATE_OK",
        "REQUEST_UPDATE from a subscriber", "REQUEST_UPDATE from a publisher", "REQUEST_UPDATE for FETCH",
        "REQUEST_UPDATE for TRACK_STATUS", "REQUEST_UPDATE for SUBSCRIBE_NAMESPACE",
        "REQUEST_UPDATE for SUBSCRIBE_TRACKS", "REQUEST_UPDATE for PUBLISH_NAMESPACE",
}

func (s ParamScope) String() string {
        var names []string
        for s != 0 {
                i := bits.TrailingZeros32(uint32(s))
                if i < len(scopeNames) {
                        names = append(names, scopeNames[i])
                }
                s &^= 1 << i
        }
        if len(names) == 0 {
                return "no message"
        }
        return strings.Join(names, "|")
}

// ScopeOfRequest is the scope of a request opener's parameters (§3.3), or 0
// for a type that does not open a request.
func ScopeOfRequest(t Type) ParamScope {
        //exhaustive:ignore // every other type maps to 0
        switch t {
        case TypeSubscribe:
                return ScopeSubscribe
        case TypePublish:
                return ScopePublish
        case TypeFetch:
                return ScopeFetch
        case TypeTrackStatus:
                return ScopeTrackStatus
        case TypeSubscribeNamespace:
                return ScopeSubscribeNamespace
        case TypeSubscribeTracks:
                return ScopeSubscribeTracks
        case TypePublishNamespace:
                return ScopePublishNamespace
        default:
                return 0
        }
}

// ScopeOfResponse is the scope of the success response to a request of type
// request: SUBSCRIBE_OK, FETCH_OK, or the REQUEST_OK form §10.5 names for it.
func ScopeOfResponse(request Type) ParamScope {
        //exhaustive:ignore // every other type maps to 0
        switch request {
        case TypeSubscribe:
                return ScopeSubscribeOK
        case TypePublish:
                return ScopePublishOK
        case TypeFetch:
                return ScopeFetchOK
        case TypeTrackStatus:
                return ScopeTrackStatusOK
        case TypeSubscribeNamespace:
                return ScopeSubscribeNamespaceOK
        case TypeSubscribeTracks:
                return ScopeSubscribeTracksOK
        case TypePublishNamespace:
                return ScopePublishNamespaceOK
        case TypeRequestUpdate:
                return ScopeRequestUpdateOK
        default:
                return 0
        }
}

// ScopeOfUpdate is the scope of a REQUEST_UPDATE sent by the sender of a
// request of type request. The subscriber of a PUBLISH may send one too; its
// scope is [ScopeUpdateFromSubscriber].
func ScopeOfUpdate(request Type) ParamScope {
        //exhaustive:ignore // every other type maps to 0
        switch request {
        case TypeSubscribe:
                return ScopeUpdateFromSubscriber
        case TypePublish:
                return ScopeUpdateFromPublisher
        case TypeFetch:
                return ScopeUpdateFetch
        case TypeTrackStatus:
                return ScopeUpdateTrackStatus
        case TypeSubscribeNamespace:
                return ScopeUpdateSubscribeNamespace
        case TypeSubscribeTracks:
                return ScopeUpdateSubscribeTracks
        case TypePublishNamespace:
                return ScopeUpdatePublishNamespace
        default:
                return 0
        }
}

// ParamScopeError reports a Message Parameter in a message form its
// definition does not list (§10.2.1), or repeated where it may not be
// (§10.2). Either is a PROTOCOL_VIOLATION.
type ParamScopeError struct {
        Type      ParamID
        Scope     ParamScope
        Duplicate bool
}

func (e *ParamScopeError) Error() string {
        if e.Duplicate {
                return fmt.Sprintf("moqt/message: duplicate %s in %s (PROTOCOL_VIOLATION §10.2)", e.Type, e.Scope)
        }
        return fmt.Sprintf("moqt/message: %s not allowed in %s (PROTOCOL_VIOLATION §10.2.1)", e.Type, e.Scope)
}

// ErrValueFormatting marks a parameter whose value does not match the
// serialization its type defines. §1.4.3: the receiver "MUST close the session
// with error code KEY_VALUE_FORMATTING_ERROR".
var ErrValueFormatting = errors.New("moqt/message: value does not match its type's serialization")

// CheckScope reports the first parameter of ps not allowed in a message of the
// given scope (§10.2.1) or repeated where it may not be (§10.2), and validates
// the values a receiver must close the session for, in ps and inside its
// FILL_PARAMETERS (§10.2.15): GROUP_ORDER (§10.2.8), FORWARD (§10.2.18),
// INCLUDE_PROPERTIES (§10.2.21) and LOCATION_FILTER (§5.1.2). Every error is a
// session-level PROTOCOL_VIOLATION, except one wrapping [ErrValueFormatting],
// a KEY_VALUE_FORMATTING_ERROR.
func (ps Parameters) CheckScope(scope ParamScope) error {
        for _, p := range ps {
                allowed := paramScopes[p.Type]
                // §10.20.1: SUBSCRIBE parameters are valid in SUBSCRIBE_TRACKS.
                if allowed&ScopeSubscribe != 0 {
                        allowed |= ScopeSubscribeTracks
                }
                if allowed&scope == 0 {
                        return &ParamScopeError{Type: p.Type, Scope: scope}
                }
        }
        if t, dup := ps.firstDuplicate(); dup {
                return &ParamScopeError{Type: t, Scope: scope, Duplicate: true}
        }
        inner, _, err := FillParametersFromParam(ps)
        if err != nil {
                return err
        }
        if err := inner.checkValues(); err != nil {
                return fmt.Errorf("FILL_PARAMETERS: %w", err)
        }
        return ps.checkValues()
}

// checkValues validates the values in ps a receiver must close the session
// for; see [Parameters.CheckScope].
func (ps Parameters) checkValues() error {
        if p, ok := ps.Find(ParamGroupOrder); ok &&
                GroupOrder(p.Byte) != GroupOrderAscending && GroupOrder(p.Byte) != GroupOrderDescending {
                return fmt.Errorf("moqt/message: GROUP_ORDER value %d outside {1,2} (PROTOCOL_VIOLATION §10.2.8)", p.Byte)
        }
        if p, ok := ps.Find(ParamForward); ok && p.Byte > 1 {
                return fmt.Errorf("moqt/message: FORWARD value %d outside {0,1} (PROTOCOL_VIOLATION §10.2.18)", p.Byte)
        }
        if _, err := IncludePropertiesFromParam(ps); err != nil {
                return err
        }
        switch _, err := LocationFilterFromParam(ps); {
        case err == nil, errors.Is(err, errEndGroupOverflow):
                return err
        default:
                return fmt.Errorf("%w: LOCATION_FILTER: %w", ErrValueFormatting, err)
        }
}

// firstDuplicate reports the first parameter type repeated in ps where its
// definition does not allow it (§10.2). The Range Filters (§5.1.4) and
// AUTHORIZATION_TOKEN (§10.2.2) may repeat.
func (ps Parameters) firstDuplicate() (ParamID, bool) {
        for i, p := range ps {
                if IsRangeFilterParam(p.Type) || p.Type == ParamAuthorizationToken {
                        continue
                }
                if slices.ContainsFunc(ps[:i], func(q Parameter) bool { return q.Type == p.Type }) {
                        return p.Type, true
                }
        }
        return 0, false
}

// ParamsOf returns m's Message Parameters, and false for a message type that
// has none.
func ParamsOf(m Message) (Parameters, bool) {
        switch m := m.(type) {
        case *Subscribe:
                return m.Parameters, true
        case *SubscribeOK:
                return m.Parameters, true
        case *Publish:
                return m.Parameters, true
        case *Fetch:
                return m.Parameters, true
        case *FetchOK:
                return m.Parameters, true
        case *TrackStatus:
                return m.Parameters, true
        case *RequestOK:
                return m.Parameters, true
        case *RequestUpdate:
                return m.Parameters, true
        case *SubscribeNamespace:
                return m.Parameters, true
        case *SubscribeTracks:
                return m.Parameters, true
        case *PublishNamespace:
                return m.Parameters, true
        case *PublishStateNotify:
                return m.Parameters, true
        default:
                return nil, false
        }
}

package message

import (
        "cmp"
        "errors"
        "fmt"
        "slices"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// ParamID is a MoQT Message Parameter type ID (§10.2). Distinct from
// SetupOption because the two code spaces overlap (parameter 0x03 vs option
// 0x03 are both AUTHORIZATION_TOKEN but in different contexts with different
// parse rules) and from session/request error codes which overlap numerically.
type ParamID uint64

// Parameter wire type IDs from §10.2.
const (
        ParamObjectDeliveryTimeout   ParamID = 0x02
        ParamAuthorizationToken      ParamID = 0x03
        ParamRendezvousTimeout       ParamID = 0x04
        ParamSubgroupDeliveryTimeout ParamID = 0x06
        ParamExpires                 ParamID = 0x08
        ParamLargestObject           ParamID = 0x09
        ParamFillTimeout             ParamID = 0x0A
        ParamForward                 ParamID = 0x10
        ParamSubscriberPriority      ParamID = 0x20
        ParamLocationFilter          ParamID = 0x21
        ParamGroupOrder              ParamID = 0x22
        ParamFillParameters          ParamID = 0x23
        // Range Filter parameters (§5.1.4, §10.2.10-14), see rangefilter.go. All
        // five are length-prefixed (KindBytes) whatever their type parity: Message
        // Parameter encodings are per definition (§10.2), not Key-Value-Pairs.
        ParamSubgroupFilter       ParamID = 0x25
        ParamObjectIDFilter       ParamID = 0x26
        ParamPriorityFilter       ParamID = 0x27
        ParamObjectPropertyFilter ParamID = 0x28
        ParamTrackPropertyFilter  ParamID = 0x29
        ParamNewGroupRequest      ParamID = 0x32
        ParamTrackNamespacePrefix ParamID = 0x34
        ParamIncludeProperties    ParamID = 0x35
)

// String returns a short name for known parameter types; unknown values render
// as hex.
func (p ParamID) String() string {
        switch p {
        case ParamObjectDeliveryTimeout:
                return "OBJECT_DELIVERY_TIMEOUT"
        case ParamAuthorizationToken:
                return "AUTHORIZATION_TOKEN"
        case ParamRendezvousTimeout:
                return "RENDEZVOUS_TIMEOUT"
        case ParamSubgroupDeliveryTimeout:
                return "SUBGROUP_DELIVERY_TIMEOUT"
        case ParamExpires:
                return "EXPIRES"
        case ParamLargestObject:
                return "LARGEST_OBJECT"
        case ParamFillTimeout:
                return "FILL_TIMEOUT"
        case ParamForward:
                return "FORWARD"
        case ParamSubscriberPriority:
                return "SUBSCRIBER_PRIORITY"
        case ParamLocationFilter:
                return "LOCATION_FILTER"
        case ParamGroupOrder:
                return "GROUP_ORDER"
        case ParamFillParameters:
                return "FILL_PARAMETERS"
        case ParamSubgroupFilter:
                return "SUBGROUP_FILTER"
        case ParamObjectIDFilter:
                return "OBJECTID_FILTER"
        case ParamPriorityFilter:
                return "PRIORITY_FILTER"
        case ParamObjectPropertyFilter:
                return "OBJECT_PROPERTY_FILTER"
        case ParamTrackPropertyFilter:
                return "TRACK_PROPERTY_FILTER"
        case ParamNewGroupRequest:
                return "NEW_GROUP_REQUEST"
        case ParamTrackNamespacePrefix:
                return "TRACK_NAMESPACE_PREFIX"
        case ParamIncludeProperties:
                return "INCLUDE_PROPERTIES"
        }
        return fmt.Sprintf("ParamID(%#x)", uint64(p))
}

// ParamKind describes a parameter's value encoding (§10.2).
type ParamKind uint8

const (
        // kindUnset is the zero value: the value kind was never set, i.e. the
        // Parameter was built as a bare struct literal rather than via a
        // constructor or parse. appendParamValue panics on it instead of silently
        // emitting a varint, so the mistake surfaces immediately.
        kindUnset    ParamKind = iota
        KindVarint             // single varint
        KindByte               // single byte (uint8)
        KindBytes              // varint-length-prefixed bytes
        KindLocation           // two varints: Group, Object
)

var paramKinds = map[ParamID]ParamKind{
        ParamObjectDeliveryTimeout:   KindVarint,
        ParamAuthorizationToken:      KindBytes,
        ParamRendezvousTimeout:       KindVarint,
        ParamSubgroupDeliveryTimeout: KindVarint,
        ParamExpires:                 KindVarint,
        ParamLargestObject:           KindLocation,
        ParamFillTimeout:             KindVarint,
        ParamForward:                 KindByte,
        ParamSubscriberPriority:      KindByte,
        ParamLocationFilter:          KindBytes,
        ParamGroupOrder:              KindByte,
        ParamFillParameters:          KindBytes,
        ParamSubgroupFilter:          KindBytes,
        ParamObjectIDFilter:          KindBytes,
        ParamPriorityFilter:          KindBytes,
        ParamObjectPropertyFilter:    KindBytes,
        ParamTrackPropertyFilter:     KindBytes,
        ParamNewGroupRequest:         KindVarint,
        ParamTrackNamespacePrefix:    KindBytes,
        ParamIncludeProperties:       KindByte,
}

// kindOf returns the registered kind for a parameter type, or an error if the
// type is unknown. Unknown parameters are a session-level PROTOCOL_VIOLATION
// per §10.2.
func kindOf(t ParamID) (ParamKind, error) {
        k, ok := paramKinds[t]
        if !ok {
                return 0, fmt.Errorf("%w %s", ErrUnknownParameter, t)
        }
        return k, nil
}

// Parameter is a single MoQT message parameter (§10.2). Exactly one of the
// value fields holds data, determined by the value kind the Parameter was
// constructed with (see [ParamKind] and the constructors below).
type Parameter struct {
        Type   ParamID
        Varint uint64
        Byte   uint8
        Bytes  []byte
        Group  uint64 // KindLocation: Group ID
        Object uint64 // KindLocation: Object ID

        // kind records how the value is encoded. It is set by every constructor
        // and by parse, so encoding is self-describing and does not consult the
        // kind registry — an extension Parameter built via a generic helper
        // encodes per the helper used, not per a (possibly absent) registry entry.
        kind ParamKind
}

// Generic construction helpers — keyed by ParamID and value kind. Use the
// typed helpers below for known parameters; reach for these only when
// constructing a parameter the registry doesn't have a dedicated helper for
// (e.g. while experimenting with extensions).

func VarintParam(t ParamID, v uint64) Parameter {
        return Parameter{Type: t, Varint: v, kind: KindVarint}
}
func ByteParam(t ParamID, v uint8) Parameter { return Parameter{Type: t, Byte: v, kind: KindByte} }
func BytesParam(t ParamID, v []byte) Parameter {
        return Parameter{Type: t, Bytes: v, kind: KindBytes}
}
func LocationParam(t ParamID, g, o uint64) Parameter {
        return Parameter{Type: t, Group: g, Object: o, kind: KindLocation}
}

// Typed helpers for the well-known parameters. Each bakes in the right
// ParamID and value kind, and enforces the constraints the spec puts on the
// value (bool for FORWARD, an enum for GROUP_ORDER, time.Duration for the
// millisecond-valued timeouts, etc.).

// ObjectDeliveryTimeoutParam builds OBJECT_DELIVERY_TIMEOUT (§10.2.4): the
// maximum duration the publisher holds a single object before declaring
// failure.
func ObjectDeliveryTimeoutParam(d time.Duration) Parameter {
        //nolint:gosec // G115: d is a non-negative timeout Duration; whole ms fits a varint.
        return VarintParam(ParamObjectDeliveryTimeout, uint64(d/time.Millisecond))
}

// RendezvousTimeoutParam builds RENDEZVOUS_TIMEOUT (§10.2.6): how long the
// subscriber is willing to wait for a publisher to become available. A zero
// duration tells the relay to respond immediately with DOES_NOT_EXIST when
// no publisher exists.
func RendezvousTimeoutParam(d time.Duration) Parameter {
        //nolint:gosec // G115: d is a non-negative timeout Duration; whole ms fits a varint.
        return VarintParam(ParamRendezvousTimeout, uint64(d/time.Millisecond))
}

// SubgroupDeliveryTimeoutParam builds SUBGROUP_DELIVERY_TIMEOUT (§10.2.3).
func SubgroupDeliveryTimeoutParam(d time.Duration) Parameter {
        //nolint:gosec // G115: d is a non-negative timeout Duration; whole ms fits a varint.
        return VarintParam(ParamSubgroupDeliveryTimeout, uint64(d/time.Millisecond))
}

// FillTimeoutParam builds FILL_TIMEOUT (§10.2.5): the maximum total duration
// a relay should spend waiting for upstream sources to provide objects that
// are not immediately available. A zero duration means the subscriber only
// wants objects that are immediately available.
func FillTimeoutParam(d time.Duration) Parameter {
        //nolint:gosec // G115: d is a non-negative timeout Duration; whole ms fits a varint.
        return VarintParam(ParamFillTimeout, uint64(d/time.Millisecond))
}

// ExpiresParam builds EXPIRES (§10.2.16): the time after which the sender
// will terminate the subscription. Zero means the subscription does not
// expire (or expires at an unknown time).
func ExpiresParam(d time.Duration) Parameter {
        //nolint:gosec // G115: d is a non-negative timeout Duration; whole ms fits a varint.
        return VarintParam(ParamExpires, uint64(d/time.Millisecond))
}

// LargestObjectParam builds LARGEST_OBJECT (§10.2.17): the largest Location
// {Group, Object} observed in the track by the sender.
func LargestObjectParam(group, object uint64) Parameter {
        return LocationParam(ParamLargestObject, group, object)
}

// ForwardParam builds FORWARD (§10.2.18). The wire value is restricted to
// 0/1 per the spec, so the helper takes a bool.
func ForwardParam(forward bool) Parameter {
        var v uint8
        if forward {
                v = 1
        }
        return ByteParam(ParamForward, v)
}

// SubscriberPriorityParam builds SUBSCRIBER_PRIORITY (§10.2.7). Lower numbers
// get higher priority; the implicit default when omitted is 128.
func SubscriberPriorityParam(priority uint8) Parameter {
        return ByteParam(ParamSubscriberPriority, priority)
}

// LocationFilterParam builds LOCATION_FILTER (§10.2.9) from a typed
// LocationFilter. The filter is serialised to bytes and stored as a
// length-prefixed KindBytes parameter per §10.2.9.
func LocationFilterParam(f *LocationFilter) Parameter {
        return BytesParam(ParamLocationFilter, f.Bytes())
}

// LocationFilterFromParam extracts and parses a LOCATION_FILTER
// parameter from a Parameters list. Returns nil, nil if the parameter is
// absent (unfiltered subscription). Returns an error if the parameter is
// present but malformed.
func LocationFilterFromParam(ps Parameters) (*LocationFilter, error) {
        p, ok := ps.Find(ParamLocationFilter)
        if !ok {
                return nil, nil //nolint:nilnil // absent optional parameter: (nil filter, nil error) is the documented contract.
        }
        return ParseLocationFilter(p.Bytes)
}

// GroupOrder is the value of the GROUP_ORDER parameter (§10.2.8). The spec
// restricts the wire value to Ascending or Descending; anything else is a
// session-level PROTOCOL_VIOLATION.
type GroupOrder uint8

const (
        GroupOrderAscending  GroupOrder = 0x1
        GroupOrderDescending GroupOrder = 0x2
)

// GroupOrderParam builds GROUP_ORDER (§10.2.8).
func GroupOrderParam(order GroupOrder) Parameter {
        return ByteParam(ParamGroupOrder, uint8(order))
}

// NewGroupRequestParam builds NEW_GROUP_REQUEST (§10.2.19): the largest known
// Group ID plus 1, or 0 if the subscriber has no Group information.
func NewGroupRequestParam(largestGroupPlusOne uint64) Parameter {
        return VarintParam(ParamNewGroupRequest, largestGroupPlusOne)
}

// TrackNamespacePrefixParam builds TRACK_NAMESPACE_PREFIX (§10.2.20): a
// namespace prefix used for namespace subscription updates. The value is a
// TrackNamespace structure serialized per §2.4.1.
func TrackNamespacePrefixParam(prefix wire.TrackNamespace) Parameter {
        // Serialize the TrackNamespace to bytes
        var buf []byte
        w := wire.NewWriter(buf)
        w.TrackNamespace(prefix)
        return BytesParam(ParamTrackNamespacePrefix, w.Bytes())
}

// TrackNamespacePrefixFromParam decodes a TRACK_NAMESPACE_PREFIX value
// (§10.2.20). A value that is not exactly one Track Namespace is an error.
func TrackNamespacePrefixFromParam(p Parameter) (wire.TrackNamespace, error) {
        r := wire.NewReader(p.Bytes)
        prefix, err := r.TrackNamespace()
        if err != nil {
                return nil, fmt.Errorf("moqt/message: TRACK_NAMESPACE_PREFIX: %w", err)
        }
        if !r.Empty() {
                return nil, fmt.Errorf("moqt/message: TRACK_NAMESPACE_PREFIX: %d trailing bytes", r.Remaining())
        }
        return prefix, nil
}

// Parameters is a list of message parameters.
//
//nolint:recvcheck // value receivers for reads, pointer receiver for in-place mutation — intentional.
type Parameters []Parameter

// Find returns the first parameter with the given type, plus a bool indicating
// presence.
func (ps Parameters) Find(t ParamID) (Parameter, bool) {
        for _, p := range ps {
                if p.Type == t {
                        return p, true
                }
        }
        return Parameter{}, false
}

// FindAll returns every parameter with the given type, in list order. Range
// Filter parameters (§5.1.4) legitimately repeat within one message (multiple
// SetIDs / Property Types), so callers that handle them must iterate all
// occurrences rather than rely on [Parameters.Find]'s first-only result.
func (ps Parameters) FindAll(t ParamID) []Parameter {
        var out []Parameter
        for _, p := range ps {
                if p.Type == t {
                        out = append(out, p)
                }
        }
        return out
}

// IsRangeFilterParam reports whether t is one of the five Range Filter
// parameter types (§5.1.4, 0x25-0x29).
func IsRangeFilterParam(t ParamID) bool {
        return t >= ParamSubgroupFilter && t <= ParamTrackPropertyFilter
}

// append writes count + sorted, delta-encoded entries to w. Duplicate types
// are written in input order; callers should de-duplicate where required.
func (ps Parameters) append(w *wire.Writer) {
        w.Varint(uint64(len(ps)))
        sorted := make(Parameters, len(ps))
        copy(sorted, ps)
        slices.SortStableFunc(sorted, func(a, b Parameter) int { return cmp.Compare(a.Type, b.Type) })
        var prev uint64
        for _, p := range sorted {
                t := uint64(p.Type)
                w.Varint(t - prev)
                appendParamValue(w, p)
                prev = t
        }
}

// errParamTypeOverflow is a parameter Type delta taking the Type past 2^64-1:
// a PROTOCOL_VIOLATION (§10.2), unlike a value that does not parse.
var errParamTypeOverflow = errors.New("moqt/message: parameter type delta overflow")

// parse reads a Number-of-Parameters varint followed by that many parameters
// from r.
func (ps *Parameters) parse(r *wire.Reader) error {
        count, err := r.Varint()
        if err != nil {
                return err
        }
        // count is untrusted (up to 2^64-1, §1.4.1): cap the preallocation by the
        // remaining bytes, since each parameter takes at least one.
        //nolint:gosec // G115: Reader.Remaining() = len(buf)-off is always >= 0.
        out := make(Parameters, 0, min(count, uint64(r.Remaining())))
        var prev uint64
        for range count {
                delta, err := r.Varint()
                if err != nil {
                        return err
                }
                if delta > ^uint64(0)-prev {
                        return errParamTypeOverflow
                }
                t := prev + delta
                p := Parameter{Type: ParamID(t)}
                if err := parseParamValue(r, &p); err != nil {
                        return err
                }
                out = append(out, p)
                prev = t
        }
        *ps = out
        return nil
}

func appendParamValue(w *wire.Writer, p Parameter) {
        switch p.kind {
        case KindVarint:
                w.Varint(p.Varint)
        case KindByte:
                w.UInt8(p.Byte)
        case KindBytes:
                w.VarintBytes(p.Bytes)
        case KindLocation:
                w.Varint(p.Group)
                w.Varint(p.Object)
        case kindUnset:
                // The Parameter was built as a bare struct literal rather than via a
                // typed helper or VarintParam/ByteParam/BytesParam/LocationParam (or
                // parse). There is no value kind to encode — surface the programming
                // error loudly instead of silently writing varint(0).
                panic(fmt.Sprintf(
                        "moqt/message: Parameter type %s has no value kind; build it with a constructor, not a bare literal",
                        p.Type,
                ))
        }
}

func parseParamValue(r *wire.Reader, p *Parameter) error {
        k, err := kindOf(p.Type)
        if err != nil {
                return err
        }
        p.kind = k
        switch k {
        case KindVarint:
                v, err := r.Varint()
                if err != nil {
                        return err
                }
                p.Varint = v
        case KindByte:
                v, err := r.UInt8()
                if err != nil {
                        return err
                }
                p.Byte = v
        case KindBytes:
                v, err := r.VarintBytes()
                if err != nil {
                        return err
                }
                p.Bytes = v
        case KindLocation:
                g, err := r.Varint()
                if err != nil {
                        return err
                }
                o, err := r.Varint()
                if err != nil {
                        return err
                }
                p.Group = g
                p.Object = o
        case kindUnset:
                // kindOf never returns kindUnset for a registered type, so this is
                // unreachable; enumerated to keep the switch exhaustive.
                return fmt.Errorf("moqt/message: parameter %s has no value kind", p.Type)
        }
        return nil
}

package message

import (
        "fmt"
        "math"
        "slices"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// PropertyType identifies a MoQT property per §12 and the IANA 'MOQ Properties'
// registry.  Types are used as absolute values in the KVPair.Type field; the
// delta encoding is handled by the wire layer.
type PropertyType = uint64

// Property type constants from §12 and the IANA registry (Table 14).
// All types listed here are from draft-ietf-moq-transport-20.
const (
        // PropertySubgroupDeliveryTimeout (0x06) is a Track or Object Property
        // (§12.1). Value: varint (milliseconds). Semantics defined in §8. As an
        // Object Property on the first object in a subgroup it overrides the
        // Track-level value for that subgroup; it is ignored on any other object.
        PropertySubgroupDeliveryTimeout PropertyType = 0x06

        // PropertyObjectDeliveryTimeout (0x02) is a Track or Object Property
        // (§12.2). Value: varint (milliseconds). Semantics defined in §8. As an
        // Object Property on the first object in a subgroup it overrides the
        // Track-level value for that subgroup; it is ignored on any other object.
        PropertyObjectDeliveryTimeout PropertyType = 0x02

        // PropertyMaxCacheDuration (0x04) is a Track Property (§12.3).
        // Value: varint (milliseconds).
        PropertyMaxCacheDuration PropertyType = 0x04

        // PropertyDefaultPublisherPriority (0x0E) is a Track Property (§12.4).
        // Value: varint 0–255. Default: 128.
        PropertyDefaultPublisherPriority PropertyType = 0x0E

        // PropertyDefaultPublisherGroupOrder (0x22) is a Track Property (§12.5).
        // Value: varint; 0x1 = Ascending (default), 0x2 = Descending.
        PropertyDefaultPublisherGroupOrder PropertyType = 0x22

        // PropertyDynamicGroups (0x30) is a Track Property (§12.6).
        // Value: varint 0 or 1.
        PropertyDynamicGroups PropertyType = 0x30

        // PropertyImmutableProperties (0x0B) is a Track or Object Property (§12.7).
        // Value: bytes containing a nested sequence of KV pairs.
        PropertyImmutableProperties PropertyType = 0x0B

        // PropertyPriorGroupIDGap (0x3C) is an Object Property (§12.8).
        // Value: varint.
        PropertyPriorGroupIDGap PropertyType = 0x3C

        // PropertyPriorObjectIDGap (0x3E) is an Object Property (§12.9).
        // Value: varint.
        PropertyPriorObjectIDGap PropertyType = 0x3E
)

// DefaultPublisherPriority is the Publisher Priority a track has when its
// DEFAULT_PUBLISHER_PRIORITY property is omitted (§12.4).
const DefaultPublisherPriority uint8 = 128

// TrackDefaultPublisherPriority returns the DEFAULT_PUBLISHER_PRIORITY (§12.4)
// in a raw Track Properties block, or [DefaultPublisherPriority] when it is
// omitted. Immutable Properties are searched too, the mutable value winning
// (§12.7). §12.4 prescribes no reaction to a value above 255; like a malformed
// block, it is read as omitted.
func TrackDefaultPublisherPriority(trackProperties []byte) uint8 {
        pairs, err := ParseTrackProperties(trackProperties)
        if err != nil {
                return DefaultPublisherPriority
        }
        if p, ok := findDefaultPublisherPriority(pairs); ok {
                return p
        }
        for _, kv := range pairs {
                if kv.Type != PropertyImmutableProperties {
                        continue
                }
                nested, err := ParseTrackProperties(kv.ByteVal)
                if err != nil {
                        return DefaultPublisherPriority
                }
                if p, ok := findDefaultPublisherPriority(nested); ok {
                        return p
                }
        }
        return DefaultPublisherPriority
}

func findDefaultPublisherPriority(pairs []wire.KVPair) (uint8, bool) {
        for _, kv := range pairs {
                if kv.Type == PropertyDefaultPublisherPriority && kv.IntVal <= math.MaxUint8 {
                        return uint8(kv.IntVal), true
                }
        }
        return 0, false
}

// ExpandImmutable returns pairs followed by the contents of each Immutable
// Properties among them, for the lookup §12.7 requires. The first match is the
// mutable value; a last-wins loop must range over the result backwards.
//
// pairs is returned as is, without allocating, when it holds no Immutable
// Properties. Contents that do not parse are an error (§12.7).
func ExpandImmutable(pairs []wire.KVPair) ([]wire.KVPair, error) {
        out := pairs
        for _, kv := range pairs {
                if kv.Type != PropertyImmutableProperties {
                        continue
                }
                nested, err := ParseTrackProperties(kv.ByteVal)
                if err != nil {
                        return nil, fmt.Errorf("moqt/message: immutable properties: %w", err)
                }
                if len(out) == len(pairs) {
                        out = slices.Clip(out) // append must not write into the caller's array
                }
                out = append(out, nested...)
        }
        return out, nil
}

// parseSearchable parses raw Properties and expands their Immutable
// Properties, see [ExpandImmutable].
func parseSearchable(raw []byte) ([]wire.KVPair, error) {
        pairs, err := ParseTrackProperties(raw)
        if err != nil {
                return nil, err
        }
        return ExpandImmutable(pairs)
}

// TrackMaxCacheDuration returns the MAX_CACHE_DURATION (§12.3) in a raw Track
// Properties block and whether it is present, the mutable value winning over
// Immutable Properties (§12.7). A block that does not parse has none.
func TrackMaxCacheDuration(trackProperties []byte) (time.Duration, bool) {
        pairs, err := parseSearchable(trackProperties)
        if err != nil {
                return 0, false
        }
        for _, kv := range pairs {
                if kv.Type == PropertyMaxCacheDuration {
                        return MillisecondTimeout(kv.IntVal), true
                }
        }
        return 0, false
}

// MandatoryTrackPropertyMin and MandatoryTrackPropertyMax define the range of
// Mandatory Track Property types per §2.5.1.  Properties in [0x4000, 0x7FFF]
// MUST have Track scope; receiving one as an Object Property is malformed.
// An endpoint that does not understand a Mandatory Track Property in PUBLISH,
// SUBSCRIBE_OK, or FETCH_OK MUST NOT process or forward that track.
const (
        MandatoryTrackPropertyMin PropertyType = 0x4000
        MandatoryTrackPropertyMax PropertyType = 0x7FFF
)

// IsMandatoryTrackProperty reports whether t is in the mandatory range
// [0x4000, 0x7FFF] per §2.5.1.
func IsMandatoryTrackProperty(t PropertyType) bool {
        return t >= MandatoryTrackPropertyMin && t <= MandatoryTrackPropertyMax
}

// ParseTrackProperties parses raw Track Properties bytes (the trailing field
// in PUBLISH, SUBSCRIBE_OK, FETCH_OK, etc.) as a sequence of KV pairs.
// Track Properties have no explicit length prefix — they are bounded by the
// outer message frame (§2.5).  The raw bytes are typically obtained via
// wire.Reader.RemainingBytes().
//
// Returns an error if any pair cannot be parsed. Mandatory Track Property
// screening (§2.5.1) is the caller's job — see
// [FirstUnknownMandatoryTrackProperty].
func ParseTrackProperties(raw []byte) ([]wire.KVPair, error) {
        if len(raw) == 0 {
                return nil, nil
        }
        r := wire.NewReader(raw)
        pairs, err := r.KVPairsRemaining()
        if err != nil {
                return nil, fmt.Errorf("moqt/message: track properties: %w", err)
        }
        return pairs, nil
}

// AppendTrackProperties serialises a slice of KV pairs as raw Track Properties
// bytes (no length prefix).  The result is suitable for appending directly to
// a message writer via w.FixedBytes().
func AppendTrackProperties(pairs []wire.KVPair) []byte {
        var w wire.Writer
        w.KVPairs(pairs)
        return w.Bytes()
}

// CheckTrackPropertyValues reports a Track Property in pairs (Immutable
// Properties expanded, see [ExpandImmutable]) whose value the draft makes
// session-fatal: DEFAULT_PUBLISHER_GROUP_ORDER outside {1, 2} (§12.5) or
// DYNAMIC_GROUPS above 1 (§12.6). The caller MUST close the session with
// PROTOCOL_VIOLATION.
func CheckTrackPropertyValues(pairs []wire.KVPair) error {
        for _, kv := range pairs {
                switch kv.Type {
                case PropertyDefaultPublisherGroupOrder:
                        if kv.IntVal != uint64(GroupOrderAscending) && kv.IntVal != uint64(GroupOrderDescending) {
                                return fmt.Errorf(
                                        "moqt/message: DEFAULT_PUBLISHER_GROUP_ORDER value %d outside {1,2} (PROTOCOL_VIOLATION §12.5)",
                                        kv.IntVal)
                        }
                case PropertyDynamicGroups:
                        if kv.IntVal > 1 {
                                return fmt.Errorf("moqt/message: DYNAMIC_GROUPS value %d above 1 (PROTOCOL_VIOLATION §12.6)", kv.IntVal)
                        }
                }
        }
        return nil
}

// FirstUnknownMandatoryTrackProperty returns the first Mandatory Track
// Property (range 0x4000–0x7FFF) in pairs whose type is not in knownTypes,
// and whether one was found — the offending type is what callers need to
// build their rejection error. A nil knownTypes treats every mandatory
// property as unknown.
//
// Per §2.5.1, an endpoint that receives Track Properties containing an
// unknown Mandatory Track Property MUST NOT process or forward that track.
func FirstUnknownMandatoryTrackProperty(
        pairs []wire.KVPair,
        knownTypes map[PropertyType]struct{},
) (PropertyType, bool) {
        for _, kv := range pairs {
                if !IsMandatoryTrackProperty(kv.Type) {
                        continue
                }
                if _, known := knownTypes[kv.Type]; !known {
                        return kv.Type, true
                }
        }
        return 0, false
}

package message

import (
        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// Publish is the PUBLISH message (§10.11).
type Publish struct {
        RequestID       uint64
        Namespace       wire.TrackNamespace
        Name            []byte
        TrackAlias      uint64
        Parameters      Parameters
        TrackProperties []byte
}

func (m *Publish) Type() Type             { return TypePublish }
func (m *Publish) GetRequestID() uint64   { return m.RequestID }
func (m *Publish) SetRequestID(id uint64) { m.RequestID = id }

func (m *Publish) Append(w *wire.Writer) {
        w.Varint(m.RequestID)
        w.TrackNamespace(m.Namespace)
        w.VarintBytes(m.Name)
        w.Varint(m.TrackAlias)
        m.Parameters.append(w)
        w.FixedBytes(m.TrackProperties)
}

func (m *Publish) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.Varint(&m.RequestID)
        s.TrackNamespace(&m.Namespace)
        s.VarintBytes(&m.Name)
        s.Varint(&m.TrackAlias)
        if err := s.Err(); err != nil {
                return err
        }
        if err := m.Parameters.parse(r); err != nil {
                return err
        }
        m.TrackProperties = r.RemainingBytes()
        return nil
}

// PublishDone is the PUBLISH_DONE message (§10.12).
type PublishDone struct {
        StatusCode  moqt.PublishDoneCode
        StreamCount uint64
        ErrorReason string
}

func (m *PublishDone) Type() Type { return TypePublishDone }

func (m *PublishDone) Append(w *wire.Writer) {
        w.Varint(uint64(m.StatusCode))
        w.Varint(m.StreamCount)
        w.ReasonPhrase(m.ErrorReason)
}

func (m *PublishDone) Parse(r *wire.Reader) error {
        s := r.Scanner()
        var code uint64
        s.Varint(&code)
        s.Varint(&m.StreamCount)
        s.ReasonPhrase(&m.ErrorReason)
        if err := s.Err(); err != nil {
                return err
        }
        m.StatusCode = moqt.PublishDoneCode(code)
        return nil
}

// Validate enforces the §2.4.1 Full Track Name size limit; ParsePayload
// invokes it automatically after decoding a PUBLISH frame.
func (m *Publish) Validate() error {
        return validateFullTrackName(m.Namespace, m.Name)
}

package message

import "github.com/floatdrop/moq-go/pkg/moqt/wire"

// PublishStateNotify is a PUBLISH_STATE_NOTIFY message per §10.10, new in
// draft-20.
//
// The publisher sends it on a subscription's bidi stream to report that the
// subscription's state changed for some reason other than a subscriber
// REQUEST_UPDATE. It is unilateral: the receiver sends no REQUEST_OK or
// REQUEST_ERROR, and it does not count against MAX_REQUEST_UPDATES
// (§10.3.1.7). It is informative — no action is required of the recipient.
//
// It carries no Request ID: the stream it arrives on names the subscription.
// That is also why it does not implement [WithRequestID], unlike the other
// request-stream messages.
//
//        PUBLISH_STATE_NOTIFY Message {
//          Type (vi64) = 0x22,
//          Length (16),
//          Number of Parameters (vi64),
//          Parameters (..) ...
//        }
//
// Only the parameters whose values changed are present; an absent parameter is
// unchanged. §10.10 requires LARGEST_OBJECT when known, so the subscriber can
// tell where in the Track the change took effect.
//
// It applies only to subscriptions and only in the publisher-to-subscriber
// direction: receiving one for another request type, or from the subscriber,
// is a session-level PROTOCOL_VIOLATION.
type PublishStateNotify struct {
        Parameters Parameters
}

// Append serializes the PUBLISH_STATE_NOTIFY message to w.
func (m *PublishStateNotify) Append(w *wire.Writer) { m.Parameters.append(w) }

// Parse deserializes the PUBLISH_STATE_NOTIFY message from r.
func (m *PublishStateNotify) Parse(r *wire.Reader) error { return m.Parameters.parse(r) }

// Type returns the wire type ID for PUBLISH_STATE_NOTIFY.
func (m *PublishStateNotify) Type() Type { return TypePublishStateNotify }

package message

import (
        "errors"
        "fmt"
        "math"
        "slices"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// ErrInvalidFilter marks a malformed Range Filter (§5.1.4). The session/relay
// layer maps it to REQUEST_ERROR with code INVALID_FILTER (§10.6, 0x36). It is
// returned for a delta that overflows 2^64-1, an out-of-range PRIORITY value
// (§10.2.12), an odd Property Type on the Object/Track Property filters
// (§10.2.13/§10.2.14), and — at the session layer — a duplicate
// (Type, SetID, Property Type) combination or a total range count exceeding the
// negotiated MAX_FILTER_RANGES (§10.3.1.6).
var ErrInvalidFilter = errors.New("moqt/message: invalid range filter (INVALID_FILTER §5.1.4)")

// Range is one inclusive [Start, End] band of a Range Filter (§5.1.4). Open
// marks the final, open-ended range — its End is omitted on the wire and it
// matches any value >= Start. End is ignored when Open is set.
type Range struct {
        Start uint64
        End   uint64
        Open  bool
}

// RangeFilter is one Range Filter parameter (§5.1.4): SUBGROUP_FILTER (0x25),
// OBJECTID_FILTER (0x26), PRIORITY_FILTER (0x27), OBJECT_PROPERTY_FILTER (0x28),
// or TRACK_PROPERTY_FILTER (0x29). Type is the parameter ID; SetID groups
// filters for AND/OR combination (§5.1.4); PropertyType is meaningful only for
// the Object/Track Property filters (0x28/0x29) and is 0 otherwise. Ranges is
// the ordered, non-overlapping set of value bands the filter selects.
type RangeFilter struct {
        Type         ParamID
        SetID        uint8
        PropertyType PropertyType // only for ParamObjectPropertyFilter / ParamTrackPropertyFilter
        Ranges       []Range
}

// hasPropertyType reports whether this filter type carries a Property Type
// prefix on the wire — only the Object/Track Property filters do (§5.1.4).
func (f *RangeFilter) hasPropertyType() bool {
        return f.Type == ParamObjectPropertyFilter || f.Type == ParamTrackPropertyFilter
}

// Append serialises the filter's value blob to w: SetID, optional Property
// Type, then the delta-encoded Ranges (§5.1.4 — Start delta from the prior
// Range's End or 0, End delta from the current Start; the final End is omitted
// for an Open range). It assumes a validated filter; a mid-list Open range
// would truncate the blob, so call [RangeFilter.Validate] first.
func (f *RangeFilter) Append(w *wire.Writer) {
        w.UInt8(f.SetID)
        if f.hasPropertyType() {
                w.Varint(f.PropertyType)
        }
        var prevEnd uint64
        for _, rg := range f.Ranges {
                w.Varint(rg.Start - prevEnd) // Start delta from prior End (0 for the first)
                if rg.Open {
                        return // final End omitted → open-ended range
                }
                w.Varint(rg.End - rg.Start) // End delta from this Start
                prevEnd = rg.End
        }
}

// Bytes serialises the filter to a fresh byte slice — the value of the
// [RangeFilterParam] parameter.
func (f *RangeFilter) Bytes() []byte {
        var w wire.Writer
        f.Append(&w)
        return w.Bytes()
}

// RangeFilterParam builds the message Parameter (§10.2) carrying f. The value
// is a length-prefixed blob (KindBytes) for all five filter types, see the
// note in params.go.
func RangeFilterParam(f *RangeFilter) Parameter {
        return BytesParam(f.Type, f.Bytes())
}

// ParseRangeFilter decodes a Range Filter parameter's value blob (raw) for
// parameter type t (§5.1.4), resolving the delta-encoded Ranges to absolute
// [Start, End] bands. The open-ended final range is detected when the blob is
// exhausted immediately after a Start. Any delta that overflows 2^64-1 is
// rejected with [ErrInvalidFilter]. Per-type value checks (PRIORITY bound, odd
// Property Type) are applied by [RangeFilter.Validate], not here.
func ParseRangeFilter(t ParamID, raw []byte) (*RangeFilter, error) {
        r := wire.NewReader(raw)
        f := &RangeFilter{Type: t}

        setID, err := r.UInt8()
        if err != nil {
                return nil, fmt.Errorf("%w: SetID: %w", ErrInvalidFilter, err)
        }
        f.SetID = setID

        if f.hasPropertyType() {
                pt, err := r.Varint()
                if err != nil {
                        return nil, fmt.Errorf("%w: property type: %w", ErrInvalidFilter, err)
                }
                f.PropertyType = pt
        }

        var prevEnd uint64
        for !r.Empty() {
                sd, err := r.Varint()
                if err != nil {
                        return nil, fmt.Errorf("%w: range start delta: %w", ErrInvalidFilter, err)
                }
                if sd > math.MaxUint64-prevEnd {
                        return nil, fmt.Errorf("%w: range start delta overflows 2^64-1", ErrInvalidFilter)
                }
                start := prevEnd + sd

                // A Start with no following End is the omitted-final-End open range.
                if r.Empty() {
                        f.Ranges = append(f.Ranges, Range{Start: start, Open: true})
                        break
                }

                ed, err := r.Varint()
                if err != nil {
                        return nil, fmt.Errorf("%w: range end delta: %w", ErrInvalidFilter, err)
                }
                if ed > math.MaxUint64-start {
                        return nil, fmt.Errorf("%w: range end delta overflows 2^64-1", ErrInvalidFilter)
                }
                f.Ranges = append(f.Ranges, Range{Start: start, End: start + ed})
                prevEnd = start + ed
        }
        return f, nil
}

// Validate applies the §5.1.4 per-filter value checks that need no session
// state: the Object/Track Property filters require an even Property Type
// (§10.2.13/§10.2.14), PRIORITY_FILTER values must fit 8 bits (§10.2.12), and
// only the final Range may be open-ended (a mid-list Open cannot round-trip).
// Duplicate-combination and MAX_FILTER_RANGES checks need session state and
// live in [RangeFiltersFromParams] / [RangeFilterSet.Validate].
func (f *RangeFilter) Validate() error {
        if f.hasPropertyType() && f.PropertyType%2 != 0 {
                return fmt.Errorf("%w: %s property type 0x%X must be even", ErrInvalidFilter, f.Type, f.PropertyType)
        }
        for i, rg := range f.Ranges {
                if rg.Open && i != len(f.Ranges)-1 {
                        return fmt.Errorf("%w: only the final range may be open-ended", ErrInvalidFilter)
                }
                if f.Type == ParamPriorityFilter && (rg.Start > 255 || (!rg.Open && rg.End > 255)) {
                        return fmt.Errorf("%w: PRIORITY value exceeds 255 (§10.2.12)", ErrInvalidFilter)
                }
        }
        return nil
}

// matchValue reports whether v falls in any of the filter's Ranges (inclusive;
// an Open range matches v >= Start). A filter with no Ranges matches nothing.
func (f *RangeFilter) matchValue(v uint64) bool {
        for _, rg := range f.Ranges {
                if rg.Open {
                        if v >= rg.Start {
                                return true
                        }
                        continue
                }
                if v >= rg.Start && v <= rg.End {
                        return true
                }
        }
        return false
}

// RangeFilterSet is the collection of Range Filters (§5.1.4) on one request,
// grouped by SetID. A value passes a group when it satisfies every filter in
// that group (AND); it passes the set when it passes any group (OR) — §5.1.4's
// "SetID=0 OR SetID=1 OR ...". A nil or empty set imposes no restriction. Build
// it with [RangeFiltersFromParams].
type RangeFilterSet struct {
        groups            []rangeGroup
        totalRanges       int
        hasObjectProperty bool // any group holds an OBJECT_PROPERTY_FILTER
        hasTrackProperty  bool // any group holds a TRACK_PROPERTY_FILTER
}

// rangeGroup holds every filter sharing one SetID (AND-combined).
type rangeGroup struct {
        setID   uint8
        filters []RangeFilter
}

type filterKey struct {
        typ    ParamID
        setID  uint8
        propTy PropertyType
}

// RangeFiltersFromParams extracts every Range Filter parameter (§5.1.4) from ps,
// validates each, rejects a duplicate (Type, SetID, Property Type) combination
// (§5.1.4), and groups them by SetID. Returns (nil, nil) when ps carries no
// range filters, matching [LocationFilterFromParam]. A zero-length parameter
// is no filter (§5.1.4). The MAX_FILTER_RANGES limit is enforced separately by
// [RangeFilterSet.Validate].
func RangeFiltersFromParams(ps Parameters) (*RangeFilterSet, error) {
        filters, err := parseRangeFilters(ps)
        if err != nil {
                return nil, err
        }
        return buildRangeFilterSet(filters)
}

// Update applies a REQUEST_UPDATE's Range Filter parameters to s (§5.1.4):
// Length 0 removes a filter, non-zero replaces it, omitted leaves it. A
// zero-length parameter carries no SetID, so filters are replaced by Parameter
// Type: every existing filter of a type the update carries is dropped. Returns
// nil when no filter remains. s may be nil.
func (s *RangeFilterSet) Update(ps Parameters) (*RangeFilterSet, error) {
        added, err := parseRangeFilters(ps)
        if err != nil {
                return nil, err
        }
        var kept []RangeFilter
        if s != nil {
                for _, g := range s.groups {
                        for _, f := range g.filters {
                                if !slices.ContainsFunc(ps, func(p Parameter) bool { return p.Type == f.Type }) {
                                        kept = append(kept, f)
                                }
                        }
                }
        }
        return buildRangeFilterSet(append(kept, added...))
}

// parseRangeFilters parses and validates the non-empty Range Filter
// parameters of ps, in order.
func parseRangeFilters(ps Parameters) ([]RangeFilter, error) {
        var out []RangeFilter
        for _, p := range ps {
                if !IsRangeFilterParam(p.Type) || len(p.Bytes) == 0 {
                        continue
                }
                f, err := ParseRangeFilter(p.Type, p.Bytes)
                if err != nil {
                        return nil, err
                }
                if err := f.Validate(); err != nil {
                        return nil, err
                }
                out = append(out, *f)
        }
        return out, nil
}

// buildRangeFilterSet rejects a duplicate (Type, SetID, Property Type)
// combination (§5.1.4) and groups filters by SetID; nil for no filters.
func buildRangeFilterSet(filters []RangeFilter) (*RangeFilterSet, error) {
        if len(filters) == 0 {
                return nil, nil //nolint:nilnil // no filters: (nil set, nil error) is the documented contract.
        }
        set := &RangeFilterSet{}
        seen := make(map[filterKey]struct{}, len(filters))
        groupIdx := make(map[uint8]int)
        for _, f := range filters {
                key := filterKey{typ: f.Type, setID: f.SetID, propTy: f.PropertyType}
                if _, dup := seen[key]; dup {
                        return nil, fmt.Errorf("%w: duplicate filter (type=%s setID=%d propertyType=0x%X)",
                                ErrInvalidFilter, f.Type, f.SetID, f.PropertyType)
                }
                seen[key] = struct{}{}
                set.totalRanges += len(f.Ranges)
                // SUBGROUP/OBJECTID/PRIORITY carry no property blob; only these two do.
                if f.Type == ParamObjectPropertyFilter {
                        set.hasObjectProperty = true
                }
                if f.Type == ParamTrackPropertyFilter {
                        set.hasTrackProperty = true
                }
                gi, ok := groupIdx[f.SetID]
                if !ok {
                        gi = len(set.groups)
                        set.groups = append(set.groups, rangeGroup{setID: f.SetID})
                        groupIdx[f.SetID] = gi
                }
                set.groups[gi].filters = append(set.groups[gi].filters, f)
        }
        return set, nil
}

// Validate enforces the MAX_FILTER_RANGES setup option (§10.3.1.6): a limit of
// 0 prohibits range filters entirely, and the total number of Ranges across all
// filters must not exceed maxFilterRanges. Returns [ErrInvalidFilter] on breach.
// A nil set (no filters) is always valid.
func (s *RangeFilterSet) Validate(maxFilterRanges uint64) error {
        if s == nil {
                return nil
        }
        if maxFilterRanges == 0 {
                return fmt.Errorf("%w: range filters not permitted (MAX_FILTER_RANGES=0)", ErrInvalidFilter)
        }
        //nolint:gosec // G115: totalRanges is a non-negative sum of len(Ranges).
        if uint64(s.totalRanges) > maxFilterRanges {
                return fmt.Errorf("%w: %d ranges exceed MAX_FILTER_RANGES=%d",
                        ErrInvalidFilter, s.totalRanges, maxFilterRanges)
        }
        return nil
}

// propertyValue returns the first value of property t in pairs: the mutable
// one when [ExpandImmutable] also found it in Immutable Properties (§12.7).
// Validate admits only even (varint) Property Types.
func propertyValue(pairs []wire.KVPair, t PropertyType) (uint64, bool) {
        for _, kv := range pairs {
                if kv.Type == t {
                        return kv.IntVal, true
                }
        }
        return 0, false
}

// matchObject reports whether the group's object-scoped filters
// (SUBGROUP/OBJECTID/PRIORITY/OBJECT_PROPERTY) all match — the AND within a
// SetID. Track-property filters in the group are not object constraints and are
// skipped here (they gate the track via trackPassPerGroup).
func (g *rangeGroup) matchObject(subgroupID, objectID uint64, priority uint8, props []wire.KVPair) bool {
        for i := range g.filters {
                f := &g.filters[i]
                //nolint:exhaustive // only the four object-scoped filter types constrain
                // an object; TRACK_PROPERTY (the default) is gated via TrackPassPerGroup,
                // and no other ParamID reaches a rangeGroup.
                switch f.Type {
                case ParamSubgroupFilter:
                        if !f.matchValue(subgroupID) {
                                return false
                        }
                case ParamObjectIDFilter:
                        if !f.matchValue(objectID) {
                                return false
                        }
                case ParamPriorityFilter:
                        if !f.matchValue(uint64(priority)) {
                                return false
                        }
                case ParamObjectPropertyFilter:
                        v, ok := propertyValue(props, f.PropertyType)
                        if !ok || !f.matchValue(v) {
                                return false
                        }
                default:
                        // ParamTrackPropertyFilter is a track constraint, gated separately
                        // via TrackPassPerGroup — not an object constraint here.
                }
        }
        return true
}

// matchTrack reports whether the group's TRACK_PROPERTY filters all match — the
// AND within a SetID for the track scope. A group with no track filter passes
// vacuously.
func (g *rangeGroup) matchTrack(props []wire.KVPair) bool {
        for i := range g.filters {
                f := &g.filters[i]
                if f.Type == ParamTrackPropertyFilter {
                        v, ok := propertyValue(props, f.PropertyType)
                        if !ok || !f.matchValue(v) {
                                return false
                        }
                }
        }
        return true
}

// MatchesObject reports whether an object with the given Subgroup ID, Object ID,
// Publisher Priority, and Object-Properties blob passes the set's object-scoped
// filters (§5.1.4): OR over SetID of (AND of the group's SUBGROUP/OBJECTID/
// PRIORITY/OBJECT_PROPERTY filters). A nil/empty set matches everything.
//
// This ignores TRACK_PROPERTY filters, so it is exact for SUBSCRIBE/FETCH (which
// carry no track filters). When a SetID mixes object and track filters (possible
// in SUBSCRIBE_TRACKS), use [RangeFilterSet.MatchesObjectInSets] with a
// [RangeFilterSet.TrackPassPerGroup] vector instead.
func (s *RangeFilterSet) MatchesObject(subgroupID, objectID uint64, priority uint8, objProps []byte) bool {
        return s.MatchesObjectInSets(subgroupID, objectID, priority, objProps, nil)
}

// MatchesObjectInSets is [RangeFilterSet.MatchesObject] with per-SetID track
// gating: group i is eligible only when trackPass[i] is true (nil trackPass =
// all groups eligible). This implements the exact §5.1.4 semantics
// OR_i(trackPass[i] AND objectFilters_i) for the mixed object+track-in-one-SetID
// case, where a naive MatchesTrack() && MatchesObject() would be wrong.
func (s *RangeFilterSet) MatchesObjectInSets(
        subgroupID, objectID uint64, priority uint8, objProps []byte, trackPass []bool,
) bool {
        if s == nil || len(s.groups) == 0 {
                return true
        }
        var props []wire.KVPair
        if s.hasObjectProperty {
                props, _ = parseSearchable(objProps) // malformed → nil → property filters miss
        }
        for i := range s.groups {
                if trackPass != nil && !trackPass[i] {
                        continue
                }
                if s.groups[i].matchObject(subgroupID, objectID, priority, props) {
                        return true
                }
        }
        return false
}

// TrackPassPerGroup returns, for each SetID group (in the same order as
// [RangeFilterSet.MatchesObjectInSets] evaluates), whether the group's
// TRACK_PROPERTY filters all pass for a track with the given Track Properties.
// Computed once per (track, subscription) and reused across that track's
// objects. Returns nil for a nil set.
func (s *RangeFilterSet) TrackPassPerGroup(trackProps []byte) []bool {
        if s == nil {
                return nil
        }
        var props []wire.KVPair
        if s.hasTrackProperty {
                props, _ = parseSearchable(trackProps)
        }
        pass := make([]bool, len(s.groups))
        for i := range s.groups {
                pass[i] = s.groups[i].matchTrack(props)
        }
        return pass
}

// MatchesTrack reports whether a track with the given Track Properties passes
// the set's TRACK_PROPERTY filters (§5.1.4 / §10.2.14) — the PUBLISH-forwarding
// gate for SUBSCRIBE_TRACKS: OR over SetID of (AND of the group's track-property
// filters). A group with no track filter passes vacuously, so a set with only
// object filters matches every track; a nil set matches everything.
func (s *RangeFilterSet) MatchesTrack(trackProps []byte) bool {
        if s == nil || len(s.groups) == 0 {
                return true
        }
        var props []wire.KVPair
        if s.hasTrackProperty {
                props, _ = parseSearchable(trackProps)
        }
        for i := range s.groups {
                if s.groups[i].matchTrack(props) {
                        return true
                }
        }
        return false
}

package message

import (
        "errors"
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// RequestUpdate is the REQUEST_UPDATE message (§10.9).
type RequestUpdate struct {
        RequestID  uint64
        Parameters Parameters
}

func (m *RequestUpdate) Type() Type             { return TypeRequestUpdate }
func (m *RequestUpdate) GetRequestID() uint64   { return m.RequestID }
func (m *RequestUpdate) SetRequestID(id uint64) { m.RequestID = id }

func (m *RequestUpdate) Append(w *wire.Writer) {
        w.Varint(m.RequestID)
        m.Parameters.append(w)
}

func (m *RequestUpdate) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.Varint(&m.RequestID)
        if err := s.Err(); err != nil {
                return err
        }
        return m.Parameters.parse(r)
}

// RequestOK is the REQUEST_OK message (§10.5). Track Properties are populated
// when used as a TRACK_STATUS_OK response and empty otherwise (PUBLISH, REQUEST_UPDATE).
type RequestOK struct {
        Parameters      Parameters
        TrackProperties []byte
}

func (m *RequestOK) Type() Type { return TypeRequestOK }

func (m *RequestOK) Append(w *wire.Writer) {
        m.Parameters.append(w)
        w.FixedBytes(m.TrackProperties)
}

func (m *RequestOK) Parse(r *wire.Reader) error {
        if err := m.Parameters.parse(r); err != nil {
                return err
        }
        m.TrackProperties = r.RemainingBytes()
        return nil
}

// Redirect carries the optional redirect payload of REQUEST_ERROR (§10.6.1).
type Redirect struct {
        ConnectURI []byte
        Namespace  wire.TrackNamespace
        TrackName  []byte
}

// RequestError is the REQUEST_ERROR message (§10.6.2). Redirect is non-nil
// when ErrorCode is REDIRECT.
type RequestError struct {
        ErrorCode     moqt.RequestErrorCode
        RetryInterval uint64
        ErrorReason   string
        Redirect      *Redirect
}

func (m *RequestError) Type() Type { return TypeRequestError }

func (m *RequestError) Append(w *wire.Writer) {
        w.Varint(uint64(m.ErrorCode))
        w.Varint(m.RetryInterval)
        w.ReasonPhrase(m.ErrorReason)
        if m.Redirect != nil {
                w.VarintBytes(m.Redirect.ConnectURI)
                w.TrackNamespace(m.Redirect.Namespace)
                w.VarintBytes(m.Redirect.TrackName)
        }
}

func (m *RequestError) Parse(r *wire.Reader) error {
        s := r.Scanner()
        var code uint64
        s.Varint(&code)
        s.Varint(&m.RetryInterval)
        s.ReasonPhrase(&m.ErrorReason)
        if err := s.Err(); err != nil {
                return err
        }
        m.ErrorCode = moqt.RequestErrorCode(code)
        if r.Empty() {
                return nil
        }
        var rd Redirect
        s.VarintBytes(&rd.ConnectURI)
        s.TrackNamespace(&rd.Namespace)
        s.VarintBytes(&rd.TrackName)
        if err := s.Err(); err != nil {
                return err
        }
        m.Redirect = &rd
        return nil
}

// Validate enforces the §10.6.2 REQUEST_ERROR invariants. It is invoked
// automatically by [ParsePayload] after decode, so a malformed REQUEST_ERROR
// (REDIRECT code without a Redirect block, or vice versa) is rejected at the
// parse boundary rather than reaching the session layer.
func (m *RequestError) Validate() error {
        return m.ValidateRedirect()
}

// ValidateRedirect enforces the §10.6.2 constraints: the Redirect block MUST
// be present when ErrorCode is REDIRECT, and MUST NOT be present otherwise.
// It is the implementation behind [RequestError.Validate]; callers may also
// invoke it directly.
func (m *RequestError) ValidateRedirect() error {
        if m.ErrorCode == moqt.RequestRedirect && m.Redirect == nil {
                return errors.New("moqt/message: ErrorCode is REDIRECT but Redirect block is absent")
        }
        if m.ErrorCode != moqt.RequestRedirect && m.Redirect != nil {
                return fmt.Errorf("moqt/message: Redirect present but ErrorCode %#x is not REDIRECT", uint64(m.ErrorCode))
        }
        return nil
}

package message

import (
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// SetupOption is a MoQT SETUP option type ID (§10.3.1). Distinct from ParamID
// because the two code spaces overlap: option 0x03 is AUTHORIZATION_TOKEN at
// the session level and parameter 0x03 is AUTHORIZATION_TOKEN at the request
// level — different message contexts, different parsing rules. The underlying
// wire field (wire.KVPair.Type) stays uint64 because KVPair is wire-generic.
type SetupOption uint64

const (
        SetupOptionPath               SetupOption = 0x01
        SetupOptionAuthorizationToken SetupOption = 0x03
        SetupOptionMaxAuthTokenCache  SetupOption = 0x04
        SetupOptionAuthority          SetupOption = 0x05
        SetupOptionMaxFilterRanges    SetupOption = 0x06
        SetupOptionMOQTImplementation SetupOption = 0x07
        SetupOptionMaxRequestUpdates  SetupOption = 0x08
)

// Setup carries the SETUP message payload (§10.3). Setup Options span the
// remainder of the message payload as a delta-encoded sequence of KVPairs.
type Setup struct {
        Options []wire.KVPair
}

// PathOption builds a PATH setup option (§10.3.1.2). Client-only, native-QUIC
// only: a PATH option received by a server, on a WebTransport session, or with
// an unsupported path triggers an INVALID_PATH session close. pathAndQuery is
// the path-abempty portion of the moqt URI, optionally followed by "?" and
// the query.
func PathOption(pathAndQuery string) wire.KVPair {
        return wire.KVPair{Type: uint64(SetupOptionPath), ByteVal: []byte(pathAndQuery)}
}

// AuthorityOption builds an AUTHORITY setup option (§10.3.1.1). Client-only,
// native-QUIC only: a server-sent or WebTransport-sent AUTHORITY triggers
// INVALID_AUTHORITY. authority is the authority portion of the moqt URI.
func AuthorityOption(authority string) wire.KVPair {
        return wire.KVPair{Type: uint64(SetupOptionAuthority), ByteVal: []byte(authority)}
}

// MOQTImplementationOption builds a MOQT_IMPLEMENTATION setup option
// (§10.3.1.5). Optional; intended for debugging and interop tracking. nameAndVersion
// SHOULD be the implementation name plus version (e.g. "mediamesh/0.1.0").
func MOQTImplementationOption(nameAndVersion string) wire.KVPair {
        return wire.KVPair{Type: uint64(SetupOptionMOQTImplementation), ByteVal: []byte(nameAndVersion)}
}

// MaxAuthTokenCacheSizeOption builds a MAX_AUTH_TOKEN_CACHE_SIZE option
// (§10.3.1.3). maxBytes is the peer-allowed total size in bytes of registered
// authorization tokens. The default if omitted is 0, which prohibits the use
// of token Aliases.
func MaxAuthTokenCacheSizeOption(maxBytes uint64) wire.KVPair {
        return wire.KVPair{Type: uint64(SetupOptionMaxAuthTokenCache), IntVal: maxBytes}
}

// MaxRequestUpdatesOption builds a MAX_REQUEST_UPDATES option (§10.3.1.7).
// maxUpdates is the maximum number of unacknowledged REQUEST_UPDATE messages
// the peer may have outstanding on any single request stream; the receiver of
// a REQUEST_UPDATE that exceeds it MUST close the session with
// TOO_MANY_REQUEST_UPDATES. The default if omitted is 0, which means the
// endpoint does not limit REQUEST_UPDATE concurrency.
func MaxRequestUpdatesOption(maxUpdates uint64) wire.KVPair {
        return wire.KVPair{Type: uint64(SetupOptionMaxRequestUpdates), IntVal: maxUpdates}
}

// MaxFilterRangesOption builds a MAX_FILTER_RANGES option (§10.3.1.6).
// maxRanges is the maximum total number of Ranges (Start/End pairs) the peer
// may send across all Range Filter parameters (§5.1.4) for a single
// subscription or fetch. The default if omitted is 0, which prohibits the peer
// from sending any Range Filter parameters.
func MaxFilterRangesOption(maxRanges uint64) wire.KVPair {
        return wire.KVPair{Type: uint64(SetupOptionMaxFilterRanges), IntVal: maxRanges}
}

func (m *Setup) Type() Type { return TypeSetup }

func (m *Setup) Append(w *wire.Writer) {
        w.KVPairs(m.Options)
}

func (m *Setup) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.KVPairsRemaining(&m.Options)
        return s.Err()
}

package message

import (
        "fmt"
        "io"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// SubgroupIDMode is the 2-bit SUBGROUP_ID_MODE sub-field of the
// SUBGROUP_HEADER Type byte (bits 1-2, §11.4.2). It controls whether and
// how the Subgroup ID is transmitted in the header.
type SubgroupIDMode uint8

const (
        // SubgroupIDImplicitZero: Subgroup ID is omitted; receiver MUST treat
        // it as 0.
        SubgroupIDImplicitZero SubgroupIDMode = 0b00
        // SubgroupIDImplicitFirstObject: Subgroup ID is omitted; receiver MUST
        // treat it as equal to the first Object ID transmitted in this
        // subgroup.
        SubgroupIDImplicitFirstObject SubgroupIDMode = 0b01
        // SubgroupIDExplicit: Subgroup ID is present in the header.
        SubgroupIDExplicit SubgroupIDMode = 0b10
        // 0b11 is reserved (§11.4.2); receiving it MUST cause a session-level
        // PROTOCOL_VIOLATION, and constructing it is a programmer error.
)

// SubgroupHeader is the alias-bearing prefix of a SUBGROUP_HEADER stream
// (§11.4.2). The flag fields correspond directly to the bits of the wire
// Type byte; Type() encodes them and DecodeSubgroupHeaderType parses them
// back.
type SubgroupHeader struct {
        // Properties: when true, every Object on this stream carries an
        // Object Properties structure (§11.2.1.2). Wire bit 0.
        Properties bool
        // SubgroupIDMode controls how the Subgroup ID is conveyed (wire
        // bits 1-2).
        SubgroupIDMode SubgroupIDMode
        // EndOfGroup: when true, this subgroup contains the largest Object
        // in the Group. Wire bit 3.
        EndOfGroup bool
        // InlinePriority: when true, the subgroup body begins with a one-byte
        // Publisher Priority value that overrides the subscription default.
        // When false (zero value, common case), the body starts directly with
        // the first Object and the subgroup inherits the Publisher Priority
        // from the SUBSCRIBE/PUBLISH control message. Wire bit 5 (the spec's
        // DEFAULT_PRIORITY bit) — set on the wire when this field is false.
        InlinePriority bool
        // ReplayingSubgroup: when true, the first Object on this stream is
        // NOT the first object the original publisher pushed for this
        // subgroup — i.e. the stream is a partial replay from a relay or
        // cache. When false (zero value, common case), the first Object on
        // the stream is the first Object of the subgroup. Wire bit 6 (the
        // spec's FIRST_OBJECT bit) — set on the wire when this field is
        // false.
        ReplayingSubgroup bool

        // TrackAlias identifies the track this subgroup belongs to within
        // the publisher → subscriber direction of the session (§11.1).
        TrackAlias uint64

        // GroupID is the Group ID of this subgroup (§11.4.2). Always present
        // on the wire after TrackAlias.
        GroupID uint64

        // SubgroupID is the Subgroup ID of this subgroup. Present on the wire
        // only when SubgroupIDMode == SubgroupIDExplicit (0b10). When the mode
        // is SubgroupIDImplicitZero the receiver treats it as 0; when the mode
        // is SubgroupIDImplicitFirstObject the receiver treats it as equal to
        // the first Object ID on the stream.
        SubgroupID uint64

        // PublisherPriority is the per-subgroup publisher priority byte.
        // Present on the wire only when InlinePriority == true. When
        // InlinePriority is false the subgroup inherits the priority from the
        // enclosing SUBSCRIBE/PUBLISH control message.
        PublisherPriority uint8
}

// Wire-byte bit layout (§11.4.2): 0b0XX1XXXX where bit 4 is always set
// and bit 7 is always clear.
const (
        subgroupBitProperties      uint64 = 0x01 // bit 0
        subgroupModeMask           uint64 = 0x06 // bits 1-2
        subgroupModeShift                 = 1
        subgroupBitEndOfGroup      uint64 = 0x08 // bit 3
        subgroupBitMandatory       uint64 = 0x10 // bit 4
        subgroupBitDefaultPriority uint64 = 0x20 // bit 5
        subgroupBitFirstObject     uint64 = 0x40 // bit 6
)

// Type returns the wire Type byte encoding the flag fields (§11.4.2).
// The mandatory bit-4 sanity bit is always set. SubgroupIDMode is masked
// to 2 bits — callers that pass an out-of-range value get the bottom
// two bits.
//
// Note that InlinePriority and ReplayingSubgroup are inverted relative
// to the wire bits: a false (zero) field sets the corresponding wire
// bit. This makes the zero value of SubgroupHeader produce the typical
// "inherit priority, original publish" Type byte (0x70).
func (h SubgroupHeader) Type() uint64 {
        t := subgroupBitMandatory
        if h.Properties {
                t |= subgroupBitProperties
        }
        t |= (uint64(h.SubgroupIDMode) & 0b11) << subgroupModeShift
        if h.EndOfGroup {
                t |= subgroupBitEndOfGroup
        }
        if !h.InlinePriority {
                t |= subgroupBitDefaultPriority
        }
        if !h.ReplayingSubgroup {
                t |= subgroupBitFirstObject
        }
        return t
}

// DecodeSubgroupHeaderType parses a wire Type byte (§11.4.2) into the
// flag fields of a SubgroupHeader. TrackAlias is left zero — the header
// parser reads it from the following varint. Returns an error if t is
// not a valid SUBGROUP_HEADER Type (i.e. IsSubgroupHeaderType(t) is
// false).
func DecodeSubgroupHeaderType(t uint64) (SubgroupHeader, error) {
        if !IsSubgroupHeaderType(t) {
                return SubgroupHeader{}, fmt.Errorf("moqt/message: invalid SUBGROUP_HEADER type %#x", t)
        }
        return SubgroupHeader{
                Properties:        t&subgroupBitProperties != 0,
                SubgroupIDMode:    SubgroupIDMode((t & subgroupModeMask) >> subgroupModeShift),
                EndOfGroup:        t&subgroupBitEndOfGroup != 0,
                InlinePriority:    t&subgroupBitDefaultPriority == 0,
                ReplayingSubgroup: t&subgroupBitFirstObject == 0,
        }, nil
}

// IsSubgroupHeaderType reports whether t is one of the valid SUBGROUP_HEADER
// type values per §11.4.2: the four ranges 0x10..0x1F, 0x30..0x3F, 0x50..0x5F,
// 0x70..0x7F, excluding values where SUBGROUP_ID_MODE (bits 1-2) is 0b11.
func IsSubgroupHeaderType(t uint64) bool {
        if t > 0x7F {
                return false
        }
        // Bit 4 must be set, bit 7 clear (0b0XX1XXXX).
        if t&subgroupBitMandatory == 0 || t&0x80 != 0 {
                return false
        }
        // SUBGROUP_ID_MODE = 0b11 is reserved.
        if t&subgroupModeMask == subgroupModeMask {
                return false
        }
        return true
}

// IsReservedSubgroupHeaderType reports whether t looks like a SUBGROUP_HEADER
// type byte (bit 4 set, bit 7 clear) but has the reserved SUBGROUP_ID_MODE
// value 0b11 in bits 1-2. Per §11.4.2, receiving such a value MUST close the
// session with a PROTOCOL_VIOLATION. §3.4 also requires closing the session on
// a truly unknown stream type; this lets the caller report the reserved mode
// distinctly from an unknown type.
func IsReservedSubgroupHeaderType(t uint64) bool {
        if t > 0x7F {
                return false
        }
        // Must look like a subgroup header: bit 4 set, bit 7 clear.
        if t&subgroupBitMandatory == 0 || t&0x80 != 0 {
                return false
        }
        // Reserved: SUBGROUP_ID_MODE bits 1-2 are both set (0b11).
        return t&subgroupModeMask == subgroupModeMask
}

// WriteSubgroupHeader writes the full SUBGROUP_HEADER wire encoding (§11.4.2):
// Type, Track Alias, Group ID, optional Subgroup ID (when
// SubgroupIDMode == SubgroupIDExplicit), and optional Publisher Priority (when
// InlinePriority == true).
func WriteSubgroupHeader(w io.Writer, h SubgroupHeader) error {
        buf := wire.AppendVarint(nil, h.Type())
        buf = wire.AppendVarint(buf, h.TrackAlias)
        buf = wire.AppendVarint(buf, h.GroupID)
        if h.SubgroupIDMode == SubgroupIDExplicit {
                buf = wire.AppendVarint(buf, h.SubgroupID)
        }
        if h.InlinePriority {
                buf = append(buf, h.PublisherPriority)
        }
        _, err := w.Write(buf)
        return err
}

// ReadSubgroupHeader reads a complete SUBGROUP_HEADER from r (§11.4.2).
// The caller must have already read the leading Type varint via
// ReadDataStreamType and verified it with IsSubgroupHeaderType; pass that
// raw type value as typ. ReadSubgroupHeader decodes the flag fields from typ
// and then reads Track Alias, Group ID, optional Subgroup ID (when
// SubgroupIDMode == SubgroupIDExplicit), and optional Publisher Priority
// (when InlinePriority is set).
func ReadSubgroupHeader(r io.Reader, typ uint64) (SubgroupHeader, error) {
        h, err := DecodeSubgroupHeaderType(typ)
        if err != nil {
                return SubgroupHeader{}, err
        }

        br := wire.NewByteReader(r)

        alias, err := wire.ReadVarint(br)
        if err != nil {
                return SubgroupHeader{}, fmt.Errorf("moqt/message: SUBGROUP_HEADER track alias: %w", err)
        }
        h.TrackAlias = alias

        groupID, err := wire.ReadVarint(br)
        if err != nil {
                return SubgroupHeader{}, fmt.Errorf("moqt/message: SUBGROUP_HEADER group ID: %w", err)
        }
        h.GroupID = groupID

        if h.SubgroupIDMode == SubgroupIDExplicit {
                subgroupID, err := wire.ReadVarint(br)
                if err != nil {
                        return SubgroupHeader{}, fmt.Errorf("moqt/message: SUBGROUP_HEADER subgroup ID: %w", err)
                }
                h.SubgroupID = subgroupID
        }

        if h.InlinePriority {
                // Publisher Priority is a single byte (§11.4.2).
                var buf [1]byte
                if _, err := io.ReadFull(r, buf[:]); err != nil {
                        return SubgroupHeader{}, fmt.Errorf("moqt/message: SUBGROUP_HEADER publisher priority: %w", err)
                }
                h.PublisherPriority = buf[0]
        }

        return h, nil
}

package message

import (
        "errors"
        "fmt"
        "math/bits"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// ErrIDOverflow reports a Group or Object ID reconstructed from a delta outside
// 0..2^64-1, a session-level PROTOCOL_VIOLATION (§11.4.2, §11.4.4.1).
var ErrIDOverflow = errors.New("moqt/message: Group or Object ID outside 0..2^64-1")

// NextSubgroupObjectID returns prev + delta + 1, the next Object ID on a
// Subgroup stream (§11.4.2), or [ErrIDOverflow] past 2^64-1.
func NextSubgroupObjectID(prev, delta uint64) (uint64, error) {
        id, carry := bits.Add64(prev, delta, 1)
        if carry != 0 {
                return 0, fmt.Errorf("%w: Object ID %d + delta %d + 1", ErrIDOverflow, prev, delta)
        }
        return id, nil
}

// Object Status values for objects with an empty payload (§11.2.1.1).
const (
        ObjectStatusNormal     uint64 = 0x0 // a normal object (carries a payload)
        ObjectStatusEndOfGroup uint64 = 0x3 // last object in the Group
        ObjectStatusEndOfTrack uint64 = 0x4 // last object in the Track
)

// SubgroupObject represents a single object serialized on a SUBGROUP_HEADER
// stream after the SubgroupHeader (§11.4.2, Figure 26).
type SubgroupObject struct {
        // ObjectIDDelta is always present on the wire. For the first object in
        // the stream it is the absolute Object ID; for subsequent objects it is
        // (currentID - previousID - 1), so sequential IDs all encode as 0.
        ObjectIDDelta uint64

        // Properties is present when SubgroupHeader.Properties == true.
        // Encoded as a length-prefixed blob (§11.2.1.2).
        // Must be non-nil (even if empty) when the header has Properties == true.
        Properties []byte

        // Payload is the object body. When non-empty, ObjectStatus is ignored.
        // Encoded on the wire as: Object Payload Length (vi64) + bytes.
        Payload []byte

        // ObjectStatus is only written when len(Payload) == 0.
        // Values: 0x0 Normal, 0x3 EndOfGroup, 0x4 EndOfTrack (§11.2.1.1).
        ObjectStatus uint64
}

// Append serializes the SubgroupObject to the wire writer.
// The hasProperties parameter indicates whether the parent SubgroupHeader
// had the Properties bit set, which determines if Properties are included.
func (o *SubgroupObject) Append(w *wire.Writer, hasProperties bool) {
        // Object ID Delta is always present (§11.4.2)
        w.Varint(o.ObjectIDDelta)

        // Properties are present only if the stream header has Properties == true
        if hasProperties {
                w.VarintBytes(o.Properties)
        }

        w.Varint(uint64(len(o.Payload)))

        // Object Status is present only when Payload Length == 0
        if len(o.Payload) == 0 {
                w.Varint(o.ObjectStatus)
        } else {
                w.FixedBytes(o.Payload)
        }
}

// Parse deserializes a SubgroupObject from r.
// r may be a *wire.Reader (in-memory) or a *wire.StreamReader (streaming).
// The hasProperties parameter indicates whether the parent SubgroupHeader
// had the Properties bit set, which determines if Properties are included.
//
// io.EOF is returned only for a stream that ends before the object's first
// byte; a FIN mid-Object (§11.4) is io.ErrUnexpectedEOF.
func (o *SubgroupObject) Parse(r wire.Decoder, hasProperties bool) error {
        delta, err := r.Varint()
        if err != nil {
                return fmt.Errorf("moqt/message: object ID delta: %w", err)
        }
        o.ObjectIDDelta = delta

        // Properties are present only if the stream header has Properties == true
        if hasProperties {
                props, err := r.VarintBytes()
                if err != nil {
                        return fmt.Errorf("moqt/message: properties: %w", truncated(err))
                }
                o.Properties = props
        } else {
                o.Properties = nil
        }

        payloadLength, err := r.Varint()
        if err != nil {
                return fmt.Errorf("moqt/message: payload length: %w", truncated(err))
        }

        // Object Status is present only when Payload Length == 0
        if payloadLength == 0 {
                status, err := r.Varint()
                if err != nil {
                        return fmt.Errorf("moqt/message: object status: %w", truncated(err))
                }
                o.ObjectStatus = status
                o.Payload = nil
        } else {
                //nolint:gosec // G115: a payloadLength >= 2^63 wraps negative; both FixedBytes implementations reject it.
                payload, err := r.FixedBytes(int(payloadLength))
                if err != nil {
                        return fmt.Errorf("moqt/message: payload: %w", truncated(err))
                }
                o.Payload = payload
                o.ObjectStatus = 0
        }

        return nil
}

// Validate checks the SubgroupObject for protocol violations.
func (o *SubgroupObject) Validate() error {
        // Object Status can only be 0x0 (Normal), 0x3 (EndOfGroup), or 0x4 (EndOfTrack)
        if len(o.Payload) == 0 {
                switch o.ObjectStatus {
                case ObjectStatusNormal, ObjectStatusEndOfGroup, ObjectStatusEndOfTrack:
                        // Valid status values
                default:
                        return fmt.Errorf("moqt/message: invalid object status 0x%X", o.ObjectStatus)
                }
                // §11.2.1.2: no properties on a non-Normal status; a Properties Length
                // of 0 carries none (§11.4.2).
                if o.ObjectStatus != ObjectStatusNormal && len(o.Properties) > 0 {
                        return fmt.Errorf("moqt/message: object status 0x%X carries properties", o.ObjectStatus)
                }
        }

        return nil
}

// IsEndOfGroup reports whether this object signals End of Group (status 0x3).
func (o *SubgroupObject) IsEndOfGroup() bool {
        return len(o.Payload) == 0 && o.ObjectStatus == ObjectStatusEndOfGroup
}

// IsEndOfTrack reports whether this object signals End of Track (status 0x4).
func (o *SubgroupObject) IsEndOfTrack() bool {
        return len(o.Payload) == 0 && o.ObjectStatus == ObjectStatusEndOfTrack
}

// IsTerminal reports whether this object is a terminal status object
// (EndOfGroup or EndOfTrack) after which no further objects may appear on the
// same Subgroup stream (§11.4.3); a later object is a malformed track (§2.4.2).
func (o *SubgroupObject) IsTerminal() bool {
        return o.IsEndOfGroup() || o.IsEndOfTrack()
}

package message

import "github.com/floatdrop/moq-go/pkg/moqt/wire"

// Subscribe is the SUBSCRIBE message (§10.7).
type Subscribe struct {
        RequestID  uint64
        Namespace  wire.TrackNamespace
        Name       []byte
        Parameters Parameters
}

func (m *Subscribe) Type() Type             { return TypeSubscribe }
func (m *Subscribe) GetRequestID() uint64   { return m.RequestID }
func (m *Subscribe) SetRequestID(id uint64) { m.RequestID = id }

func (m *Subscribe) Append(w *wire.Writer) {
        w.Varint(m.RequestID)
        w.TrackNamespace(m.Namespace)
        w.VarintBytes(m.Name)
        m.Parameters.append(w)
}

func (m *Subscribe) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.Varint(&m.RequestID)
        s.TrackNamespace(&m.Namespace)
        s.VarintBytes(&m.Name)
        if err := s.Err(); err != nil {
                return err
        }
        return m.Parameters.parse(r)
}

// SubscribeOK is the SUBSCRIBE_OK message (§10.8). Track Properties span the
// remaining bytes; we currently treat them as opaque.
type SubscribeOK struct {
        TrackAlias      uint64
        Parameters      Parameters
        TrackProperties []byte
}

func (m *SubscribeOK) Type() Type { return TypeSubscribeOK }

func (m *SubscribeOK) Append(w *wire.Writer) {
        w.Varint(m.TrackAlias)
        m.Parameters.append(w)
        w.FixedBytes(m.TrackProperties)
}

func (m *SubscribeOK) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.Varint(&m.TrackAlias)
        if err := s.Err(); err != nil {
                return err
        }
        if err := m.Parameters.parse(r); err != nil {
                return err
        }
        m.TrackProperties = r.RemainingBytes()
        return nil
}

// Validate enforces the §2.4.1 Full Track Name size limit; ParsePayload
// invokes it automatically after decoding a SUBSCRIBE frame.
func (m *Subscribe) Validate() error {
        return validateFullTrackName(m.Namespace, m.Name)
}

package message

import (
        "slices"
        "time"
)

// DeliveryTimeouts holds the effective delivery timeout pair for one
// subscription per §8. Zero values mean "no timeout".
//
// Both values are expressed as time.Duration (internally milliseconds on the
// wire). A value of 0 means the timeout is disabled for that dimension.
type DeliveryTimeouts struct {
        Object   time.Duration // OBJECT_DELIVERY_TIMEOUT (§10.2.4)
        Subgroup time.Duration // SUBGROUP_DELIVERY_TIMEOUT (§10.2.3)
}

// FillTimeoutFromParam extracts the FILL_TIMEOUT parameter (§10.2.5) from ps
// and converts it from milliseconds to time.Duration. Returns 0 if the
// parameter is absent. FILL_TIMEOUT MAY appear in a FETCH message; it is the
// maximum total duration a relay should spend waiting for upstream sources to
// provide objects that are not immediately available.
func FillTimeoutFromParam(ps Parameters) time.Duration {
        d, _ := FillTimeoutFromParamOK(ps)
        return d
}

// FillTimeoutFromParamOK is [FillTimeoutFromParam] with presence reported
// separately. §10.2.5 gives an explicit value of 0 a meaning of its own — "the
// relay MUST NOT wait for upstream delivery and MUST report any unavailable
// Objects as Timed-Out gaps" — which a bare zero return cannot distinguish
// from the parameter being absent.
func FillTimeoutFromParamOK(ps Parameters) (d time.Duration, ok bool) {
        p, ok := ps.Find(ParamFillTimeout)
        if !ok {
                return 0, false
        }
        return MillisecondTimeout(p.Varint), true
}

// MillisecondTimeout converts a varint millisecond count — the form every §8
// timeout takes on the wire, whether it arrives as a Message Parameter
// (§10.2.3 / §10.2.4 / §10.2.5) or a Track/Object Property (§12.1 / §12.2) — to
// a time.Duration. Exported so every decoder of these values agrees by
// construction rather than by copies of the same multiplication.
//
//nolint:gosec // G115: a timeout in ms; an out-of-range value yields a wrong duration, not a memory-safety issue.
func MillisecondTimeout(ms uint64) time.Duration { return time.Duration(ms) * time.Millisecond }

// ObjectDeliveryTimeoutFromParam extracts OBJECT_DELIVERY_TIMEOUT (§10.2.4)
// from ps, converting from milliseconds. Returns 0 when absent (§8: 0 disables
// the timeout).
func ObjectDeliveryTimeoutFromParam(ps Parameters) time.Duration {
        p, ok := ps.Find(ParamObjectDeliveryTimeout)
        if !ok {
                return 0
        }
        return MillisecondTimeout(p.Varint)
}

// SubgroupDeliveryTimeoutFromParam extracts SUBGROUP_DELIVERY_TIMEOUT (§10.2.3)
// from ps, converting from milliseconds. Returns 0 when absent.
func SubgroupDeliveryTimeoutFromParam(ps Parameters) time.Duration {
        p, ok := ps.Find(ParamSubgroupDeliveryTimeout)
        if !ok {
                return 0
        }
        return MillisecondTimeout(p.Varint)
}

// DeliveryTimeoutsFromParams extracts both delivery timeouts (§10.2.3/§10.2.4)
// from ps — the form a subscriber communicates them in (§8).
func DeliveryTimeoutsFromParams(ps Parameters) DeliveryTimeouts {
        return DeliveryTimeouts{
                Object:   ObjectDeliveryTimeoutFromParam(ps),
                Subgroup: SubgroupDeliveryTimeoutFromParam(ps),
        }
}

// effectiveDim combines one timeout dimension per §8: "If both the publisher's
// value and the subscriber's value are non-zero, the smaller of the two is
// used." A zero value means "no timeout", so it never wins over a non-zero one.
func effectiveDim(publisher, subscriber time.Duration) time.Duration {
        switch {
        case publisher == 0:
                return subscriber
        case subscriber == 0:
                return publisher
        default:
                return min(publisher, subscriber)
        }
}

// Effective resolves the timeouts a publisher enforces for a subscription per
// §8: the receiver holds the publisher's values (Track Property, or the
// first-object Object Property override — see [DeliveryTimeouts.ApplyObjectProperties]),
// sub holds the subscriber's Message-Parameter values, and each dimension is
// the smaller of the two non-zero values.
func (d DeliveryTimeouts) Effective(sub DeliveryTimeouts) DeliveryTimeouts {
        return DeliveryTimeouts{
                Object:   effectiveDim(d.Object, sub.Object),
                Subgroup: effectiveDim(d.Subgroup, sub.Subgroup),
        }
}

// ApplyObjectProperties returns d with any OBJECT_DELIVERY_TIMEOUT (§12.2) or
// SUBGROUP_DELIVERY_TIMEOUT (§12.1) present in rawProps overriding the
// corresponding dimension. rawProps is the Object-Properties blob of the FIRST
// object in a subgroup (§12.1/§12.2: on the first object these override the
// Track-level value for that subgroup; on any other object they are ignored, so
// callers must invoke this only for the first object). A property present with
// value 0 overrides to "disabled"; an absent property leaves d's dimension
// unchanged. Malformed props leave d unchanged. Per §12.7 the contents of
// Immutable Properties are searched too, the mutable list winning.
func (d DeliveryTimeouts) ApplyObjectProperties(rawProps []byte) DeliveryTimeouts {
        if len(rawProps) == 0 {
                return d
        }
        pairs, err := parseSearchable(rawProps) // generic KV-pair decode; scope is the caller's
        if err != nil {
                return d
        }
        for _, kv := range slices.Backward(pairs) { // the mutable value wins (§12.7)
                switch kv.Type {
                case PropertyObjectDeliveryTimeout:
                        d.Object = MillisecondTimeout(kv.IntVal)
                case PropertySubgroupDeliveryTimeout:
                        d.Subgroup = MillisecondTimeout(kv.IntVal)
                }
        }
        return d
}

package message

import (
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// AliasType identifies the serialization and processing behavior of a Token
// per §10.2.2.
type AliasType uint64

const (
        // AliasTypeDelete (0x0): Alias only. Retire the alias and its associated
        // token from the cache.
        AliasTypeDelete AliasType = 0x0

        // AliasTypeRegister (0x1): Alias + Type + Value. Register the alias in
        // the token cache for the duration of the session (or until deleted).
        AliasTypeRegister AliasType = 0x1

        // AliasTypeUseAlias (0x2): Alias only. Resolve to the (Type, Value)
        // previously registered under this alias.
        AliasTypeUseAlias AliasType = 0x2

        // AliasTypeUseValue (0x3): Type + Value only. Use the token directly;
        // no alias is stored.
        AliasTypeUseValue AliasType = 0x3
)

// String returns a human-readable name for the alias type.
func (a AliasType) String() string {
        switch a {
        case AliasTypeDelete:
                return "DELETE"
        case AliasTypeRegister:
                return "REGISTER"
        case AliasTypeUseAlias:
                return "USE_ALIAS"
        case AliasTypeUseValue:
                return "USE_VALUE"
        }
        return fmt.Sprintf("AliasType(0x%X)", uint64(a))
}

// Token is the Token structure from §10.2.2.
//
// Wire format (within the outer KindBytes length-prefixed parameter value):
//
//        Token {
//          Alias Type (vi64),
//          [Token Alias (vi64),]   -- DELETE, REGISTER, USE_ALIAS
//          [Token Type (vi64),]    -- REGISTER, USE_VALUE
//          [Token Value (..)]      -- REGISTER, USE_VALUE; raw bytes to end of value
//        }
//
// TokenValue has no inner length prefix; it occupies the remainder of the
// outer KindBytes parameter value.
type Token struct {
        AliasType  AliasType
        TokenAlias uint64 // present for DELETE, REGISTER, USE_ALIAS
        TokenType  uint64 // present for REGISTER, USE_VALUE
        TokenValue []byte // present for REGISTER, USE_VALUE
}

// Append serialises t into w. The caller is responsible for the outer
// KindBytes length prefix (handled by params.go via VarintBytes).
func (t *Token) Append(w *wire.Writer) {
        w.Varint(uint64(t.AliasType))
        switch t.AliasType {
        case AliasTypeDelete:
                w.Varint(t.TokenAlias)
        case AliasTypeRegister:
                w.Varint(t.TokenAlias)
                w.Varint(t.TokenType)
                w.FixedBytes(t.TokenValue)
        case AliasTypeUseAlias:
                w.Varint(t.TokenAlias)
        case AliasTypeUseValue:
                w.Varint(t.TokenType)
                w.FixedBytes(t.TokenValue)
        }
}

// Bytes returns the serialised Token as a byte slice, suitable for use as the
// value of a KindBytes AUTHORIZATION_TOKEN parameter.
func (t *Token) Bytes() []byte {
        var w wire.Writer
        t.Append(&w)
        return w.Bytes()
}

// Parse deserialises a Token from raw — the raw bytes of a KindBytes parameter
// value. Returns an error (caller should map to KEY_VALUE_FORMATTING_ERROR) if
// the bytes are malformed.
func (t *Token) Parse(raw []byte) error {
        r := wire.NewReader(raw)

        at, err := r.Varint()
        if err != nil {
                return fmt.Errorf("moqt/message: token alias type: %w", err)
        }
        t.AliasType = AliasType(at)

        switch t.AliasType {
        case AliasTypeDelete:
                alias, err := r.Varint()
                if err != nil {
                        return fmt.Errorf("moqt/message: token alias (DELETE): %w", err)
                }
                t.TokenAlias = alias

        case AliasTypeRegister:
                alias, err := r.Varint()
                if err != nil {
                        return fmt.Errorf("moqt/message: token alias (REGISTER): %w", err)
                }
                tokenType, err := r.Varint()
                if err != nil {
                        return fmt.Errorf("moqt/message: token type (REGISTER): %w", err)
                }
                // TokenValue occupies the remainder of the parameter value.
                t.TokenAlias = alias
                t.TokenType = tokenType
                t.TokenValue = r.RemainingBytes()

        case AliasTypeUseAlias:
                alias, err := r.Varint()
                if err != nil {
                        return fmt.Errorf("moqt/message: token alias (USE_ALIAS): %w", err)
                }
                t.TokenAlias = alias

        case AliasTypeUseValue:
                tokenType, err := r.Varint()
                if err != nil {
                        return fmt.Errorf("moqt/message: token type (USE_VALUE): %w", err)
                }
                t.TokenType = tokenType
                t.TokenValue = r.RemainingBytes()

        default:
                return fmt.Errorf("moqt/message: unknown token alias type 0x%X", at)
        }

        // DELETE and USE_ALIAS carry only an Alias (§10.2.2).
        if n := r.Remaining(); n > 0 {
                return fmt.Errorf("moqt/message: %d trailing bytes after token alias", n)
        }
        return nil
}

// AuthorizationTokenParam builds a typed AUTHORIZATION_TOKEN parameter
// (§10.2.2) from a Token. The Token is serialised to bytes and stored as a
// KindBytes parameter.
func AuthorizationTokenParam(t Token) Parameter {
        return BytesParam(ParamAuthorizationToken, t.Bytes())
}

// TokensFromParam extracts and parses all AUTHORIZATION_TOKEN parameters from
// ps. The spec allows the parameter to be repeated within a message (§10.2.2:
// "The AUTHORIZATION TOKEN parameter MAY be repeated within a message as long
// as the combination of Token Type and Token Value are unique after resolving
// any aliases"). Returns an error if any Token is malformed.
func TokensFromParam(ps Parameters) ([]Token, error) {
        var tokens []Token
        for _, p := range ps {
                if p.Type != ParamAuthorizationToken {
                        continue
                }
                var t Token
                if err := t.Parse(p.Bytes); err != nil {
                        return nil, fmt.Errorf("moqt/message: AUTHORIZATION_TOKEN: %w", err)
                }
                tokens = append(tokens, t)
        }
        return tokens, nil
}

package message

import "github.com/floatdrop/moq-go/pkg/moqt/wire"

// TrackStatus is the TRACK_STATUS message (§10.15). It queries the status
// of a track without creating a subscription. The message format is identical
// to SUBSCRIBE, but subscriber-specific parameters (like SUBSCRIBER_PRIORITY)
// must not be included.
type TrackStatus struct {
        RequestID  uint64
        Namespace  wire.TrackNamespace
        Name       []byte
        Parameters Parameters
}

// Type returns the wire type ID for TRACK_STATUS.
func (m *TrackStatus) Type() Type             { return TypeTrackStatus }
func (m *TrackStatus) GetRequestID() uint64   { return m.RequestID }
func (m *TrackStatus) SetRequestID(id uint64) { m.RequestID = id }

// Append serializes the TRACK_STATUS message to w.
func (m *TrackStatus) Append(w *wire.Writer) {
        w.Varint(m.RequestID)
        w.TrackNamespace(m.Namespace)
        w.VarintBytes(m.Name)
        m.Parameters.append(w)
}

// Parse deserializes the TRACK_STATUS message from r.
func (m *TrackStatus) Parse(r *wire.Reader) error {
        s := r.Scanner()
        s.Varint(&m.RequestID)
        s.TrackNamespace(&m.Namespace)
        s.VarintBytes(&m.Name)
        if err := s.Err(); err != nil {
                return err
        }
        return m.Parameters.parse(r)
}

// TrackStatusOK is the TRACK_STATUS_OK response (§10.15). Per the spec,
// TRACK_STATUS_OK is a REQUEST_OK (type 0x07) sent in response to TRACK_STATUS.
// It carries the same parameters and Track Properties as SUBSCRIBE_OK, but
// without a Track Alias since no subscription is created.
//
// Use RequestOK directly when sending; TrackStatusOK is a convenience alias
// that wraps RequestOK for clarity at call sites.
type TrackStatusOK = RequestOK

// Validate enforces the §2.4.1 Full Track Name size limit; ParsePayload
// invokes it automatically after decoding a TRACK_STATUS frame.
func (m *TrackStatus) Validate() error {
        return validateFullTrackName(m.Namespace, m.Name)
}

// Package message implements MoQT control- and request-stream message types
// per draft-ietf-moq-transport-20. Each Message exposes a wire Type and
// Append/Parse methods over wire.Writer/Reader.
//
// Marshal writes a complete control-message frame (Type + Length + Payload).
// Parse reads the payload only; the caller is expected to have already read
// the frame header via wire.ReadFrame and dispatched on Type.
package message

import (
        "errors"
        "fmt"
        "io"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// Type is the wire type ID for a MoQT message (§10, table 5).
type Type uint64

const (
        TypeSetup              Type = 0x2F00
        TypeGoaway             Type = 0x10
        TypeSubscribe          Type = 0x03
        TypeSubscribeOK        Type = 0x04
        TypePublish            Type = 0x1D
        TypePublishStateNotify Type = 0x22
        TypePublishDone        Type = 0x0B
        TypeRequestUpdate      Type = 0x02
        TypeRequestOK          Type = 0x07
        TypeRequestError       Type = 0x05
        TypeFetch              Type = 0x16
        TypeFetchOK            Type = 0x18
        TypeTrackStatus        Type = 0x0D
        TypePublishNamespace   Type = 0x06
        TypeNamespace          Type = 0x08
        TypeNamespaceDone      Type = 0x0E
        TypeSubscribeNamespace Type = 0x50
        TypeSubscribeTracks    Type = 0x51
        TypePublishSkipped     Type = 0x0F
)

// Message is the interface implemented by all in-scope MoQT control- and
// request-stream messages.
type Message interface {
        // Type returns the wire type ID.
        Type() Type
        // Append serializes the message payload to w.
        Append(w *wire.Writer)
        // Parse deserializes the message payload from r. r is expected to be
        // bounded to the payload length (i.e. the wire-level frame length).
        Parse(r *wire.Reader) error
}

// WithRequestID is implemented by messages that carry a Request ID as their
// first field (§10.1). These are the messages that can appear as the first
// message on a request stream: SUBSCRIBE, PUBLISH, FETCH, TRACK_STATUS,
// PUBLISH_NAMESPACE, SUBSCRIBE_NAMESPACE, SUBSCRIBE_TRACKS, and
// REQUEST_UPDATE.
type WithRequestID interface {
        Message
        // GetRequestID returns the Request ID carried by this message.
        GetRequestID() uint64
        // SetRequestID overwrites the Request ID carried by this message. The
        // session uses it to assign a freshly allocated ID (§10.1) after a
        // request stream is opened, so a failed open consumes no ID.
        SetRequestID(uint64)
}

// Marshal writes m as a complete control-message frame to dst.
func Marshal(dst io.Writer, m Message) error {
        w := wire.NewWriter(nil)
        m.Append(w)
        return wire.WriteFrame(dst, uint64(m.Type()), w.Bytes())
}

// Parse reads a single control-message frame from src and returns a typed
// Message. Once the frame is read whole, every error wraps
// [ErrMalformedMessage]; an unknown type also wraps ErrUnknownType.
func Parse(src io.Reader) (Message, error) {
        t, payload, err := wire.ReadFrame(src)
        if err != nil {
                return nil, err
        }
        return ParsePayload(Type(t), payload)
}

// ErrMalformedMessage is wrapped by every error [ParsePayload] returns, and by
// every error [Parse] returns once it has read a whole frame: an unknown type,
// a Length mismatch, or a field that fails validation, each session-fatal
// (§10). A frame that could not be read whole does not wrap it.
var ErrMalformedMessage = errors.New("moqt/message: malformed message")

// ParsePayload constructs a Message of the given Type and parses payload into
// it. Use when the caller has already read the frame header. Every error it
// returns wraps [ErrMalformedMessage].
func ParsePayload(t Type, payload []byte) (Message, error) {
        m, err := newMessage(t)
        if err != nil {
                return nil, fmt.Errorf("%w: %w", ErrMalformedMessage, err)
        }
        r := wire.NewReader(payload)
        if err := m.Parse(r); err != nil {
                return nil, fmt.Errorf("%w: parsing %s: %w", ErrMalformedMessage, t, err)
        }
        if !r.Empty() {
                return nil, fmt.Errorf("%w: %s has %d trailing bytes", ErrMalformedMessage, t, r.Remaining())
        }
        if v, ok := m.(validator); ok {
                if err := v.Validate(); err != nil {
                        return nil, fmt.Errorf("%w: validating %s: %w", ErrMalformedMessage, t, err)
                }
        }
        return m, nil
}

// validator is implemented by messages that enforce field-level invariants the
// wire decoder cannot catch on its own (e.g. a FETCH whose full track name
// exceeds §2.4.1's 4,096-byte cap). ParsePayload invokes Validate after a
// successful decode so a malformed-but-decodable control message is rejected at
// the message boundary — the session layer treats the resulting error as a
// PROTOCOL_VIOLATION — rather than propagating bad state inward.
type validator interface {
        Validate() error
}

// ErrUnknownType is returned for a message type not implemented by this
// package. Per §10 the receiver MUST close the session; §10 names no error
// code for this, and callers close with PROTOCOL_VIOLATION.
type ErrUnknownType Type

func (e ErrUnknownType) Error() string {
        return fmt.Sprintf("moqt/message: unknown type %#x", uint64(e))
}

func newMessage(t Type) (Message, error) {
        switch t {
        case TypeSetup:
                return &Setup{}, nil
        case TypeGoaway:
                return &Goaway{}, nil
        case TypeSubscribe:
                return &Subscribe{}, nil
        case TypeSubscribeOK:
                return &SubscribeOK{}, nil
        case TypePublish:
                return &Publish{}, nil
        case TypePublishStateNotify:
                return &PublishStateNotify{}, nil
        case TypePublishDone:
                return &PublishDone{}, nil
        case TypeRequestUpdate:
                return &RequestUpdate{}, nil
        case TypeRequestOK:
                return &RequestOK{}, nil
        case TypeRequestError:
                return &RequestError{}, nil
        case TypeFetch:
                return &Fetch{}, nil
        case TypeFetchOK:
                return &FetchOK{}, nil
        case TypeTrackStatus:
                return &TrackStatus{}, nil
        case TypePublishNamespace:
                return &PublishNamespace{}, nil
        case TypeNamespace:
                return &Namespace{}, nil
        case TypeNamespaceDone:
                return &NamespaceDone{}, nil
        case TypeSubscribeNamespace:
                return &SubscribeNamespace{}, nil
        case TypeSubscribeTracks:
                return &SubscribeTracks{}, nil
        case TypePublishSkipped:
                return &PublishSkipped{}, nil
        }
        return nil, ErrUnknownType(t)
}

// String returns a short identifier for the message type.
func (t Type) String() string {
        switch t {
        case TypeSetup:
                return "SETUP"
        case TypeGoaway:
                return "GOAWAY"
        case TypeSubscribe:
                return "SUBSCRIBE"
        case TypeSubscribeOK:
                return "SUBSCRIBE_OK"
        case TypePublish:
                return "PUBLISH"
        case TypePublishStateNotify:
                return "PUBLISH_STATE_NOTIFY"
        case TypePublishDone:
                return "PUBLISH_DONE"
        case TypeRequestUpdate:
                return "REQUEST_UPDATE"
        case TypeRequestOK:
                return "REQUEST_OK"
        case TypeRequestError:
                return "REQUEST_ERROR"
        case TypeFetch:
                return "FETCH"
        case TypeFetchOK:
                return "FETCH_OK"
        case TypeTrackStatus:
                return "TRACK_STATUS"
        case TypePublishNamespace:
                return "PUBLISH_NAMESPACE"
        case TypeNamespace:
                return "NAMESPACE"
        case TypeNamespaceDone:
                return "NAMESPACE_DONE"
        case TypeSubscribeNamespace:
                return "SUBSCRIBE_NAMESPACE"
        case TypeSubscribeTracks:
                return "SUBSCRIBE_TRACKS"
        case TypePublishSkipped:
                return "PUBLISH_SKIPPED"
        }
        return fmt.Sprintf("Type(%#x)", uint64(t))
}

package msf

import (
        "bytes"
        "encoding/json"
        "fmt"
)

// Catalog is an MSF catalog document (§5). A Catalog is either:
//
//   - An independent catalog: Version is set and Tracks lists the full
//     output of the publisher (§5.1).
//   - A delta update: DeltaUpdate carries an ordered list of operations
//     and Version / Tracks MUST be absent (§5.3).
//
// Catalog preserves producer-defined fields not described by the draft
// in Extras (§5.1). Unknown fields round-trip verbatim. Fields not
// listed in the draft and not present in Extras are silently dropped on
// re-serialisation.
//
// As of draft-01 a delta update is expressed as the deltaUpdate array
// (§5.1.6): an ordered sequence of [DeltaOp] objects, each naming an
// "op" ("add"/"remove"/"clone") and a list of track objects. [Apply]
// replays the operations in order per §5.3.
//
//nolint:recvcheck // MarshalJSON must stay on a value receiver so non-pointer Catalog values marshal; the mutating/validating methods must stay on pointers.
type Catalog struct {
        Version     string `json:"version,omitempty"`
        GeneratedAt int64  `json:"generatedAt,omitempty"`
        IsComplete  bool   `json:"isComplete,omitempty"`
        // Tracks is required in independent catalogs (§5.1.4). Empty
        // slices (terminator catalogs per §11.3) and nil slices are
        // emitted differently: see [Catalog.MarshalJSON].
        Tracks []Track `json:"tracks"`
        // PublishTracks declares tracks the subscriber may publish to,
        // such as logs or metrics (§5.1.5).
        PublishTracks []Track `json:"publishTracks,omitempty"`
        // DeltaUpdate, when non-nil, marks this catalog as a delta update
        // (§5.1.6). It is an ordered list of operations applied by [Apply].
        DeltaUpdate []DeltaOp `json:"deltaUpdate,omitempty"`
        // InitDataList holds initialization payloads referenced by tracks
        // via Track.InitRef (§5.1.7). Per §5.1.7 it SHOULD appear after the
        // tracks array in the document.
        InitDataList []InitData `json:"initDataList,omitempty"`
        // ContentProtections declares DRM/CENC key-acquisition metadata
        // referenced by tracks via Track.ContentProtectionRefIDs
        // (draft-ietf-moq-cmsf-01 §4.1.1). Per §4.1.1, content protection
        // information MUST NOT be duplicated at the track level.
        ContentProtections []ContentProtection `json:"contentProtections,omitempty"`

        // Extras holds producer-defined catalog-root fields. Keys MUST NOT
        // collide with the known field names; this is the producer's
        // responsibility (§5.1).
        Extras map[string]any `json:"-"`
}

// IsDelta reports whether the catalog is a delta update (§5.1.6) rather
// than an independent catalog.
func (c *Catalog) IsDelta() bool {
        return c.DeltaUpdate != nil
}

// DeltaOp is one entry in a catalog's deltaUpdate array (§5.1.6). Op is
// one of [DeltaOpAdd], [DeltaOpRemove] or [DeltaOpClone]; Tracks is the
// list of track objects the operation applies, in document order.
type DeltaOp struct {
        Op     string  `json:"op"`
        Tracks []Track `json:"tracks"`
}

// InitData is one entry in a catalog's initDataList (§5.1.7). Type is
// the reference type ([InitDataTypeInline] is the only one defined) and
// Data carries the payload as defined by that type.
type InitData struct {
        ID   string `json:"id"`
        Type string `json:"type"`
        Data string `json:"data"`
}

// Buffers describes a track's target jitter/forward buffers (§5.2.9).
// All keys are optional; absent keys leave the player free to choose.
type Buffers struct {
        Target *uint32 `json:"target,omitempty"`
        Min    *uint32 `json:"min,omitempty"`
        Max    *uint32 `json:"max,omitempty"`
}

// Accessibility is one accessibility descriptor embedded in a track
// (§5.2.44): a scheme identifier and a scheme-specific value.
type Accessibility struct {
        Scheme string `json:"scheme"`
        Value  string `json:"value"`
}

// Track is a single entry in a Catalog's Tracks / PublishTracks array
// or in a [DeltaOp]'s Tracks list (§5.2.1). Most fields are optional;
// required fields depend on the role this track plays:
//
//   - Independent catalog tracks: Name, Packaging, IsLive required.
//   - add / clone operation entries: Name required.
//   - remove operation entries: Name required, all other fields MUST be
//     absent (§5.1.6).
//   - clone operation entries: ParentName required (§5.1.6).
//
// Pointer-typed fields (IsLive, TargetLatency, RenderGroup, AltGroup,
// TemporalID, SpatialID, Buffers, Template, MaxGrpSapStartingType,
// MaxObjSapStartingType) distinguish "field absent"
// from "field set to zero/false". The remaining numeric / string fields
// use omitempty because zero is never a valid catalog value (e.g.
// bitrate=0).
type Track struct {
        Name             string                 `json:"name,omitempty"`
        Namespace        string                 `json:"namespace,omitempty"`
        Packaging        string                 `json:"packaging,omitempty"`
        EventType        string                 `json:"eventType,omitempty"`
        IsLive           *bool                  `json:"isLive,omitempty"`
        TargetLatency    *uint32                `json:"targetLatency,omitempty"`
        Buffers          *Buffers               `json:"buffers,omitempty"`
        Role             string                 `json:"role,omitempty"`
        Label            string                 `json:"label,omitempty"`
        RenderGroup      *int                   `json:"renderGroup,omitempty"`
        AltGroup         *int                   `json:"altGroup,omitempty"`
        InitRef          string                 `json:"initRef,omitempty"`
        Depends          []string               `json:"depends,omitempty"`
        Template         *MediaTimelineTemplate `json:"template,omitempty"`
        TemporalID       *int                   `json:"temporalId,omitempty"`
        SpatialID        *int                   `json:"spatialId,omitempty"`
        Codec            string                 `json:"codec,omitempty"`
        Mimetype         string                 `json:"mimetype,omitempty"`
        Framerate        float64                `json:"framerate,omitempty"`
        Timescale        uint32                 `json:"timescale,omitempty"`
        Bitrate          uint64                 `json:"bitrate,omitempty"`
        AvgBitrate       uint64                 `json:"avgBitrate,omitempty"`
        MaxGopDuration   uint64                 `json:"maxGopDuration,omitempty"`
        MaxGroupDuration uint64                 `json:"maxGroupDuration,omitempty"`
        Width            uint32                 `json:"width,omitempty"`
        Height           uint32                 `json:"height,omitempty"`
        Samplerate       uint32                 `json:"samplerate,omitempty"`
        ChannelConfig    string                 `json:"channelConfig,omitempty"`
        DisplayWidth     uint32                 `json:"displayWidth,omitempty"`
        DisplayHeight    uint32                 `json:"displayHeight,omitempty"`
        Lang             string                 `json:"lang,omitempty"`
        ParentName       string                 `json:"parentName,omitempty"`
        ParentNamespace  string                 `json:"parentNamespace,omitempty"`
        TrackDuration    uint64                 `json:"trackDuration,omitempty"`
        ConnectionURI    string                 `json:"connectionUri,omitempty"`
        Token            string                 `json:"token,omitempty"`
        EncryptionScheme string                 `json:"encryptionScheme,omitempty"`
        CipherSuite      string                 `json:"cipherSuite,omitempty"`
        KeyID            string                 `json:"keyId,omitempty"`
        TrackBaseKey     string                 `json:"trackBaseKey,omitempty"`
        AuthInfo         map[string]any         `json:"authInfo,omitempty"`
        Accessibility    []Accessibility        `json:"accessibility,omitempty"`

        // MaxGrpSapStartingType and MaxObjSapStartingType bound the stream
        // access point type a Group / Object may start with
        // (draft-ietf-moq-cmsf-01 §3.5.2.1, §3.5.2.2). Valid range 0-3.
        MaxGrpSapStartingType *int `json:"maxGrpSapStartingType,omitempty"`
        MaxObjSapStartingType *int `json:"maxObjSapStartingType,omitempty"`
        // ContentProtectionRefIDs references Catalog.ContentProtections
        // entries by RefID (CMSF §4.1.2). Presence means the track is
        // CENC-encrypted and a subscriber MUST acquire licenses before
        // decryption.
        ContentProtectionRefIDs []string `json:"contentProtectionRefIDs,omitempty"`

        // Extras holds producer-defined per-track fields (§5.6.6 example).
        // Keys MUST NOT collide with known field names.
        Extras map[string]any `json:"-"`
}

// knownCatalogFields lists every JSON key produced by Catalog's typed
// fields. Used during UnmarshalJSON to separate known fields from
// Extras.
var knownCatalogFields = map[string]struct{}{
        "version":            {},
        "generatedAt":        {},
        "isComplete":         {},
        "tracks":             {},
        "publishTracks":      {},
        "deltaUpdate":        {},
        "initDataList":       {},
        "contentProtections": {},
}

// knownTrackFields lists every JSON key produced by Track's typed
// fields.
var knownTrackFields = map[string]struct{}{
        "name":                    {},
        "namespace":               {},
        "packaging":               {},
        "eventType":               {},
        "isLive":                  {},
        "targetLatency":           {},
        "buffers":                 {},
        "role":                    {},
        "label":                   {},
        "renderGroup":             {},
        "altGroup":                {},
        "initRef":                 {},
        "depends":                 {},
        "template":                {},
        "temporalId":              {},
        "spatialId":               {},
        "codec":                   {},
        "mimetype":                {},
        "framerate":               {},
        "timescale":               {},
        "bitrate":                 {},
        "avgBitrate":              {},
        "maxGopDuration":          {},
        "maxGroupDuration":        {},
        "width":                   {},
        "height":                  {},
        "samplerate":              {},
        "channelConfig":           {},
        "displayWidth":            {},
        "displayHeight":           {},
        "lang":                    {},
        "parentName":              {},
        "parentNamespace":         {},
        "trackDuration":           {},
        "connectionUri":           {},
        "token":                   {},
        "encryptionScheme":        {},
        "cipherSuite":             {},
        "keyId":                   {},
        "trackBaseKey":            {},
        "authInfo":                {},
        "accessibility":           {},
        "maxGrpSapStartingType":   {},
        "maxObjSapStartingType":   {},
        "contentProtectionRefIDs": {},
}

// trackAlias decouples the JSON tag-driven marshaller from the
// Catalog/Track methods so MarshalJSON / UnmarshalJSON do not recurse.
type trackAlias Track

// catalogAlias plays the same role as trackAlias for Catalog.
type catalogAlias Catalog

// MarshalJSON emits the typed Track fields and merges Extras. If a key
// in Extras shadows a typed field the typed field wins; the collision
// is silently resolved in favour of the typed value because §5.1 makes
// collision the producer's responsibility.
func (t Track) MarshalJSON() ([]byte, error) {
        return mergeMarshal(trackAlias(t), t.Extras, knownTrackFields)
}

// UnmarshalJSON parses the typed Track fields and stores any other
// keys in Extras.
func (t *Track) UnmarshalJSON(data []byte) error {
        var alias trackAlias
        if err := strictUnmarshal(data, &alias); err != nil {
                return fmt.Errorf("moqt/msf: track: %w", err)
        }
        *t = Track(alias)

        extras, err := extractExtras(data, knownTrackFields)
        if err != nil {
                return fmt.Errorf("moqt/msf: track extras: %w", err)
        }
        t.Extras = extras
        return nil
}

// MarshalJSON emits the typed Catalog fields and merges Extras.
//
// MarshalJSON enforces the §5.1.4 / §5.3 rules around the tracks key:
//
//   - Independent catalogs (DeltaUpdate==nil) always include "tracks";
//     a nil slice is emitted as the empty array [] expected by the
//     §11.3 terminator example.
//   - Delta updates (DeltaUpdate!=nil) MUST NOT include "tracks" (§5.3);
//     MarshalJSON drops the key.
func (c Catalog) MarshalJSON() ([]byte, error) {
        if c.IsDelta() {
                // Strip Tracks so the alias marshaller emits "tracks": null,
                // then post-process to drop the key entirely. Using a custom
                // post-process keeps the typed-fields path symmetrical with
                // independent catalogs.
                c.Tracks = nil
                data, err := mergeMarshal(catalogAlias(c), c.Extras, knownCatalogFields)
                if err != nil {
                        return nil, err
                }
                return stripNullTracks(data)
        }
        if c.Tracks == nil {
                c.Tracks = []Track{}
        }
        return mergeMarshal(catalogAlias(c), c.Extras, knownCatalogFields)
}

// stripNullTracks removes a "tracks": null entry from the top-level
// JSON object. Used by MarshalJSON for delta catalogs.
func stripNullTracks(data []byte) ([]byte, error) {
        var m map[string]json.RawMessage
        if err := json.Unmarshal(data, &m); err != nil {
                return nil, err
        }
        if raw, ok := m["tracks"]; ok && bytes.Equal(raw, []byte("null")) {
                delete(m, "tracks")
        }
        return json.Marshal(m)
}

// UnmarshalJSON parses the typed Catalog fields and stores unknown
// catalog-root keys in Extras.
func (c *Catalog) UnmarshalJSON(data []byte) error {
        var alias catalogAlias
        if err := strictUnmarshal(data, &alias); err != nil {
                return fmt.Errorf("moqt/msf: catalog: %w", err)
        }
        *c = Catalog(alias)

        extras, err := extractExtras(data, knownCatalogFields)
        if err != nil {
                return fmt.Errorf("moqt/msf: catalog extras: %w", err)
        }
        c.Extras = extras
        return nil
}

// mergeMarshal serialises v (which must have JSON tags matching the
// known field set) and merges entries from extras whose keys are not
// already produced by v. Keys in extras that collide with v's known
// fields are silently dropped.
func mergeMarshal(v any, extras map[string]any, known map[string]struct{}) ([]byte, error) {
        base, err := json.Marshal(v)
        if err != nil {
                return nil, err
        }
        if len(extras) == 0 {
                return base, nil
        }

        // Decode base back into a map so we can re-emit in deterministic
        // order. The size cost is acceptable for catalog documents.
        var merged map[string]json.RawMessage
        if err := json.Unmarshal(base, &merged); err != nil {
                return nil, err
        }
        for k, val := range extras {
                if _, isKnown := known[k]; isKnown {
                        continue
                }
                raw, err := json.Marshal(val)
                if err != nil {
                        return nil, fmt.Errorf("moqt/msf: marshal extras[%q]: %w", k, err)
                }
                merged[k] = raw
        }
        return json.Marshal(merged)
}

// extractExtras returns the entries in data whose keys are not in known.
// Returns nil (not empty map) when there are no extras, so callers can
// omit the field entirely on a fresh struct.
func extractExtras(data []byte, known map[string]struct{}) (map[string]any, error) {
        var raw map[string]json.RawMessage
        if err := json.Unmarshal(data, &raw); err != nil {
                return nil, err
        }
        var extras map[string]any
        for k, v := range raw {
                if _, isKnown := known[k]; isKnown {
                        continue
                }
                var decoded any
                if err := json.Unmarshal(v, &decoded); err != nil {
                        return nil, fmt.Errorf("extras[%q]: %w", k, err)
                }
                if extras == nil {
                        extras = map[string]any{}
                }
                extras[k] = decoded
        }
        return extras, nil
}

// strictUnmarshal decodes data into v.  It uses a Decoder so future
// additions (e.g. DisallowUnknownFields) can be enabled without
// touching every call site.
func strictUnmarshal(data []byte, v any) error {
        dec := json.NewDecoder(bytes.NewReader(data))
        return dec.Decode(v)
}

package msf

import "fmt"

// Scheme values for ContentProtection.Scheme (draft-ietf-moq-cmsf-01
// §4.1.1.3, Table 3). SchemeCBCS is RECOMMENDED for better hardware
// decoder compatibility.
const (
        SchemeCENC = "cenc"
        SchemeCBCS = "cbcs"
)

// Well-known DRM system IDs for DRMSystem.SystemID (CMSF §4.1.1.4.1,
// Table 4).
const (
        DRMSystemIDWidevine  = "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed"
        DRMSystemIDPlayReady = "9a04f079-9840-4286-ab92-e65be0885f95"
        DRMSystemIDFairPlay  = "94ce86fb-07ff-4f43-adb8-93d2fa968ca2"
        DRMSystemIDClearKey  = "1077efec-c0b2-4d02-ace3-3c1e52e2fb4b"
)

// EventTypeCMSFSAP is the eventType value for a SAP Type timeline
// track (CMSF §3.6.1): a track with Packaging ==
// [PackagingEventTimeline] and EventType == EventTypeCMSFSAP. Its
// records convey the distribution of Stream Access Point types and
// their earliest presentation times; use [SAPRecord] to encode and
// decode them.
const EventTypeCMSFSAP = "org.ietf.moq.cmsf.sap"

// ContentProtection is one root-level entry in a catalog's
// contentProtections array (CMSF §4.1.1). Tracks reference an entry by
// RefID via Track.ContentProtectionRefIDs; per §4.1.1, content
// protection information MUST NOT be duplicated at the track level —
// all tracks reference these root-level entries.
//
// CMSF §4.2 additionally requires the initialization data of a
// protected track to carry the 'sinf'/'schm'/'schi'/'tenc' boxes. That
// data is opaque Base64 ISO BMFF in InitData.Data, which this package
// does not parse, so the requirement is the producer's to meet.
type ContentProtection struct {
        RefID      string    `json:"refID"`
        DefaultKID []string  `json:"defaultKID"`
        Scheme     string    `json:"scheme"`
        DRMSystem  DRMSystem `json:"drmSystem"`

        // Extras holds producer-defined fields on this entry, mirroring
        // [Catalog].Extras and [Track].Extras so a re-serialised catalog
        // preserves keys this implementation does not know.
        Extras map[string]any `json:"-"`
}

// DRMSystem describes one DRM system's key-acquisition metadata within
// a ContentProtection entry (CMSF §4.1.1.4).
type DRMSystem struct {
        SystemID string `json:"systemID"`
        // LAURL and CertURL are §4.1.1.4.2 and §4.1.1.4.3, whose JSON keys
        // the §5 examples fix.
        //
        // §4.1.1.4.4's Authorization URL has no typed field: the section
        // never names its JSON key and no example carries one, and since
        // laURL/certURL are abbreviations rather than section-title camel
        // case, any spelling this package chose would be a guess a future
        // revision could contradict. Extras carries it losslessly under
        // whatever key the producer used, so nothing is dropped; a typed
        // field can be added once a draft names one.
        LAURL      *URLRef `json:"laURL,omitempty"`
        CertURL    *URLRef `json:"certURL,omitempty"`
        PSSH       string  `json:"pssh,omitempty"`
        Robustness string  `json:"robustness,omitempty"`

        // Extras holds producer-defined fields on this object.
        Extras map[string]any `json:"-"`
}

// URLRef is a {url, type} pair used by DRMSystem.LAURL and CertURL
// (CMSF §4.1.1.4.2, §4.1.1.4.3). URL is required whenever the enclosing
// object is present; Type is optional and its meaning is per-field
// (license protocol, certificate MIME type).
type URLRef struct {
        URL  string `json:"url"`
        Type string `json:"type,omitempty"`
}

// knownContentProtectionFields lists every JSON key produced by
// ContentProtection's typed fields.
var knownContentProtectionFields = map[string]struct{}{
        "refID":      {},
        "defaultKID": {},
        "scheme":     {},
        "drmSystem":  {},
}

// knownDRMSystemFields lists every JSON key produced by DRMSystem's
// typed fields.
var knownDRMSystemFields = map[string]struct{}{
        "systemID":   {},
        "laURL":      {},
        "certURL":    {},
        "pssh":       {},
        "robustness": {},
}

// contentProtectionAlias and drmSystemAlias play the same role as
// [trackAlias]: they decouple the JSON tag-driven marshaller from the
// MarshalJSON / UnmarshalJSON methods so the calls do not recurse.
type (
        contentProtectionAlias ContentProtection
        drmSystemAlias         DRMSystem
)

// MarshalJSON emits the typed fields and merges Extras.
func (p ContentProtection) MarshalJSON() ([]byte, error) {
        return mergeMarshal(contentProtectionAlias(p), p.Extras, knownContentProtectionFields)
}

// UnmarshalJSON parses the typed fields and stores any other keys in
// Extras.
func (p *ContentProtection) UnmarshalJSON(data []byte) error {
        var alias contentProtectionAlias
        if err := strictUnmarshal(data, &alias); err != nil {
                return fmt.Errorf("moqt/msf: contentProtection: %w", err)
        }
        *p = ContentProtection(alias)

        extras, err := extractExtras(data, knownContentProtectionFields)
        if err != nil {
                return fmt.Errorf("moqt/msf: contentProtection extras: %w", err)
        }
        p.Extras = extras
        return nil
}

// MarshalJSON emits the typed fields and merges Extras.
func (d DRMSystem) MarshalJSON() ([]byte, error) {
        return mergeMarshal(drmSystemAlias(d), d.Extras, knownDRMSystemFields)
}

// UnmarshalJSON parses the typed fields and stores any other keys in
// Extras.
func (d *DRMSystem) UnmarshalJSON(data []byte) error {
        var alias drmSystemAlias
        if err := strictUnmarshal(data, &alias); err != nil {
                return fmt.Errorf("moqt/msf: drmSystem: %w", err)
        }
        *d = DRMSystem(alias)

        extras, err := extractExtras(data, knownDRMSystemFields)
        if err != nil {
                return fmt.Errorf("moqt/msf: drmSystem extras: %w", err)
        }
        d.Extras = extras
        return nil
}

package msf

import (
        "encoding/json"
        "fmt"
)

// SAPRecord is one decoded record of a SAP Type timeline track
// (CMSF §3.6.1). On the wire it is an Event Timeline record indexed by
// Location ('l') whose data field is a two-integer JSON array:
//
//        { "l": [GroupID, ObjectID], "data": [SAPType, EPT] }
type SAPRecord struct {
        GroupID  uint64
        ObjectID uint64
        // SAPType is 0-3. 0 means the Object does not start with an ISOBMFF
        // stream access point; 1, 2 and 3 mean it begins with a SAP of that
        // type. When the Object is the first in its Group the value MUST be
        // 1 or 2.
        SAPType int
        // EPT is the earliest media presentation timestamp, rounded to the
        // nearest millisecond, of all media samples in the Object the
        // record's location identifies.
        EPT int64
}

// validate enforces CMSF §3.6.1's constraints on the SAP type.
func (r SAPRecord) validate() error {
        if r.SAPType < 0 || r.SAPType > 3 {
                return fmt.Errorf("moqt/msf: sap record: sapType %d out of range 0-3 (CMSF §3.6.1)", r.SAPType)
        }
        // §3.6.1: "When the Object is the first Object in the Group, the
        // value MUST be equal to 1 or 2." This restates §3.4's requirement
        // that every Group begin with a SAP type 1 or 2 Object.
        //
        // Object ID 0 is the only first-in-Group case a single record can
        // prove: [MoQTransport] §2.1 lets Object IDs start above 0 and skip
        // values, so a Group whose first Object is, say, 5 is
        // indistinguishable here from a mid-Group record. Checking it would
        // need the whole Group, and a timeline document may legitimately
        // begin mid-Group, so the stricter check would reject conformant
        // input. Producers remain responsible for §3.4 in that case.
        if r.ObjectID == 0 && r.SAPType != 1 && r.SAPType != 2 {
                return fmt.Errorf(
                        "moqt/msf: sap record: group %d starts with sapType %d, MUST be 1 or 2 (CMSF §3.6.1)",
                        r.GroupID, r.SAPType)
        }
        return nil
}

// EventRecord encodes r as the Event Timeline record CMSF §3.6.1
// defines. It reports an error if r violates the section's SAP-type
// constraints.
func (r SAPRecord) EventRecord() (EventRecord, error) {
        if err := r.validate(); err != nil {
                return EventRecord{}, err
        }
        data, err := json.Marshal([2]int64{int64(r.SAPType), r.EPT})
        if err != nil {
                return EventRecord{}, err
        }
        return EventRecord{
                Index:    EventIndexLocation,
                GroupID:  r.GroupID,
                ObjectID: r.ObjectID,
                Data:     data,
        }, nil
}

// ParseSAPRecord decodes one record of a SAP Type timeline track,
// enforcing CMSF §3.6.1: the record MUST be indexed by Location and
// its data field MUST be two integers whose first is a valid SAP type.
func ParseSAPRecord(rec EventRecord) (SAPRecord, error) {
        if rec.Index != EventIndexLocation {
                return SAPRecord{}, fmt.Errorf(
                        "moqt/msf: sap record: index must be 'l' for Location, got %d (CMSF §3.6.1)", rec.Index)
        }
        var pair []int64
        if err := json.Unmarshal(rec.Data, &pair); err != nil {
                return SAPRecord{}, fmt.Errorf("moqt/msf: sap record data: %w", err)
        }
        if len(pair) != 2 {
                return SAPRecord{}, fmt.Errorf(
                        "moqt/msf: sap record data: expected 2 items, got %d (CMSF §3.6.1)", len(pair))
        }
        out := SAPRecord{
                GroupID:  rec.GroupID,
                ObjectID: rec.ObjectID,
                SAPType:  int(pair[0]),
                EPT:      pair[1],
        }
        if err := out.validate(); err != nil {
                return SAPRecord{}, err
        }
        return out, nil
}

package msf

import (
        "errors"
        "fmt"
        "maps"
        "reflect"
        "slices"
)

// Apply replays delta against base and returns the resulting catalog
// per §5.3. Apply does not mutate base or delta.
//
// Operations are processed in the order they appear in
// delta.DeltaUpdate; within each operation, its Tracks are applied in
// order. This matches the document order §5.3 requires.
//
// Errors:
//   - ErrNotDelta if delta is not a delta update.
//   - A descriptive error if any operation violates §5.3 (e.g.
//     adding a track whose Namespace+Name already exists, cloning
//     from a missing parent).
func Apply(base, delta Catalog) (Catalog, error) {
        if !delta.IsDelta() {
                return Catalog{}, ErrNotDelta
        }

        out := cloneCatalog(base)

        // §5.3 restricts deltaUpdate to track operations and forbids only
        // the tracks and version fields at the root, so a delta MAY carry
        // the root-level arrays a newly added track references. It has to:
        // CMSF §3.1 requires every CMAF track to name an initDataList entry
        // through initRef, and CMSF §4.1.2 requires a protected track's
        // contentProtectionRefIDs to resolve. Merge them before replaying
        // the operations so an added track can reference them.
        if err := mergeInitDataList(&out, delta.InitDataList); err != nil {
                return Catalog{}, err
        }
        if err := mergeContentProtections(&out, delta.ContentProtections); err != nil {
                return Catalog{}, err
        }

        for i, op := range delta.DeltaUpdate {
                switch op.Op {
                case DeltaOpAdd:
                        for j, tr := range op.Tracks {
                                if err := applyAdd(&out, tr); err != nil {
                                        return Catalog{}, fmt.Errorf("moqt/msf: deltaUpdate[%d].tracks[%d]: %w", i, j, err)
                                }
                        }
                case DeltaOpRemove:
                        for j, tr := range op.Tracks {
                                if err := applyRemove(&out, tr); err != nil {
                                        return Catalog{}, fmt.Errorf("moqt/msf: deltaUpdate[%d].tracks[%d]: %w", i, j, err)
                                }
                        }
                case DeltaOpClone:
                        for j, tr := range op.Tracks {
                                if err := applyClone(&out, tr); err != nil {
                                        return Catalog{}, fmt.Errorf("moqt/msf: deltaUpdate[%d].tracks[%d]: %w", i, j, err)
                                }
                        }
                default:
                        return Catalog{}, fmt.Errorf("moqt/msf: deltaUpdate[%d]: unknown op %q", i, op.Op)
                }
        }

        out.DeltaUpdate = nil
        if delta.GeneratedAt != 0 {
                out.GeneratedAt = delta.GeneratedAt
        }

        // Whether a track's initRef and contentProtectionRefIDs resolve is
        // cross-document state: the entries may come from the base, from
        // this delta, or from an earlier one. [Catalog.Validate] cannot see
        // that, so checking it is Apply's job.
        if err := validateTrackReferences(&out); err != nil {
                return Catalog{}, err
        }
        return out, nil
}

// ErrNotDelta is returned by [Apply] when the delta argument is not a
// delta update (deltaUpdate absent).
var ErrNotDelta = errors.New("moqt/msf: catalog is not a delta update")

// mergeInitDataList folds a delta's initDataList entries (§5.1.7) into
// out. Re-sending an identical entry is a no-op; redefining an existing
// id with different content is rejected, because §5.1.7 requires the id
// to be unique within the scope of the catalog.
func mergeInitDataList(out *Catalog, add []InitData) error {
        for _, entry := range add {
                if entry.ID == "" {
                        return errors.New("moqt/msf: delta initDataList: id is required (§5.1.7)")
                }
                i := slices.IndexFunc(out.InitDataList, func(e InitData) bool { return e.ID == entry.ID })
                if i < 0 {
                        out.InitDataList = append(out.InitDataList, entry)
                        continue
                }
                if out.InitDataList[i] != entry {
                        return fmt.Errorf("moqt/msf: delta initDataList: id %q redefined (§5.1.7)", entry.ID)
                }
        }
        return nil
}

// mergeContentProtections folds a delta's contentProtections entries
// (CMSF §4.1.1) into out under the same rule as [mergeInitDataList]:
// identical re-sends are idempotent, conflicting redefinitions of a
// refID are rejected.
func mergeContentProtections(out *Catalog, add []ContentProtection) error {
        for _, entry := range cloneContentProtections(add) {
                if entry.RefID == "" {
                        return errors.New("moqt/msf: delta contentProtections: refID is required (CMSF §4.1.1.1)")
                }
                i := slices.IndexFunc(out.ContentProtections, func(e ContentProtection) bool {
                        return e.RefID == entry.RefID
                })
                if i < 0 {
                        out.ContentProtections = append(out.ContentProtections, entry)
                        continue
                }
                if !reflect.DeepEqual(out.ContentProtections[i], entry) {
                        return fmt.Errorf(
                                "moqt/msf: delta contentProtections: refID %q redefined (CMSF §4.1.1.1)", entry.RefID)
                }
        }
        return nil
}

// applyAdd processes one add-operation track. §5.3 — adding a track
// whose (Namespace, Name) already exists is rejected; the registry has
// a fixed-attribute invariant per §5.3.
func applyAdd(out *Catalog, add Track) error {
        if add.Name == "" {
                return errors.New("add entry missing name")
        }
        for _, existing := range out.Tracks {
                if sameTrackID(existing, add) {
                        return fmt.Errorf("track %q (ns=%q) already exists", add.Name, add.Namespace)
                }
        }
        out.Tracks = append(out.Tracks, add)
        return nil
}

// applyRemove drops the named track from out.Tracks. §5.1.6 — only
// Name is required, Namespace is optional. The match is exact when
// Namespace is provided, else by Name alone.
func applyRemove(out *Catalog, rm Track) error {
        if rm.Name == "" {
                return errors.New("remove entry missing name")
        }
        for i, existing := range out.Tracks {
                if rm.Namespace != "" && existing.Namespace != rm.Namespace {
                        continue
                }
                if existing.Name != rm.Name {
                        continue
                }
                out.Tracks = append(out.Tracks[:i], out.Tracks[i+1:]...)
                return nil
        }
        // §5.1.6 doesn't explicitly require erroring on a missing remove,
        // but rejecting it surfaces producer mistakes early.
        return fmt.Errorf("no such track %q (ns=%q)", rm.Name, rm.Namespace)
}

// applyClone creates a new track that inherits the attributes of its
// parent (looked up by ParentName and optional ParentNamespace) and
// overrides any explicitly-set fields on the clone entry. §5.3 — the
// clone MUST have a different Track Name.
func applyClone(out *Catalog, clone Track) error {
        if clone.ParentName == "" {
                return errors.New("clone entry: parentName required")
        }
        if clone.Name == "" {
                return errors.New("clone entry: name required")
        }

        parent, ok := findTrack(out, clone.ParentName, clone.ParentNamespace)
        if !ok {
                return fmt.Errorf("clone entry: parent %q not found", clone.ParentName)
        }

        // Start from a deep copy of the parent then overlay non-zero fields
        // from the clone definition. ParentName/ParentNamespace are consumed
        // and not carried onto the resulting track. The copy has to be deep:
        // a clone entry that omits a slice field inherits the parent's, and
        // the two tracks must not share its backing array.
        merged := cloneTrack(parent)
        overlayTrack(&merged, clone)
        merged.Name = clone.Name
        merged.ParentName = ""
        merged.ParentNamespace = ""

        if sameTrackID(parent, merged) {
                return fmt.Errorf("clone entry: clone name %q matches parent", merged.Name)
        }
        for _, existing := range out.Tracks {
                if sameTrackID(existing, merged) {
                        return fmt.Errorf("clone entry: resulting track %q (ns=%q) already exists",
                                merged.Name, merged.Namespace)
                }
        }
        out.Tracks = append(out.Tracks, merged)
        return nil
}

// overlayTrack copies set fields from src onto dst. Pointer fields and
// slices/maps are taken if non-nil; scalars are taken if non-zero. The
// rule is "if the producer set it on the clone, prefer the clone".
// Split across three helpers to keep each within gocyclo's bound.
func overlayTrack(dst *Track, src Track) {
        overlayTrackStrings(dst, src)
        overlayTrackNumbers(dst, src)
        overlayTrackComposite(dst, src)
}

func overlayTrackStrings(dst *Track, src Track) {
        if src.Namespace != "" {
                dst.Namespace = src.Namespace
        }
        if src.Packaging != "" {
                dst.Packaging = src.Packaging
        }
        if src.EventType != "" {
                dst.EventType = src.EventType
        }
        if src.Role != "" {
                dst.Role = src.Role
        }
        if src.Label != "" {
                dst.Label = src.Label
        }
        if src.InitRef != "" {
                dst.InitRef = src.InitRef
        }
        if src.Codec != "" {
                dst.Codec = src.Codec
        }
        if src.Mimetype != "" {
                dst.Mimetype = src.Mimetype
        }
        if src.ChannelConfig != "" {
                dst.ChannelConfig = src.ChannelConfig
        }
        if src.Lang != "" {
                dst.Lang = src.Lang
        }
        if src.ConnectionURI != "" {
                dst.ConnectionURI = src.ConnectionURI
        }
        if src.Token != "" {
                dst.Token = src.Token
        }
        if src.EncryptionScheme != "" {
                dst.EncryptionScheme = src.EncryptionScheme
        }
        if src.CipherSuite != "" {
                dst.CipherSuite = src.CipherSuite
        }
        if src.KeyID != "" {
                dst.KeyID = src.KeyID
        }
        if src.TrackBaseKey != "" {
                dst.TrackBaseKey = src.TrackBaseKey
        }
}

func overlayTrackNumbers(dst *Track, src Track) {
        if src.Framerate != 0 {
                dst.Framerate = src.Framerate
        }
        if src.Timescale != 0 {
                dst.Timescale = src.Timescale
        }
        if src.Bitrate != 0 {
                dst.Bitrate = src.Bitrate
        }
        if src.AvgBitrate != 0 {
                dst.AvgBitrate = src.AvgBitrate
        }
        if src.MaxGopDuration != 0 {
                dst.MaxGopDuration = src.MaxGopDuration
        }
        if src.MaxGroupDuration != 0 {
                dst.MaxGroupDuration = src.MaxGroupDuration
        }
        if src.Width != 0 {
                dst.Width = src.Width
        }
        if src.Height != 0 {
                dst.Height = src.Height
        }
        if src.Samplerate != 0 {
                dst.Samplerate = src.Samplerate
        }
        if src.DisplayWidth != 0 {
                dst.DisplayWidth = src.DisplayWidth
        }
        if src.DisplayHeight != 0 {
                dst.DisplayHeight = src.DisplayHeight
        }
        if src.TrackDuration != 0 {
                dst.TrackDuration = src.TrackDuration
        }
}

func overlayTrackComposite(dst *Track, src Track) {
        if src.IsLive != nil {
                v := *src.IsLive
                dst.IsLive = &v
        }
        if src.TargetLatency != nil {
                v := *src.TargetLatency
                dst.TargetLatency = &v
        }
        if src.Buffers != nil {
                b := *src.Buffers
                dst.Buffers = &b
        }
        if src.RenderGroup != nil {
                v := *src.RenderGroup
                dst.RenderGroup = &v
        }
        if src.AltGroup != nil {
                v := *src.AltGroup
                dst.AltGroup = &v
        }
        if src.Depends != nil {
                dst.Depends = slices.Clone(src.Depends)
        }
        if src.Template != nil {
                v := *src.Template
                dst.Template = &v
        }
        if src.TemporalID != nil {
                v := *src.TemporalID
                dst.TemporalID = &v
        }
        if src.SpatialID != nil {
                v := *src.SpatialID
                dst.SpatialID = &v
        }
        if src.MaxGrpSapStartingType != nil {
                v := *src.MaxGrpSapStartingType
                dst.MaxGrpSapStartingType = &v
        }
        if src.MaxObjSapStartingType != nil {
                v := *src.MaxObjSapStartingType
                dst.MaxObjSapStartingType = &v
        }
        if src.ContentProtectionRefIDs != nil {
                dst.ContentProtectionRefIDs = slices.Clone(src.ContentProtectionRefIDs)
        }
        if src.AuthInfo != nil {
                dst.AuthInfo = cloneExtras(src.AuthInfo)
        }
        if src.Accessibility != nil {
                dst.Accessibility = slices.Clone(src.Accessibility)
        }
        if src.Extras != nil {
                dst.Extras = cloneExtras(src.Extras)
        }
}

func sameTrackID(a, b Track) bool {
        return a.Name == b.Name && a.Namespace == b.Namespace
}

func findTrack(c *Catalog, name, namespace string) (Track, bool) {
        for _, t := range c.Tracks {
                if t.Name != name {
                        continue
                }
                if namespace != "" && t.Namespace != namespace {
                        continue
                }
                return t, true
        }
        return Track{}, false
}

func cloneCatalog(c Catalog) Catalog {
        out := c
        out.Tracks = cloneTracks(c.Tracks)
        out.PublishTracks = cloneTracks(c.PublishTracks)
        out.InitDataList = slices.Clone(c.InitDataList)
        out.ContentProtections = cloneContentProtections(c.ContentProtections)
        out.Extras = cloneExtras(c.Extras)
        out.DeltaUpdate = nil
        return out
}

func cloneTracks(in []Track) []Track {
        if in == nil {
                return nil
        }
        out := slices.Clone(in)
        for i := range out {
                out[i] = cloneTrack(out[i])
        }
        return out
}

// cloneTrack deep-copies a track's maps and slices so the copy shares
// no backing storage with the original.
func cloneTrack(in Track) Track {
        out := in
        out.Extras = cloneExtras(in.Extras)
        out.AuthInfo = cloneExtras(in.AuthInfo)
        out.Depends = slices.Clone(in.Depends)
        out.Accessibility = slices.Clone(in.Accessibility)
        out.ContentProtectionRefIDs = slices.Clone(in.ContentProtectionRefIDs)
        return out
}

// cloneContentProtections deep-copies a catalog's contentProtections
// array (CMSF §4.1.1) for [cloneCatalog].
func cloneContentProtections(in []ContentProtection) []ContentProtection {
        if in == nil {
                return nil
        }
        out := slices.Clone(in)
        for i := range out {
                out[i].DefaultKID = slices.Clone(out[i].DefaultKID)
                out[i].Extras = cloneExtras(out[i].Extras)
                ds := &out[i].DRMSystem
                ds.LAURL = cloneURLRef(ds.LAURL)
                ds.CertURL = cloneURLRef(ds.CertURL)
                ds.Extras = cloneExtras(ds.Extras)
        }
        return out
}

// cloneURLRef copies a [DRMSystem] URL object so the clone does not
// alias the original.
func cloneURLRef(in *URLRef) *URLRef {
        if in == nil {
                return nil
        }
        return new(*in)
}

func cloneExtras(in map[string]any) map[string]any {
        if in == nil {
                return nil
        }
        out := make(map[string]any, len(in))
        maps.Copy(out, in)
        return out
}

package msf

import (
        "encoding/json"
        "fmt"
)

// EventIndex identifies which time/location field anchors an Event
// Timeline record per §8.1. Exactly one of the three index fields
// ('t', 'l', 'm') MUST be present in each record.
type EventIndex uint8

const (
        // EventIndexWallclock means the record is anchored by 't': a
        // wallclock time in milliseconds since the Unix epoch.
        EventIndexWallclock EventIndex = iota + 1
        // EventIndexLocation means the record is anchored by 'l': a
        // [Group ID, Object ID] tuple.
        EventIndexLocation
        // EventIndexMediaPTS means the record is anchored by 'm': a
        // media PTS value in milliseconds.
        EventIndexMediaPTS
)

// EventRecord is one entry in an Event Timeline track (§8.1).
//
// Only the fields relevant to Index are read on encode and populated
// on decode. Time carries the value for EventIndexWallclock and
// EventIndexMediaPTS; GroupID + ObjectID carry the location for
// EventIndexLocation.
//
// Data is the opaque application-defined payload whose schema is
// declared by the catalog's EventType field for this track (§5.2.5).
type EventRecord struct {
        Index    EventIndex
        Time     int64
        GroupID  uint64
        ObjectID uint64
        Data     json.RawMessage
}

// EventTimeline is the array of records produced by an Event Timeline
// track (§8.1).
type EventTimeline []EventRecord

// MarshalJSON encodes the timeline per §8.1.
func (e EventTimeline) MarshalJSON() ([]byte, error) {
        out := make([]map[string]json.RawMessage, len(e))
        for i, r := range e {
                entry := map[string]json.RawMessage{}
                switch r.Index {
                case EventIndexWallclock:
                        b, err := json.Marshal(r.Time)
                        if err != nil {
                                return nil, err
                        }
                        entry["t"] = b
                case EventIndexLocation:
                        b, err := json.Marshal([2]uint64{r.GroupID, r.ObjectID})
                        if err != nil {
                                return nil, err
                        }
                        entry["l"] = b
                case EventIndexMediaPTS:
                        b, err := json.Marshal(r.Time)
                        if err != nil {
                                return nil, err
                        }
                        entry["m"] = b
                default:
                        return nil, fmt.Errorf("moqt/msf: event record %d: unknown Index %d", i, r.Index)
                }
                if r.Data == nil {
                        entry["data"] = json.RawMessage("null")
                } else {
                        entry["data"] = r.Data
                }
                out[i] = entry
        }
        return json.Marshal(out)
}

// UnmarshalJSON parses an Event Timeline document. Each record MUST
// have exactly one of t/l/m and a data field.
func (e *EventTimeline) UnmarshalJSON(data []byte) error {
        var raw []map[string]json.RawMessage
        if err := json.Unmarshal(data, &raw); err != nil {
                return fmt.Errorf("moqt/msf: event timeline: %w", err)
        }
        out := make(EventTimeline, 0, len(raw))
        for i, entry := range raw {
                rec := EventRecord{}
                nIndexes := 0
                if t, ok := entry["t"]; ok {
                        if err := json.Unmarshal(t, &rec.Time); err != nil {
                                return fmt.Errorf("moqt/msf: event record %d t: %w", i, err)
                        }
                        rec.Index = EventIndexWallclock
                        nIndexes++
                }
                if l, ok := entry["l"]; ok {
                        var loc []uint64
                        if err := json.Unmarshal(l, &loc); err != nil {
                                return fmt.Errorf("moqt/msf: event record %d l: %w", i, err)
                        }
                        if len(loc) != 2 {
                                return fmt.Errorf(
                                        "moqt/msf: event record %d: l must have 2 items, got %d", i, len(loc))
                        }
                        rec.GroupID = loc[0]
                        rec.ObjectID = loc[1]
                        rec.Index = EventIndexLocation
                        nIndexes++
                }
                if m, ok := entry["m"]; ok {
                        if err := json.Unmarshal(m, &rec.Time); err != nil {
                                return fmt.Errorf("moqt/msf: event record %d m: %w", i, err)
                        }
                        rec.Index = EventIndexMediaPTS
                        nIndexes++
                }
                if nIndexes != 1 {
                        return fmt.Errorf(
                                "moqt/msf: event record %d: must have exactly one of t/l/m, got %d", i, nIndexes)
                }
                if d, ok := entry["data"]; ok {
                        rec.Data = append(json.RawMessage(nil), d...)
                }
                out = append(out, rec)
        }
        *e = out
        return nil
}

package msf

import (
        "errors"
        "fmt"
        "sync/atomic"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// GroupSequencer issues monotonically increasing MOQT Group IDs for a
// single track per §6.1. The initial value is the current Unix
// millisecond, which makes Group IDs across application restarts
// non-decreasing and avoids collisions as long as a publisher emits
// fewer than 1000 groups per second.
//
// GroupSequencer is safe for concurrent use.
type GroupSequencer struct {
        next atomic.Uint64
}

// NewGroupSequencer returns a sequencer seeded with the current
// wallclock as Unix milliseconds. The first call to [Next] returns
// that seed and increments internal state.
func NewGroupSequencer() *GroupSequencer {
        s := &GroupSequencer{}
        s.next.Store(uint64(time.Now().UnixMilli()))
        return s
}

// NewGroupSequencerAt returns a sequencer seeded at the given start ID.
// Useful for tests and for callers that maintain their own time source.
func NewGroupSequencerAt(start uint64) *GroupSequencer {
        s := &GroupSequencer{}
        s.next.Store(start)
        return s
}

// Next returns the next Group ID and advances the sequencer.
func (s *GroupSequencer) Next() uint64 {
        // atomic.Uint64.Add returns the new value, so to mirror the
        // "return current, then increment" semantics we subtract 1.
        return s.next.Add(1) - 1
}

// Peek returns the value that the next call to Next would produce
// without advancing the sequencer.
func (s *GroupSequencer) Peek() uint64 {
        return s.next.Load()
}

// PriorGapHeader returns the KV pair a publisher attaches to its first
// Object after a republish so subscribers can distinguish an
// intentional Group ID gap (e.g. encoder restart) from missing data.
// See §6.1 of the MSF draft and PRIOR_GROUP_ID_GAP in §12.8 of
// MoQ Transport.
//
// prev is the last Group ID the publisher emitted before the gap;
// curr is the first Group ID after the gap. Returns an error if
// curr <= prev (no gap) or if prev+1 == curr (no gap, just the next
// sequential ID).
func PriorGapHeader(prev, curr uint64) (wire.KVPair, error) {
        if curr <= prev {
                return wire.KVPair{}, fmt.Errorf(
                        "moqt/msf: PriorGapHeader: curr (%d) must be > prev (%d)", curr, prev)
        }
        gap := curr - prev - 1
        if gap == 0 {
                return wire.KVPair{}, errors.New("moqt/msf: PriorGapHeader: curr is the immediate successor of prev (no gap)")
        }
        return wire.KVPair{
                Type:   message.PropertyPriorGroupIDGap,
                IntVal: gap,
        }, nil
}

package msf

import (
        "encoding/json"
        "fmt"
)

// MediaTimelineRecord is one entry in a Media Timeline track (§7.1).
// On the wire each record is a JSON array of three items:
//
//        [ MediaPTS, [GroupID, ObjectID], Wallclock ]
//
// MediaPTS is the media presentation timestamp, rounded to the nearest
// millisecond, of the first media sample in the referenced Object.
// Wallclock is the time of encoding in milliseconds since the Unix
// epoch; for VOD or unknown wallclocks it is 0 (§7.1).
type MediaTimelineRecord struct {
        MediaPTS  int64
        GroupID   uint64
        ObjectID  uint64
        Wallclock int64
}

// MediaTimeline is the array of records produced by a Media Timeline
// track. Independent Objects MUST carry the full history since the
// start of the track (§7.3); incremental updates MAY carry only the
// records since the last Object in the same Group.
type MediaTimeline []MediaTimelineRecord

// MarshalJSON encodes the timeline per §7.1.
func (m MediaTimeline) MarshalJSON() ([]byte, error) {
        if m == nil {
                return []byte("[]"), nil
        }
        out := make([][3]json.RawMessage, len(m))
        for i, r := range m {
                pts, err := json.Marshal(r.MediaPTS)
                if err != nil {
                        return nil, err
                }
                loc, err := json.Marshal([2]uint64{r.GroupID, r.ObjectID})
                if err != nil {
                        return nil, err
                }
                wc, err := json.Marshal(r.Wallclock)
                if err != nil {
                        return nil, err
                }
                out[i] = [3]json.RawMessage{pts, loc, wc}
        }
        return json.Marshal(out)
}

// UnmarshalJSON parses a Media Timeline document.
func (m *MediaTimeline) UnmarshalJSON(data []byte) error {
        var raw []json.RawMessage
        if err := json.Unmarshal(data, &raw); err != nil {
                return fmt.Errorf("moqt/msf: media timeline: %w", err)
        }
        out := make(MediaTimeline, 0, len(raw))
        for i, rec := range raw {
                var triple []json.RawMessage
                if err := json.Unmarshal(rec, &triple); err != nil {
                        return fmt.Errorf("moqt/msf: media timeline record %d: %w", i, err)
                }
                if len(triple) != 3 {
                        return fmt.Errorf(
                                "moqt/msf: media timeline record %d: expected 3 items, got %d", i, len(triple))
                }
                var (
                        pts int64
                        loc []uint64
                        wc  int64
                )
                if err := json.Unmarshal(triple[0], &pts); err != nil {
                        return fmt.Errorf("moqt/msf: media timeline record %d pts: %w", i, err)
                }
                if err := json.Unmarshal(triple[1], &loc); err != nil {
                        return fmt.Errorf("moqt/msf: media timeline record %d location: %w", i, err)
                }
                if len(loc) != 2 {
                        return fmt.Errorf(
                                "moqt/msf: media timeline record %d: location must have 2 items, got %d", i, len(loc))
                }
                if err := json.Unmarshal(triple[2], &wc); err != nil {
                        return fmt.Errorf("moqt/msf: media timeline record %d wallclock: %w", i, err)
                }
                out = append(out, MediaTimelineRecord{
                        MediaPTS:  pts,
                        GroupID:   loc[0],
                        ObjectID:  loc[1],
                        Wallclock: wc,
                })
        }
        *m = out
        return nil
}

// Since returns the records in m whose MediaPTS is strictly greater
// than afterPTS. This produces the "records since the last media
// timeline Object" body for the second-and-later Objects in a Group
// per §7.3.
func (m MediaTimeline) Since(afterPTS int64) MediaTimeline {
        for i, r := range m {
                if r.MediaPTS > afterPTS {
                        return m[i:]
                }
        }
        return nil
}

// MediaTimelineTemplate is the inline media timeline template carried
// by the Track.Template field (§5.2.15, §7.4). It describes a regular,
// predictable relationship between media time, MOQT Location and
// wallclock time for fixed-duration segments, replacing an explicit
// media timeline track.
//
// On the wire it is a JSON array of six mandatory values in fixed order
// (§7.4.1):
//
//        [ startMediaTime, deltaMediaTime,
//          [startGroupID, startObjectID], [deltaGroupID, deltaObjectID],
//          startWallclock, deltaWallclock ]
type MediaTimelineTemplate struct {
        StartMediaTime int64
        DeltaMediaTime int64
        StartGroupID   uint64
        StartObjectID  uint64
        DeltaGroupID   int64
        DeltaObjectID  int64
        StartWallclock int64
        DeltaWallclock int64
}

// At computes the media timeline entry for the zero-based index n using
// the formulas in §7.4.1.
func (t MediaTimelineTemplate) At(n int64) MediaTimelineRecord {
        return MediaTimelineRecord{
                MediaPTS:  t.StartMediaTime + n*t.DeltaMediaTime,
                GroupID:   addDelta(t.StartGroupID, n*t.DeltaGroupID),
                ObjectID:  addDelta(t.StartObjectID, n*t.DeltaObjectID),
                Wallclock: t.StartWallclock + n*t.DeltaWallclock,
        }
}

// addDelta adds a signed delta to a MOQT Group/Object ID. Both operands
// are bounded by the MOQT wire format (62-bit varints), so the
// round-trip through int64 cannot overflow in practice.
//
//nolint:gosec // G115: Group/Object IDs and deltas are bounded by the 62-bit MOQT varint range.
func addDelta(base uint64, delta int64) uint64 {
        return uint64(int64(base) + delta)
}

// MarshalJSON encodes the template as the six-element array of §7.4.1.
func (t MediaTimelineTemplate) MarshalJSON() ([]byte, error) {
        //nolint:gosec // G115: Group/Object IDs are bounded by the 62-bit MOQT varint range.
        loc := [2]int64{int64(t.StartGroupID), int64(t.StartObjectID)}
        delta := [2]int64{t.DeltaGroupID, t.DeltaObjectID}
        out := []any{t.StartMediaTime, t.DeltaMediaTime, loc, delta, t.StartWallclock, t.DeltaWallclock}
        return json.Marshal(out)
}

// UnmarshalJSON parses the six-element template array of §7.4.1.
func (t *MediaTimelineTemplate) UnmarshalJSON(data []byte) error {
        var raw []json.RawMessage
        if err := json.Unmarshal(data, &raw); err != nil {
                return fmt.Errorf("moqt/msf: media timeline template: %w", err)
        }
        if len(raw) != 6 {
                return fmt.Errorf("moqt/msf: media timeline template: expected 6 items, got %d", len(raw))
        }
        var (
                startMedia, deltaMedia    int64
                startLoc, deltaLoc        []int64
                startWallclock, deltaWall int64
        )
        if err := json.Unmarshal(raw[0], &startMedia); err != nil {
                return fmt.Errorf("moqt/msf: media timeline template startMediaTime: %w", err)
        }
        if err := json.Unmarshal(raw[1], &deltaMedia); err != nil {
                return fmt.Errorf("moqt/msf: media timeline template deltaMediaTime: %w", err)
        }
        if err := json.Unmarshal(raw[2], &startLoc); err != nil {
                return fmt.Errorf("moqt/msf: media timeline template startLocation: %w", err)
        }
        if len(startLoc) != 2 {
                return fmt.Errorf("moqt/msf: media timeline template startLocation: expected 2 items, got %d", len(startLoc))
        }
        if err := json.Unmarshal(raw[3], &deltaLoc); err != nil {
                return fmt.Errorf("moqt/msf: media timeline template deltaLocation: %w", err)
        }
        if len(deltaLoc) != 2 {
                return fmt.Errorf("moqt/msf: media timeline template deltaLocation: expected 2 items, got %d", len(deltaLoc))
        }
        if err := json.Unmarshal(raw[4], &startWallclock); err != nil {
                return fmt.Errorf("moqt/msf: media timeline template startWallclock: %w", err)
        }
        if err := json.Unmarshal(raw[5], &deltaWall); err != nil {
                return fmt.Errorf("moqt/msf: media timeline template deltaWallclock: %w", err)
        }
        *t = MediaTimelineTemplate{
                StartMediaTime: startMedia,
                DeltaMediaTime: deltaMedia,
                //nolint:gosec // G115: Group/Object IDs are bounded by the 62-bit MOQT varint range.
                StartGroupID: uint64(startLoc[0]),
                //nolint:gosec // G115: Group/Object IDs are bounded by the 62-bit MOQT varint range.
                StartObjectID:  uint64(startLoc[1]),
                DeltaGroupID:   deltaLoc[0],
                DeltaObjectID:  deltaLoc[1],
                StartWallclock: startWallclock,
                DeltaWallclock: deltaWall,
        }
        return nil
}

package msf

import (
        "encoding/hex"
        "errors"
        "fmt"
        "reflect"
)

// Validate enforces the rules from §5.1 and §5.2 that can be checked
// from a Catalog value alone. It does NOT validate cross-document
// state (e.g. whether a delta's parentName exists in some prior
// catalog) — that is [Apply]'s responsibility.
//
// Validate returns nil for valid catalogs and a descriptive error for
// the first violation it encounters.
func (c *Catalog) Validate() error {
        if c.IsDelta() {
                return c.validateDelta()
        }
        return c.validateIndependent()
}

func (c *Catalog) validateIndependent() error {
        if c.Version != Version {
                // Per §5.1.1: subscriber MUST NOT parse unknown versions. We
                // still emit an error so producers learn about the mismatch.
                if c.Version == "" {
                        return errors.New("moqt/msf: version is required (§5.1.1)")
                }
                return fmt.Errorf("moqt/msf: unsupported version %q (expected %q)", c.Version, Version)
        }

        // §5.1.2 — generatedAt SHOULD NOT be included when isLive is false
        // for every track. We treat the SHOULD as advisory and skip it.

        // Per-track cross-field constraints.
        for i, tr := range c.Tracks {
                if err := validateTrack(tr); err != nil {
                        return fmt.Errorf("moqt/msf: tracks[%d]: %w", i, err)
                }
        }

        if err := validateTargetLatencyGroups(
                c.Tracks,
                "renderGroup",
                func(t Track) *int { return t.RenderGroup },
        ); err != nil {
                return err
        }
        if err := validateTargetLatencyGroups(c.Tracks, "altGroup", func(t Track) *int { return t.AltGroup }); err != nil {
                return err
        }
        if err := validateContentProtections(c); err != nil {
                return err
        }
        if err := validateTrackReferences(c); err != nil {
                return err
        }
        return nil
}

// validateContentProtections enforces draft-ietf-moq-cmsf-01 §4.1.1's
// per-entry required fields and refID uniqueness.
func validateContentProtections(c *Catalog) error {
        seen := make(map[string]struct{}, len(c.ContentProtections))
        for i, cp := range c.ContentProtections {
                if cp.RefID == "" {
                        return fmt.Errorf("moqt/msf: contentProtections[%d]: refID is required (CMSF §4.1.1.1)", i)
                }
                if _, dup := seen[cp.RefID]; dup {
                        return fmt.Errorf("moqt/msf: contentProtections[%d]: duplicate refID %q (CMSF §4.1.1.1)", i, cp.RefID)
                }
                seen[cp.RefID] = struct{}{}
                if err := validateContentProtection(cp); err != nil {
                        return fmt.Errorf("moqt/msf: contentProtections[%d]: %w", i, err)
                }
        }
        return nil
}

// validateContentProtection checks one entry's required fields, the
// §4.1.1.3 scheme enumeration and the UUID forms §4.1.1.2 / §4.1.1.4.1
// require.
func validateContentProtection(cp ContentProtection) error {
        if len(cp.DefaultKID) == 0 {
                return errors.New("defaultKID is required (CMSF §4.1.1.2)")
        }
        for j, kid := range cp.DefaultKID {
                if !isUUID(kid) {
                        return fmt.Errorf("defaultKID[%d] %q is not a UUID string (CMSF §4.1.1.2)", j, kid)
                }
        }
        switch cp.Scheme {
        case "":
                return errors.New("scheme is required (CMSF §4.1.1.3)")
        case SchemeCENC, SchemeCBCS:
        default:
                return fmt.Errorf("unknown scheme %q (CMSF §4.1.1.3, Table 3)", cp.Scheme)
        }
        return validateDRMSystem(cp.DRMSystem)
}

// validateDRMSystem checks the §4.1.1.4 DRM System object: the
// systemID UUID (§4.1.1.4.1) and the required url of every URL object
// that is present (§4.1.1.4.2, §4.1.1.4.3).
func validateDRMSystem(ds DRMSystem) error {
        if ds.SystemID == "" {
                return errors.New("drmSystem.systemID is required (CMSF §4.1.1.4.1)")
        }
        if !isUUID(ds.SystemID) {
                return fmt.Errorf("drmSystem.systemID %q is not a UUID string (CMSF §4.1.1.4.1)", ds.SystemID)
        }
        urls := []struct {
                name string
                ref  *URLRef
                sec  string
        }{
                {"laURL", ds.LAURL, "§4.1.1.4.2"},
                {"certURL", ds.CertURL, "§4.1.1.4.3"},
        }
        for _, u := range urls {
                if u.ref != nil && u.ref.URL == "" {
                        return fmt.Errorf("drmSystem.%s: url is required (CMSF %s)", u.name, u.sec)
                }
        }
        return nil
}

// validateTrackReferences enforces the catalog's two referential
// integrity rules: Track.InitRef names an initDataList entry (§5.2.13,
// CMSF §3.1) and every Track.ContentProtectionRefIDs entry names a
// contentProtections entry (CMSF §4.1.2).
func validateTrackReferences(c *Catalog) error {
        initIDs := make(map[string]struct{}, len(c.InitDataList))
        for _, init := range c.InitDataList {
                initIDs[init.ID] = struct{}{}
        }
        refIDs := make(map[string]struct{}, len(c.ContentProtections))
        for _, cp := range c.ContentProtections {
                refIDs[cp.RefID] = struct{}{}
        }
        for i, tr := range c.Tracks {
                if tr.InitRef != "" {
                        if _, ok := initIDs[tr.InitRef]; !ok {
                                return fmt.Errorf(
                                        "moqt/msf: tracks[%d]: initRef %q has no initDataList entry (§5.2.13)", i, tr.InitRef)
                        }
                }
                for _, ref := range tr.ContentProtectionRefIDs {
                        if _, ok := refIDs[ref]; !ok {
                                return fmt.Errorf(
                                        "moqt/msf: tracks[%d]: contentProtectionRefIDs references unknown refID %q (CMSF §4.1.1, §4.1.2)",
                                        i, ref)
                        }
                }
        }
        return nil
}

// isUUID reports whether s has the
// "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx" form that CMSF §4.1.1.2 and
// §4.1.1.4.1 require of key and DRM system identifiers.
func isUUID(s string) bool {
        if len(s) != 36 || s[8] != '-' || s[13] != '-' || s[18] != '-' || s[23] != '-' {
                return false
        }
        _, err := hex.DecodeString(s[:8] + s[9:13] + s[14:18] + s[19:23] + s[24:])
        return err == nil
}

func (c *Catalog) validateDelta() error {
        if c.Version != "" {
                return errors.New("moqt/msf: delta update must not contain version (§5.3)")
        }
        if c.Tracks != nil {
                return errors.New("moqt/msf: delta update must not contain tracks (§5.3)")
        }
        if len(c.DeltaUpdate) == 0 {
                return errors.New("moqt/msf: delta update must contain at least one operation (§5.3)")
        }
        // A delta MAY carry the root-level arrays its added tracks
        // reference; [Apply] merges them. Their own field rules still hold.
        if err := validateContentProtections(c); err != nil {
                return err
        }
        for i, op := range c.DeltaUpdate {
                if err := validateDeltaOp(op); err != nil {
                        return fmt.Errorf("moqt/msf: deltaUpdate[%d]: %w", i, err)
                }
        }
        return nil
}

func validateDeltaOp(op DeltaOp) error {
        switch op.Op {
        case DeltaOpAdd:
                for i, tr := range op.Tracks {
                        if err := validateTrack(tr); err != nil {
                                return fmt.Errorf("tracks[%d]: %w", i, err)
                        }
                }
        case DeltaOpRemove:
                for i, tr := range op.Tracks {
                        if tr.Name == "" {
                                return fmt.Errorf("tracks[%d]: name required (§5.1.6)", i)
                        }
                        if !isOnlyNameAndNamespace(tr) {
                                return fmt.Errorf("tracks[%d]: only name and namespace allowed (§5.1.6)", i)
                        }
                }
        case DeltaOpClone:
                for i, tr := range op.Tracks {
                        if tr.ParentName == "" {
                                return fmt.Errorf("tracks[%d]: parentName required (§5.1.6)", i)
                        }
                        if tr.Name == "" {
                                return fmt.Errorf("tracks[%d]: name required (§5.3)", i)
                        }
                }
        default:
                return fmt.Errorf("unknown op %q (§5.1.6)", op.Op)
        }
        return nil
}

func validateTrack(t Track) error {
        if t.Name == "" {
                return errors.New("name is required (§5.2.3)")
        }
        if t.Packaging == "" {
                return errors.New("packaging is required (§5.2.4)")
        }
        switch t.Packaging {
        case PackagingLOC, PackagingMediaTimeline, PackagingEventTimeline,
                PackagingMoQLog, PackagingMoQMetrics, PackagingCMAF:
        default:
                return fmt.Errorf("unknown packaging %q (§5.2.4)", t.Packaging)
        }
        if t.IsLive == nil {
                return errors.New("isLive is required (§5.2.7)")
        }
        live := *t.IsLive

        if t.Packaging == PackagingEventTimeline {
                if t.EventType == "" {
                        return errors.New("eventType is required when packaging=eventtimeline (§5.2.5)")
                }
        } else if t.EventType != "" {
                return errors.New("eventType MUST NOT be used unless packaging=eventtimeline (§5.2.5)")
        }

        // §5.2.8 / §5.2.9 — targetLatency and buffers are mutually
        // exclusive within a single track.
        if t.TargetLatency != nil && t.Buffers != nil {
                return errors.New("targetLatency and buffers are mutually exclusive (§5.2.8, §5.2.9)")
        }

        if live && t.TrackDuration != 0 {
                return errors.New("trackDuration MUST NOT be present when isLive is true (§5.2.35)")
        }

        return validateSapStartingTypes(t)
}

// validateSapStartingTypes bounds the two CMSF §3.5.2 track fields.
// §3.5.2.1 and §3.5.2.2 define them as plain numbers without stating a
// range, so the bounds come from the sections that constrain the SAP
// types themselves.
func validateSapStartingTypes(t Track) error {
        // §3.6.1 defines the SAP type value space CMSF signals as 0-3, so
        // neither field can name a type outside it.
        if t.MaxGrpSapStartingType != nil && (*t.MaxGrpSapStartingType < 0 || *t.MaxGrpSapStartingType > 3) {
                return fmt.Errorf("maxGrpSapStartingType %d out of range 0-3 (CMSF §3.6.1)", *t.MaxGrpSapStartingType)
        }
        if t.MaxObjSapStartingType != nil && (*t.MaxObjSapStartingType < 0 || *t.MaxObjSapStartingType > 3) {
                return fmt.Errorf("maxObjSapStartingType %d out of range 0-3 (CMSF §3.6.1)", *t.MaxObjSapStartingType)
        }
        // §3.4 additionally requires every Group to "begin with an Object
        // containing a stream access point (SAP) type 1 or 2", so on a CMAF
        // track the maximum type a Group starts with is 1 or 2. §3.4 sits
        // under §3 "CMAF Packaging" and says nothing about the Groups of a
        // track using any other packaging, so the tighter bound applies
        // only here.
        if t.Packaging == PackagingCMAF && t.MaxGrpSapStartingType != nil &&
                (*t.MaxGrpSapStartingType < 1 || *t.MaxGrpSapStartingType > 2) {
                return fmt.Errorf(
                        "maxGrpSapStartingType %d out of range 1-2 for packaging %q (CMSF §3.4)",
                        *t.MaxGrpSapStartingType, PackagingCMAF)
        }
        return nil
}

// validateTargetLatencyGroups enforces §5.2.8's requirement that all
// tracks belonging to the same render/altGroup share the same
// targetLatency (treating nil as a distinct "absent" value).
func validateTargetLatencyGroups(tracks []Track, groupName string, accessor func(Track) *int) error {
        groups := map[int]*uint32{}
        groupsSeen := map[int]bool{}
        for i, tr := range tracks {
                gp := accessor(tr)
                if gp == nil {
                        continue
                }
                g := *gp
                if !groupsSeen[g] {
                        groups[g] = tr.TargetLatency
                        groupsSeen[g] = true
                        continue
                }
                if !targetLatencyEqual(groups[g], tr.TargetLatency) {
                        return fmt.Errorf(
                                "moqt/msf: tracks[%d] %s=%d targetLatency mismatch within group (§5.2.8)",
                                i, groupName, g)
                }
        }
        return nil
}

func targetLatencyEqual(a, b *uint32) bool {
        if a == nil && b == nil {
                return true
        }
        if a == nil || b == nil {
                return false
        }
        return *a == *b
}

// isOnlyNameAndNamespace reports whether tr has no fields set beyond
// Name and Namespace. Used by §5.1.6: remove-operation entries MUST
// hold only Name and may hold Namespace. Clearing those two fields and
// comparing against the zero Track keeps this robust as Track grows.
func isOnlyNameAndNamespace(tr Track) bool {
        tr.Name = ""
        tr.Namespace = ""
        return reflect.DeepEqual(tr, Track{})
}

package msf

import "time"

// BeginBroadcast returns the initial independent catalog a publisher
// emits before any media-track objects per §11.2. The Version, the
// GeneratedAt wallclock, and the supplied tracks make up the catalog;
// callers serialise it (via [encoding/json.Marshal]) and write the
// result as the first Object on the catalog track.
//
// generatedAt is the wallclock the publisher wants recorded. Pass a
// zero [time.Time] to use [time.Now]. For VOD catalogs §5.1.2 says
// generatedAt SHOULD NOT be included if isLive is false; the
// VOD-conversion helper [EndBroadcastToVOD] honours that.
func BeginBroadcast(tracks []Track, generatedAt time.Time) Catalog {
        if generatedAt.IsZero() {
                generatedAt = time.Now()
        }
        out := Catalog{
                Version:     Version,
                GeneratedAt: generatedAt.UnixMilli(),
        }
        if len(tracks) > 0 {
                out.Tracks = append([]Track(nil), tracks...)
        }
        return out
}

// EndBroadcastTerminate returns the §11.3 final independent catalog
// with isComplete=true and an empty Tracks array. After emitting this
// catalog object the publisher MUST also end each active publication
// with PUBLISH_DONE status 0x2 Track Ended (moqt.PublishDoneTrackEnded;
// the MSF draft's older text calls it SUBSCRIBE_DONE) — this helper
// only constructs the catalog body.
func EndBroadcastTerminate(generatedAt time.Time) Catalog {
        if generatedAt.IsZero() {
                generatedAt = time.Now()
        }
        return Catalog{
                Version:     Version,
                GeneratedAt: generatedAt.UnixMilli(),
                IsComplete:  true,
                Tracks:      []Track{},
        }
}

// EndBroadcastToVOD returns the §11.3 catalog that converts a previously
// live broadcast into a VOD asset. Every track in prev has its IsLive
// flipped to false and is annotated with the duration from durations
// (keyed by track Name). Tracks present in prev but missing from
// durations are passed through with IsLive=false and TrackDuration
// left unset.
//
// The returned catalog is independent (not a delta). The publisher
// emits this catalog on the catalog track to signal the live-to-VOD
// transition. Per §5.1.2 generatedAt SHOULD NOT be included when
// isLive is false; this helper omits it for that reason.
func EndBroadcastToVOD(prev Catalog, durations map[string]uint64) Catalog {
        out := cloneCatalog(prev)
        out.IsComplete = false
        out.GeneratedAt = 0
        live := false
        for i := range out.Tracks {
                out.Tracks[i].IsLive = &live
                out.Tracks[i].TargetLatency = nil
                if d, ok := durations[out.Tracks[i].Name]; ok {
                        out.Tracks[i].TrackDuration = d
                }
        }
        return out
}

package session

import (
        "context"
        "errors"
        "fmt"
        "io"
        "slices"
        "sync"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// RequestBroker owns an established request stream's read side and
// serializes its writes, so REQUEST_UPDATE (§10.9) and long-lived follow-up
// traffic can safely coexist. [Session.UpdateRequest] reads its response
// directly off the stream and therefore cannot run concurrently with any
// other reader; once a request outlives its initial response — a relay's
// upstream subscription, a publisher answering subscriber updates — exactly
// one reader must own the stream, and that reader is [RequestBroker.Serve]:
//
//   - REQUEST_OK / REQUEST_ERROR answer in-flight [RequestBroker.Update]
//     calls, including §10.9's coalescing rule (a peer may answer N
//     pipelined updates with a single REQUEST_ERROR, which fails them all).
//   - AUTHORIZATION_TOKEN parameters on follow-ups are resolved through the
//     session token cache (§10.2.2); a cache fault closes the session with
//     the mandated code.
//   - A peer REQUEST_UPDATE is answered (§10.9) by the handler installed with
//     [RequestBroker.HandleUpdates], or declined with NOT_SUPPORTED when there
//     is none, since acknowledging an unapplied update would misstate the
//     request's state.
//   - A REQUEST_OK / REQUEST_ERROR before this side sent any REQUEST_UPDATE
//     answers nothing and closes the session: on a request this side sent it
//     is a second response (§5.1, §5.2, §6.2), and on a stream this side
//     answered only the requester sends REQUEST_UPDATE, bar a PUBLISH's
//     subscriber (§10.9). So does a SUBSCRIBE_OK on a Subscription (§5.1).
//   - On a [NamespaceSubscription]'s broker, a NAMESPACE_DONE for a suffix no
//     NAMESPACE announced closes the session (§10.19).
//   - A second GOAWAY on the stream, or one with a New Session URI received
//     by a server, closes the session (§10.4; see [RequestGoaways]).
//   - Everything else (PUBLISH_DONE, other unsolicited responses, …) is
//     handed to Serve's callback.
//
// Obtain one from a typed request handle's Broker method (e.g.
// [Publication.Broker]) or [Session.NewRequestBroker]; from then on every
// write to the stream must go through the broker ([RequestBroker.Update],
// [RequestBroker.WriteMessage], or broker-aware handle methods such as
// [Publication.Done]) — session streams do not serialize concurrent writers.
type RequestBroker struct {
        stream Stream
        sess   *Session

        // mu serializes stream writes and guards the waiter queue. It is
        // deliberately held across the REQUEST_UPDATE write: §10.9 responses
        // arrive in request order, so the waiter queue order must match the
        // write order.
        mu      sync.Mutex
        waiters []updateWaiter
        // updatesClosed is latched when the stream's reader exits (or Close is
        // called); subsequent Update calls fail immediately instead of queueing
        // a waiter nothing will ever answer. Plain WriteMessage stays allowed —
        // e.g. a PUBLISH_DONE after the peer tore its side down.
        updatesClosed bool
        streamClosed  bool
        // updated records that Update sent a REQUEST_UPDATE, the only thing a
        // REQUEST_OK or REQUEST_ERROR read here can answer.
        updated bool

        // onUpdate decides each peer REQUEST_UPDATE; nil declines it.
        // onUpdateFailed runs after a declined update (§10.9.1). Both are set
        // before Serve runs.
        onUpdate       UpdateHandler
        onUpdateFailed func()

        // updateScope is the §10.2.1 scope of peer REQUEST_UPDATEs; 0 skips the
        // check.
        updateScope message.ParamScope

        // See [RequestBroker.PeerMessages].
        noPeerUpdate bool
        noPeerNotify bool

        // handle is the typed handle this broker came from, if any: Serve and
        // Close report the subscription's end to it (see requestHandle.terminated).
        handle *requestHandle

        // Serve's per-stream state, kept across calls: the stream's GOAWAYs
        // (§10.4) when there is no handle to keep them, and on a
        // SUBSCRIBE_NAMESPACE the current prefix and the full namespaces a
        // NAMESPACE announced and no NAMESPACE_DONE withdrew (§10.19, §10.9.2).
        goaways    RequestGoaways
        nsPrefix   wire.TrackNamespace
        namespaces map[track.Key]struct{}
        // nsPrefixLost is set, under mu, once any Update gave up: responses no
        // longer pair reliably, so the prefix is unknown and NAMESPACE_DONEs go
        // unchecked.
        nsPrefixLost bool
}

// PeerMessages declares whether the peer may send REQUEST_UPDATE (§10.9) and
// PUBLISH_STATE_NOTIFY (§10.10) on this stream; a disallowed one closes the
// session with PROTOCOL_VIOLATION. Typed handles set this; a broker from
// [Session.NewRequestBroker] allows both. Call it before [RequestBroker.Serve].
func (b *RequestBroker) PeerMessages(requestUpdate, publishStateNotify bool) {
        b.noPeerUpdate, b.noPeerNotify = !requestUpdate, !publishStateNotify
}

// UpdateScope sets the §10.2.1 parameter scope of the peer's REQUEST_UPDATEs;
// one carrying a parameter outside it closes the session with
// PROTOCOL_VIOLATION. Typed handles set this; a broker from
// [Session.NewRequestBroker] checks nothing. Call it before
// [RequestBroker.Serve].
func (b *RequestBroker) UpdateScope(s message.ParamScope) { b.updateScope = s }

// UpdateHandler decides a peer's REQUEST_UPDATE (§10.9). It returns the
// REQUEST_OK to send, or an error: a *[RequestRejectedError] is sent as
// REQUEST_ERROR with its code, reason and Retry Interval, any other error as
// INTERNAL_ERROR. So is a REDIRECT, which cannot answer a REQUEST_UPDATE
// (§10.6.2).
type UpdateHandler func(upd *message.RequestUpdate) (*message.RequestOK, error)

// HandleUpdates installs the handler that decides peer REQUEST_UPDATEs,
// replacing any earlier one. Call it before [RequestBroker.Serve].
func (b *RequestBroker) HandleUpdates(h UpdateHandler) { b.onUpdate = h }

// answerUpdate writes the §10.9 response to upd and reports whether the
// update was accepted.
func (b *RequestBroker) answerUpdate(upd *message.RequestUpdate) (bool, error) {
        var (
                ok  *message.RequestOK
                err error
        )
        if b.onUpdate == nil {
                err = &RequestRejectedError{Code: moqt.RequestNotSupported, Reason: "REQUEST_UPDATE not supported"}
        } else {
                ok, err = b.onUpdate(upd)
        }
        // A handler that closed the session leaves no request to answer.
        select {
        case <-b.sess.Done():
                if err == nil {
                        err = ErrRequestStreamClosed
                }
                return false, err
        default:
        }
        if err == nil {
                if ok == nil {
                        ok = &message.RequestOK{}
                }
                if len(ok.TrackProperties) > 0 {
                        // §10.5: REQUEST_UPDATE_OK's Track Properties are empty.
                        err = fmt.Errorf("%w: REQUEST_UPDATE_OK", ErrTrackPropertiesNotAllowed)
                }
        }
        if err == nil {
                if werr := b.WriteMessage(ok); werr != nil {
                        return false, fmt.Errorf("moqt/session: write REQUEST_UPDATE_OK: %w", werr)
                }
                return true, nil
        }
        rej, isRej := errors.AsType[*RequestRejectedError](err)
        if !isRej || rej.Code == moqt.RequestRedirect {
                rej = &RequestRejectedError{Code: moqt.RequestInternalError, Reason: err.Error()}
        }
        if werr := b.WriteMessage(&message.RequestError{
                ErrorCode:     rej.Code,
                RetryInterval: rej.RetryInterval,
                ErrorReason:   rej.Reason,
        }); werr != nil {
                return false, fmt.Errorf("moqt/session: write REQUEST_UPDATE error: %w", werr)
        }
        return false, nil
}

// goawayState is the stream's GOAWAY checker: the handle's, shared with its
// other readers (e.g. [TrackSubscription.ReadPublishSkipped]), or the
// broker's own.
func (b *RequestBroker) goawayState() *RequestGoaways {
        if b.handle != nil {
                return &b.handle.goaways
        }
        return &b.goaways
}

// namespaceFollowup tracks a NAMESPACE or NAMESPACE_DONE on a
// SUBSCRIBE_NAMESPACE's broker: "If a subscriber receives a NAMESPACE_DONE
// before the corresponding NAMESPACE, it MUST close the session with a
// 'PROTOCOL_VIOLATION'" (§10.19). Other brokers do not track them.
func (b *RequestBroker) namespaceFollowup(m message.Message) error {
        if b.handle == nil || !b.handle.namespaces {
                return nil
        }
        b.mu.Lock()
        lost := b.nsPrefixLost
        b.mu.Unlock()
        if lost {
                return nil
        }
        switch m := m.(type) {
        case *message.Namespace:
                if b.namespaces == nil {
                        b.namespaces = make(map[track.Key]struct{})
                }
                b.namespaces[track.NewKey(slices.Concat(b.nsPrefix, m.TrackNamespaceSuffix), nil)] = struct{}{}
        case *message.NamespaceDone:
                key := track.NewKey(slices.Concat(b.nsPrefix, m.TrackNamespaceSuffix), nil)
                if _, ok := b.namespaces[key]; !ok {
                        return b.sess.closeProtocolViolation(fmt.Errorf(
                                "moqt/session: NAMESPACE_DONE for %v before its NAMESPACE", m.TrackNamespaceSuffix))
                }
                delete(b.namespaces, key)
        }
        return nil
}

// answered is the request type (SUBSCRIBE or PUBLISH) of a broker on this
// side's request, whose response the peer has sent; zero otherwise.
func (b *RequestBroker) answered() message.Type {
        if b.handle == nil {
                return 0
        }
        return b.handle.answered
}

// updateWaiter is one sent REQUEST_UPDATE awaiting its §10.9 response and,
// when the update sets TRACK_NAMESPACE_PREFIX, the prefix it switches a
// namespace subscription to once accepted (§10.9.2).
type updateWaiter struct {
        ch         chan updateResult
        prefix     wire.TrackNamespace
        setsPrefix bool
}

// updateResult carries one §10.9 response to a waiting Update call.
type updateResult struct {
        ok  *message.RequestOK
        err error
}

// ErrRequestStreamClosed is returned by [RequestBroker.Update] when the
// request stream's reader has exited (peer FIN/reset or session shutdown) —
// no further REQUEST_UPDATE can be answered.
var ErrRequestStreamClosed = errors.New("moqt/session: request stream closed")

// NewRequestBroker builds a [RequestBroker] for an established request
// stream. Typed request handles expose a Broker method that fills this in;
// use this constructor for accept-side streams (a [Request] this endpoint
// accepted). On a request this side sent, read its response first: the
// broker reads any REQUEST_OK or REQUEST_ERROR as answering a REQUEST_UPDATE,
// and closes the session on one when none was sent.
func (s *Session) NewRequestBroker(stream Stream) *RequestBroker {
        return &RequestBroker{stream: stream, sess: s}
}

// mapUpdateResponse converts a §10.9 response: REQUEST_OK passes through,
// REQUEST_ERROR becomes a *RequestRejectedError, anything else is an error. A
// REQUEST_UPDATE_OK carrying Track Properties closes the session (§10.5).
func (s *Session) mapUpdateResponse(msg message.Message) (*message.RequestOK, error) {
        switch m := msg.(type) {
        case *message.RequestOK:
                if err := s.checkRequestOKTrackProperties(nil, m); err != nil {
                        return nil, err
                }
                if err := s.CheckPeerParams(message.ScopeRequestUpdateOK, m); err != nil {
                        return nil, err
                }
                return m, nil
        case *message.RequestError:
                return nil, s.rejection(m, message.TypeRequestUpdate)
        default:
                return nil, fmt.Errorf("moqt/session: unexpected %s in REQUEST_UPDATE response", msg.Type())
        }
}

// Update sends a REQUEST_UPDATE (§10.9) on the request stream and awaits the
// single REQUEST_OK / REQUEST_ERROR the spec mandates, delivered by the
// [RequestBroker.Serve] reader. params carries only the fields to change;
// any parameter omitted keeps its prior value on the peer.
//
// A REQUEST_ERROR is surfaced as a *RequestRejectedError. On ctx expiry the
// waiter is removed from the queue, so a peer that never answers cannot
// permanently shift response routing for later updates. (If the response is
// merely late, routing for updates written after the removal shifts by one —
// the lesser evil versus permanent poisoning; conforming peers answer.)
//
// Known limitation: the REQUEST_UPDATE write itself runs under the write
// lock and is not ctx-bounded — a peer that stalls stream flow control
// blocks Update until the session dies and errors the write.
func (b *RequestBroker) Update(ctx context.Context, params message.Parameters) (*message.RequestOK, error) {
        ch := make(chan updateResult, 1)

        b.mu.Lock()
        if b.updatesClosed {
                b.mu.Unlock()
                return nil, ErrRequestStreamClosed
        }
        // Write while holding mu: it serializes writers on the stream AND keeps
        // the waiter queue order equal to the write order, which is what lets
        // the reader pair each §10.9 response with its update. The ID is
        // allocated under the same lock so IDs appear on this stream in
        // increasing order — §10.1: REQUEST_UPDATE consumes a fresh Request ID
        // from the sender's space (the stream, not the ID, names the request
        // being updated; a reused ID is a session-fatal duplicate).
        err := message.Marshal(b.stream, &message.RequestUpdate{
                RequestID:  b.sess.AllocRequestID(),
                Parameters: params,
        })
        if err == nil {
                w := updateWaiter{ch: ch}
                if p, ok := params.Find(message.ParamTrackNamespacePrefix); ok {
                        prefix, perr := message.TrackNamespacePrefixFromParam(p)
                        w.prefix, w.setsPrefix = prefix, perr == nil
                }
                b.waiters = append(b.waiters, w)
                b.updated = true
        }
        b.mu.Unlock()
        if err != nil {
                return nil, fmt.Errorf("moqt/session: write REQUEST_UPDATE: %w", err)
        }

        select {
        case res := <-ch:
                return res.ok, res.err
        case <-ctx.Done():
                b.mu.Lock()
                if i := slices.IndexFunc(b.waiters, func(w updateWaiter) bool { return w.ch == ch }); i >= 0 {
                        // Its late answer will pair with the next update, so on a
                        // SUBSCRIBE_NAMESPACE where a prefix update's REQUEST_OK falls
                        // among the NAMESPACEs (§10.9.2) is no longer known.
                        if b.handle != nil && b.handle.namespaces {
                                b.nsPrefixLost = true
                        }
                        b.waiters = slices.Delete(b.waiters, i, i+1)
                }
                b.mu.Unlock()
                return nil, ctx.Err()
        }
}

// WriteMessage marshals a control message onto the request stream under the
// same lock that serializes Update's REQUEST_UPDATE writes.
func (b *RequestBroker) WriteMessage(msg message.Message) error {
        b.mu.Lock()
        defer b.mu.Unlock()
        return message.Marshal(b.stream, msg)
}

// WriteMessageAfterSetup runs setup and then marshals msg, both under the
// write lock. It exists for responder-side visibility ordering: setup
// typically publishes state that lets other goroutines write to this stream
// through the broker (e.g. a relay registering an upstream subscription
// that a concurrent propagation path may immediately Update). Running both
// under the lock guarantees that msg is the stream's next message — a write
// triggered by the new visibility serializes behind it — while the peer
// cannot observe msg before setup completed. A setup error aborts the
// write. setup must not use the broker or write to the stream itself, and
// must not acquire locks that stream writers hold while using the broker.
func (b *RequestBroker) WriteMessageAfterSetup(setup func() error, msg message.Message) error {
        b.mu.Lock()
        defer b.mu.Unlock()
        if err := setup(); err != nil {
                return err
        }
        return message.Marshal(b.stream, msg)
}

// writeThenClose marshals msg and FINs the send side under the write lock —
// the broker-aware backend of terminal handle methods like
// [Publication.Done].
func (b *RequestBroker) writeThenClose(msg message.Message) error {
        b.mu.Lock()
        defer b.mu.Unlock()
        if err := message.Marshal(b.stream, msg); err != nil {
                return err
        }
        return b.stream.Close()
}

// route delivers a REQUEST_OK / REQUEST_ERROR read off the stream to
// in-flight Update calls: a REQUEST_OK answers the oldest waiter; a
// REQUEST_ERROR answers ALL of them, because §10.9 lets the peer coalesce
// pipelined updates and "only a single REQUEST_ERROR will be sent" for the
// batch. It reports whether any waiter consumed the message; false means
// none was pending (an unsolicited response Serve hands to its callback).
func (b *RequestBroker) route(msg message.Message) bool {
        b.mu.Lock()
        if len(b.waiters) == 0 {
                b.mu.Unlock()
                return false
        }
        var recipients []updateWaiter
        if _, isErr := msg.(*message.RequestError); isErr {
                recipients, b.waiters = b.waiters, nil
        } else {
                recipients, b.waiters = b.waiters[:1], b.waiters[1:]
        }
        b.mu.Unlock()

        ok, err := b.sess.mapUpdateResponse(msg)
        if err == nil && recipients[0].setsPrefix {
                // §10.9.2: suffixes after this REQUEST_OK are relative to the new
                // prefix. Only Serve's goroutine reads nsPrefix.
                b.nsPrefix = recipients[0].prefix
        }
        res := updateResult{ok: ok, err: err}
        for _, w := range recipients {
                w.ch <- res
        }
        return true
}

// closeUpdates latches the broker shut for updates and fails every pending
// Update with [ErrRequestStreamClosed]. Idempotent; Serve calls it on exit
// and Close calls it as part of full teardown.
func (b *RequestBroker) closeUpdates() {
        b.mu.Lock()
        waiters := b.waiters
        b.waiters = nil
        b.updatesClosed = true
        b.mu.Unlock()
        for _, w := range waiters {
                w.ch <- updateResult{err: ErrRequestStreamClosed}
        }
}

// Close cancels the request (§3.3.3): pending and future Updates fail with
// [ErrRequestStreamClosed] and both directions are reset with code, which
// unblocks a running Serve. On a broker from a [Subscription] or
// [IncomingPublication] the subscription is Terminated (§5.1) and its Track
// Alias released (§11.1). Serialized against in-flight writes; idempotent.
// Must not be called with locks that Serve's callback might need held.
func (b *RequestBroker) Close(code moqt.StreamResetCode) {
        if b.handle != nil {
                b.handle.terminated()
        }
        b.closeUpdates()
        b.mu.Lock()
        defer b.mu.Unlock()
        if b.streamClosed {
                return
        }
        b.streamClosed = true
        b.stream.CancelRead(uint64(code))
        b.stream.CancelWrite(uint64(code))
}

// receiveUpdate checks and answers a peer REQUEST_UPDATE for Serve (§10.9).
// A non-nil error means Serve stops; a violation closed the session.
func (b *RequestBroker) receiveUpdate(m *message.RequestUpdate, updates *RequestUpdateLimiter) error {
        if b.noPeerUpdate {
                return b.sess.closeProtocolViolation(errors.New(
                        "moqt/session: REQUEST_UPDATE from a peer that may not send one"))
        }
        if b.updateScope != 0 {
                if err := b.sess.CheckPeerParams(b.updateScope, m); err != nil {
                        return err
                }
        }
        // §10.1: a REQUEST_UPDATE consumes a Request ID from the sender's space;
        // a wrong-parity or duplicate ID is session-fatal.
        if err := b.sess.CheckPeerRequestID(m.RequestID); err != nil {
                _ = b.sess.Close(moqt.SessionInvalidRequestID, err.Error())
                return err
        }
        // §10.3.1.7: reject a REQUEST_UPDATE that exceeds the per-stream
        // MAX_REQUEST_UPDATES limit before acting on it.
        if err := updates.Received(); err != nil {
                _ = b.sess.Close(moqt.SessionTooManyRequestUpdates, err.Error())
                return err
        }
        // §10.9: "MUST respond with exactly one REQUEST_OK or REQUEST_ERROR".
        // onMsg still observes the update.
        accepted, err := b.answerUpdate(m)
        if err != nil {
                return err
        }
        updates.Responded()
        if !accepted && b.onUpdateFailed != nil {
                b.onUpdateFailed()
        }
        return nil
}

// readFailed ends Serve on a read error: nil on the peer's FIN, ctx.Err() on
// cancellation, and err otherwise, after resetting the read side.
func (b *RequestBroker) readFailed(ctx context.Context, err error) error {
        switch {
        case ctx.Err() != nil:
                return ctx.Err()
        case errors.Is(err, io.EOF):
                if b.handle != nil {
                        b.handle.peerFinished()
                }
                return nil
        case errors.Is(err, message.ErrMalformedMessage):
                // §10, and §10.2 for an unknown parameter.
                b.stream.CancelRead(uint64(moqt.StreamResetInternalError))
                return b.sess.closeProtocolViolation(err)
        default:
                // Covers peer resets too (a STOP_SENDING on an already-reset stream
                // is a transport no-op).
                b.stream.CancelRead(uint64(moqt.StreamResetInternalError))
                return err
        }
}

// Serve owns every read on the request stream until the peer tears it down
// (EOF / reset), ctx is cancelled (the read side is then reset to unblock
// the parse), or onMsg returns false. On exit, pending and future Update
// calls fail with [ErrRequestStreamClosed].
//
// Responses route to Update waiters; a token cache fault closes the session
// (§10.2.2); peer REQUEST_UPDATEs are answered as described on
// [RequestBroker], and a REQUEST_OK or REQUEST_ERROR before this side sent a
// REQUEST_UPDATE closes the session (§5.1, §10.9). Every other message,
// including each REQUEST_UPDATE, is passed to onMsg (nil means "discard");
// return false from onMsg to stop serving.
//
// A read error resets the read side with INTERNAL_ERROR; a malformed follow-up
// also closes the session with PROTOCOL_VIOLATION (§10). Serve returns nil on
// a clean FIN or an onMsg stop, ctx.Err() on cancellation, and the read/token
// error otherwise.
//
// On a broker from a [Subscription] or [IncomingPublication], every exit but
// an onMsg stop Terminates the subscription (§5.1) and releases its Track
// Alias (§11.1).
func (b *RequestBroker) Serve(ctx context.Context, onMsg func(message.Message) bool) error {
        defer b.closeUpdates()
        stopped := false // by onMsg: the stream may still carry the subscription
        defer func() {
                if !stopped && b.handle != nil {
                        b.handle.terminated()
                }
        }()
        stop := context.AfterFunc(ctx, func() {
                b.stream.CancelRead(uint64(moqt.StreamResetSessionClosed))
        })
        defer stop()

        // §10.3.1.7: per-stream MAX_REQUEST_UPDATES enforcement. One limiter per
        // stream, since the limit is scoped to a single request stream.
        updates := b.sess.NewRequestUpdateLimiter()

        for {
                msg, err := message.Parse(b.stream)
                if err != nil {
                        return b.readFailed(ctx, err)
                }

                // §10.2.2: follow-ups may REGISTER/DELETE token aliases; skipping
                // this would silently desynchronize the peer's view of the token
                // cache. A cache fault is session-fatal with the mandated code.
                if _, err := b.sess.ProcessFollowupTokens(msg); err != nil {
                        if tce, ok := errors.AsType[*TokenCacheError](err); ok {
                                _ = b.sess.Close(tce.Code, tce.Error())
                        }
                        return err
                }

                switch m := msg.(type) {
                case *message.SubscribeOK:
                        // §5.1: the response to this side's SUBSCRIBE was read already,
                        // and "The peer SHOULD close the session with a protocol error if
                        // it receives more than one."
                        if b.answered() == message.TypeSubscribe {
                                return b.sess.closeProtocolViolation(errors.New("moqt/session: SUBSCRIBE_OK after the response"))
                        }
                case *message.RequestOK, *message.RequestError:
                        // The request's own response was read before the broker
                        // attached, so every REQUEST_OK here is a REQUEST_UPDATE_OK
                        // (§10.5), even one whose Update gave up. Before any Update it
                        // answers nothing: on a request this side sent it is a second
                        // response, a second PUBLISH_OK among them (§5.1, §5.2, §6.2),
                        // and on a stream this side answered the peer, as requester, was
                        // sent no REQUEST_UPDATE to answer (§10.9).
                        b.mu.Lock()
                        updated := b.updated
                        b.mu.Unlock()
                        if !updated {
                                return b.sess.closeProtocolViolation(
                                        fmt.Errorf("moqt/session: %s before any REQUEST_UPDATE", m.Type()))
                        }
                        if err := b.sess.checkRequestOKTrackProperties(nil, m); err != nil {
                                return err
                        }
                        // §10.6.1, before routing: an unsolicited REQUEST_ERROR reaches
                        // onMsg instead of an Update.
                        if e, isErr := m.(*message.RequestError); isErr {
                                if err := redirectViolation(
                                        e.Redirect,
                                        message.TypeRequestUpdate,
                                        b.sess.role == roleServer,
                                ); err != nil {
                                        return b.sess.closeProtocolViolation(fmt.Errorf("moqt/session: received %w", err))
                                }
                        }
                        if err := b.sess.CheckPeerParams(message.ScopeRequestUpdateOK, m); err != nil {
                                return err
                        }
                        if b.route(msg) {
                                continue
                        }
                        // Unsolicited response — surface via onMsg below.
                case *message.Goaway:
                        if err := b.goawayState().Received(b.sess, m); err != nil {
                                return err
                        }
                case *message.Namespace, *message.NamespaceDone:
                        if err := b.namespaceFollowup(m); err != nil {
                                return err
                        }
                case *message.PublishStateNotify:
                        if b.noPeerNotify {
                                return b.sess.closeProtocolViolation(errors.New(
                                        "moqt/session: PUBLISH_STATE_NOTIFY from a peer that may not send one"))
                        }
                        if err := b.sess.CheckPeerParams(message.ScopePublishStateNotify, m); err != nil {
                                return err
                        }
                case *message.RequestUpdate:
                        if err := b.receiveUpdate(m, updates); err != nil {
                                return err
                        }
                }

                if onMsg != nil && !onMsg(msg) {
                        stopped = true
                        return nil
                }
        }
}

package session

import (
        "errors"
        "fmt"
        "io"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
)

// sendControl queues a control message for the send loop. Blocks if the queue
// is full or the session is done.
func (s *Session) sendControl(msg message.Message) error {
        select {
        case s.controlOut <- msg:
                return nil
        case <-s.done:
                return errors.New("moqt/session: closed")
        }
}

// controlSendLoop serializes writes onto the send-control stream. It exits on
// session shutdown or on the first write error.
func (s *Session) controlSendLoop() {
        for {
                select {
                case msg := <-s.controlOut:
                        if err := message.Marshal(s.sendCtrl, msg); err != nil {
                                if s.sessionDoneAlready() {
                                        return
                                }
                                _ = s.Close(moqt.SessionInternalError, "control send failure")
                                return
                        }
                case <-s.done:
                        return
                }
        }
}

// controlRecvLoop reads framed control messages off the recv-control stream
// and dispatches them. The loop owns shutdown on read failure unless the
// session is already terminating.
func (s *Session) controlRecvLoop() {
        for {
                msg, err := message.Parse(s.recvCtrl)
                if err != nil {
                        if s.sessionDoneAlready() {
                                return
                        }
                        if errors.Is(err, io.EOF) {
                                // Peer closed the control stream cleanly. §3.3 forbids
                                // this during the session lifetime; treat it as a
                                // protocol violation.
                                _ = s.Close(moqt.SessionProtocolViolation, "peer closed control stream")
                                return
                        }
                        _ = s.Close(moqt.SessionProtocolViolation, err.Error())
                        return
                }
                if err := s.dispatchControl(msg); err != nil {
                        if s.sessionDoneAlready() {
                                return
                        }
                        // Every violation dispatchControl can report is a §3.5
                        // PROTOCOL_VIOLATION: the only messages it accepts after SETUP are
                        // GOAWAY (§10.4), a duplicate SETUP, and anything table 5 in §10
                        // disallows outright. A rule mandating a different close code would
                        // need the handler to carry it out of here.
                        _ = s.Close(moqt.SessionProtocolViolation, err.Error())
                        return
                }
        }
}

func (s *Session) sessionDoneAlready() bool {
        select {
        case <-s.done:
                return true
        default:
                return false
        }
}

// dispatchControl handles a single control-stream message after SETUP. Per
// table 5 in §10, only GOAWAY is valid on the control stream after SETUP for
// the messages in scope; anything else is a protocol violation.
func (s *Session) dispatchControl(msg message.Message) error {
        switch m := msg.(type) {
        case *message.Goaway:
                return s.handleGoaway(m)
        case *message.Setup:
                return errors.New("duplicate SETUP on control stream")
        default:
                return fmt.Errorf("unexpected %s on control stream", msg.Type())
        }
}

package session

import (
        "context"
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// paddingDatagramType is the MoQT PADDING datagram type (§11.5.2).
const paddingDatagramType uint64 = 0x132B3E29

// ReceiveDatagram blocks until a QUIC DATAGRAM frame arrives from the peer,
// parses it, and returns the contained ObjectDatagram. PADDING datagrams
// (§11.5.2) are silently consumed and the call retries. Unknown datagram types
// close the session (§11) with PROTOCOL_VIOLATION.
//
// Transport-level errors (session closed, ctx cancelled) are returned
// unwrapped so the caller can distinguish them from parse failures.
//
// An Object whose Properties make its track malformed is returned together
// with an error wrapping [ErrMalformedTrack], so the caller knows which track
// to cancel; the session stays up.
func (s *Session) ReceiveDatagram(ctx context.Context) (*message.ObjectDatagram, error) {
        for {
                raw, err := s.conn.ReceiveDatagram(ctx)
                if err != nil {
                        return nil, err
                }

                // Peek the type varint to dispatch. ObjectDatagram.Parse will re-read
                // it from a fresh Reader, so we only need the value here.
                peek := wire.NewReader(raw)
                typ, err := peek.Varint()
                if err != nil {
                        return nil, s.closeProtocolViolation(
                                fmt.Errorf("moqt/session: datagram type varint: %w", err))
                }

                switch {
                case message.IsValidDatagramType(typ):
                        obj := &message.ObjectDatagram{}
                        if err := obj.Parse(wire.NewReader(raw)); err != nil {
                                return nil, s.closeProtocolViolation(
                                        fmt.Errorf("moqt/session: parse OBJECT_DATAGRAM: %w", err))
                        }
                        if len(obj.Properties) > 0 {
                                if err := message.CheckObjectProperties(obj.Properties, obj.GroupID, obj.ObjectID); err != nil {
                                        return obj, fmt.Errorf("%w: datagram Group %d Object %d: %w",
                                                ErrMalformedTrack, obj.GroupID, obj.ObjectID, err)
                                }
                        }
                        return obj, nil

                case typ == paddingDatagramType:
                        // §11.5.2: receiver MUST discard all data in a padding datagram.
                        continue

                default:
                        return nil, s.closeProtocolViolation(
                                fmt.Errorf("moqt/session: unknown datagram type %#x", typ))
                }
        }
}

// SendDatagram serializes d and sends it as a single QUIC DATAGRAM
// frame. Returns an error if d fails validation or the payload exceeds the
// negotiated max_datagram_frame_size (the transport returns an error in that
// case; per §11.3 the object is silently dropped at the sender).
//
// SendDatagram is the publisher-side counterpart of [Session.ReceiveDatagram].
func (s *Session) SendDatagram(d *message.ObjectDatagram) error {
        if err := d.Validate(); err != nil {
                return fmt.Errorf("moqt/session: SendDatagram: %w", err)
        }
        // Reuse a pooled writer rather than allocating one per datagram. The
        // transport copies the bytes before SendDatagram returns (quic-go and the
        // in-process pipe both make a copy), so the buffer is free to recycle.
        w, _ := writerPool.Get().(*wire.Writer)
        w.Reset()
        d.Append(w)
        err := s.conn.SendDatagram(w.Bytes())
        writerPool.Put(w)
        return err
}

// closeProtocolViolation closes the session with PROTOCOL_VIOLATION and
// returns err so callers can return it directly.
func (s *Session) closeProtocolViolation(err error) error {
        _ = s.Close(moqt.SessionProtocolViolation, err.Error())
        return err
}

package session

import (
        "bufio"
        "context"
        "errors"
        "fmt"
        "io"
        "math/bits"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// ErrPaddingStream is returned by AcceptDataStream when a padding uni-stream
// (§11.5, type 0x132B3E28) is received. Callers SHOULD loop and call
// AcceptDataStream again.
var ErrPaddingStream = errors.New("moqt/session: padding stream received (ignorable)")

// ErrMalformedTrack wraps the error a read returns for an Object that makes
// its track malformed (§2.4.2), such as Object Properties that fail
// [message.CheckObjectProperties]. The session stays up; the caller MUST
// cancel the corresponding subscription or fetch (§2.4.2).
var ErrMalformedTrack = errors.New("moqt/session: malformed track")

// ---------------------------------------------------------------------------
// DataStream — sealed interface returned by AcceptDataStream
// ---------------------------------------------------------------------------

// DataStream is the sealed interface returned by AcceptDataStream. The
// concrete type is either *IncomingSubgroupStream or *IncomingFetchStream;
// callers type-switch to obtain the typed stream.
//
// Read is included so that io.Copy / io.ReadAll work directly on the
// interface without a type assertion.
type DataStream interface {
        // Read returns body bytes that follow the parsed header.
        Read(p []byte) (int, error)
        // Cancel resets the stream with the given application code (§3.3.4).
        Cancel(code moqt.StreamResetCode)
        // isDataStream seals the interface to this package.
        isDataStream()
}

// ---------------------------------------------------------------------------
// IncomingSubgroupStream
// ---------------------------------------------------------------------------

// IncomingSubgroupStream is an accepted inbound SUBGROUP_HEADER uni-stream
// whose leading header has already been parsed. The remaining bytes are the
// body, consumed via [IncomingSubgroupStream.ReadObject] (raw, delta-encoded
// ObjectID), [IncomingSubgroupStream.ReadDecoded] (absolute IDs with state
// carried across calls), or Read (raw bytes). The peer's FIN surfaces as
// io.EOF from any of these methods.
type IncomingSubgroupStream struct {
        // Header is the parsed SUBGROUP_HEADER (§11.4.2).
        Header message.SubgroupHeader
        src    ReceiveStream
        br     *bufio.Reader

        // rd is a StreamReader bound to br once at construction and reused by
        // every ReadObject call. Allocating it per-object showed up as the
        // single largest allocation site on the fanout read path (it escapes
        // to the heap because it's passed as the wire.Decoder interface to
        // SubgroupObject.Parse).
        rd *wire.StreamReader

        // Decoder state, kept by ReadObject (§11.4.2).
        decPrevObject       uint64
        decHavePrev         bool
        decSubgroupID       uint64 // resolved per §11.4.2 (zero / first-object / explicit)
        decSubgroupResolved bool

        // sess is the owning session, used by TrackKey to resolve Header.TrackAlias
        // against the inbound alias registry live, at call time.
        sess *Session
}

// TrackKey returns the track this subgroup belongs to, resolved from the
// stream's §11.1 Track Alias (Header.TrackAlias) via the inbound alias registry
// the session populates on SUBSCRIBE_OK and [Request.AcceptPublish]. The second
// result is false when the alias is not registered, in which case callers fall
// back to Header.TrackAlias and their own mapping.
//
// Resolution is live: it queries the registry at call time, not at accept time.
// So if a subgroup stream is accepted before the SUBSCRIBE_OK that binds its
// alias (a legitimate §11.1 ordering), a TrackKey call once the alias has been
// registered resolves correctly rather than being pinned to a stale snapshot.
func (s *IncomingSubgroupStream) TrackKey() (track.Key, bool) {
        return s.sess.LookupInboundTrackAlias(s.Header.TrackAlias)
}

// InboundTrack is [IncomingSubgroupStream.TrackKey] plus the rest of what the
// alias is bound to. Resolution is live, as for TrackKey.
func (s *IncomingSubgroupStream) InboundTrack() (InboundTrack, bool) {
        return s.sess.LookupInboundTrack(s.Header.TrackAlias)
}

// AwaitInboundTrack is [IncomingSubgroupStream.InboundTrack] that waits for
// the stream's Track Alias to be registered, until ctx ends or the session
// closes (§11.4.2 allows buffering "for a brief period").
//
// The unread stream holds connection flow control, and the bundled transports
// do not reserve it for control streams (§11.4.2), so early data can stall
// the SUBSCRIBE_OK being waited for. Keep ctx's bound short.
func (s *IncomingSubgroupStream) AwaitInboundTrack(ctx context.Context) (InboundTrack, bool) {
        return s.sess.awaitInboundTrack(ctx, s.Header.TrackAlias)
}

func (s *IncomingSubgroupStream) isDataStream() {}

// Read returns body bytes that follow the parsed header. Prefer ReadObject
// for correctly-framed object access.
func (s *IncomingSubgroupStream) Read(p []byte) (int, error) { return s.br.Read(p) }

// ObjectID returns the absolute Object ID (§11.4.2) of the Object last read,
// including one returned with [ErrMalformedTrack].
func (s *IncomingSubgroupStream) ObjectID() uint64 { return s.decPrevObject }

// Cancel resets the stream with the given application code (§3.3.4).
func (s *IncomingSubgroupStream) Cancel(code moqt.StreamResetCode) {
        s.src.CancelRead(uint64(code))
}

// ReadObject reads the next framed SubgroupObject from the stream body.
// Returns (nil, io.EOF) when the peer has FIN'd the stream cleanly. The
// returned [message.SubgroupObject] holds the raw §11.4.2 ObjectIDDelta;
// use [IncomingSubgroupStream.ReadDecoded] when you want absolute IDs and
// implicit SubgroupID resolution done for you.
//
// An Object whose Properties make the track malformed returns an error
// wrapping [ErrMalformedTrack].
func (s *IncomingSubgroupStream) ReadObject() (*message.SubgroupObject, error) {
        obj := &message.SubgroupObject{}
        if err := obj.Parse(s.rd, s.Header.Properties); err != nil {
                return nil, s.sess.checkFINMidObject(err)
        }
        // §11.2.1.1, §11.2.1.2.
        if err := obj.Validate(); err != nil {
                return nil, s.sess.closeProtocolViolation(fmt.Errorf("moqt/session: subgroup object: %w", err))
        }
        // §11.4.2. Resolved here so the Properties check has the absolute ID.
        objectID := obj.ObjectIDDelta
        if s.decHavePrev {
                var err error
                if objectID, err = message.NextSubgroupObjectID(s.decPrevObject, obj.ObjectIDDelta); err != nil {
                        return nil, s.sess.closeProtocolViolation(err)
                }
        }
        s.decPrevObject, s.decHavePrev = objectID, true
        // Resolve the §11.4.2 SubgroupID once per stream, before the Properties
        // check, so a malformed first Object still sets it.
        if !s.decSubgroupResolved {
                switch s.Header.SubgroupIDMode {
                case message.SubgroupIDImplicitZero:
                        s.decSubgroupID = 0
                case message.SubgroupIDImplicitFirstObject:
                        s.decSubgroupID = objectID
                case message.SubgroupIDExplicit:
                        s.decSubgroupID = s.Header.SubgroupID
                }
                s.decSubgroupResolved = true
        }
        if len(obj.Properties) > 0 {
                if err := message.CheckObjectProperties(obj.Properties, s.Header.GroupID, objectID); err != nil {
                        return nil, fmt.Errorf("%w: Group %d Object %d: %w", ErrMalformedTrack, s.Header.GroupID, objectID, err)
                }
        }
        return obj, nil
}

// DecodedSubgroupObject is the absolute-coordinates view of one §11.4.2
// SubgroupObject. ReadDecoded reconstructs the absolute ObjectID from the
// per-object delta + the running previous ObjectID, and resolves the
// SubgroupID from the header's [message.SubgroupIDMode]:
//
//   - SubgroupIDImplicitZero        → SubgroupID = 0
//   - SubgroupIDImplicitFirstObject → SubgroupID = first object's ObjectID
//   - SubgroupIDExplicit            → SubgroupID = header's value
//
// GroupID is constant for the stream (from the header) and is copied
// onto every decoded object so callers can pass the decoded value alone
// without also threading the stream header through their pipeline.
type DecodedSubgroupObject struct {
        GroupID      uint64
        SubgroupID   uint64
        ObjectID     uint64
        ObjectStatus uint64
        Properties   []byte
        Payload      []byte
}

// ReadDecoded reads the next SubgroupObject and resolves §11.4.2 deltas
// into absolute coordinates, carrying decoder state across calls.
// Returns (nil, io.EOF) on clean stream FIN.
//
// The first object's ObjectIDDelta is its absolute ObjectID; subsequent
// objects' deltas encode (currentID - prevID - 1) so consecutive IDs all
// encode as zero.
func (s *IncomingSubgroupStream) ReadDecoded() (*DecodedSubgroupObject, error) {
        raw, err := s.ReadObject()
        if err != nil {
                return nil, err
        }
        d := &DecodedSubgroupObject{ // ReadObject resolved the IDs
                GroupID:      s.Header.GroupID,
                SubgroupID:   s.decSubgroupID,
                ObjectID:     s.decPrevObject,
                ObjectStatus: raw.ObjectStatus,
                Properties:   raw.Properties,
                Payload:      raw.Payload,
        }
        return d, nil
}

// ---------------------------------------------------------------------------
// IncomingFetchStream
// ---------------------------------------------------------------------------

// IncomingFetchStream is an accepted inbound FETCH_HEADER uni-stream whose
// leading header has already been parsed. The remaining bytes are the body,
// consumed via [IncomingFetchStream.ReadObject] (raw, delta-encoded fields),
// [IncomingFetchStream.ReadDecoded] (absolute IDs, with state carried across
// calls), or Read (raw bytes). The peer's FIN surfaces as io.EOF from any of
// these methods.
type IncomingFetchStream struct {
        // Header is the parsed FETCH_HEADER (§11.4.4).
        Header message.FetchHeader
        src    ReceiveStream
        br     *bufio.Reader

        // rd is a StreamReader bound to br once at construction and reused by
        // every ReadObject call — see [IncomingSubgroupStream.rd].
        rd *wire.StreamReader

        sess *Session

        // GroupOrder tells [IncomingFetchStream.ReadDecoded] how to
        // interpret cross-group GroupIDDeltas (§11.4.4.1): ascending →
        // newGroup = prevGroup + delta + 1; descending → newGroup =
        // prevGroup - delta - 1. The §11.4.4 wire format does not encode
        // the direction; the caller knows it from the GROUP_ORDER
        // parameter it sent in FETCH (§10.2.8: Ascending when omitted) or,
        // for a fill fetch stream, from FILL_PARAMETERS, else the
        // subscription's group order (§10.2.15), which defaults to the
        // Track's publisher preference (§10.2.8). Defaults to ascending when
        // unset (zero value).
        GroupOrder message.GroupOrder

        // Decoder state used by ReadDecoded — running absolute values
        // carried across objects so each call only has to apply the
        // current object's deltas.
        //
        // decHavePrev means a prior Group/Object ID exists — a real object OR
        // a §11.4.4.2 End-of-Range marker (markers ARE the prior for the
        // Group/Object dimension). decHaveActual means a real object was
        // decoded: only then do a prior Subgroup ID and prior Priority exist —
        // §11.4.4.2: "If there was no prior Object, using a flag that
        // references the prior Subgroup ID [or Priority] is a
        // PROTOCOL_VIOLATION."
        decPrevGroup    uint64
        decPrevObject   uint64
        decPrevSubgroup uint64
        decPrevPriority uint8
        decHavePrev     bool
        decHaveActual   bool
}

func (s *IncomingFetchStream) isDataStream() {}

// Read returns body bytes that follow the parsed header. Prefer ReadObject
// for correctly-framed object access.
func (s *IncomingFetchStream) Read(p []byte) (int, error) { return s.br.Read(p) }

// Cancel resets the stream with the given application code (§3.3.4).
func (s *IncomingFetchStream) Cancel(code moqt.StreamResetCode) {
        s.src.CancelRead(uint64(code))
}

// ReadObject reads the next framed FetchObject from the stream body.
// Returns (nil, io.EOF) when the peer has FIN'd the stream cleanly. Fields
// on the returned [message.FetchObject] are raw — GroupIDDelta and
// ObjectIDDelta carry §11.4.4 wire deltas, not absolute IDs. Use
// [IncomingFetchStream.ReadDecoded] when you want absolute IDs reconstructed
// for you.
func (s *IncomingFetchStream) ReadObject() (*message.FetchObject, error) {
        obj := &message.FetchObject{}
        if err := obj.Parse(s.rd); err != nil {
                // §11.4.4: a Serialization Flags value of 128 or more that is not
                // an End of Range marker "is a PROTOCOL_VIOLATION".
                if errors.Is(err, message.ErrInvalidFetchFlags) {
                        return nil, s.sess.closeProtocolViolation(fmt.Errorf("moqt/session: fetch object: %w", err))
                }
                return nil, s.sess.checkFINMidObject(err)
        }
        return obj, nil
}

// DecodedFetchObject is the absolute-coordinates view of one §11.4.4
// FetchObject. ReadDecoded reconstructs GroupID / ObjectID / SubgroupID /
// PublisherPriority from the wire deltas + previous objects so the caller
// doesn't have to maintain state itself.
//
// End-of-range markers (§11.4.4.2) surface via EndOfNonExistentRange,
// EndOfUnknownRange or EndOfTimedOutRange; for those, GroupID / ObjectID hold the absolute range
// boundary the marker carries and the payload / properties fields are zero.
type DecodedFetchObject struct {
        GroupID           uint64
        ObjectID          uint64
        SubgroupID        uint64
        PublisherPriority uint8
        Properties        []byte
        Payload           []byte

        // Datagram reports the §11.4.4.1 Datagram bit (0x40): the object was
        // published with Forwarding Preference "Datagram" and has no Subgroup
        // ID (SubgroupID is 0).
        Datagram bool

        EndOfNonExistentRange bool // §11.4.4.2 flag 0x8C
        EndOfUnknownRange     bool // §11.4.4.2 flag 0x10C
        EndOfTimedOutRange    bool // §11.4.4.2 flag 0x20C (draft-20)
}

// IsEndOfRange reports whether this is any §11.4.4.2 end-of-range marker rather
// than a delivered Object. Markers describe a span of Locations — non-existent,
// unknown, or timed out — and carry no payload, subgroup, priority or
// properties, so most consumers want to skip them as a group. Prefer this over
// testing the three flags individually: draft-20 added the third one, and every
// place that had enumerated the first two silently started treating it as an
// Object.
func (d *DecodedFetchObject) IsEndOfRange() bool {
        return d.EndOfNonExistentRange || d.EndOfUnknownRange || d.EndOfTimedOutRange
}

// ReadDecoded reads the next FetchObject and resolves §11.4.4 deltas
// into absolute coordinates, carrying decoder state across calls.
// Returns (nil, io.EOF) on clean stream FIN.
//
// Subgroup-ID encoding modes (§11.4.4.1): Zero, Prior, PriorPlusOne, and
// Explicit are all resolved against the previous object's SubgroupID.
// Priority is inherited from the previous object when the per-object
// PRIORITY flag is absent.
//
// The first object on the stream carries absolute GroupID / ObjectID
// directly in the delta fields (per §11.4.4.1); subsequent objects' deltas
// are interpreted using [IncomingFetchStream.GroupOrder] for cross-group
// transitions.
//
// An Object whose Properties make the track malformed returns an error
// wrapping [ErrMalformedTrack]. ReadObject does not check, since it has no
// absolute IDs.
func (s *IncomingFetchStream) ReadDecoded() (*DecodedFetchObject, error) {
        raw, err := s.ReadObject()
        if err != nil {
                return nil, err
        }

        // §11.4.4.2: end-of-range markers carry absolute Group/Object IDs in
        // the otherwise-delta fields — and those values become the "prior
        // Group ID and prior Object ID" for the next object. The prior
        // Subgroup ID / Priority stay those of the last ACTUAL object
        // (decHaveActual tracks whether one exists).
        if raw.IsEndOfRange() {
                s.decPrevGroup = raw.GroupIDDelta
                s.decPrevObject = raw.ObjectIDDelta
                s.decHavePrev = true
                return &DecodedFetchObject{
                        GroupID:               raw.GroupIDDelta,
                        ObjectID:              raw.ObjectIDDelta,
                        EndOfNonExistentRange: raw.IsEndOfNonExistentRange(),
                        EndOfUnknownRange:     raw.IsEndOfUnknownRange(),
                        EndOfTimedOutRange:    raw.IsEndOfTimedOutRange(),
                }, nil
        }

        d := &DecodedFetchObject{
                Datagram:   raw.IsDatagram(),
                Properties: raw.Properties,
                Payload:    raw.ObjectPayload,
        }

        // §11.4.4.1 / §11.4.4.2: flags referencing the prior Object's Subgroup
        // ID or Priority are a PROTOCOL_VIOLATION until a real (non-End-of-Range)
        // object has been decoded.
        if !s.decHaveActual {
                if !raw.IsDatagram() {
                        if m := raw.SubgroupMode(); m == message.FetchSubgroupIDPrior ||
                                m == message.FetchSubgroupIDPriorPlusOne {
                                return nil, s.sess.closeProtocolViolation(fmt.Errorf(
                                        "moqt/session: fetch object references prior subgroup with no prior object (flags 0x%X)",
                                        raw.SerializationFlags))
                        }
                }
                if raw.SerializationFlags&message.FetchFlagPriority == 0 {
                        return nil, s.sess.closeProtocolViolation(fmt.Errorf(
                                "moqt/session: fetch object references prior priority with no prior object (flags 0x%X)",
                                raw.SerializationFlags))
                }
        }

        // Group / Object reconstruction.
        switch {
        case !s.decHavePrev:
                // §11.4.4.1: the first object carries absolute Group and Object IDs.
                // An End-of-Range marker counts as a prior here.
                if raw.SerializationFlags&message.FetchFlagGroupIDDelta == 0 ||
                        raw.SerializationFlags&message.FetchFlagObjectIDDelta == 0 {
                        return nil, s.sess.closeProtocolViolation(fmt.Errorf(
                                "moqt/session: first fetch object missing Group/Object ID delta (flags 0x%X)",
                                raw.SerializationFlags))
                }
                d.GroupID = raw.GroupIDDelta
                d.ObjectID = raw.ObjectIDDelta
        default:
                // §11.4.4.1: unlike §11.4.2, a present Object ID Delta carries no
                // implicit +1; an absent one means the prior ID plus one, in any
                // group. An ID outside 0..2^64-1 is a PROTOCOL_VIOLATION.
                var over uint64
                d.GroupID = s.decPrevGroup
                newGroup := raw.SerializationFlags&message.FetchFlagGroupIDDelta != 0
                if newGroup {
                        if s.decGroupOrder() == message.GroupOrderDescending {
                                d.GroupID, over = bits.Sub64(s.decPrevGroup, raw.GroupIDDelta, 1)
                        } else {
                                d.GroupID, over = bits.Add64(s.decPrevGroup, raw.GroupIDDelta, 1)
                        }
                }
                if over == 0 {
                        switch {
                        case raw.SerializationFlags&message.FetchFlagObjectIDDelta == 0:
                                d.ObjectID, over = bits.Add64(s.decPrevObject, 0, 1)
                        case newGroup:
                                d.ObjectID = raw.ObjectIDDelta
                        default:
                                d.ObjectID, over = bits.Add64(s.decPrevObject, raw.ObjectIDDelta, 0)
                        }
                }
                if over != 0 {
                        return nil, s.sess.closeProtocolViolation(fmt.Errorf(
                                "%w: FETCH object after {%d,%d} (flags 0x%X)",
                                message.ErrIDOverflow, s.decPrevGroup, s.decPrevObject, raw.SerializationFlags))
                }
        }

        // SubgroupID reconstruction per §11.4.4.1 modes. Datagram objects
        // (bit 0x40) have no Subgroup ID and the mode bits are ignored; they
        // also don't become the "prior Object's Subgroup ID" for later objects
        // (the spec is silent here; this mirrors the §11.4.4.2 rule that the
        // prior Subgroup ID comes from the last actual subgroup object).
        if !d.Datagram {
                switch raw.SubgroupMode() {
                case message.FetchSubgroupIDZero:
                        d.SubgroupID = 0
                case message.FetchSubgroupIDPrior:
                        d.SubgroupID = s.decPrevSubgroup
                case message.FetchSubgroupIDPriorPlusOne:
                        d.SubgroupID = s.decPrevSubgroup + 1
                case message.FetchSubgroupIDExplicit:
                        d.SubgroupID = raw.SubgroupID
                }
                s.decPrevSubgroup = d.SubgroupID
        }

        // Priority: inherit from previous unless explicitly set on this object.
        if raw.SerializationFlags&message.FetchFlagPriority != 0 {
                d.PublisherPriority = raw.PublisherPriority
        } else {
                d.PublisherPriority = s.decPrevPriority
        }

        // Advance decoder state (decPrevSubgroup advances above, with the same
        // datagram guard as its reconstruction).
        s.decPrevGroup = d.GroupID
        s.decPrevObject = d.ObjectID
        s.decPrevPriority = d.PublisherPriority
        s.decHavePrev = true
        s.decHaveActual = true

        if len(d.Properties) > 0 {
                if err := message.CheckObjectProperties(d.Properties, d.GroupID, d.ObjectID); err != nil {
                        return nil, fmt.Errorf("%w: Group %d Object %d: %w", ErrMalformedTrack, d.GroupID, d.ObjectID, err)
                }
        }
        return d, nil
}

// decGroupOrder returns the caller-configured GroupOrder, defaulting to
// ascending when the zero value is set. Ascending matches the relay's
// default FETCH response order, so most callers can ignore the field.
func (s *IncomingFetchStream) decGroupOrder() message.GroupOrder {
        if s.GroupOrder == message.GroupOrderDescending {
                return message.GroupOrderDescending
        }
        return message.GroupOrderAscending
}

// ---------------------------------------------------------------------------
// Session method — accept inbound data streams
// ---------------------------------------------------------------------------

// AcceptDataStream blocks until the peer opens the next data uni-stream,
// parses its leading header, and returns it wrapped so the caller can
// consume the body. The concrete type is either *IncomingSubgroupStream or
// *IncomingFetchStream; callers type-switch to obtain the typed stream.
//
// A stream that ends or is reset before its header is complete is skipped
// (§11.4.1). Interpretation: a FIN mid-header is treated the same, since §11.4
// asks for PROTOCOL_VIOLATION only on a FIN mid-Object. The returned errors
// are:
//   - ErrPaddingStream when a padding stream (§11.5.1) is received — callers
//     SHOULD loop and call AcceptDataStream again;
//   - *message.UnknownDataStreamTypeError (§3.4) or
//     *message.ReservedSubgroupIDModeError (§11.4.2); AcceptDataStream has
//     already closed the session with PROTOCOL_VIOLATION;
//   - transport-level errors (session closed, ctx cancelled), unwrapped from
//     the underlying conn.
//
// Every error other than ErrPaddingStream means the loop should terminate.
func (s *Session) AcceptDataStream(ctx context.Context) (DataStream, error) {
        for {
                ds, err := s.acceptDataStream(ctx)
                if errors.Is(err, errAbortedDataStream) {
                        continue
                }
                return ds, err
        }
}

// checkFINMidObject closes the session with PROTOCOL_VIOLATION when a data
// stream ended with a FIN mid-Object (§11.4), which the parsers report as
// io.ErrUnexpectedEOF; a reset is left alone (§11.4.1). err is returned
// unchanged.
func (s *Session) checkFINMidObject(err error) error {
        if errors.Is(err, io.ErrUnexpectedEOF) {
                _ = s.closeProtocolViolation(err)
        }
        return err
}

// errAbortedDataStream marks a data stream abandoned before its header was
// complete; [Session.AcceptDataStream] skips it.
var errAbortedDataStream = errors.New("moqt/session: data stream ended before its header")

func (s *Session) acceptDataStream(ctx context.Context) (DataStream, error) {
        src, err := s.nextUniStream(ctx)
        if err != nil {
                return nil, err
        }
        // The header reads below are context-free stream I/O; bridge ctx with
        // CancelRead (the readResponse pattern) so a peer that opens a stream
        // but stalls mid-header cannot wedge the accept loop past cancellation.
        stop := context.AfterFunc(ctx, func() {
                src.CancelRead(uint64(moqt.StreamResetCancelled))
        })
        defer stop()
        // aborted abandons a stream whose header read failed; header fields
        // cannot be malformed, so the only failure is an early FIN or reset.
        aborted := func() error {
                src.CancelRead(uint64(moqt.StreamResetInternalError))
                if ctx.Err() != nil {
                        return ctx.Err()
                }
                return errAbortedDataStream
        }
        br := bufio.NewReader(src)
        typ, err := message.ReadDataStreamType(br)
        if err != nil {
                return nil, aborted()
        }
        switch {
        case message.IsSubgroupHeaderType(typ):
                hdr, err := message.ReadSubgroupHeader(br, typ)
                if err != nil {
                        return nil, aborted()
                }
                in := &IncomingSubgroupStream{Header: hdr, src: src, br: br, rd: wire.NewStreamReader(br), sess: s}
                return in, nil
        case message.IsFetchHeaderType(typ):
                hdr, err := message.ReadFetchHeader(br)
                if err != nil {
                        return nil, aborted()
                }
                return &IncomingFetchStream{Header: hdr, src: src, br: br, rd: wire.NewStreamReader(br), sess: s}, nil
        case typ == message.PaddingStreamType:
                // §11.5.1: padding streams MUST be silently discarded. CancelRead
                // (STOP_SENDING) abandons the stream and frees its flow control.
                src.CancelRead(uint64(moqt.StreamResetInternalError))
                return nil, ErrPaddingStream
        case message.IsReservedSubgroupHeaderType(typ):
                // §11.4.2: SUBGROUP_ID_MODE 0b11 is reserved.
                src.CancelRead(uint64(moqt.StreamResetInternalError))
                return nil, s.closeProtocolViolation(&message.ReservedSubgroupIDModeError{Type: typ})
        default:
                // §3.4: unknown stream type.
                src.CancelRead(uint64(moqt.StreamResetInternalError))
                return nil, s.closeProtocolViolation(&message.UnknownDataStreamTypeError{Type: typ})
        }
}

package session

import (
        "context"
        "errors"
        "fmt"
        "sync"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// ErrDeliveryTimeout is returned by OutgoingSubgroupStream.WriteObject,
// WriteObjectAt or WriteObjectReceivedAt when the OBJECT_DELIVERY_TIMEOUT has
// been exceeded for the object being written. The stream is reset with
// StreamResetDeliveryTimeout before this error is returned.
//
// Not returned by Write: the raw path cannot see object boundaries or receipt
// times, so it enforces no object timeout at all — see its own doc.
var ErrDeliveryTimeout = errors.New("moqt/session: object delivery timeout exceeded")

// ErrObjectIDNotIncreasing is returned by
// [OutgoingSubgroupStream.WriteObjectAt] when the supplied absolute Object ID
// is not strictly greater than the previous object's on the same stream. The
// §11.4.2 delta encoding is (currentID - previousID - 1), so Object IDs within
// a subgroup MUST strictly increase; WriteObjectAt rejects a violation instead
// of emitting an underflowed delta. Nothing is written and the stream stays
// usable.
var ErrObjectIDNotIncreasing = errors.New("moqt/session: subgroup object ID not strictly increasing")

// writerPool is a sync.Pool for wire.Writer instances to reduce allocations
// in WriteObject calls. The benchmark shows significant allocations from
// creating new Writer instances for each object.
var writerPool = sync.Pool{
        New: func() any {
                return wire.NewWriter(nil)
        },
}

// ---------------------------------------------------------------------------
// OutgoingSubgroupStream
// ---------------------------------------------------------------------------

// OutgoingSubgroupStream is an outbound SUBGROUP_HEADER uni-stream whose
// leading header has already been written. WriteObject appends a framed
// SubgroupObject; Write appends raw body bytes; Close FINs the stream
// cleanly; Cancel resets it.
//
// If delivery timeouts are configured via WithDeliveryTimeouts:
//   - OBJECT_DELIVERY_TIMEOUT: checked before every object is passed to the
//     transport, against that object's own receipt time. If the elapsed time
//     exceeds the timeout the stream is reset with StreamResetDeliveryTimeout
//     and ErrDeliveryTimeout is returned.
//   - SUBGROUP_DELIVERY_TIMEOUT: a timer is started when Close() is called.
//     If the timer fires before the peer acknowledges all data, the stream is
//     reset. Enforced only on a [DeliveryTrackingSendStream].
type OutgoingSubgroupStream struct {
        header message.SubgroupHeader

        dst SendStream

        // pubTimeouts and subTimeouts are the two halves §8 resolves separately:
        // the publisher's Track Property values (possibly overridden by the first
        // object's Object Properties) and the subscriber's Message Parameter
        // values. They are kept apart until the first object arrives because the
        // override applies to the publisher's half ALONE — merging early and
        // overriding the merged value would silently discard a subscriber timeout
        // shorter than the publisher's override.
        pubTimeouts message.DeliveryTimeouts
        subTimeouts message.DeliveryTimeouts

        objectTimeout   time.Duration // resolved; 0 = disabled
        subgroupTimeout time.Duration // resolved; 0 = disabled

        // sawFirstObject gates the §12.1/§12.2 first-object delivery-timeout
        // override: only the first object of the subgroup may override the
        // Track-level timeouts, so the override is applied at most once.
        sawFirstObject bool

        // Encoder state for WriteObjectAt: the running absolute Object ID so each
        // call only has to apply the §11.4.2 delta. encHavePrev is false until the
        // first object is written (its delta is the absolute ID).
        encPrevObject uint64
        encHavePrev   bool

        // Set by [Publication.OpenSubgroup]: onObject is told each written
        // object's Location, paused reports a Forward State of 0 (§11.4.3), ended
        // reports that the publication ended (Done), and onEnd is told when the
        // stream is FINished or reset.
        onObject func(group, object uint64)
        paused   func() bool
        ended    func() bool
        onEnd    func()
}

// WithDeliveryTimeouts returns a shallow copy of s configured with the §8
// delivery timeouts. Zero values disable the corresponding timeout.
//
// The two halves are supplied separately because §8 does not treat them
// symmetrically: "the publisher's value is the Object Property when present on
// the first object of the subgroup, and the Track Property otherwise. If both
// the publisher's value and the subscriber's value are non-zero, the smaller
// of the two is used." The override therefore resolves within the publisher's
// half, and only the result is compared against the subscriber's. A caller
// that pre-merges the two loses that ordering: a first-object override would
// replace a subscriber timeout it was never allowed to outrank.
//
// A publisher with no subscriber-supplied values passes the zero
// DeliveryTimeouts as subscriber, which never wins over a non-zero publisher
// value.
func (s *OutgoingSubgroupStream) WithDeliveryTimeouts(
        publisher, subscriber message.DeliveryTimeouts,
) *OutgoingSubgroupStream {
        cp := *s
        cp.pubTimeouts = publisher
        cp.subTimeouts = subscriber
        // Resolved now so a subgroup whose first object carries no override — the
        // common case — enforces the right values from its very first write.
        eff := publisher.Effective(subscriber)
        cp.objectTimeout = eff.Object
        cp.subgroupTimeout = eff.Subgroup
        return &cp
}

// WriteObject serializes obj onto the stream with correct wire framing.
// The hasProperties flag is taken from the stored SubgroupHeader automatically.
//
// OBJECT_DELIVERY_TIMEOUT is measured from the moment this call is made, which
// is the correct §8 reading for an original publisher handing over an object as
// it is produced. A relay — which received the object earlier, and may have
// spent the interval blocked on some other subscriber — must use
// [OutgoingSubgroupStream.WriteObjectReceivedAt] instead, or the object's age
// is measured from the wrong end.
func (s *OutgoingSubgroupStream) WriteObject(obj *message.SubgroupObject) error {
        return s.WriteObjectReceivedAt(time.Now(), obj)
}

// WriteObjectReceivedAt is [OutgoingSubgroupStream.WriteObject] with the
// object's §8 receipt time supplied by the caller: "the time at which the last
// header byte of every object has been either received from the upstream
// subscription, or provided by the original publisher application".
//
// The clock is per object, not per stream. An object that reaches the transport
// promptly passes however long the stream has already been open, and one that
// queued behind a blocked write fails however new the stream is — which is the
// whole point, since a stream-lifetime cap would reset healthy subscribers for
// no reason other than having stayed subscribed.
func (s *OutgoingSubgroupStream) WriteObjectReceivedAt(
        receivedAt time.Time,
        obj *message.SubgroupObject,
) error {
        // §12.1/§12.2: the first object of a subgroup may carry
        // OBJECT/SUBGROUP_DELIVERY_TIMEOUT as Object Properties that override the
        // Track-level values for this subgroup; the same properties on any later
        // object "is ignored". The override applies to the publisher's half only,
        // and the subscriber's values are compared against the result — see
        // WithDeliveryTimeouts. Resolved before the timeout check so the overridden
        // value takes effect immediately, and before Close reads subgroupTimeout.
        //
        // "First object of the subgroup", not "first object on this stream": the two
        // diverge whenever a stream does not begin at the subgroup's start — a
        // subscriber that joined mid-subgroup, or a §11.4.3 gap-reopen. §11.4.2's
        // FIRST_OBJECT bit is exactly that distinction, and ReplayingSubgroup is its
        // inverse, so a replay stream must not treat whatever object it happens to
        // start with as carrying the subgroup's override.
        if !s.sawFirstObject && s.header.Properties && !s.header.ReplayingSubgroup {
                eff := s.pubTimeouts.ApplyObjectProperties(obj.Properties).Effective(s.subTimeouts)
                s.objectTimeout = eff.Object
                s.subgroupTimeout = eff.Subgroup
        }
        s.sawFirstObject = true

        // §10.12: Done reset the stream before PUBLISH_DONE.
        if s.ended != nil && s.ended() {
                return ErrPublicationEnded
        }
        // §5.1: no Objects while the Forward State is 0; §11.4.3: reset.
        if s.paused != nil && s.paused() {
                s.Cancel(moqt.StreamResetCancelled)
                return ErrForwardPaused
        }
        if err := s.checkObjectTimeout(receivedAt); err != nil {
                return err
        }
        wr, _ := writerPool.Get().(*wire.Writer)
        wr.Reset()
        obj.Append(wr, s.header.Properties)
        _, err := s.dst.Write(wr.Bytes())
        writerPool.Put(wr)
        if err != nil {
                return err
        }
        // §11.4.2: delta + 1 after the first.
        objectID := obj.ObjectIDDelta
        if s.encHavePrev {
                objectID = s.encPrevObject + obj.ObjectIDDelta + 1
        }
        s.encPrevObject, s.encHavePrev = objectID, true
        if s.onObject != nil {
                s.onObject(s.header.GroupID, objectID)
        }
        return nil
}

// WriteObjectAt writes obj with its §11.4.2 ObjectIDDelta computed from the
// absolute objectID and the stream's running previous Object ID — the encoding
// mirror of [IncomingSubgroupStream.ReadDecoded]. The caller supplies absolute
// Object IDs (the way applications think about them); obj.ObjectIDDelta is
// ignored and overwritten. For the first object on the stream the delta is the
// absolute ID; for each later object it is (objectID - previousID - 1).
//
// Object IDs within a subgroup MUST strictly increase (the delta would
// otherwise underflow). If objectID is not greater than the previous object's,
// WriteObjectAt writes nothing, leaves the stream usable, and returns
// [ErrObjectIDNotIncreasing]. OBJECT_DELIVERY_TIMEOUT is enforced exactly as in
// [OutgoingSubgroupStream.WriteObject].
//
// Use the lower-level [OutgoingSubgroupStream.WriteObject] when you want to set
// ObjectIDDelta yourself.
func (s *OutgoingSubgroupStream) WriteObjectAt(objectID uint64, obj *message.SubgroupObject) error {
        if s.encHavePrev {
                if objectID <= s.encPrevObject {
                        return fmt.Errorf("%w: object ID %d not greater than previous %d",
                                ErrObjectIDNotIncreasing, objectID, s.encPrevObject)
                }
                obj.ObjectIDDelta = objectID - s.encPrevObject - 1
        } else {
                obj.ObjectIDDelta = objectID
        }
        return s.WriteObject(obj)
}

// Write appends raw body bytes after the previously-written header. Prefer
// WriteObject for correctly-framed object access; Write is an escape hatch
// for callers that manage framing themselves.
//
// OBJECT_DELIVERY_TIMEOUT is NOT enforced here. §8 measures it per object,
// from the moment that object was received, and a caller that manages its own
// framing is the only party that knows where one object ends and the next
// begins — bytes handed to Write carry no such boundary. A caller that wants
// the timeout enforced should use
// [OutgoingSubgroupStream.WriteObjectReceivedAt], which has both facts.
// SUBGROUP_DELIVERY_TIMEOUT still applies, since Close enforces it.
func (s *OutgoingSubgroupStream) Write(p []byte) (int, error) {
        return s.dst.Write(p)
}

// checkObjectTimeout enforces OBJECT_DELIVERY_TIMEOUT (§8) against one
// object's receipt time, resetting the stream with DELIVERY_TIMEOUT.
func (s *OutgoingSubgroupStream) checkObjectTimeout(receivedAt time.Time) error {
        if s.objectTimeout <= 0 {
                return nil
        }
        elapsed := time.Since(receivedAt)
        if elapsed > s.objectTimeout {
                s.Cancel(moqt.StreamResetDeliveryTimeout)
                return fmt.Errorf("%w (elapsed %s, limit %s)",
                        ErrDeliveryTimeout, elapsed, s.objectTimeout)
        }
        return nil
}

// Close FINs the send side cleanly. Callers must have no concurrent Writes
// in flight.
//
// If SUBGROUP_DELIVERY_TIMEOUT is set and the stream is a
// [DeliveryTrackingSendStream], Close starts a goroutine that resets the
// stream if the peer has not acknowledged all data within the timeout (§8).
func (s *OutgoingSubgroupStream) Close() error {
        err := s.dst.Close()
        if s.onEnd != nil {
                s.onEnd()
        }
        tracked, ok := s.dst.(DeliveryTrackingSendStream)
        if s.subgroupTimeout > 0 && ok {
                finished := tracked.Finished()
                timeout := s.subgroupTimeout
                dst := s.dst
                go func() {
                        timer := time.NewTimer(timeout)
                        defer timer.Stop()
                        select {
                        case <-finished:
                                // All data acknowledged (or stream already reset) — nothing to do.
                        case <-timer.C:
                                // Timer fired before ACK: reset the stream per §8.
                                dst.CancelWrite(uint64(moqt.StreamResetDeliveryTimeout))
                        }
                }()
        }
        return err
}

// Cancel resets the stream with the given application code (§3.3.4).
func (s *OutgoingSubgroupStream) Cancel(code moqt.StreamResetCode) {
        s.dst.CancelWrite(uint64(code))
        if s.onEnd != nil {
                s.onEnd()
        }
}

// SetSendPriority forwards the composite §7.2 scheduling key to the underlying
// transport when it supports per-stream prioritisation (i.e. implements
// [PrioritizedSendStream]). Adapters that don't satisfy the interface
// silently no-op. See [StreamPriority] and [PrioritizedSendStream] for the
// full contract.
func (s *OutgoingSubgroupStream) SetSendPriority(priority StreamPriority) {
        if p, ok := s.dst.(PrioritizedSendStream); ok {
                p.SetSendPriority(priority)
        }
}

// MarkReliable marks the bytes written to this subgroup stream so far (the
// SUBGROUP_HEADER plus any objects) as reliably delivered even if the stream is
// later reset, when the transport supports the RESET_STREAM_AT extension (i.e.
// the underlying stream implements [ReliableResetStream]). This implements the
// §11.4.3 guidance that a reset data stream's reliable_size should cover at
// least the header. It is a no-op when the transport lacks the extension.
func (s *OutgoingSubgroupStream) MarkReliable() {
        if r, ok := s.dst.(ReliableResetStream); ok {
                r.SetReliableBoundary()
        }
}

// ---------------------------------------------------------------------------
// OutgoingFetchStream
// ---------------------------------------------------------------------------

// OutgoingFetchStream is an outbound FETCH_HEADER uni-stream whose leading
// header has already been written. WriteObject appends a framed FetchObject;
// Write appends raw body bytes; Close FINs the stream cleanly; Cancel resets
// it. Fetch streams do not carry delivery timeouts.
type OutgoingFetchStream struct {
        dst SendStream
}

// WriteObject serializes obj onto the stream with correct wire framing.
func (s *OutgoingFetchStream) WriteObject(obj *message.FetchObject) error {
        wr, _ := writerPool.Get().(*wire.Writer)
        wr.Reset()
        obj.Append(wr)
        _, err := s.dst.Write(wr.Bytes())
        writerPool.Put(wr)
        return err
}

// Write appends raw body bytes after the previously-written header. Prefer
// WriteObject for correctly-framed object access.
func (s *OutgoingFetchStream) Write(p []byte) (int, error) { return s.dst.Write(p) }

// Close FINs the send side cleanly.
func (s *OutgoingFetchStream) Close() error { return s.dst.Close() }

// Cancel resets the stream with the given application code (§3.3.4).
func (s *OutgoingFetchStream) Cancel(code moqt.StreamResetCode) {
        s.dst.CancelWrite(uint64(code))
}

// ---------------------------------------------------------------------------
// Session method — open outbound data streams
// ---------------------------------------------------------------------------

// OpenSubgroup opens an outbound SUBGROUP_HEADER uni-stream (§11.4.2),
// writes the full header (Type, Track Alias, Group ID, optional Subgroup ID,
// optional Publisher Priority), and returns the body writer. The caller MUST
// Close to FIN the stream once all objects have been written, or Cancel to
// reset.
func (s *Session) OpenSubgroup(h message.SubgroupHeader) (*OutgoingSubgroupStream, error) {
        return s.OpenSubgroupContext(context.Background(), h)
}

// OpenSubgroupContext is [Session.OpenSubgroup] with a cancellation bound on
// the header write: writing the SUBGROUP_HEADER blocks on the receiver's
// flow control, so a peer that stops reading can wedge the caller
// indefinitely. Cancelling ctx resets the nascent stream and unblocks the
// write. ctx does not govern the returned stream's later writes — bound
// those separately (e.g. a context.AfterFunc calling Cancel).
func (s *Session) OpenSubgroupContext(
        ctx context.Context,
        h message.SubgroupHeader,
) (*OutgoingSubgroupStream, error) {
        sg, reset, err := s.openSubgroup(ctx, h, false)
        if err != nil {
                return nil, err
        }
        if reset {
                // ctx was cancelled just after the header write went through, and
                // the stream is reset.
                return nil, fmt.Errorf("moqt/session: write SUBGROUP_HEADER: %w", ctx.Err())
        }
        return sg, nil
}

// openSubgroup opens a subgroup stream and writes its header, which
// cancelling ctx interrupts by resetting the stream, first marking what was
// written reliable when markReliable is set (§11.4.3). It returns the stream
// whenever the whole header was written, since the peer can then attribute
// the stream to its track, with reset reporting that ctx reset it just after.
func (s *Session) openSubgroup(
        ctx context.Context,
        h message.SubgroupHeader,
        markReliable bool,
) (sg *OutgoingSubgroupStream, reset bool, err error) {
        dst, err := s.conn.OpenUniStream()
        if err != nil {
                return nil, false, err
        }
        stop := context.AfterFunc(ctx, func() {
                if r, ok := dst.(ReliableResetStream); ok && markReliable {
                        r.SetReliableBoundary()
                }
                dst.CancelWrite(uint64(moqt.StreamResetCancelled))
        })
        if err := message.WriteSubgroupHeader(dst, h); err != nil {
                stop()
                dst.CancelWrite(uint64(moqt.StreamResetInternalError))
                if ctx.Err() != nil {
                        return nil, false, fmt.Errorf("moqt/session: write SUBGROUP_HEADER: %w", ctx.Err())
                }
                return nil, false, fmt.Errorf("moqt/session: write SUBGROUP_HEADER: %w", err)
        }
        return &OutgoingSubgroupStream{header: h, dst: dst}, !stop(), nil
}

package session

import (
        "context"
        "errors"
        "sync"

        "github.com/floatdrop/moq-go/pkg/moqt"
)

// SubgroupHandler handles one inbound SUBGROUP_HEADER stream that a [Demux]
// routed to it by §11.1 Track Alias. It is invoked synchronously — usually by
// [Demux.Run], but also by [Demux.HandleTrack], which hands over any streams
// parked before the handler existed on the goroutine that registers it. Those
// two can therefore overlap: a handler holding per-track state of its own must
// synchronise it. Spawn a goroutine inside the handler when streams must be
// processed concurrently (see [Demux.Run]).
type SubgroupHandler func(*IncomingSubgroupStream)

// FetchHandler handles one inbound FETCH_HEADER stream that a [Demux] routed to
// it by Request ID (§10.1). Invoked synchronously by [Demux.Run].
type FetchHandler func(*IncomingFetchStream)

// Demux routes the data streams accepted from a [Session] to per-track and
// per-request handlers, replacing the hand-rolled "AcceptDataStream loop +
// type-switch + Track-Alias match" a subscriber otherwise writes.
//
// Subgroup streams are dispatched by their §11.1 Track Alias — the value a
// subscriber gets from [Subscription.TrackAlias]; FETCH streams by their §11.4.4
// Request ID — the ID the subscriber's FETCH was assigned. A FETCH stream with
// no registered handler is passed to the OnUnknown callback; an unmatched
// subgroup stream is parked instead, for the reasons below.
//
// Handlers may be registered or replaced at any time, including while
// [Demux.Run] is executing: a subscriber learns a track's alias only from its
// SUBSCRIBE_OK, which can arrive after Run has started. Registration is
// safe for concurrent use.
//
// A subgroup stream whose Track Alias has no handler yet is PARKED rather than
// passed to OnUnknown, and released to the handler [Demux.HandleTrack]
// registers for that alias. This is not a nicety: a publisher may start
// sending a track's subgroup streams as soon as it has accepted the SUBSCRIBE,
// which can be before SUBSCRIBE_OK has come back and named the alias to
// register under, so the first Groups of a live broadcast routinely arrive
// with nowhere to go. Resetting them loses that media, and against at least
// one CDN it did worse — two streams reset on arrival and the subscription
// then delivered nothing for the rest of the run.
//
// §11.4.2 says exactly what the choice is — "if an endpoint receives a
// subgroup with an unknown Track Alias, it MAY abandon the stream, or choose
// to buffer it for a brief period to handle reordering with the control
// message that establishes the Track Alias" — and abandoning was measured to
// cost whole runs.
//
// "A brief period" is what [parkLimit] bounds: at most that many streams wait
// per alias, past which the oldest is reset, and [Demux.Run] resets whatever
// is still parked when it returns. Without a bound a stream for an alias
// nobody ever resolves would sit open with its flow control withheld for the
// life of the session, where resetting it at least frees the peer.
//
// OnUnknown therefore sees FETCH streams with no registered handler, not
// subgroup streams.
//
// The zero value is not ready for use — construct with [NewDemux].
type Demux struct {
        mu        sync.Mutex
        subgroup  map[uint64]SubgroupHandler // keyed by Track Alias
        fetch     map[uint64]FetchHandler    // keyed by Request ID
        parked    map[uint64][]*IncomingSubgroupStream
        parkedN   int                 // total across parked, kept in step with it
        retired   map[uint64]struct{} // aliases registered and then unregistered
        onUnknown func(DataStream)
}

// parkLimit is how many subgroup streams may wait for one Track Alias at once,
// and parkTotalLimit how many may wait across all of them. The window being
// covered is a single control-message round trip, so a few Groups' worth is
// the right order; past either bound a stream is reset, which is §11.4.2's
// other option ("MAY abandon the stream").
//
// Both are counts, where §11.4.2 says "for a brief period" — time. Nothing
// here evicts on a timer: with fewer than parkLimit streams for an alias that
// is never claimed, they wait until [Demux.Run] returns. What the counts bound
// is how much can be pinned open at once, which is the part that matters for
// the deadlock §11.4.2 warns about.
//
// parkTotalLimit exists because per-alias bounding alone is not a bound: a
// peer opening subgroup streams for many bogus aliases would park parkLimit of
// each. A parked stream is header-parsed and then unread, so its body sits in
// the transport's receive buffers and consumes the CONNECTION-level window —
// and §11.4.2 continues, past the sentence quoted on [Demux]: "To prevent
// deadlocks, endpoints MUST allocate connection flow control to the control
// streams before allocating it to any data streams. Otherwise, a receiver
// might wait for a control message containing a Track Alias to release flow
// control, while the sender waits for flow control to send the message." That
// MUST binds the transport adapter, which is below this layer and does not
// currently enforce it; parkTotalLimit is what keeps Demux from being the
// thing that walks into it.
const (
        parkLimit      = 8
        parkTotalLimit = 32
)

// NewDemux returns an empty Demux ready for handler registration.
func NewDemux() *Demux {
        return &Demux{
                subgroup: make(map[uint64]SubgroupHandler),
                fetch:    make(map[uint64]FetchHandler),
                parked:   make(map[uint64][]*IncomingSubgroupStream),
                retired:  make(map[uint64]struct{}),
        }
}

// HandleTrack registers h for inbound subgroup streams whose Track Alias is
// alias — typically the [Subscription.TrackAlias] of a subscription this side
// opened. A nil h unregisters alias; registering an alias that already has a
// handler replaces it.
//
// Streams for alias that arrived before this call were parked (see [Demux])
// and are handed to h now, in arrival order, before HandleTrack returns.
// It reports how many — a caller whose output is timing-sensitive wants to
// know that those Groups arrived earlier than they were read.
func (d *Demux) HandleTrack(alias uint64, h SubgroupHandler) (released int) {
        d.mu.Lock()
        if h == nil {
                delete(d.subgroup, alias)
                // §11.1: "Objects can arrive after a subscription has been
                // cancelled. Subscribers SHOULD retain sufficient state to quickly
                // discard these unwanted Objects, rather than treating them as
                // belonging to an unknown Track Alias." Retiring the alias is that
                // state: what is already parked goes now, and anything later for
                // this alias is discarded on arrival rather than parked, since the
                // control message parking waits for is never coming.
                d.retired[alias] = struct{}{}
                stale := d.parked[alias]
                d.dropParkedLocked(alias)
                d.mu.Unlock()
                for _, s := range stale {
                        s.Cancel(moqt.StreamResetCancelled)
                }
                return 0
        }
        delete(d.retired, alias) // re-subscribed under the same alias
        d.subgroup[alias] = h
        held := d.parked[alias]
        d.dropParkedLocked(alias)
        d.mu.Unlock()

        // Outside the lock: a handler may block for the life of the stream.
        for _, s := range held {
                h(s)
        }
        return len(held)
}

// HandleFetch registers h for the inbound FETCH stream answering the FETCH with
// the given Request ID. A nil h unregisters it; re-registering replaces.
func (d *Demux) HandleFetch(requestID uint64, h FetchHandler) {
        d.mu.Lock()
        defer d.mu.Unlock()
        if h == nil {
                delete(d.fetch, requestID)
                return
        }
        d.fetch[requestID] = h
}

// OnUnknown sets the callback invoked for an accepted FETCH stream that
// matches no registered handler. With no callback set (the default, or a nil
// f), such a stream is reset with StreamResetInternalError and dropped so it
// does not leak.
//
// It does NOT see subgroup streams. One whose Track Alias has no handler is
// parked (see [Demux]) rather than reported, because at that point it is far
// more likely to be early than unwanted. One arriving for an alias that was
// registered and then unregistered is reset immediately, per §11.1.
func (d *Demux) OnUnknown(f func(DataStream)) {
        d.mu.Lock()
        defer d.mu.Unlock()
        d.onUnknown = f
}

// parkLocked holds s until [Demux.HandleTrack] claims its alias, resetting the
// oldest once parkLimit is exceeded. Caller holds d.mu.
func (d *Demux) parkLocked(s *IncomingSubgroupStream) bool {
        alias := s.Header.TrackAlias
        if _, retired := d.retired[alias]; retired {
                return false // §11.1: discard promptly, do not buffer.
        }
        if d.parkedN >= parkTotalLimit {
                return false // see parkTotalLimit.
        }
        d.parked[alias] = append(d.parked[alias], s)
        d.parkedN++
        for len(d.parked[alias]) > parkLimit {
                d.parked[alias][0].Cancel(moqt.StreamResetCancelled)
                d.parked[alias] = d.parked[alias][1:]
                d.parkedN--
        }
        return true
}

// dropParkedLocked forgets alias's queue, keeping parkedN in step. It does not
// touch the streams; the caller owns them. Caller holds d.mu.
func (d *Demux) dropParkedLocked(alias uint64) {
        d.parkedN -= len(d.parked[alias])
        delete(d.parked, alias)
}

// discardParked resets everything still waiting for an alias nobody claimed,
// so a stream does not sit open holding its flow control for the life of the
// session. §3.3.4 asks for a relevant code, and a subscriber winding down is
// not an implementation fault: reporting one would have a peer's metrics blame
// this end for a normal end of run.
func (d *Demux) discardParked() {
        d.mu.Lock()
        defer d.mu.Unlock()
        for _, held := range d.parked {
                for _, s := range held {
                        s.Cancel(moqt.StreamResetCancelled)
                }
        }
        clear(d.parked)
        d.parkedN = 0
}

// Run accepts data streams from sess and dispatches each to its registered
// handler until ctx is cancelled or [Session.AcceptDataStream] returns a
// non-padding error, which Run returns. Padding streams (§11.5.1) are skipped.
//
// Dispatch is synchronous: a handler runs to completion before Run accepts the
// next stream, mirroring a hand-written accept loop. A handler that reads a
// long-lived stream therefore blocks the loop, so spawn a goroutine inside the
// handler when streams must be read concurrently.
func (d *Demux) Run(ctx context.Context, sess *Session) error {
        defer d.discardParked()
        for {
                ds, err := sess.AcceptDataStream(ctx)
                if err != nil {
                        if errors.Is(err, ErrPaddingStream) {
                                continue
                        }
                        return err
                }
                d.dispatch(ds)
        }
}

// dispatch routes one accepted data stream to its registered handler, or to the
// unknown path when none matches.
func (d *Demux) dispatch(ds DataStream) {
        switch s := ds.(type) {
        case *IncomingSubgroupStream:
                d.mu.Lock()
                h := d.subgroup[s.Header.TrackAlias]
                if h == nil {
                        // Early, not unwanted — unless the alias is retired or the
                        // park is full, in which case §11.4.2's other option applies.
                        parked := d.parkLocked(s)
                        d.mu.Unlock()
                        if !parked {
                                s.Cancel(moqt.StreamResetCancelled)
                        }
                        return
                }
                d.mu.Unlock()
                h(s)
                return
        case *IncomingFetchStream:
                d.mu.Lock()
                h := d.fetch[s.Header.RequestID]
                d.mu.Unlock()
                if h != nil {
                        h(s)
                        return
                }
        }
        d.unknown(ds)
}

func (d *Demux) unknown(ds DataStream) {
        d.mu.Lock()
        f := d.onUnknown
        d.mu.Unlock()
        if f != nil {
                f(ds)
                return
        }
        ds.Cancel(moqt.StreamResetInternalError)
}

package session

import (
        "context"
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// FetchRequest is a live FETCH operation. It owns the request stream
// (embedded, so Close / reads / message.Marshal work directly on it) plus the
// Request ID follow-up traffic needs, so the caller can send REQUEST_UPDATE via
// [FetchRequest.Update] without holding it separately. The response objects
// arrive on a separate FETCH_HEADER uni-stream (§11.4.4) via
// [Session.AcceptDataStream], not on the embedded stream. It is returned by
// [Session.Fetch].
type FetchRequest struct {
        // requestHandle carries the FETCH request stream — still open for
        // REQUEST_UPDATE follow-ups (Close it to cancel the fetch) — and
        // provides Update.
        requestHandle

        // OK is the parsed FETCH_OK response — EndOfTrack, EndLocation,
        // negotiated Parameters, and TrackProperties.
        OK *message.FetchOK
}

// Fetch opens a FETCH request stream (§10.13) and awaits FETCH_OK or
// REQUEST_ERROR. The session assigns m.RequestID; the caller supplies the
// track name and, in a LOCATION_FILTER parameter, the range (§5.1.2).
//
// On success a [FetchRequest] is returned whose embedded stream stays open (the
// caller may send REQUEST_UPDATE via [FetchRequest.Update]) and whose OK holds
// the parsed FETCH_OK. The publisher will open a FETCH_HEADER uni-stream (§11.4.4)
// carrying the response objects; the caller receives that via AcceptDataStream.
//
// On REQUEST_ERROR the stream is closed and a *RequestRejectedError is
// returned.
func (s *Session) Fetch(ctx context.Context, m *message.Fetch) (*FetchRequest, error) {
        return awaitRequestResponse(ctx, s, m,
                func(stream Stream, ok *message.FetchOK) (*FetchRequest, error) {
                        // §2.5.1: "the subscriber MUST cancel the fetch".
                        if err := s.validateTrackProperties(ok.TrackProperties, "FETCH_OK"); err != nil {
                                cancelRequest(stream)
                                return nil, err
                        }
                        if err := checkFetchOKEnd(m, ok); err != nil {
                                cancelRequest(stream)
                                return nil, s.closeProtocolViolation(err)
                        }
                        // The responder may send neither REQUEST_UPDATE (§10.9) nor
                        // PUBLISH_STATE_NOTIFY (§10.10).
                        return &FetchRequest{
                                Stream:    stream,
                                s:         s,
                                requestID: m.RequestID,
                                OK:        ok,
                        }, nil
                })
}

// checkFetchOKEnd enforces §10.14: an End Location "smaller than the Start
// Location" is a PROTOCOL_VIOLATION.
//
// A relative Start is not known here, but FETCH_OK's End never passes the
// Largest Object (§10.14). Next Object and a relative StartGroup of 0 start
// past it, so any End but {0, 0} (no content yet) precedes them; a larger
// relative StartGroup starts at or before it.
func checkFetchOKEnd(m *message.Fetch, ok *message.FetchOK) error {
        // Without a filter the range starts at {0, 0}. Parsed into a local value
        // to avoid an allocation per FETCH.
        p, found := m.Parameters.Find(message.ParamLocationFilter)
        if !found {
                return nil
        }
        var filter message.LocationFilter
        _ = filter.Parse(wire.NewReader(p.Bytes))
        f := &filter
        switch {
        case f.NextObject(), f.RelativeStart() && f.StartGroup == 0:
                if ok.EndLocation != (message.Location{}) {
                        return fmt.Errorf("moqt/session: FETCH_OK End Location %v precedes a FETCH that starts "+
                                "after the Largest Object", ok.EndLocation)
                }
        case f.RelativeStart():
        default:
                if start := f.Start(message.Location{}, false); ok.EndLocation.Less(start) {
                        return fmt.Errorf("moqt/session: FETCH_OK End Location %v precedes the FETCH Start %v",
                                ok.EndLocation, start)
                }
        }
        return nil
}

// FetchResponder is the publisher side of a FETCH (§10.13) this endpoint
// accepted via [Request.AcceptFetch] — the accept-side counterpart of
// [Session.Fetch]. FETCH_OK has already been written on the embedded request
// stream; the response objects are streamed on a separate FETCH_HEADER
// uni-stream (§11.4.4) opened via [FetchResponder.OpenFetchStream], which binds
// this fetch's Request ID automatically. The embedded request stream stays open
// for REQUEST_UPDATE follow-ups.
type FetchResponder struct {
        // Stream is the FETCH request stream, still open for REQUEST_UPDATE
        // follow-ups. Close it to end the fetch.
        Stream

        s         *Session
        requestID uint64
}

// OpenFetchStream opens the outbound FETCH_HEADER uni-stream (§11.4.4) carrying
// this fetch's response objects, with the Request ID bound automatically. The
// caller MUST Close the returned stream to FIN it once all objects are written,
// or Cancel to reset. It is [Session.OpenFetchStream] pre-bound to this fetch.
func (f *FetchResponder) OpenFetchStream() (*OutgoingFetchStream, error) {
        return f.s.OpenFetchStream(message.FetchHeader{RequestID: f.requestID})
}

// AcceptFetch accepts an inbound FETCH (§10.13) and returns a [FetchResponder]
// for streaming the response objects — the accept-side counterpart of
// [Session.Fetch]. r.First MUST be a *message.Fetch.
//
// ok carries the FETCH_OK fields the caller wants to set (EndOfTrack,
// EndLocation, negotiated Parameters, TrackProperties); it may be nil for the
// all-default reply. AcceptFetch writes FETCH_OK and returns a responder whose
// [FetchResponder.OpenFetchStream] is pre-bound to this fetch's Request ID.
func (r *Request) AcceptFetch(ok *message.FetchOK) (*FetchResponder, error) {
        f, isFetch := r.First.(*message.Fetch)
        if !isFetch {
                return nil, fmt.Errorf("moqt/session: AcceptFetch on a %s request", r.First.Type())
        }
        if ok == nil {
                ok = &message.FetchOK{}
        }
        if err := message.Marshal(r.Stream, ok); err != nil {
                return nil, fmt.Errorf("moqt/session: write FETCH_OK: %w", err)
        }
        return &FetchResponder{Stream: r.Stream, s: r.s, requestID: f.RequestID}, nil
}

// OpenFetchStream opens an outbound FETCH_HEADER uni-stream (§11.4.4),
// writes the header (Type + Request ID), and returns the body writer. The
// caller MUST Close to FIN the stream once all fetch objects have been
// written, or Cancel to reset.
func (s *Session) OpenFetchStream(h message.FetchHeader) (*OutgoingFetchStream, error) {
        dst, err := s.conn.OpenUniStream()
        if err != nil {
                return nil, err
        }
        if err := message.WriteFetchHeader(dst, h); err != nil {
                dst.CancelWrite(uint64(moqt.StreamResetInternalError))
                return nil, fmt.Errorf("moqt/session: write FETCH_HEADER: %w", err)
        }
        return &OutgoingFetchStream{dst: dst}, nil
}

package session

import (
        "errors"
        "fmt"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
)

// GoawayReceived returns a channel that is closed when a GOAWAY arrives from
// the peer. After the channel closes, PeerGoaway returns the parsed message.
func (s *Session) GoawayReceived() <-chan struct{} { return s.goawayCh }

// PeerGoaway returns the GOAWAY most recently received from the peer, or nil
// if none has arrived.
func (s *Session) PeerGoaway() *message.Goaway {
        s.mu.Lock()
        defer s.mu.Unlock()
        return s.goawayReceived
}

// OnGoaway registers a callback invoked exactly once when the first GOAWAY
// arrives from the peer, passing the parsed message (whose NewSessionURI and
// Timeout drive client-side session migration per §3.6/§10.4). The handler
// runs in its own goroutine so it must not assume any ordering with other
// session activity, and it may safely block (e.g. to dial a new session and
// re-issue subscriptions) without stalling the control-receive loop.
//
// OnGoaway is level-triggered: if a GOAWAY has already been received when
// OnGoaway is called, the handler fires immediately. Only the most recently
// registered handler is retained, and the at-most-once guarantee is per
// session — a handler registered after the GOAWAY has already fired the
// previously registered one will itself fire (once) on registration.
//
// Passing a nil handler clears any previously registered callback (provided
// it has not yet fired).
func (s *Session) OnGoaway(handler func(*message.Goaway)) {
        s.mu.Lock()
        // If a GOAWAY already arrived and no handler has fired yet, run this one
        // now and mark it fired so handleGoaway won't double-invoke.
        if s.goawayReceived != nil && !s.goawayFired {
                g := s.goawayReceived
                s.goawayFired = true
                s.mu.Unlock()
                if handler != nil {
                        go handler(g)
                }
                return
        }
        s.goawayHandler = handler
        s.mu.Unlock()
}

// SendGoaway sends a GOAWAY on the control stream and transitions the session
// to the draining state. newURI may be empty; timeout is the grace period
// before the local side may forcibly close the session with GOAWAY_TIMEOUT.
// Returns an error if GOAWAY has already been sent, or if the local role is
// client and newURI is non-empty (§10.4: "A client MUST send a zero-length
// New Session URI in any GOAWAY").
func (s *Session) SendGoaway(timeout time.Duration, newURI string) error {
        if s.role == roleClient && newURI != "" {
                return errors.New("moqt/session: client MUST send a zero-length New Session URI in GOAWAY")
        }
        s.mu.Lock()
        if s.goawaySent {
                s.mu.Unlock()
                return errors.New("moqt/session: GOAWAY already sent")
        }
        s.goawaySent = true
        s.mu.Unlock()

        msg := &message.Goaway{
                NewSessionURI: []byte(newURI),
                //nolint:gosec // G115: timeout is non-negative; whole ms fits a varint.
                Timeout: uint64(timeout / time.Millisecond),
        }
        return s.sendControl(msg)
}

// GoawaySent reports whether [Session.SendGoaway] was called: the session is
// draining, and §10.4 says its sender "SHOULD avoid initiating requests unless
// required by migration".
func (s *Session) GoawaySent() bool {
        s.mu.Lock()
        defer s.mu.Unlock()
        return s.goawaySent
}

// handleGoaway records a received GOAWAY and notifies any waiter on
// GoawayReceived. §10.4: a second GOAWAY on the same control stream MUST
// terminate the session with PROTOCOL_VIOLATION.
func (s *Session) handleGoaway(m *message.Goaway) error {
        s.mu.Lock()
        if s.goawayReceived != nil {
                s.mu.Unlock()
                return errors.New("duplicate GOAWAY on control stream")
        }
        if err := s.checkGoawayURI(m); err != nil {
                s.mu.Unlock()
                return err
        }
        s.goawayReceived = m
        // Snapshot the registered handler under the lock and mark it fired so a
        // later OnGoaway call won't re-invoke it. Run it in its own goroutine
        // (outside the lock) so a blocking migration handler can't stall the
        // control-receive loop.
        var handler func(*message.Goaway)
        if s.goawayHandler != nil && !s.goawayFired {
                handler = s.goawayHandler
                s.goawayFired = true
        }
        s.mu.Unlock()
        close(s.goawayCh)
        if handler != nil {
                go handler(m)
        }
        return nil
}

// checkGoawayURI enforces §10.4 "If a server receives a GOAWAY with a
// non-zero New Session URI Length it MUST close the session with a
// PROTOCOL_VIOLATION": a client cannot direct a server to migrate.
func (s *Session) checkGoawayURI(m *message.Goaway) error {
        if s.role == roleServer && len(m.NewSessionURI) > 0 {
                return errors.New("GOAWAY from client carries non-empty URI")
        }
        return nil
}

// RequestGoaways checks the GOAWAYs one request stream carries (§10.4): "The
// endpoint MUST close the session with a PROTOCOL_VIOLATION ... if it receives
// more than one GOAWAY on the control stream or on a single request stream",
// and a server closes on one carrying a New Session URI. The zero value is
// ready for a stream; callers that read a request stream with [message.Parse]
// keep one per stream.
//
// Not safe for concurrent use, which matches the single reader of a request
// stream ([RequestBroker.Serve] and the relay's per-stream readers).
type RequestGoaways struct {
        seen bool
}

// Received checks a GOAWAY read off the stream. On a violation it closes s
// with PROTOCOL_VIOLATION and returns the error.
func (g *RequestGoaways) Received(s *Session, m *message.Goaway) error {
        if g.seen {
                return s.closeProtocolViolation(errors.New("moqt/session: second GOAWAY on a request stream"))
        }
        g.seen = true
        if err := s.checkGoawayURI(m); err != nil {
                return s.closeProtocolViolation(fmt.Errorf("moqt/session: request stream %w", err))
        }
        return nil
}

package session

import (
        "context"
        "errors"
        "fmt"
        "io"

        "golang.org/x/sync/errgroup"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// handshake performs the SETUP exchange (§3.3): each side writes SETUP on its
// own control stream and reads the peer's, in parallel under an errgroup. Both
// directions are bridged to the context with context.AfterFunc, so a peer
// that stalls mid-SETUP cannot block past cancellation.
//
// Data streams that arrive before the control stream are held for
// AcceptDataStream (see acceptControlStream); request streams wait in the
// transport until AcceptRequest.
func (s *Session) handshake(ctx context.Context, options []wire.KVPair) error {
        g, gctx := errgroup.WithContext(ctx)

        var (
                sendStream SendStream
                recvStream ReceiveStream
                peerOpts   []wire.KVPair
        )

        g.Go(func() error {
                stream, err := s.conn.OpenUniStream()
                if err != nil {
                        return fmt.Errorf("open send control: %w", err)
                }
                stop := context.AfterFunc(gctx, func() {
                        stream.CancelWrite(uint64(moqt.StreamResetCancelled))
                })
                defer stop()
                if err := message.Marshal(stream, &message.Setup{Options: options}); err != nil {
                        stream.CancelWrite(uint64(moqt.StreamResetInternalError))
                        if gctx.Err() != nil {
                                return gctx.Err()
                        }
                        return fmt.Errorf("write SETUP: %w", err)
                }
                sendStream = stream
                return nil
        })

        g.Go(func() error {
                stream, err := s.acceptControlStream(gctx)
                if err != nil {
                        return fmt.Errorf("accept control: %w", err)
                }
                stop := context.AfterFunc(gctx, func() {
                        stream.CancelRead(uint64(moqt.StreamResetCancelled))
                })
                defer stop()
                msg, err := message.Parse(stream)
                if err != nil {
                        stream.CancelRead(uint64(moqt.StreamResetInternalError))
                        if gctx.Err() != nil {
                                return gctx.Err()
                        }
                        return fmt.Errorf("read SETUP: %w", err)
                }
                setup, ok := msg.(*message.Setup)
                if !ok {
                        stream.CancelRead(uint64(moqt.StreamResetInternalError))
                        return fmt.Errorf("expected SETUP, got %s", msg.Type())
                }
                recvStream = stream
                peerOpts = setup.Options
                return nil
        })

        if err := g.Wait(); err != nil {
                return err
        }
        // A cancellation racing a fully successful exchange can fire a stale
        // AfterFunc AFTER Marshal/Parse returned but BEFORE the deferred stop()
        // detached it — resetting a stream we are about to adopt as the
        // session's control stream while g.Wait still returns nil. Any stale
        // fire implies ctx is cancelled by now, so failing here closes the
        // window (the caller tears the conn down as on any handshake error).
        if err := ctx.Err(); err != nil {
                return err
        }

        s.sendCtrl = sendStream
        s.recvCtrl = recvStream
        s.peerOptions = peerOpts
        return nil
}

// maxEarlyDataStreams caps how many data streams the handshake holds before
// the peer's control stream arrives; more are refused with EXCESSIVE_LOAD.
const maxEarlyDataStreams = 32

// acceptControlStream returns the peer's control stream. Data streams that
// arrive first are held, with their type bytes replayed, for AcceptDataStream
// (§3.3: "the data SHOULD be buffered"). A padding stream is discarded and one
// reset before its type is skipped; one FINed before its type fails the
// handshake. The first other stream is returned, with its leading bytes
// replayed, for the SETUP parse to judge.
func (s *Session) acceptControlStream(ctx context.Context) (ReceiveStream, error) {
        for {
                stream, err := s.conn.AcceptUniStream(ctx)
                if err != nil {
                        return nil, err
                }
                stop := context.AfterFunc(ctx, func() { stream.CancelRead(uint64(moqt.StreamResetCancelled)) })
                rec := &recordingByteReader{r: stream}
                typ, err := wire.ReadVarint(rec)
                stop()
                switch {
                case errors.Is(err, io.EOF), errors.Is(err, io.ErrUnexpectedEOF):
                        // FIN before a whole type: no valid stream of any kind (§3.3).
                        return nil, fmt.Errorf("read stream type: %w", err)
                case err != nil:
                        // Reset before its type: skipped, as after setup (§11.4.1).
                        stream.CancelRead(uint64(moqt.StreamResetInternalError))
                        if ctx.Err() != nil {
                                return nil, ctx.Err()
                        }
                        continue
                case typ == message.PaddingStreamType:
                        // §11.5.1.
                        stream.CancelRead(uint64(moqt.StreamResetInternalError))
                        continue
                }
                replayed := &prefixedStream{ReceiveStream: stream, prefix: rec.read}
                if !isDataStreamType(typ) {
                        return replayed, nil
                }
                s.earlyMu.Lock()
                held := len(s.earlyData) < maxEarlyDataStreams
                if held {
                        s.earlyData = append(s.earlyData, replayed)
                }
                s.earlyMu.Unlock()
                if !held {
                        stream.CancelRead(uint64(moqt.StreamResetExcessiveLoad))
                }
        }
}

// nextUniStream returns a data stream held during the handshake, if any, and
// otherwise the transport's next uni stream.
func (s *Session) nextUniStream(ctx context.Context) (ReceiveStream, error) {
        s.earlyMu.Lock()
        if len(s.earlyData) > 0 {
                st := s.earlyData[0]
                s.earlyData = s.earlyData[1:]
                s.earlyMu.Unlock()
                return st, nil
        }
        s.earlyMu.Unlock()
        return s.conn.AcceptUniStream(ctx)
}

func isDataStreamType(typ uint64) bool {
        return message.IsSubgroupHeaderType(typ) || message.IsFetchHeaderType(typ)
}

// recordingByteReader reads r one byte at a time, keeping what it read.
type recordingByteReader struct {
        r    io.Reader
        read []byte
}

func (b *recordingByteReader) ReadByte() (byte, error) {
        var one [1]byte
        if _, err := io.ReadFull(b.r, one[:]); err != nil {
                return 0, err
        }
        b.read = append(b.read, one[0])
        return one[0], nil
}

// prefixedStream is a ReceiveStream whose first bytes were already read:
// Read returns them before reading on.
type prefixedStream struct {
        ReceiveStream

        prefix []byte
}

func (p *prefixedStream) Read(b []byte) (int, error) {
        if len(p.prefix) > 0 {
                n := copy(b, p.prefix)
                p.prefix = p.prefix[n:]
                return n, nil
        }
        return p.ReceiveStream.Read(b)
}

// Package conntest holds transport-test helpers shared between the
// quicconn and wtconn adapter test packages. Keeping the self-signed
// certificate boilerplate here avoids duplicating it across both.
package conntest

import (
        "crypto/ed25519"
        "crypto/rand"
        "crypto/tls"
        "crypto/x509"
        "encoding/pem"
        "math/big"
        "net"
        "testing"
        "time"
)

// TLSConfig builds a one-shot ed25519 self-signed certificate valid for
// localhost / 127.0.0.1 and returns a *tls.Config advertising the given
// ALPN protocols. ed25519 key generation is orders of magnitude faster
// than RSA, which matters when the test runs under -race -count=N.
func TLSConfig(t *testing.T, nextProtos ...string) *tls.Config {
        t.Helper()
        pub, priv, err := ed25519.GenerateKey(rand.Reader)
        if err != nil {
                t.Fatalf("ed25519.GenerateKey: %v", err)
        }
        template := x509.Certificate{
                SerialNumber: big.NewInt(1),
                DNSNames:     []string{"localhost"},
                IPAddresses:  []net.IP{net.IPv4(127, 0, 0, 1)},
                NotBefore:    time.Now().Add(-time.Hour),
                NotAfter:     time.Now().Add(time.Hour),
        }
        certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, pub, priv)
        if err != nil {
                t.Fatalf("x509.CreateCertificate: %v", err)
        }
        keyDER, err := x509.MarshalPKCS8PrivateKey(priv)
        if err != nil {
                t.Fatalf("MarshalPKCS8PrivateKey: %v", err)
        }
        certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
        keyPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER})
        cert, err := tls.X509KeyPair(certPEM, keyPEM)
        if err != nil {
                t.Fatalf("X509KeyPair: %v", err)
        }
        return &tls.Config{
                Certificates: []tls.Certificate{cert},
                NextProtos:   nextProtos,
        }
}

// SendDatagramUntilReceived sends payload repeatedly until recv yields a
// datagram or the deadline passes, and returns what arrived.
//
// Datagrams are unreliable by definition (§11.3), so a test that sent once and
// asserted receipt would be asserting something the transport never promised —
// and would fail occasionally against a correct implementation, which is the
// worst kind of test. Retrying removes the loss lottery without weakening
// anything: an adapter that never delivers still fails, just at the deadline
// rather than on the first drop.
//
// It lives here because both transport adapters need it and neither can import
// the other's test package.
func SendDatagramUntilReceived(
        t *testing.T,
        send func([]byte) error,
        recv <-chan []byte,
        payload []byte,
) []byte {
        t.Helper()
        deadline := time.Now().Add(10 * time.Second)
        for {
                if err := send(payload); err != nil {
                        t.Fatalf("SendDatagram: %v", err)
                }
                select {
                case got := <-recv:
                        return got
                case <-time.After(50 * time.Millisecond):
                }
                if time.Now().After(deadline) {
                        t.Fatal("no datagram arrived; the adapter never delivered one")
                }
        }
}

package conntest

import (
        "context"
        "errors"
        "fmt"
        "io"
        "testing"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
)

// Suite describes one [session.Conn] implementation for [RunSuite].
type Suite struct {
        // NewPair returns two connected endpoints, cleaned up via t.Cleanup.
        //
        // bidiLimit, when > 0, caps how many bidirectional streams the CLIENT may
        // open before the transport reports the peer's limit exhausted. How a
        // transport expresses that is its own business — quic-go takes
        // MaxIncomingStreams on the server, sessiontest caps the opener's credit
        // directly — but the observable contract is the same. Implementations
        // with SupportsBidiLimit false may ignore the argument.
        NewPair func(t *testing.T, bidiLimit int64) (client, server session.Conn)

        // SupportsBidiLimit reports whether NewPair can honour bidiLimit. When
        // false, the ErrNoStreamCredit subtest skips with a reason rather than
        // silently passing — see the note on that subtest.
        SupportsBidiLimit bool
}

// RunSuite drives the behaviour every [session.Conn] adapter must implement
// identically, and is the reason it exists as a shared suite rather than as
// per-adapter tests.
//
// The session layer never imports a QUIC library: it is written against Conn
// and Stream, so every guarantee it relies on is a guarantee some adapter has
// to make good on. Three do (quicconn, wtconn, sessiontest), and only one of
// them — the in-process one — is exercised by `go test`. The other two carry
// the semantics that matter in production and are otherwise covered solely by
// the interop jobs. That is how a stale flag in entrypoint-relay.sh once broke
// the WebTransport path with the whole unit suite green.
//
// So the rule in CLAUDE.md — transport behaviour added to the interface must
// land in all three adapters — is enforced here mechanically instead of by
// review: add a subtest, and every adapter is held to it at once.
//
// What is pinned, and who depends on it:
//
//   - Conn.Context ends when the connection does. The relay's per-session
//     handler goroutines hang off it.
//   - A send stream's Context ends once it is closed. It does not track
//     acknowledgement: quic-go cancels it when Close queues the FIN.
//   - CancelWrite unblocks the peer's Read rather than leaving it parked.
//   - A send stream's Context also ends on the peer's STOP_SENDING (§3.3.3).
//   - OpenStream reports an exhausted peer limit as ErrNoStreamCredit. This
//     one is a documented MUST on the interface, and PUBLISH_SKIPPED (§10.21)
//     is built on it: the relay reacts to the sentinel instead of blocking.
//     An adapter returning the raw transport error would make the relay hang
//     where the spec says to send PUBLISH_SKIPPED.
//
// One constraint the subtests are written to: nothing here may assume the
// transport buffers a write the peer has not accepted yet. Real QUIC does, so
// a write-then-accept sequence passes on both real adapters — and deadlocks on
// sessiontest, whose streams are synchronous io.Pipes. Draining concurrently
// keeps the suite testing the Conn contract rather than each transport's
// buffering, which the contract says nothing about.
func RunSuite(t *testing.T, s Suite) {
        t.Helper()

        t.Run("ConnContextEndsOnClose", func(t *testing.T) {
                client, _ := s.NewPair(t, 0)
                select {
                case <-client.Context().Done():
                        t.Fatal("Conn.Context was already done on a live connection")
                default:
                }
                if err := client.CloseWithError(uint64(moqt.SessionNoError), "bye"); err != nil {
                        t.Fatalf("CloseWithError: %v", err)
                }
                awaitDone(client.Context(), t, "Conn.Context after CloseWithError")
        })

        t.Run("BidiStreamContextEndsAfterClose", func(t *testing.T) {
                client, server := s.NewPair(t, 0)
                stream, err := client.OpenStream()
                if err != nil {
                        t.Fatalf("OpenStream: %v", err)
                }
                // The peer drains so the write can complete on every transport.
                // Drain concurrently — see the note on buffering above.
                drained := drainAsync(t, func() (io.Reader, error) { return server.AcceptStream(t.Context()) })
                if _, err := stream.Write([]byte("hello")); err != nil {
                        t.Fatalf("Write: %v", err)
                }
                if err := stream.Close(); err != nil {
                        t.Fatalf("Close: %v", err)
                }
                awaitDrain(t, drained, "hello")
                awaitDone(stream.Context(), t, "bidi SendStream.Context after Close")
        })

        t.Run("UniStreamContextEndsAfterClose", func(t *testing.T) {
                client, server := s.NewPair(t, 0)
                stream, err := client.OpenUniStream()
                if err != nil {
                        t.Fatalf("OpenUniStream: %v", err)
                }
                drained := drainAsync(t, func() (io.Reader, error) {
                        return server.AcceptUniStream(t.Context())
                })
                if _, err := stream.Write([]byte("hello")); err != nil {
                        t.Fatalf("Write: %v", err)
                }
                if err := stream.Close(); err != nil {
                        t.Fatalf("Close: %v", err)
                }
                awaitDrain(t, drained, "hello")
                awaitDone(stream.Context(), t, "uni SendStream.Context after Close")
        })

        t.Run("CancelWriteUnblocksPeerRead", func(t *testing.T) {
                client, server := s.NewPair(t, 0)
                stream, err := client.OpenStream()
                if err != nil {
                        t.Fatalf("OpenStream: %v", err)
                }
                // Accept and read the first byte on another goroutine: a transport
                // that does not surface a stream before it carries data needs the
                // write to happen, and one with synchronous streams needs the read to
                // happen for the write to return. Then park that goroutine in a second
                // Read, which is what CancelWrite has to wake.
                readErr := make(chan error, 1)
                accepted := make(chan struct{})
                go func() {
                        peer, err := server.AcceptStream(t.Context())
                        if err != nil {
                                readErr <- fmt.Errorf("accept: %w", err)
                                return
                        }
                        if _, err := io.ReadFull(peer, make([]byte, 1)); err != nil {
                                readErr <- fmt.Errorf("first read: %w", err)
                                return
                        }
                        close(accepted)
                        _, err = peer.Read(make([]byte, 1))
                        readErr <- err
                }()
                if _, err := stream.Write([]byte("x")); err != nil {
                        t.Fatalf("Write: %v", err)
                }
                select {
                case <-accepted:
                case err := <-readErr:
                        t.Fatalf("peer never parked in Read: %v", err)
                case <-time.After(5 * time.Second):
                        t.Fatal("peer never received the first byte")
                }
                // The peer is now parked in Read. A reset must wake it with an error
                // rather than an EOF: EOF would read as a clean end of data.
                stream.CancelWrite(uint64(moqt.StreamResetInternalError))
                select {
                case err := <-readErr:
                        if err == nil {
                                t.Fatal("peer Read returned success after the writer reset the stream")
                        }
                        if errors.Is(err, io.EOF) {
                                t.Errorf("peer Read saw io.EOF after a reset, want a stream error: "+
                                        "a reset must not be indistinguishable from a clean FIN (got %v)", err)
                        }
                case <-time.After(5 * time.Second):
                        t.Fatal("CancelWrite did not unblock the peer's Read")
                }
        })

        runStopSendingSubtests(t, s)

        t.Run("OpenStreamReportsNoStreamCredit", func(t *testing.T) {
                if !s.SupportsBidiLimit {
                        // Not a silent pass: this transport cannot be made to exhaust its
                        // own limit from a test. webtransport-go negotiates WebTransport
                        // stream limits over capsules and exposes no knob to lower them,
                        // so its ErrNoStreamCredit mapping is covered only by inspection
                        // and by the interop jobs.
                        t.Skip("transport cannot impose a bidi-stream limit in-test")
                }
                const limit = 2
                client, _ := s.NewPair(t, limit)

                for i := range limit {
                        if _, err := client.OpenStream(); err != nil {
                                t.Fatalf("OpenStream #%d within the limit: %v", i+1, err)
                        }
                }
                _, err := client.OpenStream()
                if err == nil {
                        t.Fatal("OpenStream past the peer's limit succeeded; the limit was not applied")
                }
                if !errors.Is(err, session.ErrNoStreamCredit) {
                        t.Errorf("OpenStream past the peer's limit = %v, want session.ErrNoStreamCredit — "+
                                "the adapter must map its transport's stream-limit error onto the sentinel, "+
                                "or PUBLISH_SKIPPED (§10.21) cannot detect the condition", err)
                }
        })
}

// runStopSendingSubtests pins that a send stream's Context ends when the peer
// sends STOP_SENDING, the only cancel signal after a requester's FIN (§3.3.3).
func runStopSendingSubtests(t *testing.T, s Suite) {
        t.Helper()
        for _, tc := range []struct {
                name string
                open func(session.Conn) (session.SendStream, error)
                peer func(ctx context.Context, c session.Conn) (session.ReceiveStream, error)
        }{
                {"BidiStreamContextEndsOnPeerStopSending",
                        func(c session.Conn) (session.SendStream, error) { return c.OpenStream() },
                        func(ctx context.Context, c session.Conn) (session.ReceiveStream, error) { return c.AcceptStream(ctx) }},
                {"UniStreamContextEndsOnPeerStopSending",
                        func(c session.Conn) (session.SendStream, error) { return c.OpenUniStream() },
                        func(ctx context.Context, c session.Conn) (session.ReceiveStream, error) {
                                return c.AcceptUniStream(ctx)
                        }},
        } {
                t.Run(tc.name, func(t *testing.T) {
                        client, server := s.NewPair(t, 0)
                        stream, err := tc.open(client)
                        if err != nil {
                                t.Fatalf("open: %v", err)
                        }
                        // Accept and read the first byte concurrently (see the note on
                        // buffering above), then stop reading.
                        stopped := make(chan error, 1)
                        go func() {
                                peer, err := tc.peer(t.Context(), server)
                                if err != nil {
                                        stopped <- fmt.Errorf("accept: %w", err)
                                        return
                                }
                                if _, err := io.ReadFull(peer, make([]byte, 1)); err != nil {
                                        stopped <- fmt.Errorf("read: %w", err)
                                        return
                                }
                                peer.CancelRead(uint64(moqt.StreamResetCancelled))
                                stopped <- nil
                        }()
                        if _, err := stream.Write([]byte("x")); err != nil {
                                t.Fatalf("Write: %v", err)
                        }
                        if err := <-stopped; err != nil {
                                t.Fatal(err)
                        }
                        awaitDone(stream.Context(), t, "SendStream.Context after the peer's STOP_SENDING")
                })
        }
}

// awaitDone fails the test unless ctx is cancelled promptly.
func awaitDone(ctx context.Context, t *testing.T, what string) {
        t.Helper()
        select {
        case <-ctx.Done():
        case <-time.After(5 * time.Second):
                t.Fatalf("%s was never cancelled", what)
        }
}

// drainAsync accepts a stream via accept and reads it to EOF on another
// goroutine, reporting the bytes (or the failure) on the returned channel.
// Concurrency is required, not stylistic: a synchronous-pipe transport blocks
// the writer until someone reads.
//
// It reports through a channel rather than calling t.Fatalf because Fatalf
// from a non-test goroutine does not stop the test.
func drainAsync(t *testing.T, accept func() (io.Reader, error)) <-chan drainResult {
        t.Helper()
        ch := make(chan drainResult, 1)
        go func() {
                r, err := accept()
                if err != nil {
                        ch <- drainResult{err: fmt.Errorf("accept: %w", err)}
                        return
                }
                b, err := io.ReadAll(r)
                ch <- drainResult{data: b, err: err}
        }()
        return ch
}

type drainResult struct {
        data []byte
        err  error
}

// awaitDrain fails the test unless the peer read exactly want.
func awaitDrain(t *testing.T, ch <-chan drainResult, want string) {
        t.Helper()
        select {
        case r := <-ch:
                if r.err != nil {
                        t.Fatalf("draining the peer stream: %v", r.err)
                }
                if string(r.data) != want {
                        t.Fatalf("peer read %q, want %q", r.data, want)
                }
        case <-time.After(5 * time.Second):
                t.Fatal("the peer never received the stream's data")
        }
}

package session

import (
        "context"
        "errors"
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
)

// NamespacePublication is an established PUBLISH_NAMESPACE request (§10.16). It
// embeds the still-open request stream (so Close / writes / message.Marshal work
// directly on it) and carries the peer's REQUEST_OK. The namespace stays
// published until the request is cancelled ([NamespacePublication.Close],
// §6.2, §3.3.3); a FIN does not withdraw it (§3.3.2). NAMESPACE and
// NAMESPACE_DONE answer a SUBSCRIBE_NAMESPACE instead (§10.17, §10.18).
type NamespacePublication struct {
        // Stream is the PUBLISH_NAMESPACE request stream, still open for
        // REQUEST_UPDATE follow-ups (§10.9). [NamespacePublication.Close]
        // withdraws the publication.
        Stream

        // OK is the REQUEST_OK the peer replied with.
        OK *message.RequestOK
}

// Close withdraws the namespace by cancelling the request (§6.2, §3.3.3).
func (p *NamespacePublication) Close() error {
        cancelRequest(p.Stream)
        return nil
}

// NamespaceSubscription is an established SUBSCRIBE_NAMESPACE request (§10.19).
// It embeds the still-open request stream and carries the peer's REQUEST_OK;
// NAMESPACE / NAMESPACE_DONE notifications arrive on the embedded stream. Read
// it with its [NamespaceSubscription.Broker]'s [RequestBroker.Serve], which
// enforces the session-level rules: §10 and §10.2.1, a NAMESPACE_DONE only
// after its NAMESPACE (§10.19), and no PUBLISH_STATE_NOTIFY or REQUEST_UPDATE
// from the publisher (§10.9, §10.10). The NAMESPACE_DONE check follows a
// TRACK_NAMESPACE_PREFIX update from the REQUEST_OK that accepts it (§10.9.2).
// A caller using message.Parse must apply these rules itself.
// [NamespaceSubscription.Update] modifies the subscription (§10.9); Close ends
// it (§6.1, §3.3.3).
type NamespaceSubscription struct {
        // requestHandle carries the SUBSCRIBE_NAMESPACE request stream, still
        // open to receive NAMESPACE / NAMESPACE_DONE notifications.
        requestHandle

        // OK is the REQUEST_OK the peer replied with.
        OK *message.RequestOK
}

// TrackSubscription is an established SUBSCRIBE_TRACKS request (§10.20). It
// embeds the still-open request stream and carries the peer's REQUEST_OK.
// Follow-up PUBLISH_SKIPPED notifications are read via
// [TrackSubscription.ReadPublishSkipped], or its [TrackSubscription.Broker]'s
// [RequestBroker.Serve], never both at once. Both close the session on a
// PUBLISH_STATE_NOTIFY or REQUEST_UPDATE from the publisher (§10.9, §10.10).
// [TrackSubscription.Update] modifies the subscription (§10.9); Close ends it
// (§6.1, §3.3.3).
type TrackSubscription struct {
        // requestHandle carries the SUBSCRIBE_TRACKS request stream, still open
        // to receive PUBLISH_SKIPPED follow-ups.
        requestHandle

        // OK is the REQUEST_OK the peer replied with.
        OK *message.RequestOK
}

// Update sends a REQUEST_UPDATE (§10.9), e.g. a new TRACK_NAMESPACE_PREFIX
// (§10.9.2), through the subscription's [NamespaceSubscription.Broker], whose
// Serve must be running to deliver the answer. Reading the answer directly
// would take a NAMESPACE for it.
func (n *NamespaceSubscription) Update(ctx context.Context, params message.Parameters) (*message.RequestOK, error) {
        return n.Broker().Update(ctx, params)
}

// Update sends a REQUEST_UPDATE (§10.9) through the subscription's
// [TrackSubscription.Broker], whose Serve must be running to deliver the
// answer, as for [NamespaceSubscription.Update].
func (t *TrackSubscription) Update(ctx context.Context, params message.Parameters) (*message.RequestOK, error) {
        return t.Broker().Update(ctx, params)
}

// PublishNamespace opens a PUBLISH_NAMESPACE request stream (§10.16) and
// awaits REQUEST_OK or REQUEST_ERROR. The session assigns m.RequestID; the
// caller supplies Namespace and optional Parameters.
//
// On success a [NamespacePublication] is returned whose embedded stream stays
// open (the caller may send REQUEST_UPDATE on it, §10.9). On REQUEST_ERROR the
// stream is closed and a *RequestRejectedError is returned.
func (s *Session) PublishNamespace(
        ctx context.Context,
        m *message.PublishNamespace,
) (*NamespacePublication, error) {
        return awaitRequestResponse(ctx, s, m,
                func(stream Stream, ok *message.RequestOK) (*NamespacePublication, error) {
                        return &NamespacePublication{Stream: stream, OK: ok}, nil
                })
}

// SubscribeNamespace opens a SUBSCRIBE_NAMESPACE request stream (§10.19) and
// awaits REQUEST_OK or REQUEST_ERROR. The session assigns m.RequestID; the
// caller supplies TrackNamespacePrefix and optional Parameters.
//
// On success a [NamespaceSubscription] is returned whose embedded stream stays
// open (the caller will receive NAMESPACE / NAMESPACE_DONE messages on it). On
// REQUEST_ERROR the stream is closed and a *RequestRejectedError is returned;
// any other first response closes the session with PROTOCOL_VIOLATION (§10.19).
func (s *Session) SubscribeNamespace(
        ctx context.Context,
        m *message.SubscribeNamespace,
) (*NamespaceSubscription, error) {
        return awaitRequestResponse(ctx, s, m,
                func(stream Stream, ok *message.RequestOK) (*NamespaceSubscription, error) {
                        return &NamespaceSubscription{
                                Stream: stream, s: s, requestID: m.RequestID,
                                namespaces: true, nsPrefix: m.TrackNamespacePrefix, OK: ok,
                        }, nil
                })
}

// SubscribeTracks opens a SUBSCRIBE_TRACKS request stream (§10.20) and awaits
// REQUEST_OK or REQUEST_ERROR. The session assigns m.RequestID; the caller
// supplies TrackNamespacePrefix and optional Parameters.
//
// On success a [TrackSubscription] is returned whose embedded stream stays open
// for PUBLISH_SKIPPED follow-ups (read via [TrackSubscription.ReadPublishSkipped]).
// On REQUEST_ERROR the stream is closed and a *RequestRejectedError is returned;
// any other first response closes the session with PROTOCOL_VIOLATION (§10.20).
func (s *Session) SubscribeTracks(ctx context.Context, m *message.SubscribeTracks) (*TrackSubscription, error) {
        return awaitRequestResponse(ctx, s, m,
                func(stream Stream, ok *message.RequestOK) (*TrackSubscription, error) {
                        if err := s.checkTrackPropertyValues(ok.TrackProperties, "SUBSCRIBE_TRACKS_OK"); err != nil {
                                cancelRequest(stream)
                                return nil, err
                        }
                        return &TrackSubscription{Stream: stream, s: s, requestID: m.RequestID, OK: ok}, nil
                })
}

// IncomingNamespacePublication is an accepted inbound PUBLISH_NAMESPACE (§10.16)
// — the receiving side of [Session.PublishNamespace]'s [NamespacePublication],
// returned by [Request.AcceptPublishNamespace]. REQUEST_OK has been sent; the
// announcer's follow-ups (REQUEST_UPDATE, GOAWAY) arrive on the embedded
// stream. Read it with a [Session.NewRequestBroker]'s [RequestBroker.Serve],
// which closes the session on a malformed message (§10) or a GOAWAY violation
// (§10.4). Before Serve, call [RequestBroker.PeerMessages](true, false), since
// PUBLISH_STATE_NOTIFY applies only to subscriptions (§10.10), and
// [RequestBroker.UpdateScope](message.ScopeUpdatePublishNamespace), so an
// update's parameters are checked (§10.2.1) before it is answered; without
// [RequestBroker.HandleUpdates] each REQUEST_UPDATE is declined with
// NOT_SUPPORTED. Cancel the request (CancelRead and CancelWrite, §3.3.3) to
// revoke acceptance (§6.2); Close only FINs this side (§3.3.2).
type IncomingNamespacePublication struct {
        // Stream is the PUBLISH_NAMESPACE request stream, still open to receive
        // the announcer's follow-ups.
        Stream
}

// IncomingNamespaceSubscription is an accepted inbound SUBSCRIBE_NAMESPACE
// (§10.19) — the announcing side of [Session.SubscribeNamespace]'s
// [NamespaceSubscription], returned by [Request.AcceptSubscribeNamespace].
// REQUEST_OK has been sent; the caller announces matching namespaces by writing
// NAMESPACE / NAMESPACE_DONE to the embedded stream (e.g. via message.Marshal).
// Close it to end the subscription.
type IncomingNamespaceSubscription struct {
        // Stream is the SUBSCRIBE_NAMESPACE request stream, still open for
        // NAMESPACE / NAMESPACE_DONE follow-ups. Close it to end the subscription.
        Stream
}

// IncomingTrackSubscription is an accepted inbound SUBSCRIBE_TRACKS (§10.20) —
// the publishing side of [Session.SubscribeTracks]'s [TrackSubscription],
// returned by [Request.AcceptSubscribeTracks]. REQUEST_OK has been sent; the
// publisher forwards matching tracks as PUBLISH requests on new streams (see
// [Session.OpenPublish]) and signals stream exhaustion with
// [IncomingTrackSubscription.WritePublishSkipped] (§6.1 / §10.21). Close it to
// end the subscription.
type IncomingTrackSubscription struct {
        // Stream is the SUBSCRIBE_TRACKS request stream, still open for
        // PUBLISH_SKIPPED follow-ups. Close it to end the subscription.
        Stream
}

// WritePublishSkipped sends a PUBLISH_SKIPPED (§6.1 / §10.21) on the
// SUBSCRIBE_TRACKS stream, telling the subscriber the publisher could not open a
// PUBLISH stream for the named track because it has no available bidirectional
// streams. It is the publisher-side counterpart of
// [TrackSubscription.ReadPublishSkipped].
func (t *IncomingTrackSubscription) WritePublishSkipped(pb *message.PublishSkipped) error {
        return message.Marshal(t.Stream, pb)
}

// acceptNamespaceRequest is the shared accept path of the three namespace
// requests (§10.16 / §10.19 / §10.20): assert the request's first message is
// of type M, reply the all-default REQUEST_OK, and hand the still-open
// stream to wrap. op names the caller for error messages.
func acceptNamespaceRequest[M message.Message, T any](r *Request, op string, wrap func(Stream) T) (T, error) {
        var zero T
        if _, ok := r.First.(M); !ok {
                return zero, fmt.Errorf("moqt/session: %s on a %s request", op, r.First.Type())
        }
        if err := r.Reply(&message.RequestOK{}); err != nil {
                return zero, fmt.Errorf("moqt/session: %s: write REQUEST_OK: %w", op, err)
        }
        return wrap(r.Stream), nil
}

// AcceptPublishNamespace accepts an inbound PUBLISH_NAMESPACE (§10.16), replies
// REQUEST_OK, and returns an [IncomingNamespacePublication] for receiving the
// announcer's follow-ups — the accept-side counterpart of
// [Session.PublishNamespace]. r.First MUST be a
// *message.PublishNamespace.
func (r *Request) AcceptPublishNamespace() (*IncomingNamespacePublication, error) {
        return acceptNamespaceRequest[*message.PublishNamespace](r, "AcceptPublishNamespace",
                func(s Stream) *IncomingNamespacePublication { return &IncomingNamespacePublication{Stream: s} })
}

// AcceptSubscribeNamespace accepts an inbound SUBSCRIBE_NAMESPACE (§10.19),
// replies REQUEST_OK, and returns an [IncomingNamespaceSubscription] for
// announcing matching namespaces via NAMESPACE / NAMESPACE_DONE — the
// accept-side counterpart of [Session.SubscribeNamespace]. r.First MUST be a
// *message.SubscribeNamespace.
func (r *Request) AcceptSubscribeNamespace() (*IncomingNamespaceSubscription, error) {
        return acceptNamespaceRequest[*message.SubscribeNamespace](r, "AcceptSubscribeNamespace",
                func(s Stream) *IncomingNamespaceSubscription { return &IncomingNamespaceSubscription{Stream: s} })
}

// AcceptSubscribeTracks accepts an inbound SUBSCRIBE_TRACKS (§10.20), replies
// REQUEST_OK, and returns an [IncomingTrackSubscription] for forwarding matching
// PUBLISHes and sending PUBLISH_SKIPPED follow-ups — the accept-side counterpart
// of [Session.SubscribeTracks]. r.First MUST be a *message.SubscribeTracks.
func (r *Request) AcceptSubscribeTracks() (*IncomingTrackSubscription, error) {
        return acceptNamespaceRequest[*message.SubscribeTracks](r, "AcceptSubscribeTracks",
                func(s Stream) *IncomingTrackSubscription { return &IncomingTrackSubscription{Stream: s} })
}

// ReadPublishSkipped reads the next follow-up message on this SUBSCRIBE_TRACKS
// response stream and returns it as a PUBLISH_SKIPPED.
//
// This is the subscriber side of §6.1 / §10.21. After the initial REQUEST_OK,
// the publisher sends PUBLISH_SKIPPED on this stream when it cannot open a
// PUBLISH stream for a matching track because it has no available
// bidirectional streams. (Forwarded PUBLISHes themselves arrive on their own
// new bidi streams via [Session.AcceptRequest], not here.) The returned
// message names the track the publisher couldn't push; the caller's sanctioned
// recovery is to issue an explicit SUBSCRIBE for it.
//
// It cannot be combined with [TrackSubscription.Update], whose answer only the
// broker's Serve delivers.
//
// It blocks until a PUBLISH_SKIPPED arrives or the stream ends. A single
// GOAWAY is skipped (§10.4); a caller that would re-issue the request at its
// New Session URI reads with the broker instead, whose Serve hands it over. A
// second GOAWAY, and a PUBLISH_STATE_NOTIFY or REQUEST_UPDATE from the
// publisher (§10.9, §10.10), close the session with PROTOCOL_VIOLATION. Any
// other message is reported as an error, as is the underlying read error (e.g.
// io.EOF when the publisher FINs the SUBSCRIBE_TRACKS stream).
func (t *TrackSubscription) ReadPublishSkipped() (*message.PublishSkipped, error) {
        for {
                m, err := message.Parse(t)
                if errors.Is(err, message.ErrMalformedMessage) {
                        return nil, t.s.closeProtocolViolation(fmt.Errorf("moqt/session: read SUBSCRIBE_TRACKS follow-up: %w", err))
                }
                if err != nil {
                        return nil, fmt.Errorf("moqt/session: read SUBSCRIBE_TRACKS follow-up: %w", err)
                }
                switch m := m.(type) {
                case *message.PublishSkipped:
                        return m, nil
                case *message.Goaway:
                        // §10.4: legal once; the request is not migrated.
                        if err := t.goaways.Received(t.s, m); err != nil {
                                return nil, err
                        }
                case *message.PublishStateNotify, *message.RequestUpdate:
                        // §10.10: PUBLISH_STATE_NOTIFY "applies only to subscriptions";
                        // §10.9: the publisher did not send this request.
                        return nil, t.s.closeProtocolViolation(
                                fmt.Errorf("moqt/session: %s on a SUBSCRIBE_TRACKS stream", m.Type()))
                default:
                        return nil, fmt.Errorf(
                                "moqt/session: unexpected %s on SUBSCRIBE_TRACKS stream, want PUBLISH_SKIPPED",
                                m.Type(),
                        )
                }
        }
}

package session

import (
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// Option configures a session opened via Client or Server. See WithPath,
// WithAuthority, WithImplementation, WithMaxAuthTokenCacheSize,
// WithTokenVerifier, and WithGrease for the available knobs.
type Option func(*config)

// config carries the resolved set of options applied to a single
// Client/Server call. Unexported so callers can only construct it via
// Option helpers.
type config struct {
        setupOptions                  []wire.KVPair
        knownMandatoryTrackProperties map[message.PropertyType]struct{}

        // maxAuthTokenCacheSize is the byte budget for the inbound
        // authorization-token alias cache (§10.2.2). It mirrors the value
        // advertised to the peer via MAX_AUTH_TOKEN_CACHE_SIZE and is captured
        // here by WithMaxAuthTokenCacheSize so open() can size the cache. The
        // default (0) prohibits alias registration per §10.3.1.3.
        maxAuthTokenCacheSize uint64

        // maxRequestUpdates is the per-request-stream limit on unacknowledged
        // inbound REQUEST_UPDATEs (§10.3.1.7). It mirrors the value advertised to
        // the peer via MAX_REQUEST_UPDATES and is captured here by
        // WithMaxRequestUpdates so the session can enforce it on the receive side.
        // The default (0) means REQUEST_UPDATE concurrency is not limited.
        maxRequestUpdates uint64

        // maxFilterRanges is the per-subscription/fetch limit on the total number
        // of Range Filter ranges (§10.3.1.6). It mirrors the value advertised via
        // MAX_FILTER_RANGES and is captured by WithMaxFilterRanges so the receive
        // side can reject over-limit filters with INVALID_FILTER. The default (0)
        // prohibits Range Filters entirely.
        maxFilterRanges uint64

        // setupTokens are the tokens WithSetupToken added to SETUP, in order.
        setupTokens []message.Token

        // tokenVerifier is the optional application policy that turns a resolved
        // (Type, Value) authorization token into an allow/deny decision. nil
        // disables verification (all tokens are accepted by the transport; the
        // application is responsible for any out-of-band checks).
        tokenVerifier TokenVerifier
}

// WithSetupToken adds an AUTHORIZATION TOKEN setup option (§10.3.1.4). Repeat
// it for several tokens; they are sent in order.
//
// Opening the session fails on DELETE, USE_ALIAS or a repeated REGISTER alias
// (§10.2.2). [Session.SetupTokenAliases] reports which REGISTERed aliases fit
// the peer's cache.
func WithSetupToken(t message.Token) Option {
        return func(c *config) {
                c.setupTokens = append(c.setupTokens, t)
                c.setupOptions = append(c.setupOptions, wire.KVPair{
                        Type: uint64(message.SetupOptionAuthorizationToken), ByteVal: t.Bytes(),
                })
        }
}

// WithImplementation sets the MOQT_IMPLEMENTATION SETUP option — a
// free-form identifier of this peer's implementation. Recommended for
// every peer; advisory in spec terms.
func WithImplementation(nameAndVersion string) Option {
        return func(c *config) {
                c.setupOptions = append(c.setupOptions, message.MOQTImplementationOption(nameAndVersion))
        }
}

// WithPath sets the PATH SETUP option (§10.3.1.2). Client-only — using
// this on Server is a protocol violation per the spec.
func WithPath(pathAndQuery string) Option {
        return func(c *config) {
                c.setupOptions = append(c.setupOptions, message.PathOption(pathAndQuery))
        }
}

// WithAuthority sets the AUTHORITY SETUP option (§10.3.1.1). Client-only.
func WithAuthority(authority string) Option {
        return func(c *config) {
                c.setupOptions = append(c.setupOptions, message.AuthorityOption(authority))
        }
}

// WithMaxAuthTokenCacheSize sets MAX_AUTH_TOKEN_CACHE_SIZE — the maximum
// byte size of the per-session authorization-token alias cache (§10.2.2).
//
// The same budget sizes the inbound TokenCache the session uses to process
// AUTHORIZATION_TOKEN aliases on request streams: maxBytes is both advertised
// to the peer in SETUP and used to bound how many REGISTER tokens the peer may
// install. The default (option absent) is 0, which prohibits alias
// registration entirely per §10.3.1.3.
func WithMaxAuthTokenCacheSize(maxBytes uint64) Option {
        return func(c *config) {
                c.maxAuthTokenCacheSize = maxBytes
                c.setupOptions = append(c.setupOptions, message.MaxAuthTokenCacheSizeOption(maxBytes))
        }
}

// WithMaxRequestUpdates sets MAX_REQUEST_UPDATES (§10.3.1.7) — the maximum
// number of unacknowledged REQUEST_UPDATE messages this endpoint is willing to
// receive on any single request stream. The value is both advertised to the
// peer in SETUP and enforced on inbound follow-ups: a REQUEST_UPDATE that
// arrives while the stream already holds max outstanding updates closes the
// session with TOO_MANY_REQUEST_UPDATES.
//
// A REQUEST_UPDATE is outstanding from receipt until this endpoint writes the
// mandated REQUEST_OK/REQUEST_ERROR. The default (option absent, or 0) does not
// limit REQUEST_UPDATE concurrency.
func WithMaxRequestUpdates(maxUpdates uint64) Option {
        return func(c *config) {
                c.maxRequestUpdates = maxUpdates
                c.setupOptions = append(c.setupOptions, message.MaxRequestUpdatesOption(maxUpdates))
        }
}

// WithMaxFilterRanges sets MAX_FILTER_RANGES (§10.3.1.6) — the maximum total
// number of Range Filter ranges this endpoint will accept across all Range
// Filter parameters on a single subscription or fetch (§5.1.4). The value is
// advertised to the peer in SETUP and used to reject over-limit or (when 0)
// any Range Filters with INVALID_FILTER. The default (option absent, or 0)
// prohibits Range Filters entirely.
func WithMaxFilterRanges(maxRanges uint64) Option {
        return func(c *config) {
                c.maxFilterRanges = maxRanges
                c.setupOptions = append(c.setupOptions, message.MaxFilterRangesOption(maxRanges))
        }
}

// WithTokenVerifier installs an application policy that authorizes resolved
// AUTHORIZATION_TOKEN tokens (§10.2.2). After the session resolves a request's
// tokens (handling REGISTER / USE_ALIAS / USE_VALUE / DELETE against the
// inbound cache), it invokes v.VerifyToken for each resolved token so the
// application can validate signatures, expiry, audience, and scope — concerns
// the transport deliberately leaves out (§13.3).
//
// Passing nil (or never calling this option) disables verification: tokens are
// still parsed and aliases are still maintained, but no allow/deny decision is
// made at the transport layer.
func WithTokenVerifier(v TokenVerifier) Option {
        return func(c *config) {
                c.tokenVerifier = v
        }
}

// WithGrease enables GREASE (§14): a random unknown SETUP option is injected
// into the outbound SETUP message to exercise the peer's tolerance of unknown
// values. GREASE values follow the pattern 0x7F * N + 0x9D and are always
// larger than all currently defined SETUP option types, so appending preserves
// the non-decreasing Type order that §1.4.3's Delta Type encoding requires
// (Setup Options are Key-Value-Pairs, not §10.2 Message Parameters).
func WithGrease() Option {
        return func(c *config) {
                c.setupOptions = append(c.setupOptions, message.GreaseSetupOption())
        }
}

// WithKnownMandatoryTrackProperties configures the set of Mandatory Track
// Property types (range 0x4000–0x7FFF per §2.5.1) that this endpoint
// understands. When the session receives Track Properties (in SUBSCRIBE_OK,
// FETCH_OK, or TRACK_STATUS_OK) containing a mandatory property not in this
// set, it returns *ErrUnsupportedMandatoryTrackProperty; [Request.AcceptPublish]
// refuses such a PUBLISH with UNSUPPORTED_EXTENSION.
//
// If this option is never called, or types is empty or nil, no Mandatory Track
// Property is known and every one is refused: an endpoint that does not
// understand one "MUST NOT process or forward that track" (§2.5.1). List the
// types this endpoint understands to accept them.
func WithKnownMandatoryTrackProperties(types map[message.PropertyType]struct{}) Option {
        return func(c *config) {
                c.knownMandatoryTrackProperties = types
        }
}

package session

import (
        "context"
        "errors"
        "fmt"
        "slices"
        "sync"
        "sync/atomic"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
)

// Publication is a live track this side publishes objects on. It owns the
// request stream (embedded, so Close / writes / message.Marshal work directly on
// it) and the Track Alias the session assigned, and it opens subgroup
// uni-streams for the track via [Publication.OpenSubgroup] without the caller
// having to thread the alias around. It is returned both by [Session.Publish]
// (publisher-initiated, the PUBLISH side) and by [Request.AcceptSubscribe]
// (answering an inbound SUBSCRIBE) — in both cases this endpoint is the one
// sending objects.
type Publication struct {
        // requestHandle carries the request stream. [Publication.Done] ends the
        // publication gracefully (PUBLISH_DONE, then FIN); Close cancels it.
        //
        // §10.9: Update is valid on a Publication from [Session.Publish] but not
        // on one from [Request.AcceptSubscribe], where this side did not send
        // the request.
        requestHandle

        alias uint64

        // The subgroups opened via OpenSubgroup, for Done (§10.12). subMu
        // orders each OpenSubgroup against Done, which waits for the opens in
        // flight (opening) and cancels their header writes (endCtx), so every
        // subgroup opened is counted in subgroupCount, and the ones still open
        // (open) are reset before PUBLISH_DONE. A subgroup's onEnd closes over
        // the pointer registered here, which its copies (WithDeliveryTimeouts)
        // share.
        subMu         sync.Mutex
        subgroupCount uint64
        open          map[*OutgoingSubgroupStream]struct{}
        opening       sync.WaitGroup
        endCtx        context.Context
        endCancel     context.CancelFunc

        // paused is the inverse of the §5.1 Forward State.
        paused atomic.Bool

        // largest is the largest Location this side announced in SUBSCRIBE_OK
        // or PUBLISH, or wrote through OpenSubgroup streams since, reported as
        // LARGEST_OBJECT in REQUEST_UPDATE_OK (§10.9.1, §10.2.17).
        largestMu  sync.Mutex
        largest    message.Location
        hasLargest bool

        // ended is latched by the first Done, under subMu, so PUBLISH_DONE is
        // sent once and no subgroup opens or writes after it.
        ended atomic.Bool

        brokerInit sync.Once
}

// ErrForwardPaused is returned while the subscription's Forward State is 0
// (§5.1): [Publication.OpenSubgroup] opens nothing, and a write to an open
// subgroup resets that stream first (§11.4.3). A REQUEST_UPDATE with FORWARD=1
// resumes.
var ErrForwardPaused = errors.New("moqt/session: Forward State is 0; not sending objects")

// ErrPublicationEnded is returned by [Publication.OpenSubgroup], and by the
// WriteObject methods of a subgroup it opened, once the publication has ended
// — by [Publication.Done], or automatically after a declined REQUEST_UPDATE
// (§10.9.1).
var ErrPublicationEnded = errors.New("moqt/session: publication ended (PUBLISH_DONE sent)")

// newPublication builds a Publication whose initial Forward State is the
// establishing message's FORWARD (§5.1), or 1 when omitted (§10.2.18), and
// whose Largest Object starts at the LARGEST_OBJECT this side announced in
// its SUBSCRIBE_OK or PUBLISH (§10.2.17), if any.
func newPublication(
        s *Session,
        stream Stream,
        requestID, alias uint64,
        establishing, announced message.Parameters,
) *Publication {
        // The subscriber may send REQUEST_UPDATE (§10.9) but not
        // PUBLISH_STATE_NOTIFY (§10.10).
        p := &Publication{
                Stream: stream, s: s, requestID: requestID, alias: alias,
                peerUpdate: true, updateScope: message.ScopeUpdateFromSubscriber,
                open: make(map[*OutgoingSubgroupStream]struct{}),
        }
        p.endCtx, p.endCancel = context.WithCancel(context.Background())
        if f, ok := establishing.Find(message.ParamForward); ok {
                p.paused.Store(f.Byte == 0)
        }
        if lo, ok := announced.Find(message.ParamLargestObject); ok {
                p.largest, p.hasLargest = message.Location{Group: lo.Group, Object: lo.Object}, true
        }
        return p
}

// Broker returns the publication's [RequestBroker] (see [requestHandle.Broker])
// with [Publication.ApplyUpdate] installed to decide REQUEST_UPDATEs. A
// declined update ends the subscription with PUBLISH_DONE UPDATE_FAILED
// (§10.9.1). Use [RequestBroker.HandleUpdates] to support more parameters.
func (p *Publication) Broker() *RequestBroker {
        b := p.requestHandle.Broker()
        p.brokerInit.Do(func() {
                b.HandleUpdates(p.ApplyUpdate)
                b.onUpdateFailed = func() {
                        _ = p.Done(moqt.PublishDoneUpdateFailed, "REQUEST_UPDATE declined")
                }
        })
        return b
}

// ApplyUpdate is the publication's built-in REQUEST_UPDATE handling (§10.9).
// FORWARD (§10.2.18) sets the Forward State; SUBSCRIBER_PRIORITY,
// NEW_GROUP_REQUEST and AUTHORIZATION_TOKEN are accepted without action (the
// Serve callback still sees the update). Any other parameter declines the
// whole update with NOT_SUPPORTED.
//
// The REQUEST_UPDATE_OK carries LARGEST_OBJECT (§10.9.1, §10.2.17): the
// larger of the one this side's SUBSCRIBE_OK or PUBLISH reported and the
// largest Object written through [Publication.OpenSubgroup] streams; other
// objects are not seen.
func (p *Publication) ApplyUpdate(upd *message.RequestUpdate) (*message.RequestOK, error) {
        forward, setForward := false, false
        for _, prm := range upd.Parameters {
                if prm.Type == message.ParamForward {
                        // §10.2.18.
                        if prm.Byte > 1 {
                                return nil, p.s.closeProtocolViolation(
                                        fmt.Errorf("moqt/session: FORWARD value %d in REQUEST_UPDATE", prm.Byte))
                        }
                        forward, setForward = prm.Byte == 1, true
                        continue
                }
                if !slices.Contains(acceptedUpdateParams, prm.Type) {
                        return nil, &RequestRejectedError{
                                Code:   moqt.RequestNotSupported,
                                Reason: fmt.Sprintf("REQUEST_UPDATE parameter %#x not supported", uint64(prm.Type)),
                        }
                }
        }
        if setForward {
                p.paused.Store(!forward)
        }
        ok := &message.RequestOK{}
        p.largestMu.Lock()
        if p.hasLargest {
                ok.Parameters = message.Parameters{message.LargestObjectParam(p.largest.Group, p.largest.Object)}
        }
        p.largestMu.Unlock()
        return ok, nil
}

// acceptedUpdateParams are the parameters, besides FORWARD, that
// [Publication.ApplyUpdate] accepts without further action.
var acceptedUpdateParams = []message.ParamID{
        message.ParamSubscriberPriority,
        message.ParamNewGroupRequest,
        message.ParamAuthorizationToken,
}

// noteObject records a written object for LARGEST_OBJECT.
func (p *Publication) noteObject(group, object uint64) {
        loc := message.Location{Group: group, Object: object}
        p.largestMu.Lock()
        if !p.hasLargest || p.largest.Less(loc) {
                p.largest, p.hasLargest = loc, true
        }
        p.largestMu.Unlock()
}

// TrackAlias reports the §11.1 Track Alias bound to this publication — the
// integer inbound subgroup streams carry to identify the track. It is the
// value the caller supplied in message.Publish.TrackAlias, or, when that was
// the zero value, the one [Session.Publish] allocated via
// [Session.AllocOutboundTrackAlias].
func (p *Publication) TrackAlias() uint64 { return p.alias }

// OpenSubgroup opens an outbound SUBGROUP_HEADER uni-stream (§11.4.2) for this
// publication's track, filling in the Track Alias automatically — h.TrackAlias
// is ignored and overwritten. It is otherwise identical to
// [Session.OpenSubgroup]: the caller MUST Close the returned stream to FIN it
// once all objects are written, or Cancel to reset. After [Publication.Done]
// it opens nothing, and the WriteObject methods of a subgroup it opened fail,
// both with [ErrPublicationEnded].
func (p *Publication) OpenSubgroup(h message.SubgroupHeader) (*OutgoingSubgroupStream, error) {
        p.subMu.Lock()
        if p.ended.Load() {
                p.subMu.Unlock()
                return nil, ErrPublicationEnded
        }
        if p.paused.Load() {
                p.subMu.Unlock()
                return nil, ErrForwardPaused
        }
        p.opening.Add(1)
        p.subMu.Unlock()
        defer p.opening.Done()

        h.TrackAlias = p.alias
        // Once its header is written the peer can attribute the stream, so it
        // counts, even if Done reset it just after; its writes then fail.
        sg, _, err := p.s.openSubgroup(p.endCtx, h, true)
        if err != nil {
                if p.endCtx.Err() != nil {
                        return nil, ErrPublicationEnded // Done reset the header write
                }
                return nil, err
        }
        sg.onObject = p.noteObject
        sg.paused = p.paused.Load
        sg.ended = p.ended.Load
        sg.onEnd = func() { p.forget(sg) }
        p.subMu.Lock()
        p.subgroupCount++
        p.open[sg] = struct{}{}
        p.subMu.Unlock()
        return sg, nil
}

// forget drops a subgroup that was FINished or reset from the ones Done
// resets.
func (p *Publication) forget(sg *OutgoingSubgroupStream) {
        p.subMu.Lock()
        delete(p.open, sg)
        p.subMu.Unlock()
}

// Done ends the publication (§10.12): "A sender MUST NOT send PUBLISH_DONE
// until it has closed all streams it will ever open", so Done stops new
// subgroups, resets the ones still open with CANCELLED, then writes a
// PUBLISH_DONE with the given status code and reason and FINs the request
// stream. It does not wait for subgroups to drain: finish them with Close
// first to deliver their objects. A subgroup's WriteObject after Done fails
// with [ErrPublicationEnded].
//
// The Stream Count is the number of subgroup streams opened via
// [Publication.OpenSubgroup], exact however those opens race Done. Subgroups
// opened via [Session.OpenSubgroup] directly are not counted — send
// PUBLISH_DONE yourself via message.Marshal if you need a different count.
//
// Only the first call sends; later ones return nil.
func (p *Publication) Done(code moqt.PublishDoneCode, reason string) error {
        p.subMu.Lock()
        if p.ended.Load() {
                p.subMu.Unlock()
                return nil
        }
        p.ended.Store(true)
        p.subMu.Unlock()
        // Opens in flight finish promptly: their header writes are cancelled.
        p.endCancel()
        p.opening.Wait()

        p.subMu.Lock()
        open, count := p.open, p.subgroupCount
        p.open = nil
        p.subMu.Unlock()
        // §11.4.3: ending the subscription early resets the subgroups it cut
        // short, keeping what was written, header first, reliable so the
        // subscriber can attribute each reset stream when handling PUBLISH_DONE.
        for sg := range open {
                sg.MarkReliable()
                sg.dst.CancelWrite(uint64(moqt.StreamResetCancelled))
        }
        if err := p.writeThenClose(&message.PublishDone{
                StatusCode:  code,
                StreamCount: count,
                ErrorReason: reason,
        }); err != nil {
                return fmt.Errorf("moqt/session: write PUBLISH_DONE: %w", err)
        }
        return nil
}

// IncomingPublication is the receiving side of a publisher-initiated PUBLISH
// (§10.11) this endpoint accepted via [Request.AcceptPublish] — the accept-side
// counterpart of [Session.Subscribe]'s [Subscription]. The objects arrive on
// subgroup uni-streams (or datagrams) keyed by [IncomingPublication.TrackAlias]
// and are consumed via [Session.AcceptDataStream]; the embedded request stream
// stays open for follow-ups — PUBLISH_DONE from the publisher, or a
// REQUEST_UPDATE this side sends via [IncomingPublication.Update] to adjust
// forwarding (§10.9). Close it to end the reception.
//
// Its Track Alias is released (§11.1) when the subscription is Terminated
// (§5.1), as for a [Subscription].
type IncomingPublication struct {
        // requestHandle carries the PUBLISH request stream — still open for
        // follow-up traffic (inbound PUBLISH_DONE, outbound REQUEST_UPDATE;
        // Close it to end the reception) — and provides Update.
        requestHandle

        alias uint64
}

// TrackAlias reports the §11.1 Track Alias the publisher assigned — the integer
// inbound subgroup and datagram streams carry to identify this track (resolve it
// via [Session.LookupInboundTrackAlias]).
func (p *IncomingPublication) TrackAlias() uint64 { return p.alias }

// Publish opens a PUBLISH request stream (§10.11) and awaits the peer's
// initial response. It is [Session.OpenPublish] plus the response wait: the
// session assigns m.RequestID (after the stream opens, so a blocked open
// consumes no ID) and, when m.TrackAlias is the zero value, a Track Alias via
// [Session.AllocOutboundTrackAlias]; the caller supplies Namespace / Name /
// Parameters / TrackProperties. On success a [Publication] is returned whose
// embedded stream stays open for PUBLISH_DONE / REQUEST_UPDATE follow-ups and
// whose [Publication.OpenSubgroup] opens subgroup uni-streams for the track.
// On REQUEST_ERROR the stream is closed and a *RequestRejectedError is
// returned.
//
// To assign the Track Alias yourself (e.g. to mirror an upstream alias), set
// m.TrackAlias before calling — any non-zero value is used as-is — or drop to
// [Session.OpenPublish] for full control over the stream lifecycle.
func (s *Session) Publish(ctx context.Context, m *message.Publish) (*Publication, error) {
        if m.TrackAlias == 0 {
                m.TrackAlias = s.AllocOutboundTrackAlias()
        }
        return awaitRequestResponse(ctx, s, m,
                func(stream Stream, _ *message.RequestOK) (*Publication, error) {
                        // The PUBLISH sets the initial Forward State (§5.1); PUBLISH_OK
                        // carries no subscription parameters.
                        p := newPublication(s, stream, m.RequestID, m.TrackAlias, m.Parameters, m.Parameters)
                        p.answered = message.TypePublish
                        return p, nil
                })
}

// OpenPublish opens a PUBLISH request stream (§10.11) without blocking on
// stream-flow-control credit and without awaiting the peer's response. It is
// the relay-side counterpart of [Publish]: relay fan-out is fire-and-continue,
// so the caller owns the stream's read side.
//
// If the peer's stream limit is currently exhausted it returns
// [ErrNoStreamCredit] and consumes NO Request ID — the ID is allocated only
// after the stream is successfully opened (see [Session.openAllocRequest]), so
// a blocked attempt leaves the session's Request ID sequence untouched. This
// lets the caller react to an exhausted limit by sending PUBLISH_SKIPPED (§6.1,
// §10.21) instead. On success it assigns m.RequestID, writes the PUBLISH as the
// stream's first message, and returns the still-open bidi stream so the caller
// can read the peer's REQUEST_OK / REQUEST_ERROR and send follow-ups (subgroup
// streams, PUBLISH_DONE, REQUEST_UPDATE).
func (s *Session) OpenPublish(m *message.Publish) (Stream, error) {
        return s.openAllocRequest(m)
}

// AwaitPublishOK reads the response to a PUBLISH sent with
// [Session.OpenPublish], with the checks [Session.Publish] applies (§10.5,
// §10.2.1). REQUEST_ERROR is returned as a *RequestRejectedError. The stream
// stays open either way.
func (s *Session) AwaitPublishOK(ctx context.Context, stream Stream) (*message.RequestOK, error) {
        resp, err := s.readResponse(ctx, stream)
        if err != nil {
                return nil, fmt.Errorf("moqt/session: read PUBLISH response: %w", err)
        }
        switch m := resp.(type) {
        case *message.RequestOK:
                if err := s.checkRequestOKTrackProperties((*message.Publish)(nil), m); err != nil {
                        return nil, err
                }
                if err := s.CheckPeerParams(message.ScopePublishOK, m); err != nil {
                        return nil, err
                }
                return m, nil
        case *message.RequestError:
                return nil, s.rejection(m, message.TypePublish)
        default:
                return nil, fmt.Errorf("moqt/session: unexpected %s in PUBLISH response", resp.Type())
        }
}

package quicconn

import (
        "context"
        "crypto/tls"
        "fmt"
        "net"
        "net/netip"

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

        "github.com/floatdrop/moq-go/pkg/moqt/session"
)

// Dial opens a raw-QUIC connection to addr ("host:port") and returns it as a
// [session.Conn], ready for a client-side MOQT SETUP. It is the client-side
// counterpart of [NewListener], and the one dial path every MOQT client in this
// repo shares — the relay's cross-relay dialer and the demo/interop CLIs alike.
//
// The name in addr is resolved here rather than left to [quic.DialAddr], which
// resolves via [net.ResolveUDPAddr]. That helper returns the first *IPv4*
// address for a bare "host:port" and reaches for IPv6 only when the string
// carries a bracketed literal — it picks the family by looking for a '[' in the
// string (net.addrList.forResolve), not by what the resolver ranked first. So a
// dual-stack peer named by hostname is always dialed over IPv4, even where only
// the IPv6 path carries traffic: the Initials leave, nothing answers, and the
// dial fails with "timeout: no recent network activity" while the host is
// plainly reachable over IPv6.
//
// [net.Resolver.LookupIPAddr] instead returns every address in RFC 6724 order —
// the same ranking getaddrinfo / `getent ahosts` report — and each is tried in
// turn, so a host whose first address is unreachable still connects on the next.
// Candidates go back to [quic.DialAddr] as literals via [net.JoinHostPort],
// which brackets IPv6: that pins the family chosen here and costs no second
// lookup. ctx bounds the whole sequence, so a caller's dial timeout applies
// across all candidates rather than per candidate.
func Dial(ctx context.Context, addr string, tlsCfg *tls.Config, quicCfg *quic.Config) (session.Conn, error) {
        candidates, err := resolveDialCandidates(ctx, addr)
        if err != nil {
                return nil, err
        }
        var lastErr error
        for _, candidate := range candidates {
                qc, err := quic.DialAddr(ctx, candidate, tlsCfg, quicCfg)
                if err == nil {
                        return New(qc), nil
                }
                lastErr = fmt.Errorf("dial %s: %w", candidate, err)
                // A cancelled/expired ctx fails every remaining candidate the same way;
                // report the first real failure instead of the derived ones.
                if ctx.Err() != nil {
                        break
                }
        }
        return nil, lastErr
}

// resolveDialCandidates expands a "host:port" into the addresses to dial, in the
// resolver's preferred order. A host that is already an IP literal resolves to
// itself: no DNS, and no chance of the family flipping under a caller that
// deliberately pinned one. [netip.ParseAddr] rather than [net.ParseIP] so a
// zone-scoped literal ("fe80::1%eth0") is recognized as one too.
func resolveDialCandidates(ctx context.Context, addr string) ([]string, error) {
        host, port, err := net.SplitHostPort(addr)
        if err != nil {
                return nil, err
        }
        if _, err := netip.ParseAddr(host); err == nil {
                return []string{addr}, nil
        }
        // LookupIPAddr reports an error rather than an empty slice when a name has
        // no addresses, so the dial loop above always has at least one candidate.
        ips, err := net.DefaultResolver.LookupIPAddr(ctx, host)
        if err != nil {
                return nil, err
        }
        candidates := make([]string, len(ips))
        for i, ip := range ips {
                // IPAddr.String carries the zone; JoinHostPort adds the brackets.
                candidates[i] = net.JoinHostPort(ip.String(), port)
        }
        return candidates, nil
}

// Package quicconn adapts github.com/quic-go/quic-go's *quic.Conn to the
// transport-neutral session.Conn interface.
//
// This is the sole boundary in the moqt tree where quic-go's concrete types
// meet the session abstraction. Putting it in a dedicated subpackage lets the
// rest of pkg/moqt (and its tests) stay independent of quic-go's surface.
//
// quic-go uses typed-uint64 aliases (quic.StreamErrorCode,
// quic.ApplicationErrorCode) for error codes; session.Conn / SendStream /
// ReceiveStream use plain uint64. The wrappers below do the lossless
// conversion at each call site.
package quicconn

import (
        "context"
        "errors"
        "net"

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

        "github.com/floatdrop/moq-go/pkg/moqt/session"
)

// New wraps c so it satisfies session.Conn.
func New(c *quic.Conn) session.Conn { return &conn{q: c} }

// Compile-time satisfaction check.
var _ session.Conn = (*conn)(nil)

// conn holds a *quic.Conn by named field rather than embedding. Embedding
// would promote quic-go's CloseWithError(quic.ApplicationErrorCode, string)
// onto the wrapper; the session.Conn interface demands
// CloseWithError(uint64, string). Two methods of the same name with different
// signatures aren't allowed on a single Go type, so we delegate explicitly.
type conn struct{ q *quic.Conn }

func (c *conn) OpenUniStream() (session.SendStream, error) {
        s, err := c.q.OpenUniStream()
        if err != nil {
                if _, ok := errors.AsType[*quic.StreamLimitReachedError](err); ok {
                        return nil, session.ErrNoStreamCredit
                }
                return nil, err
        }
        return &sendStream{s: s}, nil
}

func (c *conn) AcceptUniStream(ctx context.Context) (session.ReceiveStream, error) {
        s, err := c.q.AcceptUniStream(ctx)
        if err != nil {
                return nil, err
        }
        return &recvStream{s: s}, nil
}

// OpenStream opens a bidirectional stream without blocking. quic-go returns a
// *quic.StreamLimitReachedError when the peer's stream limit is exhausted; we
// map that onto session.ErrNoStreamCredit so callers can detect it
// transport-neutrally with errors.Is.
func (c *conn) OpenStream() (session.Stream, error) {
        s, err := c.q.OpenStream()
        if err != nil {
                if _, ok := errors.AsType[*quic.StreamLimitReachedError](err); ok {
                        return nil, session.ErrNoStreamCredit
                }
                return nil, err
        }
        return &bidiStream{s: s}, nil
}

func (c *conn) AcceptStream(ctx context.Context) (session.Stream, error) {
        s, err := c.q.AcceptStream(ctx)
        if err != nil {
                return nil, err
        }
        return &bidiStream{s: s}, nil
}

func (c *conn) CloseWithError(code uint64, reason string) error {
        return c.q.CloseWithError(quic.ApplicationErrorCode(code), reason)
}

func (c *conn) Context() context.Context { return c.q.Context() }

func (c *conn) SendDatagram(payload []byte) error {
        return c.q.SendDatagram(payload)
}

func (c *conn) ReceiveDatagram(ctx context.Context) ([]byte, error) {
        return c.q.ReceiveDatagram(ctx)
}

// sendStream wraps *quic.SendStream. Named field for the same reason as conn.
type sendStream struct{ s *quic.SendStream }

func (s *sendStream) Write(p []byte) (int, error) { return s.s.Write(p) }
func (s *sendStream) Close() error                { return s.s.Close() }
func (s *sendStream) CancelWrite(code uint64) {
        s.s.CancelWrite(quic.StreamErrorCode(code))
}

// SetReliableBoundary satisfies [session.ReliableResetStream] by forwarding to
// quic-go's RESET_STREAM_AT support. It is a no-op unless the peer enabled the
// extension (quic.Config.EnableStreamResetPartialDelivery).
func (s *sendStream) SetReliableBoundary() { s.s.SetReliableBoundary() }

// Context is cancelled when the send side is closed or reset, or the peer
// sends STOP_SENDING. quic-go's SendStream.Context() provides this directly.
func (s *sendStream) Context() context.Context { return s.s.Context() }

// recvStream wraps *quic.ReceiveStream.
type recvStream struct{ s *quic.ReceiveStream }

func (s *recvStream) Read(p []byte) (int, error) { return s.s.Read(p) }
func (s *recvStream) CancelRead(code uint64) {
        s.s.CancelRead(quic.StreamErrorCode(code))
}

// bidiStream wraps *quic.Stream.
type bidiStream struct{ s *quic.Stream }

func (s *bidiStream) Read(p []byte) (int, error)  { return s.s.Read(p) }
func (s *bidiStream) Write(p []byte) (int, error) { return s.s.Write(p) }
func (s *bidiStream) Close() error                { return s.s.Close() }
func (s *bidiStream) CancelRead(code uint64) {
        s.s.CancelRead(quic.StreamErrorCode(code))
}
func (s *bidiStream) CancelWrite(code uint64) {
        s.s.CancelWrite(quic.StreamErrorCode(code))
}

// Context is cancelled when the send side is closed or reset, or the peer
// sends STOP_SENDING. quic-go's Stream embeds SendStream which has Context().
func (s *bidiStream) Context() context.Context { return s.s.Context() }

// Listener adapts a *quic.Listener so it can be handed directly to the
// relay's accept loop. The relay's listener interface requires
// Accept(ctx) → session.Conn, Addr() → net.Addr, and Close() → error;
// this type satisfies it structurally without forcing this package to
// import pkg/relay.
//
// The caller owns the underlying *quic.Listener — its TLS config, ALPN
// selection ("moqt-20"), QUIC parameters, and listening socket. Close
// on the Listener forwards to the underlying *quic.Listener, which is
// also what the caller would call themselves on shutdown; both paths
// are equivalent.
type Listener struct{ ql *quic.Listener }

// NewListener wraps ql so it can be passed to relay.New.
//
// Typical wiring:
//
//        ql, err := quic.ListenAddr(":4433", tlsCfg, quicCfg)
//        if err != nil { … }
//        r := relay.New(quicconn.NewListener(ql), relay.Config{ … })
//        go r.Start(ctx)
func NewListener(ql *quic.Listener) *Listener { return &Listener{ql: ql} }

// Accept blocks until the next inbound *quic.Conn arrives, then wraps
// it via [New] into a session.Conn the relay can hand to session.Server.
// ctx cancellation propagates to the underlying Accept.
func (l *Listener) Accept(ctx context.Context) (session.Conn, error) {
        c, err := l.ql.Accept(ctx)
        if err != nil {
                return nil, err
        }
        return New(c), nil
}

// Addr returns the address the underlying quic-go listener is bound to.
func (l *Listener) Addr() net.Addr { return l.ql.Addr() }

// Close closes the underlying *quic.Listener. Subsequent Accept calls
// unblock with the quic-go close error.
func (l *Listener) Close() error { return l.ql.Close() }

package session

import (
        "bytes"
        "context"
        "errors"
        "fmt"
        "io"
        "math"
        "slices"
        "sync"
        "sync/atomic"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// ErrRequestIDParityViolation is returned by [Session.CheckPeerRequestID] when
// the peer sends a Request ID whose parity does not match the expected value
// per §10.1. [Session.AcceptRequest] has already closed the session with
// INVALID_REQUEST_ID; another caller of CheckPeerRequestID MUST.
type ErrRequestIDParityViolation struct {
        RequestID    uint64
        ExpectedEven bool // true = expected even (peer is client), false = expected odd (peer is server)
}

func (e *ErrRequestIDParityViolation) Error() string {
        want := "even"
        if !e.ExpectedEven {
                want = "odd"
        }
        return fmt.Sprintf(
                "moqt/session: peer Request ID %d has wrong parity (want %s) — INVALID_REQUEST_ID",
                e.RequestID,
                want,
        )
}

// ErrDuplicateRequestID is returned by [Session.CheckPeerRequestID] (and thus
// AcceptRequest) when the peer reuses a Request ID (§10.1: "a duplicate
// Request ID" MUST close the session with INVALID_REQUEST_ID). Cross-stream
// delivery reordering is tolerated — an ID below the high-water mark counts
// as a duplicate only once every unseen ID it could have been is accounted
// for. [Session.AcceptRequest] has already closed the session with
// INVALID_REQUEST_ID; another caller of CheckPeerRequestID MUST.
type ErrDuplicateRequestID struct {
        RequestID uint64
        MaxSeen   uint64
}

func (e *ErrDuplicateRequestID) Error() string {
        return fmt.Sprintf(
                "moqt/session: peer Request ID %d already consumed (high-water mark %d) — INVALID_REQUEST_ID",
                e.RequestID,
                e.MaxSeen,
        )
}

// ErrUnexpectedRequestUpdate is returned by AcceptRequest when a peer opens a
// request stream with REQUEST_UPDATE, which §10.9 allows only as a follow-up.
// AcceptRequest has already closed the session with PROTOCOL_VIOLATION.
type ErrUnexpectedRequestUpdate struct {
        RequestID uint64
}

// ErrUnexpectedPublishStateNotify is returned by AcceptRequest when a peer
// opens a request stream with PUBLISH_STATE_NOTIFY, which §10.10 allows only on
// an existing subscription's stream. AcceptRequest has already closed the
// session with PROTOCOL_VIOLATION.
var ErrUnexpectedPublishStateNotify = errors.New(
        "moqt/session: PUBLISH_STATE_NOTIFY as the first message of a request stream — PROTOCOL_VIOLATION")

func (e *ErrUnexpectedRequestUpdate) Error() string {
        return fmt.Sprintf(
                "moqt/session: REQUEST_UPDATE (Request ID %d) as the first message of a request stream — PROTOCOL_VIOLATION",
                e.RequestID,
        )
}

// ErrUnexpectedRequestOpener is returned by AcceptRequest when a peer opens a
// request stream with a message that is not a request opener (§3.3), including
// an unknown type. AcceptRequest has already closed the session with
// PROTOCOL_VIOLATION.
type ErrUnexpectedRequestOpener struct {
        Type message.Type
}

func (e *ErrUnexpectedRequestOpener) Error() string {
        return fmt.Sprintf(
                "moqt/session: %s as the first message of a request stream — PROTOCOL_VIOLATION", e.Type)
}

// isRequestOpener reports whether msg is one of the seven messages that may
// open a request stream (§3.3; marked "First" in Table 5).
func isRequestOpener(msg message.Message) bool {
        switch msg.(type) {
        case *message.Subscribe, *message.Publish, *message.Fetch, *message.TrackStatus,
                *message.PublishNamespace, *message.SubscribeNamespace, *message.SubscribeTracks:
                return true
        }
        return false
}

// ErrTooManyRequestUpdates is returned by [RequestUpdateLimiter.Received] when
// a peer exceeds the per-request-stream MAX_REQUEST_UPDATES limit it was
// advertised (§10.3.1.7). The caller MUST close the session with
// SessionTooManyRequestUpdates.
type ErrTooManyRequestUpdates struct {
        Limit uint64
}

func (e *ErrTooManyRequestUpdates) Error() string {
        return fmt.Sprintf(
                "moqt/session: peer exceeded MAX_REQUEST_UPDATES (%d) outstanding on a request stream — TOO_MANY_REQUEST_UPDATES",
                e.Limit,
        )
}

// RequestUpdateLimiter enforces the receive-side MAX_REQUEST_UPDATES limit
// (§10.3.1.7) for a single request stream. A REQUEST_UPDATE is "outstanding"
// from when it is received until this endpoint writes the mandated
// REQUEST_OK/REQUEST_ERROR; the sender may not have more than the advertised
// limit outstanding at once. Construct one per stream via
// [Session.NewRequestUpdateLimiter].
//
// A limiter is not safe for concurrent use, which matches the single-reader
// invariant of the follow-up loops ([RequestBroker.Serve] and the relay's
// per-stream readers). A limit of 0 (the default, meaning the option was not
// advertised) disables the check.
type RequestUpdateLimiter struct {
        limit       uint64
        outstanding uint64
}

// NewRequestUpdateLimiter returns a limiter seeded with the MAX_REQUEST_UPDATES
// value this session advertised to the peer.
func (s *Session) NewRequestUpdateLimiter() *RequestUpdateLimiter {
        return &RequestUpdateLimiter{limit: s.maxRequestUpdates}
}

// Received records an inbound REQUEST_UPDATE. It returns
// [*ErrTooManyRequestUpdates] when the stream already holds the advertised
// limit of outstanding updates (§10.3.1.7: the endpoint MUST then close the
// session with TOO_MANY_REQUEST_UPDATES); the caller owns that close, mirroring
// [Session.CheckPeerRequestID]. On success the update counts as outstanding
// until a matching [RequestUpdateLimiter.Responded].
func (l *RequestUpdateLimiter) Received() error {
        if l.limit != 0 && l.outstanding >= l.limit {
                return &ErrTooManyRequestUpdates{Limit: l.limit}
        }
        l.outstanding++
        return nil
}

// Responded releases the credit a successful [RequestUpdateLimiter.Received]
// took, once this endpoint has written the mandated REQUEST_OK/REQUEST_ERROR.
// Callers pair it with exactly one Received that returned nil (a Received that
// errored closes the session and never reaches here), so outstanding is always
// at least 1 on entry.
func (l *RequestUpdateLimiter) Responded() {
        l.outstanding--
}

// RequestRejectedError is returned by a request opener when the peer answers
// with REQUEST_ERROR (§10.6).
type RequestRejectedError struct {
        Code   moqt.RequestErrorCode
        Reason string
        // RetryInterval is the raw Retry Interval (§10.6.2); see
        // [RequestRejectedError.RetryAfter].
        RetryInterval uint64
        // Redirect is where to retry the request (§10.6.1), present exactly when
        // Code is REDIRECT: an empty ConnectURI means this session's URI.
        Redirect *message.Redirect
}

// rejection converts a REQUEST_ERROR answering a message of type req into a
// *RequestRejectedError; for a REQUEST_UPDATE the request it updates is not
// known, so only the Connect URI is checked. A Redirect that §10.6.1 makes
// session-fatal closes the session with PROTOCOL_VIOLATION, and that error is
// returned instead.
func (s *Session) rejection(m *message.RequestError, req message.Type) error {
        if err := redirectViolation(m.Redirect, req, s.role == roleServer); err != nil {
                return s.closeProtocolViolation(fmt.Errorf("moqt/session: received %w", err))
        }
        return &RequestRejectedError{
                Code:          m.ErrorCode,
                Reason:        m.ErrorReason,
                RetryInterval: m.RetryInterval,
                Redirect:      m.Redirect,
        }
}

// redirectViolation reports a Redirect for a request of type req that the
// receiver, a server when atServer, MUST close the session for (§10.6.1): "If
// a server receives a Redirect with a non-zero Connect URI Length", or a
// non-empty Track Name "in a Redirect for a namespace-scoped request".
func redirectViolation(rd *message.Redirect, req message.Type, atServer bool) error {
        switch {
        case rd == nil:
                return nil
        case atServer && len(rd.ConnectURI) > 0:
                return errors.New("a Redirect with a Connect URI at a server (§10.6.1)")
        case len(rd.TrackName) > 0 &&
                (req == message.TypeSubscribeNamespace || req == message.TypePublishNamespace ||
                        req == message.TypeSubscribeTracks):
                return fmt.Errorf("a Redirect with a Track Name for %s (§10.6.1)", req)
        }
        return nil
}

// RetryAfter decodes RetryInterval (§10.6.2): whether the request may be
// retried with the same parameters, and the minimum wait before doing so.
func (e *RequestRejectedError) RetryAfter() (time.Duration, bool) {
        if e.RetryInterval == 0 {
                return 0, false
        }
        ms := e.RetryInterval - 1
        if ms > uint64(math.MaxInt64/int64(time.Millisecond)) {
                return time.Duration(math.MaxInt64), true // a varint can exceed Duration
        }
        return time.Duration(ms) * time.Millisecond, true
}

func (e *RequestRejectedError) Error() string {
        return fmt.Sprintf("moqt request rejected: %s (code %#x)", e.Reason, uint64(e.Code))
}

// Request is an inbound MoQT request stream (§3.3, §10.1) after its first
// message has been parsed.
//
// "Request" here matches MoQT's terminology, not the one-shot RPC sense the
// word usually implies in Go. A request is a long-lived request-response
// interaction identified by a Request ID: the bidi stream stays open for the
// lifetime of the operation, the responder writes an initial response
// (REQUEST_OK / REQUEST_ERROR / SUBSCRIBE_OK / PUBLISH_OK), and either side
// may send follow-up messages on the same stream — REQUEST_UPDATE from the
// requester, PUBLISH_DONE from a publisher, additional REQUEST_OKs in
// response to updates, and so on — until one side FINs or resets the stream.
//
// Handlers read First to decide what to do, write responses via Reply or
// RejectError, and use Stream directly for any further messages or to close
// the send side.
type Request struct {
        Stream Stream
        First  message.Message

        // Tokens holds the AUTHORIZATION_TOKEN values (§10.2.2) carried by
        // First, fully resolved against the inbound token cache: REGISTER and
        // USE_VALUE tokens contribute their (Type, Value) directly, USE_ALIAS
        // tokens are resolved to the previously-registered value, and DELETE
        // tokens are applied to the cache without producing an entry here. It is
        // nil when the request carried no tokens. Handlers (and
        // [Session.VerifyRequestTokens]) consult it to authorize the request;
        // callers never see a bare alias.
        Tokens []ResolvedToken

        // s is the owning session, used by the AcceptSubscribe / AcceptPublish
        // helpers to allocate Track Aliases and register inbound aliases.
        s *Session

        // okSent records that Reply has sent a REQUEST_OK, so a later one answers
        // a REQUEST_UPDATE (§10.5).
        okSent atomic.Bool
}

// AcceptRequest blocks until a peer opens a bidirectional stream, reads and
// parses the first message, and returns the result. The session must be past
// SETUP (i.e. Open has returned successfully).
//
// Requests that target a reserved namespace the MOQT implementation owns
// (§3.2.1 "." and §3.2.2 ".session") are answered with REQUEST_ERROR
// DOES_NOT_EXIST and skipped transparently — the caller (and, for a relay, any
// other session) never observes them, satisfying "Relays MUST NOT forward
// requests for session-level tracks and namespaces". AcceptRequest loops until
// it has an application-visible request to return.
//
// A stream not opened by a request message (§3.3), whose first message is
// malformed (§10, wrapping [message.ErrMalformedMessage]), or a PUBLISH with a
// session-fatal Track Property value (§12.5, §12.6), closes the session with
// PROTOCOL_VIOLATION; the error is *ErrUnexpectedRequestOpener,
// *ErrUnexpectedRequestUpdate, ErrUnexpectedPublishStateNotify or the parse
// error. A Request ID violation (§10.1) closes it with INVALID_REQUEST_ID and
// returns *ErrRequestIDParityViolation or *ErrDuplicateRequestID, and a token
// cache fault (§10.2.2) with the *TokenCacheError's Code. A stream that ends
// or is reset before its first message is complete fails only that request
// (§3.3.2, §3.3.3): it is reset and AcceptRequest moves on.
func (s *Session) AcceptRequest(ctx context.Context) (*Request, error) {
        for {
                stream, err := s.conn.AcceptStream(ctx)
                if err != nil {
                        return nil, err
                }
                // readResponse bridges ctx to the otherwise context-free Parse, so a
                // peer that opens the stream but stalls mid-message cannot wedge the
                // accept loop past cancellation.
                msg, err := s.readResponse(ctx, stream)
                if err != nil {
                        if ctx.Err() != nil {
                                resetStream(stream)
                                return nil, ctx.Err()
                        }
                        // §3.3: readResponse already closed the session; this only shapes
                        // the error.
                        if typ, ok := errors.AsType[message.ErrUnknownType](err); ok {
                                resetStream(stream)
                                return nil, s.closeProtocolViolation(&ErrUnexpectedRequestOpener{Type: message.Type(typ)})
                        }
                        if errors.Is(err, message.ErrMalformedMessage) {
                                resetStream(stream)
                                return nil, fmt.Errorf("moqt/session: parse request first message: %w", err)
                        }
                        // §3.3.2, §3.3.3: the peer ended or reset the stream first, which
                        // fails that request only. A closed session fails AcceptStream.
                        cancelRequest(stream)
                        continue
                }

                // §10.9, §3.3: REQUEST_UPDATE never opens a stream.
                if upd, ok := msg.(*message.RequestUpdate); ok {
                        resetStream(stream)
                        return nil, s.closeProtocolViolation(&ErrUnexpectedRequestUpdate{RequestID: upd.RequestID})
                }

                // §10.10: PUBLISH_STATE_NOTIFY never opens a stream. It carries no
                // Request ID, so the §10.1 check below would not catch it.
                if _, ok := msg.(*message.PublishStateNotify); ok {
                        resetStream(stream)
                        return nil, s.closeProtocolViolation(ErrUnexpectedPublishStateNotify)
                }

                // §3.3: "Bidirectional streams MUST NOT begin with any other message
                // type unless negotiated."
                if !isRequestOpener(msg) {
                        resetStream(stream)
                        return nil, s.closeProtocolViolation(&ErrUnexpectedRequestOpener{Type: msg.Type()})
                }

                // §10.2.1 / §10.2: a parameter outside the opener's scope, or a
                // repeated one, closes the session.
                if err := s.CheckPeerParams(message.ScopeOfRequest(msg.Type()), msg); err != nil {
                        resetStream(stream)
                        return nil, err
                }
                // §12.5, §12.6: on receipt, whatever the application answers.
                if pub, ok := msg.(*message.Publish); ok {
                        if err := s.checkTrackPropertyValues(pub.TrackProperties, "PUBLISH"); err != nil {
                                resetStream(stream)
                                return nil, err
                        }
                }

                // §10.1 parity and duplicate check.
                if m, ok := msg.(message.WithRequestID); ok {
                        if err := s.CheckPeerRequestID(m.GetRequestID()); err != nil {
                                resetStream(stream)
                                _ = s.Close(moqt.SessionInvalidRequestID, err.Error())
                                return nil, err
                        }
                }

                // §10.2.2: REGISTER tokens commit before any rejection, so the alias
                // persists even if the request fails. A *TokenCacheError is
                // session-fatal.
                tokens, err := s.processRequestTokens(msg)
                if err != nil {
                        resetStream(stream)
                        if tce, ok := errors.AsType[*TokenCacheError](err); ok {
                                _ = s.Close(tce.Code, tce.Error())
                        }
                        return nil, err
                }

                // §3.2.1 / §3.2.2: reserved namespaces the implementation owns are
                // rejected here, after token processing; other "."-prefixed ones reach
                // the application.
                if reason, reject := reservedNamespaceRejection(msg); reject {
                        rejectStreamWithError(stream, moqt.RequestDoesNotExist, reason)
                        continue
                }

                return &Request{Stream: stream, First: msg, Tokens: tokens, s: s}, nil
        }
}

// maxTrackedRequestIDGaps bounds [Session.CheckPeerRequestID]'s memory for
// below-the-mark Request IDs that may still legitimately arrive late. A
// conforming peer creates gaps only through delivery reordering of in-flight
// requests (it allocates in +2 increments), so the bound is far above any
// realistic reorder window; when it overflows, the lowest (oldest) gaps are
// evicted first — they are the least plausible late arrivals — and a later
// arrival for an evicted one reads as a duplicate.
const maxTrackedRequestIDGaps = 1024

// evictLowestGapsLocked removes the n smallest Request IDs from gaps, keeping
// the newest entries claimable when the cap forces a choice. O(cap log cap),
// and only runs on a jump that overflows the cap. Caller holds s.mu.
func evictLowestGapsLocked(gaps map[uint64]struct{}, n int) {
        ids := make([]uint64, 0, len(gaps))
        for id := range gaps {
                ids = append(ids, id)
        }
        slices.Sort(ids)
        for _, id := range ids[:min(n, len(ids))] {
                delete(gaps, id)
        }
}

// CheckPeerRequestID validates one inbound Request ID per §10.1 and records
// it. It applies to every peer message that consumes a Request ID — the
// first message of a request stream (AcceptRequest calls this) and follow-up
// REQUEST_UPDATEs ([RequestBroker.Serve] and relay follow-up readers call it
// for those).
//
// Two violations are session-fatal per §10.1, and the caller MUST close the
// session with [moqt.SessionInvalidRequestID] (AcceptRequest does so itself):
//
//   - wrong parity for the sender (*ErrRequestIDParityViolation);
//   - a duplicate ID (*ErrDuplicateRequestID).
//
// An ID below the high-water mark is NOT automatically a duplicate: the peer
// allocates in +2 increments, but requests ride separate QUIC streams and
// can be delivered out of order, so each unseen ID below the mark stays
// claimable exactly once.
func (s *Session) CheckPeerRequestID(rid uint64) error {
        // §10.1: the client generates even Request IDs (starting at 0), the
        // server odd ones (starting at 1); peerMustBeEven is true when we are
        // the server.
        peerMustBeEven := s.role == roleServer
        if peerMustBeEven && rid%2 != 0 {
                return &ErrRequestIDParityViolation{RequestID: rid, ExpectedEven: true}
        }
        if !peerMustBeEven && rid%2 != 1 {
                return &ErrRequestIDParityViolation{RequestID: rid, ExpectedEven: false}
        }

        s.mu.Lock()
        defer s.mu.Unlock()
        if !s.peerRequestIDSeen || rid > s.peerRequestIDMax {
                s.recordRequestIDGapsLocked(rid, peerMustBeEven)
                s.peerRequestIDSeen = true
                s.peerRequestIDMax = rid
                return nil
        }
        if _, open := s.peerRequestIDGaps[rid]; open {
                delete(s.peerRequestIDGaps, rid)
                return nil
        }
        return &ErrDuplicateRequestID{RequestID: rid, MaxSeen: s.peerRequestIDMax}
}

// recordRequestIDGapsLocked records the peer Request IDs an advance of the
// high-water mark to rid skips over, as claimable reorder gaps. The peer's
// sequence starts at its parity base (§10.1: client 0, server 1), so on the
// very first observation everything below rid is potentially in flight. All
// new gaps are newer than every existing entry (which lie below the previous
// mark), so keeping the newest cap-many claimable means inserting at most
// cap new gaps (newest first) and evicting the lowest old entries to make
// room. Caller holds s.mu.
func (s *Session) recordRequestIDGapsLocked(rid uint64, peerMustBeEven bool) {
        lo := uint64(0)
        if !peerMustBeEven {
                lo = 1
        }
        if s.peerRequestIDSeen {
                lo = s.peerRequestIDMax + 2
        }
        if rid <= lo {
                return
        }
        newGaps := maxTrackedRequestIDGaps
        if d := (rid - lo) / 2; d < maxTrackedRequestIDGaps {
                newGaps = int(d)
        }
        if excess := len(s.peerRequestIDGaps) + newGaps - maxTrackedRequestIDGaps; excess > 0 {
                evictLowestGapsLocked(s.peerRequestIDGaps, excess)
        }
        if s.peerRequestIDGaps == nil {
                s.peerRequestIDGaps = make(map[uint64]struct{})
        }
        for id, n := rid, 0; n < newGaps; n++ {
                id -= 2
                s.peerRequestIDGaps[id] = struct{}{}
        }
}

// resetStream cancels both directions of a bidi request stream (§3.3.3) with
// StreamResetInternalError (§3.3.4) — the common teardown when a request stream
// is abandoned mid-parse or fails §10.1 validation.
func resetStream(s Stream) {
        s.CancelRead(uint64(moqt.StreamResetInternalError))
        s.CancelWrite(uint64(moqt.StreamResetInternalError))
}

// rejectStreamWithError applies [Request.RejectError]'s teardown on the
// pre-Request path in AcceptRequest, where no *Request value exists yet. Its
// error is dropped because there is nothing left to do with it: RejectError
// has already reset the stream if the REQUEST_ERROR could not be sent.
func rejectStreamWithError(stream Stream, code moqt.RequestErrorCode, reason string) {
        _ = (&Request{Stream: stream}).RejectError(code, reason)
}

// requestHandle is the state every typed request handle embeds: the open
// request stream, the owning session and the request's §10.1 Request ID. It
// provides the shared Close, Update and Broker methods.
type requestHandle struct {
        // Stream is the request stream. [requestHandle.Close] cancels the
        // request; Stream.Close only FINs this side (§3.3.2).
        Stream

        s         *Session
        requestID uint64

        brokerOnce sync.Once
        broker     atomic.Pointer[RequestBroker]

        // finished records that writeThenClose sent this side's final message
        // and FIN, so Close must not reset it.
        finished atomic.Bool

        // Follow-ups the peer may send (§10.9 / §10.10) and the §10.2.1 scope of
        // its REQUEST_UPDATEs; applied to the broker on creation.
        peerUpdate, peerNotify bool
        updateScope            message.ParamScope

        // inboundAlias is the inbound Track Alias this request registered
        // (§11.1), when holdsAlias: the receiving side of a subscription. It is
        // released once the subscription is Terminated (§5.1); see terminated.
        inboundAlias uint64
        holdsAlias   bool
        released     atomic.Bool
        // peerDone records that the publisher's FIN was read.
        peerDone atomic.Bool

        // namespaces marks a SUBSCRIBE_NAMESPACE, whose broker checks each
        // NAMESPACE_DONE against the NAMESPACEs before it (§10.19); nsPrefix is
        // the Track Namespace Prefix it was sent with.
        namespaces bool
        nsPrefix   wire.TrackNamespace

        // goaways checks the GOAWAYs this stream carries (§10.4), for every
        // reader of the handle.
        goaways RequestGoaways

        // answered is TypeSubscribe or TypePublish for a request this side sent,
        // whose response was read: its broker closes the session on another
        // (§5.1). Zero otherwise.
        answered message.Type
}

// Read reads the request stream. On the receiving side of a subscription
// ([Subscription], [IncomingPublication]) the publisher's FIN ends it — "by
// sending PUBLISH_DONE and closing the stream" (§5.1) — which releases the
// Track Alias (§11.1).
func (h *requestHandle) Read(p []byte) (int, error) {
        n, err := h.Stream.Read(p)
        if errors.Is(err, io.EOF) {
                h.peerFinished()
        }
        return n, err
}

// peerFinished records the publisher's FIN on the receiving side of a
// subscription, which Terminates it (§5.1).
func (h *requestHandle) peerFinished() {
        if h.holdsAlias {
                h.peerDone.Store(true)
                h.terminated()
        }
}

// terminated releases the inbound Track Alias the first time the subscription
// is Terminated: a handle can see that more than once (Close after the FIN),
// and a second release would take a registration another subscription
// sharing the alias holds (§5.1).
func (h *requestHandle) terminated() {
        if h.holdsAlias && h.released.CompareAndSwap(false, true) {
                h.s.UnregisterInboundTrackAlias(h.inboundAlias)
        }
}

// Close cancels the request by resetting both stream directions (§3.3.3). If
// this side already sent its final message and FIN (e.g. [Publication.Done]),
// only reading is stopped, so that message is not lost. A subscription is then
// Terminated (§5.1), so its Track Alias is released.
//
// Once a subscription's publisher has completed it (its FIN was read through
// the handle or its broker), Close FINs instead, as §3.3.2 asks once nothing
// further will be sent. Otherwise Close cannot know whether the peer completed the
// request; after that, use Stream.Close to FIN.
func (h *requestHandle) Close() error {
        h.terminated()
        if h.peerDone.Load() {
                return h.Stream.Close()
        }
        if h.finished.Load() {
                h.Stream.CancelRead(uint64(moqt.StreamResetCancelled))
                return nil
        }
        cancelRequest(h.Stream)
        return nil
}

// cancelRequest cancels a request stream in both directions (§3.3.3).
func cancelRequest(s Stream) {
        s.CancelRead(uint64(moqt.StreamResetCancelled))
        s.CancelWrite(uint64(moqt.StreamResetCancelled))
}

// Broker returns this request's [RequestBroker], creating it on first call.
// Run [RequestBroker.Serve] to own the stream's reads when follow-up traffic
// must coexist with updates. Once created, the handle's Update and terminal
// writes like [Publication.Done] go through the broker.
func (h *requestHandle) Broker() *RequestBroker {
        h.brokerOnce.Do(func() {
                b := h.s.NewRequestBroker(h.Stream)
                b.PeerMessages(h.peerUpdate, h.peerNotify)
                b.UpdateScope(h.updateScope)
                b.handle = h
                b.nsPrefix = h.nsPrefix
                h.broker.Store(b)
        })
        return h.broker.Load()
}

// Update sends a REQUEST_UPDATE (§10.9) and awaits its REQUEST_OK or
// REQUEST_ERROR. params carries only the fields to change.
//
// Without a [requestHandle.Broker] this is [Session.UpdateRequest] and must be
// the stream's only reader; with one it delegates to [RequestBroker.Update].
func (h *requestHandle) Update(ctx context.Context, params message.Parameters) (*message.RequestOK, error) {
        if b := h.broker.Load(); b != nil {
                return b.Update(ctx, params)
        }
        return h.s.UpdateRequest(ctx, h, params) // h, so Read sees a FIN
}

// writeThenClose writes msg and FINs the send side, through the broker's
// write lock when one exists.
func (h *requestHandle) writeThenClose(msg message.Message) error {
        if b := h.broker.Load(); b != nil {
                if err := b.writeThenClose(msg); err != nil {
                        return err
                }
                h.finished.Store(true)
                return nil
        }
        if err := message.Marshal(h.Stream, msg); err != nil {
                return err
        }
        if err := h.Stream.Close(); err != nil {
                return err
        }
        h.finished.Store(true)
        return nil
}

// openRequest opens a bidi stream and writes first as its initial message. On
// a write error the stream is reset.
func (s *Session) openRequest(first message.Message) (Stream, error) {
        stream, err := s.conn.OpenStream()
        if err != nil {
                return nil, err
        }
        return writeFirst(stream, first)
}

// openAllocRequest opens a request stream and writes m as its first message.
// m's Request ID (§10.1) is allocated only after the open succeeds, so a
// failed open consumes no ID. It does not await the response.
func (s *Session) openAllocRequest(m message.WithRequestID) (Stream, error) {
        stream, err := s.conn.OpenStream()
        if err != nil {
                return nil, err
        }
        m.SetRequestID(s.AllocRequestID())
        return writeFirst(stream, m)
}

// writeFirst marshals first as the initial message of a freshly opened request
// stream. On a write failure the stream is reset and the error is returned.
func writeFirst(stream Stream, first message.Message) (Stream, error) {
        if err := message.Marshal(stream, first); err != nil {
                resetStream(stream)
                return nil, fmt.Errorf("moqt/session: write request first message: %w", err)
        }
        return stream, nil
}

// readResponse parses one message from stream, honoring ctx by resetting the
// read side with StreamResetCancelled when ctx is done; it then returns
// ctx.Err(). A malformed message closes the session with PROTOCOL_VIOLATION.
//
// A cancellation landing between a successful Parse and stop() still resets
// the read side; callers' contexts are session-lifetime, so this only happens
// during shutdown.
func (s *Session) readResponse(ctx context.Context, stream Stream) (message.Message, error) {
        stop := context.AfterFunc(ctx, func() {
                stream.CancelRead(uint64(moqt.StreamResetCancelled))
        })
        defer stop()
        msg, err := message.Parse(stream)
        if err != nil && ctx.Err() != nil {
                return nil, ctx.Err()
        }
        // §10, §10.2: unknown type, bad Length or unknown parameter.
        if errors.Is(err, message.ErrMalformedMessage) {
                return nil, s.closeProtocolViolation(err)
        }
        return msg, err
}

// awaitRequestResponse opens a request stream for m and awaits the initial
// response. An OK is handed to onOK, which then owns the stream; REQUEST_ERROR
// (§10.6) becomes a *RequestRejectedError; anything else is an error, and for
// SUBSCRIBE_NAMESPACE and SUBSCRIBE_TRACKS also closes the session (§10.19,
// §10.20). On either failure the stream is closed. One GOAWAY ahead of the
// response is checked, the response still awaited, and the GOAWAY put back as
// the stream's first follow-up (§10.4), bar where the response must come
// first.
func awaitRequestResponse[OK message.Message, R any](
        ctx context.Context,
        s *Session,
        m message.WithRequestID,
        onOK func(stream Stream, ok OK) (R, error),
) (R, error) {
        var zero R
        stream, err := s.openAllocRequest(m)
        if err != nil {
                return zero, err
        }
        resp, err := s.readResponse(ctx, stream)
        // §10.4: "A GOAWAY MAY also be sent on a request stream to initiate
        // migration of that individual request", before its response too, save
        // where the response MUST be "the first message" (§6.1, §6.2). It is
        // checked as any on the stream is, and left for the stream's next reader,
        // as one after the response is; the request is still answered.
        if g, isGoaway := resp.(*message.Goaway); isGoaway && err == nil && !responseFirst(m.Type()) {
                var goaways RequestGoaways
                if gerr := goaways.Received(s, g); gerr != nil {
                        return zero, gerr
                }
                resp, err = s.readResponse(ctx, stream)
                if g2, again := resp.(*message.Goaway); again && err == nil {
                        return zero, goaways.Received(s, g2) // a second GOAWAY (§10.4)
                }
                stream = replayMessage(stream, g)
        }
        if err != nil {
                _ = stream.Close()
                return zero, fmt.Errorf("moqt/session: read %s response: %w", m.Type(), err)
        }
        if ok, isOK := resp.(OK); isOK {
                if err := s.checkRequestOKTrackProperties(m, resp); err != nil {
                        _ = stream.Close()
                        return zero, err
                }
                r, err := onOK(stream, ok)
                if err != nil {
                        return r, err
                }
                // §10.2.1, checked after onOK: for a SUBSCRIBE it registers the Track
                // Alias the publisher may already be sending on, so nothing may delay it.
                if err := s.CheckPeerParams(message.ScopeOfResponse(m.Type()), resp); err != nil {
                        return zero, err
                }
                return r, nil
        }
        _ = stream.Close()
        if rerr, isErr := resp.(*message.RequestError); isErr {
                return zero, s.rejection(rerr, m.Type())
        }
        err = fmt.Errorf("moqt/session: unexpected %s in %s response", resp.Type(), m.Type())
        // §10.19, §10.20: "If the subscriber receives any message other than a
        // REQUEST_OK or a REQUEST_ERROR as the first message on the response half
        // of the stream, then it MUST close the session with a PROTOCOL_VIOLATION."
        if t := m.Type(); t == message.TypeSubscribeNamespace || t == message.TypeSubscribeTracks {
                return zero, s.closeProtocolViolation(err)
        }
        return zero, err
}

// responseFirst reports whether the response to a request of type t MUST be
// the first message on its stream: to SUBSCRIBE_NAMESPACE and SUBSCRIBE_TRACKS
// (§6.1), and to PUBLISH_NAMESPACE (§6.2).
func responseFirst(t message.Type) bool {
        return t == message.TypePublishNamespace || t == message.TypeSubscribeNamespace ||
                t == message.TypeSubscribeTracks
}

// replayStream is a request stream whose reads begin with a message already
// read off it. A handle's Stream is one after an early GOAWAY.
type replayStream struct {
        Stream

        r io.Reader
}

func (s *replayStream) Read(p []byte) (int, error) { return s.r.Read(p) }

// replayMessage returns stream with m put back ahead of what it has left.
func replayMessage(stream Stream, m message.Message) Stream {
        var buf bytes.Buffer
        _ = message.Marshal(&buf, m) // a message just parsed re-encodes
        return &replayStream{Stream: stream, r: io.MultiReader(&buf, stream)}
}

// UpdateRequest sends a REQUEST_UPDATE (§10.9) with a fresh Request ID (§10.1)
// on an established request stream and awaits its REQUEST_OK or REQUEST_ERROR
// (*RequestRejectedError). params carries only the fields to change. The
// stream is left open either way.
//
// UpdateRequest reads the response directly, so it MUST NOT run concurrently
// with another reader of the stream; use [RequestBroker.Update] when the
// stream needs a standing reader.
func (s *Session) UpdateRequest(
        ctx context.Context,
        stream Stream,
        params message.Parameters,
) (*message.RequestOK, error) {
        if err := message.Marshal(stream, &message.RequestUpdate{
                RequestID:  s.AllocRequestID(),
                Parameters: params,
        }); err != nil {
                return nil, fmt.Errorf("moqt/session: write REQUEST_UPDATE: %w", err)
        }
        resp, err := s.readResponse(ctx, stream)
        if err != nil {
                return nil, fmt.Errorf("moqt/session: read REQUEST_UPDATE response: %w", err)
        }
        return s.mapUpdateResponse(resp)
}

// CheckPeerParams checks the Message Parameters of a peer message m against
// scope (§10.2.1), §10.2's duplicate rule and the values the draft makes
// session-fatal (see [message.Parameters.CheckScope]). On a violation it
// closes the session, with KEY_VALUE_FORMATTING_ERROR for a value that does
// not parse (§1.4.3) and PROTOCOL_VIOLATION otherwise, and returns the error.
//
// The session checks the messages it reads itself; callers that read a
// request stream with [message.Parse] call it for what they read.
func (s *Session) CheckPeerParams(scope message.ParamScope, m message.Message) error {
        params, ok := message.ParamsOf(m)
        if !ok {
                return nil
        }
        if err := params.CheckScope(scope); err != nil {
                if errors.Is(err, message.ErrValueFormatting) {
                        _ = s.Close(moqt.SessionKeyValueFormattingError, err.Error())
                        return err
                }
                return s.closeProtocolViolation(err)
        }
        return nil
}

// emptyPropertiesOK names the REQUEST_OK answering req when §10.5 says its
// Track Properties are empty ("they are empty in PUBLISH_OK,
// REQUEST_UPDATE_OK, SUBSCRIBE_NAMESPACE_OK and PUBLISH_NAMESPACE_OK"). A nil
// req, a SUBSCRIBE or a FETCH means a REQUEST_UPDATE_OK. SUBSCRIBE_TRACKS
// reports false: req alone cannot tell its first OK from an update's, so
// callers that can pass nil for the latter.
func emptyPropertiesOK(req message.Message) (string, bool) {
        switch req.(type) {
        case nil, *message.Subscribe, *message.Fetch:
                return "REQUEST_UPDATE_OK", true
        case *message.Publish:
                return "PUBLISH_OK", true
        case *message.PublishNamespace:
                return "PUBLISH_NAMESPACE_OK", true
        case *message.SubscribeNamespace:
                return "SUBSCRIBE_NAMESPACE_OK", true
        }
        return "", false
}

// checkRequestOKTrackProperties closes the session with PROTOCOL_VIOLATION
// when a received REQUEST_OK answering req (nil for a REQUEST_UPDATE) carries
// Track Properties §10.5 says are empty.
func (s *Session) checkRequestOKTrackProperties(req, resp message.Message) error {
        ok, isOK := resp.(*message.RequestOK)
        if !isOK || len(ok.TrackProperties) == 0 {
                return nil
        }
        name, empty := emptyPropertiesOK(req)
        if !empty {
                return nil
        }
        return s.closeProtocolViolation(fmt.Errorf("moqt/session: Track Properties in %s", name))
}

// Reply marshals a response message onto the request's stream and leaves it
// open. Use RejectError or Stream.Close to end the send direction.
//
// A REQUEST_OK carrying Track Properties where §10.5 says they are empty is
// refused with [ErrTrackPropertiesNotAllowed] and nothing is written. On a
// SUBSCRIBE_TRACKS stream, every REQUEST_OK after the first sent through Reply
// counts as a REQUEST_UPDATE_OK.
func (r *Request) Reply(msg message.Message) error {
        ok, isOK := msg.(*message.RequestOK)
        if isOK && len(ok.TrackProperties) > 0 {
                answering := r.First
                if _, st := answering.(*message.SubscribeTracks); st && r.okSent.Load() {
                        answering = nil // a REQUEST_UPDATE_OK
                }
                if name, empty := emptyPropertiesOK(answering); empty {
                        return fmt.Errorf("%w: %s", ErrTrackPropertiesNotAllowed, name)
                }
        }
        if err := message.Marshal(r.Stream, msg); err != nil {
                return err
        }
        if isOK {
                r.okSent.Store(true)
        }
        return nil
}

// RejectError writes a REQUEST_ERROR with Retry Interval 0 (§10.6.2), stops
// reading with CANCELLED (§3.3.4) and FINs the stream (§3.3.3). If the
// REQUEST_ERROR cannot be written, the stream is reset with INTERNAL_ERROR
// instead so the requester is not left waiting. Use [Request.Reject] to invite
// a retry.
func (r *Request) RejectError(code moqt.RequestErrorCode, reason string) error {
        return r.Reject(&RequestRejectedError{Code: code, Reason: reason})
}

// Reject is [Request.RejectError] with rej's Code, Reason, RetryInterval and,
// for REDIRECT, Redirect (§10.6.2). Nothing is written, and an error returned,
// when the Redirect does not match the code, or is one the peer MUST close the
// session for (§10.6.1): a Connect URI sent to a server, or a Track Name for a
// namespace-scoped request.
func (r *Request) Reject(rej *RequestRejectedError) error {
        m := &message.RequestError{
                ErrorCode:     rej.Code,
                RetryInterval: rej.RetryInterval,
                ErrorReason:   rej.Reason,
                Redirect:      rej.Redirect,
        }
        if err := m.ValidateRedirect(); err != nil {
                return fmt.Errorf("moqt/session: Reject: %w", err)
        }
        // Refuse what the peer would have to close the session for.
        if rej.Redirect != nil {
                if err := redirectViolation(rej.Redirect, r.First.Type(), r.s.role == roleClient); err != nil {
                        return fmt.Errorf("moqt/session: Reject: %w", err)
                }
        }
        if err := message.Marshal(r.Stream, m); err != nil {
                resetStream(r.Stream)
                return err
        }
        // §3.3.4: "SHOULD use a relevant error code". The request ended, so this
        // side stopped reading because it was cancelled, not because it failed.
        r.Stream.CancelRead(uint64(moqt.StreamResetCancelled))
        return r.Stream.Close()
}

// AcceptSubscribe accepts an inbound SUBSCRIBE (§10.7): it writes
// SUBSCRIBE_OK and returns a [Publication] bound to its Track Alias. r.First
// MUST be a *message.Subscribe.
//
// ok may be nil for the all-default reply; a zero TrackAlias is allocated with
// [Session.AllocOutboundTrackAlias].
func (r *Request) AcceptSubscribe(ok *message.SubscribeOK) (*Publication, error) {
        sub, isSub := r.First.(*message.Subscribe)
        if !isSub {
                return nil, fmt.Errorf("moqt/session: AcceptSubscribe on a %s request", r.First.Type())
        }
        if ok == nil {
                ok = &message.SubscribeOK{}
        }
        if ok.TrackAlias == 0 {
                ok.TrackAlias = r.s.AllocOutboundTrackAlias()
        }
        if err := message.Marshal(r.Stream, ok); err != nil {
                return nil, fmt.Errorf("moqt/session: write SUBSCRIBE_OK: %w", err)
        }
        return newPublication(r.s, r.Stream, sub.RequestID, ok.TrackAlias, sub.Parameters, ok.Parameters), nil
}

// AcceptPublish accepts an inbound PUBLISH (§10.11): it registers the Track
// Alias (§11.1), replies REQUEST_OK and returns an [IncomingPublication].
// r.First MUST be a *message.Publish.
//
// Track Properties that fail validation (see
// [WithKnownMandatoryTrackProperties]) are rejected with REQUEST_ERROR —
// UNSUPPORTED_EXTENSION for an unknown Mandatory Track Property (§2.5.1),
// INTERNAL_ERROR for ones that do not parse — and the error returned. A
// session-fatal value (§12.5, §12.6) closed the session in AcceptRequest. An
// alias collision closes the session with DUPLICATE_TRACK_ALIAS and returns
// *ErrDuplicateTrackAlias (§11.1).
func (r *Request) AcceptPublish() (*IncomingPublication, error) {
        pub, isPub := r.First.(*message.Publish)
        if !isPub {
                return nil, fmt.Errorf("moqt/session: AcceptPublish on a %s request", r.First.Type())
        }
        if err := r.s.validateTrackProperties(pub.TrackProperties, "PUBLISH"); err != nil {
                _ = r.RejectError(TrackPropertiesRejectCode(err), err.Error())
                return nil, err
        }
        key := track.NewKey(pub.Namespace, pub.Name)
        if err := r.s.RegisterInboundTrack(pub.TrackAlias, key, pub.TrackProperties); err != nil {
                return nil, err
        }
        if err := message.Marshal(r.Stream, &message.RequestOK{}); err != nil {
                // Never Established (§5.1), so the alias is not held (§11.1).
                r.s.UnregisterInboundTrackAlias(pub.TrackAlias)
                return nil, fmt.Errorf("moqt/session: write PUBLISH REQUEST_OK: %w", err)
        }
        // The publisher may send REQUEST_UPDATE (§10.9) and PUBLISH_STATE_NOTIFY
        // (§10.10).
        return &IncomingPublication{
                Stream:       r.Stream,
                s:            r.s,
                requestID:    pub.RequestID,
                peerUpdate:   true,
                peerNotify:   true,
                updateScope:  message.ScopeUpdateFromPublisher,
                inboundAlias: pub.TrackAlias,
                holdsAlias:   true,
                alias:        pub.TrackAlias,
        }, nil
}

package session

import (
        "context"
        "sync"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
)

// RequestHandler handles one inbound request that a [RequestMux] routed to it by
// the [message.Type] of its first message. It is invoked synchronously by
// [RequestMux.Run]; spawn a goroutine inside it when a request must be serviced
// concurrently with accepting the next one (see [RequestMux.Run]).
type RequestHandler func(*Request)

// RequestMux routes the requests accepted from a [Session] to per-type handlers,
// replacing the hand-rolled "AcceptRequest loop + type-switch + dispatch" a
// server otherwise writes. It is the request-stream counterpart of [Demux],
// which does the same for inbound data streams.
//
// Requests are dispatched by the [message.Type] of their first message — e.g.
// [message.TypeSubscribe] for an inbound SUBSCRIBE. A request whose type has no
// registered handler is passed to the OnUnknown callback.
//
// Handlers may be registered or replaced at any time, including while
// [RequestMux.Run] is executing. Registration is safe for concurrent use.
//
// The zero value is not ready for use — construct with [NewRequestMux].
type RequestMux struct {
        mu        sync.RWMutex
        handlers  map[message.Type]RequestHandler
        onUnknown func(*Request)
}

// NewRequestMux returns an empty RequestMux ready for handler registration.
func NewRequestMux() *RequestMux {
        return &RequestMux{handlers: make(map[message.Type]RequestHandler)}
}

// Handle registers h for inbound requests whose first message is of type t
// (e.g. [message.TypeSubscribe]). A nil h unregisters t; registering a type
// that already has a handler replaces it.
func (m *RequestMux) Handle(t message.Type, h RequestHandler) {
        m.mu.Lock()
        defer m.mu.Unlock()
        if h == nil {
                delete(m.handlers, t)
                return
        }
        m.handlers[t] = h
}

// HandleType registers h for inbound requests whose first message is the
// concrete type T (e.g. *message.Subscribe), handing h the already-asserted
// typed message alongside the [*Request]. It is the generic form of
// [RequestMux.Handle]: the [message.Type] key is derived from T, and the type
// assertion a Handle callback would otherwise repeat (req.First.(*message.X)) is
// done once, here.
//
// A nil h unregisters T's type; registering a type that already has a handler
// replaces it.
func (m *RequestMux) HandleType[T message.WithRequestID](h func(*Request, T)) {
        var zero T // nil pointer; message Type() methods are constant returns
        if h == nil {
                m.Handle(zero.Type(), nil)
                return
        }
        m.Handle(zero.Type(), func(req *Request) {
                msg, _ := req.First.(T)
                h(req, msg)
        })
}

// OnUnknown sets the callback invoked for an accepted request whose type has no
// registered handler. With no callback set (the default, or a nil f), an
// unmatched request is rejected with REQUEST_ERROR NOT_SUPPORTED and its stream
// FIN'd so it does not leak.
func (m *RequestMux) OnUnknown(f func(*Request)) {
        m.mu.Lock()
        defer m.mu.Unlock()
        m.onUnknown = f
}

// Run accepts requests from sess and dispatches each to its registered handler
// until ctx is cancelled or [Session.AcceptRequest] returns an error, which Run
// returns.
//
// Run surfaces an AcceptRequest error unchanged. Any but ctx's means the
// session has ended: a protocol violation (§3.3, §10, §10.1, §10.2.2) arrives
// already closed with the mapped code (see [Session.AcceptRequest]).
//
// Dispatch is synchronous: a handler runs to completion before Run accepts the
// next request, mirroring a hand-written accept loop and [Demux.Run]. A handler
// that keeps a request stream open for the lifetime of a subscription therefore
// blocks the loop, so spawn a goroutine inside the handler when requests must be
// serviced concurrently.
func (m *RequestMux) Run(ctx context.Context, sess *Session) error {
        for {
                req, err := sess.AcceptRequest(ctx)
                if err != nil {
                        return err
                }
                m.dispatch(req)
        }
}

// dispatch routes one accepted request to its registered handler, or to the
// unknown path when none matches.
func (m *RequestMux) dispatch(req *Request) {
        m.mu.RLock()
        h := m.handlers[req.First.Type()]
        f := m.onUnknown
        m.mu.RUnlock()
        if h != nil {
                h(req)
                return
        }
        if f != nil {
                f(req)
                return
        }
        _ = req.RejectError(moqt.RequestNotSupported, "moqt/session: no handler for request type")
}

package session

import (
        "bytes"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// reservedDot is a Track Namespace first field of exactly "." (0x2e), which
// §3.2.1 reserves for no purpose.
var reservedDot = []byte{0x2e}

// sessionNamespace is the ".session" first field (§3.2.2) MOQT reserves for
// session-level tracks and namespaces managed by the implementation.
var sessionNamespace = []byte(".session")

// requestNamespace returns the Track Namespace carried by a request's first
// message, or ok=false for a message that is not a request opener.
func requestNamespace(msg message.Message) (ns wire.TrackNamespace, ok bool) {
        switch m := msg.(type) {
        case *message.Subscribe:
                return m.Namespace, true
        case *message.Publish:
                return m.Namespace, true
        case *message.TrackStatus:
                return m.Namespace, true
        case *message.PublishNamespace:
                return m.Namespace, true
        case *message.SubscribeNamespace:
                return m.TrackNamespacePrefix, true
        case *message.SubscribeTracks:
                return m.TrackNamespacePrefix, true
        case *message.Fetch:
                return m.Namespace, true
        default:
                return nil, false
        }
}

// reservedNamespaceRejection classifies a request's Track Namespace against the
// §3.2.1 / §3.2.2 reserved-namespace rules and reports whether the request MUST
// be rejected with DOES_NOT_EXIST before the application ever sees it. The
// decision keys on the first namespace tuple field:
//
//   - exactly "." (§3.2.1): reserved for no purpose — reject.
//   - ".session" (§3.2.2): the session-level namespace, owned by the MOQT
//     implementation rather than the application. This library implements no
//     session-level tracks, so every such request is "unrecognized" and MUST be
//     rejected with DOES_NOT_EXIST — which also subsumes the §3.2.2 rule that a
//     ".session" namespace with an empty Track Name does not exist. A future
//     session-level extension would dispatch its recognized tracks here instead
//     of rejecting.
//   - any other "."-prefixed value (§3.2.1): an unrecognized reserved namespace
//     that MUST be passed to the application so future extensions don't break
//     older implementations — so it is NOT rejected here.
//   - anything else: an ordinary namespace — not rejected.
func reservedNamespaceRejection(msg message.Message) (reason string, reject bool) {
        ns, ok := requestNamespace(msg)
        if !ok || len(ns) == 0 {
                return "", false
        }
        switch first := ns[0]; {
        case bytes.Equal(first, reservedDot):
                return `reserved namespace "." (§3.2.1)`, true
        case bytes.Equal(first, sessionNamespace):
                return "unrecognized session-level namespace (§3.2.2)", true
        default:
                return "", false
        }
}

// Package session implements the MoQT session layer: SETUP handshake, control
// stream multiplexing, request-ID allocation, and graceful termination via
// GOAWAY (§3.3, §3.5, §10.3, §10.4 of draft-ietf-moq-transport-20).
//
// The package does not depend on a specific transport. It operates against the
// Conn interface, which any QUIC-like transport can satisfy.
package session

import (
        "context"
        "errors"
        "fmt"
        "sync"
        "sync/atomic"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/uri"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// role identifies whether this endpoint initiated (client) or accepted
// (server) the underlying QUIC connection. role determines Request ID parity
// per §10.1: client IDs are even, server IDs are odd. The type is unexported
// because callers select a role by calling Client or Server rather than
// passing a value.
type role uint8

const (
        roleClient role = iota
        roleServer
)

// Session represents one MoQT session over a Conn after the SETUP handshake
// has completed. The Session owns the control-stream goroutines until Close
// is called.
type Session struct {
        conn Conn
        role role

        sendCtrl SendStream
        recvCtrl ReceiveStream

        peerOptions []wire.KVPair

        // earlyData holds data streams that arrived before the peer's control
        // stream (§3.3); see acceptControlStream.
        earlyMu   sync.Mutex
        earlyData []ReceiveStream

        // See [Session.SetupTokenAliases] and [Session.SetupTokens]. Written once
        // during open.
        setupTokenAliases []uint64
        setupTokens       []ResolvedToken

        // Outgoing Request ID allocator: client starts at 0 (even), server at 1
        // (odd); each AllocRequestID advances by 2 (§10.1).
        nextRequestID atomic.Uint64

        // Outgoing Track Alias allocator. §11.1: aliases are scoped to the
        // publisher → subscriber direction of one session, so each end keeps an
        // independent counter for tracks it advertises to the peer. The first
        // allocation is 1, not 0: AllocOutboundTrackAlias reserves 0 as the
        // "unset, auto-allocate" sentinel used by Publish/OpenPublish/Reply (see
        // AllocOutboundTrackAlias). The spec does not constrain parity the way it
        // does for Request IDs.
        nextOutboundTrackAlias atomic.Uint64

        // Serialized writes onto the control stream. Producers send through
        // sendControl; the controlSendLoop drains and writes.
        controlOut chan message.Message

        mu             sync.Mutex
        goawayReceived *message.Goaway
        goawaySent     bool
        // goawayCh is closed when goawayReceived transitions from nil to set.
        goawayCh chan struct{}
        // goawayHandler is the optional callback registered via OnGoaway. It is
        // invoked exactly once, in its own goroutine, when the first GOAWAY
        // arrives from the peer. goawayFired guards the at-most-once invocation
        // across the handleGoaway and OnGoaway paths. Both are protected by mu.
        goawayHandler func(*message.Goaway)
        goawayFired   bool

        // Inbound Request ID tracking (§10.1). Protected by mu.
        // peerRequestIDSeen is false until the first inbound Request ID arrives;
        // peerRequestIDMax is the high-water mark. The peer allocates IDs in +2
        // increments, but requests ride separate QUIC streams and can be
        // DELIVERED out of order, so an ID below the mark is not automatically a
        // duplicate: peerRequestIDGaps holds the not-yet-seen IDs below the mark
        // (bounded by maxTrackedRequestIDGaps) that a late-delivered request may
        // still legitimately claim. See [Session.CheckPeerRequestID].
        peerRequestIDSeen bool
        peerRequestIDMax  uint64
        peerRequestIDGaps map[uint64]struct{}

        // Inbound Track Alias → track mapping (§11.1). Protected by mu.
        // Populated via RegisterInboundTrack when the peer assigns an alias
        // (SUBSCRIBE_OK or PUBLISH). A duplicate alias for a different track is
        // a DUPLICATE_TRACK_ALIAS session error.
        inboundAliases map[uint64]InboundTrack

        // inboundAliasRefs counts the registrations of each alias. Protected by
        // mu.
        inboundAliasRefs map[uint64]int

        // aliasRegistered is closed and replaced each time a new alias is bound,
        // waking awaitInboundTrack. Protected by mu.
        aliasRegistered chan struct{}

        // knownMandatoryTrackProperties is the set of Mandatory Track Property
        // types (range 0x4000–0x7FFF) this endpoint supports. Configured via
        // WithKnownMandatoryTrackProperties. nil means none are known.
        knownMandatoryTrackProperties map[message.PropertyType]struct{}

        // tokenCache is the inbound authorization-token alias cache (§10.2.2).
        // AcceptRequest drives Register / Delete / Resolve on it from the
        // AUTHORIZATION_TOKEN parameters of each inbound request. Always non-nil
        // (sized 0 when MAX_AUTH_TOKEN_CACHE_SIZE was not negotiated, which
        // prohibits aliasing per §10.3.1.3). The cache has its own internal
        // mutex; it is not protected by mu.
        tokenCache *TokenCache

        // tokenVerifier is the optional application policy consulted for each
        // resolved token. nil disables verification. Set once at construction
        // via WithTokenVerifier; never mutated, so no lock is required.
        tokenVerifier TokenVerifier

        // maxRequestUpdates is the per-request-stream limit on unacknowledged
        // inbound REQUEST_UPDATEs we advertised via MAX_REQUEST_UPDATES
        // (§10.3.1.7). 0 means unlimited. Set once at construction via
        // WithMaxRequestUpdates; read by NewRequestUpdateLimiter, so no lock.
        maxRequestUpdates uint64

        // maxFilterRanges is the total Range Filter range budget we advertised via
        // MAX_FILTER_RANGES (§10.3.1.6). 0 prohibits Range Filters. Set once at
        // construction via WithMaxFilterRanges; read via MaxFilterRanges(), no lock.
        maxFilterRanges uint64

        closeOnce sync.Once
        // closeErr holds the *ClosedError cause; atomic because Err may
        // be called at any time, not only after Done fires.
        closeErr atomic.Pointer[ClosedError]
        // done is closed when the session terminates for any reason.
        done chan struct{}
}

// Client performs the SETUP handshake from the client side (the initiator of
// the underlying QUIC connection) and returns a ready Session. Request IDs on
// this side are even (§10.1). If the handshake fails, conn is closed and the
// error is returned. On success the caller owns the Session and must Close it.
func Client(ctx context.Context, conn Conn, opts ...Option) (*Session, error) {
        return open(ctx, conn, opts, roleClient)
}

// Server performs the SETUP handshake from the server side (the acceptor of
// the underlying QUIC connection) and returns a ready Session. Request IDs on
// this side are odd (§10.1). Errors and ownership match Client.
func Server(ctx context.Context, conn Conn, opts ...Option) (*Session, error) {
        return open(ctx, conn, opts, roleServer)
}

func open(ctx context.Context, conn Conn, opts []Option, r role) (*Session, error) {
        var cfg config
        for _, o := range opts {
                o(&cfg)
        }

        s := &Session{
                conn: conn,
                role: r,
                // Control-stream traffic is sparse (only GOAWAY after SETUP per §10
                // table 5, and at most one outbound GOAWAY per session). A buffer of
                // 1 lets a sender hand off while the previous frame is being written
                // to the transport; anything larger just delays backpressure.
                controlOut:                    make(chan message.Message, 1),
                goawayCh:                      make(chan struct{}),
                done:                          make(chan struct{}),
                inboundAliases:                make(map[uint64]InboundTrack),
                inboundAliasRefs:              make(map[uint64]int),
                aliasRegistered:               make(chan struct{}),
                knownMandatoryTrackProperties: cfg.knownMandatoryTrackProperties,
                tokenCache:                    NewTokenCache(cfg.maxAuthTokenCacheSize),
                tokenVerifier:                 cfg.tokenVerifier,
                maxRequestUpdates:             cfg.maxRequestUpdates,
                maxFilterRanges:               cfg.maxFilterRanges,
        }
        var first uint64
        if r == roleServer {
                first = 1
        }
        s.nextRequestID.Store(first)

        // Before the handshake, so an option we must not send is never sent. The
        // conn is left untouched and un-closed: nothing has been written to it, so
        // disposing of it stays the caller's choice.
        if err := checkOutboundSetupOptions(r, conn, cfg.setupOptions); err != nil {
                return nil, err
        }
        if err := checkOutboundSetupTokens(cfg.setupTokens); err != nil {
                return nil, err
        }
        if err := s.handshake(ctx, cfg.setupOptions); err != nil {
                _ = conn.CloseWithError(uint64(moqt.SessionProtocolViolation), err.Error())
                return nil, err
        }
        if code, err := s.checkPeerSetupOptions(); err != nil {
                _ = conn.CloseWithError(uint64(code), err.Error())
                return nil, err
        }
        if err := s.processSetupTokens(); err != nil {
                tce, _ := errors.AsType[*TokenCacheError](err)
                _ = conn.CloseWithError(uint64(tce.Code), err.Error())
                return nil, err
        }
        s.setupTokenAliases = heldSetupAliases(cfg.setupTokens, s.peerOptions)

        go s.controlSendLoop()
        go s.controlRecvLoop()

        return s, nil
}

// PeerOptions returns the SETUP options the peer advertised. The returned
// slice aliases internal state and must not be mutated.
func (s *Session) PeerOptions() []wire.KVPair { return s.peerOptions }

// webTransportConn is an optional capability a Conn adapter may implement to
// report that it runs over WebTransport. An adapter that does not implement it
// is treated as native QUIC, so the checks below are best-effort — they cover
// the wtconn adapter in this repo, not a third-party WebTransport adapter that
// stays silent.
//
// Deliberately not a Conn method: only the §10.3.1 PATH/AUTHORITY rules need
// it, and per CLAUDE.md anything added to that interface must land in all three
// adapters plus any external one.
type webTransportConn interface {
        IsWebTransport() bool
}

func overWebTransport(conn Conn) bool {
        wt, ok := conn.(webTransportConn)
        return ok && wt.IsWebTransport()
}

// checkOutboundSetupOptions enforces the send side of PATH (§10.3.1.2) and
// AUTHORITY (§10.3.1.1). Each says the option "MUST NOT be used by the server,
// or when WebTransport is used" — so unlike checkPeerSetupOptions, which
// handles what a peer did to us, this catches what we are about to do to a
// peer, before the SETUP goes out.
//
// It fails the open rather than dropping the offending option, because silently
// discarding it would leave a client believing it had requested an authority
// the server never saw.
func checkOutboundSetupOptions(r role, conn Conn, opts []wire.KVPair) error {
        for _, opt := range opts {
                var name string
                switch message.SetupOption(opt.Type) {
                case message.SetupOptionPath:
                        name = "PATH"
                case message.SetupOptionAuthority:
                        name = "AUTHORITY"
                case message.SetupOptionAuthorizationToken,
                        message.SetupOptionMaxAuthTokenCache,
                        message.SetupOptionMaxFilterRanges,
                        message.SetupOptionMOQTImplementation,
                        message.SetupOptionMaxRequestUpdates:
                        continue
                default:
                        continue
                }
                if r == roleServer {
                        return fmt.Errorf("moqt/session: %s setup option is client-only (§10.3.1)", name)
                }
                if overWebTransport(conn) {
                        return fmt.Errorf(
                                "moqt/session: %s setup option must not be used on a WebTransport session (§10.3.1); "+
                                        "HTTP/3 carries the path and authority in the CONNECT request", name)
                }
        }
        // The server would close the session over a non-RFC 3986 value.
        if _, err := checkPathAndAuthoritySyntax(opts); err != nil {
                return fmt.Errorf("moqt/session: %w", err)
        }
        return nil
}

// checkPeerSetupOptions enforces the receive side of the PATH (§10.3.1.2) and
// AUTHORITY (§10.3.1.1) setup options. Each names three conditions under which
// a received option MUST close the session — it came from a server, it arrived
// on a WebTransport session, or the server does not support the value — with
// INVALID_PATH and INVALID_AUTHORITY respectively, whose numeric values come
// from the §3.5 registry. Returns that close code and the reason, or a nil
// error when the peer's options are acceptable.
//
// The first two are enforced. The third is not: it needs the server to be told
// which paths and authorities it serves, and that configuration does not exist
// — a server that never learns its own names cannot check them.
//
// Note the asymmetry with the role gate. A server ignores these options over
// native QUIC because receiving them there is precisely what they are for; it
// must still reject them over WebTransport, where §10.3.1 forbids them
// outright and HTTP/3 carries the same information in the CONNECT request.
func (s *Session) checkPeerSetupOptions() (moqt.SessionErrorCode, error) {
        overWT := overWebTransport(s.conn)
        if s.role != roleClient && !overWT {
                return checkPathAndAuthoritySyntax(s.peerOptions)
        }
        violation := func(name string) error {
                if overWT {
                        return fmt.Errorf(
                                "peer sent a %s setup option on a WebTransport session (§10.3.1)", name)
                }
                return fmt.Errorf("server sent a %s setup option, which is client-only (§10.3.1)", name)
        }
        for _, opt := range s.peerOptions {
                switch message.SetupOption(opt.Type) {
                case message.SetupOptionPath:
                        return moqt.SessionInvalidPath, violation("PATH")
                case message.SetupOptionAuthority:
                        return moqt.SessionInvalidAuthority, violation("AUTHORITY")
                case message.SetupOptionAuthorizationToken,
                        message.SetupOptionMaxAuthTokenCache,
                        message.SetupOptionMaxFilterRanges,
                        message.SetupOptionMOQTImplementation,
                        message.SetupOptionMaxRequestUpdates:
                        // Legal in both directions. Listed rather than folded into the
                        // default so that adding a §10.3.1 option fails the exhaustive
                        // linter until someone decides whether a server may send it.
                default:
                        // §10.3 requires a receiver ignore options it does not recognize.
                }
        }
        return moqt.SessionNoError, nil
}

// checkPathAndAuthoritySyntax checks PATH (§10.3.1.2) and AUTHORITY
// (§10.3.1.1) against RFC 3986, returning MALFORMED_PATH or
// MALFORMED_AUTHORITY as the close code.
func checkPathAndAuthoritySyntax(opts []wire.KVPair) (moqt.SessionErrorCode, error) {
        for _, opt := range opts {
                switch message.SetupOption(opt.Type) {
                case message.SetupOptionPath:
                        if err := uri.CheckPathAndQuery(string(opt.ByteVal)); err != nil {
                                return moqt.SessionMalformedPath, fmt.Errorf("malformed PATH setup option: %w", err)
                        }
                case message.SetupOptionAuthority:
                        if err := uri.CheckAuthority(string(opt.ByteVal)); err != nil {
                                return moqt.SessionMalformedAuthority, fmt.Errorf("malformed AUTHORITY setup option: %w", err)
                        }
                case message.SetupOptionAuthorizationToken,
                        message.SetupOptionMaxAuthTokenCache,
                        message.SetupOptionMaxFilterRanges,
                        message.SetupOptionMOQTImplementation,
                        message.SetupOptionMaxRequestUpdates:
                default:
                }
        }
        return moqt.SessionNoError, nil
}

// MaxFilterRanges returns the MAX_FILTER_RANGES value this session advertised
// (§10.3.1.6) — the total Range Filter range budget it will accept on any one
// subscription or fetch. 0 prohibits Range Filters. The relay enforces this
// when validating a request's Range Filters against [message.RangeFilterSet.Validate].
func (s *Session) MaxFilterRanges() uint64 { return s.maxFilterRanges }

// AllocRequestID returns the next outbound Request ID per §10.1.
func (s *Session) AllocRequestID() uint64 {
        return s.nextRequestID.Add(2) - 2
}

// Done returns a channel that is closed when the session has terminated.
func (s *Session) Done() <-chan struct{} { return s.done }

// Err returns the close cause — a *ClosedError carrying the §3.5
// error code and reason of the first Close call — or nil when the session
// was closed cleanly (SessionNoError) or is still open. The value is
// published before Done is closed, so the natural pattern
// <-sess.Done(); sess.Err() is race-free.
func (s *Session) Err() error {
        // Explicit nil check: returning a nil *ClosedError directly
        // would produce a non-nil error interface.
        if e := s.closeErr.Load(); e != nil {
                return e
        }
        return nil
}

// ClosedError is the close cause stored by [Session.Close] and
// returned by [Session.Err] for a non-clean close.
type ClosedError struct {
        Code   moqt.SessionErrorCode
        Reason string
}

func (e *ClosedError) Error() string {
        return fmt.Sprintf("moqt/session: closed with code %#x: %s", uint64(e.Code), e.Reason)
}

// Close terminates the session, cancelling the control streams and closing
// the underlying connection with the given code (§3.5). Calling Close more
// than once is safe; only the first call's code takes effect. The returned
// error is the transport's close error (nil in the common case), NOT the
// close cause — that is what [Session.Err] reports.
func (s *Session) Close(code moqt.SessionErrorCode, reason string) error {
        var transportErr error
        s.closeOnce.Do(func() {
                // Publish the cause BEFORE closing done: Err must be safe to call
                // the moment Done fires.
                if code != moqt.SessionNoError {
                        s.closeErr.Store(&ClosedError{Code: code, Reason: reason})
                }
                close(s.done)
                // CancelRead first so the recv loop unblocks. The control stream
                // must not be FIN'd cleanly during session lifetime (§3.3); we
                // reset both directions instead.
                if s.recvCtrl != nil {
                        s.recvCtrl.CancelRead(uint64(moqt.StreamResetSessionClosed))
                }
                if s.sendCtrl != nil {
                        s.sendCtrl.CancelWrite(uint64(moqt.StreamResetSessionClosed))
                }
                transportErr = s.conn.CloseWithError(uint64(code), reason)
        })
        return transportErr
}

package sessiontest

import (
        "io"
        "net"
        "sync"
)

// pipeReadCloser / pipeWriteCloser are the two halves of an in-process pipe.
// Both *io.PipeReader/*io.PipeWriter (synchronous) and *bufPipe's reader/writer
// (buffered) satisfy them, so [uniStream] / [bidiStream] can be backed by
// either without caring which.
type pipeReadCloser interface {
        io.Reader
        CloseWithError(error) error
}

type pipeWriteCloser interface {
        io.Writer
        Close() error
        CloseWithError(error) error
}

// newPipe returns a connected reader/writer pair. With bufSize <= 0 it returns
// a synchronous io.Pipe (the historical default — every Write blocks until a
// Read drains it). With bufSize > 0 it returns a [bufPipe] of that capacity,
// which lets the writer run ahead and is what the throughput benchmarks use to
// avoid measuring per-object goroutine scheduling instead of forwarding work.
func newPipe(bufSize int) (pipeReadCloser, pipeWriteCloser) {
        if bufSize > 0 {
                return newBufPipe(bufSize)
        }
        return io.Pipe()
}

// bufPipe is a bounded, buffered, in-memory byte pipe — a drop-in alternative
// to io.Pipe for the sessiontest transport. io.Pipe is fully synchronous: every
// Write blocks until a Read consumes it, forcing a goroutine handoff per write,
// so a relay/session benchmark over it spends ~85% of its CPU in the scheduler
// (usleep / cond_signal / cond_wait) rather than in forwarding code. bufPipe
// lets the writer run ahead by up to `cap` buffered bytes before blocking, so
// producer and consumer wake in bursts instead of lock-stepping per object.
//
// Semantics otherwise match io.Pipe:
//   - a clean writer Close surfaces as io.EOF to the reader, but only after the
//     buffered bytes have been drained;
//   - CloseWithError on either half unblocks the other half with that error.
type bufPipe struct {
        mu       sync.Mutex
        notEmpty sync.Cond
        notFull  sync.Cond

        buf  []byte // ring buffer
        r, n int    // read index, number of bytes currently buffered

        rerr error // set when the reader closes; returned to the writer
        werr error // set when the writer closes; returned to the reader once drained
}

func newBufPipe(capacity int) (pipeReadCloser, pipeWriteCloser) {
        bp := &bufPipe{buf: make([]byte, capacity)}
        bp.notEmpty.L = &bp.mu
        bp.notFull.L = &bp.mu
        return bufPipeReader{bp}, bufPipeWriter{bp}
}

// write copies all of p into the ring, blocking while the buffer is full. It
// returns early with rerr if the reader has gone away (mirrors io.Pipe writing
// to a closed reader), and with werr when the WRITE side itself was already
// closed — real QUIC rejects writes after FIN, and silently buffering them
// here would let tests pass flows production fails on.
func (bp *bufPipe) write(p []byte) (int, error) {
        bp.mu.Lock()
        defer bp.mu.Unlock()
        total := 0
        for len(p) > 0 {
                if bp.werr != nil {
                        return total, net.ErrClosed
                }
                for bp.n == len(bp.buf) && bp.rerr == nil && bp.werr == nil {
                        bp.notFull.Wait()
                }
                if bp.rerr != nil {
                        return total, bp.rerr
                }
                if bp.werr != nil {
                        return total, net.ErrClosed
                }
                // Copy into the contiguous free region starting at the write index,
                // stopping at the buffer end (the next iteration handles the wrap).
                w := (bp.r + bp.n) % len(bp.buf)
                chunk := min(len(p), len(bp.buf)-bp.n, len(bp.buf)-w)
                copy(bp.buf[w:w+chunk], p[:chunk])
                bp.n += chunk
                total += chunk
                p = p[chunk:]
                bp.notEmpty.Signal()
        }
        return total, nil
}

// read drains up to len(p) bytes, blocking while the buffer is empty. Once the
// writer has closed and the buffer is empty it returns werr (io.EOF on a clean
// close).
//
// The two half-closes are deliberately asymmetric. A writer close is reported
// only after the buffer drains — that is what makes this a buffered pipe, and
// it matches io.Pipe for a clean Close. A reader close takes effect at once,
// buffered bytes or not: it models QUIC STOP_SENDING, and a stream that kept
// handing back objects after the session reset it would let a test observe
// delivery production never performs. Checking rerr before the drain loop is
// also what gives the reader's own error precedence over a writer close that
// lands afterwards, so a reset is not reported as a clean io.EOF.
func (bp *bufPipe) read(p []byte) (int, error) {
        if len(p) == 0 {
                return 0, nil
        }
        bp.mu.Lock()
        defer bp.mu.Unlock()
        if bp.rerr != nil {
                return 0, bp.rerr
        }
        for bp.n == 0 {
                if bp.werr != nil {
                        return 0, bp.werr
                }
                if bp.rerr != nil {
                        return 0, bp.rerr
                }
                bp.notEmpty.Wait()
        }
        chunk := min(len(p), bp.n, len(bp.buf)-bp.r)
        copy(p, bp.buf[bp.r:bp.r+chunk])
        bp.r = (bp.r + chunk) % len(bp.buf)
        bp.n -= chunk
        bp.notFull.Signal()
        return chunk, nil
}

func (bp *bufPipe) closeWrite(err error) error {
        if err == nil {
                err = io.EOF
        }
        bp.mu.Lock()
        if bp.werr == nil {
                bp.werr = err
        }
        bp.notEmpty.Broadcast()
        bp.notFull.Broadcast() // wake any writer blocked on a full buffer
        bp.mu.Unlock()
        return nil
}

func (bp *bufPipe) closeRead(err error) error {
        if err == nil {
                err = io.ErrClosedPipe
        }
        bp.mu.Lock()
        if bp.rerr == nil {
                bp.rerr = err
        }
        bp.notFull.Broadcast()
        bp.notEmpty.Broadcast()
        bp.mu.Unlock()
        return nil
}

type bufPipeReader struct{ bp *bufPipe }

func (r bufPipeReader) Read(p []byte) (int, error)     { return r.bp.read(p) }
func (r bufPipeReader) CloseWithError(err error) error { return r.bp.closeRead(err) }

type bufPipeWriter struct{ bp *bufPipe }

func (w bufPipeWriter) Write(p []byte) (int, error)    { return w.bp.write(p) }
func (w bufPipeWriter) Close() error                   { return w.bp.closeWrite(io.EOF) }
func (w bufPipeWriter) CloseWithError(err error) error { return w.bp.closeWrite(err) }

package sessiontest

import (
        "context"
        "fmt"
        "sync/atomic"

        "github.com/floatdrop/moq-go/pkg/moqt/session"
)

// Op identifies the [session.Conn] or [session.Stream] operation a [FaultFunc]
// is being consulted about.
type Op int

const (
        OpOpenStream Op = iota
        OpOpenUniStream
        OpAcceptStream
        OpAcceptUniStream
        OpSendDatagram
        OpReceiveDatagram
        OpStreamWrite
        OpStreamRead
        OpStreamClose
)

// numOps sizes the per-Op counter arrays. It is deliberately an untyped
// constant rather than a trailing iota member: as an Op it would be a phantom
// enum value every exhaustive switch had to handle.
const numOps = int(OpStreamClose) + 1

func (o Op) String() string {
        switch o {
        case OpOpenStream:
                return "OpenStream"
        case OpOpenUniStream:
                return "OpenUniStream"
        case OpAcceptStream:
                return "AcceptStream"
        case OpAcceptUniStream:
                return "AcceptUniStream"
        case OpSendDatagram:
                return "SendDatagram"
        case OpReceiveDatagram:
                return "ReceiveDatagram"
        case OpStreamWrite:
                return "StreamWrite"
        case OpStreamRead:
                return "StreamRead"
        case OpStreamClose:
                return "StreamClose"
        }
        return fmt.Sprintf("Op(%d)", int(o))
}

// FaultOp describes the operation a [FaultFunc] is being consulted about.
type FaultOp struct {
        // Op is the operation about to be performed.
        Op Op

        // Stream is the ordinal of the stream the operation is on, numbered from
        // 1 in the order the conn handed streams out. Unidirectional and
        // bidirectional streams share the one sequence, and both opening and
        // accepting allocate from it. Stream is 0 for connection-level operations
        // (the opens and accepts themselves, and the datagram calls).
        //
        // Ordinals are only stable when the test controls the order streams are
        // created in. Code under test that opens streams from several goroutines
        // — a relay fanning out to subscribers, for one — does not give that;
        // match on Buf there instead.
        Stream int

        // N counts occurrences of this Op, from 1: per stream for the stream
        // operations, per conn for the connection-level ones.
        N int

        // Buf is the caller's buffer for OpStreamWrite and OpStreamRead, and the
        // payload for OpSendDatagram; nil for every other Op. On OpStreamRead it
        // is the destination buffer, which the read has not filled yet — its
        // length is the size of the read, not data. Do not retain or modify it.
        Buf []byte
}

// FaultFunc is consulted before each wrapped operation. Returning a non-nil
// error makes the operation fail with that error instead of being performed;
// returning nil lets it through untouched.
//
// It runs on whichever goroutine drives the operation — for a relay under
// test, several at once — so it must be safe for concurrent use.
type FaultFunc func(FaultOp) error

// Faulty wraps c so fault is consulted before every operation on the conn and
// on every stream the conn hands out, letting a test make a chosen write, open
// or read fail. It exists to reach the error branches a healthy in-process pipe
// never takes: the "reply failed" and "write failed" paths that in production
// run when a peer's transport goes bad.
//
// Faults are injected in front of the wrapped operation, so a failed one never
// reaches the underlying conn: a failed Write puts no bytes on the stream, and
// a failed OpenStream consumes no stream credit.
//
// Two limitations, both deliberate:
//
//   - A failed Write reports (0, err). Real QUIC can fail part-way through and
//     report a short write; nothing in this tree distinguishes the two, so the
//     wrapper does not model it.
//   - Wrapped streams do not forward the optional [session.PrioritizedSendStream],
//     [session.ReliableResetStream] and [session.DeliveryTrackingSendStream]
//     interfaces. No sessiontest stream implements any of them, and silently
//     dropping §7.2 priority, RESET_STREAM_AT or §8 SUBGROUP_DELIVERY_TIMEOUT
//     would be a confusing way to find that out, so Faulty panics rather than
//     wrap a stream that does.
func Faulty(c session.Conn, fault FaultFunc) session.Conn {
        fc := &faultyConn{Conn: c}
        fc.fault = fault
        return fc
}

// FailNth returns a [FaultFunc] failing the nth occurrence of op with err,
// counted from 1 across the whole conn. Every other operation succeeds.
func FailNth(op Op, n int, err error) FaultFunc {
        var seen atomic.Int64
        return func(f FaultOp) error {
                if f.Op != op {
                        return nil
                }
                if seen.Add(1) == int64(n) {
                        return err
                }
                return nil
        }
}

// FailAll returns a [FaultFunc] failing every occurrence of op with err.
func FailAll(op Op, err error) FaultFunc {
        return func(f FaultOp) error {
                if f.Op != op {
                        return nil
                }
                return err
        }
}

// faultyCounter holds the fault state shared by the conn and its streams: the
// hook, which stream this is (0 for the conn itself), and how many times each
// Op has been seen here. It contains atomics, so it is embedded and filled in
// place — never copied.
type faultyCounter struct {
        fault  FaultFunc
        stream int
        counts [numOps]atomic.Int64
}

func (f *faultyCounter) check(op Op, buf []byte) error {
        return f.fault(FaultOp{
                Op:     op,
                Stream: f.stream,
                N:      int(f.counts[op].Add(1)),
                Buf:    buf,
        })
}

type faultyConn struct {
        session.Conn
        faultyCounter

        streams atomic.Int64 // stream ordinal allocator
}

func (c *faultyConn) OpenStream() (session.Stream, error) {
        if err := c.check(OpOpenStream, nil); err != nil {
                return nil, err
        }
        s, err := c.Conn.OpenStream()
        if err != nil {
                return nil, err
        }
        return c.newStream(s), nil
}

func (c *faultyConn) AcceptStream(ctx context.Context) (session.Stream, error) {
        if err := c.check(OpAcceptStream, nil); err != nil {
                return nil, err
        }
        s, err := c.Conn.AcceptStream(ctx)
        if err != nil {
                return nil, err
        }
        return c.newStream(s), nil
}

func (c *faultyConn) OpenUniStream() (session.SendStream, error) {
        if err := c.check(OpOpenUniStream, nil); err != nil {
                return nil, err
        }
        s, err := c.Conn.OpenUniStream()
        if err != nil {
                return nil, err
        }
        assertPlainSend(s)
        fs := &faultySendStream{SendStream: s}
        c.initCounter(&fs.faultyCounter)
        return fs, nil
}

func (c *faultyConn) AcceptUniStream(ctx context.Context) (session.ReceiveStream, error) {
        if err := c.check(OpAcceptUniStream, nil); err != nil {
                return nil, err
        }
        s, err := c.Conn.AcceptUniStream(ctx)
        if err != nil {
                return nil, err
        }
        fs := &faultyRecvStream{ReceiveStream: s}
        c.initCounter(&fs.faultyCounter)
        return fs, nil
}

func (c *faultyConn) SendDatagram(payload []byte) error {
        if err := c.check(OpSendDatagram, payload); err != nil {
                return err
        }
        return c.Conn.SendDatagram(payload)
}

func (c *faultyConn) ReceiveDatagram(ctx context.Context) ([]byte, error) {
        if err := c.check(OpReceiveDatagram, nil); err != nil {
                return nil, err
        }
        return c.Conn.ReceiveDatagram(ctx)
}

func (c *faultyConn) newStream(s session.Stream) *faultyStream {
        assertPlainSend(s)
        fs := &faultyStream{Stream: s}
        c.initCounter(&fs.faultyCounter)
        return fs
}

// initCounter fills a stream's fault state in place and allocates its ordinal.
func (c *faultyConn) initCounter(fc *faultyCounter) {
        fc.fault = c.fault
        fc.stream = int(c.streams.Add(1))
}

// assertPlainSend panics if s implements one of the optional SendStream
// interfaces the wrappers cannot forward. See [Faulty] for why this is loud
// rather than silent.
func assertPlainSend(s any) {
        switch s.(type) {
        case session.PrioritizedSendStream:
                panic("sessiontest.Faulty: refusing to wrap a session.PrioritizedSendStream — " +
                        "the wrapper cannot forward SetSendPriority")
        case session.ReliableResetStream:
                panic("sessiontest.Faulty: refusing to wrap a session.ReliableResetStream — " +
                        "the wrapper cannot forward SetReliableBoundary")
        case session.DeliveryTrackingSendStream:
                panic("sessiontest.Faulty: refusing to wrap a session.DeliveryTrackingSendStream — " +
                        "the wrapper cannot forward Finished")
        }
}

type faultyStream struct {
        session.Stream
        faultyCounter
}

func (s *faultyStream) Write(p []byte) (int, error) {
        if err := s.check(OpStreamWrite, p); err != nil {
                return 0, err
        }
        return s.Stream.Write(p)
}

func (s *faultyStream) Read(p []byte) (int, error) {
        if err := s.check(OpStreamRead, p); err != nil {
                return 0, err
        }
        return s.Stream.Read(p)
}

func (s *faultyStream) Close() error {
        if err := s.check(OpStreamClose, nil); err != nil {
                return err
        }
        return s.Stream.Close()
}

type faultySendStream struct {
        session.SendStream
        faultyCounter
}

func (s *faultySendStream) Write(p []byte) (int, error) {
        if err := s.check(OpStreamWrite, p); err != nil {
                return 0, err
        }
        return s.SendStream.Write(p)
}

func (s *faultySendStream) Close() error {
        if err := s.check(OpStreamClose, nil); err != nil {
                return err
        }
        return s.SendStream.Close()
}

type faultyRecvStream struct {
        session.ReceiveStream
        faultyCounter
}

func (s *faultyRecvStream) Read(p []byte) (int, error) {
        if err := s.check(OpStreamRead, p); err != nil {
                return 0, err
        }
        return s.ReceiveStream.Read(p)
}

// Package sessiontest provides in-process helpers for testing MoQT session
// code without a real QUIC transport. NewConnPair returns two session.Conn
// endpoints backed by io.Pipes; streams opened on one end are accepted on
// the other. NewSessionPair goes one step further and performs the full SETUP
// handshake, returning two ready *session.Session values.
//
// The implementation deliberately mirrors a real QUIC stream's semantics
// where it matters for tests:
//
//   - Opening a stream never blocks (per the Conn contract): the stream is
//     offered to the peer's accept queue immediately, and a full queue —
//     the peer isn't accepting — surfaces as ErrNoStreamCredit rather than
//     a deadlocked opener. (Unlike quic-go, the peer can therefore see a
//     uni stream before its first byte is written.)
//   - CancelRead / CancelWrite unblock any in-flight Read / Write with an
//     error.
//   - CloseWithError cancels the shared connection context, which unblocks
//     any pending Accept on either end.
package sessiontest

import (
        "context"
        "errors"
        "sync"
        "testing"

        "github.com/floatdrop/moq-go/pkg/moqt/session"
)

// NewSessionPair performs the MoQT SETUP handshake over an in-process conn
// pair and returns two ready Sessions — client (even Request IDs) and server
// (odd Request IDs). Both sessions are closed via tb.Cleanup when the test
// or benchmark ends.
//
// The parameter is testing.TB rather than *testing.T so the helper serves
// both tests and benchmarks. Because testing.TB does not expose Context()
// (that method lives only on *testing.T / *testing.B), the handshake context
// is managed internally and cancelled via tb.Cleanup.
func NewSessionPair(tb testing.TB) (client, server *session.Session) {
        tb.Helper()
        connA, connB := NewConnPair()

        ctx, cancel := context.WithCancel(context.Background())
        tb.Cleanup(cancel)

        var (
                wg           sync.WaitGroup
                aSess, bSess *session.Session
                aErr, bErr   error
        )
        wg.Go(func() {
                aSess, aErr = session.Client(ctx, connA)
        })
        wg.Go(func() {
                bSess, bErr = session.Server(ctx, connB)
        })
        wg.Wait()
        if aErr != nil {
                tb.Fatalf("sessiontest.NewSessionPair client: %v", aErr)
        }
        if bErr != nil {
                tb.Fatalf("sessiontest.NewSessionPair server: %v", bErr)
        }
        tb.Cleanup(func() {
                _ = aSess.Close(0, "")
                _ = bSess.Close(0, "")
        })
        return aSess, bSess
}

// NewConnPair returns two session.Conn endpoints wired together in-process.
// Both endpoints have unlimited outbound bidirectional-stream credit; use
// [NewConnPairWithLimits] to cap one or both sides for PUBLISH_SKIPPED-style
// stream-exhaustion testing.
func NewConnPair() (a, b session.Conn) {
        return NewConnPairWithLimits(-1, -1)
}

// NewConnPairWithLimits is [NewConnPair] with an explicit cap on each
// endpoint's outbound bidirectional-stream credit, modelling the peer's QUIC
// MAX_STREAMS limit. aBidiLimit caps how many bidi streams endpoint a may
// open; bBidiLimit does the same for endpoint b. A negative limit means
// unlimited. Once an endpoint's credit is exhausted, [pipeConn.OpenStream]
// returns [session.ErrNoStreamCredit] immediately — mirroring real QUIC, where
// a new stream cannot be opened until the peer raises the MAX_STREAMS limit.
func NewConnPairWithLimits(aBidiLimit, bBidiLimit int) (a, b session.Conn) {
        return newConnPair(aBidiLimit, bBidiLimit, 0)
}

// NewConnPairBuffered is [NewConnPair] but with each stream backed by a
// buffered [bufPipe] of bufSize bytes instead of a synchronous io.Pipe. The
// writer can run ahead by up to bufSize bytes before blocking, which decouples
// the producer and consumer goroutines so a relay/session throughput benchmark
// measures forwarding work rather than per-object goroutine scheduling. Both
// endpoints have unlimited outbound bidi-stream credit.
func NewConnPairBuffered(bufSize int) (a, b session.Conn) {
        return newConnPair(-1, -1, bufSize)
}

func newConnPair(aBidiLimit, bBidiLimit, bufSize int) (a, b session.Conn) {
        aUniToB := make(chan *uniStream, 4)
        bUniToA := make(chan *uniStream, 4)
        aBidiToB := make(chan *bidiStream, 4)
        bBidiToA := make(chan *bidiStream, 4)
        // Datagram channels: what A sends, B receives, and vice-versa.
        aDatagramToB := make(chan []byte, 16)
        bDatagramToA := make(chan []byte, 16)
        aCtx, aCancel := context.WithCancel(context.Background())
        bCtx, bCancel := context.WithCancel(context.Background())
        return &pipeConn{
                        uniOut: aUniToB, uniIn: bUniToA,
                        bidiOut: aBidiToB, bidiIn: bBidiToA,
                        datagramOut: aDatagramToB, datagramIn: bDatagramToA,
                        ctx: aCtx, ctxCancel: aCancel,
                        bidiCredit: aBidiLimit,
                        bufSize:    bufSize,
                },
                &pipeConn{
                        uniOut: bUniToA, uniIn: aUniToB,
                        bidiOut: bBidiToA, bidiIn: aBidiToB,
                        datagramOut: bDatagramToA, datagramIn: aDatagramToB,
                        ctx: bCtx, ctxCancel: bCancel,
                        bidiCredit: bBidiLimit,
                        bufSize:    bufSize,
                }
}

var errCancelled = errors.New("sessiontest: stream cancelled")
var errConnClosed = errors.New("sessiontest: connection closed")

// uniStream is a unidirectional pipe. The opener writes via w; the acceptor
// reads via r. The same struct satisfies both SendStream and ReceiveStream;
// each side gets back the appropriate interface, which constrains which
// methods they can call.
//
// ctx / ctxCancel implement Context(): the context is cancelled when Close()
// or CancelWrite() is called, or when the acceptor's CancelRead stops reading
// (STOP_SENDING) — the same three events as quic-go.
type uniStream struct {
        r         pipeReadCloser
        w         pipeWriteCloser
        ctx       context.Context
        ctxCancel context.CancelFunc
}

func newUniStream(bufSize int) *uniStream {
        r, w := newPipe(bufSize)
        ctx, cancel := context.WithCancel(context.Background())
        return &uniStream{r: r, w: w, ctx: ctx, ctxCancel: cancel}
}

func (s *uniStream) Write(p []byte) (int, error) { return s.w.Write(p) }
func (s *uniStream) Close() error {
        err := s.w.Close()
        s.ctxCancel() // the send side is closed (not "all data committed")
        return err
}
func (s *uniStream) CancelWrite(uint64) {
        _ = s.w.CloseWithError(errCancelled)
        s.ctxCancel() // signal reset
}
func (s *uniStream) Read(p []byte) (int, error) { return s.r.Read(p) }

// CancelRead is the acceptor's STOP_SENDING; it also cancels Context.
func (s *uniStream) CancelRead(uint64) {
        _ = s.r.CloseWithError(errCancelled)
        s.ctxCancel()
}

// Context is cancelled when Close() or CancelWrite() has been called, or the
// peer called CancelRead.
func (s *uniStream) Context() context.Context { return s.ctx }

// bidiStream is two io.Pipes wired so each end reads what the other writes.
// ctx / ctxCancel implement Context() on the send side; peerCtxCancel cancels
// the other end's on this end's CancelRead (STOP_SENDING).
type bidiStream struct {
        r             pipeReadCloser
        w             pipeWriteCloser
        ctx           context.Context
        ctxCancel     context.CancelFunc
        peerCtxCancel context.CancelFunc
}

func newBidiStreamPair(bufSize int) (a, b *bidiStream) {
        aR, aW := newPipe(bufSize) // a writes, b reads
        bR, bW := newPipe(bufSize) // b writes, a reads
        aCtx, aCancel := context.WithCancel(context.Background())
        bCtx, bCancel := context.WithCancel(context.Background())
        return &bidiStream{r: bR, w: aW, ctx: aCtx, ctxCancel: aCancel, peerCtxCancel: bCancel},
                &bidiStream{r: aR, w: bW, ctx: bCtx, ctxCancel: bCancel, peerCtxCancel: aCancel}
}

func (s *bidiStream) Read(p []byte) (int, error)  { return s.r.Read(p) }
func (s *bidiStream) Write(p []byte) (int, error) { return s.w.Write(p) }
func (s *bidiStream) Close() error {
        err := s.w.Close()
        s.ctxCancel() // the send side is closed (not "all data committed")
        return err
}
func (s *bidiStream) CancelRead(uint64) {
        _ = s.r.CloseWithError(errCancelled)
        s.peerCtxCancel() // STOP_SENDING ends the peer's send side
}
func (s *bidiStream) CancelWrite(uint64) {
        _ = s.w.CloseWithError(errCancelled)
        s.ctxCancel() // signal reset
}

// Context is cancelled when Close() or CancelWrite() has been called, or the
// peer called CancelRead.
func (s *bidiStream) Context() context.Context { return s.ctx }

// cancellable is satisfied by both uniStream and bidiStream — anything the
// pipeConn hands out and needs to forcibly tear down on connection close.
type cancellable interface {
        CancelRead(uint64)
        CancelWrite(uint64)
}

type pipeConn struct {
        uniOut, uniIn           chan *uniStream
        bidiOut, bidiIn         chan *bidiStream
        datagramOut, datagramIn chan []byte

        ctx       context.Context
        ctxCancel context.CancelFunc

        mu      sync.Mutex
        closed  bool
        tracked []cancellable

        // bidiCredit caps how many outbound bidirectional streams this endpoint
        // may open, modelling the peer's QUIC MAX_STREAMS limit. A negative value
        // means unlimited. bidiUsed counts streams already opened; both are
        // guarded by mu.
        bidiCredit int
        bidiUsed   int

        // bufSize selects the per-stream pipe backing: 0 = synchronous io.Pipe,
        // >0 = a buffered bufPipe of that capacity (see [newPipe]).
        bufSize int
}

// reserveBidiCredit accounts for one outbound bidi stream against the cap.
// Returns false when the cap is set (non-negative) and already exhausted.
func (c *pipeConn) reserveBidiCredit() bool {
        c.mu.Lock()
        defer c.mu.Unlock()
        if c.bidiCredit >= 0 && c.bidiUsed >= c.bidiCredit {
                return false
        }
        c.bidiUsed++
        return true
}

// track records s so CloseWithError can cancel it. Returns false if the conn
// has already been closed, in which case the caller should cancel s itself.
func (c *pipeConn) track(s cancellable) bool {
        c.mu.Lock()
        defer c.mu.Unlock()
        if c.closed {
                return false
        }
        c.tracked = append(c.tracked, s)
        return true
}

func (c *pipeConn) OpenUniStream() (session.SendStream, error) {
        select {
        case <-c.ctx.Done():
                return nil, errConnClosed
        default:
        }
        s := newUniStream(c.bufSize)
        // Non-blocking per the Conn contract: a full accept queue means the
        // peer isn't draining opens — surface ErrNoStreamCredit instead of
        // deadlocking the opener (the transport equivalent of an exhausted
        // stream limit).
        select {
        case c.uniOut <- s:
                if !c.track(s) {
                        s.CancelRead(0)
                        s.CancelWrite(0)
                        return nil, errConnClosed
                }
                return s, nil
        case <-c.ctx.Done():
                return nil, errConnClosed
        default:
                return nil, session.ErrNoStreamCredit
        }
}

func (c *pipeConn) AcceptUniStream(ctx context.Context) (session.ReceiveStream, error) {
        select {
        case s := <-c.uniIn:
                if !c.track(s) {
                        s.CancelRead(0)
                        s.CancelWrite(0)
                        return nil, errConnClosed
                }
                return s, nil
        case <-ctx.Done():
                return nil, ctx.Err()
        case <-c.ctx.Done():
                return nil, errConnClosed
        }
}

// OpenStream is the non-blocking bidi open. When the credit cap is exhausted
// it returns session.ErrNoStreamCredit instead of blocking, mirroring
// quic-go's Conn.OpenStream / StreamLimitReachedError.
func (c *pipeConn) OpenStream() (session.Stream, error) {
        if !c.reserveBidiCredit() {
                return nil, session.ErrNoStreamCredit
        }
        mine, peers := newBidiStreamPair(c.bufSize)
        // Non-blocking per the Conn contract — see OpenUniStream.
        select {
        case c.bidiOut <- peers:
                if !c.track(mine) {
                        mine.CancelRead(0)
                        mine.CancelWrite(0)
                        return nil, errConnClosed
                }
                return mine, nil
        case <-c.ctx.Done():
                return nil, errConnClosed
        default:
                return nil, session.ErrNoStreamCredit
        }
}

func (c *pipeConn) AcceptStream(ctx context.Context) (session.Stream, error) {
        select {
        case s := <-c.bidiIn:
                if !c.track(s) {
                        s.CancelRead(0)
                        s.CancelWrite(0)
                        return nil, errConnClosed
                }
                return s, nil
        case <-ctx.Done():
                return nil, ctx.Err()
        case <-c.ctx.Done():
                return nil, errConnClosed
        }
}

// SendDatagram delivers payload to the peer's datagramIn channel. The payload
// is copied so the caller may reuse the slice immediately.
func (c *pipeConn) SendDatagram(payload []byte) error {
        cp := make([]byte, len(payload))
        copy(cp, payload)
        select {
        case c.datagramOut <- cp:
                return nil
        case <-c.ctx.Done():
                return errConnClosed
        }
}

// ReceiveDatagram blocks until a datagram arrives from the peer or ctx / the
// connection is cancelled.
func (c *pipeConn) ReceiveDatagram(ctx context.Context) ([]byte, error) {
        select {
        case p := <-c.datagramIn:
                return p, nil
        case <-ctx.Done():
                return nil, ctx.Err()
        case <-c.ctx.Done():
                return nil, errConnClosed
        }
}

// CloseWithError mirrors quic-go: cancelling the connection also forcibly
// tears down every stream the conn has handed out, so any in-flight Read or
// Write on those streams unblocks with an error.
func (c *pipeConn) CloseWithError(uint64, string) error {
        c.mu.Lock()
        if c.closed {
                c.mu.Unlock()
                return nil
        }
        c.closed = true
        tracked := c.tracked
        c.tracked = nil
        c.mu.Unlock()

        c.ctxCancel()
        for _, s := range tracked {
                s.CancelRead(0)
                s.CancelWrite(0)
        }
        return nil
}

func (c *pipeConn) Context() context.Context { return c.ctx }

package session

import (
        "context"
        "io"

        "github.com/floatdrop/moq-go/pkg/moqt"
)

// DrainAndWait keeps a request stream alive until the peer closes its send
// side (FIN or RESET_STREAM) or ctx is cancelled, whichever comes first.
// It does not expect any meaningful data on the stream — incoming bytes are
// read and discarded.
//
// This is the canonical "hold a long-lived request stream open" primitive
// for MoQT consumers. §6.1, §6.2, and §10.7 all model the request stream as
// a subscription-lifetime keepalive: post-OK there are no further wire
// messages, but the stream must stay open as long as the requester still
// wants the subscription. Both the relay's session handlers and end
// subscribers / publishers benefit from a shared implementation rather than
// re-inventing the ctx-cancel + CancelRead dance at every call site.
//
// On ctx cancellation DrainAndWait calls [ReceiveStream.CancelRead] with
// [moqt.StreamResetSessionClosed] so the underlying read unblocks promptly.
// The function does not return until the read goroutine has exited; this is
// what makes it safe to invoke from a [sync.WaitGroup.Go] without leaking.
//
// DrainAndWait is concurrency-safe in the trivial sense that the inner
// goroutine is owned by this call; do not invoke it concurrently with other
// readers of the same Stream. In particular it discards §10.9 responses, so
// it cannot be combined with Update on the same stream — when a request
// needs both a lifetime keepalive AND updates or follow-up handling, run
// [RequestBroker.Serve] (via the handle's Broker method) instead.
func DrainAndWait(ctx context.Context, s Stream) {
        done := make(chan struct{})
        go func() {
                defer close(done)
                _, _ = io.Copy(io.Discard, s)
        }()
        select {
        case <-done:
        case <-ctx.Done():
                s.CancelRead(uint64(moqt.StreamResetSessionClosed))
                <-done
        }
}

package session

import (
        "context"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
)

// Subscription is a live subscriber-initiated track subscription. It owns the
// request stream (embedded, so Close / reads / message.Marshal work directly
// on it) plus the identifiers follow-up traffic needs — the Request ID and the
// publisher-assigned Track Alias — so the caller can send REQUEST_UPDATE via
// [Subscription.Update] without holding them separately. It is returned by
// [Session.Subscribe].
//
// The Track Alias stays registered while the subscription is Established and
// is released (§11.1) once it is Terminated (§5.1): on [Subscription.Close] or
// its broker's [RequestBroker.Close], when the publisher's FIN is read through
// the Subscription or by its broker's [RequestBroker.Serve], or when Serve
// cancels the stream.
type Subscription struct {
        // requestHandle carries the SUBSCRIBE request stream — still open for
        // follow-up traffic (REQUEST_UPDATE and inbound PUBLISH_DONE; Close it
        // to end the subscription) — and provides Update.
        requestHandle

        // OK is the parsed SUBSCRIBE_OK response — the publisher-assigned Track
        // Alias, negotiated Parameters, and TrackProperties.
        OK *message.SubscribeOK
}

// TrackAlias reports the §11.1 Track Alias the publisher assigned to this
// subscription — the integer inbound subgroup and datagram streams carry to
// identify the track (see [Session.AcceptDataStream]). It is shorthand for
// sub.OK.TrackAlias.
func (sub *Subscription) TrackAlias() uint64 { return sub.OK.TrackAlias }

// Subscribe opens a SUBSCRIBE request stream (§10.7) and awaits SUBSCRIBE_OK.
// The session assigns m.RequestID; the caller supplies the rest. On success a
// [Subscription] is returned whose embedded stream stays open for follow-up
// traffic (REQUEST_UPDATE via [Subscription.Update], inbound PUBLISH_DONE) and
// whose [Subscription.TrackAlias] matches the alias on inbound subgroup
// streams. REQUEST_ERROR is surfaced as a *RequestRejectedError and the stream
// is closed.
func (s *Session) Subscribe(ctx context.Context, m *message.Subscribe) (*Subscription, error) {
        return awaitRequestResponse(ctx, s, m,
                func(stream Stream, ok *message.SubscribeOK) (*Subscription, error) {
                        // §2.5.1: "the subscriber MUST cancel the subscription".
                        if err := s.validateTrackProperties(ok.TrackProperties, "SUBSCRIBE_OK"); err != nil {
                                cancelRequest(stream)
                                return nil, err
                        }
                        // §11.1: register the alias the publisher assigned so we can detect
                        // DUPLICATE_TRACK_ALIAS if the same alias is reused for a different track.
                        key := track.NewKey(m.Namespace, m.Name)
                        if err := s.RegisterInboundTrack(ok.TrackAlias, key, ok.TrackProperties); err != nil {
                                cancelRequest(stream)
                                return nil, err
                        }
                        // The publisher may send PUBLISH_STATE_NOTIFY (§10.10) but not
                        // REQUEST_UPDATE (§10.9).
                        return &Subscription{
                                Stream:       stream,
                                s:            s,
                                requestID:    m.RequestID,
                                peerNotify:   true,
                                inboundAlias: ok.TrackAlias,
                                holdsAlias:   true,
                                answered:     message.TypeSubscribe,
                                OK:           ok,
                        }, nil
                })
}

package session

import (
        "fmt"
        "sync"

        "github.com/floatdrop/moq-go/pkg/moqt"
)

// tokenCacheEntry holds the resolved (type, value) for a registered alias.
type tokenCacheEntry struct {
        tokenType uint64
        value     []byte
        size      uint64 // 16 + len(value), per §10.3.1.3
}

// TokenCache is a per-session, per-direction alias cache per §10.2.2.
//
// Client and server each maintain independent caches (separate alias spaces).
// The cache is thread-safe.
//
// Cache size accounting per §10.3.1.3:
//   - Token size = 16 bytes + len(TokenValue)
//   - Total = Σ(registered token sizes) − Σ(deregistered token sizes)
//   - maxSize = 0 prohibits alias registration (default when MAX_AUTH_TOKEN_CACHE_SIZE
//     is not negotiated in SETUP)
type TokenCache struct {
        mu      sync.Mutex
        maxSize uint64 // 0 = aliases prohibited
        used    uint64 // current total size
        entries map[uint64]*tokenCacheEntry
}

// NewTokenCache creates a cache with the given maximum byte size.
// maxSize=0 prohibits alias registration (the default per §10.3.1.3).
func NewTokenCache(maxSize uint64) *TokenCache {
        return &TokenCache{
                maxSize: maxSize,
                entries: make(map[uint64]*tokenCacheEntry),
        }
}

// Register adds alias → (tokenType, value) to the cache per §10.2.2 REGISTER.
//
// Returns:
//   - moqt.SessionDuplicateAuthTokenAlias if alias is already registered.
//   - moqt.SessionAuthTokenCacheOverflow if adding would exceed maxSize.
//
// Per §10.2.2: even if the message fails for other reasons, a REGISTER that
// does not cause a session error MUST be stored. The caller is responsible for
// applying that rule (i.e. call Register before validating the message).
func (c *TokenCache) Register(alias, tokenType uint64, value []byte) error {
        size := uint64(16) + uint64(len(value))

        c.mu.Lock()
        defer c.mu.Unlock()

        if _, exists := c.entries[alias]; exists {
                return fmt.Errorf("moqt/session: token alias %d already registered (%w)",
                        alias, sessionErr(moqt.SessionDuplicateAuthTokenAlias))
        }

        if c.maxSize == 0 {
                // Aliases are prohibited when MAX_AUTH_TOKEN_CACHE_SIZE was not negotiated.
                return fmt.Errorf("moqt/session: token alias registration prohibited (maxSize=0) (%w)",
                        sessionErr(moqt.SessionAuthTokenCacheOverflow))
        }

        if c.used+size > c.maxSize {
                return fmt.Errorf("moqt/session: token cache overflow (used=%d size=%d max=%d) (%w)",
                        c.used, size, c.maxSize, sessionErr(moqt.SessionAuthTokenCacheOverflow))
        }

        valueCopy := make([]byte, len(value))
        copy(valueCopy, value)
        c.entries[alias] = &tokenCacheEntry{
                tokenType: tokenType,
                value:     valueCopy,
                size:      size,
        }
        c.used += size
        return nil
}

// Delete removes alias from the cache per §10.2.2 DELETE.
// Returns moqt.SessionUnknownAuthTokenAlias if alias is not registered.
func (c *TokenCache) Delete(alias uint64) error {
        c.mu.Lock()
        defer c.mu.Unlock()

        entry, exists := c.entries[alias]
        if !exists {
                return fmt.Errorf("moqt/session: token alias %d not registered (%w)",
                        alias, sessionErr(moqt.SessionUnknownAuthTokenAlias))
        }

        c.used -= entry.size
        delete(c.entries, alias)
        return nil
}

// Resolve returns the (tokenType, value) for alias per §10.2.2 USE_ALIAS.
// Returns moqt.SessionUnknownAuthTokenAlias if alias is not registered.
// The returned value slice is a copy owned by the caller.
func (c *TokenCache) Resolve(alias uint64) (tokenType uint64, value []byte, err error) {
        c.mu.Lock()
        defer c.mu.Unlock()

        entry, exists := c.entries[alias]
        if !exists {
                return 0, nil, fmt.Errorf("moqt/session: token alias %d not registered (%w)",
                        alias, sessionErr(moqt.SessionUnknownAuthTokenAlias))
        }

        out := make([]byte, len(entry.value))
        copy(out, entry.value)
        return entry.tokenType, out, nil
}

// Size returns the current total cache size in bytes.
func (c *TokenCache) Size() uint64 {
        c.mu.Lock()
        defer c.mu.Unlock()
        return c.used
}

// MaxSize returns the configured maximum cache size.
func (c *TokenCache) MaxSize() uint64 {
        return c.maxSize // immutable after construction; no lock needed
}

// sessionErr wraps a SessionErrorCode as an error so callers can use errors.Is
// to identify the specific session-level error to signal.
type sessionErrCode moqt.SessionErrorCode

func (e sessionErrCode) Error() string {
        return fmt.Sprintf("session error code 0x%X", uint64(e))
}

func (e sessionErrCode) Is(target error) bool {
        t, ok := target.(sessionErrCode)
        return ok && t == e
}

func sessionErr(code moqt.SessionErrorCode) error {
        return sessionErrCode(code)
}

package session

import (
        "bytes"
        "context"
        "errors"
        "fmt"
        "slices"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// ResolvedToken is a fully-resolved AUTHORIZATION_TOKEN (§10.2.2): the
// (Token Type, Token Value) pair an application policy needs to make an
// authorization decision, with all alias indirection already removed.
//
// The session produces a ResolvedToken for every REGISTER, USE_ALIAS, and
// USE_VALUE token on an inbound request (DELETE tokens carry no value and so
// produce none). USE_ALIAS tokens are resolved against the inbound TokenCache
// before the value is exposed, so a verifier never sees a bare alias.
//
// Value is owned by the caller (a fresh copy per resolution); mutating it does
// not affect the cache.
type ResolvedToken struct {
        // Type is the Token Type (§10.2.2) — an application-defined identifier
        // of the token scheme (e.g. a registry entry for a CAT or JWT profile).
        Type uint64
        // Value is the raw, opaque Token Value. Its interpretation is entirely
        // up to the TokenVerifier; the transport treats it as bytes (§13.3).
        Value []byte
}

// TokenVerifier is the application policy that authorizes resolved
// authorization tokens. The session invokes VerifyToken once per
// ResolvedToken carried by an inbound request, after the token has been
// resolved against the inbound cache.
//
// The transport deliberately defines no token format (§13.3); a verifier is
// where signature checking, expiry, audience, and scope validation live.
//
// Returning nil authorizes the token. Returning a non-nil error denies the
// request the token accompanied: wrap the error with [*TokenDeniedError] (or
// use [DenyToken]) to choose the REQUEST_ERROR code the peer receives —
// notably [moqt.RequestExpiredAuthToken] for an expired token per §10.2.2. A
// plain error denies with [moqt.RequestUnauthorized].
//
// VerifyToken must be safe for concurrent use: requests on a session are
// dispatched concurrently, so multiple goroutines may call it at once.
type TokenVerifier interface {
        VerifyToken(ctx context.Context, sess *Session, tok ResolvedToken) error
}

// TokenVerifierFunc adapts an ordinary function to the [TokenVerifier]
// interface, so a policy can be supplied inline without a named type.
type TokenVerifierFunc func(ctx context.Context, sess *Session, tok ResolvedToken) error

// VerifyToken calls f.
func (f TokenVerifierFunc) VerifyToken(ctx context.Context, sess *Session, tok ResolvedToken) error {
        return f(ctx, sess, tok)
}

// TokenCacheError is returned by [Session.AcceptRequest] when processing an
// inbound request's AUTHORIZATION_TOKEN parameters fails at the cache layer
// (§10.2.2). These are session-level faults: a malformed token, a duplicate
// REGISTER alias, a cache overflow, or a USE_ALIAS / DELETE referencing an
// unknown alias. Code is the SESSION_ERROR the session is closed with:
// AcceptRequest has already done so, and a caller of
// [Session.ProcessFollowupTokens] MUST.
type TokenCacheError struct {
        // Code is the §10.2.2 SESSION_ERROR code to terminate the session with.
        Code moqt.SessionErrorCode
        // Err is the underlying cache or parse error, for diagnostics.
        Err error
}

// Error implements the error interface.
func (e *TokenCacheError) Error() string {
        return fmt.Sprintf("moqt/session: token cache error (session code 0x%X): %v", uint64(e.Code), e.Err)
}

// Unwrap exposes the underlying error for errors.Is / errors.As.
func (e *TokenCacheError) Unwrap() error { return e.Err }

// TokenDeniedError is returned by token verification to deny a single request
// with an explicit MoQT REQUEST_ERROR code. Unlike [TokenCacheError] it is a
// per-request rejection, not a session-level fault: the caller should reply
// REQUEST_ERROR and leave the session running.
//
// Code MUST be one of the §10.6 REQUEST_ERROR codes (see
// [moqt.RequestErrorCode]); the zero value collapses to
// [moqt.RequestUnauthorized].
type TokenDeniedError struct {
        // Code is the REQUEST_ERROR code to send. Zero ⇒ RequestUnauthorized.
        Code moqt.RequestErrorCode
        // Reason is the human-readable reason forwarded to the peer.
        Reason string
        // Err is the underlying verifier error, for diagnostics. Optional.
        Err error
}

// Error implements the error interface.
func (e *TokenDeniedError) Error() string {
        if e.Reason != "" {
                return fmt.Sprintf("moqt/session: token denied (request code 0x%X): %s", uint64(e.RequestErrorCode()), e.Reason)
        }
        return fmt.Sprintf("moqt/session: token denied (request code 0x%X)", uint64(e.RequestErrorCode()))
}

// Unwrap exposes the underlying verifier error for errors.Is / errors.As.
func (e *TokenDeniedError) Unwrap() error { return e.Err }

// RequestErrorCode returns the REQUEST_ERROR code to send, substituting
// [moqt.RequestUnauthorized] for the zero value.
func (e *TokenDeniedError) RequestErrorCode() moqt.RequestErrorCode {
        if e.Code == 0 {
                return moqt.RequestUnauthorized
        }
        return e.Code
}

// DenyToken constructs a [*TokenDeniedError]. Use it from a [TokenVerifier]
// to reject a request with a specific REQUEST_ERROR code:
//
//        return session.DenyToken(moqt.RequestExpiredAuthToken, "token expired")
func DenyToken(code moqt.RequestErrorCode, reason string) error {
        return &TokenDeniedError{Code: code, Reason: reason}
}

// TokenCache returns the session's inbound authorization-token alias cache
// (§10.2.2). It is primarily exposed for inspection and tests; the session
// drives Register / Resolve / Delete on it automatically from inbound request
// parameters in [Session.AcceptRequest]. Always non-nil.
func (s *Session) TokenCache() *TokenCache { return s.tokenCache }

// processRequestTokens parses the AUTHORIZATION_TOKEN parameters of msg and
// applies each token to the inbound cache per §10.2.2, returning the resolved
// (Type, Value) tokens for any REGISTER / USE_ALIAS / USE_VALUE entries.
//
// A REGISTER is committed immediately, before the request is validated, so a
// later rejection does not roll it back (§10.2.2: "even if the message fails
// for other reasons").
//
// A cache-layer failure is returned as a [*TokenCacheError] carrying the code
// the caller must close the session with.
func (s *Session) processRequestTokens(msg message.Message) ([]ResolvedToken, error) {
        ps, ok := message.ParamsOf(msg)
        if !ok {
                return nil, nil
        }
        tokens, err := message.TokensFromParam(ps)
        if err != nil {
                // §10.2.2: an undecodable Token, including an unknown Alias Type,
                // closes with KEY_VALUE_FORMATTING_ERROR. This follows that MUST
                // over §3.5's MALFORMED_AUTH_TOKEN description.
                return nil, &TokenCacheError{Code: moqt.SessionKeyValueFormattingError, Err: err}
        }
        if len(tokens) == 0 {
                return nil, nil
        }

        var resolved []ResolvedToken
        for i := range tokens {
                tok, ok, err := s.applyToken(&tokens[i])
                if err != nil {
                        return nil, err
                }
                if ok {
                        resolved = append(resolved, tok)
                }
        }
        return resolved, nil
}

// applyToken applies one token to the inbound cache (§10.2.2) and returns what
// it resolves to; ok is false for a DELETE. A cache failure is a
// [*TokenCacheError].
func (s *Session) applyToken(t *message.Token) (tok ResolvedToken, ok bool, err error) {
        switch t.AliasType {
        case message.AliasTypeRegister:
                // §10.2.2: register before any further validation.
                if err := s.tokenCache.Register(t.TokenAlias, t.TokenType, t.TokenValue); err != nil {
                        return ResolvedToken{}, false, &TokenCacheError{Code: sessionCodeForCacheErr(err), Err: err}
                }
                return ResolvedToken{Type: t.TokenType, Value: bytes.Clone(t.TokenValue)}, true, nil

        case message.AliasTypeUseAlias:
                typ, val, err := s.tokenCache.Resolve(t.TokenAlias)
                if err != nil {
                        return ResolvedToken{}, false, &TokenCacheError{Code: sessionCodeForCacheErr(err), Err: err}
                }
                return ResolvedToken{Type: typ, Value: val}, true, nil

        case message.AliasTypeUseValue:
                return ResolvedToken{Type: t.TokenType, Value: bytes.Clone(t.TokenValue)}, true, nil

        case message.AliasTypeDelete:
                if err := s.tokenCache.Delete(t.TokenAlias); err != nil {
                        return ResolvedToken{}, false, &TokenCacheError{Code: sessionCodeForCacheErr(err), Err: err}
                }
        }
        // No other Alias Type: Token.Parse rejects it.
        return ResolvedToken{}, false, nil
}

// VerifyRequestTokens runs the configured [TokenVerifier] over the tokens the
// session resolved for req (see [Request.Tokens]). It returns nil when no
// verifier is configured or every token is authorized, and a
// [*TokenDeniedError] (mappable to a REQUEST_ERROR) for the first denial.
//
// The relay calls this before dispatching a request; standalone session users
// can call it from their own request loop. It is safe to call with a req whose
// Tokens slice is empty.
func (s *Session) VerifyRequestTokens(ctx context.Context, req *Request) error {
        return s.VerifyTokens(ctx, req.Tokens)
}

// VerifyTokens is [Session.VerifyRequestTokens] for tokens resolved by
// [Session.ProcessFollowupTokens], such as a REQUEST_UPDATE's (§10.2.2).
func (s *Session) VerifyTokens(ctx context.Context, toks []ResolvedToken) error {
        if s.tokenVerifier == nil || len(toks) == 0 {
                return nil
        }
        for _, tok := range toks {
                if err := s.tokenVerifier.VerifyToken(ctx, s, tok); err != nil {
                        if denied, ok := errors.AsType[*TokenDeniedError](err); ok {
                                return denied
                        }
                        return &TokenDeniedError{Code: moqt.RequestUnauthorized, Reason: err.Error(), Err: err}
                }
        }
        return nil
}

// sessionCodeForCacheErr maps a [TokenCache] error to its §10.2.2 SESSION_ERROR
// code. The cache wraps a sentinel via sessionErr, so errors.Is identifies
// which one. An unrecognised error defaults to MALFORMED_AUTH_TOKEN.
func sessionCodeForCacheErr(err error) moqt.SessionErrorCode {
        for _, c := range []moqt.SessionErrorCode{
                moqt.SessionDuplicateAuthTokenAlias,
                moqt.SessionAuthTokenCacheOverflow,
                moqt.SessionUnknownAuthTokenAlias,
        } {
                if errors.Is(err, sessionErr(c)) {
                        return c
                }
        }
        return moqt.SessionMalformedAuthToken
}

// ProcessFollowupTokens resolves the AUTHORIZATION_TOKEN parameters (§10.2.2)
// of a follow-up message, such as a REQUEST_UPDATE, read off an established
// request stream. Code that reads follow-ups with message.Parse MUST route
// them through here, or the token cache diverges from the peer's.
//
// A *TokenCacheError carries the code the caller must close the session with.
// Messages without token parameters return (nil, nil).
func (s *Session) ProcessFollowupTokens(msg message.Message) ([]ResolvedToken, error) {
        return s.processRequestTokens(msg)
}

// SetupTokens returns copies of the resolved tokens the peer sent in
// AUTHORIZATION TOKEN options of its SETUP (§10.3.1.4). The session does not
// verify them.
func (s *Session) SetupTokens() []ResolvedToken {
        out := make([]ResolvedToken, len(s.setupTokens))
        for i, t := range s.setupTokens {
                out[i] = ResolvedToken{Type: t.Type, Value: bytes.Clone(t.Value)}
        }
        return out
}

// processSetupTokens applies the AUTHORIZATION TOKEN options in the peer's
// SETUP (§10.3.1.4) and keeps the resolved tokens for [Session.SetupTokens].
// Unlike a request, a server closes with PROTOCOL_VIOLATION on DELETE or
// USE_ALIAS (§10.2.2), and a REGISTER that overflows the cache is treated as
// USE_VALUE (§10.3.1.4).
//
// Assumption: a REGISTER that both repeats an alias and overflows closes with
// DUPLICATE_AUTH_TOKEN_ALIAS; the draft does not say which rule wins.
//
// Every error is a [*TokenCacheError] carrying the code to close with.
func (s *Session) processSetupTokens() error {
        for _, opt := range s.peerOptions {
                if message.SetupOption(opt.Type) != message.SetupOptionAuthorizationToken {
                        continue
                }
                var t message.Token
                if err := t.Parse(opt.ByteVal); err != nil {
                        // §10.2.2.
                        return &TokenCacheError{Code: moqt.SessionKeyValueFormattingError,
                                Err: fmt.Errorf("moqt/session: AUTHORIZATION TOKEN setup option: %w", err)}
                }
                if s.role == roleServer &&
                        (t.AliasType == message.AliasTypeDelete || t.AliasType == message.AliasTypeUseAlias) {
                        return &TokenCacheError{Code: moqt.SessionProtocolViolation,
                                Err: fmt.Errorf("moqt/session: %s token in the client's SETUP (§10.2.2)", t.AliasType)}
                }
                tok, ok, err := s.applyToken(&t)
                if tce, isTCE := errors.AsType[*TokenCacheError](err); isTCE &&
                        t.AliasType == message.AliasTypeRegister && tce.Code == moqt.SessionAuthTokenCacheOverflow {
                        tok, ok, err = ResolvedToken{Type: t.TokenType, Value: bytes.Clone(t.TokenValue)}, true, nil
                }
                if err != nil {
                        return err
                }
                if ok {
                        s.setupTokens = append(s.setupTokens, tok)
                }
        }
        return nil
}

// SetupTokenAliases returns the aliases of the REGISTER tokens this endpoint
// sent in SETUP ([WithSetupToken]) that fit the peer's
// MAX_AUTH_TOKEN_CACHE_SIZE, in the order sent (§10.3.1.4). Only these may be
// referenced with USE_ALIAS.
func (s *Session) SetupTokenAliases() []uint64 { return slices.Clone(s.setupTokenAliases) }

// checkOutboundSetupTokens refuses setup tokens the peer would close the
// session over (§10.2.2): DELETE, USE_ALIAS, or an alias REGISTERed twice.
func checkOutboundSetupTokens(toks []message.Token) error {
        var registered []uint64
        for _, t := range toks {
                switch t.AliasType {
                case message.AliasTypeRegister:
                        if slices.Contains(registered, t.TokenAlias) {
                                return fmt.Errorf("moqt/session: AUTHORIZATION TOKEN setup option registers alias %d twice (§10.2.2)",
                                        t.TokenAlias)
                        }
                        registered = append(registered, t.TokenAlias)
                case message.AliasTypeUseValue:
                case message.AliasTypeDelete, message.AliasTypeUseAlias:
                        return fmt.Errorf("moqt/session: AUTHORIZATION TOKEN setup option with %s (§10.2.2)", t.AliasType)
                default:
                        return fmt.Errorf("moqt/session: AUTHORIZATION TOKEN setup option with %s", t.AliasType)
                }
        }
        return nil
}

// heldSetupAliases replays the peer's cache accounting ([TokenCache.Register])
// against its MAX_AUTH_TOKEN_CACHE_SIZE (§10.3.1.3; default 0) and returns
// the aliases of the REGISTERs that fit (§10.3.1.4).
func heldSetupAliases(toks []message.Token, peerOptions []wire.KVPair) []uint64 {
        var limit uint64
        for _, opt := range peerOptions {
                if message.SetupOption(opt.Type) == message.SetupOptionMaxAuthTokenCache {
                        limit = opt.IntVal
                }
        }
        var (
                used uint64
                held []uint64
        )
        for _, t := range toks {
                if t.AliasType != message.AliasTypeRegister {
                        continue
                }
                size := uint64(16) + uint64(len(t.TokenValue))
                if used+size > limit {
                        continue
                }
                used += size
                held = append(held, t.TokenAlias)
        }
        return held
}

package session

import (
        "errors"
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// ErrUnsupportedMandatoryTrackProperty is returned when Track Properties
// (received in SUBSCRIBE_OK, FETCH_OK, TRACK_STATUS_OK, or an inbound
// PUBLISH) contain a Mandatory Track Property (range 0x4000–0x7FFF per
// §2.5.1) that this endpoint does not understand. The caller MUST NOT
// process or forward the track.
//
// For outbound requests (Subscribe, Fetch, TrackStatus) the session layer
// returns this error directly; [Request.AcceptPublish] replies REQUEST_ERROR
// UNSUPPORTED_EXTENSION before returning it.
type ErrUnsupportedMandatoryTrackProperty struct {
        // PropertyType is the first unrecognised mandatory property type found.
        PropertyType message.PropertyType
        // Context describes where the property was encountered (e.g.
        // "SUBSCRIBE_OK", "FETCH_OK", "PUBLISH").
        Context string
}

func (e *ErrUnsupportedMandatoryTrackProperty) Error() string {
        return fmt.Sprintf(
                "moqt/session: unsupported mandatory track property 0x%X in %s (§2.5.1 — UNSUPPORTED_EXTENSION)",
                e.PropertyType, e.Context,
        )
}

// ErrMalformedTrackProperties is wrapped by the error [ValidateTrackProperties]
// returns when raw Track Properties do not parse: a Key-Value-Pair that
// "cannot be parsed" makes the track malformed (§12.7, §2.4.2). A request is
// refused with INTERNAL_ERROR (see [TrackPropertiesRejectCode]).
var ErrMalformedTrackProperties = errors.New("moqt/session: malformed track properties")

// ErrTrackPropertiesNotAllowed is returned, and nothing sent, when asked to
// send a REQUEST_OK with Track Properties where §10.5 says they are empty.
var ErrTrackPropertiesNotAllowed = errors.New("moqt/session: track properties not allowed in this REQUEST_OK")

// ValidateTrackProperties parses raw Track Properties bytes and checks for
// unknown Mandatory Track Properties (range 0x4000–0x7FFF per §2.5.1).
//
// knownMandatory is the set of Mandatory Track Property types this endpoint
// supports. Every mandatory property found in raw that is not in this set
// causes *ErrUnsupportedMandatoryTrackProperty to be returned. An empty or nil
// map means "I support no mandatory extensions" — any mandatory property will
// be rejected.
//
// Returns the parsed pairs on success. context is used in the error message
// to identify the source message (e.g. "SUBSCRIBE_OK").
func ValidateTrackProperties(
        raw []byte,
        knownMandatory map[message.PropertyType]struct{},
        context string,
) ([]wire.KVPair, error) {
        pairs, err := message.ParseTrackProperties(raw)
        if err != nil {
                return nil, fmt.Errorf("%w in %s: %w", ErrMalformedTrackProperties, context, err)
        }
        // §12.7: Mandatory Track Properties inside Immutable Properties count.
        all, err := message.ExpandImmutable(pairs)
        if err != nil {
                return nil, fmt.Errorf("%w in %s: %w", ErrMalformedTrackProperties, context, err)
        }
        if typ, unknown := message.FirstUnknownMandatoryTrackProperty(all, knownMandatory); unknown {
                return nil, &ErrUnsupportedMandatoryTrackProperty{
                        PropertyType: typ,
                        Context:      context,
                }
        }
        return pairs, nil
}

// CheckTrackProperties validates raw Track Properties: a value the draft makes
// session-fatal closes the session (§12.5, §12.6), and against the types
// configured with [WithKnownMandatoryTrackProperties] (§2.5.1) it returns
// *ErrUnsupportedMandatoryTrackProperty or an error wrapping
// [ErrMalformedTrackProperties]; see [TrackPropertiesRejectCode]. It is for
// callers that bypass [Request.AcceptPublish] and the outbound openers, which
// already check. Without that option no Mandatory Track Property is known.
func (s *Session) CheckTrackProperties(raw []byte, context string) error {
        return s.validateTrackProperties(raw, context)
}

// validateTrackProperties checks Track Properties received in context. A value
// the draft makes session-fatal (see [message.CheckTrackPropertyValues])
// closes the session with PROTOCOL_VIOLATION. Then, against the session's
// configured set of known mandatory track property types, an unknown one is
// refused (§2.5.1). Without WithKnownMandatoryTrackProperties none is known,
// so every Mandatory Track Property is refused: an endpoint that does not
// understand one "MUST NOT process or forward that track".
func (s *Session) validateTrackProperties(raw []byte, context string) error {
        if err := s.checkTrackPropertyValues(raw, context); err != nil {
                return err
        }
        _, err := ValidateTrackProperties(raw, s.knownMandatoryTrackProperties, context)
        return err
}

// checkTrackPropertyValues closes the session with PROTOCOL_VIOLATION when
// Track Properties received in context hold a value the draft makes
// session-fatal (see [message.CheckTrackPropertyValues]): "If an endpoint
// receives a value outside this range, it MUST close the session" (§12.5,
// §12.6). Immutable Properties are searched when they parse (§12.7); Track
// Properties that do not parse are left to the caller.
func (s *Session) checkTrackPropertyValues(raw []byte, context string) error {
        pairs, parseErr := message.ParseTrackProperties(raw)
        if parseErr == nil {
                if all, err := message.ExpandImmutable(pairs); err == nil {
                        pairs = all
                }
        }
        if err := message.CheckTrackPropertyValues(pairs); err != nil {
                return s.closeProtocolViolation(fmt.Errorf("%s: %w", context, err))
        }
        return nil
}

// TrackPropertiesRejectCode is the REQUEST_ERROR code for a Track Properties
// validation error on a PUBLISH or SUBSCRIBE: UNSUPPORTED_EXTENSION for an
// unknown Mandatory Track Property (§2.5.1), else INTERNAL_ERROR.
// MALFORMED_TRACK is defined only "In response to a FETCH" (§10.6.2).
func TrackPropertiesRejectCode(err error) moqt.RequestErrorCode {
        if _, ok := errors.AsType[*ErrUnsupportedMandatoryTrackProperty](err); ok {
                return moqt.RequestUnsupportedExtension
        }
        return moqt.RequestInternalError
}

package session

import (
        "context"
        "errors"
        "fmt"
        "io"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
)

// TrackStatusRequest is a completed TRACK_STATUS request (§10.15), returned by
// [Session.TrackStatus] with the peer's TRACK_STATUS_OK. It cannot be updated,
// so this side of the stream is already FINned. Close stops reading.
type TrackStatusRequest struct {
        Stream

        // OK is the TRACK_STATUS_OK the peer replied with.
        OK *message.TrackStatusOK
}

// Close releases the stream's receive side.
func (t *TrackStatusRequest) Close() error {
        t.Stream.CancelRead(uint64(moqt.StreamResetCancelled))
        return nil
}

// AcceptTrackStatus accepts an inbound TRACK_STATUS (§10.15) and replies
// TRACK_STATUS_OK with the given status fields — the accept-side counterpart of
// [Session.TrackStatus]. r.First MUST be a *message.TrackStatus.
//
// ok carries the TRACK_STATUS_OK fields (status, largest location, Track
// Properties — [message.TrackStatusOK] is an alias of [message.RequestOK]); it
// may be nil for the all-default reply. The stream is FINned after the reply
// (§10.15), so no handle is returned.
//
// Any later message from the requester closes the session with
// PROTOCOL_VIOLATION (§10.9, §10). Interpretation: for a well-formed message
// other than REQUEST_UPDATE this reads §10.15's "first and only message"; the
// draft names no consequence.
func (r *Request) AcceptTrackStatus(ok *message.TrackStatusOK) error {
        if _, isTS := r.First.(*message.TrackStatus); !isTS {
                return fmt.Errorf("moqt/session: AcceptTrackStatus on a %s request", r.First.Type())
        }
        if ok == nil {
                ok = &message.TrackStatusOK{}
        }
        if err := message.Marshal(r.Stream, ok); err != nil {
                // As in RejectError: do not leave the requester waiting.
                resetStream(r.Stream)
                return fmt.Errorf("moqt/session: write TRACK_STATUS_OK: %w", err)
        }
        if err := r.Stream.Close(); err != nil {
                return fmt.Errorf("moqt/session: FIN TRACK_STATUS stream: %w", err)
        }
        go r.rejectTrackStatusFollowups()
        return nil
}

// rejectTrackStatusFollowups closes the session if anything arrives on an
// answered TRACK_STATUS stream (see [Request.AcceptTrackStatus]). It returns
// quietly on the requester's FIN, a reset or session close.
func (r *Request) rejectTrackStatusFollowups() {
        src := &readErrRecorder{r: r.Stream}
        msg, err := message.Parse(src)
        switch {
        case errors.Is(err, io.EOF), src.err != nil && !errors.Is(src.err, io.EOF):
                return // requester FIN, or the stream/session went away
        case err != nil:
                _ = r.s.closeProtocolViolation(fmt.Errorf(
                        "moqt/session: malformed data on a TRACK_STATUS request stream: %w", err))
        default:
                _ = r.s.closeProtocolViolation(fmt.Errorf(
                        "moqt/session: %s on a TRACK_STATUS request stream", msg.Type()))
        }
}

// readErrRecorder remembers the last non-nil error its reader returned, so a
// parse failure can be told apart from a transport failure beneath it.
type readErrRecorder struct {
        r   io.Reader
        err error
}

func (e *readErrRecorder) Read(p []byte) (int, error) {
        n, err := e.r.Read(p)
        if err != nil {
                e.err = err
        }
        return n, err
}

// TrackStatus opens a TRACK_STATUS request stream (§10.15) and awaits
// REQUEST_OK (TRACK_STATUS_OK) or REQUEST_ERROR. The session assigns
// m.RequestID; the caller supplies Namespace, Name, and optional Parameters.
//
// On success this side of the stream is FINned (§3.3.2) and a
// [TrackStatusRequest] holding the TRACK_STATUS_OK is returned. On
// REQUEST_ERROR the stream is closed and a *RequestRejectedError is returned.
func (s *Session) TrackStatus(ctx context.Context, m *message.TrackStatus) (*TrackStatusRequest, error) {
        return awaitRequestResponse(ctx, s, m,
                func(stream Stream, ok *message.RequestOK) (*TrackStatusRequest, error) {
                        // Assumption: §2.5.1 lists only PUBLISH, SUBSCRIBE_OK and FETCH_OK,
                        // but TRACK_STATUS_OK carries the Track Properties "it would have
                        // set in a SUBSCRIBE_OK" (§10.15), so the SUBSCRIBE_OK rule applies.
                        if err := s.validateTrackProperties(ok.TrackProperties, "TRACK_STATUS_OK"); err != nil {
                                _ = stream.Close()
                                return nil, err
                        }
                        if err := stream.Close(); err != nil {
                                return nil, fmt.Errorf("moqt/session: FIN TRACK_STATUS stream: %w", err)
                        }
                        return &TrackStatusRequest{Stream: stream, OK: ok}, nil
                })
}

package session

import (
        "context"
        "fmt"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
)

// AllocOutboundTrackAlias returns the next Track Alias to use when this side
// advertises a new track to the peer (§11.1). Aliases are independent across
// sessions, so callers must remap when forwarding between two sessions.
//
// Allocation starts at 1: [Session.Publish] and [Request.AcceptSubscribe]
// treat a zero TrackAlias as "allocate one for me", so an allocated 0 would be
// silently replaced. [Session.OpenPublish] does not: its caller allocates.
func (s *Session) AllocOutboundTrackAlias() uint64 {
        return s.nextOutboundTrackAlias.Add(1)
}

// ErrDuplicateTrackAlias is returned by [Session.RegisterInboundTrack] when the
// peer assigns a Track Alias that is already in use for a different track
// (§11.1). The session is already closed with
// [moqt.SessionDuplicateTrackAlias].
type ErrDuplicateTrackAlias struct {
        Alias    uint64
        Existing track.Key
        New      track.Key
}

func (e *ErrDuplicateTrackAlias) Error() string {
        return fmt.Sprintf(
                "moqt/session: Track Alias %d already in use for a different track — DUPLICATE_TRACK_ALIAS",
                e.Alias,
        )
}

// InboundTrack is what an inbound Track Alias is bound to (§11.1).
type InboundTrack struct {
        Key track.Key

        // DefaultPublisherPriority is the DEFAULT_PUBLISHER_PRIORITY (§12.4) of
        // the message that bound the alias, or 128 when omitted. Subgroups and
        // datagrams with the DEFAULT_PRIORITY bit inherit it (§11.4.2, §11.3.1).
        DefaultPublisherPriority uint8

        // MaxCacheDuration is the MAX_CACHE_DURATION (§12.3) of the same message,
        // if HasMaxCacheDuration.
        MaxCacheDuration    time.Duration
        HasMaxCacheDuration bool
}

// RegisterInboundTrack records that the peer has assigned alias to the track
// identified by key, along with the Track Properties of the message that did
// so. This MUST be called by the subscriber when it receives a SUBSCRIBE_OK
// (whose TrackAlias field is the alias) and by the server when it receives a
// PUBLISH (whose TrackAlias field is the alias).
//
// Registering an alias again for the same track counts one more registration
// (§5.1 allows subscriptions to share an alias); it stays registered until
// each is released by [Session.UnregisterInboundTrackAlias]. Assumption: the
// latest Track Properties replace earlier ones — the draft does not say which
// a shared alias carries, and a release cannot tell which registration it
// ends, so the survivor may keep a released one's properties.
//
// If alias is registered for a different track, the session is closed with
// DUPLICATE_TRACK_ALIAS and *ErrDuplicateTrackAlias returned: "it MUST close
// the session with error DUPLICATE_TRACK_ALIAS" (§11.1).
func (s *Session) RegisterInboundTrack(alias uint64, key track.Key, trackProperties []byte) error {
        in := InboundTrack{
                Key:                      key,
                DefaultPublisherPriority: message.TrackDefaultPublisherPriority(trackProperties),
        }
        in.MaxCacheDuration, in.HasMaxCacheDuration = message.TrackMaxCacheDuration(trackProperties)
        if err := s.registerInboundTrack(alias, in); err != nil {
                _ = s.Close(moqt.SessionDuplicateTrackAlias, err.Error())
                return err
        }
        return nil
}

func (s *Session) registerInboundTrack(alias uint64, in InboundTrack) error {
        s.mu.Lock()
        defer s.mu.Unlock()
        if existing, ok := s.inboundAliases[alias]; ok {
                if existing.Key != in.Key {
                        return &ErrDuplicateTrackAlias{Alias: alias, Existing: existing.Key, New: in.Key}
                }
                s.inboundAliases[alias] = in
                s.inboundAliasRefs[alias]++
                return nil
        }
        s.inboundAliases[alias] = in
        s.inboundAliasRefs[alias] = 1
        close(s.aliasRegistered)
        s.aliasRegistered = make(chan struct{})
        return nil
}

// RegisterInboundTrackAlias is [Session.RegisterInboundTrack] for a message
// that carried no Track Properties.
func (s *Session) RegisterInboundTrackAlias(alias uint64, key track.Key) error {
        return s.RegisterInboundTrack(alias, key, nil)
}

// UnregisterInboundTrackAlias releases one registration of alias (see
// [Session.RegisterInboundTrack]); the last release removes it. Call it once
// the subscription is Terminated (§5.1), for an alias registered with
// RegisterInboundTrack directly: a [Subscription] or [IncomingPublication]
// releases its own.
//
// Unregistering an alias that was never registered is a no-op.
func (s *Session) UnregisterInboundTrackAlias(alias uint64) {
        s.mu.Lock()
        defer s.mu.Unlock()
        if s.inboundAliasRefs[alias] > 1 {
                s.inboundAliasRefs[alias]--
                return
        }
        delete(s.inboundAliases, alias)
        delete(s.inboundAliasRefs, alias)
}

// LookupInboundTrack returns what alias was bound to by an earlier
// [Session.RegisterInboundTrack] call, or (zero, false) if the alias is not
// currently registered.
func (s *Session) LookupInboundTrack(alias uint64) (InboundTrack, bool) {
        s.mu.Lock()
        defer s.mu.Unlock()
        in, ok := s.inboundAliases[alias]
        return in, ok
}

// awaitInboundTrack is [Session.LookupInboundTrack] that waits for alias to be
// registered, until ctx ends or the session closes.
func (s *Session) awaitInboundTrack(ctx context.Context, alias uint64) (InboundTrack, bool) {
        for {
                s.mu.Lock()
                in, ok := s.inboundAliases[alias]
                registered := s.aliasRegistered
                s.mu.Unlock()
                if ok {
                        return in, true
                }
                select {
                case <-registered:
                case <-ctx.Done():
                        return InboundTrack{}, false
                case <-s.done:
                        return InboundTrack{}, false
                }
        }
}

// LookupInboundTrackAlias is [Session.LookupInboundTrack] reduced to the
// track.Key.
func (s *Session) LookupInboundTrackAlias(alias uint64) (track.Key, bool) {
        in, ok := s.LookupInboundTrack(alias)
        return in.Key, ok
}

// Package wtconn adapts github.com/quic-go/webtransport-go's *webtransport.Session
// to the transport-neutral session.Conn interface.
//
// This is the WebTransport counterpart of the quicconn package. It is the sole
// boundary in the moqt tree where webtransport-go's concrete types meet the
// session abstraction. Putting it in a dedicated subpackage lets the rest of
// pkg/moqt (and its tests) stay independent of webtransport-go's surface.
//
// webtransport-go uses webtransport.StreamErrorCode (uint32) for stream-level
// error codes and webtransport.SessionErrorCode (uint32) for session-level
// error codes; session.Conn / SendStream / ReceiveStream use plain uint64.
// The wrappers below do the narrowing conversion at each call site. MoQT
// error codes fit comfortably in 32 bits (the largest defined code is 0x34),
// so no information is lost in practice.
package wtconn

import (
        "context"
        "errors"
        "log/slog"
        "net"
        "net/http"
        "sync"

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

        "github.com/floatdrop/moq-go/pkg/moqt/session"
)

// New wraps s so it satisfies session.Conn.
func New(s *webtransport.Session) session.Conn { return &conn{s: s} }

// Compile-time satisfaction check.
var _ session.Conn = (*conn)(nil)

// conn holds a *webtransport.Session by named field rather than embedding.
// Embedding would promote webtransport's CloseWithError(SessionErrorCode, string)
// onto the wrapper; the session.Conn interface demands CloseWithError(uint64, string).
// Two methods of the same name with different signatures aren't allowed on a
// single Go type, so we delegate explicitly.
type conn struct{ s *webtransport.Session }

func (c *conn) OpenUniStream() (session.SendStream, error) {
        s, err := c.s.OpenUniStream()
        if err != nil {
                if _, ok := errors.AsType[*quic.StreamLimitReachedError](err); ok {
                        return nil, session.ErrNoStreamCredit
                }
                return nil, err
        }
        return &sendStream{s: s}, nil
}

func (c *conn) AcceptUniStream(ctx context.Context) (session.ReceiveStream, error) {
        s, err := c.s.AcceptUniStream(ctx)
        if err != nil {
                return nil, err
        }
        return &recvStream{s: s}, nil
}

// OpenStream opens a bidirectional stream without blocking. webtransport-go's
// OpenStream delegates to the underlying *quic.Conn, so an exhausted stream
// limit surfaces as a *quic.StreamLimitReachedError; we map that onto
// session.ErrNoStreamCredit for transport-neutral detection.
func (c *conn) OpenStream() (session.Stream, error) {
        s, err := c.s.OpenStream()
        if err != nil {
                if _, ok := errors.AsType[*quic.StreamLimitReachedError](err); ok {
                        return nil, session.ErrNoStreamCredit
                }
                return nil, err
        }
        return &bidiStream{s: s}, nil
}

func (c *conn) AcceptStream(ctx context.Context) (session.Stream, error) {
        s, err := c.s.AcceptStream(ctx)
        if err != nil {
                return nil, err
        }
        return &bidiStream{s: s}, nil
}

func (c *conn) CloseWithError(code uint64, reason string) error {
        //nolint:gosec // G115: MoQT session error codes fit uint32 (WebTransport's error-code width).
        return c.s.CloseWithError(webtransport.SessionErrorCode(code), reason)
}

func (c *conn) Context() context.Context { return c.s.Context() }

// IsWebTransport reports that this Conn runs over WebTransport, which
// §10.3.1.1 and §10.3.1.2 make relevant: the PATH and AUTHORITY setup options
// MUST NOT be used on a WebTransport session, since HTTP/3 already carries that
// information in the CONNECT request. The session layer asserts for this
// method to refuse those options before sending them.
//
// Not part of [session.Conn] on purpose — one rule needs it, and putting it on
// the interface would tax every adapter, third-party ones included, for a check
// the QUIC adapters answer trivially.
func (c *conn) IsWebTransport() bool { return true }

func (c *conn) SendDatagram(payload []byte) error {
        return c.s.SendDatagram(payload)
}

func (c *conn) ReceiveDatagram(ctx context.Context) ([]byte, error) {
        return c.s.ReceiveDatagram(ctx)
}

// sendStream wraps *webtransport.SendStream. Named field for the same reason
// as conn.
type sendStream struct{ s *webtransport.SendStream }

func (s *sendStream) Write(p []byte) (int, error) { return s.s.Write(p) }
func (s *sendStream) Close() error                { return s.s.Close() }
func (s *sendStream) CancelWrite(code uint64) {
        //nolint:gosec // G115: MoQT stream error codes fit uint32 (WebTransport's error-code width).
        s.s.CancelWrite(webtransport.StreamErrorCode(code))
}

// Context is cancelled when the send side is closed or reset, or the peer
// sends STOP_SENDING. webtransport-go's SendStream.Context() provides this
// directly.
func (s *sendStream) Context() context.Context { return s.s.Context() }

// recvStream wraps *webtransport.ReceiveStream.
type recvStream struct{ s *webtransport.ReceiveStream }

func (s *recvStream) Read(p []byte) (int, error) { return s.s.Read(p) }
func (s *recvStream) CancelRead(code uint64) {
        //nolint:gosec // G115: MoQT stream error codes fit uint32 (WebTransport's error-code width).
        s.s.CancelRead(webtransport.StreamErrorCode(code))
}

// bidiStream wraps *webtransport.Stream.
type bidiStream struct{ s *webtransport.Stream }

func (s *bidiStream) Read(p []byte) (int, error)  { return s.s.Read(p) }
func (s *bidiStream) Write(p []byte) (int, error) { return s.s.Write(p) }
func (s *bidiStream) Close() error                { return s.s.Close() }
func (s *bidiStream) CancelRead(code uint64) {
        //nolint:gosec // G115: MoQT stream error codes fit uint32 (WebTransport's error-code width).
        s.s.CancelRead(webtransport.StreamErrorCode(code))
}
func (s *bidiStream) CancelWrite(code uint64) {
        //nolint:gosec // G115: MoQT stream error codes fit uint32 (WebTransport's error-code width).
        s.s.CancelWrite(webtransport.StreamErrorCode(code))
}

// Context is cancelled when the send side is closed or reset, or the peer
// sends STOP_SENDING. webtransport-go's Stream embeds SendStream which has
// Context().
func (s *bidiStream) Context() context.Context { return s.s.Context() }

// defaultBacklog bounds the pending-session queue used by [Listener].
// A small queue absorbs handler-invocation bursts while the relay accept
// loop catches up.
const defaultBacklog = 16

// Listener adapts a *webtransport.Server so it can be handed directly
// to the relay's accept loop. WebTransport sessions arrive via HTTP/3
// handler invocations, not via a synchronous Accept on a socket, so
// the listener registers a path handler on the caller's *http.ServeMux
// and bridges accepted sessions through a buffered channel.
//
// The caller owns:
//
//   - The *webtransport.Server (typically constructed with
//     [webtransport.ConfigureHTTP3Server] on the underlying
//     *http3.Server).
//   - The HTTP/3 server's lifecycle: ListenAndServe / Serve on the
//     desired socket, and Close on shutdown. [Listener.Close] only
//     stops Accept from yielding new sessions so the relay's accept
//     loop unwinds — it does NOT close the *webtransport.Server.
//
// The Listener type satisfies relay.Listener structurally without
// importing pkg/relay.
type Listener struct {
        server *webtransport.Server
        addr   net.Addr
        queue  chan session.Conn

        closeOnce sync.Once
        done      chan struct{}
}

// NewListener registers a WebTransport upgrade handler at path on mux
// and returns a Listener suitable for relay.New.
//
//   - server: the configured *webtransport.Server. Upgrade is called
//     on this server for every inbound request that hits path.
//   - mux: the HTTP/3 server's request mux. The Listener does NOT
//     mount its own mux so the caller can multiplex MOQT-over-
//     WebTransport with other HTTP/3 endpoints on the same server.
//   - path: the HTTP/3 path the WebTransport CONNECT must target
//     (e.g. "/moq").
//   - addr: the address the Listener reports via [Listener.Addr];
//     pass nil if you don't need it (the relay only uses it for
//     log lines).
//   - queueSize: bounds the pending-session backlog before the
//     upgrade handler starts dropping sessions on the floor. Pass
//     0 for the package default ([defaultBacklog]).
//
// Typical wiring:
//
//        h3 := &http3.Server{TLSConfig: tlsCfg}
//        webtransport.ConfigureHTTP3Server(h3)
//        wts := &webtransport.Server{H3: h3, CheckOrigin: …}
//        mux := http.NewServeMux()
//        udpConn, _ := net.ListenPacket("udp", ":4433")
//        wts.H3.Handler = mux
//
//        listener := wtconn.NewListener(wts, mux, "/moq", udpConn.LocalAddr(), 0)
//        r := relay.New(listener, relay.Config{ … })
//
//        go wts.Serve(udpConn)
//        go r.Start(ctx)
func NewListener(
        server *webtransport.Server,
        mux *http.ServeMux,
        path string,
        addr net.Addr,
        queueSize int,
) *Listener {
        if queueSize <= 0 {
                queueSize = defaultBacklog
        }
        l := &Listener{
                server: server,
                addr:   addr,
                queue:  make(chan session.Conn, queueSize),
                done:   make(chan struct{}),
        }
        mux.HandleFunc(path, l.upgrade)
        return l
}

// upgrade is the HTTP/3 handler the Listener registers on the mux.
// It performs the WebTransport upgrade and hands the resulting
// *webtransport.Session to Accept via the bounded queue. If Accept is
// not draining (closed Listener or a burst exceeding queueSize), the
// freshly-accepted session is closed immediately so the client sees
// the failure rather than hanging.
func (l *Listener) upgrade(w http.ResponseWriter, r *http.Request) {
        ctx := r.Context()
        slog.DebugContext(ctx, "wtconn: upgrade request",
                "remote", r.RemoteAddr, "method", r.Method, "path", r.URL.Path,
                "proto", r.Proto, "origin", r.Header.Get("Origin"))
        sess, err := l.server.Upgrade(w, r)
        if err != nil {
                slog.DebugContext(ctx, "wtconn: upgrade failed", "remote", r.RemoteAddr, "err", err)
                return
        }
        slog.DebugContext(ctx, "wtconn: upgrade ok",
                "remote", r.RemoteAddr, "wt_protocol", sess.SessionState().ApplicationProtocol)
        select {
        case l.queue <- New(sess):
        case <-l.done:
                _ = sess.CloseWithError(0, "listener closed")
        default:
                slog.WarnContext(ctx, "wtconn: dropping session: accept backlog full", "remote", r.RemoteAddr)
                _ = sess.CloseWithError(0, "accept backlog full")
        }
}

// Accept blocks until the next upgraded WebTransport session arrives,
// then returns it as a session.Conn. ctx cancellation and Close both
// unblock Accept.
func (l *Listener) Accept(ctx context.Context) (session.Conn, error) {
        select {
        case c := <-l.queue:
                return c, nil
        case <-ctx.Done():
                return nil, ctx.Err()
        case <-l.done:
                return nil, net.ErrClosed
        }
}

// Addr returns the address the caller passed to NewListener, or nil
// if none was provided.
func (l *Listener) Addr() net.Addr { return l.addr }

// Close signals the Listener to stop yielding new sessions. The
// underlying *webtransport.Server keeps running; closing it is the
// caller's responsibility.
//
// Close is idempotent. Subsequent Accept calls return [net.ErrClosed].
func (l *Listener) Close() error {
        l.closeOnce.Do(func() { close(l.done) })
        return nil
}

// Package track provides domain types for MoQT track identification per
// §2.4.1 of draft-ietf-moq-transport: Full Track Name and a comparable Key
// derived from it for use as a Go map key.
package track

import (
        "encoding/binary"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// FullTrackName identifies a single track. The slice fields make the value
// type non-comparable; use Key for map indexing or exact-equality checks.
type FullTrackName struct {
        Namespace wire.TrackNamespace
        Name      []byte
}

// Key is a canonical, comparable representation of a Full Track Name. The
// namespace is wire-encoded (length-prefixed tuples per §2.4.1) so distinct
// tuple lists never collide — e.g. namespace ("a","b") with name "c" and
// namespace ("a") with name "bc" map to different Keys even though a naive
// byte concatenation would tie them.
type Key struct {
        namespace string // wire.TrackNamespace bytes, as string for comparability
        name      string
}

// Key returns the canonical map-friendly representation.
func (n FullTrackName) Key() Key {
        w := wire.NewWriter(nil)
        w.TrackNamespace(n.Namespace)
        return Key{namespace: string(w.Bytes()), name: string(n.Name)}
}

// NewKey is a convenience for callers that already have the namespace + name
// as separate values (e.g. fields parsed off a SUBSCRIBE / PUBLISH message).
func NewKey(ns wire.TrackNamespace, name []byte) Key {
        return FullTrackName{Namespace: ns, Name: name}.Key()
}

// Bytes returns the canonical binary encoding of the Key: an unsigned varint
// length prefix on the wire-encoded namespace, followed by the namespace bytes
// and then the track name. The length prefix keeps the (namespace, name)
// boundary unambiguous, so distinct splits never collide — the same guarantee
// the [Key] struct gives as a map key, made available as bytes.
//
// Distributed discovery backends need this: their FindTrack only receives a
// Key (never the originating [FullTrackName]), so they derive a stable storage
// key from it here. The encoding is deterministic but one-way — reconstruct a
// FullTrackName from stored metadata, not by parsing this.
func (k Key) Bytes() []byte {
        b := make([]byte, 0, binary.MaxVarintLen64+len(k.namespace)+len(k.name))
        b = binary.AppendUvarint(b, uint64(len(k.namespace)))
        b = append(b, k.namespace...)
        b = append(b, k.name...)
        return b
}

package uri

import (
        "errors"
        "fmt"
        "net/netip"
        "strings"
)

// CheckAuthority reports whether s is an RFC 3986 authority (§3.2) with a
// non-empty host, the form the AUTHORITY Setup Option carries (§10.3.1.1,
// §3.1.1):
//
//        authority = [ userinfo "@" ] host [ ":" port ]
//
// Unlike [Parse] (net/url), it accepts nothing outside the RFC 3986 grammar:
// no raw non-ASCII, no zone identifiers.
func CheckAuthority(s string) error {
        hostport := s
        if userinfo, rest, found := strings.Cut(s, "@"); found {
                if !validChars(userinfo, ":", true) {
                        return fmt.Errorf("uri: authority %q: invalid userinfo", s)
                }
                hostport = rest
        }
        host, port, err := splitHostPort(hostport)
        if err != nil {
                return fmt.Errorf("uri: authority %q: %w", s, err)
        }
        if host == "" {
                return fmt.Errorf("uri: authority %q has an empty host", s)
        }
        if strings.Trim(port, "0123456789") != "" {
                return fmt.Errorf("uri: authority %q: invalid port", s)
        }
        return nil
}

// CheckPathAndQuery reports whether s is an RFC 3986 path-abempty optionally
// followed by "?" and a query, the form the PATH Setup Option carries
// (§10.3.1.2):
//
//        path-abempty = *( "/" segment )
//        segment      = *pchar
//        query        = *( pchar / "/" / "?" )
func CheckPathAndQuery(s string) error {
        path, query, _ := strings.Cut(s, "?")
        if path != "" && path[0] != '/' {
                return fmt.Errorf("uri: path %q does not begin with \"/\"", s)
        }
        for seg := range strings.SplitSeq(path, "/") {
                if !validChars(seg, ":@", true) {
                        return fmt.Errorf("uri: path %q: invalid segment %q", s, seg)
                }
        }
        if !validChars(query, ":@/?", true) {
                return fmt.Errorf("uri: path %q: invalid query", s)
        }
        return nil
}

// splitHostPort splits host [ ":" port ] and checks the host: an IP-literal,
// or a reg-name (which covers IPv4address: its characters are a subset).
func splitHostPort(hostport string) (host, port string, err error) {
        if !strings.HasPrefix(hostport, "[") {
                host, port, _ = strings.Cut(hostport, ":")
                if !validChars(host, "", true) {
                        return "", "", errors.New("invalid host")
                }
                return host, port, nil
        }
        end := strings.IndexByte(hostport, ']')
        if end < 0 {
                return "", "", errors.New("unterminated IP-literal")
        }
        if !validIPLiteral(hostport[1:end]) {
                return "", "", errors.New("invalid IP-literal")
        }
        host, port = hostport[:end+1], hostport[end+1:]
        if port != "" && port[0] != ':' {
                return "", "", errors.New("junk after IP-literal")
        }
        return host, strings.TrimPrefix(port, ":"), nil
}

// validChars reports whether s consists of RFC 3986 unreserved characters,
// sub-delims, the bytes in extra, and — when pct is set — well-formed
// pct-encoded triplets.
func validChars(s, extra string, pct bool) bool {
        for i := 0; i < len(s); i++ {
                c := s[i]
                switch {
                case c == '%' && pct:
                        if i+2 >= len(s) || !isHex(s[i+1]) || !isHex(s[i+2]) {
                                return false
                        }
                        i += 2
                case isUnreserved(c), strings.IndexByte("!$&'()*+,;=", c) >= 0, strings.IndexByte(extra, c) >= 0:
                default:
                        return false
                }
        }
        return true
}

// validIPLiteral checks the inside of "[...]":
//
//        IP-literal = "[" ( IPv6address / IPvFuture  ) "]"
//        IPvFuture  = "v" 1*HEXDIG "." 1*( unreserved / sub-delims / ":" )
func validIPLiteral(s string) bool {
        if s != "" && (s[0] == 'v' || s[0] == 'V') {
                ver, rest, found := strings.Cut(s[1:], ".")
                return found && ver != "" && strings.Trim(ver, "0123456789abcdefABCDEF") == "" &&
                        rest != "" && validChars(rest, ":", false)
        }
        // RFC 3986 has no zone identifiers, which netip would accept.
        if strings.IndexByte(s, '%') >= 0 {
                return false
        }
        a, err := netip.ParseAddr(s)
        return err == nil && a.Is6()
}

func isUnreserved(c byte) bool {
        return 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z' || '0' <= c && c <= '9' ||
                c == '-' || c == '.' || c == '_' || c == '~'
}

func isHex(c byte) bool {
        return '0' <= c && c <= '9' || 'a' <= c && c <= 'f' || 'A' <= c && c <= 'F'
}

// Package uri parses and validates "moqt" URIs and their fragment identifiers
// as defined by draft-ietf-moq-transport-20 §3.1.1 and §3.1.2.
//
//        moqt-URI = "moqt" "://" authority path-abempty [ "?" query ]
//
// A parsed [URI] exposes everything the connection-setup paths need:
// [URI.HostPort] for dialing (applying the §3.1.3 default port of 443),
// [URI.Authority] and [URI.PathAndQuery] for the AUTHORITY / PATH Setup
// Options carried on a native-QUIC connection (§3.1.5 / §10.3.1), and
// [URI.HTTPSURL] for the https URL a WebTransport client connects to
// (§3.1.4). Fragments are parsed and validated but, per §3.1.2, are processed
// locally by the client and never transmitted to the server.
//
// The package depends only on the standard library so it can be used from any
// layer of the stack without pulling in the session machinery.
package uri

import (
        "fmt"
        "net"
        "net/url"
        "strings"
)

// Scheme is the URI scheme defined for MOQT servers (§3.1.1).
const Scheme = "moqt"

// DefaultPort is used when the authority omits an explicit port (§3.1.3:
// "If the port is omitted in the URI, a default port of 443 is used").
const DefaultPort = "443"

// URI is a parsed, validated "moqt" URI (§3.1.1).
type URI struct {
        // Authority is the host[:port] exactly as supplied (no default port
        // filled in). The host subcomponent is guaranteed non-empty.
        Authority string

        // Host is the host subcomponent of the authority, without any port.
        Host string

        // Port is the explicit port from the authority, or [DefaultPort] when the
        // authority omitted one.
        Port string

        // Path is the path-abempty component: either empty or beginning with "/".
        Path string

        // RawQuery is the query component without its leading "?", empty when the
        // URI carried no query.
        RawQuery string

        // Fragment is the parsed fragment identifier, or nil when the URI carried
        // none. Per §3.1.2 the fragment is processed locally and never sent to
        // the server.
        Fragment *Fragment
}

// Fragment is a parsed moqt URI fragment identifier (§3.1.2):
//
//        moqt://example.com/app#<type>:<value>
type Fragment struct {
        // Type is the registered fragment type identifier: a non-empty string of
        // ASCII lowercase letters, digits, and hyphens (a-z, 0-9, -).
        Type string

        // Value is the type-specific value following the first colon. Its
        // semantics are defined by the specification that registers Type.
        Value string
}

// Parse parses and validates a "moqt" URI per §3.1.1 / §3.1.2. It returns an
// error when the scheme is not "moqt", the URI is not hierarchical, the
// authority has an empty host, or a present fragment does not match the
// "type:value" grammar.
func Parse(raw string) (*URI, error) {
        u, err := url.Parse(raw)
        if err != nil {
                return nil, fmt.Errorf("uri: parse %q: %w", raw, err)
        }
        if u.Scheme != Scheme {
                return nil, fmt.Errorf("uri: scheme %q, want %q", u.Scheme, Scheme)
        }
        // A hierarchical URI ("scheme://...") leaves Opaque empty and fills Host.
        // An opaque form like "moqt:foo" is rejected.
        if u.Opaque != "" {
                return nil, fmt.Errorf("uri: %q is not hierarchical (expected %s://)", raw, Scheme)
        }
        host := u.Hostname()
        if host == "" {
                // §3.1.1: "The authority portion MUST NOT contain an empty host
                // portion."
                return nil, fmt.Errorf("uri: %q has an empty host", raw)
        }
        port := u.Port()
        if port == "" {
                port = DefaultPort
        }

        out := &URI{
                Authority: u.Host,
                Host:      host,
                Port:      port,
                // EscapedPath, not Path: the struct carries the RAW path-abempty
                // component. url.URL.Path is percent-DECODED — using it would turn
                // "/a%3Fb" into "/a?b", making the PATH Setup Option (§3.1.5, §10.3.1.2)
                // ambiguous and String()/HTTPSURL() emit invalid URIs.
                Path:     u.EscapedPath(),
                RawQuery: u.RawQuery,
        }

        // A '#' in the raw URI introduces a fragment component (§3.1.2). When
        // present it MUST match the "type:value" grammar, so a present-but-empty
        // or colon-less fragment is an error rather than silently ignored.
        if strings.IndexByte(raw, '#') >= 0 {
                frag, err := parseFragment(u.Fragment)
                if err != nil {
                        return nil, err
                }
                out.Fragment = frag
        }

        return out, nil
}

// parseFragment validates the "type:value" grammar of §3.1.2 against the
// (percent-decoded) fragment text.
func parseFragment(frag string) (*Fragment, error) {
        typ, val, ok := strings.Cut(frag, ":")
        if !ok {
                return nil, fmt.Errorf("uri: fragment %q missing \"type:value\" colon (§3.1.2)", frag)
        }
        if !validFragmentType(typ) {
                return nil, fmt.Errorf(
                        "uri: fragment type %q must be a non-empty run of ASCII [a-z0-9-] (§3.1.2)", typ)
        }
        return &Fragment{Type: typ, Value: val}, nil
}

// validFragmentType reports whether s is a non-empty run of ASCII lowercase
// letters, digits, and hyphens (§3.1.2).
func validFragmentType(s string) bool {
        if s == "" {
                return false
        }
        for i := range len(s) {
                c := s[i]
                switch {
                case c >= 'a' && c <= 'z':
                case c >= '0' && c <= '9':
                case c == '-':
                default:
                        return false
                }
        }
        return true
}

// HostPort returns the "host:port" string for dialing, applying the §3.1.3
// default port of 443 when the URI omitted one.
func (u *URI) HostPort() string {
        return net.JoinHostPort(u.Host, u.Port)
}

// PathAndQuery returns the path-abempty with the query appended, the value to
// carry in the PATH Setup Option (§3.1.5 / §10.3.1.2). It is empty when the
// URI has neither a path nor a query.
func (u *URI) PathAndQuery() string {
        if u.RawQuery == "" {
                return u.Path
        }
        return u.Path + "?" + u.RawQuery
}

// HTTPSURL converts the moqt URI to the https URL a WebTransport client
// connects to (§3.1.4): the scheme is replaced with https and the authority,
// path, and query are preserved. The fragment is omitted because it is
// processed locally and never sent to the server (§3.1.2).
func (u *URI) HTTPSURL() string {
        var b strings.Builder
        b.WriteString("https://")
        b.WriteString(u.Authority)
        b.WriteString(u.Path)
        if u.RawQuery != "" {
                b.WriteByte('?')
                b.WriteString(u.RawQuery)
        }
        return b.String()
}

// String reconstructs the moqt URI, including any fragment. The authority is
// emitted exactly as parsed, so a URI that omitted its port round-trips
// without a default port appearing.
func (u *URI) String() string {
        var b strings.Builder
        b.WriteString(Scheme)
        b.WriteString("://")
        b.WriteString(u.Authority)
        b.WriteString(u.Path)
        if u.RawQuery != "" {
                b.WriteByte('?')
                b.WriteString(u.RawQuery)
        }
        if u.Fragment != nil {
                b.WriteByte('#')
                b.WriteString(u.Fragment.Type)
                b.WriteByte(':')
                b.WriteString(u.Fragment.Value)
        }
        return b.String()
}

package wire

import (
        "encoding/binary"
        "fmt"
        "io"
)

// MaxControlMessagePayload is the largest payload that fits in a control
// message's 16-bit Length field (§10).
const MaxControlMessagePayload = 0xFFFF

// ReadFrame reads a single MoQT control-message frame (Type + Length + Payload)
// from r, returning the message type and the payload bytes. The returned
// payload is freshly allocated; the caller owns it.
//
// ReadFrame returns io.EOF only when r reports EOF before the type byte has
// been read; once any byte has been consumed, a truncated frame surfaces as
// io.ErrUnexpectedEOF.
func ReadFrame(r io.Reader) (uint64, []byte, error) {
        msgType, err := ReadVarint(NewByteReader(r))
        if err != nil {
                return 0, nil, err
        }
        var lenBuf [2]byte
        if _, err := io.ReadFull(r, lenBuf[:]); err != nil {
                if err == io.EOF {
                        err = io.ErrUnexpectedEOF
                }
                return 0, nil, err
        }
        length := binary.BigEndian.Uint16(lenBuf[:])
        if length == 0 {
                return msgType, nil, nil
        }
        payload := make([]byte, length)
        if _, err := io.ReadFull(r, payload); err != nil {
                if err == io.EOF {
                        err = io.ErrUnexpectedEOF
                }
                return 0, nil, err
        }
        return msgType, payload, nil
}

// WriteFrame writes a MoQT control-message frame to w. It returns an error if
// the payload exceeds MaxControlMessagePayload.
func WriteFrame(w io.Writer, msgType uint64, payload []byte) error {
        if len(payload) > MaxControlMessagePayload {
                return fmt.Errorf("moqt/wire: control message payload %d exceeds %d", len(payload), MaxControlMessagePayload)
        }
        hdr := make([]byte, 0, VarintLen(msgType)+2)
        hdr = AppendVarint(hdr, msgType)
        //nolint:gosec // G115: len(payload) is checked <= MaxControlMessagePayload (0xFFFF) above, so it fits 16 bits.
        hdr = append(hdr, byte(len(payload)>>8), byte(len(payload)))
        if _, err := w.Write(hdr); err != nil {
                return err
        }
        if len(payload) == 0 {
                return nil
        }
        _, err := w.Write(payload)
        return err
}

package wire

import (
        "cmp"
        "fmt"
        "slices"
)

// KVPair is a MoQT Key-Value-Pair (§1.4.3). When Type is even, IntVal carries
// the value (encoded as a single varint). When Type is odd, ByteVal carries the
// value (length-prefixed bytes).
//
// KVPairs are used for SETUP Options (§10.3.1); they appear delta-encoded by
// Type within a list, with the running "previous type" starting at zero.
type KVPair struct {
        Type    uint64
        IntVal  uint64
        ByteVal []byte
}

// IsBytes reports whether this KVPair carries length-prefixed bytes (Type odd)
// rather than a varint (Type even).
func (p KVPair) IsBytes() bool { return p.Type&1 == 1 }

// MaxKVPairValueBytes is the per-pair byte-value cap from §1.4.3.
const MaxKVPairValueBytes = 0xFFFF

// KVPair appends a single KVPair using prev as the running previous Type, and
// returns the new previous Type. The first pair in a list passes prev=0.
func (w *Writer) KVPair(p KVPair, prev uint64) uint64 {
        w.Varint(p.Type - prev)
        if p.IsBytes() {
                w.VarintBytes(p.ByteVal)
        } else {
                w.Varint(p.IntVal)
        }
        return p.Type
}

// KVPairs appends a list of KVPairs with delta encoding starting from prev=0.
// Pairs are sorted by Type before encoding so callers do not need to order them.
// The sort is stable: pairs of one Type keep the caller's order, which a
// repeated option can depend on: which SETUP REGISTERs fit the peer's cache
// depends on their order (§10.3.1.3, §10.3.1.4).
func (w *Writer) KVPairs(pairs []KVPair) {
        slices.SortStableFunc(pairs, func(a, b KVPair) int { return cmp.Compare(a.Type, b.Type) })
        var prev uint64
        for _, p := range pairs {
                prev = w.KVPair(p, prev)
        }
}

// KVPair reads a single KVPair using prev as the running previous Type, and
// returns the pair plus the new previous Type. A byte value is a copy.
func (r *Reader) KVPair(prev uint64) (KVPair, uint64, error) {
        return r.kvPair(prev, true)
}

// KVPairView is [Reader.KVPair] without the copy, for per-Object inspection
// that must not allocate: a byte value aliases the reader's buffer.
func (r *Reader) KVPairView(prev uint64) (KVPair, uint64, error) {
        return r.kvPair(prev, false)
}

func (r *Reader) kvPair(prev uint64, copyBytes bool) (KVPair, uint64, error) {
        delta, err := r.Varint()
        if err != nil {
                return KVPair{}, prev, err
        }
        if delta > ^uint64(0)-prev {
                return KVPair{}, prev, fmt.Errorf("moqt/wire: kv pair type delta overflow (prev=%d delta=%d)", prev, delta)
        }
        t := prev + delta
        p := KVPair{Type: t}
        if p.IsBytes() {
                b, err := r.varintBytes(copyBytes)
                if err != nil {
                        return KVPair{}, prev, err
                }
                if len(b) > MaxKVPairValueBytes {
                        return KVPair{}, prev, fmt.Errorf(
                                "moqt/wire: kv pair value length %d exceeds %d",
                                len(b),
                                MaxKVPairValueBytes,
                        )
                }
                p.ByteVal = b
        } else {
                v, err := r.Varint()
                if err != nil {
                        return KVPair{}, prev, err
                }
                p.IntVal = v
        }
        return p, t, nil
}

// KVPairsRemaining reads KVPairs until the reader is empty. This is used for
// SETUP, where Setup Options span the entire control-message payload (§10.3).
func (r *Reader) KVPairsRemaining() ([]KVPair, error) {
        var (
                pairs []KVPair
                prev  uint64
        )
        for !r.Empty() {
                p, next, err := r.KVPair(prev)
                if err != nil {
                        return nil, err
                }
                pairs = append(pairs, p)
                prev = next
        }
        return pairs, nil
}

package wire

import (
        "bytes"
        "fmt"
)

// MaxTrackNamespaceFields is the upper bound on tuple count per §2.4.1.
const MaxTrackNamespaceFields = 32

// MaxFullTrackNameBytes is the upper bound on the sum of all namespace field
// lengths plus the track name length per §2.4.1.
const MaxFullTrackNameBytes = 4096

// TrackNamespace is an ordered set of 0..32 binary fields (§2.4.1).
type TrackNamespace [][]byte

// Namespace builds a TrackNamespace from string fields — the ergonomic form of
// the TrackNamespace{[]byte("a"), []byte("b")} literal. Each argument becomes
// one §2.4.1 field, in order. Namespace fields MAY contain arbitrary bytes; for
// non-UTF-8 fields use the [][]byte literal directly.
func Namespace(parts ...string) TrackNamespace {
        ns := make(TrackNamespace, len(parts))
        for i, p := range parts {
                ns[i] = []byte(p)
        }
        return ns
}

// TrackNamespace reads a TrackNamespace per §2.4.1. Each field must be at least
// one byte; the tuple count must not exceed MaxTrackNamespaceFields. The
// returned slices are owned by the caller (see Reader.FixedBytes).
func (r *Reader) TrackNamespace() (TrackNamespace, error) {
        count, err := r.Varint()
        if err != nil {
                return nil, err
        }
        if count > MaxTrackNamespaceFields {
                return nil, fmt.Errorf("moqt/wire: track namespace has %d fields, max %d", count, MaxTrackNamespaceFields)
        }
        ns := make(TrackNamespace, 0, count)
        total := 0
        for i := range count {
                field, err := r.VarintBytes()
                if err != nil {
                        return nil, err
                }
                if len(field) == 0 {
                        return nil, fmt.Errorf("moqt/wire: track namespace field %d has zero length", i)
                }
                // §2.4.1: "If an endpoint receives a Track Namespace or a Full
                // Track Name exceeding 4,096 bytes, it MUST close the session with
                // a PROTOCOL_VIOLATION." The namespace-only half is enforced here,
                // at the single parse point every message shares; messages that
                // also carry a Track Name check the combined length in Validate.
                total += len(field)
                if total > MaxFullTrackNameBytes {
                        return nil, fmt.Errorf(
                                "moqt/wire: track namespace exceeds %d bytes (§2.4.1)", MaxFullTrackNameBytes)
                }
                ns = append(ns, field)
        }
        return ns, nil
}

// TrackNamespace appends a TrackNamespace per §2.4.1.
func (w *Writer) TrackNamespace(ns TrackNamespace) {
        w.Varint(uint64(len(ns)))
        for _, field := range ns {
                w.VarintBytes(field)
        }
}

// ByteLen reports the sum of field lengths (used to enforce the 4096-byte
// Full Track Name limit alongside the Track Name's length).
func (ns TrackNamespace) ByteLen() int {
        total := 0
        for _, f := range ns {
                total += len(f)
        }
        return total
}

// HasPrefix reports whether prefix is a (non-strict) prefix of ns in the
// field-by-field sense of §2.4.1. A zero-length prefix matches every ns,
// matching the §6.1 "Either message with zero Track Namespace fields
// indicates the sender is interested in all namespaces" rule used by
// SUBSCRIBE_NAMESPACE / SUBSCRIBE_TRACKS matching.
//
// Fields are compared as opaque binary; namespace components MAY contain
// any bytes per §2.4.1.
func (ns TrackNamespace) HasPrefix(prefix TrackNamespace) bool {
        if len(prefix) > len(ns) {
                return false
        }
        for i, p := range prefix {
                if !bytes.Equal(p, ns[i]) {
                        return false
                }
        }
        return true
}

// String renders the namespace as "/comp1/comp2/..." with each
// component shown verbatim. Intended for log and error messages;
// callers that need a strict serialization should use Writer.TrackNamespace.
func (ns TrackNamespace) String() string {
        var b []byte
        b = append(b, '/')
        for i, c := range ns {
                if i > 0 {
                        b = append(b, '/')
                }
                b = append(b, c...)
        }
        return string(b)
}

package wire

// Scanner is a sticky-error decoding cursor over a [Reader]. Each accessor
// reads one field into the supplied pointer and records the first error it
// hits; once an error is recorded every later accessor is a no-op until Err is
// consulted. It removes the repetitive per-field error handling that otherwise
// dominates message Parse methods:
//
//        func (m *Subscribe) Parse(r *wire.Reader) error {
//                s := r.Scanner()
//                s.Varint(&m.RequestID)
//                s.TrackNamespace(&m.Namespace)
//                s.VarintBytes(&m.Name)
//                if err := s.Err(); err != nil {
//                        return err
//                }
//                return m.Parameters.parse(r)
//        }
//
// A Scanner delegates to its Reader and advances the same read offset, so the
// underlying Reader stays usable directly after the Scanner (e.g. for
// Parameters.parse or a RemainingBytes tail) once Err reports no error.
//
// Scanner only wraps the in-memory [Reader]; the streaming [StreamReader] /
// [Decoder] path is unaffected.
type Scanner struct {
        r   *Reader
        err error
}

// Scanner returns a sticky-error cursor over r.
func (r *Reader) Scanner() *Scanner { return &Scanner{r: r} }

// Err returns the first error any accessor recorded, or nil.
func (s *Scanner) Err() error { return s.err }

// scan runs read unless an error is already pending, storing the value into dst
// or recording the error. It is the single implementation behind every typed
// accessor below.
func scan[T any](s *Scanner, dst *T, read func() (T, error)) {
        if s.err != nil {
                return
        }
        v, err := read()
        if err != nil {
                s.err = err
                return
        }
        *dst = v
}

// Varint reads a leading-ones varint (§1.4.1) into dst.
func (s *Scanner) Varint(dst *uint64) { scan(s, dst, s.r.Varint) }

// UInt8 reads a single byte into dst.
func (s *Scanner) UInt8(dst *uint8) { scan(s, dst, s.r.UInt8) }

// VarintBytes reads a varint-length-prefixed byte slice into dst.
func (s *Scanner) VarintBytes(dst *[]byte) { scan(s, dst, s.r.VarintBytes) }

// ReasonPhrase reads a §1.4.4 reason phrase into dst.
func (s *Scanner) ReasonPhrase(dst *string) { scan(s, dst, s.r.ReasonPhrase) }

// TrackNamespace reads a §2.4.1 track namespace into dst.
func (s *Scanner) TrackNamespace(dst *TrackNamespace) { scan(s, dst, s.r.TrackNamespace) }

// KVPairsRemaining reads delta-encoded KV pairs to end-of-buffer into dst.
func (s *Scanner) KVPairsRemaining(dst *[]KVPair) { scan(s, dst, s.r.KVPairsRemaining) }

package wire

import (
        "encoding/binary"
        "errors"
        "io"
        "math/bits"
)

// MoQT variable-length integers (draft-ietf-moq-transport-20 §1.4.1).
//
// Unlike QUIC's RFC 9000 §16 varints — which use the high 2 bits of the first
// byte to select a 1/2/4/8-byte length — MoQT uses a "leading-ones" scheme: the
// number of leading 1 bits in the first byte gives the encoded length (1 to 9
// bytes). The value occupies the bits after the first 0, plus all subsequent
// bytes, in network byte order.
//
//        Leading bits | Length | First byte | Value bytes
//        0            | 1      | 0xxxxxxx   | (none)
//        10           | 2      | 10xxxxxx   | 1
//        110          | 3      | 110xxxxx   | 2
//        1110         | 4      | 1110xxxx   | 3
//        11110        | 5      | 11110xxx   | 4
//        111110       | 6      | 111110xx   | 5
//        1111110      | 7      | 1111110x   | 6
//        11111110     | 8      | 11111110   | 7
//        11111111     | 9      | 11111111   | 8
//
// §1.4.1 also notes integers "do not need to be encoded using the minimum
// number of bytes", so decoders accept non-minimal encodings; AppendVarint
// always emits the minimal form.

// VarintLen returns the number of bytes AppendVarint uses to encode v.
func VarintLen(v uint64) int {
        switch {
        case v < 1<<7:
                return 1
        case v < 1<<14:
                return 2
        case v < 1<<21:
                return 3
        case v < 1<<28:
                return 4
        case v < 1<<35:
                return 5
        case v < 1<<42:
                return 6
        case v < 1<<49:
                return 7
        case v < 1<<56:
                return 8
        default:
                return 9
        }
}

// AppendVarint appends the minimal leading-ones encoding of v to dst and
// returns the extended slice.
func AppendVarint(dst []byte, v uint64) []byte {
        n := VarintLen(v)
        if n == 9 {
                // 0xFF prefix (8 leading ones) followed by the full 64-bit value.
                dst = append(dst, 0xFF)
                return binary.BigEndian.AppendUint64(dst, v)
        }
        // For n<=8 the value uses 7n bits, so it fits in the low n bytes of its
        // big-endian form, leaving the top n bits of the first byte free for the
        // (n-1)-leading-ones prefix.
        var buf [8]byte
        binary.BigEndian.PutUint64(buf[:], v)
        out := buf[8-n:]
        out[0] |= ^(byte(0xFF) >> (n - 1))
        return append(dst, out...)
}

// varintLenFromFirst returns the total encoded length implied by the first
// byte's leading-ones count.
func varintLenFromFirst(first byte) int {
        // The encoded length is (leading ones)+1, except 0xFF (8 leading ones)
        // which denotes the 9-byte form.
        ones := bits.LeadingZeros8(^first)
        if ones == 8 {
                return 9
        }
        return ones + 1
}

// ParseVarint decodes a leading-ones varint from the front of b, returning the
// value and the number of bytes consumed. It returns ErrShortBuffer if b is
// shorter than the encoding the first byte announces.
func ParseVarint(b []byte) (uint64, int, error) {
        if len(b) == 0 {
                return 0, 0, ErrShortBuffer
        }
        n := varintLenFromFirst(b[0])
        if len(b) < n {
                return 0, 0, ErrShortBuffer
        }
        if n == 9 {
                return binary.BigEndian.Uint64(b[1:9]), 9, nil
        }
        // Low (8-n) bits of the first byte are value bits (0 for n==8).
        v := uint64(b[0] & (0xFF >> n))
        for i := 1; i < n; i++ {
                v = v<<8 | uint64(b[i])
        }
        return v, n, nil
}

// ReadVarint decodes a leading-ones varint from r, reading exactly the bytes of
// one encoding (never any look-ahead), so it is safe to call repeatedly on the
// same underlying stream.
func ReadVarint(r io.ByteReader) (uint64, error) {
        first, err := r.ReadByte()
        if err != nil {
                return 0, err
        }
        n := varintLenFromFirst(first)
        if n == 1 {
                return uint64(first), nil
        }
        var rest [8]byte
        for i := range n - 1 {
                b, err := r.ReadByte()
                if err != nil {
                        if errors.Is(err, io.EOF) {
                                err = io.ErrUnexpectedEOF
                        }
                        return 0, err
                }
                rest[i] = b
        }
        if n == 9 {
                return binary.BigEndian.Uint64(rest[:8]), nil
        }
        v := uint64(first & (0xFF >> n))
        for i := range n - 1 {
                v = v<<8 | uint64(rest[i])
        }
        return v, nil
}

// NewByteReader adapts an io.Reader to io.ByteReader by reading a single byte
// per call, with no buffering or look-ahead, so a varint read leaves the
// underlying reader positioned exactly after the varint.
func NewByteReader(r io.Reader) io.ByteReader {
        if br, ok := r.(io.ByteReader); ok {
                return br
        }
        return &byteReaderAdapter{r: r}
}

// Package wire implements MoQT wire-format primitives per
// draft-ietf-moq-transport-20: variable-length integers (§1.4.1, up to 2^64-1),
// reason phrases (§1.4.4), track namespaces (§2.4.1), key-value pairs used in
// SETUP options (§1.4.3, §10.3.1), and control-message framing (§10).
//
// Encoding follows an append-style: builders accumulate bytes into a Writer.
// Decoding uses a stateful Reader bounded by an input buffer; running past the
// buffer yields ErrShortBuffer, which the message layer maps to a session-level
// PROTOCOL_VIOLATION (§3.5).
//
// For streaming decoding (e.g. data uni-streams), use StreamReader which wraps
// an io.Reader and exposes the same Decoder interface as Reader.
package wire

import (
        "errors"
        "fmt"
        "io"
        "unicode/utf8"
)

// ErrShortBuffer is returned when a read would consume bytes past the end of
// the input buffer. Callers should treat this as a malformed message.
var ErrShortBuffer = errors.New("moqt/wire: short buffer")

// ErrFieldTooLarge is returned by StreamReader when a length-prefixed field
// claims more bytes than [MaxStreamFieldSize]. Callers should treat it as a
// malformed message (PROTOCOL_VIOLATION, §3.5).
var ErrFieldTooLarge = errors.New("moqt/wire: field exceeds maximum size")

// MaxStreamFieldSize bounds a single length-prefixed field (object payload,
// properties blob, name, …) that [StreamReader] will allocate for. Because a
// StreamReader reads from an unbounded io.Reader, FixedBytes refuses to
// pre-allocate more than this for a peer-supplied length, so a malicious peer
// cannot trigger an unbounded allocation by claiming a huge length before
// sending the bytes. (The in-memory [Reader] is already bounded by its buffer
// and is not subject to this limit.) The default is generous enough for large
// media objects such as 4K keyframes; deployments carrying larger objects can
// raise it.
var MaxStreamFieldSize = 16 << 20 // 16 MiB

// Reader consumes MoQT wire primitives from an in-memory buffer. It tracks the
// read offset; partial reads do not advance the offset.
type Reader struct {
        buf []byte
        off int
}

// NewReader returns a Reader over buf. buf is not copied; the caller must not
// mutate it while the Reader is in use.
func NewReader(buf []byte) *Reader { return &Reader{buf: buf} }

// Remaining returns the number of bytes left to consume.
func (r *Reader) Remaining() int { return len(r.buf) - r.off }

// Empty reports whether the reader has consumed all bytes.
func (r *Reader) Empty() bool { return r.off >= len(r.buf) }

// Varint reads a MoQT leading-ones varint (§1.4.1, 1–9 bytes).
func (r *Reader) Varint() (uint64, error) {
        v, n, err := ParseVarint(r.buf[r.off:])
        if err != nil {
                return 0, err
        }
        r.off += n
        return v, nil
}

// UInt8 reads a single byte.
func (r *Reader) UInt8() (uint8, error) {
        if r.Remaining() < 1 {
                return 0, ErrShortBuffer
        }
        v := r.buf[r.off]
        r.off++
        return v, nil
}

// FixedBytes reads exactly n bytes. The returned slice is a fresh copy that
// the caller owns; mutating it does not affect the Reader's buffer, and
// retaining it does not pin the buffer for GC. Zero-length reads return nil.
func (r *Reader) FixedBytes(n int) ([]byte, error) {
        // n < 0 when a caller's int conversion of a varint >= 2^63 wrapped.
        if n < 0 || r.Remaining() < n {
                return nil, ErrShortBuffer
        }
        if n == 0 {
                return nil, nil
        }
        out := make([]byte, n)
        copy(out, r.buf[r.off:r.off+n])
        r.off += n
        return out, nil
}

// RemainingBytes consumes and returns a copy of all unconsumed bytes. Used
// when a message has a trailing variable-length field bounded only by the
// outer frame length (e.g. Track Properties in SUBSCRIBE_OK / PUBLISH).
// Zero-length returns nil.
func (r *Reader) RemainingBytes() []byte {
        n := r.Remaining()
        if n == 0 {
                return nil
        }
        out := make([]byte, n)
        copy(out, r.buf[r.off:])
        r.off += n
        return out
}

// VarintBytes reads a varint length followed by that many bytes. The returned
// slice is owned by the caller (see FixedBytes).
func (r *Reader) VarintBytes() ([]byte, error) {
        return r.varintBytes(true)
}

// varintBytes is VarintBytes, returning a slice of the buffer rather than a
// copy when copyBytes is false.
func (r *Reader) varintBytes(copyBytes bool) ([]byte, error) {
        n, err := r.Varint()
        if err != nil {
                return nil, err
        }
        // §1.4.1 varints reach 2^64-1: bound n before the int conversion.
        if n > uint64(r.Remaining()) { //nolint:gosec // G115: Remaining() is len(buf)-off >= 0.
                return nil, ErrShortBuffer
        }
        if copyBytes {
                return r.FixedBytes(int(n)) //nolint:gosec // G115: n <= Remaining() above.
        }
        end := r.off + int(n) //nolint:gosec // G115: n <= Remaining() above.
        b := r.buf[r.off:end:end]
        r.off = end
        return b, nil
}

// ReasonPhrase reads a varint-length-prefixed UTF-8 string per §1.4.4. The
// maximum allowed length is 1024 bytes; exceeding this yields an error that
// the caller should map to PROTOCOL_VIOLATION.
func (r *Reader) ReasonPhrase() (string, error) {
        n, err := r.Varint()
        if err != nil {
                return "", err
        }
        if n > MaxReasonPhraseBytes {
                return "", fmt.Errorf("moqt/wire: reason phrase length %d exceeds %d", n, MaxReasonPhraseBytes)
        }
        b, err := r.FixedBytes(int(n))
        if err != nil {
                return "", err
        }
        return string(b), nil
}

// Writer accumulates encoded MoQT bytes. The zero value is ready to use.
type Writer struct {
        buf []byte
}

// NewWriter returns a Writer that appends to buf (which may be nil). Use
// Bytes to retrieve the accumulated output.
func NewWriter(buf []byte) *Writer { return &Writer{buf: buf} }

// Bytes returns the accumulated output. The returned slice aliases the
// Writer's internal buffer.
func (w *Writer) Bytes() []byte { return w.buf }

// Reset clears the writer's buffer, allowing it to be reused.
func (w *Writer) Reset() { w.buf = w.buf[:0] }

// Varint appends a MoQT leading-ones varint (§1.4.1).
func (w *Writer) Varint(v uint64) { w.buf = AppendVarint(w.buf, v) }

// UInt8 appends a single byte.
func (w *Writer) UInt8(v uint8) { w.buf = append(w.buf, v) }

// FixedBytes appends raw bytes without any length prefix.
func (w *Writer) FixedBytes(p []byte) { w.buf = append(w.buf, p...) }

// VarintBytes appends a varint length followed by the bytes themselves.
func (w *Writer) VarintBytes(p []byte) {
        w.Varint(uint64(len(p)))
        w.FixedBytes(p)
}

// MaxReasonPhraseBytes is the §1.4.4 cap on an encoded reason phrase. [Reader]
// rejects anything longer, and [Writer.ReasonPhrase] truncates to it.
const MaxReasonPhraseBytes = 1024

// ReasonPhrase appends a reason phrase per §1.4.4, truncating to
// [MaxReasonPhraseBytes].
//
// Truncating rather than encoding as-is because a Writer method has no way to
// report an error, so the alternative is emitting a frame that every conforming
// peer must treat as a PROTOCOL_VIOLATION — losing the tail of a diagnostic
// string is strictly better than losing the session that was trying to report
// it. The reason phrase is not always ours to bound: REQUEST_ERROR and
// PUBLISH_ERROR carry one built from a token verifier's error text, and a
// third-party [TokenVerifier] can return a string of any length.
//
// The cut lands on a rune boundary, since §1.4.4 specifies UTF-8 and slicing
// mid-rune would produce a phrase the peer decodes as replacement characters.
func (w *Writer) ReasonPhrase(s string) {
        if len(s) > MaxReasonPhraseBytes {
                s = s[:MaxReasonPhraseBytes]
                // Drop a rune the cut split (and any invalid trailing bytes with it).
                for len(s) > 0 {
                        if r, size := utf8.DecodeLastRuneInString(s); r == utf8.RuneError && size <= 1 {
                                s = s[:len(s)-1]
                                continue
                        }
                        break
                }
        }
        w.VarintBytes([]byte(s))
}

// Decoder is the read-side interface shared by Reader (in-memory) and
// StreamReader (streaming io.Reader). Parse methods in the message package
// accept Decoder so they work in both contexts.
type Decoder interface {
        Varint() (uint64, error)
        UInt8() (uint8, error)
        FixedBytes(n int) ([]byte, error)
        VarintBytes() ([]byte, error)
}

// StreamReader wraps an io.Reader and exposes the same Decoder interface as
// Reader. It is intended for parsing self-delimiting wire objects directly
// from a QUIC uni-stream without buffering the entire object first.
type StreamReader struct {
        r io.Reader
        // qr reads one byte at a time (no look-ahead) so a varint read leaves the
        // stream positioned exactly after the varint.
        qr io.ByteReader
}

// NewStreamReader returns a StreamReader over r. r should already be buffered
// (e.g. a *bufio.Reader) for efficiency; StreamReader does not add its own
// buffering layer.
func NewStreamReader(r io.Reader) *StreamReader {
        return &StreamReader{r: r, qr: NewByteReader(r)}
}

// Varint reads a MoQT leading-ones varint (§1.4.1) from the underlying stream.
func (s *StreamReader) Varint() (uint64, error) {
        return ReadVarint(s.qr)
}

// UInt8 reads a single byte.
func (s *StreamReader) UInt8() (uint8, error) {
        b, err := s.qr.ReadByte()
        return b, err
}

// FixedBytes reads exactly n bytes.
func (s *StreamReader) FixedBytes(n int) ([]byte, error) {
        if n == 0 {
                return nil, nil
        }
        // n is derived from a peer-supplied varint; guard the allocation so a
        // bogus length cannot OOM us. n < 0 when a caller's int conversion of a
        // varint >= 2^63 wrapped.
        if n < 0 || n > MaxStreamFieldSize {
                return nil, fmt.Errorf("%w: %d > %d", ErrFieldTooLarge, n, MaxStreamFieldSize)
        }
        buf := make([]byte, n)
        _, err := io.ReadFull(s.r, buf)
        return buf, err
}

// VarintBytes reads a varint length then that many bytes.
func (s *StreamReader) VarintBytes() ([]byte, error) {
        n, err := s.Varint()
        if err != nil {
                return nil, err
        }
        // §1.4.1 varints reach 2^64-1: bound n before the int conversion.
        if n > uint64(MaxStreamFieldSize) { //nolint:gosec // G115: a size cap, never negative.
                return nil, fmt.Errorf("%w: %d > %d", ErrFieldTooLarge, n, MaxStreamFieldSize)
        }
        return s.FixedBytes(int(n)) //nolint:gosec // G115: n <= MaxStreamFieldSize above.
}

// byteReaderAdapter wraps an io.Reader to implement io.ByteReader by reading
// one byte at a time. Used when the underlying reader does not implement
// io.ByteReader directly.
type byteReaderAdapter struct {
        r   io.Reader
        buf [1]byte
}

func (b *byteReaderAdapter) ReadByte() (byte, error) {
        _, err := io.ReadFull(b.r, b.buf[:])
        return b.buf[0], err
}

package relay

import (
        "context"
        "errors"
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
)

// Authorizer is the relay's pluggable authorization hook. Every request
// handler in [pkg/relay] consults the authorizer before performing any state
// mutation; a non-nil return causes the relay to reply REQUEST_ERROR with
// the [DeniedError]'s mapped code (see [DeniedError.RequestErrorCode]).
//
// A REQUEST_UPDATE that changes a SUBSCRIBE_NAMESPACE's or SUBSCRIBE_TRACKS's
// TRACK_NAMESPACE_PREFIX is authorized again (§10.19, §10.20): the method
// receives the subscription as updated: the new prefix, and the
// AUTHORIZATION_TOKENs of the latest request or update that carried any (a
// DELETE authorizes nothing). A denial refuses the update and ends the
// subscription (§10.9.1).
//
// The interface is split per request type for two reasons:
//
//   - It lets a policy reject categories of request without having to
//     type-switch internally — the most common case is "this peer is allowed
//     to subscribe but not publish."
//   - It surfaces the parsed message to the policy, so token-based schemes
//     can inspect things like AUTH_TOKEN parameters or the target Track
//     Namespace without re-parsing the wire form.
//
// Every method receives:
//
//   - ctx: the per-request context. Authorizers may consult it for tracing,
//     cancellation, or token-cache lookups. The relay cancels ctx when the
//     request stream or the session terminates.
//   - sess: the MOQT session the request arrived on. Policies often inspect
//     [session.Session.PeerOptions] (for AUTHORITY / PATH / implementation
//     name) or session-scoped TLS / SETUP-time attestations.
//   - msg: the parsed request message. Authorizers MUST NOT mutate it.
//
// Returning nil grants the request; returning a non-nil error denies it.
// The relay treats any non-nil return as a denial regardless of error type,
// but wrapping with [*DeniedError] (or using [Deny] / [DenyReason]) gives
// the relay an explicit REQUEST_ERROR code to forward to the peer. A plain
// error (e.g. one returned from a downstream token-validation library)
// maps to [moqt.RequestUnauthorized] by default.
type Authorizer interface {
        AuthorizeSubscribe(ctx context.Context, sess *session.Session, msg *message.Subscribe) error
        AuthorizePublish(ctx context.Context, sess *session.Session, msg *message.Publish) error
        AuthorizePublishNamespace(ctx context.Context, sess *session.Session, msg *message.PublishNamespace) error
        AuthorizeFetch(ctx context.Context, sess *session.Session, msg *message.Fetch) error
        AuthorizeSubscribeNamespace(ctx context.Context, sess *session.Session, msg *message.SubscribeNamespace) error
        AuthorizeSubscribeTracks(ctx context.Context, sess *session.Session, msg *message.SubscribeTracks) error
        AuthorizeTrackStatus(ctx context.Context, sess *session.Session, msg *message.TrackStatus) error
}

// DeniedError is returned by an [Authorizer] method to deny a request with
// an explicit MoQT REQUEST_ERROR code. The relay maps Code directly onto the
// REQUEST_ERROR it sends downstream; Reason is forwarded as the human-readable
// reason string.
//
// Code MUST be one of the §10.6 / IANA §15.11.2 REQUEST_ERROR codes (see
// [moqt.RequestErrorCode]). If Code is the zero value, the relay substitutes
// [moqt.RequestUnauthorized] when forming the REQUEST_ERROR.
type DeniedError struct {
        Code   moqt.RequestErrorCode
        Reason string
}

// Error implements the error interface.
func (e *DeniedError) Error() string {
        if e.Reason == "" {
                return fmt.Sprintf("relay: request denied (code %#x)", uint64(e.Code))
        }
        return fmt.Sprintf("relay: request denied (code %#x): %s", uint64(e.Code), e.Reason)
}

// RequestErrorCode returns the REQUEST_ERROR code the relay should use when
// rejecting the request. A zero value collapses to [moqt.RequestUnauthorized]
// — the spec's default rejection code for an authorization failure.
func (e *DeniedError) RequestErrorCode() moqt.RequestErrorCode {
        if e.Code == 0 {
                return moqt.RequestUnauthorized
        }
        return e.Code
}

// Deny is a constructor for [*DeniedError]. Use it when the policy already
// knows which REQUEST_ERROR code to surface:
//
//        return relay.Deny(moqt.RequestUnauthorized, "missing JWT")
func Deny(code moqt.RequestErrorCode, reason string) error {
        return &DeniedError{Code: code, Reason: reason}
}

// DenyReason is a convenience for the common case where the policy only
// wants to attach a human-readable reason and is happy with the default
// REQUEST_ERROR code ([moqt.RequestUnauthorized]).
func DenyReason(reason string) error {
        return &DeniedError{Code: moqt.RequestUnauthorized, Reason: reason}
}

// CodeForAuthorizerError extracts the REQUEST_ERROR code the relay should
// use when rejecting an authorization failure. If err wraps a [*DeniedError],
// its code is returned; otherwise the default [moqt.RequestUnauthorized] is
// returned so a policy that returns a plain error still gets a sensible
// MoQT reply rather than [moqt.RequestInternalError].
func CodeForAuthorizerError(err error) moqt.RequestErrorCode {
        if denied, ok := errors.AsType[*DeniedError](err); ok {
                return denied.RequestErrorCode()
        }
        return moqt.RequestUnauthorized
}

// ReasonForAuthorizerError extracts the human-readable reason for an
// authorization denial. If err wraps a [*DeniedError] and that error has a
// non-empty Reason field, the reason is returned; otherwise the error's
// Error() string is returned. This avoids leaking the internal "relay:
// request denied (code 0x1):" prefix into the wire reply.
func ReasonForAuthorizerError(err error) string {
        if denied, ok := errors.AsType[*DeniedError](err); ok && denied.Reason != "" {
                return denied.Reason
        }
        if err == nil {
                return ""
        }
        return err.Error()
}

// AllowAllAuthorizer is the package's permissive default. Every method
// returns nil. It exists so unit tests, in-process integration tests, and
// pure-relay-of-relays topologies that defer authorization to a downstream
// layer can run without writing a custom policy.
//
// Production deployments SHOULD replace this with a token- or
// session-attestation-aware implementation via [Config.Authorizer]. The relay
// invokes the authorizer once per request, and per prefix update, before any
// state mutation, so the cost of policy evaluation is bounded by the request
// rate rather than the object rate.
type AllowAllAuthorizer struct{}

var _ Authorizer = AllowAllAuthorizer{}

// AuthorizeSubscribe returns nil.
func (AllowAllAuthorizer) AuthorizeSubscribe(context.Context, *session.Session, *message.Subscribe) error {
        return nil
}

// AuthorizePublish returns nil.
func (AllowAllAuthorizer) AuthorizePublish(context.Context, *session.Session, *message.Publish) error {
        return nil
}

// AuthorizePublishNamespace returns nil.
func (AllowAllAuthorizer) AuthorizePublishNamespace(
        context.Context,
        *session.Session,
        *message.PublishNamespace,
) error {
        return nil
}

// AuthorizeFetch returns nil.
func (AllowAllAuthorizer) AuthorizeFetch(context.Context, *session.Session, *message.Fetch) error {
        return nil
}

// AuthorizeSubscribeNamespace returns nil.
func (AllowAllAuthorizer) AuthorizeSubscribeNamespace(
        context.Context,
        *session.Session,
        *message.SubscribeNamespace,
) error {
        return nil
}

// AuthorizeSubscribeTracks returns nil.
func (AllowAllAuthorizer) AuthorizeSubscribeTracks(context.Context, *session.Session, *message.SubscribeTracks) error {
        return nil
}

// AuthorizeTrackStatus returns nil.
func (AllowAllAuthorizer) AuthorizeTrackStatus(context.Context, *session.Session, *message.TrackStatus) error {
        return nil
}

// Package cache holds the relay's per-track Object Cache (§9.1 caching
// relays). Storage is a fixed-capacity circular ring buffer (FIFO) with
// an auxiliary {GroupID, ObjectID} → ring-slot index map for O(1) point
// lookup and overwrite-in-place.
//
// Eviction policy:
//
//   - Size-bounded: when the ring is full, a new Put evicts the oldest
//     entry. Re-Put of an existing key overwrites in place and does NOT
//     evict anything.
//   - Time-bounded: per-entry MaxCacheDuration is applied at read time.
//     Get / GetRange skip entries older than the configured age; the
//     ring itself does no proactive cleanup. This avoids background
//     goroutines and is exactly equivalent for callers because the
//     only consumers of stored data are the FETCH handlers, which
//     can't see anything Get / GetRange filters out.
//
// The §10.2.17 LARGEST_OBJECT watermark is maintained outside the ring,
// under the same mutex, and is monotonic — evictions and TTL expiry
// don't roll it back.
package cache

import (
        "cmp"
        "slices"
        "sync"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
)

// defaultUnboundedCapacity is the ring size used when callers pass
// maxSize <= 0. Production callers (the track registry) always pass a
// positive value; the fallback exists so test helpers that want
// "effectively unbounded" keep working without exporting a separate
// constructor. 1<<16 is large enough that no in-tree test fills it.
const defaultUnboundedCapacity = 1 << 16

// ForwardingPreference distinguishes objects that arrived on a §11.4.2
// SUBGROUP_HEADER stream from objects that arrived as §11.3 datagrams.
// A FETCH response must replay this verbatim because the subscriber's
// decoding depends on it.
type ForwardingPreference uint8

const (
        // ForwardingSubgroup: the object came from a SUBGROUP_HEADER stream.
        ForwardingSubgroup ForwardingPreference = iota
        // ForwardingDatagram: the object arrived as a §11.3 OBJECT_DATAGRAM.
        ForwardingDatagram
)

// CachedObject is the unified record stored by [ObjectCache]. It covers both
// subgroup objects (where SubgroupID + an absolute ObjectID arrive on a
// stream header + per-object delta) and datagrams (which carry their own
// absolute ObjectID and no subgroup notion).
type CachedObject struct {
        GroupID           uint64
        ObjectID          uint64 // absolute, decoded
        SubgroupID        uint64 // 0 for datagrams (no subgroup notion)
        PublisherPriority uint8
        ForwardingPref    ForwardingPreference
        Status            uint64 // §11.2.1.1; 0 for normal objects
        Properties        []byte // retained by reference; opaque to the cache
        Payload           []byte // retained by reference; opaque
        ReceivedAt        time.Time

        // MaxCacheDuration is the MAX_CACHE_DURATION (§12.3) of the upstream the
        // Object arrived through (per upstream, so it can differ within a track),
        // and HasMaxCacheDuration whether it sent one. A present 0 means never
        // serve from the cache (this implementation's reading).
        MaxCacheDuration    time.Duration
        HasMaxCacheDuration bool

        // Stitched marks an Object read from an upstream FETCH for one response
        // rather than stored: ReceivedAt is when it was read, and only a positive
        // MaxCacheDuration bounds it, not the relay's TTL (see [ObjectCache.Expired]).
        Stitched bool

        // EndOfUnknownRange marks this element as a §11.4.4.2 End of Unknown
        // Range (0x10C) FETCH marker rather than a stored object: every Location
        // from the previous element in the response stream (exclusive) through
        // {GroupID, ObjectID} (inclusive) has unknown status. Markers exist only
        // on the FETCH serve path — upstream stitching manufactures them for
        // sub-ranges no source could vouch for — and are never stored in the
        // cache ring.
        EndOfUnknownRange bool

        // EndOfTimedOutRange marks this element as a §11.4.4.2 End of Timed-Out
        // Range (0x20C) FETCH marker: same shape as EndOfUnknownRange, but the
        // Objects it covers were abandoned because FILL_TIMEOUT (§10.2.5) expired
        // rather than because no source could vouch for them. draft-20 split the
        // two so a subscriber can tell "retry might help" from "status unknown".
        EndOfTimedOutRange bool
}

// IsRangeMarker reports whether the element is any §11.4.4.2 end-of-range
// marker rather than a stored object.
func (o *CachedObject) IsRangeMarker() bool {
        return o.EndOfUnknownRange || o.EndOfTimedOutRange
}

// IsStatusMarker reports whether the object is a §11.2.1.1 status marker
// (End of Group / End of Track) rather than a real object. Markers describe
// the absence of objects: the subscription path forwards them, but FETCH
// responses never serialize them — the Object Status field "is absent in
// Objects delivered via a FETCH". Zero-length payloads with Status 0 are
// real (Normal) objects, not markers.
func (o *CachedObject) IsStatusMarker() bool {
        return len(o.Payload) == 0 && o.Status != 0
}

// cacheKey is the composite (GroupID, ObjectID) key used to deduplicate
// re-Puts of the same object onto the same ring slot.
type cacheKey struct {
        Group  uint64
        Object uint64
}

// ObjectCache is a per-track, fixed-capacity, FIFO ring-buffer object
// cache. Each [TrackEntry] holds one.
//
// Concurrency: an RWMutex guards the ring. Writes (Put / Delete) take the
// write lock and are O(1); reads (Get / GetRange / Len)
// take the read lock. FETCH reads are
// O(capacity) (default 1024), so the read lock lets concurrent FETCHes — the
// flash-crowd-of-joining-subscribers case the relay is built for — scan in
// parallel and contend only with the short live-ingest Put rather than
// serialising behind one another. Stored structs are never mutated in place
// (see [ObjectCache.insertLocked]), so readers may dereference the pointers
// they collect without holding the lock.
type ObjectCache struct {
        mu sync.RWMutex

        // ring is a fixed-length slice of pointers; nil means "empty slot".
        // head is the next write position (mod len(ring)).
        ring []*CachedObject
        head int
        size int

        // index maps live keys to their position in ring, for O(1) Get,
        // overwrite-in-place on duplicate Put, and Delete.
        index map[cacheKey]int

        // maxAge is the relay's own read-side TTL. Zero disables it. An Object's
        // MAX_CACHE_DURATION applies on top (see CachedObject).
        maxAge time.Duration
}

// effectiveMaxSize returns a non-zero capacity. Callers that pass 0
// (test helpers that want "effectively unbounded") get
// [defaultUnboundedCapacity].
func effectiveMaxSize(maxSize int) int {
        if maxSize <= 0 {
                return defaultUnboundedCapacity
        }
        return maxSize
}

// NewObjectCache constructs an empty ObjectCache.
//
//   - maxSize: maximum number of stored objects per cache. <= 0 falls
//     back to [defaultUnboundedCapacity].
//   - maxDuration: per-object TTL applied at read time. <= 0 disables
//     time-based filtering; stored objects then live until size-based
//     eviction or explicit Delete.
func NewObjectCache(maxSize int, maxDuration time.Duration) *ObjectCache {
        capacity := effectiveMaxSize(maxSize)
        return &ObjectCache{
                ring:   make([]*CachedObject, capacity),
                index:  make(map[cacheKey]int, capacity),
                maxAge: maxDuration,
        }
}

// Put stores obj in the cache, taking ownership of it: the cache keeps
// obj's pointer (and its Properties / Payload slices) by reference — it
// does NOT copy them. Callers MUST NOT mutate obj or its slices after Put
// returns, since the very same struct is later handed out by Get /
// GetRange. Storing by reference is what keeps the fanout hot path free of
// per-object copies; the one allocation is the CachedObject the caller
// builds.
//
// Put overwrites obj.ReceivedAt with the current time.
//
// If the cache already holds an entry at the same {GroupID, ObjectID},
// Put replaces it (the previous struct is dropped) and does NOT advance
// the ring head — the new entry inherits the existing slot. If the ring is
// full and the key is new, the oldest entry (the one currently at the head
// position) is evicted to make room.
func (c *ObjectCache) Put(obj *CachedObject) {
        if obj == nil {
                return
        }
        // Stamp the arrival time before taking the lock so time.Now() stays out
        // of the write critical section that contends with FETCH range scans.
        obj.ReceivedAt = time.Now()
        c.mu.Lock()
        c.insertLocked(obj)
        c.mu.Unlock()
}

// PutDatagram is a thin adapter that converts a §11.3 OBJECT_DATAGRAM
// into a CachedObject and stores it. Datagrams have no subgroup, so
// SubgroupID is 0; ForwardingPref records the wire shape so a FETCH
// response can replay it as a datagram even if the subscriber's transport
// supports both. maxAge / hasMaxAge are CachedObject.MaxCacheDuration.
func (c *ObjectCache) PutDatagram(d *message.ObjectDatagram, maxAge time.Duration, hasMaxAge bool) {
        if d == nil {
                return
        }
        c.Put(&CachedObject{
                GroupID:           d.GroupID,
                ObjectID:          d.ObjectID,
                SubgroupID:        0,
                PublisherPriority: d.PublisherPriority,
                ForwardingPref:    ForwardingDatagram,
                Status:            d.ObjectStatus,
                Properties:        d.Properties,
                Payload:           d.ObjectPayload,

                MaxCacheDuration:    maxAge,
                HasMaxCacheDuration: hasMaxAge,
        })
}

// insertLocked stores src (by reference) into the ring. On a duplicate key
// it replaces the slot's pointer — the previous struct is dropped, not
// reused; on a new key it consumes the head slot, evicting whatever struct
// occupied it. Structs are never mutated in place or recycled: that is
// exactly what lets Get / GetRange hand out the raw stored pointers without
// a torn-read hazard (an evicted struct is simply orphaned from the ring
// and stays valid for any existing holder). This is the single mutation
// point for ring + index + size; the caller MUST hold c.mu.
func (c *ObjectCache) insertLocked(src *CachedObject) {
        key := cacheKey{Group: src.GroupID, Object: src.ObjectID}

        if idx, ok := c.index[key]; ok {
                c.ring[idx] = src
                return
        }

        // New key: evict whatever currently occupies the head slot (removing it
        // from the index), overwrite the slot with src, and advance head.
        prev := c.ring[c.head]
        if prev != nil {
                delete(c.index, cacheKey{Group: prev.GroupID, Object: prev.ObjectID})
                c.size--
        }
        c.ring[c.head] = src
        c.index[key] = c.head
        c.head = (c.head + 1) % len(c.ring)
        c.size++
}

// notExpiredLocked reports whether obj may still be served: within its own
// MAX_CACHE_DURATION, if it has one, and within the relay's TTL. Caller must
// hold c.mu.
func (c *ObjectCache) notExpiredLocked(obj *CachedObject) bool {
        age := time.Since(obj.ReceivedAt)
        if obj.HasMaxCacheDuration && (obj.MaxCacheDuration <= 0 || age > obj.MaxCacheDuration) {
                return false
        }
        return c.maxAge <= 0 || age <= c.maxAge
}

// Expired reports whether obj, in a FETCH or fill response on this cache's
// track, may no longer be served (§12.3: "MUST NOT start forwarding"): one
// taken from this cache within its own MAX_CACHE_DURATION and the relay's TTL
// (see notExpiredLocked). Range markers never expire. An
// Object stitched from an upstream FETCH expires only past its own
// MAX_CACHE_DURATION ("any individual Object received through this
// subscription or fetch"); a present 0 sets no limit on it, since it is passed
// through rather than served from the cache (interpretation).
func (c *ObjectCache) Expired(obj *CachedObject) bool {
        switch {
        case obj.ReceivedAt.IsZero():
                return false
        case obj.Stitched:
                return obj.MaxCacheDuration > 0 && time.Since(obj.ReceivedAt) > obj.MaxCacheDuration
        }
        c.mu.RLock()
        defer c.mu.RUnlock()
        return !c.notExpiredLocked(obj)
}

// Get returns the stored object at {group, object}, or (nil, false) if
// nothing is recorded there (never written, evicted by size pressure,
// or filtered out by TTL).
//
// The returned *CachedObject is the cache's own stored pointer, NOT a copy
// — callers MUST treat it (and its Properties / Payload slices) as
// read-only. The pointer stays valid indefinitely, even after the entry is
// evicted: [Put] never mutates a stored struct in place, it only replaces a
// ring slot's pointer, so an evicted struct is merely orphaned from the
// ring and remains safe for any existing holder.
//
// Note: this method is O(1) and copy-free. It is not used on the relay hot
// path, but the test suite exercises it heavily.
func (c *ObjectCache) Get(group, object uint64) (*CachedObject, bool) {
        c.mu.RLock()
        defer c.mu.RUnlock()
        idx, ok := c.index[cacheKey{Group: group, Object: object}]
        if !ok {
                return nil, false
        }
        obj := c.ring[idx]
        if obj == nil || !c.notExpiredLocked(obj) {
                return nil, false
        }
        return obj, true
}

// Len returns the number of currently-stored objects (including
// non-existence markers and including TTL-expired entries that have
// not yet been overwritten). The Len is exactly the count of live ring
// slots; with TTL enabled, callers should remember that a non-zero Len
// does not guarantee a Get will return anything.
func (c *ObjectCache) Len() int {
        c.mu.RLock()
        defer c.mu.RUnlock()
        return c.size
}

// ---------------------------------------------------------------------------
// Range scan
// ---------------------------------------------------------------------------

// GetRange returns every stored object whose Location is in [start, end]
// (inclusive on both ends), sorted by (group, object) in the requested
// direction:
//
//   - [message.GroupOrderAscending]:  groups asc, objects asc within group.
//   - [message.GroupOrderDescending]: groups desc, objects asc within group.
//
// Within a group the inner order is always ascending by Object ID, matching
// §10.13: "Within each group, objects are sent in Object ID order".
//
// An empty or inverted range (end < start) returns nil.
//
// The returned slice holds the cache's own stored pointers (no copy);
// callers MUST treat the objects as read-only. They stay valid after
// eviction for the same reason as [ObjectCache.Get]'s result: Put never
// mutates a stored struct in place, it only replaces ring pointers.
//
// Implementation note: GetRange walks the entire ring once and filters
// matches. The ring is small (default 1024 entries) and FETCH is not
// the hot path, so the O(capacity) cost is acceptable. A sorted index
// could be added without changing the signature if profiling later shows
// the scan to dominate.
func (c *ObjectCache) GetRange(start, end message.Location, order message.GroupOrder) []*CachedObject {
        if end.Less(start) {
                return nil
        }
        c.mu.RLock()
        out := make([]*CachedObject, 0)
        for _, obj := range c.ring {
                if obj == nil {
                        continue
                }
                loc := message.Location{Group: obj.GroupID, Object: obj.ObjectID}
                if loc.Less(start) || end.Less(loc) {
                        continue
                }
                if !c.notExpiredLocked(obj) {
                        continue // §12.3: its state "becomes unknown"; the caller treats it so
                }
                // Append the stored pointer directly — Put never recycles or
                // mutates a stored struct, so this never aliases storage a later
                // Put could overwrite.
                out = append(out, obj)
        }
        c.mu.RUnlock()
        sortObjects(out, order)
        return out
}

// sortObjects sorts in-place by (group, object). Group direction is
// controlled by order; objects within a group are always ascending.
// An unknown GroupOrder falls back to ascending.
func sortObjects(objs []*CachedObject, order message.GroupOrder) {
        slices.SortStableFunc(objs, func(a, b *CachedObject) int {
                if a.GroupID != b.GroupID {
                        if order == message.GroupOrderDescending {
                                return cmp.Compare(b.GroupID, a.GroupID)
                        }
                        return cmp.Compare(a.GroupID, b.GroupID)
                }
                return cmp.Compare(a.ObjectID, b.ObjectID)
        })
}

// Delete removes the entry at {group, object} if any. Idempotent: a
// missing entry is a silent no-op.
//
// Note: Delete leaves a tombstone — the ring slot becomes empty but
// the head pointer is not rewound, so the freed capacity is reclaimed
// by the next Put rather than immediately. This keeps FIFO ordering
// stable across mixed Put / Delete sequences.
func (c *ObjectCache) Delete(group, object uint64) {
        key := cacheKey{Group: group, Object: object}
        c.mu.Lock()
        defer c.mu.Unlock()
        idx, ok := c.index[key]
        if !ok {
                return
        }
        c.ring[idx] = nil
        delete(c.index, key)
        c.size--
}

package relay

import (
        "bytes"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/track"
)

// CacheTTLPolicy is a per-track override for [Config.MaxCacheDuration]. The
// relay invokes it once per track at TrackEntry creation time (never on the
// fanout hot path) to decide that track's object-cache retention.
//
// Semantics:
//
//   - Return a positive duration to use that TTL for the matching track.
//   - Return [CacheTTLInfinite] to disable time-based eviction entirely for
//     the matching track (the FIFO size cap from [Config.MaxCacheSize] still
//     applies).
//   - Return 0 (the zero value) to fall through to [Config.MaxCacheDuration].
//
// The policy MUST be safe for concurrent invocation, MUST NOT block, and
// SHOULD be free of side effects — it runs inside the registry's write lock.
// Implementations are typically small predicates on Name (e.g. for an MSF
// catalog track). The relay deliberately exposes only a function-shaped hook
// rather than coupling to any specific Track-Name vocabulary: the binary that
// builds the policy (cmd/relay, an embedded app, …) owns the protocol-specific
// rules.
type CacheTTLPolicy func(name track.FullTrackName) time.Duration

// CacheTTLInfinite is the sentinel a [CacheTTLPolicy] returns to request
// "no time-based eviction" for the matching track. It exists so policy authors
// can express "retain indefinitely" without knowing the object cache's
// internal "non-positive means TTL disabled" convention, and so a return value
// of 0 keeps its natural meaning ("use the default").
const CacheTTLInfinite = time.Duration(-1)

// TrackNameTTL returns a [CacheTTLPolicy] giving every track whose Name equals
// name the retention ttl, and leaving every other track on
// [Config.MaxCacheDuration]. A ttl of 0 is read as "retain indefinitely" and
// becomes [CacheTTLInfinite]; any positive ttl is honoured verbatim. An empty
// name returns nil, which disables the override entirely.
//
// Matching is namespace-agnostic: every publisher's track of that Name gets the
// same retention. That fits the MSF per-broadcaster catalog model, where each
// participant owns a namespace but they all share one catalog Name.
//
// An MSF relay needs it for catalogs: published once on join, they would
// otherwise expire under the default retention and a late joiner's fill fetch
// stream would find nothing.
//
// The choice of which Name and how long still belongs to the binary; only the
// shape of the predicate is shared.
func TrackNameTTL(name string, ttl time.Duration) CacheTTLPolicy {
        if name == "" {
                return nil
        }
        want := []byte(name)
        override := ttl
        if override == 0 {
                override = CacheTTLInfinite
        }
        return func(n track.FullTrackName) time.Duration {
                if bytes.Equal(n.Name, want) {
                        return override
                }
                return 0 // fall through to Config.MaxCacheDuration
        }
}

// Package discovery is the relay's cross-instance track + namespace
// advertisement abstraction.
//
// The interface answers two questions: "which relay instance hosts a
// publisher for this track?" and "which relay instance serves this
// namespace prefix?" — both essential for routing in a multi-relay
// deployment.
//
// The default implementation is [MemoryStore], which keeps state local
// to a single relay process. Watch channels only see events emitted by
// the same MemoryStore, so a single relay with MemoryStore behaves
// identically to a relay with no discovery at all. Production
// deployments swap in a distributed backend (NATS JetStream KV, Redis,
// etc.) behind the same interface; the relay code does not change.
//
// The relay's [TrackRegistry] and [NamespaceRegistry] use this
// abstraction so multi-instance support is a backend swap rather than
// a rewrite.
package discovery

import (
        "context"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// Op is the kind of a discovery event. Publish announces availability;
// Unpublish announces removal. A single backend may emit either kind on
// the same key over the entry's lifetime. SnapshotDone carries no Info: it
// ends a watch's initial snapshot (see [DiscoveryStore.WatchTracks]).
type Op int

const (
        // OpPublish — a track or namespace became available on a relay.
        OpPublish Op = iota
        // OpUnpublish — a track or namespace is no longer available.
        OpUnpublish
        // OpSnapshotDone — the watch has delivered its whole initial snapshot;
        // every event after it is a live change.
        OpSnapshotDone
)

// String returns "publish", "unpublish" or "snapshot-done".
func (o Op) String() string {
        switch o {
        case OpPublish:
                return "publish"
        case OpUnpublish:
                return "unpublish"
        case OpSnapshotDone:
                return "snapshot-done"
        }
        return "unknown"
}

// TrackInfo describes a track available on a relay instance.
type TrackInfo struct {
        // Key uniquely identifies the track per §2.4.1. Used as the map
        // index by all backends.
        Key track.Key

        // FullName retains the unhashed {namespace, name} tuple so
        // downstream subscribers can echo it on the wire and humans can
        // read it in logs.
        FullName track.FullTrackName

        // Properties is the opaque Track Properties blob the upstream
        // publisher attached (see §9.6). Stored by reference; callers
        // MUST NOT mutate after handing it to the store.
        Properties []byte

        // RelayAddr identifies the relay instance hosting this track. For
        // MemoryStore it is whatever the local relay registered itself as
        // (typically empty in single-relay deployments). NATS/Redis
        // backends use it to route upstream connections to the right peer.
        RelayAddr string

        // PublishedAt records when the entry was last written. Backends
        // MAY use this for TTL eviction.
        PublishedAt time.Time
}

// NamespaceInfo describes a namespace prefix available on a relay
// instance. Multiple TrackInfos share a NamespaceInfo iff their full
// names start with the same Prefix.
type NamespaceInfo struct {
        // Prefix is the namespace tuple advertised by PUBLISH_NAMESPACE
        // (§6.2 / §10.16). A zero-length tuple matches every track — used
        // by SUBSCRIBE_NAMESPACE with no filter.
        Prefix wire.TrackNamespace

        // RelayAddr — see [TrackInfo.RelayAddr].
        RelayAddr string

        // PublishedAt — see [TrackInfo.PublishedAt].
        PublishedAt time.Time
}

// TrackEvent is what [DiscoveryStore.WatchTracks] yields. Op tells
// callers whether the entry is being added or removed.
type TrackEvent struct {
        Op   Op
        Info TrackInfo
}

// NamespaceEvent is what [DiscoveryStore.WatchNamespaces] yields.
type NamespaceEvent struct {
        Op   Op
        Info NamespaceInfo
}

// DiscoveryStore is the relay's cross-instance metadata fabric.
//
// All methods are safe for concurrent use. Backends SHOULD treat
// repeated Publish for the same (Key|Prefix, RelayAddr) tuple as
// idempotent updates rather than duplicates — multiple sessions on
// the same relay can independently advertise the same track or
// namespace and the store should collapse them.
//
// Find operations return a snapshot; callers must not rely on
// subsequent reads observing the same set. Watch is the right
// primitive for "tell me when this changes".
//
// Implementations MUST honor ctx cancellation and deadlines on every
// call: the relay's registries invoke Publish/Unpublish while holding
// their internal locks (keeping store order consistent with registry
// state), so a backend that ignores ctx stalls the whole registry.
//
// Close releases backend resources (network connections, goroutines).
// Watch channels MUST be drained or their owning context cancelled
// before Close to avoid backend-side blocking. After Close all methods
// return [ErrClosed].
type DiscoveryStore interface {
        // PublishTrack advertises a track. RelayAddr / Properties /
        // PublishedAt come from info; the backend uses Key as the unique
        // index.
        PublishTrack(ctx context.Context, info TrackInfo) error

        // UnpublishTrack removes the track advertisement keyed by
        // (key, relayAddr). Unknown entries are silent no-ops.
        UnpublishTrack(ctx context.Context, key track.Key, relayAddr string) error

        // FindTrack returns every advertisement of this track across all
        // relay instances. A zero-length slice with no error means
        // "nobody hosts this track right now."
        FindTrack(ctx context.Context, key track.Key) ([]TrackInfo, error)

        // PublishNamespace advertises a namespace prefix.
        PublishNamespace(ctx context.Context, info NamespaceInfo) error

        // UnpublishNamespace removes the namespace advertisement keyed by
        // (prefix, relayAddr).
        UnpublishNamespace(ctx context.Context, prefix wire.TrackNamespace, relayAddr string) error

        // FindNamespace returns every namespace advertisement whose
        // Prefix is a prefix of namespace (in the §9.5 / wire.TrackNamespace
        // HasPrefix sense). A query for ["a","b","c"] matches advertised
        // prefixes ["a"], ["a","b"], and ["a","b","c"]; advertised
        // prefix ["a","b","c","d"] does NOT match. This is the ancestor
        // direction: "which relays serve a covering prefix for this track?"
        FindNamespace(ctx context.Context, namespace wire.TrackNamespace) ([]NamespaceInfo, error)

        // FindNamespacesUnder is the descendant complement of FindNamespace:
        // it returns every advertisement whose Prefix extends (is at or below)
        // prefix. A query for ["a"] matches advertised prefixes ["a"], ["a","b"],
        // and ["a","b","c"]; ["x"] does NOT match. A zero-length prefix matches
        // every advertisement. The relay itself seeds from WatchNamespaces instead.
        FindNamespacesUnder(ctx context.Context, prefix wire.TrackNamespace) ([]NamespaceInfo, error)

        // WatchTracks returns a channel that first delivers the current set of
        // track advertisements as OpPublish events (the snapshot) followed by one
        // OpSnapshotDone, then streams every subsequent Publish / Unpublish the
        // backend observes (local + remote), until ctx is cancelled or the store
        // is closed. The channel is closed when the watch ends.
        //
        // The snapshot→follow handoff is gapless: no event is missed or
        // duplicated across it, so no separate Find is needed.
        //
        // The snapshot is delivered in full. After it, a slow consumer must not
        // block other watchers, and a backend that cannot deliver a live event
        // MUST end that watch (close its channel) rather than drop the event, so
        // the consumer notices and re-watches.
        WatchTracks(ctx context.Context) (<-chan TrackEvent, error)

        // WatchNamespaces streams namespace events. Same snapshot-then-follow
        // contract as WatchTracks.
        WatchNamespaces(ctx context.Context) (<-chan NamespaceEvent, error)

        // Withdraw removes every advertisement published for relayAddr, so peers
        // stop resolving it as an upstream while it drains (§3.6). Peers observe
        // the removals as OpUnpublish events.
        //
        // Withdraw is terminal for that address's advertising side: afterwards
        // PublishTrack / PublishNamespace for relayAddr MUST return [ErrWithdrawn]
        // without restoring anything. Find / Watch stay usable, and Unpublish
        // calls for relayAddr become no-ops.
        //
        // Withdrawing an address that advertised nothing, or withdrawing twice, is a
        // silent no-op. Unlike Close, Withdraw releases no backend resources.
        Withdraw(ctx context.Context, relayAddr string) error

        // Close releases backend resources.
        Close() error
}

package discovery

import (
        "context"
        "errors"
        "log/slog"
        "slices"
        "sync"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// nowFunc is the time source used by [MemoryStore.PublishTrack] /
// [MemoryStore.PublishNamespace] to stamp PublishedAt when the caller
// leaves it zero. Overridable from tests if deterministic timestamps
// ever become useful.
var nowFunc = time.Now

// ErrClosed is returned by [DiscoveryStore] methods after Close has run.
var ErrClosed = errors.New("discovery: store closed")

// ErrWithdrawn is returned by the Publish calls for a relay address that has
// been withdrawn (see [DiscoveryStore.Withdraw]). It is not a failure: the
// relay is shutting down, and re-advertising it would undo the withdrawal.
var ErrWithdrawn = errors.New("discovery: relay withdrawn")

// defaultWatchBufferSize is how many live events a watcher may fall behind
// before its watch is ended (see [DiscoveryStore.WatchTracks]).
const defaultWatchBufferSize = 32

// MemoryStore is the in-process [DiscoveryStore] for single-relay
// deployments. All state is local; Watch channels only see events the
// MemoryStore itself emitted, so a single relay using MemoryStore
// behaves identically to one with no discovery at all. Distributed
// backends (NATS / Redis) replace this without touching relay internals.
//
// The store is safe for concurrent use. Watch delivery is non-blocking, so
// a slow consumer cannot stall a publisher.
type MemoryStore struct {
        mu         sync.RWMutex
        tracks     map[trackEntryKey]TrackInfo
        namespaces map[namespaceEntryKey]NamespaceInfo
        trackWatch []*watcher[TrackEvent]
        nsWatch    []*watcher[NamespaceEvent]
        closed     bool
        // withdrawn records relay addresses that called Withdraw, so a late
        // Publish cannot re-advertise a relay that is draining.
        withdrawn  map[string]struct{}
        log        *slog.Logger
        bufferSize int
}

// trackEntryKey indexes a TrackInfo by (key, relayAddr): the same track
// hosted on different relays produces distinct entries.
type trackEntryKey struct {
        key  track.Key
        addr string
}

// namespaceEntryKey indexes a NamespaceInfo. The prefix is stored as
// its wire-encoded byte string (canonical key for nested tuples — see
// [track.Key.namespace] for the same trick).
type namespaceEntryKey struct {
        prefix string
        addr   string
}

// NewMemoryStore constructs an empty in-memory store. It logs a warning
// when a slow watcher's watch is ended, to [slog.Default] unless a logger
// is set.
func NewMemoryStore(opts ...MemoryStoreOption) *MemoryStore {
        s := &MemoryStore{
                tracks:     make(map[trackEntryKey]TrackInfo),
                namespaces: make(map[namespaceEntryKey]NamespaceInfo),
                withdrawn:  make(map[string]struct{}),
                bufferSize: defaultWatchBufferSize,
        }
        for _, opt := range opts {
                opt(s)
        }
        if s.log == nil {
                s.log = slog.Default()
        }
        return s
}

// MemoryStoreOption tweaks a [MemoryStore] at construction time.
type MemoryStoreOption func(*MemoryStore)

// WithWatchBufferSize overrides the per-watcher channel capacity.
// Values <= 0 fall back to the package default.
func WithWatchBufferSize(n int) MemoryStoreOption {
        return func(s *MemoryStore) {
                if n > 0 {
                        s.bufferSize = n
                }
        }
}

var _ DiscoveryStore = (*MemoryStore)(nil)

// PublishTrack stores info; an existing entry with the same
// (Key, RelayAddr) is replaced atomically. The store ignores PublishedAt
// if zero (caller-friendly default).
func (s *MemoryStore) PublishTrack(_ context.Context, info TrackInfo) error {
        s.mu.Lock()
        if s.closed {
                s.mu.Unlock()
                return ErrClosed
        }
        if _, ok := s.withdrawn[info.RelayAddr]; ok {
                s.mu.Unlock()
                return ErrWithdrawn
        }
        if info.PublishedAt.IsZero() {
                info.PublishedAt = nowFunc()
        }
        s.tracks[trackEntryKey{key: info.Key, addr: info.RelayAddr}] = info
        // Send under the lock so it cannot race a watcher's close; log after
        // unlocking so a slow logger cannot stall the store.
        dropped := fanout(&s.trackWatch, TrackEvent{Op: OpPublish, Info: info})
        s.mu.Unlock()
        s.warnDropped(dropped, OpPublish, "key", info.Key)
        return nil
}

// UnpublishTrack removes the (key, relayAddr) entry. Missing entries
// are no-ops; no event is emitted in that case.
func (s *MemoryStore) UnpublishTrack(_ context.Context, key track.Key, relayAddr string) error {
        s.mu.Lock()
        if s.closed {
                s.mu.Unlock()
                return ErrClosed
        }
        idx := trackEntryKey{key: key, addr: relayAddr}
        info, ok := s.tracks[idx]
        if !ok {
                s.mu.Unlock()
                return nil
        }
        delete(s.tracks, idx)
        // Send under the lock, log after — see [MemoryStore.PublishTrack].
        dropped := fanout(&s.trackWatch, TrackEvent{Op: OpUnpublish, Info: info})
        s.mu.Unlock()
        s.warnDropped(dropped, OpUnpublish, "key", key)
        return nil
}

// FindTrack returns every advertisement of key across all RelayAddrs.
func (s *MemoryStore) FindTrack(_ context.Context, key track.Key) ([]TrackInfo, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()
        if s.closed {
                return nil, ErrClosed
        }
        var out []TrackInfo
        for k, v := range s.tracks {
                if k.key == key {
                        out = append(out, v)
                }
        }
        return out, nil
}

// PublishNamespace stores info; identical (Prefix, RelayAddr) replaces.
func (s *MemoryStore) PublishNamespace(_ context.Context, info NamespaceInfo) error {
        s.mu.Lock()
        if s.closed {
                s.mu.Unlock()
                return ErrClosed
        }
        if _, ok := s.withdrawn[info.RelayAddr]; ok {
                s.mu.Unlock()
                return ErrWithdrawn
        }
        if info.PublishedAt.IsZero() {
                info.PublishedAt = nowFunc()
        }
        s.namespaces[namespaceEntryKey{prefix: namespaceWireKey(info.Prefix), addr: info.RelayAddr}] = info
        // Send under the lock, log after — see [MemoryStore.PublishTrack].
        dropped := fanout(&s.nsWatch, NamespaceEvent{Op: OpPublish, Info: info})
        s.mu.Unlock()
        s.warnDropped(dropped, OpPublish, "prefix", info.Prefix)
        return nil
}

// UnpublishNamespace removes the (prefix, relayAddr) entry.
func (s *MemoryStore) UnpublishNamespace(_ context.Context, prefix wire.TrackNamespace, relayAddr string) error {
        s.mu.Lock()
        if s.closed {
                s.mu.Unlock()
                return ErrClosed
        }
        idx := namespaceEntryKey{prefix: namespaceWireKey(prefix), addr: relayAddr}
        info, ok := s.namespaces[idx]
        if !ok {
                s.mu.Unlock()
                return nil
        }
        delete(s.namespaces, idx)
        // Send under the lock, log after — see [MemoryStore.PublishTrack].
        dropped := fanout(&s.nsWatch, NamespaceEvent{Op: OpUnpublish, Info: info})
        s.mu.Unlock()
        s.warnDropped(dropped, OpUnpublish, "prefix", prefix)
        return nil
}

// FindNamespace returns every advertisement whose Prefix is a non-strict
// prefix of namespace. A query for ["a","b","c"] matches stored prefixes
// ["a"], ["a","b"], ["a","b","c"]; ["a","b","c","d"] does NOT match.
func (s *MemoryStore) FindNamespace(_ context.Context, namespace wire.TrackNamespace) ([]NamespaceInfo, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()
        if s.closed {
                return nil, ErrClosed
        }
        var out []NamespaceInfo
        for _, v := range s.namespaces {
                if namespace.HasPrefix(v.Prefix) {
                        out = append(out, v)
                }
        }
        return out, nil
}

// FindNamespacesUnder returns every advertisement whose Prefix extends prefix
// (the descendant direction — see [DiscoveryStore.FindNamespacesUnder]).
func (s *MemoryStore) FindNamespacesUnder(_ context.Context, prefix wire.TrackNamespace) ([]NamespaceInfo, error) {
        s.mu.RLock()
        defer s.mu.RUnlock()
        if s.closed {
                return nil, ErrClosed
        }
        var out []NamespaceInfo
        for _, v := range s.namespaces {
                if v.Prefix.HasPrefix(prefix) {
                        out = append(out, v)
                }
        }
        return out, nil
}

// WatchTracks implements [DiscoveryStore.WatchTracks]. Snapshotting and
// registering happen under one lock, which makes the handoff gapless; the
// channel holds the whole snapshot plus the live headroom.
func (s *MemoryStore) WatchTracks(ctx context.Context) (<-chan TrackEvent, error) {
        s.mu.Lock()
        if s.closed {
                s.mu.Unlock()
                return nil, ErrClosed
        }
        w := newWatcher[TrackEvent](len(s.tracks) + 1 + s.bufferSize)
        for _, v := range s.tracks {
                w.ch <- TrackEvent{Op: OpPublish, Info: v} // fits: capacity includes the snapshot
        }
        w.ch <- TrackEvent{Op: OpSnapshotDone}
        s.trackWatch = append(s.trackWatch, w)
        s.mu.Unlock()

        go watchLifecycle(ctx, &s.mu, &s.trackWatch, w)
        return w.ch, nil
}

// WatchNamespaces — see [MemoryStore.WatchTracks].
func (s *MemoryStore) WatchNamespaces(ctx context.Context) (<-chan NamespaceEvent, error) {
        s.mu.Lock()
        if s.closed {
                s.mu.Unlock()
                return nil, ErrClosed
        }
        w := newWatcher[NamespaceEvent](len(s.namespaces) + 1 + s.bufferSize)
        for _, v := range s.namespaces {
                w.ch <- NamespaceEvent{Op: OpPublish, Info: v} // fits: capacity includes the snapshot
        }
        w.ch <- NamespaceEvent{Op: OpSnapshotDone}
        s.nsWatch = append(s.nsWatch, w)
        s.mu.Unlock()

        go watchLifecycle(ctx, &s.mu, &s.nsWatch, w)
        return w.ch, nil
}

// Withdraw implements [DiscoveryStore.Withdraw].
func (s *MemoryStore) Withdraw(_ context.Context, relayAddr string) error {
        s.mu.Lock()
        if s.closed {
                s.mu.Unlock()
                return ErrClosed
        }
        s.withdrawn[relayAddr] = struct{}{}
        // Events go out under the lock — see [MemoryStore.PublishTrack].
        dropped := 0
        for idx, info := range s.tracks {
                if idx.addr != relayAddr {
                        continue
                }
                delete(s.tracks, idx)
                dropped += fanout(&s.trackWatch, TrackEvent{Op: OpUnpublish, Info: info})
        }
        for idx, info := range s.namespaces {
                if idx.addr != relayAddr {
                        continue
                }
                delete(s.namespaces, idx)
                dropped += fanout(&s.nsWatch, NamespaceEvent{Op: OpUnpublish, Info: info})
        }
        s.mu.Unlock()
        s.warnDropped(dropped, OpUnpublish, "relay_addr", relayAddr)
        return nil
}

// Close closes every active watch channel and rejects further
// operations with [ErrClosed].
func (s *MemoryStore) Close() error {
        s.mu.Lock()
        defer s.mu.Unlock()
        if s.closed {
                return nil
        }
        s.closed = true
        // End under the lock so the closes cannot race a concurrent fanout send
        // (which also holds s.mu); ending also releases each lifecycle goroutine.
        for _, w := range s.trackWatch {
                w.end()
        }
        for _, w := range s.nsWatch {
                w.end()
        }
        s.trackWatch = nil
        s.nsWatch = nil
        return nil
}

// watcher is one watch. done closes with ch, releasing the lifecycle
// goroutine however the watch ends: ctx, Close, or overflow.
type watcher[T any] struct {
        ch   chan T
        done chan struct{}
}

func newWatcher[T any](buffer int) *watcher[T] {
        return &watcher[T]{ch: make(chan T, buffer), done: make(chan struct{})}
}

// end closes the watch. The caller holds the store's lock and has removed w
// from its list, so it is ended exactly once.
func (w *watcher[T]) end() {
        close(w.ch)
        close(w.done)
}

// watchLifecycle ends w when ctx is cancelled, unless it ended first.
func watchLifecycle[T any](ctx context.Context, mu sync.Locker, watchers *[]*watcher[T], w *watcher[T]) {
        select {
        case <-w.done:
                return // overflow or Close ended it
        case <-ctx.Done():
        }
        mu.Lock()
        defer mu.Unlock()
        i := slices.Index(*watchers, w)
        if i < 0 {
                return // ended while this goroutine waited for the lock
        }
        *watchers = slices.Delete(*watchers, i, i+1)
        w.end()
}

// fanout delivers ev to each watcher with a non-blocking send, ending (not
// skipping) any whose buffer is full, and returns how many it ended. It MUST
// be called with s.mu held, so its sends and closes exclude the other ends.
func fanout[T any](watchers *[]*watcher[T], ev T) int {
        ended := 0
        *watchers = slices.DeleteFunc(*watchers, func(w *watcher[T]) bool {
                select {
                case w.ch <- ev:
                        return false
                default:
                        w.end()
                        ended++
                        return true
                }
        })
        return ended
}

// warnDropped logs that n slow watchers had their watch ended, if any. Call
// it after releasing s.mu.
func (s *MemoryStore) warnDropped(n int, op Op, keyAttr string, keyVal any) {
        if n == 0 {
                return
        }
        s.log.Warn("discovery: ended the watch of slow watcher(s)",
                "op", op.String(), keyAttr, keyVal, "watchers", n)
}

// namespaceWireKey serialises a TrackNamespace into a canonical byte
// string suitable for use as a map key. The same trick is used by
// track.Key so callers don't have to worry about field-count vs.
// concatenated-bytes collisions.
func namespaceWireKey(ns wire.TrackNamespace) string {
        w := wire.NewWriter(nil)
        w.TrackNamespace(ns)
        return string(w.Bytes())
}

package relay

import (
        "cmp"
        "math"
        "slices"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/relay/cache"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// A FETCH response asserts what it does not carry: "Any gaps in the Group and
// Object IDs in the response stream indicate objects that do not exist"
// (§10.13), and "All signals that an Object does not exist are
// authoritative" (§2.1). The relay knows an Object does not exist only from a
// signal — a Prior Group or Object ID Gap (§12.8, §12.9), a Group's or the
// Track's end (§11.2.1.1, §11.4.2), or an upstream's FETCH response — never from a gap in
// what it happened to receive: "A gap in the observed Object IDs does not by
// itself convey any information about the skipped Objects" (§2.1). So a
// response is built from the requested range classified into Objects, known
// absent Locations, and the rest, whose status is unknown and is either asked
// of an upstream (§10.13) or marked with an End of Unknown Range (§11.4.4.2).
//
// Interpretation: an End of Range marker covers "Locations between the last
// serialized Object, if any, and this Location" (§11.4.4.2) in the order the
// response carries them: Groups in its Group Order, Object IDs ascending
// within a Group. In Descending order that is not Location order.

// locSucc returns the Location right after l, and false when l is the last.
func locSucc(l message.Location) (message.Location, bool) {
        switch {
        case l.Object < math.MaxUint64:
                return message.Location{Group: l.Group, Object: l.Object + 1}, true
        case l.Group < math.MaxUint64:
                return message.Location{Group: l.Group + 1}, true
        }
        return message.Location{}, false
}

// uncovered returns the parts of [start, end] no range in known covers, in
// ascending order. It sorts known.
func uncovered(start, end message.Location, known []registry.LocRange) []registry.LocRange {
        slices.SortFunc(known, func(a, b registry.LocRange) int { return a.Lo.Compare(b.Lo) })
        var out []registry.LocRange
        cur := start
        for _, k := range known {
                if k.Hi.Less(cur) {
                        continue
                }
                if end.Less(k.Lo) {
                        break
                }
                if cur.Less(k.Lo) {
                        pred, _ := fetchPredecessor(k.Lo) // k.Lo > cur, so it has one
                        out = append(out, registry.LocRange{Lo: cur, Hi: pred})
                }
                next, ok := locSucc(k.Hi)
                if !ok || end.Less(next) {
                        return out
                }
                cur = next
        }
        return append(out, registry.LocRange{Lo: cur, Hi: end})
}

// knownFromCache is what objs, Objects the cache holds, establish: each
// Object's Location; from an END_OF_GROUP or END_OF_TRACK status, the rest of
// its Group or of the Track (§11.2.1.1); and the gaps its Prior Group and
// Object ID Gaps announce (§12.8, §12.9). [message.CheckObjectProperties]
// rejected a gap above an Object's ID on receipt.
func knownFromCache(objs []*cache.CachedObject) []registry.LocRange {
        out := make([]registry.LocRange, 0, len(objs))
        for _, o := range objs {
                loc := message.Location{Group: o.GroupID, Object: o.ObjectID}
                r := registry.LocRange{Lo: loc, Hi: loc}
                switch o.Status {
                case message.ObjectStatusEndOfGroup:
                        r.Hi.Object = math.MaxUint64
                case message.ObjectStatusEndOfTrack:
                        r.Hi = message.Location{Group: math.MaxUint64, Object: math.MaxUint64}
                }
                out = append(out, r)
                gaps := message.ObjectPriorGaps(o.Properties)
                if gaps.HasObject && gaps.Object > 0 {
                        out = append(out, registry.LocRange{
                                Lo: message.Location{Group: o.GroupID, Object: o.ObjectID - gaps.Object},
                                Hi: message.Location{Group: o.GroupID, Object: o.ObjectID - 1},
                        })
                }
                if gaps.HasGroup && gaps.Group > 0 {
                        out = append(out, registry.LocRange{
                                Lo: message.Location{Group: o.GroupID - gaps.Group},
                                Hi: message.Location{Group: o.GroupID - 1, Object: math.MaxUint64},
                        })
                }
        }
        return out
}

// unknownIn returns the Locations of [start, end] whose status the relay does
// not know: neither among objs, the cached Objects of the range, nor known
// absent from them or entry's ledger.
func unknownIn(
        entry *registry.TrackEntry,
        objs []*cache.CachedObject,
        start, end message.Location,
) []registry.LocRange {
        return uncovered(start, end, append(knownFromCache(objs), entry.KnownAbsent()...))
}

// streamCompare orders Locations as a FETCH response in order carries them
// (§10.13): Groups in the Group Order, Object IDs ascending within a Group.
func streamCompare(a, b message.Location, order message.GroupOrder) int {
        if order == message.GroupOrderDescending && a.Group != b.Group {
                return cmp.Compare(b.Group, a.Group)
        }
        return a.Compare(b)
}

// runEnds returns where each part of r a response in order carries without
// interruption ends, in stream order. Ascending, r is one run. Descending, a
// range across Groups is up to three: the start of its highest Group, the
// Groups strictly between, and the rest of its lowest Group.
func runEnds(r registry.LocRange, order message.GroupOrder) []message.Location {
        if order != message.GroupOrderDescending || r.Lo.Group == r.Hi.Group {
                return []message.Location{r.Hi}
        }
        ends := []message.Location{r.Hi}
        if r.Hi.Group-r.Lo.Group > 1 {
                ends = append(ends, message.Location{Group: r.Lo.Group + 1, Object: math.MaxUint64})
        }
        return append(ends, message.Location{Group: r.Lo.Group, Object: math.MaxUint64})
}

// fetchElements orders objs for a response in order and adds, at the end of
// each run of unknown and timedOut Locations, an End of Unknown or Timed-Out
// Range marker (§11.4.4.2), so that a plain gap in the response only ever
// covers Locations known not to exist. unknown and timedOut must be disjoint
// and hold no Object of objs. A marker's run starts after the previous
// element, so it may also cover known-absent Locations: weaker, never false.
func fetchElements(
        objs []*cache.CachedObject,
        unknown, timedOut []registry.LocRange,
        order message.GroupOrder,
) []*cache.CachedObject {
        elems := slices.Clone(objs)
        for _, r := range unknown {
                for _, at := range runEnds(r, order) {
                        elems = append(elems, unknownRangeMarker(at))
                }
        }
        for _, r := range timedOut {
                for _, at := range runEnds(r, order) {
                        elems = append(elems, timedOutRangeMarker(at))
                }
        }
        slices.SortStableFunc(elems, func(a, b *cache.CachedObject) int {
                return streamCompare(
                        message.Location{Group: a.GroupID, Object: a.ObjectID},
                        message.Location{Group: b.GroupID, Object: b.ObjectID},
                        order)
        })
        // A marker's coverage starts after the previous element, so a marker
        // right after one of its own kind covers both: keep only the later.
        out := elems[:0]
        for i, e := range elems {
                if i+1 < len(elems) && e.IsRangeMarker() && elems[i+1].IsRangeMarker() &&
                        e.EndOfTimedOutRange == elems[i+1].EndOfTimedOutRange {
                        continue
                }
                out = append(out, e)
        }
        return out
}

package relay

import (
        "context"
        "errors"
        "log/slog"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/relay/cache"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// runDatagramLoop forwards each received [message.ObjectDatagram] to the
// downstream subscribers, until a session-level receive error, which it
// returns. Send failures and lookup misses drop the datagram (§11.3); a
// malformed track is ended (§2.4.2) and the loop reads on.
func (h *sessionHandler) runDatagramLoop(ctx context.Context) error {
        for {
                d, err := h.sess.ReceiveDatagram(ctx)
                if errors.Is(err, session.ErrMalformedTrack) {
                        // Per-track, not per-session: end that track and read on.
                        if in, ok := h.sess.LookupInboundTrack(d.TrackAlias); ok {
                                if entry, ok := h.tracks.Get(in.Key); ok {
                                        h.endMalformedTrack(ctx, entry, h.sess, err)
                                }
                        }
                        continue
                }
                if err != nil {
                        if errors.Is(err, context.Canceled) {
                                return ctx.Err()
                        }
                        return err
                }
                h.handleDatagram(ctx, d)
        }
}

// handleDatagram is the per-datagram counterpart of [runFanout]'s
// per-object body.
func (h *sessionHandler) handleDatagram(ctx context.Context, d *message.ObjectDatagram) {
        in, ok := h.sess.LookupInboundTrack(d.TrackAlias)
        if !ok {
                // §11.3: an unknown Track Alias MAY be dropped.
                h.log.LogAttrs(ctx, slog.LevelDebug, "datagram: unknown inbound Track Alias",
                        slog.Uint64("alias", d.TrackAlias))
                return
        }

        entry, ok := h.tracks.Get(in.Key)
        if !ok {
                h.log.LogAttrs(ctx, slog.LevelDebug, "datagram: track entry gone",
                        slog.Uint64("alias", d.TrackAlias))
                return
        }

        // §11.3.1: resolve the inherited DEFAULT_PUBLISHER_PRIORITY (§12.4) for
        // the cache and PRIORITY_FILTER; the forwarded Type still omits the byte.
        if d.HasDefaultPriority() {
                d.PublisherPriority = in.DefaultPublisherPriority
        }

        // §9.3: the first copy of {GroupID, ObjectID} wins, unless an announced
        // gap says it does not exist (§2.1, §9.1).
        info := registry.ObjectInfo{
                Group:      d.GroupID,
                Object:     d.ObjectID,
                Datagram:   true,
                Priority:   d.PublisherPriority,
                Status:     d.ObjectStatus,
                EndOfGroup: d.HasEndOfGroup(),
                Gaps:       message.ObjectPriorGaps(d.Properties),
        }
        claim, err := entry.ClaimDelivered(info)
        if err != nil {
                h.endMalformedTrack(ctx, entry, h.sess, err)
                return
        }
        // An aged-out datagram is dropped like a redundant one: it has no stream
        // to reset (§9.4 covers multi-object streams).
        if claim != registry.ClaimFresh {
                if err := recordRedundant(entry, info, &cache.CachedObject{
                        GroupID:           d.GroupID,
                        ObjectID:          d.ObjectID,
                        PublisherPriority: d.PublisherPriority,
                        ForwardingPref:    cache.ForwardingDatagram,
                        Status:            d.ObjectStatus,
                        Properties:        d.Properties,
                        Payload:           d.ObjectPayload,
                }); err != nil {
                        h.endMalformedTrack(ctx, entry, h.sess, err)
                }
                return
        }

        // §10.2.17
        entry.UpdateLargest(message.Location{Group: d.GroupID, Object: d.ObjectID})

        // The cache keeps the buffers by reference; nothing mutates them after.
        entry.Cache.PutDatagram(d, in.MaxCacheDuration, in.HasMaxCacheDuration)

        downstream := entry.CopyDownstream()
        for _, sub := range downstream {
                // A datagram belongs to no Subgroup (§2.2), and §5.1.4 does not say
                // how a SUBGROUP_FILTER treats one; the relay filters it as Subgroup 0.
                if sub.ForwardDecision(d.GroupID, d.ObjectID, 0, d.PublisherPriority, d.Properties) != registry.Forward {
                        continue
                }
                // §9.7: only the Track Alias changes, bar the exception below.
                out := *d
                out.TrackAlias = sub.TrackAlias
                // A subscriber without Track Properties (§10.2.21) cannot inherit
                // DEFAULT_PUBLISHER_PRIORITY, so the priority is written out.
                if !sub.IncludesProperties() {
                        out.Type &^= message.DatagramDefaultPriorityBit
                }
                // §10.12: PUBLISH_DONE waits for a send in progress, and none starts
                // after it.
                if !sub.BeginDatagram() {
                        continue
                }
                err := sub.Session.SendDatagram(&out)
                sub.EndDatagram()
                if err != nil {
                        // §11.3: datagrams may be dropped.
                        h.log.LogAttrs(ctx, slog.LevelDebug, "datagram: SendDatagram failed",
                                slog.Uint64("sub_id", sub.ID),
                                slog.String("err", err.Error()))
                }
        }
}

package relay

import (
        "bytes"
        "fmt"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/relay/cache"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// recordRedundant handles dup, a copy [registry.TrackEntry.ClaimDelivered]
// found redundant, described to the ledger by info: it is checked against the
// first copy ([checkDuplicate]) and then recorded
// ([registry.TrackEntry.RecordDuplicate]). An error wraps
// [session.ErrMalformedTrack].
func recordRedundant(entry *registry.TrackEntry, info registry.ObjectInfo, dup *cache.CachedObject) error {
        if err := checkDuplicate(entry.Cache, dup); err != nil {
                return err
        }
        return entry.RecordDuplicate(info)
}

// checkDuplicate compares dup, a copy that lost the §9.3 dedup claim, with the
// first copy in c. A different Forwarding Preference, Subgroup ID, Priority or
// Payload (§9.1), or different Immutable Properties (§2.4.2, §12.7; one copy
// having them and the other not counts), makes the track malformed: the error
// wraps [session.ErrMalformedTrack]. Mutable Properties may differ (§9.1).
//
// Normal becoming End of Group or End of Track is the existing-to-not-existing
// change §9.1 allows, and the reverse a late Object (§2.1). Only a Normal
// Object has a Payload or Properties (§11.2.1.1, §11.2.1.2), so those are
// compared only when both copies are Normal.
//
// A copy an announced gap says does not exist is compared too, when the
// first copy is cached: §2.1 excuses its arrival, not a different content.
//
// On a difference the first copy is removed from the cache: it triggered the
// Malformed Track status too, and such Objects MUST NOT be cached (§2.4.2).
//
// Limitation: nothing is compared when the first copy is not in the cache:
// evicted, expired (§12.3), or not yet put there by a concurrent contributor.
func checkDuplicate(c *cache.ObjectCache, dup *cache.CachedObject) error {
        first, ok := c.Get(dup.GroupID, dup.ObjectID)
        if !ok || first.IsRangeMarker() {
                return nil
        }
        var field string
        switch {
        case first.ForwardingPref != dup.ForwardingPref:
                field = "Forwarding Preference"
        case first.ForwardingPref == cache.ForwardingSubgroup && first.SubgroupID != dup.SubgroupID:
                field = "Subgroup ID"
        case first.PublisherPriority != dup.PublisherPriority:
                field = "Priority"
        case first.Status != message.ObjectStatusNormal || dup.Status != message.ObjectStatusNormal:
                return nil
        case !bytes.Equal(first.Payload, dup.Payload):
                field = "Payload"
        case !sameImmutableProperties(first.Properties, dup.Properties):
                field = "Immutable Properties"
        default:
                return nil
        }
        c.Delete(dup.GroupID, dup.ObjectID)
        return fmt.Errorf("%w: a duplicate of Object %d in Group %d has a different %s (§9.1, §2.4.2)",
                session.ErrMalformedTrack, dup.ObjectID, dup.GroupID, field)
}

// sameImmutableProperties reports whether a and b, raw Object Properties,
// carry the same Immutable Properties, byte for byte, or neither has any.
func sameImmutableProperties(a, b []byte) bool {
        av, aok := message.ImmutableProperties(a)
        bv, bok := message.ImmutableProperties(b)
        return aok == bok && bytes.Equal(av, bv)
}

package relay

import (
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "math"
        "sync"
        "sync/atomic"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/relay/cache"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// fwdObject pairs a SubgroupObject with its absolute Object ID: filtering
// punches holes in the forwarded sequence, so the writer re-encodes the
// §11.4.2 ObjectIDDelta against its own outbound stream.
type fwdObject struct {
        obj        *message.SubgroupObject
        absID      uint64
        enqueuedAt time.Time // stamped in publish; used for the §8 lag window

        // maxCacheAge is the upstream's MAX_CACHE_DURATION (§12.3), zero for no
        // limit: the Object is not forwarded once older than that.
        maxCacheAge time.Duration

        // first marks the subgroup's true first object (§11.4.2 FIRST_OBJECT);
        // only an outbound stream beginning with it sets the bit.
        first bool

        // follows reports that the Object handed to the writer before this one
        // was read just before it from the same inbound stream, with only
        // Objects the subscriber's filters rejected between (see
        // [subgroupWriter.admit]).
        follows bool

        // pubTimeouts is [subgroupWriterSet.pubTimeouts] when this Object was
        // published: nil until the Subgroup's first Object was forwarded. Never
        // written through.
        pubTimeouts *message.DeliveryTimeouts
}

// inboundPos is an Object's place on its inbound subgroup stream: the stream,
// and how many Objects had been read from it, this one included.
type inboundPos struct {
        src *session.IncomingSubgroupStream
        seq uint64
}

// subgroupWriterSet is the payload of a [registry.SharedSubgroup]: one
// outbound writer per downstream subscriber for a (GroupID, SubgroupID),
// shared by every inbound stream contributing that Subgroup. All access holds
// [registry.SharedSubgroup.Mu].
//
// A nil writer records a sub that was not Established when scanned, so it is
// not retried.
type subgroupWriterSet struct {
        writers map[*registry.DownstreamSub]*subgroupWriter
        // hdr is the first contributor's SUBGROUP_HEADER, reused for every writer;
        // TrackAlias is overwritten per subscriber.
        hdr message.SubgroupHeader
        // gen is the downstream generation at the last joiner scan; the scan is
        // skipped while it is unchanged.
        gen uint64

        // lowest is the lowest Object ID forwarded, if forwarded. Objects are
        // published in ascending order (§2.2), so a FIRST_OBJECT claim (§11.4.2)
        // for a higher ID is wrong, whether or not a given subscriber got the
        // lower one. A new set starts from the ledger's
        // ([registry.TrackEntry.LowestForwarded]), so it holds across contributors
        // within the ledger's window; for the same reason
        // [subgroupWriterSet.outcome] measures a clean contributor's coverage
        // from it.
        lowest    uint64
        forwarded bool
        // runLo and runHi are the lowest and highest Object IDs forwarded
        // through this set, if hasRun; unbroken reports that every ID between
        // them was (see [subgroupWriterSet.admitAgedOut]).
        runLo, runHi     uint64
        hasRun, unbroken bool

        // cleanFrom is the lowest Object ID from which a contributor that ended
        // cleanly delivered every Object, if sawClean: 0 for one whose stream
        // starts the Subgroup (§11.4.2 FIRST_OBJECT), else its first Object's
        // (see [subgroupWriterSet.outcome]). resetCode is a reset contributor's
        // code, used when none ended cleanly.
        cleanFrom uint64
        sawClean  bool
        resetCode moqt.StreamResetCode

        // pubTimeouts points at firstTimeouts once the Subgroup's first Object
        // was forwarded: the publisher's §8 delivery timeouts for the Subgroup,
        // "the Object Property when present on the first object of the subgroup,
        // and the Track Property otherwise" (§12.1, §12.2). A replay stream (a
        // joiner's, or one reopened after a gap) starts past that Object, so its
        // session cannot find the override itself (see
        // [session.OutgoingSubgroupStream.WriteObjectReceivedAt]). Set once, and
        // handed to writers on every later [fwdObject]. While nil, as when the
        // relay never saw the first Object, writers keep the Track's.
        pubTimeouts   *message.DeliveryTimeouts
        firstTimeouts message.DeliveryTimeouts
}

// claimFirst records that the Object at objectID is forwarded and reports
// whether it starts the Subgroup: its contributor claims so (claimed) and no
// lower ID was forwarded (see subgroupWriterSet.lowest). Callers hold sg.Mu.
func (s *subgroupWriterSet) claimFirst(objectID uint64, claimed bool) bool {
        lowest := !s.forwarded || objectID < s.lowest
        if lowest {
                s.lowest, s.forwarded = objectID, true
        }
        switch {
        case !s.hasRun:
                s.runLo, s.runHi, s.hasRun, s.unbroken = objectID, objectID, true, true
        case s.runHi < math.MaxUint64 && objectID == s.runHi+1:
                s.runHi = objectID
        case s.runLo > 0 && objectID == s.runLo-1:
                s.runLo = objectID
        default:
                s.runLo, s.runHi, s.unbroken = min(s.runLo, objectID), max(s.runHi, objectID), false
        }
        return claimed && lowest
}

// leave records how one contributor ended: reset with code, or cleanly after
// delivering every Object from coverFrom on. covers is false for a replay
// stream that ended before its first Object, which vouches for none. Callers
// hold sg.Mu.
func (s *subgroupWriterSet) leave(reset bool, code moqt.StreamResetCode, covers bool, coverFrom uint64) {
        switch {
        case reset:
                s.resetCode = code
        case covers && (!s.sawClean || coverFrom < s.cleanFrom):
                s.cleanFrom, s.sawClean = coverFrom, true
        }
}

// outcome reports whether the merged streams end with a reset, and with which
// code, once every contributor has left. §11.4.3: "If a sender closes the
// stream before delivering all such objects to the QUIC stream, it MUST reset
// the stream." A clean contributor delivered every Object from cleanFrom on,
// so they FIN only if every Object forwarded below cleanFrom was followed by
// the next one up to it: an unbroken run from the lowest one forwarded
// (subgroupWriterSet.lowest) reaching cleanFrom - 1. Otherwise a reset
// contributor may have held Objects between that nobody forwarded, and they
// reset with CANCELLED.
//
// Deviation: lowest also counts what the ledger saw forwarded through an
// earlier set of the Subgroup, released when its last contributor left, whose
// run is forgotten; so a contributor arriving after that resets unless it
// covers from lowest, even if it did continue the earlier run, where §11.4.3
// says a sender that "has delivered all objects in a Subgroup ... MUST close
// the stream with a FIN".
//
// Interpretation: Objects below the lowest one forwarded count as before the
// Start Location, as for a joiner (see subgroupWriter.incomplete), so a lone
// replay upstream's FIN still FINs. The run is broken, and the streams reset,
// whenever Object IDs are forwarded out of order or are not consecutive (the
// Group split across Subgroups): the relay cannot tell a skipped ID from one
// that does not exist. Callers hold sg.Mu.
func (s *subgroupWriterSet) outcome() (reset bool, code moqt.StreamResetCode) {
        switch {
        case !s.sawClean:
                return true, s.resetCode
        case s.cleanFrom <= s.lowest,
                s.unbroken && s.runLo == s.lowest && s.runHi >= s.cleanFrom-1:
                return false, 0
        }
        return true, moqt.StreamResetCancelled
}

// admitAgedOut reports whether an Object at objectID of the Subgroup hdr
// names, which [registry.TrackEntry.ClaimDelivered] returned as claim, may be
// forwarded. Only a [registry.ClaimAgedOut] one may not: with its Group out of
// the dedup window, it goes on only above every Object forwarded through this
// set, where it can neither repeat one nor arrive out of order (§2.2). That
// covers a single upstream's open stream. One this set is known to have
// forwarded (an end of its run, or inside an unbroken one) is a redundant
// copy, dropped as [registry.ClaimRedundant] would be.
//
// Any other is dropped, and each stream it would have gone on is marked
// incomplete as [subgroupWriter.admit] would mark it, so it resets rather
// than FIN (§11.4.3). Deviation: §9.4 says a relay "MUST NOT reorder or drop
// objects received on a multi-object stream"; the relay drops such an Object
// rather than risk forwarding a duplicate it can no longer detect (§9.3). A
// Subgroup whose Object IDs are not consecutive (its Group split across
// Subgroups) has no unbroken run, so there a redundant copy resets too.
// Callers hold sg.Mu.
func (s *subgroupWriterSet) admitAgedOut(
        claim registry.Claim,
        hdr message.SubgroupHeader,
        objectID uint64,
        props []byte,
) bool {
        switch {
        case claim != registry.ClaimAgedOut, s.hasRun && objectID > s.runHi:
                return true
        case s.hasRun && (objectID == s.runLo || objectID == s.runHi ||
                s.unbroken && s.runLo < objectID && objectID < s.runHi):
                return false
        }
        for _, w := range s.writers {
                if w == nil {
                        continue
                }
                v := w.sub.ForwardDecision(hdr.GroupID, objectID, hdr.SubgroupID, hdr.PublisherPriority, props)
                if v == registry.Forward {
                        w.markIncomplete(moqt.StreamResetInternalError)
                } else {
                        w.skip(v, objectID)
                }
        }
        return false
}

// resolveImplicitSubgroupID handles §11.4.2 SUBGROUP_ID_MODE 0b01 (Subgroup ID
// = first Object ID): it reads the first object and rewrites hdr to the
// explicit form. The returned pending object must be processed as the
// stream's first. Other modes return (nil, true).
//
// ok=false means the stream ended first; the caller just returns. A malformed
// first Object ends the track (§2.4.2).
func (h *sessionHandler) resolveImplicitSubgroupID(
        ctx context.Context,
        entry *registry.TrackEntry,
        stream *session.IncomingSubgroupStream,
        hdr *message.SubgroupHeader,
) (pending *message.SubgroupObject, ok bool) {
        if hdr.SubgroupIDMode != message.SubgroupIDImplicitFirstObject {
                return nil, true
        }
        if hdr.ReplayingSubgroup {
                // On a replay the first object need not be the subgroup's first,
                // so the implied ID is only as reliable as the sender.
                h.log.LogAttrs(ctx, slog.LevelDebug,
                        "fanout: implicit-first-object Subgroup ID on a replay stream",
                        slog.Uint64("group", hdr.GroupID))
        }
        obj, err := stream.ReadObject()
        if err != nil {
                switch {
                case errors.Is(err, io.EOF):
                case errors.Is(err, session.ErrMalformedTrack):
                        stream.Cancel(moqt.StreamResetMalformedTrack)
                        h.endMalformedTrack(ctx, entry, h.sess, err)
                default:
                        h.log.LogAttrs(ctx, slog.LevelDebug,
                                "fanout: inbound stream ended before first-object Subgroup ID resolved",
                                slog.String("err", err.Error()))
                        // Stop a publisher still writing into a stream nobody reads.
                        stream.Cancel(moqt.StreamResetInternalError)
                }
                return nil, false
        }
        hdr.SubgroupID = obj.ObjectIDDelta // first object: the delta IS the absolute ID
        hdr.SubgroupIDMode = message.SubgroupIDExplicit
        return obj, true
}

// A subgroup stream can arrive before the SUBSCRIBE_OK binding its Track Alias
// (§11.1). §11.4.2 lets the receiver "abandon the stream, or choose to buffer
// it for a brief period"; the relay leaves it unread for up to earlyAliasWait,
// then abandons it.
//
// The deadline also breaks a flow-control deadlock: the bundled transports do
// not reserve connection credit for control streams (§11.4.2), so unread
// early data can hold back the SUBSCRIBE_OK itself. maxEarlyStreams caps the
// waiting streams per session; past it they are abandoned at once.
const (
        earlyAliasWait  = time.Second
        maxEarlyStreams = 32
)

// testHookEarlyStreamWaiting, when set, runs as a subgroup stream starts
// waiting for its Track Alias, so a test can send the SUBSCRIBE_OK only once
// the stream is known to have arrived first.
var testHookEarlyStreamWaiting atomic.Pointer[func(alias uint64)]

// resolveInboundTrack returns what stream's Track Alias is bound to, waiting
// within the bounds above. An unresolved alias abandons the stream and reports
// false: EXCESSIVE_LOAD (§3.3.4) past maxEarlyStreams, else INTERNAL_ERROR.
func (h *sessionHandler) resolveInboundTrack(
        ctx context.Context,
        stream *session.IncomingSubgroupStream,
) (session.InboundTrack, bool) {
        if in, ok := stream.InboundTrack(); ok {
                return in, true
        }
        defer h.earlyStreams.Add(-1)
        if h.earlyStreams.Add(1) > maxEarlyStreams {
                h.log.LogAttrs(ctx, slog.LevelWarn, "fanout: too many streams waiting for their Track Alias",
                        slog.Uint64("alias", stream.Header.TrackAlias))
                stream.Cancel(moqt.StreamResetExcessiveLoad)
                return session.InboundTrack{}, false
        }
        waitCtx, cancel := context.WithTimeout(ctx, earlyAliasWait)
        defer cancel()
        if hook := testHookEarlyStreamWaiting.Load(); hook != nil {
                (*hook)(stream.Header.TrackAlias)
        }
        in, ok := stream.AwaitInboundTrack(waitCtx)
        if !ok {
                h.log.LogAttrs(ctx, slog.LevelWarn, "fanout: Track Alias still unknown, abandoning stream",
                        slog.Uint64("alias", stream.Header.TrackAlias))
                stream.Cancel(moqt.StreamResetInternalError)
        }
        return in, ok
}

// runFanout forwards one inbound subgroup stream to every downstream
// subscriber, remapping the Track Alias per subscriber.
//
// §9.3: inbound streams carrying the same (GroupID, SubgroupID) share one
// outbound writer per subscriber (§2.2: a Subgroup is not split across
// streams), and [registry.TrackEntry.ClaimDelivered] drops duplicate objects
// (§9.3) and ones an announced gap says do not exist (§2.1, §9.1). Each writer is a [subgroupWriter] goroutine behind a bounded queue.
// The inbound FIN-vs-reset reaches the outbound streams only when the last
// contributor leaves.
func (h *sessionHandler) runFanout(ctx context.Context, stream *session.IncomingSubgroupStream) {
        hdr := stream.Header

        in, ok := h.resolveInboundTrack(ctx, stream)
        if !ok {
                return // abandoned: the stream is reset, the session stays up
        }
        key := in.Key

        entry, ok := h.tracks.Get(key)
        if !ok {
                // The subscription ended after the alias was registered.
                h.log.LogAttrs(ctx, slog.LevelDebug, "fanout: track entry gone, dropping stream",
                        slog.Uint64("alias", hdr.TrackAlias))
                stream.Cancel(moqt.StreamResetInternalError)
                return
        }

        // §11.4.2: a header without a Priority byte inherits the alias's
        // DEFAULT_PUBLISHER_PRIORITY (§12.4). Resolve it once so the cache, FETCH,
        // PRIORITY_FILTER and §7.2 scheduling see it; the outbound header still
        // omits the byte (see openWriterForSub for the exception).
        if !hdr.InlinePriority {
                hdr.PublisherPriority = in.DefaultPublisherPriority
        }

        // One TrackRef per stream: it allocates, and is reported per object.
        ref := h.trackRef(entry.FullName)

        // §12.3: a MAX_CACHE_DURATION of 0 limits only serving from the cache.
        var liveMaxAge time.Duration
        if in.HasMaxCacheDuration {
                liveMaxAge = in.MaxCacheDuration
        }

        // Everything below keys on hdr.SubgroupID, so resolve it first.
        pending, ok := h.resolveImplicitSubgroupID(ctx, entry, stream, &hdr)
        if !ok {
                return
        }

        sgKey := registry.SubgroupKey{Group: hdr.GroupID, Subgroup: hdr.SubgroupID}
        sg, created := entry.AcquireSubgroup(sgKey, func() any {
                return &subgroupWriterSet{
                        writers: make(map[*registry.DownstreamSub]*subgroupWriter),
                        hdr:     hdr,
                }
        })
        set, _ := sg.Set.(*subgroupWriterSet)

        if created {
                // Under sg.Mu so a concurrent contributor's joiner scan can't
                // double-open. The stream is drained even with no subscribers, so its
                // Objects still reach the cache (§9.1) and any later joiner.
                initialSubs, gen := entry.CopyDownstreamWithGen()
                pubTimeouts := entry.DeliveryTimeouts()
                lowest, forwarded := entry.LowestForwarded(hdr.GroupID, hdr.SubgroupID)
                sg.Mu.Lock()
                set.gen = gen
                // The minimum, not an assignment: a contributor that joined the set
                // before this lock may already have forwarded a lower Object.
                if forwarded && (!set.forwarded || lowest < set.lowest) {
                        set.lowest, set.forwarded = lowest, true
                }
                for _, sub := range initialSubs {
                        h.openWriterForSub(ctx, set.hdr, sub, set.writers, pubTimeouts, ref)
                }
                sg.Mu.Unlock()
        }

        // This contributor's termination, applied outbound only if it is the last
        // to leave the Subgroup (§9.3).
        var (
                inboundReset     bool
                inboundResetCode = moqt.StreamResetCancelled
                // See [subgroupWriterSet.leave].
                covers    = !hdr.ReplayingSubgroup
                coverFrom uint64
        )
        defer func() {
                // Record the outcome before releasing, so the last contributor decides
                // FIN vs reset over all of them.
                sg.Mu.Lock()
                set.leave(inboundReset, inboundResetCode, covers, coverFrom)
                last := entry.ReleaseSubgroup(sgKey)
                if !last {
                        sg.Mu.Unlock()
                        return // other upstreams still feed this Subgroup — leave writers up.
                }
                reset, code := set.outcome()
                ws := make([]*subgroupWriter, 0, len(set.writers))
                for _, w := range set.writers {
                        if w == nil {
                                continue
                        }
                        wReset, wCode := reset, code
                        // §11.4.3: a group now outside the subscription's range is
                        // reset, not FIN'd.
                        if !wReset && registry.GroupOutOfRange(hdr.GroupID, w.sub.GetFilter()) {
                                wReset, wCode = true, moqt.StreamResetCancelled
                        }
                        w.close(wReset, wCode)
                        ws = append(ws, w)
                }
                sg.Mu.Unlock()
                joinWriters(ws)
        }()

        var (
                firstObj = true
                // last is the last Object read, which a FIN ends the Subgroup after.
                last *message.SubgroupObject
                // pos counts every Object read, dedup losers included: each is an
                // Object between its neighbours (§11.4.3).
                pos = inboundPos{src: stream}
                // terminalSeen: an EndOfGroup/EndOfTrack was read on this inbound
                // stream; any later object makes the track malformed (§11.4.3,
                // §2.4.2).
                terminalSeen bool
        )
        // malformed ends the track for an Object read from stream (§2.4.2); the
        // Object is neither cached nor forwarded.
        malformed := func(err error) {
                stream.Cancel(moqt.StreamResetMalformedTrack)
                inboundReset, inboundResetCode = true, moqt.StreamResetMalformedTrack
                h.endMalformedTrack(ctx, entry, h.sess, err)
        }

        for {
                obj, err := pending, error(nil)
                pending = nil
                if obj == nil {
                        obj, err = stream.ReadObject()
                }
                if err != nil {
                        inboundReset, inboundResetCode = h.inboundEnded(ctx, entry, stream, hdr, last, err)
                        return
                }
                last = obj

                if terminalSeen {
                        h.log.LogAttrs(ctx, slog.LevelDebug,
                                "fanout: object after EndOfGroup/EndOfTrack — malformed track",
                                slog.Uint64("group", hdr.GroupID), slog.Uint64("subgroup", hdr.SubgroupID))
                        malformed(fmt.Errorf("object after END_OF_GROUP / END_OF_TRACK in Group %d", hdr.GroupID))
                        return
                }

                isTrueFirst := firstObj && !hdr.ReplayingSubgroup
                firstObj = false
                pos.seq++
                objectID := stream.ObjectID() // resolved by ReadObject (§11.4.2)
                if !covers {
                        coverFrom, covers = objectID, true
                }

                // Whether or not this copy wins the dedup claim below; the next
                // iteration acts on it, so it can be set now.
                terminalSeen = obj.IsTerminal()

                // §9.3: the first upstream to deliver {GroupID, ObjectID} forwards it,
                // unless an announced gap says it does not exist (§2.1, §9.1). Outside
                // sg.Mu, so dedup losers never touch the writer set.
                info := registry.ObjectInfo{
                        Group:    hdr.GroupID,
                        Object:   objectID,
                        Subgroup: hdr.SubgroupID,
                        Priority: hdr.PublisherPriority,
                        Status:   obj.ObjectStatus,
                        Gaps:     message.ObjectPriorGaps(obj.Properties),
                }
                claim, err := entry.ClaimDelivered(info)
                if err != nil {
                        malformed(err)
                        return
                }
                if claim == registry.ClaimRedundant {
                        if err := recordRedundant(entry, info, &cache.CachedObject{
                                GroupID:           hdr.GroupID,
                                ObjectID:          objectID,
                                SubgroupID:        hdr.SubgroupID,
                                PublisherPriority: hdr.PublisherPriority,
                                ForwardingPref:    cache.ForwardingSubgroup,
                                Status:            obj.ObjectStatus,
                                Properties:        obj.Properties,
                                Payload:           obj.Payload,
                        }); err != nil {
                                malformed(err)
                                return
                        }
                        continue // redundant copy already forwarded by a peer upstream.
                }

                // Under sg.Mu: joiner detection, writer open and publish are atomic
                // against other contributors and the last-contributor teardown.
                sg.Mu.Lock()
                if !set.admitAgedOut(claim, hdr, objectID, obj.Properties) {
                        sg.Mu.Unlock()
                        continue
                }

                // Counted after the dedup claim, so redundant copies don't count.
                h.metrics.ObjectReceived(ref, hdr.SubgroupID)

                // Cache before bumping LARGEST_OBJECT, so a FETCH that snapshots the
                // new watermark finds the object cached.
                entry.Cache.Put(&cache.CachedObject{
                        GroupID:           hdr.GroupID,
                        ObjectID:          objectID,
                        SubgroupID:        hdr.SubgroupID,
                        PublisherPriority: hdr.PublisherPriority,
                        ForwardingPref:    cache.ForwardingSubgroup,
                        Status:            obj.ObjectStatus,
                        Properties:        obj.Properties,
                        Payload:           obj.Payload,

                        MaxCacheDuration:    in.MaxCacheDuration,
                        HasMaxCacheDuration: in.HasMaxCacheDuration,
                })

                // Serialised with AddDownstreamSnapshotLargest: a new sub either saw
                // the old Largest and appears in newSubs (delivered live), or saw the
                // new one (its fill fetch stream covers this object).
                loc := message.Location{Group: hdr.GroupID, Object: objectID}
                var newSubs []*registry.DownstreamSub
                newSubs, set.gen = entry.UpdateLargestAndDetectNew(loc,
                        func(s *registry.DownstreamSub) bool { _, ok := set.writers[s]; return ok }, set.gen)
                for _, sub := range newSubs {
                        h.openWriterForSub(ctx, set.hdr, sub, set.writers, entry.DeliveryTimeouts(), ref)
                }

                set.forward(entry, pos, hdr, objectID, obj, isTrueFirst, liveMaxAge)
                sg.Mu.Unlock()
        }
}

// forward hands obj, the Object at objectID of the Subgroup hdr names, read at
// pos, to every writer whose subscriber takes it. claimed reports that its
// contributor claims it starts the Subgroup (see [subgroupWriterSet.claimFirst]);
// maxCacheAge is [fwdObject.maxCacheAge]. Callers hold sg.Mu.
func (s *subgroupWriterSet) forward(
        entry *registry.TrackEntry,
        pos inboundPos,
        hdr message.SubgroupHeader,
        objectID uint64,
        obj *message.SubgroupObject,
        claimed bool,
        maxCacheAge time.Duration,
) {
        first := s.claimFirst(objectID, claimed)
        // §8: the first Object's Properties settle the publisher's timeouts for
        // the Subgroup.
        if first && s.pubTimeouts == nil {
                s.firstTimeouts = entry.DeliveryTimeouts().ApplyObjectProperties(obj.Properties)
                s.pubTimeouts = &s.firstTimeouts
        }

        // §5.1.2 filters run before enqueue, so a miss takes no queue slot.
        for _, w := range s.writers {
                if w == nil {
                        continue
                }
                if take, follows := w.admit(pos, hdr, objectID, obj.Properties); take {
                        w.publish(fwdObject{
                                obj:         obj,
                                absID:       objectID,
                                first:       first,
                                maxCacheAge: maxCacheAge,
                                follows:     follows,
                                pubTimeouts: s.pubTimeouts,
                        })
                }
        }
}

// inboundEnded handles err, which ended the reads of stream (whose header is
// hdr, and whose last Object read was last, nil if none), and reports whether
// the stream counts as reset, and with which code, for the Subgroup's FIN
// (§9.3).
func (h *sessionHandler) inboundEnded(
        ctx context.Context,
        entry *registry.TrackEntry,
        stream *session.IncomingSubgroupStream,
        hdr message.SubgroupHeader,
        last *message.SubgroupObject,
        err error,
) (reset bool, code moqt.StreamResetCode) {
        switch {
        case errors.Is(err, io.EOF):
                // Clean end of stream; the last contributor will FIN. It ends the
                // Subgroup after the last Object read (§2.4.2).
                if last == nil {
                        return false, 0
                }
                err = entry.SubgroupEnded(registry.ObjectInfo{
                        Group:    hdr.GroupID,
                        Object:   stream.ObjectID(),
                        Subgroup: hdr.SubgroupID,
                        Priority: hdr.PublisherPriority,
                        Status:   last.ObjectStatus,
                }, hdr.EndOfGroup)
                if err == nil {
                        return false, 0
                }
        case errors.Is(err, context.Canceled):
                // The session is going away: reset, never FIN.
                return true, moqt.StreamResetCancelled
        case !errors.Is(err, session.ErrMalformedTrack):
                h.log.LogAttrs(ctx, slog.LevelDebug, "fanout: inbound ReadObject failed",
                        slog.String("err", err.Error()))
                // An unparseable object leaves the publisher writing; stop it.
                stream.Cancel(moqt.StreamResetInternalError)
                return true, moqt.StreamResetCancelled
        }
        stream.Cancel(moqt.StreamResetMalformedTrack)
        h.endMalformedTrack(ctx, entry, h.sess, err)
        return true, moqt.StreamResetMalformedTrack
}

// openWriterForSub starts a subgroupWriter for sub and records it in writers
// (nil when sub is not Established, so it is not retried).
//
// No transport I/O here: callers hold sg.Mu, and a header write blocked on
// one subscriber's flow control would stall the whole subgroup. The writer
// opens its stream lazily, before the first object it forwards.
func (h *sessionHandler) openWriterForSub(
        ctx context.Context,
        hdr message.SubgroupHeader,
        sub *registry.DownstreamSub,
        writers map[*registry.DownstreamSub]*subgroupWriter,
        pubTimeouts message.DeliveryTimeouts,
        ref TrackRef,
) {
        if _, already := writers[sub]; already {
                return
        }
        if !sub.IsEstablished() {
                writers[sub] = nil
                return
        }
        subHdr := hdr
        subHdr.TrackAlias = sub.TrackAlias
        // A subscriber without Track Properties (§10.2.21) cannot inherit
        // DEFAULT_PUBLISHER_PRIORITY (§12.4), so the priority is written out.
        if !sub.IncludesProperties() {
                subHdr.InlinePriority = true
        }
        // cancelIO unblocks a writer wedged on a subscriber that stopped
        // reading (see [joinWriters]), and resets the open stream of one that
        // cancelled (§5.1.1).
        ioCtx, cancelIO := context.WithCancel(ctx)
        w := &subgroupWriter{
                sub:                 sub,
                ctx:                 ioCtx,
                cancelIO:            cancelIO,
                unwatchSub:          context.AfterFunc(sub.Cancelled(), cancelIO),
                hdr:                 subHdr,
                inbox:               make(chan fwdObject, h.sendQueueSize),
                done:                make(chan struct{}),
                log:                 h.log,
                metrics:             h.metrics,
                ref:                 ref,
                maxDropsBeforeReset: h.maxDropsBeforeReset,
                maxLag:              h.maxFanoutLag,
                // §8: kept apart, since the §12.1/§12.2 first-object override
                // applies to the publisher's half alone.
                pubTimeouts: pubTimeouts,
                subTimeouts: sub.GetDeliveryTimeouts(),
        }
        writers[sub] = w
        h.spawn(w.run)
}

// subgroupWriter is the per-subscriber writer goroutine: it drains an inbox
// onto outbound subgroup streams on the subscriber's session.
//
//   - An Object that is not the next Object resets the stream and opens a
//     fresh one (§11.4.3, see [isNextObject]).
//   - A clean inbound EOF FINs the stream; an inbound error resets it.
//   - A full inbox drops the object. An object that waited longer than
//     maxLag resets with TOO_FAR_BEHIND and terminates the subscription
//     (§3.3.4); the optional maxDropsBeforeReset cap does so with
//     EXCESSIVE_LOAD.
//   - An elapsed §8 delivery timeout resets only that stream with
//     DELIVERY_TIMEOUT; the subscription survives (§3.3.4).
type subgroupWriter struct {
        sub *registry.DownstreamSub
        // ctx bounds every blocking stream operation; cancelIO resets the
        // in-flight stream, unwedging a writer blocked on a stalled subscriber.
        // The subscriber's cancellation calls it too (§5.1.1), until unwatchSub.
        ctx        context.Context
        cancelIO   context.CancelFunc
        unwatchSub func() bool
        hdr        message.SubgroupHeader          // template; TrackAlias already remapped
        out        *session.OutgoingSubgroupStream // nil until run opens it lazily
        unbridge   func() bool                     // stops the current stream's ctx→Cancel bridge
        inbox      chan fwdObject
        done       chan struct{}
        log        *slog.Logger
        metrics    Metrics
        // ref labels every Metrics call; built once, since it allocates.
        ref                 TrackRef
        maxDropsBeforeReset int
        maxLag              time.Duration
        // pubTimeouts and subTimeouts are the §8 delivery-timeout halves, resolved
        // per outbound stream; zero disables a dimension. run replaces
        // pubTimeouts with [fwdObject.pubTimeouts] once that is set.
        pubTimeouts message.DeliveryTimeouts
        subTimeouts message.DeliveryTimeouts

        closeOnce        sync.Once
        dropsMu          sync.Mutex
        drops            int
        closed           bool                 // set under dropsMu inside close
        inboundReset     bool                 // set under dropsMu inside close
        inboundResetCode moqt.StreamResetCode // §3.3.4 reset code when inboundReset; set inside close
        // incomplete records that this subscription skipped an Object after its
        // Start Location (filter, Forward State 0, overflow, expiry, or an Object
        // dropped as aged out of the dedup window), so its streams end with a
        // reset, not a FIN (§11.4.3), with incompleteCode: EXCESSIVE_LOAD after any
        // overflow, else the first recorded (INTERNAL_ERROR for an aged-out drop,
        // else CANCELLED). Set under dropsMu.
        //
        // Interpretation: Objects published before the subscription joined count
        // as before its Start Location, so a joiner's stream may still FIN.
        incomplete     bool
        incompleteCode moqt.StreamResetCode
        // lastAdmitted is the last Object ID admit let through: a SkipBeforeStart
        // above it means the Start was raised past sent Objects; one below is a
        // straggler from another upstream. Only touched by admit, under sg.Mu.
        lastAdmitted uint64
        hasAdmitted  bool
        // lastPos is where the last Object admit let through was read, moved
        // past the filtered Objects read after it; zero once the relay drops
        // one. Only touched by admit and publish, under sg.Mu.
        lastPos inboundPos
        // withProps: the streams run opens set PROPERTIES (§11.4.2). Only
        // touched by run.
        withProps bool
}

// admit decides whether w takes the Object at objectID of the subgroup hdr
// names, read at pos, closing w when it will take none again. follows is
// [fwdObject.follows] for a taken Object.
func (w *subgroupWriter) admit(
        pos inboundPos,
        hdr message.SubgroupHeader,
        objectID uint64,
        props []byte,
) (take, follows bool) {
        follows = w.lastPos.src == pos.src && w.lastPos.seq+1 == pos.seq
        v := w.sub.ForwardDecision(hdr.GroupID, objectID, hdr.SubgroupID, hdr.PublisherPriority, props)
        if v == registry.Forward {
                w.lastAdmitted, w.hasAdmitted = objectID, true
                w.lastPos = pos
                return true, follows
        }
        w.skip(v, objectID)
        // §11.4.3: an Object that "did not pass the subscriber's filters" does
        // not separate the ones either side of it. Forward State is one of them
        // (§5.1.5).
        if follows {
                w.lastPos = pos
        }
        return false, false
}

// skip records that w omits the Object at objectID for verdict v, other than
// [registry.Forward], closing w when it will take none again.
func (w *subgroupWriter) skip(v registry.ForwardVerdict, objectID uint64) {
        switch v {
        case registry.SkipObject, registry.SkipPaused:
                // The stream stays open for later Objects, but the Subgroup is now
                // incomplete (§11.4.3).
                w.markIncomplete(moqt.StreamResetCancelled)
        case registry.SkipGroup, registry.SkipEnded:
                // No further Object will pass: reset once the queue is written
                // (§11.4.3), so a PUBLISH_DONE waiting on the streams can follow.
                w.close(true, moqt.StreamResetCancelled)
        case registry.SkipBeforeStart:
                // §11.4.3 allows a FIN after omitting these, unless a REQUEST_UPDATE
                // raised the Start past an admitted Object.
                if w.hasAdmitted && objectID > w.lastAdmitted {
                        w.markIncomplete(moqt.StreamResetCancelled)
                }
        case registry.Forward:
        }
}

// markIncomplete sets subgroupWriter.incomplete; the first code recorded
// stands, except that EXCESSIVE_LOAD overrides.
func (w *subgroupWriter) markIncomplete(code moqt.StreamResetCode) {
        w.dropsMu.Lock()
        w.markIncompleteLocked(code)
        w.dropsMu.Unlock()
}

func (w *subgroupWriter) markIncompleteLocked(code moqt.StreamResetCode) {
        // The subscriber expects its own filter's omissions, not the relay's
        // load (§3.3.4).
        if !w.incomplete || code == moqt.StreamResetExcessiveLoad {
                w.incomplete = true
                w.incompleteCode = code
        }
}

// resetCode is incompleteCode when incomplete, else CANCELLED.
func (w *subgroupWriter) resetCode() moqt.StreamResetCode {
        w.dropsMu.Lock()
        defer w.dropsMu.Unlock()
        if w.incomplete {
                return w.incompleteCode
        }
        return moqt.StreamResetCancelled
}

// publish enqueues fwd without blocking, stamping its enqueue time for the
// lag check. On overflow the object is dropped, and past maxDropsBeforeReset
// the writer is closed in reset mode. It is a no-op after close. Callers hold
// sg.Mu.
func (w *subgroupWriter) publish(fwd fwdObject) {
        w.dropsMu.Lock()
        if w.closed {
                w.dropsMu.Unlock()
                return
        }
        w.dropsMu.Unlock()

        fwd.enqueuedAt = time.Now()
        select {
        case w.inbox <- fwd:
                w.metrics.ObjectForwarded(w.ref, w.hdr.SubgroupID)
        default:
                w.metrics.ObjectDropped(w.ref, w.hdr.SubgroupID)
                w.lastPos = inboundPos{} // the next Object does not follow a sent one
                w.dropsMu.Lock()
                w.drops++
                w.markIncompleteLocked(moqt.StreamResetExcessiveLoad)
                drops := w.drops
                capped := w.maxDropsBeforeReset > 0 && w.drops > w.maxDropsBeforeReset
                w.dropsMu.Unlock()

                w.log.Debug("fanout: dropped object on full inbox",
                        "sub_id", w.sub.ID, "drops", drops)
                if capped {
                        w.log.Warn("fanout: subscriber hit MaxDropsBeforeReset cap, terminating",
                                "sub_id", w.sub.ID, "drops", drops)
                        w.close(true, moqt.StreamResetExcessiveLoad)
                }
        }
}

// lagging reports whether fwd waited in the queue longer than the §8 lag
// window allows.
func (w *subgroupWriter) lagging(fwd fwdObject) bool {
        return w.maxLag > 0 && time.Since(fwd.enqueuedAt) > w.maxLag
}

// expired reports whether fwd is older than its MAX_CACHE_DURATION (§12.3).
// Deviation: the age runs from when the relay finished reading the Object,
// not from "the beginning of the Object".
func expired(fwd fwdObject) bool {
        return fwd.maxCacheAge > 0 && time.Since(fwd.enqueuedAt) > fwd.maxCacheAge
}

// dropExpired handles an Object skipped by [subgroupWriter.expired]. A
// header-only stream may claim FIRST_OBJECT for it, so it is reset and the
// next Object opens a replay stream (§11.4.2).
func (w *subgroupWriter) dropExpired(hasWritten bool) {
        w.markIncomplete(moqt.StreamResetCancelled)
        if hasWritten || w.out == nil {
                return
        }
        if w.unbridge != nil {
                w.unbridge()
                w.unbridge = nil
        }
        w.closeOut(false, moqt.StreamResetCancelled)
}

// closeOut FINs or resets the current outbound stream and reports it closed
// to the subscription, whose PUBLISH_DONE waits on its streams (§10.12).
func (w *subgroupWriter) closeOut(fin bool, code moqt.StreamResetCode) {
        if w.out == nil {
                return
        }
        if fin {
                _ = w.out.Close()
        } else {
                w.out.Cancel(code)
        }
        w.dropOut()
}

// dropOut is closeOut for an outbound stream the session has already reset.
func (w *subgroupWriter) dropOut() {
        if w.out == nil {
                return
        }
        w.out = nil
        w.sub.StreamClosed()
}

// run drains the inbox onto outbound streams until close, then FINs or
// resets the stream from what close recorded. After a write failure it keeps
// draining, so publish never blocks.
func (w *subgroupWriter) run() {
        defer close(w.done)
        defer w.unwatchSub()

        var (
                prevID      uint64
                hasWritten  bool
                writeFailed bool
                // dropped: an Object was dropped since the last one written, so the
                // next is not known to follow it.
                dropped bool
        )
        w.withProps = w.hdr.Properties

        // reopen resets the current outbound stream (if any) and opens a fresh
        // one, for the lazy first open, after a §11.4.3 gap, and to carry Object
        // Properties the old header could not. first sets the §11.4.2
        // FIRST_OBJECT bit; otherwise the stream is a replay. All its blocking
        // I/O is bounded by w.ctx.
        reopen := func(first bool) bool {
                if w.unbridge != nil {
                        w.unbridge()
                        w.unbridge = nil
                }
                w.closeOut(false, w.resetCode())
                hdr := w.hdr
                hdr.Properties = w.withProps
                hdr.ReplayingSubgroup = !first
                if !first && hdr.SubgroupIDMode == message.SubgroupIDImplicitFirstObject {
                        // A replay stream's first object would imply the wrong ID.
                        hdr.SubgroupIDMode = message.SubgroupIDExplicit
                }
                fresh, err := w.openCounted(hdr)
                if err != nil {
                        w.log.Debug("fanout: OpenSubgroup (reopen) failed",
                                "sub_id", w.sub.ID, "err", err.Error())
                        return false
                }
                // §8: WithDeliveryTimeouts returns a copy; the bridge must cancel it.
                fresh = fresh.WithDeliveryTimeouts(w.pubTimeouts, w.subTimeouts)
                w.out = fresh
                w.unbridge = context.AfterFunc(w.ctx, func() {
                        fresh.Cancel(moqt.StreamResetCancelled)
                })
                hasWritten = false
                w.applyPriority()
                // §11.4.3: keep the new stream's header reliable across resets.
                w.out.MarkReliable()
                return true
        }
        defer func() {
                if w.unbridge != nil {
                        w.unbridge()
                }
                // Guard: an unreported stream would hold PUBLISH_DONE forever (§10.12).
                w.closeOut(false, moqt.StreamResetCancelled)
        }()

        // failWrites latches this writer broken and stops contributors
        // enqueueing. Only close, under sg.Mu, closes the inbox.
        var writeFailedLatched bool
        failWrites := func() {
                writeFailed = true
                if !writeFailedLatched {
                        writeFailedLatched = true
                        w.dropsMu.Lock()
                        w.closed = true
                        w.dropsMu.Unlock()
                }
        }

        var lagExceeded bool
        for fwd := range w.inbox {
                if w.lagging(fwd) {
                        w.log.Warn("fanout: subscriber exceeded MaxFanoutLag, terminating",
                                "sub_id", w.sub.ID, "lag", time.Since(fwd.enqueuedAt).String())
                        lagExceeded = true
                        break
                }

                if writeFailed {
                        continue
                }

                // §8: the Subgroup's timeouts, for the streams reopen opens below.
                if fwd.pubTimeouts != nil {
                        w.pubTimeouts = *fwd.pubTimeouts
                }

                cause, stale := w.reopenCause(fwd, prevID, hasWritten, dropped)

                // Lazy first open, off sg.Mu (see openWriterForSub).
                if w.out == nil {
                        if !reopen(fwd.first) {
                                failWrites()
                                continue
                        }
                }

                if stale {
                        w.metrics.SubgroupStreamReset(w.ref, w.hdr.SubgroupID, cause)
                        if !reopen(fwd.first) {
                                failWrites()
                                continue
                        }
                }

                // §12.3: "MUST NOT start forwarding" an expired Object. Checked after
                // the open above, which can block.
                if expired(fwd) {
                        w.dropExpired(hasWritten)
                        dropped = true
                        continue
                }

                // Re-encode ObjectIDDelta against this outbound stream (§11.4.2).
                out := *fwd.obj
                if !hasWritten {
                        out.ObjectIDDelta = fwd.absID
                } else {
                        out.ObjectIDDelta = fwd.absID - prevID - 1
                }
                // §8 measures OBJECT_DELIVERY_TIMEOUT from when the object was
                // received. Deviation: enqueuedAt is stamped after the whole object
                // was read, not at its first byte, so the timeout is lenient.
                if err := w.out.WriteObjectReceivedAt(fwd.enqueuedAt, &out); err != nil {
                        // The stream is already reset with DELIVERY_TIMEOUT (§3.3.4);
                        // resetting again would overwrite that code.
                        if errors.Is(err, session.ErrDeliveryTimeout) {
                                w.log.Debug("fanout: delivery timeout, abandoning subgroup stream",
                                        "sub_id", w.sub.ID, "group", w.hdr.GroupID,
                                        "subgroup", w.hdr.SubgroupID)
                                w.metrics.SubgroupStreamReset(w.ref, w.hdr.SubgroupID, ResetCauseDeliveryTimeout)
                                w.dropOut()
                                failWrites()
                                continue
                        }
                        w.log.Debug("fanout: WriteObject failed",
                                "sub_id", w.sub.ID, "err", err.Error())
                        w.metrics.SubgroupStreamReset(w.ref, w.hdr.SubgroupID, ResetCauseWriteError)
                        w.closeOut(false, moqt.StreamResetInternalError)
                        failWrites()
                        continue
                }
                prevID = fwd.absID
                hasWritten = true
                dropped = false
                // §11.4.3: a later reset still delivers what was written.
                w.out.MarkReliable()
        }

        w.dropsMu.Lock()
        dropCapped := w.maxDropsBeforeReset > 0 && w.drops > w.maxDropsBeforeReset
        inboundReset := w.inboundReset
        inboundResetCode := w.inboundResetCode
        incomplete, incompleteCode := w.incomplete, w.incompleteCode
        w.dropsMu.Unlock()

        if lagExceeded || dropCapped {
                // Slow reader: reset and terminate the subscription. §3.3.4:
                // TOO_FAR_BEHIND for the lag window, EXCESSIVE_LOAD for the drop cap.
                resetCode := moqt.StreamResetTooFarBehind
                cause := ResetCauseTooFarBehind
                if dropCapped && !lagExceeded {
                        resetCode = moqt.StreamResetExcessiveLoad
                        cause = ResetCauseExcessiveLoad
                }
                w.metrics.SubscriptionResetSlowReader(w.ref, cause)
                // Refuse further enqueues; only close may close the inbox.
                w.dropsMu.Lock()
                w.closed = true
                w.dropsMu.Unlock()
                w.closeOut(false, resetCode)

                // Cancel the request stream so handleSubscribe unregisters the sub.
                // Only if this writer ended it: otherwise a PUBLISH_DONE may be under
                // way, and the reset could discard it.
                if w.sub.Terminate() && w.sub.Stream != nil {
                        w.sub.Stream.CancelRead(uint64(resetCode))
                        w.sub.Stream.CancelWrite(uint64(resetCode))
                }
                return
        }

        if writeFailed {
                return
        }

        if w.out == nil {
                return
        }

        if inboundReset {
                // §11.4.3: "A relay might immediately reset the corresponding
                // downstream stream".
                w.metrics.SubgroupStreamReset(w.ref, w.hdr.SubgroupID, ResetCauseInboundReset)
                w.closeOut(false, inboundResetCode)
                return
        }

        if incomplete {
                // §11.4.3: FIN only after "all objects in a Subgroup".
                w.closeOut(false, incompleteCode)
                return
        }

        w.closeOut(true, 0)
}

// reopenCause reports whether fwd needs a fresh outbound stream after one
// whose last Object is prevID, and why. §11.4.3: only "the next Object" may go
// on an existing stream. And the header is the first contributor's (§9.3), so
// a later one's Object Properties, which MUST be forwarded (§2.5), turn
// PROPERTIES on for this and every later stream.
func (w *subgroupWriter) reopenCause(
        fwd fwdObject,
        prevID uint64,
        hasWritten, dropped bool,
) (ResetCause, bool) {
        needProps := !w.withProps && len(fwd.obj.Properties) > 0
        if needProps {
                w.withProps = true
        }
        switch {
        case !hasWritten:
                return 0, false
        case needProps:
                return ResetCauseProperties, true
        case !isNextObject(fwd, prevID, dropped):
                return ResetCauseGap, true
        }
        return 0, false
}

// isNextObject reports whether fwd is "the next Object" (§11.4.3) on a stream
// whose last Object is prevID. Of the draft's ways to tell, the relay uses:
// the Object ID is one greater; fwd follows the last Object on its inbound
// stream; or its Prior Object ID Gap (§12.9) says the IDs between do not
// exist. Knowing from the cache or the subscriber's filters that they are in
// other Subgroups or filtered out is not used (a choice). A gap that also
// covers prevID says prevID no longer exists (§2.1), which shows nothing about
// the order.
func isNextObject(fwd fwdObject, prevID uint64, dropped bool) bool {
        if fwd.absID == prevID+1 {
                return true
        }
        if dropped || fwd.absID <= prevID {
                return false
        }
        if fwd.follows {
                return true
        }
        gap, ok := message.PriorObjectIDGap(fwd.obj.Properties)
        return ok && fwd.absID-gap == prevID+1
}

// openCounted opens a subgroup stream, counting it for the §10.12 Stream
// Count; it fails once the subscription has terminated.
func (w *subgroupWriter) openCounted(hdr message.SubgroupHeader) (*session.OutgoingSubgroupStream, error) {
        if !w.sub.BeginStream() {
                return nil, errSubscriptionTerminated
        }
        out, err := w.sub.Session.OpenSubgroupContext(w.ctx, hdr)
        w.sub.EndStream(err == nil)
        return out, err
}

// applyPriority sets the §7.2 effective priority on the current outbound
// stream. It runs on each (re)open, so a SUBSCRIBER_PRIORITY change applies
// from the next stream.
func (w *subgroupWriter) applyPriority() {
        if w.out == nil {
                return
        }
        w.out.SetSendPriority(w.sub.EffectiveStreamPriority(
                w.hdr.PublisherPriority, w.hdr.GroupID, w.hdr.SubgroupID,
        ))
}

// close closes the inbox, recording whether the writer ends its stream with
// a reset (and code) or a FIN. The first call wins. Queued objects are still
// written; a wedged writer is bounded by [joinWriters].
func (w *subgroupWriter) close(reset bool, code moqt.StreamResetCode) {
        w.closeOnce.Do(func() {
                w.dropsMu.Lock()
                w.closed = true
                w.inboundReset = reset
                w.inboundResetCode = code
                w.dropsMu.Unlock()
                close(w.inbox)
        })
}

// defaultWriterJoinTimeout bounds [joinWriters] when no MaxFanoutLag is
// configured.
const defaultWriterJoinTimeout = 5 * time.Second

// joinTimeout is the deadline for [joinWriters]: a healthy writer drains
// within MaxFanoutLag or terminates itself.
func (w *subgroupWriter) joinTimeout() time.Duration {
        if w.maxLag > 0 {
                return w.maxLag
        }
        return defaultWriterJoinTimeout
}

// joinWriters waits for every writer to finish after close. A writer wedged
// in a stream write never dequeues again, so at one shared deadline every
// still-running writer's I/O is cancelled: N stalled subscribers cost one
// timeout, not N.
func joinWriters(ws []*subgroupWriter) {
        if len(ws) == 0 {
                return
        }
        t := time.NewTimer(ws[0].joinTimeout()) // same handler config across ws
        defer t.Stop()
        for i, w := range ws {
                select {
                case <-w.done:
                        continue
                case <-t.C:
                        for _, u := range ws[i:] {
                                select {
                                case <-u.done:
                                        continue
                                default:
                                }
                                u.log.Warn("fanout: writer did not finish draining, cancelling its stream I/O",
                                        "sub_id", u.sub.ID)
                                u.cancelIO()
                        }
                        for _, u := range ws[i:] {
                                <-u.done
                        }
                        return
                }
        }
}

package relay

import (
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "math"
        "slices"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/relay/cache"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// defaultUpstreamFetchTimeout bounds an upstream stitch FETCH when the
// downstream supplied no FILL_TIMEOUT. It keeps a fetch-capable upstream that
// nonetheless stalls (or never answers FETCH) from wedging the downstream
// handler: once it elapses, the cache is served with the unknown Locations
// marked Timed-Out.
const defaultUpstreamFetchTimeout = 5 * time.Second

// trackKnown reports whether entry stands for a track the relay actually knows
// of. Bare existence does not say so: subscribeUpstreamOnSession creates the
// entry before the upstream round trip that would confirm the track, because
// it must be in place before the §11.1 Track Alias in SUBSCRIBE_OK can route
// (#85). Between those two points the entry describes a track nobody has
// vouched for yet.
//
// The distinction is visible on the wire. Answering a FETCH from such an entry
// falls through to the §10.13 "no Objects have been published" rule and
// returns INVALID_RANGE — "the range you asked for cannot be satisfied" —
// where §10.6 DOES_NOT_EXIST, "the track or namespace is not available at the
// publisher", is the truthful answer. A client deciding whether to retry, and
// with what, needs them kept apart.
//
// A watermark means a publisher has vouched for the track even if the
// subscription that carried it has since gone; a registered subscription means
// one is vouching for it now.
func trackKnown(entry *registry.TrackEntry) bool {
        if _, ok := entry.GetLargest(); ok {
                return true
        }
        // FETCH is not a hot path (see GetRange), so the copies are fine.
        return len(entry.CopyUpstream()) > 0 || len(entry.CopyDownstream()) > 0
}

// handleFetch implements FETCH (§10.13): validate the requested range, reply
// FETCH_OK, open a FETCH_HEADER uni-stream, and serialise the cached objects
// in the requested group order. Gaps in the response stream are how the spec
// signals "objects do not exist" (§10.13), so what the cache cannot vouch for
// is asked of an upstream FETCH when one is reachable, or covered by §11.4.4.2
// End of Range markers; see [sessionHandler.stitchedFetchObjects].
func (h *sessionHandler) handleFetch(ctx context.Context, req *session.Request, msg *message.Fetch) {
        if err := h.auth.AuthorizeFetch(ctx, h.sess, msg); err != nil {
                h.rejectAuth(ctx, req, "Fetch", err)
                return
        }

        fullName := track.FullTrackName{Namespace: msg.Namespace, Name: msg.Name}
        entry, ok := h.tracks.Get(fullName.Key())
        if !ok || !trackKnown(entry) {
                _ = req.RejectError(moqt.RequestDoesNotExist, "relay: track not known")
                return
        }

        largest, hasLargest := entry.GetLargest()
        if !hasLargest {
                // §10.13: "If no Objects have been published for the track or Start
                // Location is greater than the Largest Object the publisher MUST
                // return REQUEST_ERROR with error code INVALID_RANGE."
                _ = req.RejectError(moqt.RequestInvalidRange, "relay: no objects published")
                return
        }

        // draft-20 moved the FETCH range out of the message and into the
        // LOCATION_FILTER parameter (§5.1.2), inclusive at both ends. An absent
        // filter fetches the whole track up to Largest Object. AcceptRequest has
        // validated it (see [message.Parameters.CheckScope]).
        filter, _ := message.LocationFilterFromParam(msg.Parameters)
        if filter == nil {
                filter = &message.LocationFilter{}
        }
        start := filter.Start(largest, hasLargest)

        // §10.13: Start > Largest is INVALID_RANGE.
        if largest.Less(start) {
                _ = req.RejectError(moqt.RequestInvalidRange, "relay: start beyond largest object")
                return
        }

        // A 4-field filter can name an end below its own start (EndGroupDelta 0 with
        // EndObject < StartObject), which §5.1.2 does not itself forbid. Answering it
        // would put us in violation of §10.14 — "If End Location is smaller than the
        // Start Location in the corresponding FETCH the receiver MUST close the
        // session with a PROTOCOL_VIOLATION" — so one malformed FETCH would tear down
        // every other subscription on the session. Reject the request instead.
        if end, ok := filter.End(); ok && end.Less(start) {
                _ = req.RejectError(moqt.RequestInvalidRange, "relay: end before start")
                return
        }

        order := fetchGroupOrder(msg.Parameters)
        fillTimeout := resolveFillBudget(msg.Parameters)
        rangeFilters, ok := h.fetchRangeFilters(ctx, req, msg.Parameters)
        if !ok {
                return
        }

        // The response EndLocation is fixed by the watermark (§10.14) and is
        // independent of which objects we end up streaming, so reply FETCH_OK
        // before doing any (possibly slow) upstream stitching.
        endLocation := capFetchEndLocation(filter, largest)
        // §10.14: End Of Track is "1 if all Objects have been published on this
        // Track, and the End Location is the final Object in the Track". The
        // END_OF_TRACK Object raised the watermark, so a FETCH running past it
        // ends there.
        trackEnd, ended := entry.TrackEnd()
        var properties []byte
        if includeProperties(msg.Parameters) { // §10.2.21
                properties = entry.GetProperties()
        }
        if err := req.Reply(&message.FetchOK{
                EndOfTrack:      ended && endLocation == trackEnd,
                EndLocation:     endLocation,
                TrackProperties: properties,
        }); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "FETCH_OK reply failed",
                        slog.String("err", err.Error()))
                return
        }

        // Serve (and account for) only the range FETCH_OK announced: everything
        // past the capped EndLocation is outside the response by definition, so
        // neither objects nor §11.4.4.2 unknown markers may reference it.
        h.serveFetchObjects(ctx, req, "fetch", msg.RequestID, entry, fullName,
                start, endLocation, order, fillTimeout, rangeFilters)
}

// resolveFillBudget reads FILL_TIMEOUT (§10.2.5) and resolves the "absent"
// case to the local default, so downstream a zero means only what §10.2.5 says
// it means: do not wait for upstream at all.
func resolveFillBudget(ps message.Parameters) time.Duration {
        if d, ok := message.FillTimeoutFromParamOK(ps); ok {
                return d
        }
        return defaultUpstreamFetchTimeout
}

// fetchRangeFilters parses and validates the §5.1.4 Range Filters on a FETCH's
// parameters against the negotiated MAX_FILTER_RANGES. On an invalid or
// over-limit filter it answers REQUEST_ERROR INVALID_FILTER (§10.6) and returns
// ok=false, so the caller aborts before replying FETCH_OK.
func (h *sessionHandler) fetchRangeFilters(
        ctx context.Context, req *session.Request, ps message.Parameters,
) (*message.RangeFilterSet, bool) {
        rf, err := message.RangeFiltersFromParams(ps)
        if err == nil && rf != nil {
                err = rf.Validate(h.sess.MaxFilterRanges())
        }
        if err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "FETCH range filter rejected",
                        slog.String("err", err.Error()))
                _ = req.RejectError(moqt.RequestInvalidFilter, err.Error())
                return nil, false
        }
        return rf, true
}

// readFetchUpdates is the follow-up dispatch loop for an established FETCH
// whose data stream is out: REQUEST_UPDATE (§10.9) routes to
// [sessionHandler.handleFetchUpdate]; any other follow-up is ignored. On the
// requester's FIN the relay FINs back (§3.3.2).
func (h *sessionHandler) readFetchUpdates(ctx context.Context, req *session.Request, out *session.OutgoingFetchStream) {
        updates := h.sess.NewRequestUpdateLimiter()
        fin := readRequestStream(ctx, h.sess, req.Stream, func(m message.Message) bool {
                if h.isPeerStateNotify(m) {
                        return false
                }
                upd, ok := m.(*message.RequestUpdate)
                if !ok {
                        return true
                }
                // §10.2.1: out-of-scope parameters are session-fatal.
                if h.sess.CheckPeerParams(message.ScopeUpdateFetch, upd) != nil {
                        return false
                }
                // §10.1: the update consumes a Request ID; a parity or duplicate
                // violation is session-fatal.
                if !h.handleFollowupRequestID(ctx, upd) {
                        return false
                }
                // §10.3.1.7: enforce the per-stream MAX_REQUEST_UPDATES limit.
                if !h.handleRequestUpdateLimit(ctx, updates) {
                        return false
                }
                // §10.2.2: an update may REGISTER/DELETE token aliases; a cache
                // fault there is session-fatal.
                toks, ok := h.handleFollowupTokens(ctx, upd)
                if !ok || !h.handleFetchUpdate(ctx, req, out, toks) {
                        return false
                }
                updates.Responded()
                return true
        })
        if fin {
                _ = req.Stream.Close()
        }
}

// handleFetchUpdate answers a REQUEST_UPDATE (§10.9) to an in-flight FETCH
// with REQUEST_OK: the in-scope parameters have nothing to change on a
// finished snapshot. An update whose tokens (toks) the TokenVerifier denies
// fails, and false is returned: the FETCH is over.
func (h *sessionHandler) handleFetchUpdate(
        ctx context.Context,
        req *session.Request,
        out *session.OutgoingFetchStream,
        toks []session.ResolvedToken,
) bool {
        if rej := h.refuseUpdateTokens(ctx, toks); rej != nil {
                _ = req.RejectError(rej.ErrorCode, rej.ErrorReason)
                // §10.9.1: "When a REQUEST_UPDATE fails for a FETCH, the publisher
                // MUST reset the FETCH data stream." It was FINed already, so this
                // only aborts delivery of what the requester has not acknowledged.
                code := moqt.StreamResetCancelled
                if rej.ErrorCode == moqt.RequestExpiredAuthToken {
                        code = moqt.StreamResetExpiredAuthToken // §3.3.4
                }
                out.Cancel(code)
                return false
        }
        if err := req.Reply(&message.RequestOK{}); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "FETCH REQUEST_UPDATE_OK write failed",
                        slog.String("err", err.Error()))
        }
        return true
}

// fetchGroupOrder is a FETCH's GROUP_ORDER (§10.2.8): Ascending when omitted.
// AcceptRequest has closed the session on a value outside {1, 2}.
func fetchGroupOrder(ps message.Parameters) message.GroupOrder {
        if p, ok := ps.Find(message.ParamGroupOrder); ok {
                return message.GroupOrder(p.Byte)
        }
        return message.GroupOrderAscending
}

// capFetchEndLocation resolves a FETCH's end from its Location filter and
// caps it at Largest Object per §10.14: "This is the End Location from the
// FETCH request Location Filter parameter unless the requested range extends
// beyond Largest Object at the time the request was processed."
//
// draft-20 made both the request range and FETCH_OK's End Location inclusive
// (§5.1.2), so — unlike draft-19's "last Object plus 1, or 0 for the whole
// group" encoding — no exclusive/inclusive conversion is involved.
func capFetchEndLocation(filter *message.LocationFilter, largest message.Location) message.Location {
        end, ok := filter.End()
        if !ok || largest.Less(end) {
                // §5.1.2: "When they are omitted from a Fetch, the EndGroup and
                // EndObject are Largest Object."
                return largest
        }
        return end
}

// stitchedFetchObjects answers a FETCH range [start, end] from the relay's
// cache, asking an upstream about the Locations whose status it does not know
// (§10.13: "If it encounters an object in the requested range that is not
// cached and has unknown status, the relay MUST pause subsequent delivery
// until it has confirmed the object's status upstream"). See fetch_ranges.go
// for what the relay knows.
//
// With a fetch-capable upstream, one FETCH covers the span from the first
// unknown Location to the last, within the FILL_TIMEOUT budget (§10.2.5). Its
// Objects fill the holes, and what it marks unknown or timed out stays so where
// the relay does not know better; the cached Objects are served either way,
// and what the relay knows does not exist stays a gap. With no such upstream, or
// when its FETCH fails or times out, the unknown Locations are marked End of
// Unknown or Timed-Out Range (§11.4.4.2) and the cached Objects served: the
// relay can "indicate the range of unknown Objects and continue serving other
// known Objects" (§10.13). Upstream-fetched objects are NOT cached back: the
// FIFO ring is keyed by arrival, so old backfill would evict live objects.
//
// A non-nil refusal (see fetchUpstreamRange) means the track must not be
// forwarded; no objects are returned.
func (h *sessionHandler) stitchedFetchObjects(
        ctx context.Context,
        entry *registry.TrackEntry,
        fullName track.FullTrackName,
        start, end message.Location,
        order message.GroupOrder,
        fillTimeout time.Duration,
) (objs []*cache.CachedObject, refusal error) {
        // An expired Object (§12.3) is not returned: its status is unknown.
        cached := entry.Cache.GetRange(start, end, message.GroupOrderAscending)
        unknown := unknownIn(entry, cached, start, end)
        if len(unknown) == 0 {
                return fetchElements(cached, nil, nil, order), nil
        }
        up := h.pickFetchUpstream(entry)
        if up == nil {
                return fetchElements(cached, unknown, nil, order), nil
        }
        span := registry.LocRange{Lo: unknown[0].Lo, Hi: unknown[len(unknown)-1].Hi}
        done := h.tracks.BeginRequest(message.TypeFetch, up.Session, fullName.Key())
        ans, refusal := h.fetchUpstreamRange(ctx, up, fullName, span, order, fillTimeout)
        done()
        if errors.Is(refusal, session.ErrMalformedTrack) {
                h.endMalformedTrack(ctx, entry, up.Session, refusal)
        }
        if refusal != nil {
                return nil, refusal
        }
        switch ans.failed {
        case upstreamUnknown:
                return fetchElements(cached, unknown, nil, order), nil
        case upstreamTimedOut:
                return fetchElements(cached, nil, unknown, order), nil
        case upstreamAnswered:
        }

        // The upstream answered for the span: its Objects fill the holes, and
        // under its FIN the rest does not exist, except what it marked unknown or
        // timed out and the relay has no signal for either.
        have := make(map[message.Location]bool, len(cached))
        for _, o := range cached {
                have[message.Location{Group: o.GroupID, Object: o.ObjectID}] = true
        }
        merged := cached
        for _, o := range ans.objs {
                if !have[message.Location{Group: o.GroupID, Object: o.ObjectID}] {
                        merged = append(merged, o)
                }
        }
        return fetchElements(merged, intersect(ans.unknown, unknown), intersect(ans.timedOut, unknown), order), nil
}

// intersect returns the Locations both a and b hold, each a set of disjoint
// ranges.
func intersect(a, b []registry.LocRange) []registry.LocRange {
        var out []registry.LocRange
        for _, x := range a {
                for _, y := range b {
                        lo, hi := x.Lo, x.Hi
                        if lo.Less(y.Lo) {
                                lo = y.Lo
                        }
                        if y.Hi.Less(hi) {
                                hi = y.Hi
                        }
                        if !hi.Less(lo) {
                                out = append(out, registry.LocRange{Lo: lo, Hi: hi})
                        }
                }
        }
        return out
}

// pickFetchUpstream returns an Established, fetch-capable upstream the relay
// can issue a stitch FETCH to, or nil.
//
// Only upstreams the relay reached via an on-demand SUBSCRIBE (a relay/origin,
// marked FetchCapable in subscribeUpstream) are eligible: a directly-connected
// leaf publisher pushes live objects and is not expected to answer FETCH, so
// stitching to it would only stall. The requester's own session is eligible
// like any other, as FETCH follows SUBSCRIBE's matching rules (§9.5) and a
// self-subscription is "identical" to any other (§5.1). Deviation: not while
// a stitch FETCH for the track to it is in flight, since this request may be
// that FETCH routed back, and a second one would loop (§6.2); a concurrent
// FETCH looks the same. Either way the hole is marked unknown (§10.13).
func (h *sessionHandler) pickFetchUpstream(entry *registry.TrackEntry) *registry.UpstreamSub {
        for _, u := range entry.CopyUpstream() {
                if !u.FetchCapable || !u.IsEstablished() || u.Session == nil || goingAway(u.Session) {
                        continue
                }
                if u.Session == h.sess && h.tracks.RequestPending(message.TypeFetch, u.Session, entry.FullName.Key()) {
                        continue
                }
                return u
        }
        return nil
}

// upstreamFailure is how an upstream FETCH failed to answer at all.
type upstreamFailure uint8

const (
        upstreamAnswered upstreamFailure = iota
        // upstreamUnknown: refused, reset, malformed or out of order; nothing it
        // sent is vouched for.
        upstreamUnknown
        // upstreamTimedOut: the FILL_TIMEOUT budget ran out (§10.2.5).
        upstreamTimedOut
)

// upstreamAnswer is an upstream FETCH response for a span, as Location ranges:
// its Objects; the parts its End of Unknown and Timed-Out Range markers
// covered, and any past a capped FETCH_OK End Location; every other Location
// of the span, a gap under a clean FIN, is known not to exist (§10.13).
type upstreamAnswer struct {
        objs              []*cache.CachedObject
        unknown, timedOut []registry.LocRange
        failed            upstreamFailure
}

// fetchUpstreamRange issues a standalone FETCH for span on the upstream's
// session, awaits the response stream via the relay's fetch router, and reads
// it into an upstreamAnswer. End of Non-Existent Range markers need no
// record: under a clean FIN a gap already says so.
//
// It returns a refusal instead when the track MUST NOT be forwarded: a
// FETCH_OK with unacceptable Track Properties (§2.5.1), or a response Object
// that makes the track malformed (§2.4.2, wrapping [session.ErrMalformedTrack]).
func (h *sessionHandler) fetchUpstreamRange(
        ctx context.Context,
        up *registry.UpstreamSub,
        fullName track.FullTrackName,
        span registry.LocRange,
        order message.GroupOrder,
        fillTimeout time.Duration,
) (ans upstreamAnswer, refusal error) {
        // §10.2.5: an explicit 0 means "MUST NOT wait for upstream delivery"
        // (fillTimeout is already resolved, see [resolveFillBudget]).
        if fillTimeout == 0 {
                return upstreamAnswer{failed: upstreamTimedOut}, nil
        }

        params := message.Parameters{}
        if order == message.GroupOrderDescending {
                params = append(params, message.GroupOrderParam(message.GroupOrderDescending))
        }
        // §5.1.2: the range rides in LOCATION_FILTER. EndGroupDelta is delta-encoded
        // from the start group, and EndObject makes the end Object-precise.
        params = append(params, message.AbsoluteRangeObjectFilter(
                span.Lo, span.Hi.Group-span.Lo.Group, span.Hi.Object))
        fmsg := &message.Fetch{
                Namespace:  fullName.Namespace,
                Name:       fullName.Name,
                Parameters: params,
        }

        // Bound the upstream round-trip so a silent or non-FETCH-answering
        // upstream degrades to cache-plus-marked-unknown instead of wedging the
        // downstream handler. FILL_TIMEOUT, when present, is the subscriber's
        // explicit budget; otherwise fall back to a default.
        fctx, cancel := context.WithTimeout(ctx, fillTimeout)
        defer cancel()

        fr, err := up.Session.Fetch(fctx, fmsg)
        if err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "upstream FETCH failed",
                        slog.String("err", err.Error()))
                // §2.5.1: Session.Fetch has cancelled it; the caller resets the
                // downstream stream.
                if isTrackPropertiesErr(err) {
                        return upstreamAnswer{}, err
                }
                if fctx.Err() != nil {
                        return upstreamAnswer{failed: upstreamTimedOut}, nil
                }
                return upstreamAnswer{failed: upstreamUnknown}, nil
        }
        defer fr.Close()

        // The upstream echoes our Request ID in the response's FETCH_HEADER, so
        // the body stream lands on the upstream session's data loop keyed by
        // fmsg.RequestID. Register after Fetch (the ID is only assigned there);
        // the router tolerates a response that races ahead of registration.
        ch, cleanup := h.fetch.Register(up.Session, fmsg.RequestID)
        defer cleanup()

        var fs *session.IncomingFetchStream
        select {
        case fs = <-ch:
        case <-fctx.Done():
                h.log.LogAttrs(ctx, slog.LevelDebug, "upstream FETCH response timed out")
                return upstreamAnswer{failed: upstreamTimedOut}, nil
        }
        if fs == nil {
                return upstreamAnswer{failed: upstreamUnknown}, nil
        }
        // ReadDecoded needs the response's group order to resolve cross-group
        // deltas (§11.4.4.1); the upstream serves in the order our FETCH asked
        // for.
        fs.GroupOrder = order
        // §10.2.5: the budget covers the response too; when it runs out, what
        // has arrived is kept and the rest reported Timed-Out. When the track is
        // found malformed, up is cancelled with MALFORMED_TRACK if it sent the
        // Object (§2.4.2; see endMalformedTrack), and the caller resets the
        // downstream stream.
        defer context.AfterFunc(fctx, func() {
                code := moqt.StreamResetCancelled
                if mt, ok := errors.AsType[*malformedTrackCause](context.Cause(fctx)); ok && mt.src == up.Session {
                        code = moqt.StreamResetMalformedTrack
                }
                fs.Cancel(code)
        })()
        // §12.3: this FETCH's MAX_CACHE_DURATION bounds each Object it delivers
        // (see [cache.ObjectCache.Expired]). Deviation: the age runs from when the
        // relay read the Object whole, not from "the beginning of the Object".
        maxAge, hasMaxAge := message.TrackMaxCacheDuration(fr.OK.TrackProperties)

        var prev *message.Location
        for {
                obj, err := fs.ReadDecoded()
                if errors.Is(err, io.EOF) {
                        break // clean FIN: the upstream's gaps are authoritative (§10.13)
                }
                if errors.Is(err, session.ErrMalformedTrack) {
                        // §2.4.2: fr.Close (deferred) cancels the fetch; the caller
                        // resets the downstream stream.
                        fs.Cancel(moqt.StreamResetMalformedTrack)
                        return upstreamAnswer{}, err
                }
                if err != nil && fctx.Err() != nil {
                        // Without a FIN its gaps assert nothing (§10.13): all of the span
                        // it did not send or mark is Timed-Out.
                        h.log.LogAttrs(ctx, slog.LevelDebug, "upstream FETCH response timed out mid-read")
                        known := slices.Concat(ans.unknown, ans.timedOut)
                        for _, o := range ans.objs {
                                loc := message.Location{Group: o.GroupID, Object: o.ObjectID}
                                known = append(known, registry.LocRange{Lo: loc, Hi: loc})
                        }
                        ans.timedOut = append(ans.timedOut, uncovered(span.Lo, span.Hi, known)...)
                        return ans, nil
                }
                if err != nil {
                        // No FIN (or a FIN mid-object), so the gaps in what arrived
                        // assert nothing.
                        h.log.LogAttrs(ctx, slog.LevelDebug, "upstream FETCH stream failed mid-read",
                                slog.String("err", err.Error()))
                        return upstreamAnswer{failed: upstreamUnknown}, nil
                }
                loc := message.Location{Group: obj.GroupID, Object: obj.ObjectID}
                // §10.14: nothing past its own End Location.
                if !upstreamFetchElemOK(loc, prev, span, order) || fr.OK.EndLocation.Less(loc) {
                        h.log.LogAttrs(ctx, slog.LevelDebug, "upstream FETCH element out of range or order",
                                slog.Uint64("group", loc.Group), slog.Uint64("object", loc.Object))
                        return upstreamAnswer{failed: upstreamUnknown}, nil
                }
                switch {
                case obj.EndOfUnknownRange:
                        ans.unknown = append(ans.unknown, streamCovered(prev, loc, span, order)...)
                case obj.EndOfTimedOutRange:
                        ans.timedOut = append(ans.timedOut, streamCovered(prev, loc, span, order)...)
                case obj.EndOfNonExistentRange:
                default:
                        // The §11.4.4.1 Datagram bit carries the original wire shape
                        // across this relay hop.
                        pref := cache.ForwardingSubgroup
                        if obj.Datagram {
                                pref = cache.ForwardingDatagram
                        }
                        ans.objs = append(ans.objs, &cache.CachedObject{
                                GroupID:             obj.GroupID,
                                ObjectID:            obj.ObjectID,
                                SubgroupID:          obj.SubgroupID,
                                PublisherPriority:   obj.PublisherPriority,
                                ForwardingPref:      pref,
                                Properties:          obj.Properties,
                                Payload:             obj.Payload,
                                ReceivedAt:          time.Now(),
                                MaxCacheDuration:    maxAge,
                                HasMaxCacheDuration: hasMaxAge,
                                Stitched:            true,
                        })
                }
                prev = &loc
        }

        // A clean FIN asserts gaps only up to the FETCH_OK End Location (§10.13).
        // If the upstream capped it below the span (§10.13: End beyond its
        // Largest), what lies past it has unknown status.
        // No element lay past it, and Session.Fetch refused one before the span.
        if authEnd := fr.OK.EndLocation; authEnd.Less(span.Hi) {
                next, _ := locSucc(authEnd) // below span.Hi, so it has one
                ans.unknown = append(ans.unknown, registry.LocRange{Lo: next, Hi: span.Hi})
        }
        return ans, nil
}

// upstreamFetchElemOK validates one element of an upstream FETCH response for
// span before it is re-served downstream: it lies inside span, and after the
// previous element prev (nil for the first) in the order the response carries
// them (see [streamCompare]), as §11.4.4's delta encoding requires. A
// violation means the upstream is nonconformant; trusting the element would
// corrupt the downstream stream, so the caller discards the response.
func upstreamFetchElemOK(
        loc message.Location,
        prev *message.Location,
        span registry.LocRange,
        order message.GroupOrder,
) bool {
        if loc.Less(span.Lo) || span.Hi.Less(loc) {
                return false
        }
        return prev == nil || streamCompare(*prev, loc, order) < 0
}

// streamCovered returns, as Location ranges, what an End of Range marker at at
// covers in a response to a FETCH of span in order: the Locations after the
// previous element prev (nil for the first) up to at, in the order the
// response carries them (see fetch_ranges.go). at is after prev.
func streamCovered(
        prev *message.Location,
        at message.Location,
        span registry.LocRange,
        order message.GroupOrder,
) []registry.LocRange {
        if order != message.GroupOrderDescending {
                from := span.Lo
                if prev != nil {
                        from, _ = locSucc(*prev) // at is after prev, so it has one
                }
                return []registry.LocRange{{Lo: from, Hi: at}}
        }
        // Descending: Group g of span carries Objects lo(g) through hi(g).
        lo := func(g uint64) uint64 {
                if g == span.Lo.Group {
                        return span.Lo.Object
                }
                return 0
        }
        hi := func(g uint64) uint64 {
                if g == span.Hi.Group {
                        return span.Hi.Object
                }
                return math.MaxUint64
        }
        var from message.Location
        switch {
        case prev == nil:
                from = message.Location{Group: span.Hi.Group, Object: lo(span.Hi.Group)}
        case prev.Object < hi(prev.Group):
                from = message.Location{Group: prev.Group, Object: prev.Object + 1}
        default:
                from = message.Location{Group: prev.Group - 1, Object: lo(prev.Group - 1)}
        }
        if from.Group == at.Group {
                return []registry.LocRange{{Lo: from, Hi: at}}
        }
        out := []registry.LocRange{{Lo: message.Location{Group: at.Group, Object: lo(at.Group)}, Hi: at}}
        if from.Group-at.Group > 1 {
                out = append(out, registry.LocRange{
                        Lo: message.Location{Group: at.Group + 1},
                        Hi: message.Location{Group: from.Group - 1, Object: math.MaxUint64},
                })
        }
        return append(out, registry.LocRange{Lo: from, Hi: message.Location{Group: from.Group, Object: hi(from.Group)}})
}

// unknownRangeMarker returns the serve-path element that streamFetchObjects
// serializes as a §11.4.4.2 End of Unknown Range (0x10C) marker at loc.
func unknownRangeMarker(loc message.Location) *cache.CachedObject {
        return &cache.CachedObject{
                GroupID:           loc.Group,
                ObjectID:          loc.Object,
                EndOfUnknownRange: true,
        }
}

// timedOutRangeMarker is [unknownRangeMarker] for the §10.2.5 case: the
// FILL_TIMEOUT budget ran out, so the Objects are reported as Timed-Out rather
// than unknown-status gaps.
func timedOutRangeMarker(loc message.Location) *cache.CachedObject {
        return &cache.CachedObject{
                GroupID:            loc.Group,
                ObjectID:           loc.Object,
                EndOfTimedOutRange: true,
        }
}

// fetchPredecessor returns the Location immediately below loc in (group,
// object) order, and false when loc is {0, 0} (nothing precedes it). The
// object-underflow case rolls back to the end of the previous group.
func fetchPredecessor(loc message.Location) (message.Location, bool) {
        switch {
        case loc.Object > 0:
                return message.Location{Group: loc.Group, Object: loc.Object - 1}, true
        case loc.Group > 0:
                return message.Location{Group: loc.Group - 1, Object: math.MaxUint64}, true
        default:
                return message.Location{}, false
        }
}

// streamFetchObjects writes the cached objects to the FETCH response
// stream with §11.4.4 delta encoding:
//
//   - The first object includes both GroupIDDelta and ObjectIDDelta
//     flags; the values are absolute (§11.4.4.1).
//   - Subsequent objects in the same group omit ObjectIDDelta when
//     consecutive (the subscriber reconstructs ObjectID = prior + 1);
//     otherwise ObjectIDDelta = ObjectID - prior, with no +1 unlike the
//     §11.4.2 subgroup rule (§11.4.4.1).
//   - Subsequent objects in a different group set GroupIDDelta:
//     ascending → newGroup - priorGroup - 1, descending →
//     priorGroup - newGroup - 1 (§11.4.4.1). ObjectIDDelta is then the
//     absolute Object ID in the new group.
//   - Datagram-flavoured objects set bit 0x40 (§11.4.4.1); subscriber
//     ignores the subgroup bits.
//   - [cache.CachedObject.EndOfUnknownRange] elements serialize as §11.4.4.2
//     End of Unknown Range markers (0x10C) with absolute Group/Object IDs,
//     and become the prior Location for the delta encoding of what follows.
//
// The returned count is the number of real objects written (markers are
// serialized but not counted — they carry no payload).
func streamFetchObjects(
        out *session.OutgoingFetchStream,
        objs []*cache.CachedObject,
        expired func(*cache.CachedObject) bool,
) (int, error) {
        var (
                written      int
                prevGroup    uint64
                prevObject   uint64
                prevPriority uint8
                // havePrev: a prior Group/Object ID exists — a real object or a
                // §11.4.4.2 End-of-Range marker. haveActual: a real object was
                // written — only then do a prior Subgroup ID / Priority exist
                // (mirror of ReadDecoded's decHavePrev / decHaveActual).
                havePrev   bool
                haveActual bool
                // Inferred from the ordering of the first vs second object.
                // Without a second object we don't need the direction.
                descending bool
        )

        for _, o := range objs {
                // §12.3: a slow reader can hold the stream until a cached Object
                // expires; mark it unknown, since a gap asserts non-existence.
                if !o.IsRangeMarker() && !o.IsStatusMarker() && expired != nil && expired(o) {
                        o = &cache.CachedObject{GroupID: o.GroupID, ObjectID: o.ObjectID, EndOfUnknownRange: true}
                }
                if o.IsRangeMarker() {
                        // §11.4.4.2 End of Unknown / Timed-Out Range: the Group/Object ID fields
                        // carry the absolute range boundary, and the marker becomes
                        // the prior Location for subsequent delta encoding — but not
                        // a prior *actual* object, so the next object still spells
                        // out its Priority (and never references the prior Subgroup).
                        flags := uint64(message.FetchEndOfUnknownRange)
                        if o.EndOfTimedOutRange {
                                flags = message.FetchEndOfTimedOutRange
                        }
                        if err := out.WriteObject(&message.FetchObject{
                                SerializationFlags: flags,
                                GroupIDDelta:       o.GroupID,
                                ObjectIDDelta:      o.ObjectID,
                        }); err != nil {
                                return written, err
                        }
                        prevGroup, prevObject = o.GroupID, o.ObjectID
                        havePrev = true
                        continue
                }

                // §11.2.1.1: the Object Status field "is absent in Objects
                // delivered via a FETCH". Cached status markers describe absence,
                // so they are simply not serialized — their knowledge still reaches
                // the fetcher: the marker bumped the LARGEST_OBJECT watermark on
                // ingest, FETCH_OK's EndLocation extends through it
                // (capFetchEndLocation), and §11.4.4's gap rule makes the trailing
                // gap of a FIN-terminated response authoritative non-existence.
                // Emitting End of Non-Existent Range (0x8C) instead would be
                // redundant: §11.4.4.2 reserves it for splitting non-serialized
                // ranges into known-non-existent and unknown parts.
                if o.IsStatusMarker() {
                        continue
                }

                fo := &message.FetchObject{}

                switch {
                case !havePrev:
                        // §11.4.4.1: first object MUST include both
                        // GroupIDDelta and ObjectIDDelta flags; values are
                        // absolute.
                        fo.SerializationFlags |= message.FetchFlagGroupIDDelta | message.FetchFlagObjectIDDelta
                        fo.GroupIDDelta = o.GroupID
                        fo.ObjectIDDelta = o.ObjectID

                case o.GroupID != prevGroup:
                        // Cross-group. Detect direction from the first such
                        // transition: descending iff new group < prior.
                        if !descending && o.GroupID < prevGroup {
                                descending = true
                        } else if descending && o.GroupID > prevGroup {
                                // Direction reversed mid-stream — should
                                // never happen because GetRange returns
                                // stably sorted output, but if it did the
                                // safest action is to abandon the optimised
                                // delta encoding and reset the GroupIDDelta
                                // using ascending convention.
                                descending = false
                        }
                        fo.SerializationFlags |= message.FetchFlagGroupIDDelta | message.FetchFlagObjectIDDelta
                        if descending {
                                fo.GroupIDDelta = prevGroup - o.GroupID - 1
                        } else {
                                fo.GroupIDDelta = o.GroupID - prevGroup - 1
                        }
                        fo.ObjectIDDelta = o.ObjectID

                default:
                        // Same group. §11.4.4 cannot express a non-ascending Object ID
                        // here — the delta only ever adds. The inputs are sorted in
                        // stream order (fetchElements), so hitting this is an
                        // internal invariant violation; fail rather than emit a wrapped
                        // delta the subscriber must treat as a session-fatal overflow.
                        if o.ObjectID <= prevObject {
                                return written, fmt.Errorf(
                                        "relay: fetch serialization order violation: {%d,%d} after {%d,%d}",
                                        o.GroupID, o.ObjectID, prevGroup, prevObject)
                        }
                        // §11.4.4.1: no +1, unlike the §11.4.2 subgroup rule.
                        if o.ObjectID != prevObject+1 {
                                fo.SerializationFlags |= message.FetchFlagObjectIDDelta
                                fo.ObjectIDDelta = o.ObjectID - prevObject
                        }
                }

                switch o.ForwardingPref {
                case cache.ForwardingDatagram:
                        // §11.4.4.1: bit 0x40 marks the object as a
                        // Datagram-flavoured object; the subscriber ignores
                        // the two subgroup bits.
                        fo.SerializationFlags |= message.FetchFlagDatagram
                case cache.ForwardingSubgroup:
                        // Subgroup: encode the SubgroupID explicitly. The
                        // "prior + 0/1" subgroup modes are micro-optimisations
                        // over the explicit form; we always emit explicit for
                        // simplicity.
                        fo.SerializationFlags = (fo.SerializationFlags &^ message.FetchFlagSubgroupIDMode) |
                                uint64(message.FetchSubgroupIDExplicit)
                        fo.SubgroupID = o.SubgroupID
                }

                // Publisher priority: emit when it differs from the prior actual
                // object's — or when there is none (the first object, and the
                // first object after a leading marker, §11.4.4.2).
                if !haveActual || o.PublisherPriority != prevPriority {
                        fo.SerializationFlags |= message.FetchFlagPriority
                        fo.PublisherPriority = o.PublisherPriority
                }

                if len(o.Properties) > 0 {
                        fo.SerializationFlags |= message.FetchFlagProperties
                        fo.Properties = o.Properties
                }

                fo.ObjectPayload = o.Payload

                if err := out.WriteObject(fo); err != nil {
                        return written, err
                }
                written++

                prevGroup = o.GroupID
                prevObject = o.ObjectID
                prevPriority = o.PublisherPriority
                havePrev = true
                haveActual = true
        }

        return written, nil
}

package relay

import (
        "context"
        "errors"
        "log/slog"
        "slices"
        "sync/atomic"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// testHookBeforeFill, when set by a test, runs as a fill is about to be
// evaluated, after the response that reported its Largest Object.
var testHookBeforeFill atomic.Pointer[func(track.FullTrackName)]

// maybeServeFill opens and serves a fill fetch stream for a subscription when
// the SUBSCRIBE or REQUEST_UPDATE carried FILL_PARAMETERS (§5.1.3).
//
// requestID is the Request ID of the message that asked for the fill; the
// FETCH_HEADER carries it, so one subscription can have several fills open.
// largest is the Largest Object the subscriber was told and the live filter
// is anchored on (the one SUBSCRIBE_OK or REQUEST_UPDATE_OK reported), and
// hasLargest whether there was one: the fill ends there (§5.1.3), so an Object
// arriving since goes out live only. For a forwarded PUBLISH that is the
// registration snapshot, not the PUBLISH's own LARGEST_OBJECT, which can be
// older: ending there would leave the Objects between neither live nor filled.
//
// A failure resets the fill stream, leaves the subscription unaffected, and
// is returned for the log. AcceptRequest has closed the session on a malformed
// FILL_PARAMETERS (§10.2.15; see [message.Parameters.CheckScope]).
func (h *sessionHandler) maybeServeFill(
        ctx context.Context,
        sub *registry.DownstreamSub,
        entry *registry.TrackEntry,
        fullName track.FullTrackName,
        requestID uint64,
        ps message.Parameters,
        largest message.Location,
        hasLargest bool,
) error {
        inner, requested, _ := message.FillParametersFromParam(ps)
        if !requested {
                return nil
        }
        if hook := testHookBeforeFill.Load(); hook != nil {
                (*hook)(fullName)
        }

        // §5.1.3.1: from here a failure MUST open the fill stream and reset it.
        fail := func(err error) error {
                h.resetFillStream(ctx, sub, requestID)
                return err
        }

        // §5.1.3.1: only "while Forward State is 1"; a later unpause does not
        // open one retroactively.
        if sub.ForwardState() != 1 {
                return nil
        }

        // The fill range is evaluated with Fetch rules (§5.1.2), so it never
        // extends past the Largest Object reported. With nothing published then
        // there is nothing to fill.
        if !hasLargest {
                return nil
        }

        // §5.1.3: the fill range comes from the LOCATION_FILTER inside
        // FILL_PARAMETERS, falling back to the subscription's own filter, and to
        // the whole track when neither is present.
        filter, _ := message.LocationFilterFromParam(inner)
        if filter == nil {
                filter = sub.GetFilter()
        }
        if filter == nil {
                filter = &message.LocationFilter{}
        }

        start := filter.Start(largest, hasLargest)
        end := capFetchEndLocation(filter, largest)
        // §5.1.3: "If the fill range is empty, or starts after Largest Object, the
        // publisher does not open a fill fetch stream."
        if largest.Less(start) || end.Less(start) {
                return nil
        }

        // §10.2.15: a parameter omitted from FILL_PARAMETERS keeps the value it
        // has for the subscription, so the inner list only carries the overrides.
        order := message.GroupOrder(sub.GroupOrder)
        if p, ok := inner.Find(message.ParamGroupOrder); ok {
                order = message.GroupOrder(p.Byte)
        }
        fillTimeout := resolveFillBudget(inner)

        // §5.1.3: the fill "inherits the subscription's parameters". A filter
        // type inside FILL_PARAMETERS overrides that type as a REQUEST_UPDATE
        // would (§5.1.4); the other types are inherited.
        rangeFilters := sub.GetRangeFilters()
        if slices.ContainsFunc(inner, func(p message.Parameter) bool { return message.IsRangeFilterParam(p.Type) }) {
                var err error
                rangeFilters, err = rangeFilters.Update(inner)
                if err == nil && rangeFilters != nil {
                        err = rangeFilters.Validate(h.sess.MaxFilterRanges())
                }
                if err != nil {
                        return fail(err)
                }
        }

        h.relayGo(func() {
                // §5.1.3.1: "When the subscription is cancelled, the publisher MUST
                // reset any open fill fetch streams".
                fillCtx, cancel := context.WithCancelCause(ctx)
                defer cancel(nil)
                defer context.AfterFunc(sub.Cancelled(), func() { cancel(errRequestCancelled) })()
                // §2.4.2: a malformed track resets its fill fetch streams too.
                defer entry.AddFetch(cancel)()
                h.serveFill(fillCtx, sub, requestID, entry, fullName, start, end, order, fillTimeout, rangeFilters)
        })
        return nil
}

// serveFill writes one fill fetch stream; the FIN signals completion
// (§5.1.3.1), and [sessionHandler.streamFetchRange] resets it on a write error.
func (h *sessionHandler) serveFill(
        ctx context.Context,
        sub *registry.DownstreamSub,
        requestID uint64,
        entry *registry.TrackEntry,
        fullName track.FullTrackName,
        start, end message.Location,
        order message.GroupOrder,
        fillTimeout time.Duration,
        rangeFilters *message.RangeFilterSet,
) {
        h.log.LogAttrs(ctx, slog.LevelDebug, "serving fill fetch stream",
                slog.Uint64("request_id", requestID),
                slog.Uint64("start_group", start.Group),
                slog.Uint64("end_group", end.Group))

        h.streamFetchRange(ctx, "fill", sub, requestID, entry, fullName,
                start, end, order, fillTimeout, rangeFilters)
}

// resetFillStream signals a fill failure the only way §5.1.3.1 allows: open the
// fill fetch stream and reset it right after the FETCH_HEADER. Otherwise the
// subscriber cannot tell it from an empty fill range, which opens no stream.
func (h *sessionHandler) resetFillStream(ctx context.Context, sub *registry.DownstreamSub, requestID uint64) {
        out, err := openFillOrFetchStream(h.sess, sub, requestID)
        if err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "could not open fill stream to reset it",
                        slog.Uint64("request_id", requestID), slog.String("err", err.Error()))
                return
        }
        out.Cancel(moqt.StreamResetInternalError)
        sub.StreamClosed()
}

// errSubscriptionTerminated reports a subgroup or fill stream not opened
// because its subscription already ended: §10.12 forbids streams after
// PUBLISH_DONE.
var errSubscriptionTerminated = errors.New("relay: subscription terminated before the stream opened")

// openFillOrFetchStream opens a FETCH_HEADER stream. A fill fetch stream
// counts toward sub's §10.12 PUBLISH_DONE Stream Count; a standalone FETCH
// response passes a nil sub.
func openFillOrFetchStream(
        sess *session.Session,
        sub *registry.DownstreamSub,
        requestID uint64,
) (*session.OutgoingFetchStream, error) {
        if sub == nil {
                return sess.OpenFetchStream(message.FetchHeader{RequestID: requestID})
        }
        if !sub.BeginStream() {
                return nil, errSubscriptionTerminated
        }
        out, err := sess.OpenFetchStream(message.FetchHeader{RequestID: requestID})
        sub.EndStream(err == nil)
        return out, err
}

package relay

import (
        "context"
        "errors"
        "log/slog"
        "slices"
        "sync"
        "sync/atomic"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// testHookBeforeForwardClaim, when set by a test, runs as forwardTrack is
// about to claim the track, to hold a forward there while another one runs.
var testHookBeforeForwardClaim atomic.Pointer[func(track.FullTrackName)]

// forwardTrack is a SUBSCRIBE_TRACKS subscriber's [registry.SubscriberEntry.ForwardTrack]:
// it sends the subscriber a PUBLISH for te (§6.1) with the current
// SUBSCRIBE_TRACKS parameters (§10.20.1) and serves the resulting
// subscription. At most one PUBLISH per track, and none for a track the
// subscriber publishes, already receives, or is SUBSCRIBing to.
func (h *sessionHandler) forwardTrack(ctx context.Context) func(*registry.SubscriberEntry, *registry.TrackEntry) {
        return func(sub *registry.SubscriberEntry, te *registry.TrackEntry) {
                if goingAway(h.sess) {
                        return // §10.4: no new PUBLISH on a session with a GOAWAY
                }
                fullName := te.FullName
                if !fullName.Namespace.HasPrefix(sub.Prefix()) {
                        return // a TRACK_NAMESPACE_PREFIX update moved the subscription away
                }
                // §5.1.4: PUBLISHes that fail the Range Filters are not forwarded.
                tp := sub.TracksParams()
                if !tp.RangeFilters.MatchesTrack(te.GetProperties()) {
                        return
                }
                // §6.1: "excluding tracks published by the subscriber". Relay policy,
                // not §6.1: nor a track it receives on its own SUBSCRIBE. The in-flight
                // check comes before HasDownstreamOn, as the SUBSCRIBE registers its
                // downstream before it stops being in flight; and before the claim, so
                // the held forward's replay does not find the claim still taken.
                key := fullName.Key()
                if te.HasUpstreamOn(h.sess) || h.holdForward(key, sub) {
                        return
                }
                if hook := testHookBeforeForwardClaim.Load(); hook != nil {
                        (*hook)(fullName)
                }
                // The claim covers a forward whose downstream is not registered yet,
                // and is refused while a PUBLISH_SKIPPED holds for this upstream epoch.
                epoch := te.UpstreamEpoch()
                if !sub.ClaimForward(key, epoch) {
                        return
                }
                // Checked under the claim: a forward releases it only once its
                // downstream is registered, so whichever forward claims second sees
                // that downstream and sends no second PUBLISH.
                if te.HasDownstreamOn(h.sess) {
                        sub.ReleaseForward(key)
                        return
                }
                var properties []byte
                if includeProperties(tp.Params) { // §10.2.21
                        properties = te.GetProperties()
                }
                fwd := &message.Publish{
                        Namespace: fullName.Namespace,
                        Name:      fullName.Name,
                        // §11.1: aliases are per session.
                        TrackAlias:      h.sess.AllocOutboundTrackAlias(),
                        Parameters:      publishParamsForSubscriber(tp, te),
                        TrackProperties: properties,
                }
                // §6.1: without bidi-stream credit, send PUBLISH_SKIPPED instead.
                stream, err := h.sess.OpenPublish(fwd)
                if err != nil {
                        if errors.Is(err, session.ErrNoStreamCredit) {
                                h.emitPublishSkipped(ctx, sub, fullName, epoch)
                        } else {
                                h.log.LogAttrs(ctx, slog.LevelDebug, "PUBLISH forward failed", slog.String("err", err.Error()))
                        }
                        sub.ReleaseForward(key)
                        return
                }
                h.relayGo(func() {
                        h.serveForwardedPublish(ctx, stream, fwd, tp.Params, te, func() { sub.ReleaseForward(key) })
                })
        }
}

// serveForwardedPublish serves the subscription a forwarded PUBLISH opened,
// as a downstream on te registered before PUBLISH_OK, since objects may flow
// before it (§10.11). A REQUEST_ERROR ends it; otherwise it
// is served like a SUBSCRIBE's.
func (h *sessionHandler) serveForwardedPublish(
        ctx context.Context,
        stream session.Stream,
        fwd *message.Publish,
        params message.Parameters,
        te *registry.TrackEntry,
        registered func(),
) {
        fullName := te.FullName
        sub := registry.NewDownstreamSub(h.allocSubID(), h.sess, stream, fwd.TrackAlias)
        sub.OpenedByPublish()
        // handleSubscribeTracks already refused parameters this would reject.
        _ = installSubscribeParams(sub, params)
        // The Group Order the PUBLISH stated (see publishParamsForSubscriber), so
        // the subscription's fills follow what it was told (§10.20.1).
        if p, ok := fwd.Parameters.Find(message.ParamGroupOrder); ok {
                sub.SetGroupOrder(p.Byte)
        }
        _, largest, has, added := h.tracks.AddDownstreamSnapshotLargest(fullName, sub)
        registered()
        if !added {
                sub.TerminateWithPublishDone(moqt.PublishDoneTrackEnded, "relay: upstream gone")
                return
        }
        sub.SetLargestAtSubscribe(largest, has)
        defer h.tracks.RemoveDownstream(fullName, sub.ID)
        ref := h.trackRef(fullName)
        h.metrics.SubscriptionOpened(ref)
        defer h.metrics.SubscriptionClosed(ref)
        // §5.1.1: once the subscriber cancels, or the session ends, reset the
        // streams still open for it.
        defer sub.Cancel()
        // §9.2: a forwarding subscriber resumes a paused upstream.
        if sub.ForwardState() == 1 {
                h.propagateForwardUpstream(ctx, fullName)
        }

        if _, err := h.sess.AwaitPublishOK(ctx, stream); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "forwarded PUBLISH refused",
                        slog.String("name", string(fullName.Name)), slog.String("err", err.Error()))
                // §5.1, §5.1.1: the subscriber's REQUEST_ERROR terminates the
                // subscription, so no PUBLISH_DONE follows.
                sub.EndRefused()
                return
        }
        // §10.20.1: each forwarded subscription gets its own fill fetch stream,
        // named by the PUBLISH's Request ID (§10.1; §5.1.3 names only SUBSCRIBE
        // and REQUEST_UPDATE).
        // The registration snapshot, which the live filter is anchored on; see
        // maybeServeFill.
        if err := h.maybeServeFill(ctx, sub, te, fullName, fwd.RequestID, params, largest, has); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "fill fetch stream not opened",
                        slog.String("err", err.Error()))
        }
        // §10.2.19
        if p, ok := params.Find(message.ParamNewGroupRequest); ok {
                h.propagateNewGroupUpstream(ctx, fullName, p.Varint)
        }
        h.readSubscribeUpdates(ctx, stream, sub, fullName)
}

// inflightSubscribe is a track's SUBSCRIBEs in flight on one session.
type inflightSubscribe struct {
        n int
        // held are the SUBSCRIBE_TRACKS entries whose forward of the track was
        // held back meanwhile; see [sessionHandler.holdForward].
        held []*registry.SubscriberEntry
}

// beginSubscribe marks a SUBSCRIBE for key in flight on this session until
// end, which may run more than once. From before it establishes an upstream
// until its downstream is registered, the track can have an upstream and no
// downstream here, so [sessionHandler.forwardTrack] would otherwise offer the
// SUBSCRIBE_TRACKS holders on this session the track it is SUBSCRIBing to.
//
// When the last SUBSCRIBE for key ends, the forwards held back are offered
// again: forwardTrack declines them if a SUBSCRIBE registered its downstream,
// and sends them if none did, since the track is then received no other way
// (§10.20).
func (h *sessionHandler) beginSubscribe(key track.Key) (end func()) {
        h.subscribingMu.Lock()
        defer h.subscribingMu.Unlock()
        if h.subscribing == nil {
                h.subscribing = make(map[track.Key]*inflightSubscribe)
        }
        f := h.subscribing[key]
        if f == nil {
                f = &inflightSubscribe{}
                h.subscribing[key] = f
        }
        f.n++
        return sync.OnceFunc(func() {
                h.subscribingMu.Lock()
                f.n--
                var held []*registry.SubscriberEntry
                if f.n == 0 {
                        delete(h.subscribing, key)
                        held = f.held
                }
                h.subscribingMu.Unlock()
                if len(held) == 0 {
                        return
                }
                if te, ok := h.tracks.Get(key); ok && hasEstablishedUpstream(te) {
                        for _, sub := range held {
                                sub.ForwardTrack(sub, te)
                        }
                }
        })
}

// holdForward reports whether a SUBSCRIBE for key is in flight on this session
// (see [sessionHandler.beginSubscribe]), recording sub to be offered the track
// again once none is.
func (h *sessionHandler) holdForward(key track.Key, sub *registry.SubscriberEntry) bool {
        h.subscribingMu.Lock()
        defer h.subscribingMu.Unlock()
        f := h.subscribing[key]
        if f == nil {
                return false
        }
        if !slices.Contains(f.held, sub) {
                f.held = append(f.held, sub)
        }
        return true
}

package relay

import (
        "context"
        "log/slog"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// endMalformedTrack handles a malformed track (§2.4.2) detected in an Object
// src sent on entry's track: every downstream subscription ends with
// PUBLISH_DONE MALFORMED_TRACK, every downstream fetch stream is reset with
// MALFORMED_TRACK, and only the upstream subscriptions on src are cancelled,
// so a redundant publisher (§9.3) keeps serving. An upstream FETCH ends with
// the downstream fetch stream it fills: its data stream gets STOP_SENDING
// MALFORMED_TRACK when it is on src, else CANCELLED, and its request stream
// is cancelled with CANCELLED. Callers never cache the Object. Open subgroup
// streams are reset now, since PUBLISH_DONE waits for them (§10.12).
//
// Downstreams are terminated before the upstream is cancelled: the first
// termination wins, and the upstream's teardown would use its own code.
//
// Interpretation: the upstream subscription's cancel and the upstream FETCH
// data stream's STOP_SENDING also use MALFORMED_TRACK, which §3.3.4 defines
// for the downstream direction.
func (h *sessionHandler) endMalformedTrack(
        ctx context.Context,
        entry *registry.TrackEntry,
        src *session.Session,
        cause error,
) {
        downstream := entry.CopyDownstream()
        for _, sub := range downstream {
                sub.TerminateWithPublishDone(moqt.PublishDoneMalformedTrack, "relay: malformed track")
        }
        h.resetWriters(entry, downstream)
        // An upstream FETCH runs only under a downstream fetch stream, so this
        // cancels it too; see fetchUpstreamRange.
        entry.CancelFetches(&malformedTrackCause{src: src, err: cause})

        cancelled := 0
        for _, up := range entry.CopyUpstream() {
                if up.Session == src {
                        up.Cancel(moqt.StreamResetMalformedTrack)
                        cancelled++
                }
        }
        // Info once per upstream; later detections are Debug.
        level := slog.LevelDebug
        if cancelled > 0 {
                level = slog.LevelInfo
        }
        h.log.LogAttrs(ctx, level, "malformed track: ended it downstream, cancelled the upstream",
                slog.String("name", string(entry.FullName.Name)), slog.String("err", cause.Error()))
}

// malformedTrackCause is what a fetch stream on a malformed track is cancelled
// with: the detection error, wrapping [session.ErrMalformedTrack], and the
// session that sent the Object, whose FETCH is cancelled (§2.4.2: "MUST
// cancel any corresponding subscription or fetches for that Track from that
// publisher").
type malformedTrackCause struct {
        src *session.Session
        err error
}

func (c *malformedTrackCause) Error() string { return c.err.Error() }
func (c *malformedTrackCause) Unwrap() error { return c.err }

// resetWriters resets the open subgroup writers of subs on entry with
// MALFORMED_TRACK. The nil slot keeps the joiner scan from reopening one.
func (h *sessionHandler) resetWriters(entry *registry.TrackEntry, subs []*registry.DownstreamSub) {
        if len(subs) == 0 {
                return
        }
        var ws []*subgroupWriter
        for _, sg := range entry.CopySubgroups() {
                set, _ := sg.Set.(*subgroupWriterSet)
                sg.Mu.Lock()
                for _, sub := range subs {
                        if w := set.writers[sub]; w != nil {
                                w.close(true, moqt.StreamResetMalformedTrack)
                                set.writers[sub] = nil
                                ws = append(ws, w)
                        }
                }
                sg.Mu.Unlock()
        }
        if len(ws) > 0 {
                h.spawn(func() { joinWriters(ws) })
        }
}

package relay

import (
        "bytes"
        "context"
        "errors"
        "log/slog"
        "slices"
        "sync"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// handlePublishNamespace implements the PUBLISH_NAMESPACE flow (§6.2, §10.16):
//
//  1. Authorize.
//  2. Register the namespace in [registry.NamespaceRegistry].
//  3. Reply REQUEST_OK on the request stream.
//  4. Announce the namespace to matching SUBSCRIBE_NAMESPACE holders
//     ([registry.NamespaceRegistry.AnnouncePublisher]); unregistration
//     withdraws it.
//  5. SUBSCRIBE the publisher for every existing track the namespace covers
//     (§9.5; see [sessionHandler.subscribeExistingTracks]).
//  6. Serve follow-ups until the publisher cancels the request (§6.2), then
//     unregister.
func (h *sessionHandler) handlePublishNamespace(
        ctx context.Context,
        req *session.Request,
        msg *message.PublishNamespace,
) {
        if err := h.auth.AuthorizePublishNamespace(ctx, h.sess, msg); err != nil {
                h.rejectAuth(ctx, req, "PublishNamespace", err)
                return
        }

        entry := h.names.RegisterPublisher(msg.Namespace, h.sess, req.Stream)
        defer h.names.UnregisterPublisher(entry)

        if err := req.Reply(&message.RequestOK{}); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "PublishNamespace REQUEST_OK write failed",
                        slog.String("err", err.Error()))
                return
        }
        h.names.AnnouncePublisher(entry)

        // §9.5: no further SUBSCRIBEs once the publisher withdraws.
        nsCtx, cancel := context.WithCancel(ctx)
        defer cancel()
        h.spawn(func() { h.subscribeExistingTracks(nsCtx, entry) })

        h.serveNamespaceFollowups(ctx, req, h.publishNamespaceUpdate(req))
}

// publishNamespaceUpdate answers a REQUEST_UPDATE on a PUBLISH_NAMESPACE with
// REQUEST_OK. One whose tokens the TokenVerifier denies fails, and the relay
// closes the stream: "When a REQUEST_UPDATE fails for a ... PUBLISH_NAMESPACE,
// the responder MUST close the bidi stream" (§10.9.1). This goroutine is the
// only writer on the stream after REQUEST_OK, so the replies write directly.
func (h *sessionHandler) publishNamespaceUpdate(
        req *session.Request,
) func(context.Context, *message.RequestUpdate, []session.ResolvedToken) bool {
        return func(ctx context.Context, _ *message.RequestUpdate, toks []session.ResolvedToken) bool {
                if rej := h.refuseUpdateTokens(ctx, toks); rej != nil {
                        _ = req.RejectError(rej.ErrorCode, rej.ErrorReason)
                        return false
                }
                if err := req.Reply(&message.RequestOK{}); err != nil {
                        h.log.LogAttrs(ctx, slog.LevelDebug, "namespace REQUEST_UPDATE_OK write failed",
                                slog.String("err", err.Error()))
                        // Reset the read side so the peer learns reads stopped.
                        req.Stream.CancelRead(uint64(moqt.StreamResetInternalError))
                        return false
                }
                return true
        }
}

// subscribeExistingTracks SUBSCRIBEs pub for every existing track its
// namespace covers (§9.5: "it MUST send a SUBSCRIBE to the publisher").
// Tracks it skips, and a racing downstream SUBSCRIBE, are covered by
// [sessionHandler.subscribeMissingPublishers].
func (h *sessionHandler) subscribeExistingTracks(ctx context.Context, pub *registry.PublisherEntry) {
        for _, e := range h.tracks.MatchNamespace(pub.Namespace) {
                if ctx.Err() != nil {
                        return
                }
                h.subscribeLatePublisher(ctx, e.FullName, pub)
        }
}

// subscribeMissingPublishers runs once a downstream is on the track's entry
// and SUBSCRIBEs the covering publishers that have no upstream for it: all of
// them when the upstream set was reused, otherwise only those registered
// after seq, since a fresh set just tried the rest. Refused publishers are
// skipped (see [sessionHandler.subscribeLatePublisher]).
func (h *sessionHandler) subscribeMissingPublishers(
        ctx context.Context,
        entry *registry.TrackEntry,
        reused bool,
        seq uint64,
) {
        pubs := h.names.MatchPublishers(entry.FullName.Namespace)
        now := time.Now()
        entry.RetainRefusals(pubs, now)
        for _, pub := range pubs {
                if (!reused && pub.Seq <= seq) ||
                        entry.HasUpstreamOn(pub.Session) || entry.Refused(pub, now) {
                        continue
                }
                h.spawn(func() { h.subscribeLatePublisher(ctx, entry.FullName, pub) })
        }
}

// subscribeLatePublisher opens an upstream subscription for an existing track
// on pubEntry, a publisher whose PUBLISH_NAMESPACE covers it but which is not
// yet among the track's upstreams (§9.5).
//
// A refusal is recorded on the track entry: a REQUEST_ERROR until its Retry
// Interval passes, or for good when it is 0 (§10.6.2), and a Track Properties
// mismatch (§2.5.1) for good, while the entry and registration last.
//
// Deviation (§9.5 does not qualify "each matching subscription"): skipped
// while the track has no downstream, since the upstream would never be
// released. A pub that itself receives the track is SUBSCRIBEd too (§5.1:
// self-subscriptions "are identical to subscriptions initiated by other
// endpoints").
func (h *sessionHandler) subscribeLatePublisher(
        ctx context.Context,
        fullName track.FullTrackName,
        pubEntry *registry.PublisherEntry,
) {
        pub := pubEntry.Session
        e, ok := h.tracks.Get(fullName.Key())
        if !ok || !hasEstablishedUpstream(e) || len(e.CopyDownstream()) == 0 || e.Refused(pubEntry, time.Now()) {
                return
        }
        release, claimed := h.tracks.ClaimUpstream(pub, fullName.Key())
        if !claimed {
                return // already subscribed there, or being subscribed
        }
        defer release()
        entry, up, err := h.subscribeUpstreamOnSession(ctx, pub, fullName, nil, anyDownstreamForwards(e))
        if err != nil {
                if rej, ok := errors.AsType[*session.RequestRejectedError](err); ok {
                        var retryAt time.Time // zero: never
                        if after, retry := rej.RetryAfter(); retry {
                                retryAt = time.Now().Add(after)
                        }
                        e.NoteRefusal(pubEntry, retryAt)
                } else if isTrackPropertiesErr(err) {
                        e.NoteRefusal(pubEntry, time.Time{})
                }
                h.log.LogAttrs(ctx, slog.LevelDebug, "late publisher: SUBSCRIBE for existing track failed",
                        slog.String("name", string(fullName.Name)),
                        slog.String("err", err.Error()))
                return
        }
        // The track's last downstream may have left during the round trip.
        if h.tracks.ReleaseIfUnsubscribed(fullName, up) {
                h.log.LogAttrs(ctx, slog.LevelDebug, "late publisher: track lost its subscribers, released",
                        slog.String("name", string(fullName.Name)))
                return
        }
        // Or switched to Forward=1 before this upstream was registered for §9.2
        // propagation. Not bound to ctx: a withdrawal (§9.5) stops only new
        // subscriptions.
        if up.ForwardState() == 0 && anyDownstreamForwards(entry) {
                h.propagateForwardUpstream(context.WithoutCancel(ctx), fullName)
        }
}

// handleSubscribeNamespace implements SUBSCRIBE_NAMESPACE (§6.1, §10.19):
//
//  1. Authorize and reserve the prefix (PREFIX_OVERLAP).
//  2. Reply REQUEST_OK.
//  3. Register in [registry.NamespaceRegistry], which queues NAMESPACE /
//     NAMESPACE_DONE for existing and later namespaces (§6.1).
//  4. Serve REQUEST_UPDATEs, including TRACK_NAMESPACE_PREFIX (§10.9.2),
//     until the subscriber cancels.
func (h *sessionHandler) handleSubscribeNamespace(
        ctx context.Context,
        req *session.Request,
        msg *message.SubscribeNamespace,
) {
        if err := h.auth.AuthorizeSubscribeNamespace(ctx, h.sess, msg); err != nil {
                h.rejectAuth(ctx, req, "SubscribeNamespace", err)
                return
        }

        if !h.nsPrefixes.reserve(msg.TrackNamespacePrefix) {
                _ = req.RejectError(moqt.RequestPrefixOverlap,
                        "prefix overlaps an established SUBSCRIBE_NAMESPACE in this session")
                return
        }
        // prefix follows TRACK_NAMESPACE_PREFIX updates (§10.9.2).
        prefix := msg.TrackNamespacePrefix
        defer func() { h.nsPrefixes.release(prefix) }()

        // §6.1: REQUEST_OK before any NAMESPACE, so reply before registering.
        if err := req.Reply(&message.RequestOK{}); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "SubscribeNamespace REQUEST_OK write failed",
                        slog.String("err", err.Error()))
                return
        }

        entry := h.names.RegisterSubscriber(
                msg.TrackNamespacePrefix,
                h.sess,
                req.Stream,
                false, /* wantsTracks */
                nil,
                nil,
        )
        defer h.names.UnregisterSubscriber(entry)

        h.spawn(entry.RunWriter)

        // Replies share the entry's queue, keeping their order with NAMESPACE.
        h.serveNamespaceFollowups(ctx, req, h.namespaceUpdate(entry, &prefix, msg))
}

// handleSubscribeTracks implements SUBSCRIBE_TRACKS (§6.1, §10.20):
//
//  1. Authorize, validate its subscription parameters (§10.20.1), and
//     reserve the prefix (PREFIX_OVERLAP).
//  2. Reply REQUEST_OK.
//  3. Register in [registry.NamespaceRegistry] with forwardTrack, then
//     forward the tracks that already exist under the prefix (§10.20).
//     Later tracks are forwarded by handlePublish.
//  4. Serve REQUEST_UPDATEs until the subscriber cancels.
func (h *sessionHandler) handleSubscribeTracks(
        ctx context.Context,
        req *session.Request,
        msg *message.SubscribeTracks,
) {
        if err := h.auth.AuthorizeSubscribeTracks(ctx, h.sess, msg); err != nil {
                h.rejectAuth(ctx, req, "SubscribeTracks", err)
                return
        }

        // §10.20.1: refused on a SUBSCRIBE's terms.
        params, err := h.resolveTracksParams(msg.Parameters)
        if err != nil {
                h.refuseSubscriptionParams(ctx, req, err)
                return
        }

        if !h.trackPrefixes.reserve(msg.TrackNamespacePrefix) {
                _ = req.RejectError(moqt.RequestPrefixOverlap,
                        "prefix overlaps an established SUBSCRIBE_TRACKS in this session")
                return
        }
        prefix := msg.TrackNamespacePrefix
        defer func() { h.trackPrefixes.release(prefix) }()

        // Reply before registering, so the OK cannot race a PUBLISH_SKIPPED
        // written once the entry is visible.
        if err := req.Reply(&message.RequestOK{}); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "SubscribeTracks REQUEST_OK write failed",
                        slog.String("err", err.Error()))
                return
        }

        entry := h.names.RegisterSubscriber(
                msg.TrackNamespacePrefix,
                h.sess,
                req.Stream,
                true, /* wantsTracks */
                params,
                h.forwardTrack(ctx),
        )
        defer h.names.UnregisterSubscriber(entry)

        h.spawn(entry.RunWriter)
        // §10.20: forward the tracks that already exist under the prefix.
        for _, te := range h.tracks.MatchNamespace(msg.TrackNamespacePrefix) {
                if hasEstablishedUpstream(te) {
                        entry.ForwardTrack(entry, te)
                }
        }
        // Replies share the entry's queue with PUBLISH_SKIPPED, so each
        // PUBLISH_SKIPPED suffix matches the prefix the subscriber last saw.
        h.serveNamespaceFollowups(ctx, req, h.tracksUpdate(entry, &prefix, msg))
}

// subscribeTracksForwarding resolves a SUBSCRIBE_TRACKS's FORWARD (§10.2.18,
// default true) and GROUP_ORDER (§10.2.8, 0 when omitted), which §10.20.1
// copies onto forwarded PUBLISHes.
func subscribeTracksForwarding(ps message.Parameters) (forward bool, groupOrder byte) {
        forward = true
        if p, ok := ps.Find(message.ParamForward); ok {
                forward = p.Byte != 0
        }
        if p, ok := ps.Find(message.ParamGroupOrder); ok {
                groupOrder = p.Byte
        }
        return forward, groupOrder
}

// serveNamespaceFollowups holds a namespace request stream open and answers
// each REQUEST_UPDATE (§10.9), validating its Request ID (§10.1) and resolving
// its tokens; update authorizes and applies it, and replies. Other follow-ups
// are ignored.
func (h *sessionHandler) serveNamespaceFollowups(
        ctx context.Context,
        req *session.Request,
        update func(context.Context, *message.RequestUpdate, []session.ResolvedToken) bool,
) {
        stream := req.Stream
        scope := message.ScopeOfUpdate(req.First.Type())
        updates := h.sess.NewRequestUpdateLimiter()
        fin := readRequestStream(ctx, h.sess, stream, func(m message.Message) bool {
                if h.isPeerStateNotify(m) {
                        return false
                }
                upd, ok := m.(*message.RequestUpdate)
                if !ok {
                        return true
                }
                // §10.2.1: parameters outside this request's update scope are
                // session-fatal.
                if h.sess.CheckPeerParams(scope, upd) != nil {
                        return false
                }
                if !h.handleFollowupRequestID(ctx, upd) {
                        return false
                }
                // §10.3.1.7: enforce the per-stream MAX_REQUEST_UPDATES limit.
                if !h.handleRequestUpdateLimit(ctx, updates) {
                        return false
                }
                toks, ok := h.handleFollowupTokens(ctx, upd)
                if !ok {
                        return false
                }
                // false from update ends the request.
                if !update(ctx, upd, toks) {
                        return false
                }
                updates.Responded()
                return true
        })
        if fin {
                // §3.3.2: a FIN is not a cancellation (§6.1, §6.2); keep the state
                // until the peer resets or sends STOP_SENDING.
                awaitRequestEnd(ctx, stream)
        }
}

// updatePrefixParam reads a REQUEST_UPDATE's TRACK_NAMESPACE_PREFIX
// (§10.2.20), if any. ok is false when it is malformed: the session is then
// closed with PROTOCOL_VIOLATION.
func (h *sessionHandler) updatePrefixParam(upd *message.RequestUpdate) (prefix wire.TrackNamespace, found, ok bool) {
        p, found := upd.Parameters.Find(message.ParamTrackNamespacePrefix)
        if !found {
                return nil, false, true
        }
        prefix, err := message.TrackNamespacePrefixFromParam(p)
        if err != nil {
                _ = h.sess.Close(moqt.SessionProtocolViolation, err.Error())
                return nil, false, false
        }
        return prefix, true, true
}

// namespaceUpdate answers a REQUEST_UPDATE on the SUBSCRIBE_NAMESPACE msg: a
// TRACK_NAMESPACE_PREFIX is authorized and applied (§10.9.2); anything else is
// acknowledged. A refused update ends the request (see [endAfterFinish]).
func (h *sessionHandler) namespaceUpdate(
        e *registry.SubscriberEntry,
        cur *wire.TrackNamespace,
        msg *message.SubscribeNamespace,
) func(context.Context, *message.RequestUpdate, []session.ResolvedToken) bool {
        updatePrefix := h.prefixUpdater(e, &h.nsPrefixes, cur)
        tokens := authorizingTokens(msg.Parameters)
        return func(ctx context.Context, upd *message.RequestUpdate, toks []session.ResolvedToken) bool {
                prefix, found, ok := h.updatePrefixParam(upd)
                if !ok {
                        return false
                }
                updTokens := updatedTokens(tokens, upd.Parameters)
                if rej := h.refuseUpdate(ctx, toks, found, func() error {
                        return h.auth.AuthorizeSubscribeNamespace(ctx, h.sess, &message.SubscribeNamespace{
                                RequestID:            msg.RequestID,
                                TrackNamespacePrefix: prefix,
                                Parameters:           withTokens(msg.Parameters, updTokens),
                        })
                }); rej != nil {
                        e.Finish(rej)
                        return endAfterFinish(ctx, e)
                }
                tokens = updTokens
                if !found {
                        e.Enqueue(&message.RequestOK{})
                        return true
                }
                if !updatePrefix(prefix) {
                        return endAfterFinish(ctx, e)
                }
                return true
        }
}

// tracksUpdate answers a REQUEST_UPDATE on a SUBSCRIBE_TRACKS. Its parameters
// are merged (see [mergeTracksUpdate]) and apply only to future forwards
// (§10.2.18: "Existing subscriptions are unaffected"). A refused update
// changes nothing and ends the request (see [endAfterFinish]).
//
// Existing tracks that newly match, by prefix or Range Filter, are forwarded
// (§10.20); a track that matched before is not offered again.
func (h *sessionHandler) tracksUpdate(
        e *registry.SubscriberEntry,
        cur *wire.TrackNamespace,
        msg *message.SubscribeTracks,
) func(context.Context, *message.RequestUpdate, []session.ResolvedToken) bool {
        tokens := authorizingTokens(msg.Parameters)
        return func(ctx context.Context, upd *message.RequestUpdate, toks []session.ResolvedToken) bool {
                prefix, hasPrefix, ok := h.updatePrefixParam(upd)
                if !ok {
                        return false
                }
                before := e.TracksParams()
                merged := mergeTracksUpdate(before.Params, upd.Parameters)
                updTokens := updatedTokens(tokens, upd.Parameters)
                if rej := h.refuseUpdate(ctx, toks, hasPrefix, func() error {
                        return h.auth.AuthorizeSubscribeTracks(ctx, h.sess, &message.SubscribeTracks{
                                RequestID:            msg.RequestID,
                                TrackNamespacePrefix: prefix,
                                Parameters:           withTokens(merged, updTokens),
                        })
                }); rej != nil {
                        e.Finish(rej)
                        return endAfterFinish(ctx, e)
                }
                tokens = updTokens
                params, err := h.resolveTracksParams(merged)
                if err != nil {
                        e.Finish(&message.RequestError{ErrorCode: moqt.RequestInvalidFilter, ErrorReason: err.Error()})
                        return endAfterFinish(ctx, e)
                }
                oldPrefix := *cur
                if hasPrefix {
                        if !h.trackPrefixes.replace(*cur, prefix) {
                                e.Finish(&message.RequestError{
                                        ErrorCode:   moqt.RequestPrefixOverlap,
                                        ErrorReason: "updated prefix overlaps another subscription in this session",
                                })
                                return endAfterFinish(ctx, e)
                        }
                        *cur = prefix
                }
                e.SetTracksParams(params)
                if hasPrefix {
                        h.names.UpdatePrefix(e, prefix, &message.RequestOK{})
                } else {
                        e.Enqueue(&message.RequestOK{})
                }
                for _, te := range h.tracks.MatchNamespace(*cur) {
                        matchedBefore := te.FullName.Namespace.HasPrefix(oldPrefix) &&
                                before.RangeFilters.MatchesTrack(te.GetProperties())
                        if !matchedBefore && hasEstablishedUpstream(te) {
                                e.ForwardTrack(e, te)
                        }
                }
                return true
        }
}

// refuseUpdate authorizes a REQUEST_UPDATE on a namespace subscription,
// returning the REQUEST_ERROR that refuses it, or nil. Its tokens (§10.2.2)
// go through the TokenVerifier as an opener's do. When it changes the prefix,
// authorize runs the Authorizer on the subscription it would become: §10.19
// and §10.20 require that "the subscriber is authorized to perform this
// namespace subscription".
func (h *sessionHandler) refuseUpdate(
        ctx context.Context,
        toks []session.ResolvedToken,
        prefixChanged bool,
        authorize func() error,
) *message.RequestError {
        if rej := h.refuseUpdateTokens(ctx, toks); rej != nil {
                return rej
        }
        if !prefixChanged {
                return nil
        }
        if err := authorize(); err != nil {
                return &message.RequestError{
                        ErrorCode:   CodeForAuthorizerError(err),
                        ErrorReason: ReasonForAuthorizerError(err),
                }
        }
        return nil
}

// authorizingTokens is the AUTHORIZATION_TOKENs in ps that can authorize a
// request: all but DELETEs, which only retire an alias (§10.2.2).
func authorizingTokens(ps message.Parameters) message.Parameters {
        var out message.Parameters
        for _, p := range ps {
                if p.Type != message.ParamAuthorizationToken {
                        continue
                }
                var tok message.Token
                if tok.Parse(p.Bytes) == nil && tok.AliasType != message.AliasTypeDelete {
                        out = append(out, p)
                }
        }
        return out
}

// updatedTokens is the tokens a namespace subscription holds after an update
// carrying upd: its authorizing tokens replace cur when it has any; otherwise
// cur "remains unchanged" (§10.9).
func updatedTokens(cur, upd message.Parameters) message.Parameters {
        if toks := authorizingTokens(upd); len(toks) > 0 {
                return toks
        }
        return cur
}

// withTokens is ps with its AUTHORIZATION_TOKENs replaced by toks: the
// subscription an update's Authorizer call judges.
func withTokens(ps, toks message.Parameters) message.Parameters {
        out := slices.DeleteFunc(slices.Clone(ps), func(p message.Parameter) bool {
                return p.Type == message.ParamAuthorizationToken
        })
        return append(out, toks...)
}

// mergeTracksUpdate applies a REQUEST_UPDATE's parameters to a
// SUBSCRIBE_TRACKS's: each type present in upd replaces every stored one of
// that type (§10.9); a zero-length Range Filter removes it (§5.1.4).
// TRACK_NAMESPACE_PREFIX and AUTHORIZATION_TOKEN belong to the update and
// are not kept; a token in upd drops the stored one.
func mergeTracksUpdate(stored, upd message.Parameters) message.Parameters {
        out := slices.DeleteFunc(slices.Clone(stored), func(p message.Parameter) bool {
                return slices.ContainsFunc(upd, func(u message.Parameter) bool { return u.Type == p.Type })
        })
        for _, p := range upd {
                switch {
                case p.Type == message.ParamTrackNamespacePrefix, p.Type == message.ParamAuthorizationToken:
                case message.IsRangeFilterParam(p.Type) && len(p.Bytes) == 0:
                default:
                        out = append(out, p)
                }
        }
        return out
}

// resolveTracksParams validates a SUBSCRIBE_TRACKS's Range Filters on a
// SUBSCRIBE's terms (§10.20.1) and MAX_FILTER_RANGES (§5.1.4), and resolves
// what forwarding needs. An error is INVALID_FILTER.
func (h *sessionHandler) resolveTracksParams(ps message.Parameters) (*registry.TracksParams, error) {
        forward, groupOrder := subscribeTracksForwarding(ps)
        rangeFilters, err := message.RangeFiltersFromParams(ps)
        if err == nil && rangeFilters != nil {
                err = rangeFilters.Validate(h.sess.MaxFilterRanges())
        }
        if err != nil {
                return nil, err
        }
        if err := installSubscribeParams(registry.NewDownstreamSub(0, h.sess, nil, 0), ps); err != nil {
                return nil, err
        }
        return &registry.TracksParams{Params: ps, Forward: forward, GroupOrder: groupOrder, RangeFilters: rangeFilters}, nil
}

// prefixUpdater applies a TRACK_NAMESPACE_PREFIX update (§10.9.2) to e and
// replies. An overlapping prefix is refused with PREFIX_OVERLAP (§10.2.20),
// and the updater reports false.
func (h *sessionHandler) prefixUpdater(
        e *registry.SubscriberEntry,
        reserved *prefixSet,
        cur *wire.TrackNamespace,
) func(wire.TrackNamespace) bool {
        return func(prefix wire.TrackNamespace) bool {
                if !reserved.replace(*cur, prefix) {
                        e.Finish(&message.RequestError{
                                ErrorCode:   moqt.RequestPrefixOverlap,
                                ErrorReason: "updated prefix overlaps another subscription in this session",
                        })
                        return false
                }
                *cur = prefix
                h.names.UpdatePrefix(e, prefix, &message.RequestOK{})
                return true
        }
}

// endAfterFinish ends a namespace subscription whose REQUEST_UPDATE was
// refused (§10.9.1: "MUST close the bidi stream"). It waits for the writer to
// send the queued REQUEST_ERROR and FIN, then reports false.
func endAfterFinish(ctx context.Context, e *registry.SubscriberEntry) bool {
        select {
        case <-e.WriterDone():
        case <-ctx.Done():
        }
        return false
}

// prefixSet holds one session's established namespace-subscription prefixes
// of one type, for PREFIX_OVERLAP (§10.19 / §10.20). The zero value is ready.
type prefixSet struct {
        mu       sync.Mutex
        prefixes []wire.TrackNamespace
}

// reserve records prefix and reports true, or reports false when it overlaps
// an established one. Interpretation: "shares a common prefix" means one is a
// prefix of the other (read literally, every pair shares the empty prefix).
func (p *prefixSet) reserve(prefix wire.TrackNamespace) bool {
        p.mu.Lock()
        defer p.mu.Unlock()
        if slices.ContainsFunc(p.prefixes, func(have wire.TrackNamespace) bool {
                return have.HasPrefix(prefix) || prefix.HasPrefix(have)
        }) {
                return false
        }
        p.prefixes = append(p.prefixes, prefix)
        return true
}

// replace swaps the reservation of old for prefix and reports true, or false
// when prefix overlaps a reservation other than old (§10.9.2).
func (p *prefixSet) replace(old, prefix wire.TrackNamespace) bool {
        p.mu.Lock()
        defer p.mu.Unlock()
        i := slices.IndexFunc(p.prefixes, func(have wire.TrackNamespace) bool {
                return slices.EqualFunc(have, old, bytes.Equal)
        })
        for j, have := range p.prefixes {
                if j != i && (have.HasPrefix(prefix) || prefix.HasPrefix(have)) {
                        return false
                }
        }
        p.prefixes[i] = prefix // old is the caller's own reservation
        return true
}

// release forgets a prefix reserve accepted.
func (p *prefixSet) release(prefix wire.TrackNamespace) {
        p.mu.Lock()
        defer p.mu.Unlock()
        if i := slices.IndexFunc(p.prefixes, func(have wire.TrackNamespace) bool {
                return slices.EqualFunc(have, prefix, bytes.Equal)
        }); i >= 0 {
                p.prefixes = slices.Delete(p.prefixes, i, i+1)
        }
}

package relay

import (
        "context"
        "fmt"
        "log/slog"
        "slices"
        "sync/atomic"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// handlePublish implements PUBLISH (§9.5, §10.11):
//
//  1. Authorize.
//  2. Register an [registry.UpstreamSub] in the registry.TrackRegistry (born
//     [registry.SubEstablished]).
//  3. Capture the publisher's Track Properties on the entry (§9.6).
//  4. Reply REQUEST_OK.
//  5. Forward the PUBLISH to every downstream SUBSCRIBE_TRACKS holder
//     whose prefix matches the track's namespace (§9.5: relay MUST send
//     PUBLISH to each matching SUBSCRIBE_TRACKS holder).
//  6. Register the publisher's Track Alias as an inbound alias so the
//     fanout path can map it back to the track.
//  7. Block reading the request stream until the publisher cancels;
//     unregister on exit.
//
// testHookAfterAliasRegistered, when set by a test, runs once a Track Alias is
// routable and before the upstream is registered, to hold that window open.
var testHookAfterAliasRegistered atomic.Pointer[func(track.FullTrackName)]

func (h *sessionHandler) handlePublish(ctx context.Context, req *session.Request, msg *message.Publish) {
        h.log.LogAttrs(ctx, slog.LevelDebug, "PUBLISH received",
                slog.String("namespace", fmt.Sprintf("%v", msg.Namespace)),
                slog.String("name", string(msg.Name)),
                slog.Uint64("alias", msg.TrackAlias))

        // §2.5.1: refuse an unknown Mandatory Track Property. Unparseable Track
        // Properties are refused with INTERNAL_ERROR, since §10.6.2 defines
        // MALFORMED_TRACK only for FETCH (see [session.TrackPropertiesRejectCode]).
        // AcceptRequest has closed the session on a session-fatal value (§12.5,
        // §12.6).
        if err := h.sess.CheckTrackProperties(msg.TrackProperties, "PUBLISH"); err != nil {
                _ = req.RejectError(session.TrackPropertiesRejectCode(err), err.Error())
                return
        }

        if err := h.auth.AuthorizePublish(ctx, h.sess, msg); err != nil {
                h.rejectAuth(ctx, req, "Publish", err)
                return
        }

        fullName := track.FullTrackName{Namespace: msg.Namespace, Name: msg.Name}

        // Create the entry before the alias becomes routable: objects may arrive
        // "possibly before PUBLISH_OK" (§10.11), and runFanout resets streams for
        // a track with no entry.
        _, createdEntry := h.tracks.GetOrCreateNew(fullName)

        // §11.1: register the publisher's chosen alias so the fanout path can map
        // it back to the track. A duplicate closed the session.
        if err := h.sess.RegisterInboundTrack(msg.TrackAlias, fullName.Key(), msg.TrackProperties); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "PUBLISH alias registration failed",
                        slog.String("err", err.Error()))
                if createdEntry {
                        h.tracks.DeleteIfUnused(fullName)
                }
                return
        }
        if hook := testHookAfterAliasRegistered.Load(); hook != nil {
                (*hook)(fullName)
        }

        // The publisher sent the PUBLISH, so it may send REQUEST_UPDATE (§10.9),
        // in the scope an accepted PUBLISH's publisher has.
        broker := h.sess.NewRequestBroker(req.Stream)
        broker.PeerMessages(true, true)
        broker.UpdateScope(message.ScopeUpdateFromPublisher)
        sub := registry.NewUpstreamSub(h.allocSubID(), h.sess, req.Stream, broker, msg.TrackAlias, msg.RequestID)
        // §5.1: the PUBLISH sets the initial Forward State (default 1).
        if f, ok := msg.Parameters.Find(message.ParamForward); ok && f.Byte == 0 {
                sub.SetForwardState(0)
        }

        // Register the upstream and reply REQUEST_OK atomically under the
        // stream's broker write lock: registration must precede the OK (a prompt
        // SUBSCRIBE elsewhere must find the track), and the OK must be the
        // stream's next message ahead of any propagated REQUEST_UPDATE.
        var entry *registry.TrackEntry
        if err := sub.Broker.WriteMessageAfterSetup(func() error {
                entry, _ = h.tracks.AddUpstream(fullName, sub, registry.WithProperties(msg.TrackProperties))
                // §10.2.17: PUBLISH may carry LARGEST_OBJECT.
                saveLargestLocation(entry, msg.Parameters)
                return nil
        }, &message.RequestOK{}); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "PUBLISH REQUEST_OK write failed",
                        slog.String("err", err.Error()))
                h.tracks.RemoveUpstream(fullName, sub.ID)
                h.sess.UnregisterInboundTrackAlias(msg.TrackAlias)
                return
        }
        defer func() {
                h.log.LogAttrs(ctx, slog.LevelDebug, "PUBLISH stream ended, removing upstream",
                        slog.String("name", string(msg.Name)))
                h.tracks.RemoveUpstream(fullName, sub.ID)
                h.sess.UnregisterInboundTrackAlias(msg.TrackAlias)
        }()
        h.log.LogAttrs(ctx, slog.LevelDebug, "PUBLISH accepted, waiting for publisher",
                slog.String("name", string(msg.Name)))

        // §9.5: resume a paused upstream if any downstream forwards. Spawned: the
        // response is read by the Serve loop serveUpstreamStream starts below.
        if sub.ForwardState() == 0 && anyDownstreamForwards(entry) {
                h.spawn(func() { h.propagateForwardUpstream(ctx, fullName) })
        }

        h.forwardToTrackSubscribers(entry)

        // Block until the publisher tears the stream down, routing §10.9
        // responses to any upstream REQUEST_UPDATE the relay sends meanwhile
        // (e.g. NEW_GROUP_REQUEST propagation).
        h.serveUpstreamStream(ctx, sub)
}

// forwardToTrackSubscribers offers entry's track to every SUBSCRIBE_TRACKS
// holder whose prefix matches (§6.1, §10.20); see forwardTrack.
func (h *sessionHandler) forwardToTrackSubscribers(entry *registry.TrackEntry) {
        for _, sub := range h.names.MatchSubscribers(entry.FullName.Namespace) {
                if sub.WantsTracks && sub.ForwardTrack != nil {
                        sub.ForwardTrack(sub, entry)
                }
        }
}

// notEchoedInPublish are the SUBSCRIBE_TRACKS parameters
// [publishParamsForSubscriber] does not copy: AUTHORIZATION_TOKEN (§10.2.2), and
// the ones it sets itself.
var notEchoedInPublish = []message.ParamID{
        message.ParamAuthorizationToken, message.ParamForward, message.ParamGroupOrder, message.ParamLargestObject,
}

// publishParamsForSubscriber builds the Parameters of a PUBLISH forwarded for
// a SUBSCRIBE_TRACKS. None come from the upstream (§10.2.1: "not forwarded by
// Relays"):
//   - SUBSCRIBE_TRACKS parameters valid on PUBLISH are echoed (§10.20.1),
//     except AUTHORIZATION_TOKEN (§10.2.2);
//   - FORWARD=0 only when the subscriber set it;
//   - GROUP_ORDER always: the subscriber's, else the publisher's preference
//     (§10.2.8), which the subscription's fills follow; §10.20.1 has these
//     "explicitly communicated in PUBLISH";
//   - LARGEST_OBJECT is the relay's own watermark (§10.2.17).
func publishParamsForSubscriber(tp *registry.TracksParams, entry *registry.TrackEntry) message.Parameters {
        var out message.Parameters
        for _, p := range tp.Params {
                if slices.Contains(notEchoedInPublish, p.Type) {
                        continue
                }
                if (message.Parameters{p}).CheckScope(message.ScopePublish) == nil {
                        out = append(out, p)
                }
        }
        if !tp.Forward {
                out = append(out, message.ForwardParam(false))
        }
        order := message.GroupOrder(tp.GroupOrder)
        if order == 0 {
                order = entry.DefaultGroupOrder()
        }
        out = append(out, message.GroupOrderParam(order))
        if largest, ok := entry.GetLargest(); ok {
                out = append(out, message.LargestObjectParam(largest.Group, largest.Object))
        }
        return out
}

// emitPublishSkipped queues a PUBLISH_SKIPPED (§10.21) for fullName on sub's
// SUBSCRIBE_TRACKS stream, the §6.1 response to an exhausted bidi-stream
// limit, and records it at the track's upstream epoch: the skip is "scoped to
// a single PUBLISH", so only a new upstream offers the track again, not a
// prefix update that moves away and back (see
// [registry.SubscriberEntry.NoteSkipped]).
func (h *sessionHandler) emitPublishSkipped(
        ctx context.Context,
        sub *registry.SubscriberEntry,
        fullName track.FullTrackName,
        epoch uint64,
) {
        if !h.names.PublishSkipped(sub, fullName.Namespace, fullName.Name) {
                return // a TRACK_NAMESPACE_PREFIX update moved the subscription away
        }
        sub.NoteSkipped(fullName.Key(), epoch)
        h.log.LogAttrs(ctx, slog.LevelDebug, "PUBLISH_SKIPPED queued",
                slog.String("name", string(fullName.Name)))
}

package relay

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "slices"
        "sync/atomic"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// testHookBeforeDownstreamRegistered, when set by a test, runs once a
// SUBSCRIBE has an upstream for its track and before its downstream is
// registered, to hold that window open.
var testHookBeforeDownstreamRegistered atomic.Pointer[func(track.FullTrackName)]

// handleSubscribe implements the SUBSCRIBE flow (§9.4, §10.7): authorize,
// serve from an Established upstream or establish one on demand (see
// [sessionHandler.acquireUpstream]), else reject with
// [moqt.RequestDoesNotExist], or with [moqt.RequestTimeout] once a
// RENDEZVOUS_TIMEOUT hold expires (§10.2.6); then register a [registry.DownstreamSub], reply
// SUBSCRIBE_OK, and serve the request stream until it ends.
func (h *sessionHandler) handleSubscribe(ctx context.Context, req *session.Request, msg *message.Subscribe) {
        h.log.LogAttrs(ctx, slog.LevelDebug, "SUBSCRIBE received",
                slog.String("namespace", fmt.Sprintf("%v", msg.Namespace)),
                slog.String("name", string(msg.Name)))

        if err := h.auth.AuthorizeSubscribe(ctx, h.sess, msg); err != nil {
                h.rejectAuth(ctx, req, "Subscribe", err)
                return
        }

        fullName := track.FullTrackName{Namespace: msg.Namespace, Name: msg.Name}

        // §10.2.19: a NEW_GROUP_REQUEST rides a new upstream SUBSCRIBE (rule 1;
        // see acquireUpstream), or is evaluated against an existing upstream below.
        newGroupReqParam, hasNewGroupReq := msg.Parameters.Find(message.ParamNewGroupRequest)

        // §11.1: outbound aliases are independent of the peer's inbound ones.
        alias := h.sess.AllocOutboundTrackAlias()

        sub := registry.NewDownstreamSub(h.allocSubID(), h.sess, req.Stream, alias)
        if err := installSubscribeParams(sub, msg.Parameters); err != nil {
                h.refuseSubscriptionParams(ctx, req, err)
                return
        }

        // Two attempts: registration fails if the last upstream vanished after
        // the establish check, and the retry re-establishes.
        var (
                entry           *registry.TrackEntry
                snapshotLargest message.Location
                snapshotHas     bool
                reusedUpstream  bool
                added           bool
        )
        // Publishers that register after acquireUpstream looked are picked up
        // below, once the downstream is on the entry.
        var pubSeq uint64
        // In flight until the downstream is registered; see
        // [sessionHandler.beginSubscribe].
        settled := h.beginSubscribe(fullName.Key())
        defer settled()
        rv := h.newRendezvous(msg.Parameters)
        for range 2 {
                var ok bool
                if reusedUpstream, pubSeq, ok = h.acquireUpstream(ctx, req, msg, sub, rv); !ok {
                        return
                }

                // Before registration, so the first stream opened for it already
                // schedules in its Group Order (§7.2).
                if cur, ok := h.tracks.Get(fullName.Key()); ok {
                        resolveGroupOrder(sub, cur)
                }
                if hook := testHookBeforeDownstreamRegistered.Load(); hook != nil {
                        (*hook)(fullName)
                }
                // Register and snapshot Largest atomically, so no object falls between
                // live delivery and the fill fetch stream.
                entry, snapshotLargest, snapshotHas, added = h.tracks.AddDownstreamSnapshotLargest(fullName, sub)
                if added {
                        break
                }
        }
        settled()
        if !added {
                h.log.LogAttrs(ctx, slog.LevelDebug, "SUBSCRIBE rejected: upstream vanished during registration")
                _ = req.RejectError(moqt.RequestDoesNotExist, "relay: upstream vanished")
                return
        }
        sub.SetLargestAtSubscribe(snapshotLargest, snapshotHas)
        // §9.5: "Relays MUST send SUBSCRIBE messages to all matching publishers".
        h.subscribeMissingPublishers(ctx, entry, reusedUpstream, pubSeq)
        // §10.20: a newly upstreamed track is offered to SUBSCRIBE_TRACKS holders;
        // after registration, so this subscriber is not offered its own track. The
        // forwards racing registration from other paths are stopped by
        // beginSubscribe.
        if !reusedUpstream {
                h.forwardToTrackSubscribers(entry)
        }

        subRef := h.trackRef(fullName)
        h.metrics.SubscriptionOpened(subRef)
        defer h.metrics.SubscriptionClosed(subRef)

        defer h.tracks.RemoveDownstream(fullName, sub.ID)
        // §5.1.1: once the subscriber cancels, or the session ends, reset the
        // streams still open for it.
        defer sub.Cancel()

        // §10.2.17: "If Objects have been published on this Track the Publisher
        // MUST include this parameter."
        var okParams message.Parameters
        if sub.HasLargestAtSubscribe {
                okParams = message.Parameters{
                        message.LargestObjectParam(
                                sub.LargestAtSubscribe.Group,
                                sub.LargestAtSubscribe.Object,
                        ),
                }
        }
        var properties []byte
        if sub.IncludesProperties() { // §10.2.21
                properties = entry.GetProperties()
        }
        // Not req.Reply: the sub is registered, so every write must go through
        // its write lock, and a racing termination yields exactly one response
        // (see [registry.DownstreamSub.WriteSubscribeOK]).
        if err := sub.WriteSubscribeOK(&message.SubscribeOK{
                TrackAlias:      alias,
                Parameters:      okParams,
                TrackProperties: properties,
        }); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "SUBSCRIBE_OK write failed",
                        slog.String("err", err.Error()))
                return
        }
        h.log.LogAttrs(ctx, slog.LevelDebug, "SUBSCRIBE_OK sent, waiting for subscriber",
                slog.String("name", string(msg.Name)),
                slog.Uint64("alias", alias))

        // §5.1.3: FILL_PARAMETERS asks for a fill fetch stream; AcceptRequest
        // has closed the session on a malformed one.
        if err := h.maybeServeFill(ctx, sub, entry, fullName, msg.RequestID, msg.Parameters,
                snapshotLargest, snapshotHas); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "fill fetch stream not opened",
                        slog.String("err", err.Error()))
        }

        if hasNewGroupReq && reusedUpstream {
                h.propagateNewGroupUpstream(ctx, fullName, newGroupReqParam.Varint)
        }

        // §9.2: a Forward=1 subscriber resumes a paused upstream it reuses.
        if reusedUpstream && sub.ForwardState() == 1 {
                h.propagateForwardUpstream(ctx, fullName)
        }

        h.readSubscribeUpdates(ctx, req.Stream, sub, fullName)
        h.log.LogAttrs(ctx, slog.LevelDebug, "SUBSCRIBE stream ended",
                slog.String("name", string(msg.Name)))
}

// readSubscribeUpdates routes REQUEST_UPDATEs (§10.9) on a downstream
// SUBSCRIBE's stream to [sessionHandler.handleSubscribeUpdate] until the
// subscriber cancels, the stream turns undecodable (see [readRequestStream])
// or ctx ends. A subscriber FIN is not a cancellation (§3.3.2): the
// subscription lives on in [awaitRequestEnd]. The stream is a SUBSCRIBE's or
// a forwarded PUBLISH's.
func (h *sessionHandler) readSubscribeUpdates(
        ctx context.Context,
        stream session.Stream,
        sub *registry.DownstreamSub,
        fullName track.FullTrackName,
) {
        updates := h.sess.NewRequestUpdateLimiter()
        fin := readRequestStream(ctx, h.sess, stream, func(m message.Message) bool {
                if h.isPeerStateNotify(m) {
                        return false
                }
                if upd, ok := m.(*message.RequestUpdate); ok {
                        // §10.2.1: parameters outside the scope of a subscriber's
                        // update are session-fatal.
                        if h.sess.CheckPeerParams(message.ScopeUpdateFromSubscriber, upd) != nil {
                                return false
                        }
                        // §10.1: the update consumes a Request ID; a parity or
                        // duplicate violation is session-fatal.
                        if !h.handleFollowupRequestID(ctx, upd) {
                                return false
                        }
                        // §10.3.1.7: enforce the per-stream MAX_REQUEST_UPDATES limit.
                        if !h.handleRequestUpdateLimit(ctx, updates) {
                                return false
                        }
                        // §10.2.2: an update may REGISTER/DELETE token aliases;
                        // a cache fault there is session-fatal.
                        toks, ok := h.handleFollowupTokens(ctx, upd)
                        if !ok {
                                return false
                        }
                        h.handleSubscribeUpdate(ctx, sub, fullName, upd, toks)
                        updates.Responded()
                }
                return true
        })
        if fin {
                awaitRequestEnd(ctx, stream)
        }
}

// handleSubscribeUpdate applies a REQUEST_UPDATE (§10.9) to a downstream
// subscription: present parameters override, omitted ones are kept. An update
// whose tokens (toks) the TokenVerifier denies, or that is malformed, gets
// REQUEST_ERROR and PUBLISH_DONE / UPDATE_FAILED.
func (h *sessionHandler) handleSubscribeUpdate(
        ctx context.Context,
        sub *registry.DownstreamSub,
        fullName track.FullTrackName,
        upd *message.RequestUpdate,
        toks []session.ResolvedToken,
) {
        // §10.9.1: REQUEST_ERROR, then PUBLISH_DONE / UPDATE_FAILED. Writes go
        // through the sub's lock.
        fail := func(rej *message.RequestError) {
                h.log.LogAttrs(ctx, slog.LevelDebug, "REQUEST_UPDATE refused",
                        slog.String("err", rej.ErrorReason))
                _ = sub.WriteMessage(rej)
                sub.TerminateWithPublishDone(moqt.PublishDoneUpdateFailed, rej.ErrorReason)
        }
        if rej := h.refuseUpdateTokens(ctx, toks); rej != nil {
                fail(rej)
                return
        }
        prevForward := sub.ForwardState()
        if err := installSubscribeParams(sub, upd.Parameters); err != nil {
                fail(&message.RequestError{ErrorCode: moqt.RequestInvalidFilter, ErrorReason: err.Error()})
                return
        }

        // §10.2.17: LARGEST_OBJECT in REQUEST_UPDATE_OK too.
        reply := &message.RequestOK{}
        var (
                largest    message.Location
                hasLargest bool
        )
        if entry, ok := h.tracks.Get(fullName.Key()); ok {
                if largest, hasLargest = entry.GetLargest(); hasLargest {
                        reply.Parameters = message.Parameters{message.LargestObjectParam(largest.Group, largest.Object)}
                }
        }
        if err := sub.WriteMessage(reply); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "REQUEST_UPDATE_OK write failed",
                        slog.String("err", err.Error()))
                return
        }

        // §9.2: a 0→1 Forward flip resumes paused upstreams.
        if prevForward == 0 && sub.ForwardState() == 1 {
                h.propagateForwardUpstream(ctx, fullName)
        }

        // §10.2.19
        if p, ok := upd.Parameters.Find(message.ParamNewGroupRequest); ok {
                h.propagateNewGroupUpstream(ctx, fullName, p.Varint)
        }

        // §5.1.3: a further fill fetch stream, named by the REQUEST_UPDATE's own
        // Request ID; fills already in flight continue.
        if entry, ok := h.tracks.Get(fullName.Key()); ok {
                if err := h.maybeServeFill(ctx, sub, entry, fullName, upd.RequestID, upd.Parameters,
                        largest, hasLargest); err != nil {
                        h.log.LogAttrs(ctx, slog.LevelDebug, "fill fetch stream not opened",
                                slog.String("err", err.Error()))
                }
        }
}

// propagateNewGroupUpstream forwards a downstream NEW_GROUP_REQUEST to each
// upstream as a REQUEST_UPDATE when §10.2.19 calls for it (see
// [registry.TrackEntry.ConsiderNewGroupRequest]).
func (h *sessionHandler) propagateNewGroupUpstream(
        ctx context.Context,
        fullName track.FullTrackName,
        value uint64,
) {
        entry, ok := h.tracks.Get(fullName.Key())
        if !ok {
                return
        }

        dynamic, err := entry.DynamicGroups()
        if err != nil {
                // Unparseable Track Properties only decline the request here.
                h.log.LogAttrs(ctx, slog.LevelDebug, "NEW_GROUP_REQUEST: unparseable Track Properties",
                        slog.String("err", err.Error()))
                return
        }
        if !entry.ConsiderNewGroupRequest(value, dynamic) {
                return
        }

        for _, up := range entry.CopyUpstream() {
                if !up.IsEstablished() {
                        continue
                }
                resp, err := up.Update(ctx, message.Parameters{message.NewGroupRequestParam(value)})
                if err != nil {
                        h.log.LogAttrs(ctx, slog.LevelDebug, "upstream NEW_GROUP_REQUEST REQUEST_UPDATE failed",
                                slog.String("err", err.Error()))
                        continue
                }
                // §10.2.17 item 1 includes REQUEST_UPDATE_OK.
                saveLargestLocation(entry, resp.Parameters)
        }
}

// propagateForwardUpstream sends REQUEST_UPDATE Forward=1 to each paused
// upstream of fullName (§9.2), saving any LARGEST_OBJECT in the reply.
func (h *sessionHandler) propagateForwardUpstream(ctx context.Context, fullName track.FullTrackName) {
        entry, ok := h.tracks.Get(fullName.Key())
        if !ok {
                return
        }
        for _, up := range entry.CopyUpstream() {
                if up.ForwardState() == 1 || !up.IsEstablished() {
                        continue
                }
                resp, err := up.Update(ctx, message.Parameters{message.ForwardParam(true)})
                if err != nil {
                        h.log.LogAttrs(ctx, slog.LevelDebug, "upstream REQUEST_UPDATE failed",
                                slog.String("err", err.Error()))
                        continue
                }
                up.SetForwardState(1)
                saveLargestLocation(entry, resp.Parameters)
        }
}

// acquireUpstream makes sure msg's track has an Established upstream, reusing
// one or establishing one on demand (see [sessionHandler.subscribeUpstream]),
// and reports whether it reused one and the publisher Seq it looked at. When
// there is none it answers req with REQUEST_ERROR and reports ok=false, after
// holding the SUBSCRIBE for a publisher up to rv's deadline if rv is non-nil
// (§10.2.6).
func (h *sessionHandler) acquireUpstream(
        ctx context.Context,
        req *session.Request,
        msg *message.Subscribe,
        sub *registry.DownstreamSub,
        rv *rendezvous,
) (reused bool, pubSeq uint64, ok bool) {
        fullName := track.FullTrackName{Namespace: msg.Namespace, Name: msg.Name}
        for {
                // Begun before looking, so a publisher arriving meanwhile is not
                // missed; it also cuts short an upstream relay's hold (see
                // subscribeUpstream) once a publisher arrives here.
                lookCtx := ctx
                var look *holdLook
                if rv != nil {
                        look = h.beginHoldLook(ctx, req, fullName)
                        lookCtx = look.ctx
                }
                pubSeq = h.names.Seq()
                e, found := h.tracks.Get(fullName.Key())
                if found && hasEstablishedUpstream(e) {
                        look.end()
                        h.log.LogAttrs(ctx, slog.LevelDebug, "SUBSCRIBE serving from existing upstream")
                        return true, pubSeq, true
                }
                h.log.LogAttrs(ctx, slog.LevelDebug, "SUBSCRIBE no established upstream, trying on-demand",
                        slog.Bool("entry_exists", found))
                var extra message.Parameters
                if p, ok := msg.Parameters.Find(message.ParamNewGroupRequest); ok {
                        extra = message.Parameters{message.NewGroupRequestParam(p.Varint)}
                }
                // §9.2: Forward=1 upstream only if some downstream forwards; sub
                // is not on the entry yet, so it is checked directly.
                wantForward := sub.ForwardState() == 1 || anyDownstreamForwards(e)
                _, established, err := h.subscribeUpstream(lookCtx, fullName, extra, wantForward, rv)
                if established {
                        look.end()
                        return false, pubSeq, true
                }
                // A relay draining this session holds nothing (§10.4).
                if look != nil && !goingAway(h.sess) && (awaitsPublisher(err) || lookCtx.Err() != nil) {
                        arrived := look.wait(rv.deadline)
                        look.end()
                        if arrived {
                                continue
                        }
                        if ctx.Err() != nil || req.Stream.Context().Err() != nil {
                                return false, 0, false // the subscriber or the session is gone
                        }
                        h.log.LogAttrs(ctx, slog.LevelInfo, "SUBSCRIBE rejected: no publisher within RENDEZVOUS_TIMEOUT",
                                slog.String("namespace", fmt.Sprintf("%v", msg.Namespace)),
                                slog.String("name", string(msg.Name)),
                                slog.Uint64("request_id", msg.RequestID))
                        // §10.2.6: "If the timeout expires without a publisher, the relay
                        // SHOULD respond with REQUEST_ERROR with error code TIMEOUT."
                        _ = req.RejectError(moqt.RequestTimeout, "relay: no publisher within RENDEZVOUS_TIMEOUT")
                        return false, 0, false
                }
                look.end()
                if err != nil {
                        h.log.LogAttrs(ctx, slog.LevelInfo, "SUBSCRIBE rejected: upstream subscribe failed",
                                slog.String("namespace", fmt.Sprintf("%v", msg.Namespace)),
                                slog.String("name", string(msg.Name)),
                                slog.Uint64("request_id", msg.RequestID),
                                slog.String("err", err.Error()))
                        rej := upstreamRejection(err)
                        rej.Reason = "relay: no upstream for track: " + err.Error()
                        _ = req.Reject(rej)
                        return false, 0, false
                }
                h.log.LogAttrs(ctx, slog.LevelInfo, "SUBSCRIBE rejected: no publisher for namespace",
                        slog.String("namespace", fmt.Sprintf("%v", msg.Namespace)),
                        slog.String("name", string(msg.Name)),
                        slog.Uint64("request_id", msg.RequestID))
                // §10.2.6: without RENDEZVOUS_TIMEOUT, or with 0, "The relay MUST
                // immediately return REQUEST_ERROR with error code DOES_NOT_EXIST".
                _ = req.RejectError(moqt.RequestDoesNotExist, "relay: no publisher for namespace")
                return false, 0, false
        }
}

// rendezvous is a SUBSCRIBE held for a publisher (§10.2.6).
type rendezvous struct {
        deadline time.Time
        // asked are the sessions tried for the track during the hold, those
        // already serving it included; one that answered is not asked again when
        // another publisher arrives.
        asked map[*session.Session]bool
}

// newRendezvous returns the hold ps's RENDEZVOUS_TIMEOUT asks for, capped at
// [Config.MaxRendezvousTimeout] ("The relay MAY use a shorter timeout than
// requested", §10.2.6), or nil for none: absent, 0, or a cap of 0.
func (h *sessionHandler) newRendezvous(ps message.Parameters) *rendezvous {
        p, ok := ps.Find(message.ParamRendezvousTimeout)
        if !ok || p.Varint == 0 || h.maxRendezvous <= 0 {
                return nil
        }
        d := h.maxRendezvous
        if p.Varint < uint64(d.Milliseconds()) { //nolint:gosec // G115: maxRendezvous is positive.
                d = message.MillisecondTimeout(p.Varint)
        }
        return &rendezvous{deadline: time.Now().Add(d), asked: make(map[*session.Session]bool)}
}

// errPublisherArrived ends a [holdLook] once a publisher arrives.
var errPublisherArrived = errors.New("relay: publisher arrived")

// holdLook is one look for a publisher during a hold: ctx ends once a
// publisher arrives for the track (errPublisherArrived), or the subscriber
// cancels the SUBSCRIBE, or the session ends.
type holdLook struct {
        ctx     context.Context
        cancel  context.CancelCauseFunc
        stopped func()
}

func (h *sessionHandler) beginHoldLook(
        ctx context.Context,
        req *session.Request,
        fullName track.FullTrackName,
) *holdLook {
        trackArrived, stopTrack := h.tracks.AwaitUpstream(fullName.Key())
        nsArrived, stopNS := h.names.AwaitPublisher(fullName.Namespace)
        lookCtx, cancel := context.WithCancelCause(ctx)
        go func() {
                select {
                case <-trackArrived:
                        cancel(errPublisherArrived)
                case <-nsArrived:
                        cancel(errPublisherArrived)
                case <-req.Stream.Context().Done():
                        cancel(nil)
                case <-lookCtx.Done():
                }
        }()
        return &holdLook{ctx: lookCtx, cancel: cancel, stopped: func() { stopTrack(); stopNS() }}
}

// wait blocks until the look ends or deadline passes, and reports whether a
// publisher arrived.
func (l *holdLook) wait(deadline time.Time) bool {
        timer := time.NewTimer(time.Until(deadline))
        defer timer.Stop()
        select {
        case <-l.ctx.Done():
        case <-timer.C:
        }
        return errors.Is(context.Cause(l.ctx), errPublisherArrived)
}

// end releases the look; a nil look is none.
func (l *holdLook) end() {
        if l == nil {
                return
        }
        l.cancel(nil)
        l.stopped()
}

// awaitsPublisher reports whether err, from [sessionHandler.subscribeUpstream],
// leaves the track without a current publisher: none matched (nil), or every
// one that did answered DOES_NOT_EXIST, TIMEOUT for an upstream relay's own
// hold, or GOING_AWAY for its own drain, failed at the transport without an
// answer (see [transportFailure]), or is draining (§10.4). subscribeUpstream
// reports such an error only when no other kind occurred.
func awaitsPublisher(err error) bool {
        if err == nil || errors.Is(err, errGoingAway) || transportFailure(err) {
                return true
        }
        rej, ok := errors.AsType[*session.RequestRejectedError](err)
        return ok && (rej.Code == moqt.RequestDoesNotExist || rej.Code == moqt.RequestTimeout ||
                rej.Code == moqt.RequestGoingAway)
}

// transportFailure reports whether err is a candidate's request failing at
// the transport without any answer about the track: a stream reset, a FIN,
// or the session ending. Not a request the relay could not open for want of
// stream credit (session.ErrNoStreamCredit), which says the publisher is
// there, only busy.
func transportFailure(err error) bool {
        if errors.Is(err, errGoingAway) || isTrackPropertiesErr(err) || errors.Is(err, session.ErrNoStreamCredit) {
                return false
        }
        _, rejected := errors.AsType[*session.RequestRejectedError](err)
        return !rejected
}

// preferCandidateErr returns whichever of err and last, two candidates'
// errors, the subscriber is answered with (see [upstreamRejection]), so the
// answer does not depend on the order candidates answer in: the higher
// ranked ([candidateErrRank]), and of equal rank the one allowing the soonest
// retry, "SHOULD NOT be retried" (Retry Interval 0, §10.6.2) only when both
// say so.
func preferCandidateErr(last, err error) error {
        if r, l := candidateErrRank(err), candidateErrRank(last); r != l {
                if r > l {
                        return err
                }
                return last
        }
        if a, b := candidateRetry(err), candidateRetry(last); a != b {
                if a != 0 && (b == 0 || a < b) {
                        return err
                }
                return last
        }
        // Of equal rank and retry, the answers differ only among "any other
        // refusal" (rank 4): prefer a specific code to INTERNAL_ERROR, then the
        // lower code.
        a, b := upstreamRejection(err).Code, upstreamRejection(last).Code
        if a != b && (b == moqt.RequestInternalError || (a != moqt.RequestInternalError && a < b)) {
                return err
        }
        return last
}

// candidateRetry is the Retry Interval err is answered with: for a
// GOING_AWAY without one, or a request not opened for want of stream credit,
// the most the relay's own jittered one can be (see [upstreamRejection]), else
// the upstream's.
func candidateRetry(err error) uint64 {
        if errors.Is(err, errGoingAway) {
                return goingAwayRetry
        }
        if errors.Is(err, session.ErrNoStreamCredit) {
                return excessiveLoadRetryMax
        }
        rej, ok := errors.AsType[*session.RequestRejectedError](err)
        if !ok {
                return 0
        }
        if rej.Code == moqt.RequestGoingAway && rej.RetryInterval == 0 {
                return goingAwayRetry
        }
        return rej.RetryInterval
}

// candidateErrRank orders candidates' errors for [preferCandidateErr]: an
// unknown Mandatory Track Property (§2.5.1: UNSUPPORTED_EXTENSION, a MUST),
// then Track Properties that do not parse, then any other refusal, which ends
// a RENDEZVOUS_TIMEOUT hold (see awaitsPublisher), then the answers saying
// the track has no publisher yet, the most actionable first: GOING_AWAY,
// which says to retry (§10.6.2), then TIMEOUT, then DOES_NOT_EXIST and a
// request that failed at the transport. GOING_AWAY is taken to outrank
// §10.2.6's DOES_NOT_EXIST for "no publisher is available": the publisher is
// known, only draining.
func candidateErrRank(err error) int {
        if err == nil {
                return 0
        }
        if _, ok := errors.AsType[*session.ErrUnsupportedMandatoryTrackProperty](err); ok {
                return 6
        }
        if isTrackPropertiesErr(err) {
                return 5
        }
        if !awaitsPublisher(err) {
                return 4
        }
        if errors.Is(err, errGoingAway) {
                return 3
        }
        rej, _ := errors.AsType[*session.RequestRejectedError](err)
        if rej != nil && rej.Code == moqt.RequestGoingAway {
                return 3
        }
        if rej != nil && rej.Code == moqt.RequestTimeout {
                return 2
        }
        return 1 // DOES_NOT_EXIST, or no REQUEST_ERROR at all
}

// subscribeUpstream subscribes fullName on every matching source (§9.5):
// each local publisher of a covering namespace and each remote relay
// Discovery resolves (capped by Config.UpstreamFanIn), skipping sessions
// already subscribed. It returns (entry, true, nil) when any upstream was
// established, (nil, false, nil) when there is no source, and
// (nil, false, err) when every candidate failed. extra is added to each
// upstream SUBSCRIBE.
func (h *sessionHandler) subscribeUpstream(
        ctx context.Context,
        fullName track.FullTrackName,
        extra message.Parameters,
        wantForward bool,
        rv *rendezvous,
) (*registry.TrackEntry, bool, error) {
        // Never subscribe twice on one session, nor, while rv holds the
        // SUBSCRIBE, on one asked before. The requester's own session is a
        // candidate like any other: "An endpoint MAY SUBSCRIBE to a Track it is
        // publishing ... Such self-subscriptions are identical to subscriptions
        // initiated by other endpoints" (§5.1).
        subscribed := map[*session.Session]bool{}
        remoteExtra := extra
        if rv != nil {
                subscribed = rv.asked
                // §10.2.6: an upstream relay holds it for what is left of the budget.
                if left := time.Until(rv.deadline); left > 0 {
                        remoteExtra = append(slices.Clip(extra), message.RendezvousTimeoutParam(left))
                }
        }
        if entry, ok := h.tracks.Get(fullName.Key()); ok {
                for _, u := range entry.CopyUpstream() {
                        subscribed[u.Session] = true
                }
        }

        var (
                resultEntry *registry.TrackEntry
                anyEstab    bool
                lastErr     error
                retry       []*session.Session
        )
        establish := func(sess *session.Session, src string, params message.Parameters) {
                if subscribed[sess] || ctx.Err() != nil {
                        return
                }
                subscribed[sess] = true // even on failure: don't retry the same source here
                // Hold the claim when free, so a late-publisher SUBSCRIBE skips. Never
                // wait on another holder: its SUBSCRIBE may fail for its own reasons.
                release, claimed := h.tracks.ClaimUpstream(sess, fullName.Key())
                if claimed {
                        defer release()
                } else if sess == h.sess {
                        // Deviation (§5.1 "identical"): while a SUBSCRIBE for the track
                        // to the requester is pending, this request may be that SUBSCRIBE
                        // routed back to the relay, and a second one to it would loop
                        // (§6.2). A genuine concurrent self-subscription looks the same
                        // and is declined too.
                        return
                }
                h.log.LogAttrs(ctx, slog.LevelDebug, "subscribeUpstream: issuing upstream SUBSCRIBE",
                        slog.String("source", src))
                entry, _, err := h.subscribeUpstreamOnSession(ctx, sess, fullName, params, wantForward)
                if err != nil {
                        if ctx.Err() != nil {
                                // A held SUBSCRIBE's look was cut short (see
                                // acquireUpstream): sess is asked again on the next.
                                delete(subscribed, sess)
                                return
                        }
                        // A request that failed at the transport said nothing about
                        // the track: sess is asked again on a held SUBSCRIBE's next
                        // look, though not twice in this one.
                        if transportFailure(err) {
                                retry = append(retry, sess)
                        }
                        // Keep going, with the error the subscriber is to get (see
                        // preferCandidateErr).
                        lastErr = preferCandidateErr(lastErr, err)
                        h.log.LogAttrs(ctx, slog.LevelDebug, "subscribeUpstream: candidate failed, continuing",
                                slog.String("source", src), slog.String("err", err.Error()))
                        return
                }
                anyEstab = true
                if resultEntry == nil {
                        resultEntry = entry
                }
        }

        publishers := h.names.MatchPublishers(fullName.Namespace)
        h.log.LogAttrs(ctx, slog.LevelDebug, "subscribeUpstream: namespace registry lookup",
                slog.String("namespace", fmt.Sprintf("%v", fullName.Namespace)),
                slog.Int("publishers_found", len(publishers)))
        for _, pub := range publishers {
                establish(pub.Session, "local-publisher", extra)
        }

        remotes, draining := h.upstreams.resolveUpstreams(ctx, fullName.Namespace)
        // A draining relay was sent no request (§10.4); it answers as a draining
        // publisher would, ranked with the other candidates' errors.
        if draining {
                lastErr = preferCandidateErr(lastErr, errGoingAway)
        }
        for _, remote := range remotes {
                establish(remote, "discovery-remote", remoteExtra)
        }

        for _, sess := range retry {
                delete(subscribed, sess)
        }
        if anyEstab {
                return resultEntry, true, nil
        }
        logAttrs := []slog.Attr{
                slog.String("namespace", fmt.Sprintf("%v", fullName.Namespace)),
                slog.String("name", string(fullName.Name)),
                slog.Int("local_publishers", len(publishers)),
                slog.Int("remote_candidates", len(remotes)),
        }
        if lastErr != nil {
                logAttrs = append(logAttrs, slog.String("last_err", lastErr.Error()))
        }
        h.log.LogAttrs(ctx, slog.LevelInfo, "subscribeUpstream: no upstream established", logAttrs...)
        return nil, false, lastErr
}

// subscribeUpstreamOnSession issues the upstream SUBSCRIBE on sess and
// registers the resulting [registry.UpstreamSub] on the track entry.
func (h *sessionHandler) subscribeUpstreamOnSession(
        ctx context.Context,
        sess *session.Session,
        fullName track.FullTrackName,
        extra message.Parameters,
        wantForward bool,
) (*registry.TrackEntry, *registry.UpstreamSub, error) {
        if goingAway(sess) {
                return nil, nil, errGoingAway
        }
        // Always Next Object (§5.1.2: StartGroup and StartObject both 0): every
        // downstream SUBSCRIBE aggregates onto this one upstream (§9.4 MAY), and
        // the fanout applies each downstream filter. The upstream passes only
        // Objects after Largest Object as of when it is processed, so on its own
        // it cannot serve a downstream range that starts earlier.
        filter := &message.LocationFilter{Fields: 2}

        params := message.Parameters{message.LocationFilterParam(filter)}
        // §9.2: with no forwarding downstream, pause the upstream (Forward=0).
        if !wantForward {
                params = append(params, message.ForwardParam(false))
        }
        params = append(params, extra...)
        subMsg := &message.Subscribe{
                Namespace:  fullName.Namespace,
                Name:       fullName.Name,
                Parameters: params,
        }
        // Create the entry before Subscribe: the alias resolves as soon as
        // Subscribe returns and streams may already be arriving, which runFanout
        // can only route to an existing entry. Not inside the round trip, which
        // would widen the window streams wait for their alias. handleFetch's
        // trackKnown keeps the empty entry off the wire meanwhile.
        var entryCreated bool
        _, entryCreated = h.tracks.GetOrCreateNew(fullName)

        upstreamStream, err := sess.Subscribe(ctx, subMsg)
        if err != nil {
                if entryCreated {
                        // Don't let unresolved names grow the registry.
                        h.tracks.DeleteIfUnused(fullName)
                }
                return nil, nil, err
        }
        if hook := testHookAfterAliasRegistered.Load(); hook != nil {
                (*hook)(fullName)
        }

        // The Subscription's broker: the publisher may not send REQUEST_UPDATE
        // (§10.9), and the session releases the alias when it ends (§11.1).
        upstreamSub := registry.NewUpstreamSub(h.allocSubID(), sess, upstreamStream, upstreamStream.Broker(),
                upstreamStream.OK.TrackAlias, subMsg.RequestID)
        upstreamSub.SetFilter(filter)
        // Match the Forward=0 sent upstream: NewUpstreamSub starts at 1, and a
        // later §9.2 resume skips upstreams already at 1.
        if !wantForward {
                upstreamSub.SetForwardState(0)
        }
        // Eligible for §9.4 stitch backfill.
        upstreamSub.FetchCapable = true
        // Torn down with its last downstream.
        upstreamSub.OnDemand = true
        entry, _ := h.tracks.AddUpstream(fullName, upstreamSub, registry.WithProperties(upstreamStream.OK.TrackProperties))
        // §10.2.17 item 1; unconditional, see [saveLargestLocation].
        saveLargestLocation(entry, upstreamStream.OK.Parameters)

        // Relay-scoped, on the upstream stream's context: other sessions'
        // subscribers share it (§9.4), so it outlives this handler.
        h.relayGo(func() {
                h.serveUpstreamStream(upstreamStream.Context(), upstreamSub)
                h.tracks.RemoveUpstream(fullName, upstreamSub.ID)
        })

        return entry, upstreamSub, nil
}

// serveUpstreamStream owns all reads on an upstream request stream (the
// relay's SUBSCRIBE, or an accepted PUBLISH) via the sub's
// [session.RequestBroker], which routes §10.9 responses to
// [registry.UpstreamSub.Update]. It returns when the stream ends or ctx is
// cancelled.
//
// Nothing else may read the stream while this runs: a second reader races
// the broker for the §10.9 responses.
func (h *sessionHandler) serveUpstreamStream(ctx context.Context, up *registry.UpstreamSub) {
        err := up.Broker.Serve(ctx, func(m message.Message) bool {
                switch m := m.(type) {
                case *message.PublishDone:
                        // Kept for the downstream PUBLISH_DONE code (§10.12).
                        up.SetPublishDone(m)
                case *message.RequestOK, *message.RequestError:
                        // Serve only hands responses here when no Update was pending.
                        h.log.LogAttrs(ctx, slog.LevelDebug,
                                "unsolicited response on upstream request stream",
                                slog.Uint64("sub_id", up.ID))
                }
                return true
        })
        if err != nil && ctx.Err() == nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "upstream request stream reader ended",
                        slog.Uint64("sub_id", up.ID), slog.String("err", err.Error()))
        }
}

// hasEstablishedUpstream reports whether the entry has an upstream
// subscription in [registry.SubEstablished].
func hasEstablishedUpstream(entry *registry.TrackEntry) bool {
        for _, u := range entry.CopyUpstream() {
                if u.IsEstablished() {
                        return true
                }
        }
        return false
}

// anyDownstreamForwards reports whether a downstream on entry (which may be
// nil) has Forward=1.
func anyDownstreamForwards(entry *registry.TrackEntry) bool {
        return entry != nil && slices.ContainsFunc(entry.CopyDownstream(),
                func(d *registry.DownstreamSub) bool { return d.ForwardState() == 1 })
}

// resolveGroupOrder gives sub, when its request omitted GROUP_ORDER, the
// publisher's preference from entry's Track Properties (§10.2.8: "If omitted
// from SUBSCRIBE or SUBSCRIBE_TRACKS, the publisher's preference from the
// Track is used"; §12.5). Its fills (§10.2.15) and stream scheduling (§7.2)
// then follow it. GROUP_ORDER cannot appear in REQUEST_UPDATE, so this holds
// for the subscription's life.
func resolveGroupOrder(sub *registry.DownstreamSub, entry *registry.TrackEntry) {
        if sub.GroupOrder == 0 {
                sub.SetGroupOrder(uint8(entry.DefaultGroupOrder()))
        }
}

// installSubscribeParams records the subscription parameters present in ps
// (§10.2) on sub, leaving absent ones unchanged. The Largest snapshot is the
// caller's (see [registry.TrackRegistry.AddDownstreamSnapshotLargest]). The
// session has already closed on a value the draft makes session-fatal (see
// [message.Parameters.CheckScope]); an error here is a Range Filter's, which
// is INVALID_FILTER (§5.1.4).
func installSubscribeParams(sub *registry.DownstreamSub, ps message.Parameters) error {
        if filter, _ := message.LocationFilterFromParam(ps); filter != nil {
                sub.SetFilter(filter)
        }
        if p, ok := ps.Find(message.ParamForward); ok {
                sub.SetForwardState(int(p.Byte))
        }

        if p, ok := ps.Find(message.ParamSubscriberPriority); ok {
                sub.SetPriority(p.Byte)
        }
        if p, ok := ps.Find(message.ParamGroupOrder); ok {
                sub.SetGroupOrder(p.Byte)
        }
        // §10.9: absent from a REQUEST_UPDATE (it cannot appear there,
        // §10.2.21), INCLUDE_PROPERTIES "remains unchanged".
        if p, ok := ps.Find(message.ParamIncludeProperties); ok {
                sub.SetIncludeProperties(p.Byte != 0)
        }

        // §10.2.3 / §10.2.4: each timeout separately, so an update of one does
        // not zero ("no timeout", §8) the other.
        timeouts := sub.GetDeliveryTimeouts()
        if p, ok := ps.Find(message.ParamObjectDeliveryTimeout); ok {
                timeouts.Object = message.MillisecondTimeout(p.Varint)
        }
        if p, ok := ps.Find(message.ParamSubgroupDeliveryTimeout); ok {
                timeouts.Subgroup = message.MillisecondTimeout(p.Varint)
        }
        sub.SetDeliveryTimeouts(timeouts)

        // §5.1.4: a named Range Filter type is replaced, others are kept.
        if !slices.ContainsFunc(ps, func(p message.Parameter) bool { return message.IsRangeFilterParam(p.Type) }) {
                return nil
        }
        rf, err := sub.GetRangeFilters().Update(ps)
        if err != nil {
                return err
        }
        if err := rf.Validate(sub.Session.MaxFilterRanges()); err != nil {
                return err
        }
        sub.SetRangeFilters(rf)
        return nil
}

// refuseSubscriptionParams answers a SUBSCRIBE or SUBSCRIBE_TRACKS whose
// Range Filters [installSubscribeParams] rejected: malformed or over the limit,
// INVALID_FILTER (§5.1.4, §10.6).
func (h *sessionHandler) refuseSubscriptionParams(ctx context.Context, req *session.Request, err error) {
        h.log.LogAttrs(ctx, slog.LevelDebug, "subscription range filter rejected",
                slog.String("err", err.Error()))
        _ = req.RejectError(moqt.RequestInvalidFilter, err.Error())
}

// errGoingAway reports a request the relay did not send because of a GOAWAY on
// the session, in either direction (§10.4; see [goingAway]).
var errGoingAway = errors.New("relay: GOAWAY on the session; no new requests on it (§10.4)")

// includeProperties reports whether INCLUDE_PROPERTIES (§10.2.21) asks for
// Track Properties: yes unless it is 0.
func includeProperties(ps message.Parameters) bool {
        p, ok := ps.Find(message.ParamIncludeProperties)
        return !ok || p.Byte != 0
}

// upstreamRejection is the REQUEST_ERROR for a downstream SUBSCRIBE whose
// upstream SUBSCRIBE failed with err.
//
// An unknown Mandatory Track Property is UNSUPPORTED_EXTENSION (§2.5.1);
// unparseable Track Properties make the track malformed (§12.7, §2.4.2), and
// INTERNAL_ERROR answers them, since MALFORMED_TRACK is defined only "In
// response to a FETCH" (§10.6.2). An upstream REQUEST_ERROR code about the
// track or the publisher's load passes through with its Retry Interval
// (§10.6.2); one about the relay's own hop or its Next Object filter, one not
// defined for SUBSCRIBE (MALFORMED_TRACK among them), or an unknown one,
// becomes INTERNAL_ERROR. If the relay ever combines downstream filters
// upstream (§9.4), INVALID_RANGE must pass through too. Any other failure
// reads as DOES_NOT_EXIST.
func upstreamRejection(err error) *session.RequestRejectedError {
        if isTrackPropertiesErr(err) {
                return &session.RequestRejectedError{Code: session.TrackPropertiesRejectCode(err)}
        }
        if errors.Is(err, errGoingAway) {
                // GOING_AWAY: "The endpoint has received a GOAWAY and MAY reject new
                // requests" (§10.6.2); on the relay's own drain, it "has sent or
                // received a GOAWAY" (§3.3.4). The publisher may return, here or
                // elsewhere.
                return &session.RequestRejectedError{
                        Code:          moqt.RequestGoingAway,
                        RetryInterval: retryIntervalAfter(goingAwayRetryAfter),
                }
        }
        if errors.Is(err, session.ErrNoStreamCredit) {
                // A live publisher the relay could not open a request to for want
                // of stream credit: it "cannot process the request at this time"
                // (§10.6.2), and may shortly.
                return &session.RequestRejectedError{
                        Code:          moqt.RequestExcessiveLoad,
                        RetryInterval: retryIntervalAfter(excessiveLoadRetry),
                }
        }
        up, ok := errors.AsType[*session.RequestRejectedError](err)
        if !ok {
                return &session.RequestRejectedError{Code: moqt.RequestDoesNotExist}
        }
        rej := &session.RequestRejectedError{Code: moqt.RequestInternalError, RetryInterval: up.RetryInterval}
        switch up.Code {
        // GOING_AWAY: an upstream relay draining before its GOAWAY reached this
        // one, which answers a draining publisher the same way.
        case moqt.RequestDoesNotExist, moqt.RequestTimeout, moqt.RequestExcessiveLoad,
                moqt.RequestUnsupportedExtension, moqt.RequestGoingAway:
                rej.Code = up.Code
                // Relay policy: a GOING_AWAY saying not to retry (Retry Interval
                // 0, §10.6.2) is taken to speak for the upstream's own draining
                // hop, and the track may be reached another way.
                if up.Code == moqt.RequestGoingAway && up.RetryInterval == 0 {
                        rej.RetryInterval = retryIntervalAfter(goingAwayRetryAfter)
                }
        case moqt.RequestInternalError, moqt.RequestUnauthorized, moqt.RequestNotSupported,
                moqt.RequestMalformedAuthToken, moqt.RequestExpiredAuthToken,
                moqt.RequestInvalidRange, moqt.RequestInvalidFilter, moqt.RequestRedirect,
                moqt.RequestMalformedTrack, moqt.RequestUninterested, moqt.RequestPrefixOverlap,
                moqt.RequestNamespaceTooLarge:
                // about the relay's hop or request, or not a SUBSCRIBE answer at all
        }
        return rej
}

// goingAwayRetryAfter is how long the relay tells a subscriber to wait before
// retrying past a draining upstream; goingAwayRetry is the largest Retry
// Interval its jitter can make of it (see [retryIntervalAfter]).
const (
        goingAwayRetryAfter = time.Second
        goingAwayRetry      = uint64(goingAwayRetryAfter/time.Millisecond) * 3 / 2
)

// excessiveLoadRetryMax is the largest Retry Interval the jitter can make of
// excessiveLoadRetry.
const excessiveLoadRetryMax = uint64(excessiveLoadRetry/time.Millisecond) * 3 / 2

// isTrackPropertiesErr reports whether err is a Track Properties validation
// failure from [session.Session.Subscribe] or [session.Session.Fetch]: an
// unknown Mandatory Track Property, or Track Properties that do not parse.
func isTrackPropertiesErr(err error) bool {
        _, ok := errors.AsType[*session.ErrUnsupportedMandatoryTrackProperty](err)
        return ok || errors.Is(err, session.ErrMalformedTrackProperties)
}

package relay

import (
        "context"
        "errors"
        "log/slog"
        "slices"
        "sync"
        "sync/atomic"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
)

// trackStatusTimeout, when set by a test, replaces trackStatusUpstreamTimeout.
var trackStatusTimeout atomic.Int64

// trackStatusUpstreamTimeout bounds a forwarded TRACK_STATUS's upstream round
// trip, as FILL_TIMEOUT's default bounds a stitch FETCH's (§13.6): a
// candidate that does not answer within it counts as TIMEOUT.
const trackStatusUpstreamTimeout = defaultUpstreamFetchTimeout

// handleTrackStatus implements TRACK_STATUS (§10.15): a metadata-only query for
// a track's Properties and existence, which the relay "treats ... identically
// as if it had received a SUBSCRIBE", without creating a subscription. The
// reply is TRACK_STATUS_OK (a REQUEST_OK, [message.TrackStatusOK]) carrying
// what the relay's SUBSCRIBE_OK would: the Track Properties and §10.2.17
// LARGEST_OBJECT.
//
// A track with an Established subscription is answered from the track
// registry. Otherwise, as SUBSCRIBE would go upstream, TRACK_STATUS is
// forwarded to every candidate SUBSCRIBE would try (§10.15: relays "MAY
// forward TRACK_STATUS to one or more publishers"; see
// [sessionHandler.trackStatusUpstream]), and their answers are combined as a
// SUBSCRIBE_OK's would be: the entry's Track Properties if it has any, else
// the first answer's (§9.6), and the largest LARGEST_OBJECT of all, the
// entry's included (§10.2.17). If none answers OK, the refusal is the one
// SUBSCRIBE would give.
func (h *sessionHandler) handleTrackStatus(ctx context.Context, req *session.Request, msg *message.TrackStatus) {
        if err := h.auth.AuthorizeTrackStatus(ctx, h.sess, msg); err != nil {
                h.rejectAuth(ctx, req, "TrackStatus", err)
                return
        }

        fullName := track.FullTrackName{Namespace: msg.Namespace, Name: msg.Name}
        var (
                properties []byte
                largest    message.Location
                hasLargest bool
        )
        entry, known := h.tracks.Get(fullName.Key())
        if known {
                properties = entry.GetProperties()
                largest, hasLargest = entry.GetLargest()
        }
        if !known || !hasEstablishedUpstream(entry) {
                // §13.1: a forwarded TRACK_STATUS holds upstream requests open, so
                // it counts against the subscription cap while it does.
                if !h.limiter.acquireSub() {
                        h.rejectExcessiveLoad(ctx, req, "subscription")
                        return
                }
                oks, err := h.forwardTrackStatus(ctx, req, fullName)
                h.limiter.releaseSub()
                if len(oks) == 0 {
                        h.rejectTrackStatus(ctx, req, err)
                        return
                }
                properties, largest, hasLargest = mergeTrackStatus(oks, properties, largest, hasLargest)
        }

        reply := &message.TrackStatusOK{}
        // §10.2.21: INCLUDE_PROPERTIES=0 empties the Track Properties only.
        if includeProperties(msg.Parameters) {
                reply.TrackProperties = properties
        }
        // §10.2.17: LARGEST_OBJECT only once Objects have been published.
        if hasLargest {
                reply.Parameters = message.Parameters{message.LargestObjectParam(largest.Group, largest.Object)}
        }
        // AcceptTrackStatus FINs after the reply (§10.15).
        if err := req.AcceptTrackStatus(reply); err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "TRACK_STATUS_OK write failed",
                        slog.String("err", err.Error()))
        }
}

// mergeTrackStatus combines TRACK_STATUS_OKs with what the relay already
// knows of the track, as a SUBSCRIBE_OK's would be: its Track Properties if
// any, else the first answer's (§9.6), and the largest LARGEST_OBJECT
// (§10.2.17).
func mergeTrackStatus(
        oks []*message.TrackStatusOK,
        properties []byte,
        largest message.Location,
        hasLargest bool,
) ([]byte, message.Location, bool) {
        for _, ok := range oks {
                if len(properties) == 0 {
                        properties = ok.TrackProperties
                }
                p, found := ok.Parameters.Find(message.ParamLargestObject)
                if l := (message.Location{Group: p.Group, Object: p.Object}); found && (!hasLargest || largest.Less(l)) {
                        largest, hasLargest = l, true
                }
        }
        return properties, largest, hasLargest
}

// rejectTrackStatus answers a TRACK_STATUS no candidate accepted: with the
// refusal SUBSCRIBE would give for err (see [upstreamRejection]), or
// DOES_NOT_EXIST when there was no candidate.
func (h *sessionHandler) rejectTrackStatus(ctx context.Context, req *session.Request, err error) {
        rej := &session.RequestRejectedError{Code: moqt.RequestDoesNotExist, Reason: "relay: track not known"}
        if err != nil {
                rej = upstreamRejection(err)
                rej.Reason = "relay: no upstream for track: " + err.Error()
        }
        if werr := req.Reject(rej); werr != nil && !errors.Is(werr, context.Canceled) {
                h.log.LogAttrs(ctx, slog.LevelDebug, "TRACK_STATUS reject write failed",
                        slog.String("err", werr.Error()))
        }
}

// testHookTrackStatusJoined, when set by a test, runs once a request has
// joined, or started, the forwarded TRACK_STATUS round for its track.
var testHookTrackStatusJoined atomic.Pointer[func(track.FullTrackName)]

// trackStatusRounds is a relay's forwarded TRACK_STATUS rounds in flight, by
// track; see [sessionHandler.forwardTrackStatus]. Its ctx ends when the relay
// stops.
type trackStatusRounds struct {
        ctx context.Context
        end context.CancelFunc

        mu     sync.Mutex
        rounds map[track.Key]*trackStatusRound
}

// trackStatusRound is one round's result, set before done closes.
type trackStatusRound struct {
        done chan struct{}
        oks  []*message.TrackStatusOK
        err  error
}

func newTrackStatusRounds() *trackStatusRounds {
        ctx, end := context.WithCancel(context.Background())
        return &trackStatusRounds{ctx: ctx, end: end, rounds: make(map[track.Key]*trackStatusRound)}
}

// forwardTrackStatus answers TRACK_STATUS for fullName from a round of
// [sessionHandler.trackStatusUpstream], shared with the other requests for the
// track that arrive while it is in flight (relay policy, so N requests do not
// make N upstream ones). The round runs on a relay-scoped goroutine, bounded
// by its timeout and by the relay stopping, not by any one requester, who
// stops waiting on its own STOP_SENDING.
//
// A request from a session the relay is asking about the track right now may
// be the relay's own TRACK_STATUS routed back to it on that session (§6.2),
// which the round waits on: it runs a round of its own instead, where the loop
// guard skips that session. A loop through other sessions (relay A's request
// reaching it back over a different session) is not detected; such a round
// waits out its timeout.
func (h *sessionHandler) forwardTrackStatus(
        ctx context.Context,
        req *session.Request,
        fullName track.FullTrackName,
) ([]*message.TrackStatusOK, error) {
        waitCtx, cancel := context.WithCancel(ctx)
        defer cancel()
        defer context.AfterFunc(req.Stream.Context(), cancel)()
        key := fullName.Key()
        if h.tracks.RequestPending(message.TypeTrackStatus, h.sess, key) {
                return h.trackStatusUpstream(waitCtx, fullName)
        }

        rs := h.statusRounds
        rs.mu.Lock()
        round := rs.rounds[key]
        if round == nil {
                round = &trackStatusRound{done: make(chan struct{})}
                rs.rounds[key] = round
                // Started from this tracked handler, so Stop joins it.
                h.relayGo(func() {
                        round.oks, round.err = h.trackStatusUpstream(rs.ctx, fullName)
                        rs.mu.Lock()
                        delete(rs.rounds, key)
                        rs.mu.Unlock()
                        close(round.done)
                })
        }
        rs.mu.Unlock()
        if hook := testHookTrackStatusJoined.Load(); hook != nil {
                (*hook)(fullName)
        }
        select {
        case <-round.done:
                return round.oks, round.err
        case <-waitCtx.Done():
                return nil, waitCtx.Err()
        }
}

// trackStatusUpstream forwards TRACK_STATUS for fullName, concurrently, to
// every candidate SUBSCRIBE would try (see [sessionHandler.subscribeUpstream]):
// each local publisher of a covering namespace and each relay Discovery
// resolves. It returns their TRACK_STATUS_OKs in that order and, when there
// are none, the refusal SUBSCRIBE would give ([preferCandidateErr]), nil when there
// was no candidate.
//
// trackStatusUpstreamTimeout bounds the whole round, resolving the Discovery
// candidates included. A draining candidate is sent nothing
// (§10.4). Relay policy, as for SUBSCRIBE and FETCH: the requester's own
// session is skipped while a TRACK_STATUS for the track to it is in flight,
// since this request may be that one routed back, and a second would loop
// (§6.2).
func (h *sessionHandler) trackStatusUpstream(
        ctx context.Context,
        fullName track.FullTrackName,
) ([]*message.TrackStatusOK, error) {
        timeout := trackStatusUpstreamTimeout
        if d := time.Duration(trackStatusTimeout.Load()); d > 0 {
                timeout = d
        }
        upCtx, cancel := context.WithTimeout(ctx, timeout)
        defer cancel()

        key := fullName.Key()
        var lastErr error
        fail := func(err error) {
                lastErr = preferCandidateErr(lastErr, err)
        }
        var candidates []*session.Session
        seen := map[*session.Session]bool{}
        add := func(sess *session.Session) {
                if seen[sess] {
                        return
                }
                seen[sess] = true
                switch {
                case goingAway(sess):
                        fail(errGoingAway)
                case sess == h.sess && h.tracks.RequestPending(message.TypeTrackStatus, sess, key):
                default:
                        candidates = append(candidates, sess)
                }
        }
        for _, pub := range h.names.MatchPublishers(fullName.Namespace) {
                add(pub.Session)
        }
        remotes, draining := h.upstreams.resolveUpstreams(upCtx, fullName.Namespace)
        if draining {
                fail(errGoingAway)
        }
        for _, remote := range remotes {
                add(remote)
        }

        oks := make([]*message.TrackStatusOK, len(candidates))
        errs := make([]error, len(candidates))
        var wg sync.WaitGroup
        for i, sess := range candidates {
                done := h.tracks.BeginRequest(message.TypeTrackStatus, sess, key)
                wg.Go(func() {
                        defer done()
                        ts, err := sess.TrackStatus(upCtx, &message.TrackStatus{Namespace: fullName.Namespace, Name: fullName.Name})
                        if err != nil {
                                if errors.Is(err, context.DeadlineExceeded) {
                                        err = &session.RequestRejectedError{
                                                Code:   moqt.RequestTimeout,
                                                Reason: "relay: TRACK_STATUS timed out",
                                        }
                                }
                                errs[i] = err
                                return
                        }
                        _ = ts.Close()
                        oks[i] = ts.OK
                })
        }
        wg.Wait()
        for _, err := range errs {
                if err != nil {
                        fail(err)
                }
        }
        return slices.DeleteFunc(oks, func(ok *message.TrackStatusOK) bool { return ok == nil }), lastErr
}

package registry

import (
        "sync"

        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// arrivals wakes the waiters on a key once a publisher arrives for it. It is
// guarded by the owning registry's mutex.
type arrivals[K comparable] map[K]*arrival

// arrival is the waiters on one key: ch closes when a publisher arrives.
type arrival struct {
        ch      chan struct{}
        ns      wire.TrackNamespace // for [arrivals.notifyCoveredLocked]
        waiters int
}

// waitLocked registers a waiter on k and returns its channel and the stop
// that unregisters it under mu, the owning registry's mutex.
func (a *arrivals[K]) waitLocked(k K, ns wire.TrackNamespace, mu sync.Locker) (<-chan struct{}, func()) {
        if *a == nil {
                *a = make(arrivals[K])
        }
        w := (*a)[k]
        if w == nil {
                w = &arrival{ch: make(chan struct{}), ns: ns}
                (*a)[k] = w
        }
        w.waiters++
        return w.ch, func() {
                mu.Lock()
                defer mu.Unlock()
                if w.waiters--; w.waiters == 0 && (*a)[k] == w {
                        delete(*a, k)
                }
        }
}

// notifyLocked wakes the waiters on k.
func (a arrivals[K]) notifyLocked(k K) {
        if w := a[k]; w != nil {
                close(w.ch)
                delete(a, k)
        }
}

// notifyCoveredLocked wakes the waiters whose namespace ns is a prefix of.
func (a arrivals[K]) notifyCoveredLocked(ns wire.TrackNamespace) {
        for k, w := range a {
                if w.ns.HasPrefix(ns) {
                        close(w.ch)
                        delete(a, k)
                }
        }
}

package registry

import (
        "sync"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
)

// fetchResponseGrace bounds how long an upstream FETCH response stream waits
// for its requesting reader to register before the router gives up and resets
// it. It only matters when the response data stream is dispatched by the
// upstream session's data loop before the downstream handler has registered:
// the Request ID is known only after [session.Session.Fetch] returns, so a
// fast upstream can race the registration. Generous relative to the in-process
// and LAN round-trips it guards against.
const fetchResponseGrace = 5 * time.Second

// fetchKey identifies an in-flight upstream FETCH by the session it was issued
// on and the Request ID the session assigned.
type fetchKey struct {
        sess  *session.Session
        reqID uint64
}

// FetchRouter rendezvouses upstream FETCH response streams with the downstream
// handler that issued the FETCH. The two sides run on different goroutines —
// the requester (a downstream FETCH handler) calls [FetchRouter.Register] and
// awaits, while the upstream session's data loop calls [FetchRouter.Deliver] —
// and they may arrive in either order, so each side get-or-creates the
// rendezvous and a buffered slot holds the stream until the reader takes it.
//
// One FetchRouter is shared per [Relay] and injected into every session
// handler.
type FetchRouter struct {
        mu      sync.Mutex
        pending map[fetchKey]chan *session.IncomingFetchStream
}

func NewFetchRouter() *FetchRouter {
        return &FetchRouter{pending: make(map[fetchKey]chan *session.IncomingFetchStream)}
}

// chanLocked returns the rendezvous channel for key, creating it if absent.
// created reports whether this call created the entry. The caller holds r.mu.
func (r *FetchRouter) chanLocked(key fetchKey) (ch chan *session.IncomingFetchStream, created bool) {
        ch, ok := r.pending[key]
        if !ok {
                ch = make(chan *session.IncomingFetchStream, 1)
                r.pending[key] = ch
                created = true
        }
        return ch, created
}

// Register reserves the rendezvous for an upstream FETCH the caller is about
// to issue (or just issued) on sess with the assigned reqID. It returns the
// channel the response stream will arrive on and a cleanup func the caller
// MUST defer: cleanup removes the entry and resets any stream that arrived but
// was never consumed (e.g. the caller timed out waiting).
func (r *FetchRouter) Register(
        sess *session.Session,
        reqID uint64,
) (<-chan *session.IncomingFetchStream, func()) {
        key := fetchKey{sess: sess, reqID: reqID}
        r.mu.Lock()
        ch, _ := r.chanLocked(key)
        r.mu.Unlock()

        cleanup := func() {
                r.mu.Lock()
                if cur, ok := r.pending[key]; ok && cur == ch {
                        delete(r.pending, key)
                }
                r.mu.Unlock()
                // Reset a stream that landed after the reader gave up.
                select {
                case s := <-ch:
                        if s != nil {
                                s.Cancel(moqt.StreamResetInternalError)
                        }
                default:
                }
        }
        return ch, cleanup
}

// Deliver hands an upstream FETCH response stream to its waiting reader. It
// reports whether the stream was accepted into the rendezvous. When Deliver
// creates the rendezvous (the response arrived before the reader registered),
// it schedules a grace timer that resets the stream if no reader claims it, so
// a stray response can't leak. It returns false only when a stream is already
// parked for the same key (a duplicate or unexpected response); the caller
// resets the incoming stream in that case.
func (r *FetchRouter) Deliver(sess *session.Session, reqID uint64, stream *session.IncomingFetchStream) bool {
        key := fetchKey{sess: sess, reqID: reqID}
        r.mu.Lock()
        ch, created := r.chanLocked(key)
        r.mu.Unlock()

        select {
        case ch <- stream:
        default:
                return false // a stream is already parked for this key
        }

        if created {
                time.AfterFunc(fetchResponseGrace, func() {
                        r.mu.Lock()
                        cur, ok := r.pending[key]
                        if !ok || cur != ch {
                                r.mu.Unlock()
                                return
                        }
                        delete(r.pending, key)
                        r.mu.Unlock()
                        select {
                        case s := <-ch:
                                if s != nil {
                                        s.Cancel(moqt.StreamResetInternalError)
                                }
                        default:
                        }
                })
        }
        return true
}

package registry

import (
        "context"
        "errors"
        "log/slog"
        "slices"
        "sync"
        "sync/atomic"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
        "github.com/floatdrop/moq-go/pkg/relay/discovery"
)

// PublisherEntry records a single PUBLISH_NAMESPACE advertisement received
// from a publisher (or upstream relay). The relay holds onto the bidi
// Stream because §6.2 / §10.16 require the same stream to stay open for the
// lifetime of the advertisement — that's also where REQUEST_OK / REQUEST_ERROR
// and the eventual cancellation FIN flow.
type PublisherEntry struct {
        // Namespace is the exact tuple the publisher advertised (§2.4.1).
        Namespace wire.TrackNamespace

        // Session is the MOQT session that owns the PUBLISH_NAMESPACE.
        Session *session.Session

        // Stream is the bidi request stream the PUBLISH_NAMESPACE arrived on.
        // The session handler reads further control messages from it and is
        // the owner that closes/cancels it on teardown.
        Stream session.Stream

        // announced is set by [NamespaceRegistry.AnnouncePublisher], making the
        // entry a source for SUBSCRIBE_NAMESPACE subscribers. Guarded by the
        // registry's mu.
        announced bool

        // Seq orders registrations: Seq > [NamespaceRegistry.Seq] read earlier
        // means the publisher registered since.
        Seq uint64
}

// TracksParams are a SUBSCRIBE_TRACKS's parameters, the initial parameters of
// each PUBLISH it causes (§10.20.1). A REQUEST_UPDATE replaces the whole value
// and affects only later PUBLISHes, so a value is never modified once stored.
type TracksParams struct {
        // Params are the parameters as sent, merged with each update.
        Params message.Parameters

        // Forward and GroupOrder are the resolved FORWARD (§10.2.18) and
        // GROUP_ORDER (§10.2.8); GroupOrder is 0 when omitted.
        Forward    bool
        GroupOrder byte

        // RangeFilters are the §5.1.4 Range Filters. nil = no restriction.
        RangeFilters *message.RangeFilterSet
}

// defaultTracksParams is what a SUBSCRIBE_TRACKS with no parameters means.
var defaultTracksParams = &TracksParams{Forward: true}

// SubscriberEntry records a single SUBSCRIBE_NAMESPACE or SUBSCRIBE_TRACKS
// announcement received from a subscriber (or downstream relay). §6.1 says
// these are open-ended subscriptions to a *prefix*: the relay must echo any
// matching PUBLISH_NAMESPACE / PUBLISH back to the subscriber as long as the
// subscription is alive.
type SubscriberEntry struct {
        // prefix is stored only under the registry lock; see
        // [SubscriberEntry.Prefix].
        prefix atomic.Pointer[wire.TrackNamespace]

        // Session is the MOQT session that owns the SUBSCRIBE_NAMESPACE /
        // SUBSCRIBE_TRACKS.
        Session *session.Session

        // Stream is the bidi request stream the subscription arrived on. After
        // the REQUEST_OK, every write to it goes through the entry's queue and
        // [SubscriberEntry.RunWriter].
        Stream session.Stream

        // WantsTracks distinguishes SUBSCRIBE_TRACKS (true: forward PUBLISH
        // messages for matching tracks) from SUBSCRIBE_NAMESPACE (false: only
        // NAMESPACE / NAMESPACE_DONE). The two share a registry entry because
        // they share the prefix-matching semantics — the session handler
        // dispatches on this flag.
        WantsTracks bool

        // tracks is meaningful only when WantsTracks; see
        // [SubscriberEntry.TracksParams].
        tracks atomic.Pointer[TracksParams]

        // ForwardTrack forwards a PUBLISH for a track to a SUBSCRIBE_TRACKS
        // subscriber (§10.20), on the subscriber's session.
        ForwardTrack func(sub *SubscriberEntry, track *TrackEntry)

        // forwarding holds the tracks with a forwarded PUBLISH in flight; see
        // [SubscriberEntry.ClaimForward].
        fwdMu      sync.Mutex
        forwarding map[track.Key]struct{}
        fwdClosed  bool // the entry is unregistered: no more forwards
        // skipped holds, for each track sent a PUBLISH_SKIPPED, the
        // [TrackEntry.UpstreamEpoch] it was sent at; see
        // [SubscriberEntry.NoteSkipped].
        skipped map[track.Key]uint64

        // announced counts the sources of each announced namespace, by wire key
        // (see namespace_state.go). Guarded by the owning registry's mu.
        announced map[string]int

        // outbox holds the messages queued for [SubscriberEntry.RunWriter], in
        // order; outReady wakes it, and closed stops it once the entry is
        // unregistered.
        outMu    sync.Mutex
        outbox   []queuedMessage
        stopped  bool
        outReady chan struct{}
        // writing is the message RunWriter is sending now; guarded by outMu.
        writing queuedMessage
        // writerDone is closed when RunWriter returns.
        writerDone chan struct{}
        closed     chan struct{}
        closeOnce  sync.Once
}

// ClaimForward reserves key while a forwarded PUBLISH for it is being opened
// and registered, and reports whether it was free. It reports false once the
// entry is unregistered (§6.1: "MUST NOT send any further PUBLISH messages"),
// and while the track was skipped at epoch, the [TrackEntry.UpstreamEpoch]
// the caller read, or a later one (see [SubscriberEntry.NoteSkipped]): a
// forward decided on an older view of the track is covered by the skip.
func (e *SubscriberEntry) ClaimForward(key track.Key, epoch uint64) bool {
        e.fwdMu.Lock()
        defer e.fwdMu.Unlock()
        if _, busy := e.forwarding[key]; busy || e.fwdClosed {
                return false
        }
        if at, ok := e.skipped[key]; ok && at >= epoch {
                return false
        }
        if e.forwarding == nil {
                e.forwarding = make(map[track.Key]struct{})
        }
        e.forwarding[key] = struct{}{}
        return true
}

// NoteSkipped records that a PUBLISH_SKIPPED was sent for key at upstream
// epoch: the relay "MUST NOT send a PUBLISH for a Track for a given
// SUBSCRIBE_TRACKS after PUBLISH_SKIPPED has been sent, scoped to a single
// PUBLISH" (§6.1), so the track is not offered again until a new upstream
// PUBLISH or SUBSCRIBE changes the epoch. Call it holding the key's claim.
func (e *SubscriberEntry) NoteSkipped(key track.Key, epoch uint64) {
        e.fwdMu.Lock()
        defer e.fwdMu.Unlock()
        if e.skipped == nil {
                e.skipped = make(map[track.Key]uint64)
        }
        e.skipped[key] = epoch
}

// ReleaseForward frees a key [SubscriberEntry.ClaimForward] reserved.
func (e *SubscriberEntry) ReleaseForward(key track.Key) {
        e.fwdMu.Lock()
        defer e.fwdMu.Unlock()
        delete(e.forwarding, key)
}

// Prefix is the namespace prefix the subscriber asked to be notified about,
// which a TRACK_NAMESPACE_PREFIX update (§10.9.2) changes. A zero-field prefix
// means "all namespaces" (§6.1).
func (e *SubscriberEntry) Prefix() wire.TrackNamespace { return *e.prefix.Load() }

// TracksParams returns the SUBSCRIBE_TRACKS parameters now in effect.
func (e *SubscriberEntry) TracksParams() *TracksParams { return e.tracks.Load() }

// SetTracksParams replaces them, for the PUBLISHes sent from now on.
func (e *SubscriberEntry) SetTracksParams(p *TracksParams) {
        if p == nil {
                p = defaultTracksParams
        }
        e.tracks.Store(p)
}

func (e *SubscriberEntry) close() {
        e.closeOnce.Do(func() {
                close(e.closed)
                e.fwdMu.Lock()
                e.fwdClosed = true
                e.fwdMu.Unlock()
        })
}

// write sends one message on the subscriber's request stream. Only
// [SubscriberEntry.RunWriter] calls it, so writes never interleave.
func (e *SubscriberEntry) write(m message.Message) error {
        return message.Marshal(e.Stream, m)
}

// NamespaceRegistry maintains the relay's view of who advertises which
// namespaces and who has subscribed to which prefixes (§9.5 / §6.1).
//
// Two slices, two queries:
//
//   - publishers, populated by [Register…] / drained by [Unregister…],
//     queried by [MatchPublishers] when an inbound SUBSCRIBE arrives and
//     the relay needs to find upstream(s) for its track.
//   - subscribers, queried by [MatchSubscribers] when an inbound
//     PUBLISH_NAMESPACE / PUBLISH arrives and the relay needs to forward
//     notifications downstream.
//
// Linear scans for both queries are intentional: namespace cardinality is
// far lower than per-track cardinality and these matches are not on the
// object fanout hot path. If profiling later shows otherwise, swapping in a
// trie behind the same API is a contained change.
//
// Cross-instance namespace advertisement is delegated to the optional
// [discovery.DiscoveryStore]; when configured, the registry mirrors
// publish / unpublish events into it.
type NamespaceRegistry struct {
        mu          sync.RWMutex
        publishers  []*PublisherEntry
        subscribers []*SubscriberEntry

        // pubCount refs each distinct namespace by its wire-encoded key.
        // Used to coalesce Discovery PublishNamespace / UnpublishNamespace
        // calls: only the 0→1 and 1→0 transitions fire events, so two
        // publishers advertising the same namespace from the same relay
        // produce one Discovery entry, not two.
        pubCount map[string]int

        // seq is the Seq of the last registered publisher. Guarded by mu.
        seq uint64
        // arrivals are the waiters for a publisher of a namespace, by its wire
        // key; see [NamespaceRegistry.AwaitPublisher]. Guarded by mu.
        arrivals arrivals[string]

        // remote is the namespaces Discovery reports other relays advertise, by
        // wire key; see [NamespaceRegistry.RemoteNamespace]. Guarded by mu.
        remote map[string]*remoteNamespace

        // discovery / relayAddr / log mirror [TrackRegistry] — see those
        // docs. nil discovery means "do not advertise"; failures log at
        // Warn and are not propagated.
        discovery discovery.DiscoveryStore
        relayAddr string
        log       *slog.Logger
}

// NamespaceRegistryOption tweaks a [NamespaceRegistry] at construction.
type NamespaceRegistryOption func(*NamespaceRegistry)

// WithNamespaceDiscovery installs a [discovery.DiscoveryStore] for
// cross-instance namespace advertisement. relayAddr is stamped into
// every [discovery.NamespaceInfo] this registry emits.
func WithNamespaceDiscovery(d discovery.DiscoveryStore, relayAddr string) NamespaceRegistryOption {
        return func(r *NamespaceRegistry) {
                r.discovery = d
                r.relayAddr = relayAddr
        }
}

// WithNamespaceRegistryLogger sets the logger used for Discovery
// warnings.
func WithNamespaceRegistryLogger(l *slog.Logger) NamespaceRegistryOption {
        return func(r *NamespaceRegistry) { r.log = l }
}

// NewNamespaceRegistry constructs an empty registry.
func NewNamespaceRegistry(opts ...NamespaceRegistryOption) *NamespaceRegistry {
        r := &NamespaceRegistry{
                pubCount: make(map[string]int),
                remote:   make(map[string]*remoteNamespace),
                log:      slog.Default(),
        }
        for _, opt := range opts {
                opt(r)
        }
        return r
}

// RegisterPublisher records a publisher's PUBLISH_NAMESPACE. The returned
// pointer is the canonical record; callers should keep it for the eventual
// [NamespaceRegistry.UnregisterPublisher] call rather than rebuilding it.
//
// Duplicate registrations from the same session for the same namespace are
// not deduplicated — §9.3 explicitly permits a relay to see PUBLISH_NAMESPACE
// for the same namespace from multiple publishers, and even a single session
// can in principle re-advertise after withdrawing. Callers that want
// duplicate-suppression policy enforce it at the request-handler layer.
func (r *NamespaceRegistry) RegisterPublisher(
        ns wire.TrackNamespace,
        sess *session.Session,
        stream session.Stream,
) *PublisherEntry {
        entry := &PublisherEntry{Namespace: ns, Session: sess, Stream: stream}
        key := namespaceWireKey(ns)
        r.mu.Lock()
        r.seq++
        entry.Seq = r.seq
        r.publishers = append(r.publishers, entry)
        r.pubCount[key]++
        if r.pubCount[key] == 1 {
                // Under r.mu, so publish/unpublish reach the Discovery store in
                // exactly the order pubCount crossed the 0 boundary — see
                // [NamespaceRegistry.unpublishNamespaceFromDiscovery].
                r.publishNamespaceToDiscovery(ns)
        }
        r.arrivals.notifyCoveredLocked(ns)
        r.mu.Unlock()
        return entry
}

// AwaitPublisher returns a channel closed once a publisher of a namespace
// covering ns next registers, or a remote relay newly advertises one, for a
// SUBSCRIBE held until its track has a publisher (§10.2.6), and the stop that
// ends the wait.
func (r *NamespaceRegistry) AwaitPublisher(ns wire.TrackNamespace) (arrived <-chan struct{}, stop func()) {
        r.mu.Lock()
        defer r.mu.Unlock()
        return r.arrivals.waitLocked(namespaceWireKey(ns), ns, &r.mu)
}

// AnnouncePublisher makes entry a source of its namespace for
// SUBSCRIBE_NAMESPACE subscribers, announcing the namespace to those that had
// no source for it. Call it only once the publisher has its REQUEST_OK;
// registration alone already makes it routable for SUBSCRIBEs.
func (r *NamespaceRegistry) AnnouncePublisher(entry *PublisherEntry) {
        r.mu.Lock()
        defer r.mu.Unlock()
        if !slices.Contains(r.publishers, entry) {
                return // RemoveSession got there first
        }
        entry.announced = true
        r.addSourceLocked(entry.Namespace)
}

// UnregisterPublisher removes a previously registered publisher entry. The
// caller passes the exact pointer that RegisterPublisher returned — this
// avoids any ambiguity when the same (session, namespace) pair has multiple
// concurrent registrations.
//
// Returns true if the entry was found and removed.
func (r *NamespaceRegistry) UnregisterPublisher(entry *PublisherEntry) bool {
        r.mu.Lock()
        before := len(r.publishers)
        r.publishers = slices.DeleteFunc(r.publishers, func(e *PublisherEntry) bool {
                return e == entry
        })
        removed := len(r.publishers) < before
        if removed {
                if entry.announced {
                        r.removeSourceLocked(entry.Namespace)
                }
                key := namespaceWireKey(entry.Namespace)
                r.pubCount[key]--
                if r.pubCount[key] <= 0 {
                        delete(r.pubCount, key)
                        // Under r.mu — see
                        // [NamespaceRegistry.unpublishNamespaceFromDiscovery].
                        r.unpublishNamespaceFromDiscovery(entry.Namespace)
                }
        }
        r.mu.Unlock()
        return removed
}

// RegisterSubscriber records a subscriber's SUBSCRIBE_NAMESPACE (when
// wantsTracks is false) or SUBSCRIBE_TRACKS (when true). params are the
// SUBSCRIBE_TRACKS parameters (§10.20.1), ignored unless wantsTracks; nil
// means none. Returns the canonical pointer for use with
// [NamespaceRegistry.UnregisterSubscriber].
func (r *NamespaceRegistry) RegisterSubscriber(
        prefix wire.TrackNamespace,
        sess *session.Session,
        stream session.Stream,
        wantsTracks bool,
        params *TracksParams,
        forwardTrack func(*SubscriberEntry, *TrackEntry),
) *SubscriberEntry {
        entry := &SubscriberEntry{
                ForwardTrack: forwardTrack,
                Session:      sess,
                Stream:       stream,
                WantsTracks:  wantsTracks,
                outReady:     make(chan struct{}, 1),
                closed:       make(chan struct{}),
                writerDone:   make(chan struct{}),
        }
        entry.prefix.Store(&prefix)
        entry.SetTracksParams(params)
        r.mu.Lock()
        defer r.mu.Unlock()
        r.subscribers = append(r.subscribers, entry)
        if !wantsTracks {
                // §6.1: announce every namespace already known under the prefix,
                // under the same lock that orders later changes to them. Local
                // publishers first, in registration order, then remote-only ones.
                counts, names := r.namespaceSources(prefix)
                entry.announced = counts
                for _, p := range r.publishers {
                        k := namespaceWireKey(p.Namespace)
                        if _, pending := names[k]; pending && p.announced {
                                entry.enqueue(namespaceMessage(p.Namespace, prefix))
                                delete(names, k)
                        }
                }
                for _, ns := range names {
                        entry.enqueue(namespaceMessage(ns, prefix))
                }
        }
        return entry
}

// UnregisterSubscriber removes a previously registered subscriber entry.
// Returns true if the entry was found and removed.
func (r *NamespaceRegistry) UnregisterSubscriber(entry *SubscriberEntry) bool {
        r.mu.Lock()
        defer r.mu.Unlock()
        before := len(r.subscribers)
        r.subscribers = slices.DeleteFunc(r.subscribers, func(e *SubscriberEntry) bool {
                return e == entry
        })
        entry.close()
        return len(r.subscribers) < before
}

// RemoveSession removes every publisher and subscriber entry owned by sess.
// This is the bulk-cleanup path session handlers take when the underlying
// transport dies — they cannot iterate the registry themselves without
// risking a stale view, so the registry does it under its own lock.
//
// Returns the number of publisher entries and subscriber entries removed,
// in that order. The split lets callers distinguish what Discovery
// unpublish calls were warranted by the cleanup.
func (r *NamespaceRegistry) RemoveSession(sess *session.Session) (publishers, subscribers int) {
        r.mu.Lock()
        beforeP := len(r.publishers)

        // Collect the namespaces this session owned so we can decrement
        // pubCount under the same lock — without it we'd lose the
        // "last publisher leaving the relay" signal that Discovery needs.
        var toUnadvertise []wire.TrackNamespace
        for _, e := range r.publishers {
                if e.Session != sess {
                        continue
                }
                key := namespaceWireKey(e.Namespace)
                if e.announced {
                        r.removeSourceLocked(e.Namespace)
                }
                r.pubCount[key]--
                if r.pubCount[key] <= 0 {
                        delete(r.pubCount, key)
                        toUnadvertise = append(toUnadvertise, e.Namespace)
                }
        }
        r.publishers = slices.DeleteFunc(r.publishers, func(e *PublisherEntry) bool {
                return e.Session == sess
        })

        beforeS := len(r.subscribers)
        r.subscribers = slices.DeleteFunc(r.subscribers, func(e *SubscriberEntry) bool {
                if e.Session == sess {
                        e.close()
                        return true
                }
                return false
        })

        // Capture the final lengths under the lock; reading them after
        // Unlock races with concurrent RemoveSession calls.
        pubsRemoved := beforeP - len(r.publishers)
        subsRemoved := beforeS - len(r.subscribers)
        for _, ns := range toUnadvertise {
                // Under r.mu — see
                // [NamespaceRegistry.unpublishNamespaceFromDiscovery].
                r.unpublishNamespaceFromDiscovery(ns)
        }
        r.mu.Unlock()
        return pubsRemoved, subsRemoved
}

// publishNamespaceToDiscovery advertises ns to the Discovery store
// (best-effort). See [TrackRegistry.publishTrackToDiscovery] for the
// rationale around synchronous calls + log-and-swallow errors.
func (r *NamespaceRegistry) publishNamespaceToDiscovery(ns wire.TrackNamespace) {
        if r.discovery == nil {
                return
        }
        ctx, cancel := context.WithTimeout(context.Background(), discoveryCallTimeout)
        defer cancel()
        if err := r.discovery.PublishNamespace(ctx, discovery.NamespaceInfo{
                Prefix:    ns,
                RelayAddr: r.relayAddr,
        }); err != nil && !errors.Is(err, discovery.ErrWithdrawn) {
                // See [TrackRegistry.publishTrackToDiscovery]: a withdrawn store is
                // shutting down, not broken.
                r.log.Warn("discovery: PublishNamespace failed", "err", err.Error(), "namespace", ns)
        }
}

// unpublishNamespaceFromDiscovery is the counterpart called when the last
// publisher of a namespace leaves.
//
// The caller MUST hold r.mu. Both this and [publishNamespaceToDiscovery]
// run under the registry lock so the store receives publish/unpublish in
// exactly the order pubCount crossed 0 — a late unpublish issued after
// releasing r.mu could race a concurrent RegisterPublisher's publish and
// erase the re-advertised namespace's record. The Discovery call is bounded
// by [discoveryCallTimeout] — and the interface requires backends to honor
// ctx deadlines — so the lock hold is bounded too.
func (r *NamespaceRegistry) unpublishNamespaceFromDiscovery(ns wire.TrackNamespace) {
        if r.discovery == nil {
                return
        }
        ctx, cancel := context.WithTimeout(context.Background(), discoveryCallTimeout)
        defer cancel()
        if err := r.discovery.UnpublishNamespace(ctx, ns, r.relayAddr); err != nil {
                r.log.Warn("discovery: UnpublishNamespace failed", "err", err.Error(), "namespace", ns)
        }
}

// namespaceWireKey serialises a wire.TrackNamespace into a canonical
// byte string suitable for use as a map key. The same trick the
// discovery package uses; we re-implement here to keep the relay
// package free of an internal dependency in case the discovery
// package's helper ever becomes test-only.
func namespaceWireKey(ns wire.TrackNamespace) string {
        w := wire.NewWriter(nil)
        w.TrackNamespace(ns)
        return string(w.Bytes())
}

// Seq returns the Seq of the most recently registered publisher; a publisher
// registered after the call has a larger one.
func (r *NamespaceRegistry) Seq() uint64 {
        r.mu.RLock()
        defer r.mu.RUnlock()
        return r.seq
}

// MatchPublishers returns every publisher entry whose advertised namespace
// is a prefix of (or equal to) ns. This implements the §9.5 rule the
// SUBSCRIBE handler uses:
//
//        "the Relay MUST send a SUBSCRIBE request to each publisher that has
//         published the subscription's namespace or prefix thereof."
//
// Example: a publisher that advertised PUBLISH_NAMESPACE ("video",) matches
// a SUBSCRIBE for ("video", "cam1"). A publisher that advertised
// ("video", "cam1") matches a SUBSCRIBE for ("video", "cam1") but NOT a
// SUBSCRIBE for ("video",) — the stored namespace must be a prefix of, or
// equal to, the queried namespace.
//
// The returned slice is a fresh allocation; callers may iterate it without
// holding the registry lock.
func (r *NamespaceRegistry) MatchPublishers(ns wire.TrackNamespace) []*PublisherEntry {
        r.mu.RLock()
        defer r.mu.RUnlock()
        var out []*PublisherEntry
        for _, e := range r.publishers {
                if ns.HasPrefix(e.Namespace) {
                        out = append(out, e)
                }
        }
        return out
}

// MatchSubscribers returns every subscriber entry whose stored prefix is a
// prefix of (or equal to) ns. This implements the §6.2 / §6.1 forwarding
// rules the PUBLISH_NAMESPACE and PUBLISH handlers use to find which
// downstream subscribers want to be notified of a newly-advertised
// namespace or a newly-published track.
//
// Example: a subscriber that sent SUBSCRIBE_NAMESPACE ("video",) matches a
// PUBLISH_NAMESPACE ("video", "cam1"). A subscriber that sent
// SUBSCRIBE_NAMESPACE ("video", "cam1") matches a PUBLISH_NAMESPACE
// ("video", "cam1") but not ("video",).
//
// Note that a SUBSCRIBE_NAMESPACE with zero fields (§6.1: "the sender is
// interested in all namespaces") matches every PUBLISH_NAMESPACE — that
// case falls out of isPrefixOf naturally.
//
// The returned slice is a fresh allocation; callers may iterate it without
// holding the registry lock.
func (r *NamespaceRegistry) MatchSubscribers(ns wire.TrackNamespace) []*SubscriberEntry {
        r.mu.RLock()
        defer r.mu.RUnlock()
        var out []*SubscriberEntry
        for _, e := range r.subscribers {
                if ns.HasPrefix(e.Prefix()) {
                        out = append(out, e)
                }
        }
        return out
}

// CopyPublishers returns a snapshot of all publisher entries, for callers that
// iterate without holding the registry lock (tests, metrics).
func (r *NamespaceRegistry) CopyPublishers() []*PublisherEntry {
        r.mu.RLock()
        defer r.mu.RUnlock()
        out := make([]*PublisherEntry, len(r.publishers))
        copy(out, r.publishers)
        return out
}

// CopySubscribers returns a snapshot of all subscriber entries. See
// [NamespaceRegistry.CopyPublishers].
func (r *NamespaceRegistry) CopySubscribers() []*SubscriberEntry {
        r.mu.RLock()
        defer r.mu.RUnlock()
        out := make([]*SubscriberEntry, len(r.subscribers))
        copy(out, r.subscribers)
        return out
}

package registry

import (
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
        "github.com/floatdrop/moq-go/pkg/relay/discovery"
)

// A SUBSCRIBE_NAMESPACE subscriber's view of the namespaces announced to it.
//
// §10.18: NAMESPACE_DONE is per namespace, not per source, so each subscriber
// counts the sources (local PUBLISH_NAMESPACEs and Discovery's remote relays)
// of each namespace under its prefix: NAMESPACE on a count's 0→1,
// NAMESPACE_DONE on its 1→0. Counts change only under the registry lock, which
// also orders the messages; [SubscriberEntry.RunWriter] sends them, so no
// stream write happens under the lock.

// remoteNamespace is one namespace Discovery reports other relays advertise.
type remoteNamespace struct {
        ns     wire.TrackNamespace
        relays map[string]struct{}
}

// namespaceSources returns every namespace under prefix with its number of
// sources, keyed by wire key. The caller holds r.mu.
func (r *NamespaceRegistry) namespaceSources(
        prefix wire.TrackNamespace,
) (map[string]int, map[string]wire.TrackNamespace) {
        counts := make(map[string]int)
        names := make(map[string]wire.TrackNamespace)
        for _, p := range r.publishers {
                if p.announced && p.Namespace.HasPrefix(prefix) {
                        k := namespaceWireKey(p.Namespace)
                        counts[k]++
                        names[k] = p.Namespace
                }
        }
        for k, rn := range r.remote {
                if rn.ns.HasPrefix(prefix) {
                        counts[k] += len(rn.relays)
                        names[k] = rn.ns
                }
        }
        return counts, names
}

// addSourceLocked records one more source of ns for every SUBSCRIBE_NAMESPACE
// subscriber whose prefix covers it, announcing ns to those that had none. The
// caller holds r.mu.
func (r *NamespaceRegistry) addSourceLocked(ns wire.TrackNamespace) {
        k := namespaceWireKey(ns)
        for _, s := range r.subscribers {
                prefix := s.Prefix()
                if s.WantsTracks || !ns.HasPrefix(prefix) {
                        continue
                }
                s.announced[k]++
                if s.announced[k] == 1 {
                        s.enqueue(namespaceMessage(ns, prefix))
                }
        }
}

// removeSourceLocked drops one source of ns for every subscriber counting it,
// sending NAMESPACE_DONE to those left with none. The caller holds r.mu.
func (r *NamespaceRegistry) removeSourceLocked(ns wire.TrackNamespace) {
        k := namespaceWireKey(ns)
        for _, s := range r.subscribers {
                if s.WantsTracks || s.announced[k] == 0 {
                        continue
                }
                s.announced[k]--
                if s.announced[k] == 0 {
                        delete(s.announced, k)
                        s.enqueue(namespaceDoneMessage(ns, s.Prefix()))
                }
        }
}

// RemoteNamespace records that the relay at relayAddr started (published) or
// stopped advertising ns, and announces the change to SUBSCRIBE_NAMESPACE
// subscribers. A repeated report changes nothing.
func (r *NamespaceRegistry) RemoteNamespace(ns wire.TrackNamespace, relayAddr string, published bool) {
        k := namespaceWireKey(ns)
        r.mu.Lock()
        defer r.mu.Unlock()
        rn := r.remote[k]
        if published {
                if rn == nil {
                        rn = &remoteNamespace{ns: ns, relays: make(map[string]struct{})}
                        r.remote[k] = rn
                }
                if _, have := rn.relays[relayAddr]; have {
                        return
                }
                rn.relays[relayAddr] = struct{}{}
                r.addSourceLocked(ns)
                r.arrivals.notifyCoveredLocked(ns)
                return
        }
        if rn == nil {
                return
        }
        if _, have := rn.relays[relayAddr]; !have {
                return
        }
        delete(rn.relays, relayAddr)
        if len(rn.relays) == 0 {
                delete(r.remote, k)
        }
        r.removeSourceLocked(ns)
}

// UpdatePrefix applies a TRACK_NAMESPACE_PREFIX update (§10.9.2) to e and
// queues ok, the update's REQUEST_OK, in order with e's other messages.
//
// For a SUBSCRIBE_NAMESPACE, namespaces no longer covered are done before ok
// (suffixes relative to the old prefix) and newly covered ones announced after
// it, relative to the new one. A SUBSCRIBE_TRACKS only changes which later
// PUBLISHes match.
func (r *NamespaceRegistry) UpdatePrefix(e *SubscriberEntry, prefix wire.TrackNamespace, ok message.Message) {
        r.mu.Lock()
        defer r.mu.Unlock()
        old := e.Prefix()
        e.prefix.Store(&prefix)
        if e.WantsTracks {
                e.enqueue(ok)
                return
        }
        counts, names := r.namespaceSources(prefix)
        // e.announced holds exactly the namespaces with sources under old, so
        // oldNames names each one to be done.
        _, oldNames := r.namespaceSources(old)
        for k := range e.announced {
                if _, still := counts[k]; !still {
                        e.enqueue(namespaceDoneMessage(oldNames[k], old))
                }
        }
        e.enqueue(ok)
        for k, ns := range names {
                if _, had := e.announced[k]; !had {
                        e.enqueue(namespaceMessage(ns, prefix))
                }
        }
        e.announced = counts
}

// Enqueue queues m on e's stream behind every message already queued, for
// replies that must keep their order relative to NAMESPACE / NAMESPACE_DONE.
func (e *SubscriberEntry) Enqueue(m message.Message) {
        e.push(m, false)
}

// Finish queues m as the last message of the request, after which the writer
// FINs the stream (§10.9.1, for a failed REQUEST_UPDATE). The owner then waits
// for [SubscriberEntry.WriterDone] and unregisters e.
func (e *SubscriberEntry) Finish(m message.Message) {
        e.push(m, true)
}

func (e *SubscriberEntry) enqueue(m message.Message) { e.push(m, false) }

// maxQueuedMessages and maxUnsentWait bound a namespace subscription's queue
// (§10.19: a blocked response stream "MAY" be reset). The relay counts a
// stream as blocked when at least maxQueuedMessages are unsent and the oldest
// has waited longer than maxUnsentWait. A slow subscriber still draining a
// large seed can therefore be reset. The check runs only when a message is
// queued. The same bound holds a SUBSCRIBE_TRACKS stream's PUBLISH_SKIPPEDs.
const (
        maxQueuedMessages = 1024
        maxUnsentWait     = time.Second
)

// queuedMessage is a message waiting for RunWriter, with its (monotonic)
// queue time. A nil m is the finish marker.
type queuedMessage struct {
        m  message.Message
        at time.Time
}

// push appends m, then the finish marker when last. Nothing is queued once
// the request is finishing or its stream failed. A push to a blocked stream
// resets both halves with EXCESSIVE_LOAD instead, which unblocks a stuck write
// and ends the request's reader.
func (e *SubscriberEntry) push(m message.Message, last bool) {
        now := time.Now()
        e.outMu.Lock()
        if e.stopped {
                e.outMu.Unlock()
                return
        }
        if e.blockedLocked(now) {
                e.stopped = true
                e.outbox = nil
                e.outMu.Unlock()
                e.Stream.CancelWrite(uint64(moqt.StreamResetExcessiveLoad))
                e.Stream.CancelRead(uint64(moqt.StreamResetExcessiveLoad))
                return
        }
        e.outbox = append(e.outbox, queuedMessage{m: m, at: now})
        if last {
                e.outbox = append(e.outbox, queuedMessage{at: now})
                e.stopped = true
        }
        e.outMu.Unlock()
        select {
        case e.outReady <- struct{}{}:
        default:
        }
}

// blockedLocked reports whether at least maxQueuedMessages are unsent and the
// oldest has waited longer than maxUnsentWait. e.outMu must be held.
func (e *SubscriberEntry) blockedLocked(now time.Time) bool {
        unsent := len(e.outbox)
        oldest := time.Time{}
        if !e.writing.at.IsZero() {
                unsent++
                oldest = e.writing.at
        } else if len(e.outbox) > 0 {
                oldest = e.outbox[0].at
        }
        return unsent >= maxQueuedMessages && now.Sub(oldest) > maxUnsentWait
}

// RunWriter sends e's queued messages in order until e is unregistered, the
// request finishes, a write fails, or the queue bound resets the stream. Its
// owner runs it once, for the subscription's lifetime. It takes one message
// at a time, so what it has not sent stays counted by the queue bound (see
// maxQueuedMessages). After a failed write it also stops reading the stream,
// so a peer's STOP_SENDING-only cancel (§3.3.3) ends the subscription.
func (e *SubscriberEntry) RunWriter() {
        defer close(e.writerDone)
        for {
                select {
                case <-e.closed:
                        return
                case <-e.outReady:
                }
                for {
                        e.outMu.Lock()
                        if len(e.outbox) == 0 {
                                stopped := e.stopped
                                e.outMu.Unlock()
                                if stopped {
                                        // Only the queue bound's reset gets here; the owner may
                                        // be waiting on WriterDone.
                                        return
                                }
                                break
                        }
                        q := e.outbox[0]
                        e.outbox[0] = queuedMessage{}
                        e.outbox = e.outbox[1:]
                        e.writing = q
                        e.outMu.Unlock()
                        if q.m == nil { // the finish marker
                                _ = e.Stream.Close()
                                return
                        }
                        err := e.write(q.m)
                        e.outMu.Lock()
                        e.writing = queuedMessage{}
                        if err != nil {
                                e.stopped = true
                                e.outbox = nil
                                e.outMu.Unlock()
                                e.Stream.CancelRead(uint64(moqt.StreamResetInternalError))
                                return
                        }
                        e.outMu.Unlock()
                }
        }
}

// WriterDone is closed once RunWriter has returned.
func (e *SubscriberEntry) WriterDone() <-chan struct{} { return e.writerDone }

// PublishSkipped queues a PUBLISH_SKIPPED (§10.21) for the track (ns, name) on
// a SUBSCRIBE_TRACKS subscriber, its suffix relative to the prefix in force at
// that point of the stream. It reports false, queuing nothing, when an update
// moved the prefix off ns.
func (r *NamespaceRegistry) PublishSkipped(e *SubscriberEntry, ns wire.TrackNamespace, name []byte) bool {
        r.mu.Lock()
        defer r.mu.Unlock()
        prefix := e.Prefix()
        if !ns.HasPrefix(prefix) {
                return false
        }
        e.enqueue(&message.PublishSkipped{TrackNamespaceSuffix: suffixAfter(ns, prefix), TrackName: name})
        return true
}

// ReplaceRemote makes the remote namespaces exactly those in ads, for a
// Discovery watch's snapshot (see [discovery.DiscoveryStore.WatchNamespaces]).
// Only differences reach subscribers. Sources are added before any is removed,
// so a namespace whose only advertising relay changed is not done and
// announced again.
func (r *NamespaceRegistry) ReplaceRemote(ads []discovery.NamespaceInfo) {
        want := make(map[string]*remoteNamespace)
        for _, ad := range ads {
                k := namespaceWireKey(ad.Prefix)
                w := want[k]
                if w == nil {
                        w = &remoteNamespace{ns: ad.Prefix, relays: make(map[string]struct{})}
                        want[k] = w
                }
                w.relays[ad.RelayAddr] = struct{}{}
        }
        r.mu.Lock()
        defer r.mu.Unlock()
        for k, w := range want {
                rn := r.remote[k]
                if rn == nil {
                        rn = &remoteNamespace{ns: w.ns, relays: make(map[string]struct{})}
                        r.remote[k] = rn
                }
                for addr := range w.relays {
                        if _, have := rn.relays[addr]; !have {
                                rn.relays[addr] = struct{}{}
                                r.addSourceLocked(rn.ns)
                                r.arrivals.notifyCoveredLocked(rn.ns)
                        }
                }
        }
        for k, rn := range r.remote {
                w := want[k]
                for addr := range rn.relays {
                        if w != nil {
                                if _, keep := w.relays[addr]; keep {
                                        continue
                                }
                        }
                        delete(rn.relays, addr)
                        r.removeSourceLocked(rn.ns)
                }
                if len(rn.relays) == 0 {
                        delete(r.remote, k)
                }
        }
}

// namespaceMessage is the NAMESPACE announcing ns to a subscriber of prefix
// (§10.17).
func namespaceMessage(ns, prefix wire.TrackNamespace) *message.Namespace {
        return &message.Namespace{TrackNamespaceSuffix: suffixAfter(ns, prefix)}
}

// namespaceDoneMessage is the NAMESPACE_DONE counterpart (§10.18).
func namespaceDoneMessage(ns, prefix wire.TrackNamespace) *message.NamespaceDone {
        return &message.NamespaceDone{TrackNamespaceSuffix: suffixAfter(ns, prefix)}
}

func suffixAfter(ns, prefix wire.TrackNamespace) wire.TrackNamespace {
        return append(wire.TrackNamespace(nil), ns[len(prefix):]...)
}

package registry

import (
        "slices"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
)

// decodedProperties holds the Track Properties the relay acts on, as opposed
// to the raw Properties block it forwards opaquely downstream per §9.6. The
// values are decoded once when an entry's Properties are set (see
// [TrackEntry.setPropertiesLocked]) so the §10.2.19 / §12 hot paths read a
// cached field instead of re-walking the block.
//
// To cache another property: add a field here, a branch in
// [decodeTrackProperties], and an accessor on [TrackEntry]. The raw block is
// still parsed only once, so a new property costs a branch, not a second pass
// over the bytes.
type decodedProperties struct {
        // parseErr is a structural failure parsing the raw block (a malformed
        // upstream Properties field). It is nil for a well-formed block. When
        // set, no field below is meaningful, so every accessor reports it.
        parseErr error

        // dynamicGroups is DYNAMIC_GROUPS=1 (§12.6). The session closed on a
        // value above 1.
        dynamicGroups bool

        // deliveryTimeouts is the publisher's Track-level OBJECT_DELIVERY_TIMEOUT
        // (§12.2) and SUBGROUP_DELIVERY_TIMEOUT (§12.1) pair. Per §8 a zero value
        // in either dimension means "no timeout", which is also what an absent
        // property decodes to — so the zero DeliveryTimeouts is the correct
        // reading of a track that declares neither.
        deliveryTimeouts message.DeliveryTimeouts

        // groupOrder is DEFAULT_PUBLISHER_GROUP_ORDER (§12.5), zero when omitted.
        // The session closed on a value outside {1, 2}.
        groupOrder message.GroupOrder
}

// decodeTrackProperties parses the raw Track Properties block once and pulls
// out the fields the relay acts on. A structural parse failure short-circuits
// to a parseErr that every accessor surfaces. The session has closed on a
// session-fatal value (§12.5, §12.6) before Properties reach here.
func decodeTrackProperties(raw []byte) decodedProperties {
        pairs, err := message.ParseTrackProperties(raw)
        if err == nil {
                // §12.7: search the contents of Immutable Properties too.
                pairs, err = message.ExpandImmutable(pairs)
        }
        if err != nil {
                return decodedProperties{parseErr: err}
        }
        var d decodedProperties
        for _, kv := range slices.Backward(pairs) { // the mutable value wins (§12.7)
                // Dispatch each property the relay acts on to its decoder. Add a
                // branch here for each new property.
                switch kv.Type {
                case message.PropertyDynamicGroups:
                        d.dynamicGroups = kv.IntVal == 1
                case message.PropertyObjectDeliveryTimeout:
                        d.deliveryTimeouts.Object = message.MillisecondTimeout(kv.IntVal)
                case message.PropertySubgroupDeliveryTimeout:
                        d.deliveryTimeouts.Subgroup = message.MillisecondTimeout(kv.IntVal)
                case message.PropertyDefaultPublisherGroupOrder:
                        d.groupOrder = message.GroupOrderAscending
                        if kv.IntVal == uint64(message.GroupOrderDescending) {
                                d.groupOrder = message.GroupOrderDescending
                        }
                }
        }
        return d
}

// DynamicGroups reports whether the track advertised DYNAMIC_GROUPS=1 (§12.6),
// using the value decoded once when Properties was set. The error is a
// structural failure parsing the Properties block; the §10.2.19 caller then
// declines the NEW_GROUP_REQUEST rather than acting on it.
func (e *TrackEntry) DynamicGroups() (bool, error) {
        e.mu.RLock()
        defer e.mu.RUnlock()
        if e.decoded.parseErr != nil {
                return false, e.decoded.parseErr
        }
        return e.decoded.dynamicGroups, nil
}

// DeliveryTimeouts returns the publisher's Track-level delivery timeouts (§8),
// using the values decoded once when Properties was set. The fanout resolves
// these against each subscriber's own §10.2.3 / §10.2.4 parameters before
// applying them to the subgroup streams it opens.
//
// A malformed Properties block reports the zero pair — "no timeout" — rather
// than an error: unlike §12.6, where acting on a bad value would mean honouring
// a NEW_GROUP_REQUEST the publisher never authorised, the safe reading of an
// undecodable timeout is not to enforce one.
func (e *TrackEntry) DeliveryTimeouts() message.DeliveryTimeouts {
        e.mu.RLock()
        defer e.mu.RUnlock()
        if e.decoded.parseErr != nil {
                return message.DeliveryTimeouts{}
        }
        return e.decoded.deliveryTimeouts
}

// DefaultGroupOrder is the publisher's Group Order preference, its
// DEFAULT_PUBLISHER_GROUP_ORDER Track Property (§12.5): Ascending when it is
// omitted ("If omitted, the publisher's preference is Ascending"), and when
// the Properties are malformed.
func (e *TrackEntry) DefaultGroupOrder() message.GroupOrder {
        e.mu.RLock()
        defer e.mu.RUnlock()
        if e.decoded.parseErr != nil || e.decoded.groupOrder == 0 {
                return message.GroupOrderAscending
        }
        return e.decoded.groupOrder
}

package registry

import (
        "context"
        "errors"
        "fmt"
        "sync"
        "sync/atomic"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
)

// SubState is the lifecycle phase of an upstream or downstream subscription
// as managed by the relay. The relay only ever observes two phases, so the
// model is deliberately just those two:
//
//   - SubEstablished: peer accepted; objects may flow.
//   - SubTerminated:  closed cleanly or by error; no further transitions.
//
// The relay constructs an UpstreamSub / DownstreamSub only once the peer has
// already accepted (it sends SUBSCRIBE_OK for a downstream sub; an upstream
// sub is built from the SUBSCRIBE_OK it received), so there is no observable
// "constructed but not yet established" phase to model — subs are born
// Established and the only transition is the one-way move to Terminated (see
// [Subscription.Terminate]).
//
// The state intentionally does NOT track per-object forwarding decisions.
// Those are fanout concerns expressed via the [message.LocationFilter]
// / Forward-state fields on the concrete [UpstreamSub] / [DownstreamSub]
// structs.
type SubState int

const (
        // SubTerminated is the absorbing state. Either the peer ended the
        // subscription (UNSUBSCRIBE / SUBSCRIBE_DONE / PUBLISH_DONE /
        // SUBSCRIBE_ERROR / PUBLISH_ERROR), the underlying request stream
        // died, or the relay tore the subscription down (auth failure,
        // session close, Stop). Once here, the registry slot can be removed
        // safely by the owning goroutine. It is the zero value so a
        // bare-struct subscription is never mistaken for live; the
        // constructors set SubEstablished explicitly.
        SubTerminated SubState = iota

        // SubEstablished means the subscription is live: objects can be
        // forwarded and REQUEST_UPDATE / UNSUBSCRIBE can be sent.
        SubEstablished
)

// String returns "Established" or "Terminated".
func (s SubState) String() string {
        switch s {
        case SubEstablished:
                return "Established"
        case SubTerminated:
                return "Terminated"
        default:
                return fmt.Sprintf("SubState(%d)", int(s))
        }
}

// Subscription is the embedded common state for [UpstreamSub] and
// [DownstreamSub]. It centralises the mutex, the state field, and the
// terminate latch so the two concrete types only have to add their
// direction-specific fields.
//
// Locking discipline:
//
//   - State, ForwardState, and Filter are guarded by mu.
//   - The Session and Stream references are set once at construction and
//     are read-only thereafter; they are not protected.
//   - Callers that read multiple fields together (e.g. State + Filter
//     during fanout) should hold the lock themselves rather than reading
//     fields individually.
type Subscription struct {
        mu sync.RWMutex

        // state is the current lifecycle phase. Set to SubEstablished by the
        // constructors and moved one-way to SubTerminated via Terminate.
        state SubState

        // ID is unique within the relay process. It serves as the stable
        // removal handle in the Track Registry; see [TrackRegistry.RemoveUpstream].
        // Set once at construction; read-only.
        ID uint64

        // RequestID is the MOQT Request ID (§10.1) of the SUBSCRIBE / PUBLISH
        // that opened this subscription's request stream, kept for identity
        // and diagnostics. (A REQUEST_UPDATE rides the same stream but consumes
        // a fresh ID from the sender's space, §10.1 — the stream, not the ID,
        // names the request being updated.) Set once at construction; read-only.
        RequestID uint64

        // Session is the MOQT session that owns this subscription's request
        // stream. Read-only after construction.
        Session *session.Session

        // Stream is the bidi request stream the SUBSCRIBE / PUBLISH was
        // issued on. The owning goroutine (the session handler's request
        // loop) is the sole writer; the relay reads from it to observe
        // peer-side updates (REQUEST_UPDATE, UNSUBSCRIBE, PUBLISH_DONE,
        // etc.). Read-only after construction; the goroutine that owns the
        // stream closes it.
        Stream session.Stream

        // TrackAlias is the alias the relay assigned to this subscription on
        // its side of the wire (§11.1). For UpstreamSub it is the alias the
        // publisher uses when sending objects to us; for DownstreamSub it is
        // the alias we use when sending objects to the subscriber. Aliases
        // are per-direction, per-session — the fanout remaps between them.
        // Set once at construction; read-only.
        TrackAlias uint64

        // forwardState is the §9.2 Forward flag the peer most recently
        // requested. 1 = deliver objects, 0 = pause delivery. The session
        // handler updates it on REQUEST_UPDATE and the fanout consults it to
        // decide whether to write objects out.
        forwardState int
}

// Terminate moves the subscription to [SubTerminated], returning true on the
// first call and false on every subsequent call. The one-shot latch lets a
// caller run teardown that must happen exactly once (e.g. emitting a single
// PUBLISH_DONE) without coordinating with other goroutines; it is safe to
// call concurrently from any goroutine.
func (s *Subscription) Terminate() bool {
        s.mu.Lock()
        defer s.mu.Unlock()
        if s.state == SubTerminated {
                return false
        }
        s.state = SubTerminated
        return true
}

// State returns the current lifecycle phase.
func (s *Subscription) State() SubState {
        s.mu.RLock()
        defer s.mu.RUnlock()
        return s.state
}

// IsEstablished reports whether the subscription is in [SubEstablished].
// Convenience wrapper for fanout / handler code that only cares whether
// objects may flow.
func (s *Subscription) IsEstablished() bool {
        return s.State() == SubEstablished
}

// IsTerminated reports whether the subscription is in [SubTerminated].
// Convenience wrapper for cleanup paths.
func (s *Subscription) IsTerminated() bool {
        return s.State() == SubTerminated
}

// SetForwardState updates the §9.2 Forward flag. The value is not validated
// here (§10.2.18's FORWARD is 0 or 1).
func (s *Subscription) SetForwardState(v int) {
        s.mu.Lock()
        s.forwardState = v
        s.mu.Unlock()
}

// ForwardState returns the most recently set §9.2 Forward flag.
func (s *Subscription) ForwardState() int {
        s.mu.RLock()
        defer s.mu.RUnlock()
        return s.forwardState
}

// ---------------------------------------------------------------------------
// UpstreamSub / DownstreamSub
// ---------------------------------------------------------------------------

// UpstreamSub represents one subscription the relay holds against a
// publisher: the relay issued a SUBSCRIBE upstream after either a local
// downstream SUBSCRIBE or an explicit PUBLISH / PUBLISH_NAMESPACE from a
// publishing peer.
//
// The Filter is the upstream-side §5.1.2 filter the relay chose for this
// subscription. Per §9.4 the relay typically subscribes upstream with the
// "Largest Object" filter so disparate downstream filters don't churn the
// upstream subscription.
//
// Embedding [Subscription] gives UpstreamSub its state machine, ID, Session,
// Stream, TrackAlias, and ForwardState fields for free.
type UpstreamSub struct {
        Subscription

        // Filter is the §5.1.2 filter the relay used in its upstream
        // SUBSCRIBE. nil means "filter unset" (i.e. the subscription has not
        // been sent yet); once set, the value is owned by the subscription
        // and must not be mutated externally.
        Filter *message.LocationFilter

        // FetchCapable marks an upstream the relay reached via an on-demand
        // SUBSCRIBE (a relay/origin, set in subscribeUpstream) — one expected to
        // answer FETCH, so the FETCH responder may stitch evicted ranges from it.
        // It stays false for a directly-connected leaf publisher, which pushes
        // live objects and does not serve FETCH.
        FetchCapable bool

        // OnDemand marks an upstream subscription the relay itself opened via
        // SUBSCRIBE to serve downstream subscribers (§9.4 aggregation). Such a
        // subscription exists only for its downstreams: when the last one
        // leaves, the registry tears it down ([UpstreamSub.CloseOnDemand]) so
        // the publisher stops streaming into a void. It stays false for
        // PUBLISH-fed upstreams, whose stream is owned by the publisher.
        OnDemand bool

        // Broker owns the request stream's read side (via
        // [session.RequestBroker.Serve], run by the relay's per-upstream reader
        // goroutine) and serializes every relay write on the stream — §10.9
        // REQUEST_UPDATEs via [UpstreamSub.Update] and other control messages
        // via [UpstreamSub.WriteMessage] must not interleave. nil only for
        // literal-constructed test fixtures; [NewUpstreamSub] always builds one.
        Broker *session.RequestBroker

        // done is the PUBLISH_DONE the upstream sent, nil until it does. Guarded
        // by the embedded Subscription's mu.
        done *message.PublishDone
}

// SetPublishDone records the PUBLISH_DONE (§10.12) the upstream ended this
// subscription with.
func (u *UpstreamSub) SetPublishDone(pd *message.PublishDone) {
        u.mu.Lock()
        u.done = pd
        u.mu.Unlock()
}

// publishDone returns what [UpstreamSub.SetPublishDone] recorded.
func (u *UpstreamSub) publishDone() *message.PublishDone {
        u.mu.RLock()
        defer u.mu.RUnlock()
        return u.done
}

// DownstreamDoneCode is the PUBLISH_DONE status code the relay sends its
// subscribers when a track's last upstream ended with upstream (nil: without
// a PUBLISH_DONE, which counts as TRACK_ENDED). §10.12: "SHOULD use a relevant
// status code". A code about the track passes through; one about the relay's
// own upstream subscription, or an unknown one, becomes INTERNAL_ERROR.
func DownstreamDoneCode(upstream *message.PublishDone) moqt.PublishDoneCode {
        if upstream == nil {
                return moqt.PublishDoneTrackEnded
        }
        switch upstream.StatusCode {
        case moqt.PublishDoneTrackEnded, moqt.PublishDoneMalformedTrack:
                return upstream.StatusCode
        case moqt.PublishDoneInternalError, moqt.PublishDoneUnauthorized, moqt.PublishDoneGoingAway,
                moqt.PublishDoneTooFarBehind, moqt.PublishDoneExpired, moqt.PublishDoneUpdateFailed,
                moqt.PublishDoneExcessiveLoad:
                return moqt.PublishDoneInternalError
        }
        return moqt.PublishDoneInternalError // a code this relay does not know
}

// updateResponseTimeout bounds the wait for the §10.9 REQUEST_OK /
// REQUEST_ERROR after Update writes a REQUEST_UPDATE. A conforming peer
// always answers; the bound keeps a peer that never does from wedging the
// dispatch loop the Update call runs on.
const updateResponseTimeout = 5 * time.Second

// Update sends a REQUEST_UPDATE (§10.9) on the upstream request stream and
// awaits the single REQUEST_OK / REQUEST_ERROR the spec mandates, bounded
// by [updateResponseTimeout] (tightened further by any earlier deadline on
// ctx). It delegates to the sub's [session.RequestBroker]; the response is
// delivered by the relay's per-upstream Serve loop. A REQUEST_ERROR is
// surfaced as a [session.RequestRejectedError]; a closed stream as
// [session.ErrRequestStreamClosed].
func (u *UpstreamSub) Update(ctx context.Context, params message.Parameters) (*message.RequestOK, error) {
        if u.Broker == nil {
                return nil, session.ErrRequestStreamClosed
        }
        ctx, cancel := context.WithTimeout(ctx, updateResponseTimeout)
        defer cancel()
        return u.Broker.Update(ctx, params)
}

// WriteMessage marshals a control message onto the upstream request stream
// under the broker's write lock — the same lock that serializes Update's
// REQUEST_UPDATE writes. session.Stream does not serialize concurrent
// writers, so every relay write on this stream after the request is
// accepted must go through here or Update.
func (u *UpstreamSub) WriteMessage(msg message.Message) error {
        if u.Broker == nil {
                return session.ErrRequestStreamClosed
        }
        return u.Broker.WriteMessage(msg)
}

// CloseOnDemand tears down an on-demand upstream subscription after its
// last downstream left by cancelling the request (§5.1: "by sending
// STOP_SENDING"). Idempotent; must be called without registry locks held
// (stream I/O).
func (u *UpstreamSub) CloseOnDemand() {
        u.Cancel(moqt.StreamResetCancelled)
}

// Cancel ends the relay's subscription to this upstream by resetting both
// directions of its request stream with code (§3.3.3). Idempotent; must be
// called without registry locks held (stream I/O).
func (u *UpstreamSub) Cancel(code moqt.StreamResetCode) {
        u.Terminate()
        if u.Broker == nil {
                return
        }
        u.Broker.Close(code)
}

// NewUpstreamSub constructs an UpstreamSub in [SubEstablished] with the given
// identity fields. The relay only builds an UpstreamSub once the upstream
// SUBSCRIBE_OK has arrived (the TrackAlias comes from it), so the
// subscription is live from construction.
//
// requestID is the §10.1 Request ID of the SUBSCRIBE / PUBLISH that opened
// the request stream, recorded for identity and diagnostics.
//
// The Forward State starts at 1: an omitted FORWARD means 1 (§10.2.18), and
// the relay's upstream requests never carry it.
//
// broker owns the request stream's reads, set up by the caller: the
// [session.Subscription]'s own for the relay's SUBSCRIBE, so the session
// releases its Track Alias when it ends (§11.1).
func NewUpstreamSub(
        id uint64,
        sess *session.Session,
        stream session.Stream,
        broker *session.RequestBroker,
        trackAlias, requestID uint64,
) *UpstreamSub {
        return &UpstreamSub{
                state:        SubEstablished,
                ID:           id,
                RequestID:    requestID,
                Session:      sess,
                Stream:       stream,
                TrackAlias:   trackAlias,
                forwardState: 1,
                Broker:       broker,
        }
}

// SetFilter installs the upstream filter. Callers must not mutate the filter
// after handing it over.
func (u *UpstreamSub) SetFilter(f *message.LocationFilter) {
        u.mu.Lock()
        u.Filter = f
        u.mu.Unlock()
}

// GetFilter returns the currently installed filter (or nil).
func (u *UpstreamSub) GetFilter() *message.LocationFilter {
        u.mu.RLock()
        defer u.mu.RUnlock()
        return u.Filter
}

// DownstreamSub represents one subscription the relay holds for a
// subscriber: the relay accepted a SUBSCRIBE from the peer and is now
// responsible for forwarding objects, applying the §5.1.2 filter, honouring
// priority (§7) and group order (§10.2.8), and respecting the Forward flag
// (§9.2).
type DownstreamSub struct {
        Subscription

        // writeMu serializes control-message writes on Stream.
        // session.Stream does not serialize concurrent writers and one
        // Marshal is multiple stream Writes, but two goroutines legitimately
        // write here: the subscriber's request handler (via WriteMessage) and
        // termination (PUBLISH_DONE via TerminateWithPublishDone).
        writeMu sync.Mutex

        // okSent records that the §10.8 SUBSCRIBE_OK went out, or that the
        // relay's own PUBLISH opened the stream (OpenedByPublish); guarded by
        // writeMu. A termination consults it: without a prior response it
        // answers with REQUEST_ERROR instead of PUBLISH_DONE.
        okSent bool

        // Filter is the §5.1.2 filter the subscriber declared. The fanout
        // consults it on every object to decide whether to forward. nil
        // means "no filter installed" — the relay treats the subscription
        // as unfiltered (delivers every object on the track).
        Filter *message.LocationFilter

        // rangeFilters holds the §5.1.4 Range Filters the subscriber declared
        // (Subgroup ID / Object ID / Publisher Priority / Object Property). The
        // fanout ANDs them with Filter per object. nil = no range restriction.
        // Guarded by mu; access via SetRangeFilters / GetRangeFilters.
        rangeFilters *message.RangeFilterSet

        // deliveryTimeouts holds the §10.2.3 / §10.2.4 values the subscriber asked
        // for. The fanout resolves them against the publisher's Track-level pair
        // (§8: the smaller of the two non-zero values) once per subgroup stream it
        // opens. Guarded by mu; access via SetDeliveryTimeouts / GetDeliveryTimeouts.
        deliveryTimeouts message.DeliveryTimeouts

        // LargestAtSubscribe is the largest object the relay had observed on
        // this track when the SUBSCRIBE was accepted (§5.1.2). Filters resolve
        // their start against it, not the live watermark, so the start does not
        // drift as objects arrive.
        LargestAtSubscribe message.Location

        // HasLargestAtSubscribe is false when no objects had been delivered
        // on the track at SUBSCRIBE time. Per §5.1.2, the Next Object and
        // relative-start filters fall back to {0,0} in that case.
        HasLargestAtSubscribe bool

        // Priority is the §7 Subscriber Priority the peer asked for. Lower
        // numeric values mean higher delivery priority. Folded into the §7.2
        // stream-scheduling key by [DownstreamSub.EffectiveStreamPriority].
        // Default per §7 / §10.2.7 is 128 (mid-range), set in NewDownstreamSub.
        Priority uint8

        // GroupOrder is the Group Order preference (§7, §10.2.8), encoded as
        // 0x1 = ascending, 0x2 = descending. It drives the group-order
        // tie-breaker in both reorder-capable paths (FETCH responses) and the
        // §7.2 rule-3 GroupKey of the subgroup-stream scheduling priority.
        // Default per §7.1: the publisher's preference, which the relay fills in
        // from the Track Properties before registering the subscription; zero
        // until then, read as Ascending.
        GroupOrder uint8

        // omitProperties records INCLUDE_PROPERTIES=0; see
        // [DownstreamSub.IncludesProperties].
        omitProperties atomic.Bool

        // streamsOpened counts the data streams opened for this subscription,
        // for the §10.12 Stream Count; streamsOpening counts opens in flight and
        // streamsOpen the opened streams not yet closed. pendingDone is a
        // PUBLISH_DONE waiting for them. Guarded by mu, the same lock as the
        // lifecycle state, so no stream is counted after termination.
        streamsOpened  uint64
        streamsOpening int
        streamsOpen    int
        // datagramsSending counts datagram sends in flight.
        datagramsSending int
        pendingDone      *pendingPublishDone

        // cancelled is done once the request ends; see [DownstreamSub.Cancel].
        cancelled context.Context
        cancel    context.CancelFunc
}

// pendingPublishDone is a termination's PUBLISH_DONE, held until the last of
// the subscription's streams closes.
type pendingPublishDone struct {
        code   moqt.PublishDoneCode
        reason string
}

// BeginStream reserves the open of one data stream for this subscription. It
// returns false once the subscription is terminated: the caller must not open
// the stream. Each true must be paired with one [DownstreamSub.EndStream].
func (d *DownstreamSub) BeginStream() bool {
        d.mu.Lock()
        defer d.mu.Unlock()
        if d.state == SubTerminated {
                return false
        }
        d.streamsOpening++
        return true
}

// EndStream completes a [DownstreamSub.BeginStream], counting the stream if
// it was opened. An opened stream must later be reported to
// [DownstreamSub.StreamClosed].
func (d *DownstreamSub) EndStream(opened bool) {
        d.mu.Lock()
        d.streamsOpening--
        if opened {
                d.streamsOpened++
                d.streamsOpen++
        }
        done, count := d.takeReadyDoneLocked()
        d.mu.Unlock()
        d.sendPublishDone(done, count)
}

// StreamClosed reports that one of the subscription's opened data streams has
// been closed (FIN) or reset.
func (d *DownstreamSub) StreamClosed() {
        d.mu.Lock()
        d.streamsOpen--
        done, count := d.takeReadyDoneLocked()
        d.mu.Unlock()
        d.sendPublishDone(done, count)
}

// BeginDatagram reserves one datagram send for this subscription (§10.12:
// PUBLISH_DONE waits until the sender "has no further datagrams to send"). It
// returns false once the subscription is terminated, and the caller must not
// send. Each true must be paired with one [DownstreamSub.EndDatagram].
func (d *DownstreamSub) BeginDatagram() bool {
        d.mu.Lock()
        defer d.mu.Unlock()
        if d.state == SubTerminated {
                return false
        }
        d.datagramsSending++
        return true
}

// EndDatagram completes a [DownstreamSub.BeginDatagram].
func (d *DownstreamSub) EndDatagram() {
        d.mu.Lock()
        d.datagramsSending--
        done, count := d.takeReadyDoneLocked()
        d.mu.Unlock()
        d.sendPublishDone(done, count)
}

// takeReadyDoneLocked returns and clears the pending PUBLISH_DONE and its
// Stream Count once no stream is open or opening and no datagram is being
// sent; nil otherwise. The caller holds mu.
func (d *DownstreamSub) takeReadyDoneLocked() (*pendingPublishDone, uint64) {
        if d.pendingDone == nil || d.streamsOpen > 0 || d.streamsOpening > 0 || d.datagramsSending > 0 {
                return nil, 0
        }
        done := d.pendingDone
        d.pendingDone = nil
        return done, d.streamsOpened
}

// NewDownstreamSub constructs a DownstreamSub in [SubEstablished]: the relay
// accepts the subscriber's SUBSCRIBE (replying SUBSCRIBE_OK) before building
// the sub, so it is live from construction.
//
// Forward State defaults to 1: an omitted FORWARD means 1 (§10.2.18).
// installSubscribeParams overrides this to 0 only when the peer explicitly
// sends FORWARD=0, and REQUEST_UPDATE can flip it later (§9.2 / §10.9).
func NewDownstreamSub(id uint64, sess *session.Session, stream session.Stream, trackAlias uint64) *DownstreamSub {
        cancelled, cancel := context.WithCancel(context.Background())
        return &DownstreamSub{
                cancelled:    cancelled,
                cancel:       cancel,
                state:        SubEstablished,
                ID:           id,
                Session:      sess,
                Stream:       stream,
                TrackAlias:   trackAlias,
                forwardState: 1,
                // §10.2.7: SUBSCRIBER_PRIORITY defaults to 128 (mid-range) when
                // the peer omits the parameter. installSubscribeParams overrides
                // this only when the SUBSCRIBE / REQUEST_UPDATE carries an explicit
                // value (including an explicit 0, the highest priority).
                Priority: 128,
        }
}

// SetFilter installs the downstream filter. Callers must not mutate the
// filter after handing it over.
func (d *DownstreamSub) SetFilter(f *message.LocationFilter) {
        d.mu.Lock()
        d.Filter = f
        d.mu.Unlock()
}

// GetFilter returns the currently installed filter (or nil).
func (d *DownstreamSub) GetFilter() *message.LocationFilter {
        d.mu.RLock()
        defer d.mu.RUnlock()
        return d.Filter
}

// SetDeliveryTimeouts records the §10.2.3 / §10.2.4 delivery timeouts the
// subscriber asked for. A zero dimension means "no timeout" per §8.
func (d *DownstreamSub) SetDeliveryTimeouts(t message.DeliveryTimeouts) {
        d.mu.Lock()
        d.deliveryTimeouts = t
        d.mu.Unlock()
}

// GetDeliveryTimeouts returns the delivery timeouts the subscriber asked for.
// The zero value means the subscriber requested none, which leaves the
// publisher's Track-level values to stand on their own (§8).
func (d *DownstreamSub) GetDeliveryTimeouts() message.DeliveryTimeouts {
        d.mu.RLock()
        defer d.mu.RUnlock()
        return d.deliveryTimeouts
}

// GetRangeFilters returns the subscription's current §5.1.4 Range Filters, nil
// when it has none.
func (d *DownstreamSub) GetRangeFilters() *message.RangeFilterSet {
        d.mu.RLock()
        defer d.mu.RUnlock()
        return d.rangeFilters
}

// SetRangeFilters installs the subscription's Range Filters (§5.1.4), which the
// fanout ANDs with the Location filter and Forward gate per object (read
// directly under mu by ForwardDecision). nil clears them (no range restriction).
func (d *DownstreamSub) SetRangeFilters(f *message.RangeFilterSet) {
        d.mu.Lock()
        d.rangeFilters = f
        d.mu.Unlock()
}

// SetPriority records the §7 Subscriber Priority. Updated when the peer
// sends a REQUEST_UPDATE.
func (d *DownstreamSub) SetPriority(p uint8) {
        d.mu.Lock()
        d.Priority = p
        d.mu.Unlock()
}

// SetIncludeProperties records INCLUDE_PROPERTIES (§10.2.21), set once from
// the SUBSCRIBE or SUBSCRIBE_TRACKS.
func (d *DownstreamSub) SetIncludeProperties(include bool) { d.omitProperties.Store(!include) }

// IncludesProperties reports whether the subscriber wants Track Properties
// (INCLUDE_PROPERTIES omitted or 1). One that does not also lacks the track's
// DEFAULT_PUBLISHER_PRIORITY (§12.4), so objects carry the priority inline.
func (d *DownstreamSub) IncludesProperties() bool { return !d.omitProperties.Load() }

// SetGroupOrder records the Group Order (§10.2.8), set once from the SUBSCRIBE
// (§7.1: it "cannot be changed").
func (d *DownstreamSub) SetGroupOrder(o uint8) {
        d.mu.Lock()
        d.GroupOrder = o
        d.mu.Unlock()
}

// SetLargestAtSubscribe records the largest-object snapshot captured when
// the subscription was accepted. The fanout feeds this into the §5.1.2
// filter evaluator so LargestObject / NextGroupStart filters resolve
// against a stable subscribe-time anchor rather than the live watermark.
func (d *DownstreamSub) SetLargestAtSubscribe(loc message.Location, hasLargest bool) {
        d.mu.Lock()
        d.LargestAtSubscribe = loc
        d.HasLargestAtSubscribe = hasLargest
        d.mu.Unlock()
}

// EffectiveStreamPriority builds the composite §7.2 scheduling key for one
// subgroup stream of this subscription, which the relay pushes down to the
// transport via [session.PrioritizedSendStream.SetSendPriority].
//
// All four §7.2 rules are encoded in the returned [session.StreamPriority],
// compared lexicographically (lower is higher priority):
//
//  1. Subscriber: this subscription's SUBSCRIBER_PRIORITY (default 128).
//  2. Publisher:  publisherPriority — the byte the subgroup carries
//     (SubgroupHeader.PublisherPriority), already resolved to the track's
//     §12.4 default when the header omits it.
//  3. GroupKey:   groupID with this subscription's GROUP_ORDER applied —
//     bitwise-complemented for Descending so a "lower is higher priority"
//     comparison sends higher Group IDs first.
//  4. Subgroup:   subgroupID — lowest Subgroup ID in a group goes first.
//
// Rules 3+4 only define an ordering between streams of the same request, but
// the transport sees streams from every subscription; §7.2 leaves the
// cross-subscription tie-break implementation-defined, so feeding it the full
// key is conformant and degrades gracefully when the transport projects the
// key onto a coarser knob.
func (d *DownstreamSub) EffectiveStreamPriority(
        publisherPriority uint8,
        groupID, subgroupID uint64,
) session.StreamPriority {
        d.mu.RLock()
        sub := d.Priority
        order := message.GroupOrder(d.GroupOrder)
        d.mu.RUnlock()

        // §7.2 rule 3: Descending order means higher Group IDs are scheduled
        // first. Complementing the Group ID flips the numeric comparison so the
        // same "lower GroupKey is higher priority" rule yields that direction.
        // An unresolved GROUP_ORDER (zero value) is read as Ascending.
        groupKey := groupID
        if order == message.GroupOrderDescending {
                groupKey = ^groupID
        }

        return session.StreamPriority{
                Subscriber: sub,
                Publisher:  publisherPriority,
                GroupKey:   groupKey,
                Subgroup:   subgroupID,
        }
}

// ForwardVerdict is [DownstreamSub.ForwardDecision]'s answer for one Object.
// §11.4.3 allows a FIN only when the stream carried every Object of the
// Subgroup "except any Objects with Locations smaller than the subscription's
// Start Location"; every other skip means the stream must end with a reset.
type ForwardVerdict uint8

const (
        // Forward: enqueue the Object.
        Forward ForwardVerdict = iota
        // SkipBeforeStart: the Object lies before the subscription's Start
        // Location, which still allows a FIN unless the Start was raised past
        // Objects already sent (the caller can tell).
        SkipBeforeStart
        // SkipObject: a filter drops this Object only; a later one in the group
        // may still pass.
        SkipObject
        // SkipPaused: Forward State 0 omits the Object (§5.1.5).
        SkipPaused
        // SkipGroup: the Location filter puts this whole group permanently out of
        // range (§11.4.3); the stream can be reset promptly.
        SkipGroup
        // SkipEnded: the subscription is terminated (§10.12).
        SkipEnded
)

// ForwardDecision decides whether an Object goes to this subscription, under
// one lock acquisition (it runs per Object per subscriber). §5.1.5: "Pass =
// Forward AND Location Filters AND Range Filters". The Location filter uses
// the subscribe-time LargestObject snapshot, not the live watermark.
func (d *DownstreamSub) ForwardDecision(
        group, object, subgroupID uint64, priority uint8, objProps []byte,
) ForwardVerdict {
        d.mu.RLock()
        ended := d.state == SubTerminated
        paused := d.forwardState == 0
        f := d.Filter
        rf := d.rangeFilters
        largest := d.LargestAtSubscribe
        has := d.HasLargestAtSubscribe
        d.mu.RUnlock()

        switch {
        case ended:
                return SkipEnded
        case paused:
                return SkipPaused
        }
        // Location filter first, so its group-exhaustion signal (§11.4.3) governs.
        loc := message.Location{Group: group, Object: object}
        if f != nil && !f.Matches(loc, largest, has) {
                switch {
                case GroupOutOfRange(group, f):
                        return SkipGroup
                case loc.Less(f.Start(largest, has)):
                        return SkipBeforeStart
                }
                return SkipObject
        }
        // Range Filters (§5.1.4): per-object AND; a miss drops the object only.
        if rf != nil && !rf.MatchesObject(subgroupID, object, priority, objProps) {
                return SkipObject
        }
        return Forward
}

// GroupOutOfRange reports whether a Subgroup belonging to group is entirely
// outside the subscription's filter range — i.e. no object in that group can
// ever pass — which makes its in-flight stream eligible for a §11.4.3 reset
// (e.g. after a REQUEST_UPDATE narrowed the End Group or raised the Start
// Location to a higher group). Only the absolute filters carry a fixed range;
// the dynamic (Next Object / relative-start) and unset filters never put a
// whole group permanently out of range, so they return false.
//
// A group equal to the Start Location's group is NOT out of range even when
// the Start Location's Object rose — objects at or above it still pass, so the
// stream stays relevant and object-level filtering handles the boundary.
func GroupOutOfRange(group uint64, f *message.LocationFilter) bool {
        if f == nil {
                return false
        }
        if f.Unfiltered() || f.RelativeStart() || f.NextObject() {
                // Start derived from the largest object (or absent entirely); no fixed
                // range that puts a whole group permanently out of range.
                return false
        }
        if group < f.StartGroup {
                return true
        }
        end, ok := f.End()
        return ok && group > end.Group
}

// Cancel ends the subscription once its request has ended: the subscriber
// cancelled it (§3.3.3), or the session is ending. [DownstreamSub.Cancelled]
// is done, so the streams still open for it are reset ("It MUST reset any open
// streams associated with the SUBSCRIBE", §5.1.1; its fill fetch streams,
// §5.1.3.1). No new PUBLISH_DONE follows; one a termination left pending still
// goes out once those streams close, since it may be the request's only
// response (§3.3.2). Safe to call more than once, and after another
// termination.
func (d *DownstreamSub) Cancel() {
        d.mu.Lock()
        d.state = SubTerminated
        d.mu.Unlock()
        d.cancel()
}

// Cancelled is done once [DownstreamSub.Cancel] has run.
func (d *DownstreamSub) Cancelled() context.Context { return d.cancelled }

// TerminateWithPublishDone ends this downstream subscription (§10.12): the
// relay writes PUBLISH_DONE on the subscriber's request stream and FINs the
// send side. If SUBSCRIBE_OK never went out, it answers with REQUEST_ERROR
// (DOES_NOT_EXIST) instead, and [DownstreamSub.WriteSubscribeOK] then refuses
// the stale OK.
//
// First termination wins; later calls do nothing. Safe to call concurrently,
// and it does no I/O itself (see [DownstreamSub.sendPublishDone]).
//
// §10.12: "MUST NOT send PUBLISH_DONE until it has closed all streams". The
// answer waits for every stream opened or opening to be reported through
// [DownstreamSub.StreamClosed], so its Stream Count is exact.
func (d *DownstreamSub) TerminateWithPublishDone(code moqt.PublishDoneCode, reason string) {
        d.mu.Lock()
        if d.state == SubTerminated {
                d.mu.Unlock()
                return // already terminated
        }
        d.state = SubTerminated
        d.pendingDone = &pendingPublishDone{code: code, reason: reason}
        done, count := d.takeReadyDoneLocked()
        d.mu.Unlock()
        d.sendPublishDone(done, count)
}

// sendPublishDone answers the terminated request on its own goroutine, so a
// subscriber that does not read its request stream delays only its own
// answer. A nil done is a no-op.
func (d *DownstreamSub) sendPublishDone(done *pendingPublishDone, streamCount uint64) {
        if done == nil || d.Stream == nil {
                return
        }
        go func() {
                d.writeMu.Lock()
                defer d.writeMu.Unlock()
                if !d.okSent {
                        // Mirror [session.Request.RejectError] (§3.3.4): an answer that
                        // cannot be written resets the stream as a failure; once written,
                        // nothing reads this stream any more, so the read side is
                        // cancelled.
                        if err := message.Marshal(d.Stream, &message.RequestError{
                                ErrorCode:   moqt.RequestDoesNotExist,
                                ErrorReason: done.reason,
                        }); err != nil {
                                d.Stream.CancelRead(uint64(moqt.StreamResetInternalError))
                                d.Stream.CancelWrite(uint64(moqt.StreamResetInternalError))
                                return
                        }
                        d.Stream.CancelRead(uint64(moqt.StreamResetCancelled))
                } else {
                        _ = message.Marshal(d.Stream, &message.PublishDone{
                                StatusCode:  done.code,
                                StreamCount: streamCount,
                                ErrorReason: done.reason,
                        })
                }
                _ = d.Stream.Close()
        }()
}

// WriteSubscribeOK writes the §10.8 SUBSCRIBE_OK response under the write
// lock and records that the request now has its response, so a later
// termination emits PUBLISH_DONE (§10.12) rather than a second response.
// If a termination won the race first, it returns
// [ErrSubscriptionTerminated] without writing.
func (d *DownstreamSub) WriteSubscribeOK(msg *message.SubscribeOK) error {
        d.writeMu.Lock()
        defer d.writeMu.Unlock()
        if d.IsTerminated() {
                return ErrSubscriptionTerminated
        }
        if err := message.Marshal(d.Stream, msg); err != nil {
                return err
        }
        d.okSent = true
        return nil
}

// OpenedByPublish records that the relay's own PUBLISH opened this
// subscription, so it ends with PUBLISH_DONE (§10.12) rather than
// REQUEST_ERROR. Call it before registering the subscription.
func (d *DownstreamSub) OpenedByPublish() {
        d.writeMu.Lock()
        d.okSent = true
        d.writeMu.Unlock()
}

// EndRefused ends a subscription its subscriber refused (REQUEST_ERROR to the
// relay's PUBLISH, §10.11): it is terminated without a PUBLISH_DONE, even one
// already pending, and the stream is closed in both directions under writeMu.
func (d *DownstreamSub) EndRefused() {
        d.mu.Lock()
        ended := d.state != SubTerminated || d.pendingDone != nil
        d.state = SubTerminated
        d.pendingDone = nil
        d.mu.Unlock()
        if !ended {
                return // already answered
        }
        d.writeMu.Lock()
        defer d.writeMu.Unlock()
        _ = d.Stream.Close()
        d.Stream.CancelRead(uint64(moqt.StreamResetCancelled))
}

// WriteMessage marshals a control message onto the downstream request stream
// under the same lock TerminateWithPublishDone uses. Every relay write on
// this stream after the DownstreamSub is registered must go through here —
// registration makes the sub reachable by registry teardown goroutines, so
// even the SUBSCRIBE_OK reply can otherwise interleave with a PUBLISH_DONE.
//
// A write after termination fails with ErrSubscriptionTerminated: the
// termination's PUBLISH_DONE + FIN is the last thing on this stream, even
// while it waits on the subscription's data streams.
func (d *DownstreamSub) WriteMessage(msg message.Message) error {
        d.writeMu.Lock()
        defer d.writeMu.Unlock()
        if d.IsTerminated() {
                return ErrSubscriptionTerminated
        }
        return message.Marshal(d.Stream, msg)
}

// ErrSubscriptionTerminated is returned by [DownstreamSub.WriteMessage] when
// the subscription has been terminated.
var ErrSubscriptionTerminated = errors.New("registry: subscription terminated")

package registry

import (
        "context"
        "fmt"
        "maps"
        "math"
        "slices"
        "sync"
        "sync/atomic"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/relay/cache"
)

// TrackEntry is the central per-track control block (§9 of
// draft-ietf-moq-transport-20). One entry exists for every track the relay
// currently knows about — created on the first SUBSCRIBE or
// PUBLISH/PUBLISH_NAMESPACE for that track, destroyed when the last upstream
// and the last downstream subscription have both gone.
//
// The Upstream slice is intentionally a list (not a single value) so the
// relay can represent the three cases §9.3 / §9.5.1 explicitly allow:
//
//   - multiple independent publishers claiming the same Full Track Name,
//   - graceful publisher relay switchover where a publisher holds two
//     overlapping sessions while migrating WiFi → cellular,
//   - redundant origins (N-redundant encoders) used for live-media
//     reliability — the relay deduplicates objects by {GroupID, ObjectID}
//     (§9.3) via [TrackEntry.ClaimDelivered] so each object is forwarded
//     downstream exactly once.
//
// Concurrency:
//
//   - TrackEntry.mu is held in read mode for the fanout hot path (every
//     incoming object reads Downstream to dispatch), and in write mode for
//     the rare mutations that add or remove subscriptions or update the
//     largest-object watermark.
//   - The registry-level lock ([TrackRegistry.mu]) protects only the
//     track map; per-entry state lives behind TrackEntry.mu so fanouts on
//     different tracks can run fully in parallel.
type TrackEntry struct {
        mu sync.RWMutex

        // Key is the canonical map identity for this track (§2.4.1).
        Key track.Key

        // FullName retains the unhashed {namespace, name} tuple because some
        // outgoing messages (PUBLISH, SUBSCRIBE_OK, TRACK_STATUS_OK, FETCH_OK)
        // must echo it back verbatim — the Key alone cannot reproduce it.
        FullName track.FullTrackName

        // Properties are the raw Track Properties the relay learned from the
        // upstream publisher (in SUBSCRIBE_OK / PUBLISH / TRACK_STATUS_OK /
        // FETCH_OK). §9.6 requires the relay to forward them on every reply
        // it generates downstream, so they are captured once and replayed.
        // The bytes are the on-the-wire encoding of the Track Properties
        // block; the relay treats them opaquely.
        Properties []byte

        // decoded holds the Track Properties the relay acts on, extracted once
        // from the raw Properties block (which is otherwise forwarded opaquely
        // per §9.6). Properties are immutable for the entry's lifetime (§9.6,
        // first-setter-wins), so decoding happens once when Properties is set —
        // see [decodeTrackProperties] for how to add a field. Set together with
        // Properties by setPropertiesLocked.
        decoded decodedProperties

        // LargestObject is the (Group, Object) high-water mark observed for
        // this track, updated by the fanout path on every incoming object and
        // by upstream control messages that carry a LARGEST_OBJECT value. §10.2.17
        // requires the relay to advertise the *maximum* of these in any
        // outbound message that includes LARGEST_OBJECT.
        //
        // The companion HasLargestObject flag distinguishes "no objects
        // observed yet" from "the first object was published at Location
        // {0, 0}" — §10.2.17 reserves wire-level omission for the former
        // and the in-memory mirror needs the same distinction. Callers
        // SHOULD read via [TrackEntry.GetLargest] rather than touching
        // these fields directly so the lock is honoured.
        LargestObject    message.Location
        HasLargestObject bool

        // Upstream is the set of publisher subscriptions feeding this track.
        // See the type-level comment above for why this is a slice.
        Upstream []*UpstreamSub

        // upstreamEpoch identifies the latest upstream added; see
        // [TrackEntry.UpstreamEpoch]. Guarded by mu.
        upstreamEpoch uint64

        // Downstream is the set of subscriber subscriptions to fan out to.
        Downstream []*DownstreamSub

        // refusals are the publisher registrations that refused a §9.5
        // late-publisher SUBSCRIBE for this track, each with the time it may be
        // asked again (zero: never); see [TrackEntry.NoteRefusal]. Guarded by mu.
        refusals map[*PublisherEntry]time.Time

        // fetches cancel the fetch streams the relay serves on this track, fill
        // fetch streams included; see [TrackEntry.AddFetch]. Guarded by mu.
        fetches map[*fetchCancel]struct{}

        // downstreamGen counts appends to Downstream. The per-object fanout
        // (UpdateLargestAndDetectNew) snapshots it alongside its initial
        // CopyDownstream and skips the O(len(Downstream)) joiner scan on every
        // object whose generation is unchanged — joiners are rare, so the common
        // case becomes a watermark bump with no scan. Bumped only on append
        // (removals introduce no joiner to detect). Guarded by mu.
        downstreamGen uint64

        // Cache is the per-track [cache.ObjectCache] the fanout writes every
        // forwarded object into. It is constructed eagerly when the entry is
        // created (via [TrackRegistry.getOrCreateLocked]) so the fanout
        // never has to nil-check.
        Cache *cache.ObjectCache

        // newGroupOutstanding records whether the relay has a NEW_GROUP_REQUEST
        // (§10.2.19) in flight upstream for this track. It stays outstanding until
        // the Largest Group advances past newGroupReqGroup, at which point the
        // publisher is deemed to have honoured the request. Guarded by mu; see
        // [TrackEntry.ConsiderNewGroupRequest].
        newGroupOutstanding bool
        newGroupReqValue    uint64 // the value last forwarded upstream
        newGroupReqGroup    uint64 // Largest Group at the moment we forwarded

        // sgMu guards subgroups. It is a separate, finer-grained lock than mu so
        // the per-(group, subgroup) fan-out bookkeeping (Acquire/Release on every
        // inbound subgroup stream) does not contend with the mu-guarded control
        // mutations or the per-object UpdateLargestAndDetectNew hot path.
        sgMu sync.Mutex

        // deliveredMu guards the §2.1 dedup ledger below. It is separate from mu so
        // the per-object dedup claim on the fanout hot path does not contend with the
        // mu-guarded control mutations.
        deliveredMu sync.Mutex

        // delivered is the dedup ledger across multiple upstream publishers (§9.3):
        // GroupID → set of Object IDs already forwarded downstream. The first upstream
        // to reach a {GroupID, ObjectID} forwards it; later copies from redundant or
        // lagging peers are dropped (§2.1 — SubgroupID is not part of object
        // identity). It lives on the entry (not on a SharedSubgroup) so peers whose
        // streams do not temporally overlap — e.g. one origin's subgroup FINs before
        // the redundant origin's arrives — still dedup. Memory is bounded by
        // [deliveredGroupWindow]: it holds that many Groups, the lowest ID pruned
        // first, and an Object of a pruned or lower Group is [ClaimAgedOut].
        // deliveredFloor is the lowest Group held, while delivered is non-empty.
        //
        // It also holds the Prior Group and Object ID Gaps (§12.8, §12.9) seen:
        // each Group its Object ID gaps and Group gap value, and groupGaps the Group
        // ID ranges announced absent, pruned with the same window.
        delivered map[uint64]*deliveredGroup
        groupGaps []idRange
        // trackEnd is where the Track ends, its END_OF_TRACK (§2.4.2), if
        // hasTrackEnd; kept past the window.
        trackEnd       message.Location
        hasTrackEnd    bool
        deliveredFloor uint64

        // subgroups holds the shared outbound fan-out state for each
        // (GroupID, SubgroupID) currently being produced by one or more upstreams.
        // §9.3 lets N redundant upstreams feed one track; §2.2 requires that the
        // objects of a single Subgroup go out on exactly ONE downstream stream per
        // subscriber. Sharing this state across every inbound runFanout goroutine
        // (each of which carries one (group, subgroup)) is what lets the relay merge
        // the upstreams into one clean outbound subgroup stream per subscriber
        // instead of one stream per (upstream × subscriber). Created lazily on the
        // first contributor and removed when the last contributor leaves
        // ([TrackEntry.AcquireSubgroup] / [TrackEntry.ReleaseSubgroup]). The payload
        // is parent-managed and opaque here, keeping the fanout's writer type out of
        // the registry layer (same one-way dependency rule as the seen predicate in
        // [TrackEntry.UpdateLargestAndDetectNew]).
        subgroups map[SubgroupKey]*SharedSubgroup
}

// SubgroupKey identifies a Subgroup within a track by its (GroupID, SubgroupID)
// pair (§2.2). It is the merge key for fanning multiple upstream publishers into
// a single downstream stream per subscriber.
type SubgroupKey struct {
        Group    uint64
        Subgroup uint64
}

// SharedSubgroup is the per-(group, subgroup) fan-out state shared across every
// inbound runFanout goroutine producing that Subgroup for a track. The Set field
// is the parent package's writer set (opaque here); Mu guards the parent's
// manipulation of it. refs counts the live inbound contributors and is guarded
// by the owning entry's sgMu, not Mu.
type SharedSubgroup struct {
        // Mu guards the parent-managed Set during writer open/close/deliver. Held
        // across outbound stream I/O, so it is deliberately distinct from the
        // entry's sgMu (which is only ever held for O(1) map/refcount edits).
        Mu sync.Mutex

        // Set is the parent package's writer set for this Subgroup
        // (a *subgroupWriterSet in pkg/relay). Opaque to the registry.
        Set any

        refs int
}

// AcquireSubgroup registers the caller as a contributor to (group, subgroup) on
// this entry, creating the shared state via newSet on the first contributor.
// Returns the shared state and whether this call created it (so the creator can
// open the initial downstream writers; later contributors reuse the existing
// writer set). Every successful Acquire must be balanced by a
// [TrackEntry.ReleaseSubgroup].
func (e *TrackEntry) AcquireSubgroup(key SubgroupKey, newSet func() any) (sg *SharedSubgroup, created bool) {
        e.sgMu.Lock()
        defer e.sgMu.Unlock()
        if e.subgroups == nil {
                e.subgroups = make(map[SubgroupKey]*SharedSubgroup)
        }
        if sg, ok := e.subgroups[key]; ok {
                sg.refs++
                return sg, false
        }
        sg = &SharedSubgroup{Set: newSet(), refs: 1}
        e.subgroups[key] = sg
        return sg, true
}

// CopySubgroups returns the Subgroups currently being fanned out. The caller
// takes each one's Mu itself.
func (e *TrackEntry) CopySubgroups() []*SharedSubgroup {
        e.sgMu.Lock()
        defer e.sgMu.Unlock()
        return slices.Collect(maps.Values(e.subgroups))
}

// deliveredGroupWindow bounds the §2.1 dedup ledger ([TrackEntry.delivered])
// to that many distinct Groups, counted rather than spanned by ID, since Group
// IDs need not be consecutive (§2.3.1, §12.8). The window must comfortably
// exceed any realistic inter-publisher group lag (a redundant origin or relay
// running a few groups behind) while keeping per-track dedup memory bounded.
//
// It assumes Group IDs mostly increase, as §2.3.1 lets a publisher choose:
// the lowest ID is pruned first, and a Group below all 32 held is taken as
// old. A publisher whose IDs decrease loses its Objects past the first 32
// Groups to [ClaimAgedOut].
const deliveredGroupWindow = 32

// Claim is [TrackEntry.ClaimDelivered]'s verdict on an Object.
type Claim int

const (
        // ClaimFresh: the first copy; forward it.
        ClaimFresh Claim = iota
        // ClaimRedundant: a copy of an Object already forwarded, or one inside a
        // gap announced earlier; drop it.
        ClaimRedundant
        // ClaimAgedOut: its Group is older than every Group the window holds, so
        // whether it was forwarded is unknown, and nothing is recorded. The
        // caller forwards it only where it cannot be a repeat.
        ClaimAgedOut
)

// ClaimDelivered is the dedup gate across multiple upstream publishers (§9.3).
// It records (group, object) as forwarded and reports whether the caller is
// the first to do so ([ClaimFresh]) or it was already forwarded by a peer
// upstream ([ClaimRedundant]). The ledger persists on the entry (not on a
// per-Subgroup structure) and is independent of the size-bounded Object Cache,
// so redundant streams that do not temporally overlap, or peers lagging by more
// than the cache capacity, still dedup correctly. Memory is bounded to
// [deliveredGroupWindow] Groups; an Object of an older one is [ClaimAgedOut].
//
// gaps are the Object's Prior Group and Object ID Gaps (§12.8, §12.9), which
// the ledger records for the whole track. An Object inside a gap announced
// earlier is known not to exist, and that is permanent (§2.1): ClaimRedundant,
// since a caching relay "SHOULD NOT cache or forward" it (§9.1). A gap covering an
// Object already received is accepted: an Object may go from existing to not
// existing (§2.1).
//
// Interpretation: §12.8 and §12.9 list both cases as making the track
// malformed, but §2.1 says the first "is not a protocol error and the Track is
// not malformed"; the relay follows §2.1 and §9.1 for both. It also takes
// §9.1's specific SHOULD NOT forward over §9.4's general "MUST NOT reorder or
// drop objects received on a multi-object stream". A cached copy of an Object
// a gap later covers is kept (§9.1 makes updating the cache a MAY). The one gap rule
// it does report, as an error wrapping [session.ErrMalformedTrack], is a Group
// carrying two Prior Group ID Gap values (§12.8). An Object ClaimDelivered
// rejects leaves no state in the ledger, and a duplicate records only its
// gaps, since the caller's §9.1 check may still reject it.
func (e *TrackEntry) ClaimDelivered(o ObjectInfo) (Claim, error) {
        group, object, gaps := o.Group, o.Object, o.Gaps
        e.deliveredMu.Lock()
        defer e.deliveredMu.Unlock()

        g := e.delivered[group]
        if g == nil && len(e.delivered) >= deliveredGroupWindow && group < e.deliveredFloor {
                return ClaimAgedOut, nil
        }
        if gaps.HasGroup && g != nil && g.hasGroupGap && g.groupGap != gaps.Group {
                return ClaimRedundant, fmt.Errorf("%w: Group %d carries Prior Group ID Gaps %d and %d (§12.8)",
                        session.ErrMalformedTrack, group, g.groupGap, gaps.Group)
        }
        if err := e.checkEndsLocked(g, o); err != nil {
                return ClaimRedundant, err
        }
        if e.announcedAbsentLocked(g, group, object) {
                return ClaimRedundant, nil
        }

        if g == nil {
                g = &deliveredGroup{objects: make(map[uint64]struct{})}
                e.addDeliveredGroupLocked(group, g)
        }
        e.recordGapsLocked(g, group, object, gaps)
        if _, ok := g.objects[object]; ok {
                return ClaimRedundant, nil
        }
        e.recordEndsLocked(g, o)
        g.objects[object] = struct{}{}
        return ClaimFresh, nil
}

// addDeliveredGroupLocked adds group's ledger entry g, first pruning the
// lowest Group, and the gaps announced below the new lowest, when the window
// is full. Runs once per Group, so its scan of the window is not per Object.
func (e *TrackEntry) addDeliveredGroupLocked(group uint64, g *deliveredGroup) {
        if e.delivered == nil {
                e.delivered = make(map[uint64]*deliveredGroup, deliveredGroupWindow)
        }
        if len(e.delivered) >= deliveredGroupWindow {
                delete(e.delivered, e.deliveredFloor)
                e.deliveredFloor = math.MaxUint64
                for id := range e.delivered {
                        e.deliveredFloor = min(e.deliveredFloor, id)
                }
                e.groupGaps = slices.DeleteFunc(e.groupGaps, func(r idRange) bool {
                        return r.hi < e.deliveredFloor
                })
        }
        if len(e.delivered) == 0 || group < e.deliveredFloor {
                e.deliveredFloor = group
        }
        e.delivered[group] = g
}

// ObjectInfo is what [TrackEntry.ClaimDelivered] checks of an Object against
// the earlier Objects of the track.
type ObjectInfo struct {
        Group, Object uint64
        // Subgroup is the Object's Subgroup ID, unless Datagram (§11.2.1).
        Subgroup uint64
        Datagram bool
        // Priority is the resolved Publisher Priority (§7).
        Priority uint8
        // Status is the Object Status (§11.2.1.1).
        Status uint64
        // EndOfGroup is a datagram's END_OF_GROUP bit (§11.3.1).
        EndOfGroup bool
        // Gaps are the Object's Prior Group and Object ID Gaps (§12.8, §12.9).
        Gaps message.PriorGaps
}

// end is where a Subgroup, Group or Track ends: its first missing Object ID
// (§2.4.2). A status Object at M ends it at M (§11.2.1.1); a FIN after Object N
// (§11.4.3), or an END_OF_GROUP bit on it (§11.4.2, §11.3.1), at N+1. The two
// agree when M = N+1, as §9.1 lets a relay turn one into the other.
//
// Interpretation: §2.4.2 calls both the "final Object", the status Object at
// M and the Object N, which would make M = N+1 two different finals.
type end struct {
        at  uint64
        set bool
        // hard reports that a FIN or END_OF_GROUP bit set it. Only then is a
        // Normal Object at at past the end: with status Objects alone, it is the
        // Object going from existing to not existing (§9.1), or the late Object of
        // §2.1, in either order.
        hard bool
}

// past reports whether an Object at id is past e: a status Object strictly
// after it, a Normal one also at it if e is hard.
func (e end) past(id uint64, normal bool) bool {
        return e.set && (id > e.at || id == e.at && normal && e.hard)
}

// conflicts reports whether e already ends somewhere other than at (set by a
// FIN or bit if hard). A status end at M and a hard end at M+1 agree: Object M
// went from existing to not existing (§9.1), and the end is M.
func (e end) conflicts(at uint64, hard bool) bool {
        switch {
        case !e.set, e.at == at:
                return false
        case hard && !e.hard:
                return at != e.at+1
        case !hard && e.hard:
                return at+1 != e.at
        }
        return true
}

// with is e also ending at at, set by a FIN or bit if hard, once checked with
// conflicts.
func (e end) with(at uint64, hard bool) end {
        switch {
        case !e.set:
                return end{at: at, set: true, hard: hard}
        case at == e.at:
                return end{at: at, set: true, hard: e.hard || hard}
        case at < e.at:
                return end{at: at, set: true} // a status end below a hard one
        }
        return e
}

// subgroupLedger is one Subgroup of a [deliveredGroup].
type subgroupLedger struct {
        priority uint8
        // maxNormal and maxStatus are the largest Normal and status Object IDs
        // received, if hasNormal and hasStatus.
        maxNormal, maxStatus uint64
        hasNormal, hasStatus bool
        end                  end
        // minObject is the lowest Object ID forwarded, if hasObject: a
        // FIRST_OBJECT claim above it is wrong (§11.4.2, §2.2). A duplicate
        // whose first copy was in another Subgroup or a datagram counts too;
        // that can only clear FIRST_OBJECT, never set it.
        minObject uint64
        hasObject bool
}

// SubgroupEnded records that an inbound subgroup stream ended with a FIN after
// lastObj: that ends its Subgroup (see [end]; on a status Object, at it) and,
// if the stream's header set END_OF_GROUP (§11.4.2), the Group. An error
// wrapping [session.ErrMalformedTrack] reports a §2.4.2 condition: another
// stream of the Subgroup ended elsewhere, or an Object past this end was
// received; the Objects past the lower end are then removed from the cache.
func (e *TrackEntry) SubgroupEnded(lastObj ObjectInfo, endOfGroup bool) error {
        group, subgroup := lastObj.Group, lastObj.Subgroup
        e.deliveredMu.Lock()
        defer e.deliveredMu.Unlock()
        g := e.delivered[group]
        if g == nil {
                return nil // aged out of the window, or every Object dropped
        }
        hard := lastObj.Status == message.ObjectStatusNormal
        at := lastObj.Object
        if hard {
                if at == math.MaxUint64 {
                        return nil // nothing lies past it
                }
                at++
        }
        ends := end{at: at, set: true, hard: hard}
        sg, seen := g.subgroups[subgroup]
        if !seen {
                sg.priority = lastObj.Priority // every Object of it was dropped
        }
        sgEnd, gEnd := sg.end.with(at, hard), g.end.with(at, hard)
        switch {
        case sg.end.conflicts(at, hard):
                e.purgePastLocked(group, &subgroup, lower(sg.end, ends))
                return fmt.Errorf("%w: Subgroup %d of Group %d ends at Objects %d and %d (§2.4.2)",
                        session.ErrMalformedTrack, subgroup, group, sg.end.at, at)
        case sg.hasStatus && sgEnd.past(sg.maxStatus, false),
                sg.hasNormal && sgEnd.past(sg.maxNormal, true):
                e.purgePastLocked(group, &subgroup, sgEnd)
                return fmt.Errorf("%w: Subgroup %d of Group %d ends at Object %d below an Object received (§2.4.2)",
                        session.ErrMalformedTrack, subgroup, group, at)
        case !endOfGroup:
        case g.end.conflicts(at, hard):
                e.purgePastLocked(group, nil, lower(g.end, ends))
                return fmt.Errorf("%w: Group %d ends at Objects %d and %d (§2.4.2)",
                        session.ErrMalformedTrack, group, g.end.at, at)
        case g.pastEnd(gEnd):
                e.purgePastLocked(group, nil, gEnd)
                return fmt.Errorf("%w: Group %d ends at Object %d below an Object received (§2.4.2)",
                        session.ErrMalformedTrack, group, at)
        }
        sg.end = sgEnd
        if g.subgroups == nil {
                g.subgroups = make(map[uint64]subgroupLedger)
        }
        g.subgroups[subgroup] = sg
        if endOfGroup {
                g.end = gEnd
        }
        return nil
}

// lower is whichever of a and b ends first.
func lower(a, b end) end {
        if b.at < a.at {
                return b
        }
        return a
}

// RecordDuplicate records what o, a copy [TrackEntry.ClaimDelivered] reported
// as already delivered, says about its Subgroup, Group and Track, once the
// caller's §9.1 check found it consistent with the first copy: an
// END_OF_GROUP arriving after a Normal Object at its ID still ends the Group,
// and a Normal copy of a status Object still counts as received.
//
// Objects claimed since o's [TrackEntry.ClaimDelivered] are checked too: an
// error wrapping [session.ErrMalformedTrack] reports a §2.4.2 condition.
func (e *TrackEntry) RecordDuplicate(o ObjectInfo) error {
        e.deliveredMu.Lock()
        defer e.deliveredMu.Unlock()
        g := e.delivered[o.Group]
        if g == nil {
                return nil
        }
        if _, ok := g.objects[o.Object]; !ok { // one dropped as announced absent
                return nil
        }
        if err := e.checkEndsLocked(g, o); err != nil {
                return err
        }
        e.recordEndsLocked(g, o)
        return nil
}

// TrackEnd reports where the Track ends, if an END_OF_TRACK Object said so.
// §2.4.2: "The final Object in a Track is the Object with Status END_OF_TRACK
// or the last Object sent in a FETCH whose response indicated End of Track";
// the relay does not record the FETCH half.
func (e *TrackEntry) TrackEnd() (message.Location, bool) {
        e.deliveredMu.Lock()
        defer e.deliveredMu.Unlock()
        return e.trackEnd, e.hasTrackEnd
}

// LowestForwarded reports the lowest Object ID forwarded in Subgroup
// (group, subgroup), if any, within the window. The writers of a Subgroup
// forget it when its last contributor leaves; this keeps it for a later
// contributor's FIRST_OBJECT claim (§11.4.2, §2.2).
func (e *TrackEntry) LowestForwarded(group, subgroup uint64) (uint64, bool) {
        e.deliveredMu.Lock()
        defer e.deliveredMu.Unlock()
        g := e.delivered[group]
        if g == nil {
                return 0, false
        }
        sg := g.subgroups[subgroup]
        return sg.minObject, sg.hasObject
}

// LocRange is an inclusive range of Locations, Lo through Hi.
type LocRange struct{ Lo, Hi message.Location }

// KnownAbsent reports the Locations the ledger knows do not exist (§2.1:
// "All signals that an Object does not exist are authoritative"), within the
// Groups it holds: each Group's end onward (§11.2.1.1, §11.4.2; a status end
// at M from M, since a status Object is not one a FETCH serializes), and the Object
// and Group ID gaps announced (§12.8, §12.9). The ranges may overlap and are
// in no particular order. A FIN or END_OF_GROUP bit that ended a Group is
// known only here, not from the cache.
func (e *TrackEntry) KnownAbsent() []LocRange {
        e.deliveredMu.Lock()
        defer e.deliveredMu.Unlock()
        var out []LocRange
        for id, g := range e.delivered {
                if g.end.set {
                        out = append(out, LocRange{
                                Lo: message.Location{Group: id, Object: g.end.at},
                                Hi: message.Location{Group: id, Object: math.MaxUint64},
                        })
                }
                for _, r := range g.objectGaps {
                        out = append(out, LocRange{
                                Lo: message.Location{Group: id, Object: r.lo},
                                Hi: message.Location{Group: id, Object: r.hi},
                        })
                }
        }
        for _, r := range e.groupGaps {
                out = append(out, LocRange{
                        Lo: message.Location{Group: r.lo},
                        Hi: message.Location{Group: r.hi, Object: math.MaxUint64},
                })
        }
        return out
}

// groupEnd reports where o ends its Group (see [end]), if it does: an
// END_OF_GROUP or END_OF_TRACK status at M at M (§11.2.1.1); a datagram's
// END_OF_GROUP bit on Object N at N+1, which §2.4.2's non-exhaustive list does
// not name but §11.3.1 defines alike. After Object 2^64-1 nothing lies past,
// so that bit ends nothing.
func groupEnd(o ObjectInfo) (at uint64, hard, ok bool) {
        switch {
        case o.Status == message.ObjectStatusEndOfGroup, o.Status == message.ObjectStatusEndOfTrack:
                return o.Object, false, true
        case o.Datagram && o.EndOfGroup && o.Object < math.MaxUint64:
                return o.Object + 1, true, true
        }
        return 0, false, false
}

// checkEndsLocked reports the §2.4.2 conditions an Object o makes against the
// earlier ones, g being its Group's ledger entry: a Publisher Priority other
// than its Subgroup's, an Object past its Subgroup's, Group's or Track's end,
// or an end placed elsewhere or below an Object received; the Objects past the
// end are then removed from the cache.
func (e *TrackEntry) checkEndsLocked(g *deliveredGroup, o ObjectInfo) error {
        normal := o.Status == message.ObjectStatusNormal
        loc := message.Location{Group: o.Group, Object: o.Object}
        switch {
        case e.hasTrackEnd && e.trackEnd.Less(loc):
                return fmt.Errorf("%w: Object %d of Group %d is past END_OF_TRACK at %d/%d (§2.4.2)",
                        session.ErrMalformedTrack, o.Object, o.Group, e.trackEnd.Group, e.trackEnd.Object)
        case o.Status != message.ObjectStatusEndOfTrack:
        case e.hasTrackEnd && e.trackEnd != loc:
                e.purgeTrackPastLocked(loc) // below the earlier one: not past it
                return fmt.Errorf("%w: END_OF_TRACK at Object %d of Group %d and at %d/%d (§2.4.2)",
                        session.ErrMalformedTrack, o.Object, o.Group, e.trackEnd.Group, e.trackEnd.Object)
        case e.purgeTrackPastLocked(loc):
                return fmt.Errorf("%w: END_OF_TRACK at Object %d of Group %d is below an Object received (§2.4.2)",
                        session.ErrMalformedTrack, o.Object, o.Group)
        }
        if g == nil {
                return nil
        }
        if g.end.past(o.Object, normal) {
                return fmt.Errorf("%w: Object %d of Group %d is past its end at %d (§2.4.2)",
                        session.ErrMalformedTrack, o.Object, o.Group, g.end.at)
        }
        if at, hard, ok := groupEnd(o); ok {
                gEnd := g.end.with(at, hard)
                switch {
                case g.end.conflicts(at, hard):
                        e.purgePastLocked(o.Group, nil, lower(g.end, end{at: at, set: true, hard: hard}))
                        return fmt.Errorf("%w: Group %d ends at Objects %d and %d (§2.4.2)",
                                session.ErrMalformedTrack, o.Group, g.end.at, at)
                case g.pastEnd(gEnd):
                        e.purgePastLocked(o.Group, nil, gEnd)
                        return fmt.Errorf("%w: Group %d ends at Object %d below an Object received (§2.4.2)",
                                session.ErrMalformedTrack, o.Group, at)
                }
        }
        if o.Datagram {
                return nil
        }
        sg, seen := g.subgroups[o.Subgroup]
        switch {
        case !seen:
        case sg.priority != o.Priority:
                return fmt.Errorf("%w: Subgroup %d of Group %d has Publisher Priorities %d and %d (§2.4.2)",
                        session.ErrMalformedTrack, o.Subgroup, o.Group, sg.priority, o.Priority)
        case sg.end.past(o.Object, normal):
                return fmt.Errorf("%w: Object %d of Subgroup %d in Group %d is past its end at %d (§2.4.2)",
                        session.ErrMalformedTrack, o.Object, o.Subgroup, o.Group, sg.end.at)
        }
        return nil
}

// purgeTrackPastLocked reports whether an Object past END_OF_TRACK at loc was
// received, in the window, and removes those from the cache. END_OF_TRACK is a
// status end (see [end]): a Normal Object at loc is the late Object of §2.1.
func (e *TrackEntry) purgeTrackPastLocked(loc message.Location) bool {
        found := false
        ended := end{at: loc.Object, set: true}
        for id, g := range e.delivered {
                switch {
                case id > loc.Group:
                        found = true
                        e.purgeGroupLocked(id)
                case id == loc.Group && g.pastEnd(ended):
                        found = true
                        e.purgePastLocked(id, nil, ended)
                }
        }
        return found
}

// purgeGroupLocked removes every Object of group from the cache, status
// Objects included (§2.4.2).
func (e *TrackEntry) purgeGroupLocked(group uint64) {
        objs := e.Cache.GetRange(
                message.Location{Group: group},
                message.Location{Group: group, Object: math.MaxUint64},
                message.GroupOrderAscending,
        )
        for _, o := range objs {
                e.Cache.Delete(group, o.ObjectID)
        }
}

// purgePastLocked removes from the cache the Objects of group, in subgroup if
// not nil, past ended: Object(s) triggering Malformed Track status MUST NOT be
// cached (§2.4.2).
func (e *TrackEntry) purgePastLocked(group uint64, subgroup *uint64, ended end) {
        objs := e.Cache.GetRange(
                message.Location{Group: group, Object: ended.at},
                message.Location{Group: group, Object: math.MaxUint64},
                message.GroupOrderAscending,
        )
        for _, o := range objs {
                if !ended.past(o.ObjectID, o.Status == message.ObjectStatusNormal) {
                        continue
                }
                if subgroup != nil && (o.ForwardingPref != cache.ForwardingSubgroup || o.SubgroupID != *subgroup) {
                        continue
                }
                e.Cache.Delete(group, o.ObjectID)
        }
}

// recordEndsLocked records what o, once checked, says about its Subgroup,
// Group and Track.
func (e *TrackEntry) recordEndsLocked(g *deliveredGroup, o ObjectInfo) {
        normal := o.Status == message.ObjectStatusNormal
        if normal {
                g.maxNormal, g.hasNormal = max(g.maxNormal, o.Object), true
        } else {
                g.maxStatus, g.hasStatus = max(g.maxStatus, o.Object), true
        }
        if at, hard, ok := groupEnd(o); ok {
                g.end = g.end.with(at, hard)
        }
        if o.Status == message.ObjectStatusEndOfTrack {
                e.trackEnd, e.hasTrackEnd = message.Location{Group: o.Group, Object: o.Object}, true
        }
        if o.Datagram {
                return
        }
        if g.subgroups == nil {
                g.subgroups = make(map[uint64]subgroupLedger)
        }
        sg, seen := g.subgroups[o.Subgroup]
        if !seen {
                sg.priority = o.Priority
        }
        if normal {
                sg.maxNormal, sg.hasNormal = max(sg.maxNormal, o.Object), true
        } else {
                sg.maxStatus, sg.hasStatus = max(sg.maxStatus, o.Object), true
        }
        if !sg.hasObject || o.Object < sg.minObject {
                sg.minObject, sg.hasObject = o.Object, true
        }
        g.subgroups[o.Subgroup] = sg
}

// deliveredGroup is one Group of the [TrackEntry.ClaimDelivered] ledger; g is
// nil for a Group not yet seen.
type deliveredGroup struct {
        objects map[uint64]struct{}
        // maxNormal and maxStatus are the largest Normal and status Object IDs
        // received, if hasNormal and hasStatus.
        maxNormal, maxStatus uint64
        hasNormal, hasStatus bool
        subgroups            map[uint64]subgroupLedger // created on first use
        end                  end
        // objectGaps are the Object ID ranges announced absent (§12.9).
        objectGaps []idRange
        // groupGap is the Group's Prior Group ID Gap (§12.8), if hasGroupGap.
        groupGap    uint64
        hasGroupGap bool
}

// pastEnd reports whether an Object received is past ended.
func (g *deliveredGroup) pastEnd(ended end) bool {
        return g.hasNormal && ended.past(g.maxNormal, true) || g.hasStatus && ended.past(g.maxStatus, false)
}

// idRange is an inclusive range of Group or Object IDs.
type idRange struct{ lo, hi uint64 }

func (r idRange) contains(id uint64) bool { return r.lo <= id && id <= r.hi }

// gapRange is the IDs a Prior Group or Object ID Gap of n on id announces
// absent (§12.8, §12.9); none for n == 0. The session layer has already
// rejected an n above id ([message.CheckObjectProperties]).
func gapRange(id, n uint64) (idRange, bool) {
        return idRange{lo: id - n, hi: id - 1}, n > 0
}

// announcedAbsentLocked reports whether the Object at (group, object), g
// being its Group's ledger entry, is inside a gap announced earlier. The
// Group's scan grows with its distinct Object ID gaps, bounded only by the
// window; publishers rarely send many, so the cost is accepted.
func (e *TrackEntry) announcedAbsentLocked(g *deliveredGroup, group, object uint64) bool {
        if slices.ContainsFunc(e.groupGaps, func(r idRange) bool { return r.contains(group) }) {
                return true
        }
        return g != nil && slices.ContainsFunc(g.objectGaps, func(r idRange) bool { return r.contains(object) })
}

// recordGapsLocked records the gaps an Object at (group, object) announced.
func (e *TrackEntry) recordGapsLocked(g *deliveredGroup, group, object uint64, gaps message.PriorGaps) {
        if r, ok := gapRange(object, gaps.Object); gaps.HasObject && ok && !slices.Contains(g.objectGaps, r) {
                g.objectGaps = append(g.objectGaps, r)
        }
        if gaps.HasGroup && !g.hasGroupGap {
                g.groupGap, g.hasGroupGap = gaps.Group, true
                if r, ok := gapRange(group, gaps.Group); ok {
                        e.groupGaps = append(e.groupGaps, r)
                }
        }
}

// ReleaseSubgroup drops one contributor from (group, subgroup) and reports
// whether that was the last one (in which case the shared state has been removed
// from the entry and the caller owns tearing down its downstream writers).
func (e *TrackEntry) ReleaseSubgroup(key SubgroupKey) (last bool) {
        e.sgMu.Lock()
        defer e.sgMu.Unlock()
        sg, ok := e.subgroups[key]
        if !ok {
                return false
        }
        sg.refs--
        if sg.refs <= 0 {
                delete(e.subgroups, key)
                return true
        }
        return false
}

// UpdateLargest moves the entry's LargestObject forward. The very first
// call flips the "has any object been observed" bit regardless of value,
// so that a publisher whose first Object is at Location {0, 0} is still
// distinguishable from "no objects observed yet" — §10.2.17 reserves the
// wire-level omission of LARGEST_OBJECT for the latter, and the
// in-memory mirror needs the same distinction. Subsequent calls advance
// the watermark only when loc is strictly greater than the current
// value.
//
// Returns true when the watermark changed (advanced or first-set);
// callers can use it to avoid redundant LARGEST_OBJECT-property
// emission downstream.
func (e *TrackEntry) UpdateLargest(loc message.Location) bool {
        e.mu.Lock()
        defer e.mu.Unlock()
        if !e.HasLargestObject || e.LargestObject.Less(loc) {
                e.LargestObject = loc
                e.HasLargestObject = true
                return true
        }
        return false
}

// GetLargest returns the current largest-object watermark and a bool that is
// true iff at least one object has been observed on this track. The bool
// distinguishes "no objects observed yet" from "first object was published at
// Location {0, 0}" — §10.2.17 reserves wire-level omission of LARGEST_OBJECT
// for the former, so the in-memory mirror needs the same distinction.
func (e *TrackEntry) GetLargest() (message.Location, bool) {
        e.mu.RLock()
        defer e.mu.RUnlock()
        return e.LargestObject, e.HasLargestObject
}

// ConsiderNewGroupRequest applies the §10.2.19 relay rules for a
// NEW_GROUP_REQUEST received on an Established subscription and reports whether
// the relay must forward it upstream (via an upstream REQUEST_UPDATE). When it
// returns true the request is recorded as outstanding.
//
// value is the downstream NEW_GROUP_REQUEST (largest known Group + 1, or 0 for
// "no Group information"). dynamicGroups reports whether the track advertised
// DYNAMIC_GROUPS=1 (§12.6). The rules:
//
//   - The Track must support dynamic Groups (unless-clause 1).
//   - The request is forwarded only when value is 0 or larger than the current
//     Largest Group; a non-zero value at or below the Largest Group is not
//     forwarded.
//   - An outstanding request with a value greater than or equal to this one
//     already covers it (unless-clause 2). An outstanding request is cleared
//     once the Largest Group advances past where it was sent.
func (e *TrackEntry) ConsiderNewGroupRequest(value uint64, dynamicGroups bool) bool {
        if !dynamicGroups {
                return false
        }

        e.mu.Lock()
        defer e.mu.Unlock()

        var largest uint64
        if e.HasLargestObject {
                largest = e.LargestObject.Group
        }

        // "After sending a NEW_GROUP_REQUEST upstream, the request is considered
        // outstanding until the Largest Group increases."
        if e.newGroupOutstanding && largest > e.newGroupReqGroup {
                e.newGroupOutstanding = false
        }

        // A non-zero value at or below the Largest Group needs no new Group.
        if value != 0 && value <= largest {
                return false
        }

        // An outstanding request of equal or greater value already covers this.
        if e.newGroupOutstanding && e.newGroupReqValue >= value {
                return false
        }

        e.newGroupOutstanding = true
        e.newGroupReqValue = value
        e.newGroupReqGroup = largest
        return true
}

// UpdateLargestAndDetectNew advances LargestObject and, under the same
// e.mu acquisition, returns any Downstream subs for which seen reports
// false. seen is consulted (the fanout passes a membership test over the
// writers it has already opened); the entry never mutates it. Used by
// runFanout per-object so a downstream sub that joined the entry's
// Downstream after the initial CopyDownstream is detected and given a
// writer for the current (and subsequent) objects on the in-flight
// subgroup stream.
//
// Atomic with [TrackRegistry.AddDownstreamSnapshotLargest]: a new sub
// either snapshots the pre-update LargestObject AND appears in newSubs
// (delivered live), or snapshots the post-update LargestObject (covered
// by its fill fetch stream). The lock pair guarantees no in-between.
// lastGen is the downstreamGen the caller observed on its previous call (or
// at its initial CopyDownstreamWithGen snapshot). When the generation is
// unchanged no sub has joined since, so the joiner scan is skipped entirely;
// gen (returned) should be fed back as lastGen on the next call.
//
// seen is a predicate rather than a concrete map so this (registry) layer
// need not know the fanout's writer type, keeping the dependency edge
// pointing one way (fanout → registry).
func (e *TrackEntry) UpdateLargestAndDetectNew(
        loc message.Location,
        seen func(*DownstreamSub) bool,
        lastGen uint64,
) (newSubs []*DownstreamSub, gen uint64) {
        e.mu.Lock()
        defer e.mu.Unlock()
        if !e.HasLargestObject || e.LargestObject.Less(loc) {
                e.LargestObject = loc
                e.HasLargestObject = true
        }
        if e.downstreamGen == lastGen {
                return nil, e.downstreamGen
        }
        for _, sub := range e.Downstream {
                if !seen(sub) {
                        newSubs = append(newSubs, sub)
                }
        }
        return newSubs, e.downstreamGen
}

// SetProperties stores the Track Properties learned from the upstream. The
// caller hands over ownership of props; callers MUST NOT mutate props after
// this call. (Properties are immutable once captured — §9.6 expects them to
// be replayed verbatim.)
func (e *TrackEntry) SetProperties(props []byte) {
        e.mu.Lock()
        e.setPropertiesLocked(props)
        e.mu.Unlock()
}

// setPropertiesLocked stores the raw Properties bytes and decodes the fields
// the relay acts on in the same step, so the decoded values never drift from
// the raw bytes. Callers must hold e.mu.
func (e *TrackEntry) setPropertiesLocked(raw []byte) {
        e.Properties = raw
        e.decoded = decodeTrackProperties(raw)
}

// GetProperties returns the raw Track Properties captured from the upstream
// publisher. The returned slice is the same byte buffer stored on the entry;
// callers MUST NOT mutate it.
func (e *TrackEntry) GetProperties() []byte {
        e.mu.RLock()
        defer e.mu.RUnlock()
        return e.Properties
}

// fetchCancel is one [TrackEntry.AddFetch] registration.
type fetchCancel struct{ cancel context.CancelCauseFunc }

// AddFetch registers cancel as how to end a fetch stream the relay serves on
// this track until remove is called; see [TrackEntry.CancelFetches].
func (e *TrackEntry) AddFetch(cancel context.CancelCauseFunc) (remove func()) {
        f := &fetchCancel{cancel}
        e.mu.Lock()
        if e.fetches == nil {
                e.fetches = make(map[*fetchCancel]struct{})
        }
        e.fetches[f] = struct{}{}
        e.mu.Unlock()
        return func() {
                e.mu.Lock()
                delete(e.fetches, f)
                e.mu.Unlock()
        }
}

// CancelFetches cancels every fetch stream registered with
// [TrackEntry.AddFetch] with cause: §2.4.2, a relay that detects a malformed
// track MUST "reset any fetch streams".
func (e *TrackEntry) CancelFetches(cause error) {
        e.mu.RLock()
        fs := slices.Collect(maps.Keys(e.fetches))
        e.mu.RUnlock()
        for _, f := range fs {
                f.cancel(cause)
        }
}

// CopyUpstream returns a snapshot of the current upstream slice. Callers
// that want to iterate without holding the entry lock for the whole
// iteration use this so they don't have to coordinate with mutators.
func (e *TrackEntry) CopyUpstream() []*UpstreamSub {
        e.mu.RLock()
        defer e.mu.RUnlock()
        out := make([]*UpstreamSub, len(e.Upstream))
        copy(out, e.Upstream)
        return out
}

// HasUpstreamOn reports whether one of the entry's upstream subscriptions is
// on sess.
func (e *TrackEntry) HasUpstreamOn(sess *session.Session) bool {
        e.mu.RLock()
        defer e.mu.RUnlock()
        return slices.ContainsFunc(e.Upstream, func(u *UpstreamSub) bool { return u.Session == sess })
}

// upstreamEpochs numbers upstream additions process-wide, so epochs only grow
// and never repeat on an entry created again for the same track.
var upstreamEpochs atomic.Uint64

// UpstreamEpoch identifies the latest upstream PUBLISH or SUBSCRIBE added to
// the track (see [TrackRegistry.AddUpstream]); it is never 0 once one was.
// A PUBLISH_SKIPPED holds until it grows (see [SubscriberEntry.ClaimForward]).
func (e *TrackEntry) UpstreamEpoch() uint64 {
        e.mu.RLock()
        defer e.mu.RUnlock()
        return e.upstreamEpoch
}

// NoteRefusal records that pub refused a late-publisher SUBSCRIBE for this
// track and may not be asked again before retryAt; a zero retryAt means not
// while this entry and that registration both last.
func (e *TrackEntry) NoteRefusal(pub *PublisherEntry, retryAt time.Time) {
        e.mu.Lock()
        defer e.mu.Unlock()
        if e.refusals == nil {
                e.refusals = make(map[*PublisherEntry]time.Time)
        }
        e.refusals[pub] = retryAt
}

// Refused reports whether a refusal from pub still stands at now.
func (e *TrackEntry) Refused(pub *PublisherEntry, now time.Time) bool {
        e.mu.RLock()
        defer e.mu.RUnlock()
        retryAt, ok := e.refusals[pub]
        return ok && (retryAt.IsZero() || now.Before(retryAt))
}

// RetainRefusals forgets refusals that no longer stand at now or whose
// publisher is not in current, the registrations still covering the track.
func (e *TrackEntry) RetainRefusals(current []*PublisherEntry, now time.Time) {
        stale := func(p *PublisherEntry, retryAt time.Time) bool {
                return (!retryAt.IsZero() && !now.Before(retryAt)) || !slices.Contains(current, p)
        }
        e.mu.RLock()
        anyStale := false
        for p, retryAt := range e.refusals {
                if stale(p, retryAt) {
                        anyStale = true
                        break
                }
        }
        e.mu.RUnlock()
        if !anyStale {
                return
        }
        e.mu.Lock()
        defer e.mu.Unlock()
        maps.DeleteFunc(e.refusals, stale)
}

// HasDownstreamOn reports whether one of the entry's live downstream
// subscriptions is on sess. A terminated one, lingering until its subscriber
// closes its side, does not count.
func (e *TrackEntry) HasDownstreamOn(sess *session.Session) bool {
        e.mu.RLock()
        defer e.mu.RUnlock()
        return slices.ContainsFunc(e.Downstream, func(d *DownstreamSub) bool {
                return d.Session == sess && !d.IsTerminated()
        })
}

// CopyDownstream returns a snapshot of the current downstream slice. See
// [TrackEntry.CopyUpstream] for rationale.
func (e *TrackEntry) CopyDownstream() []*DownstreamSub {
        e.mu.RLock()
        defer e.mu.RUnlock()
        out := make([]*DownstreamSub, len(e.Downstream))
        copy(out, e.Downstream)
        return out
}

// CopyDownstreamWithGen is [TrackEntry.CopyDownstream] plus the matching
// downstreamGen, captured under the same lock so the fanout can seed its
// joiner-scan skip with a generation that is exactly consistent with the
// snapshot (a sub joining after this returns bumps the generation and so is
// still detected on the next per-object call).
func (e *TrackEntry) CopyDownstreamWithGen() ([]*DownstreamSub, uint64) {
        e.mu.RLock()
        defer e.mu.RUnlock()
        out := make([]*DownstreamSub, len(e.Downstream))
        copy(out, e.Downstream)
        return out, e.downstreamGen
}

// Package registry holds the relay's process-wide shared state: the track
// registry (object routing + per-track cache), the namespace registry
// (PUBLISH_NAMESPACE / SUBSCRIBE_NAMESPACE bookkeeping), the fetch router
// (rendezvous for upstream FETCH response streams), and the subscription
// state machine (UpstreamSub / DownstreamSub).
//
// It is the bottom layer of the relay: the parent pkg/relay session handlers
// depend on it, but it never imports the parent — the dependency edge only
// ever points handler → registry. Living under internal/ also keeps these
// types out of pkg/relay's public API; they are exported for the package's own
// white-box tests, not for external consumers. See the pkg/relay package doc
// for the full layer map.
package registry

import (
        "context"
        "errors"
        "log/slog"
        "slices"
        "sync"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
        "github.com/floatdrop/moq-go/pkg/relay/cache"
        "github.com/floatdrop/moq-go/pkg/relay/discovery"
)

// Default per-track object-cache bounds. Used by [NewTrackRegistry] when
// the caller does not supply [WithCacheConfig]. The relay's Config overrides
// these (relay.New reads them to fill unset Config fields); tests that
// construct registries directly inherit the defaults.
const (
        DefaultCacheMaxSize     = 1024
        DefaultCacheMaxDuration = 30 * time.Second
)

// discoveryCallTimeout bounds each best-effort call into the DiscoveryStore
// (publish/unpublish of tracks and namespaces). Discovery is off the critical
// path, so a short timeout keeps a slow store from stalling registry
// bookkeeping; failures are logged and swallowed.
const discoveryCallTimeout = 100 * time.Millisecond

// CacheTTLPolicy is the registry's view of the per-track Object Cache TTL
// override: given a track's Full Track Name, return the TTL to use. It is
// the structural twin of the public relay.CacheTTLPolicy; relay converts
// its exported type to this one at the registry boundary, which keeps the
// dependency pointing one way (registry never imports its parent) while
// still giving each layer a named, self-documenting type rather than a
// bare function signature. [resolveCacheTTL] documents the return-value
// contract (negative disables eviction, 0 falls through to the default).
type CacheTTLPolicy func(name track.FullTrackName) time.Duration

// TrackRegistry indexes [TrackEntry] values by [track.Key]. It is the single
// rendezvous point for everything in the relay that needs to address a track
// — request handlers, fanout, the cache, and the discovery store.
//
// Locking strategy: the registry-level RWMutex protects only the tracks map.
// All entry mutation happens under TrackEntry.mu, which the helpers below
// acquire in the appropriate mode. This keeps the registry-level critical
// sections O(1) and lets per-track work proceed in parallel.
type TrackRegistry struct {
        mu     sync.RWMutex
        tracks map[track.Key]*TrackEntry

        // cacheMaxSize / cacheMaxDuration are the per-track Object Cache
        // bounds applied to every entry created by this registry.
        cacheMaxSize     int
        cacheMaxDuration time.Duration

        // cacheTTLPolicy, when non-nil, may override cacheMaxDuration on a
        // per-track basis. See [CacheTTLPolicy] for the contract and
        // [resolveCacheTTL] for how its result is interpreted.
        cacheTTLPolicy CacheTTLPolicy

        // discovery is the cross-instance track advertisement fabric. nil
        // means "do not advertise" — the relay still works as a local
        // single-instance setup. When non-nil, the registry publishes a
        // [discovery.TrackInfo] on the first AddUpstream for a track and
        // unpublishes on the last RemoveUpstream.
        discovery discovery.DiscoveryStore

        // relayAddr is the address the relay registers itself as in
        // Discovery entries. Empty for single-relay deployments.
        relayAddr string

        // log is used for warn-level reports when a Discovery call fails.
        // Discovery failures are NOT propagated to the caller — the
        // registry is the source of truth for local state, Discovery is
        // best-effort.
        log *slog.Logger

        // claims holds the upstream SUBSCRIBEs in flight, one per (session,
        // track); see [TrackRegistry.ClaimUpstream]. Guarded by mu.
        claims map[upstreamClaim]struct{}
        // inflight counts the relay's own FETCHes and TRACK_STATUSes in flight,
        // per (type, session, track); see [TrackRegistry.BeginRequest]. Guarded
        // by mu.
        inflight map[inflightRequest]int
        // arrivals are the waiters for a track's next upstream; see
        // [TrackRegistry.AwaitUpstream]. Guarded by mu.
        arrivals arrivals[track.Key]
}

type upstreamClaim struct {
        sess *session.Session
        key  track.Key
}

// TrackRegistryOption tweaks a [TrackRegistry] at construction time.
type TrackRegistryOption func(*TrackRegistry)

// WithCacheConfig sets the per-track object-cache bounds applied to every
// new entry the registry constructs. maxSize is the per-track upper bound
// on stored objects; maxDuration is the maximum age before time-based
// eviction. Values <= 0 fall back to the package defaults
// ([DefaultCacheMaxSize], [DefaultCacheMaxDuration]).
func WithCacheConfig(maxSize int, maxDuration time.Duration) TrackRegistryOption {
        return func(r *TrackRegistry) {
                if maxSize > 0 {
                        r.cacheMaxSize = maxSize
                }
                if maxDuration > 0 {
                        r.cacheMaxDuration = maxDuration
                }
        }
}

// WithCacheTTLPolicy installs a per-track TTL override hook. See
// [CacheTTLPolicy] for the contract and [resolveCacheTTL] for how the
// returned duration is interpreted. Passing a nil policy is allowed and
// equivalent to not calling this option — every track uses the default
// TTL from [WithCacheConfig].
//
// Typical use is to give one well-known track (e.g. an MSF catalog
// track) infinite retention while every other track keeps the default
// 30-second bound — the operator wires the rule into the policy at the
// binary layer so the relay stays protocol-agnostic.
func WithCacheTTLPolicy(policy CacheTTLPolicy) TrackRegistryOption {
        return func(r *TrackRegistry) {
                r.cacheTTLPolicy = policy
        }
}

// WithTrackDiscovery installs a [discovery.DiscoveryStore] for
// cross-instance track advertisement. relayAddr is the value stamped
// into every [discovery.TrackInfo] this registry emits.
func WithTrackDiscovery(d discovery.DiscoveryStore, relayAddr string) TrackRegistryOption {
        return func(r *TrackRegistry) {
                r.discovery = d
                r.relayAddr = relayAddr
        }
}

// WithTrackRegistryLogger sets the logger used for Discovery warnings.
func WithTrackRegistryLogger(l *slog.Logger) TrackRegistryOption {
        return func(r *TrackRegistry) { r.log = l }
}

// NewTrackRegistry constructs an empty registry. Default per-track cache
// bounds are [DefaultCacheMaxSize] / [DefaultCacheMaxDuration]; callers
// override them with [WithCacheConfig].
func NewTrackRegistry(opts ...TrackRegistryOption) *TrackRegistry {
        r := &TrackRegistry{
                tracks:           make(map[track.Key]*TrackEntry),
                cacheMaxSize:     DefaultCacheMaxSize,
                cacheMaxDuration: DefaultCacheMaxDuration,
                log:              slog.Default(),
        }
        for _, opt := range opts {
                opt(r)
        }
        return r
}

// Get returns the entry for key, or (nil, false) if no such track is known.
// The returned pointer is valid until the entry is destroyed (last Remove*
// call) — readers that want to keep it across long operations should
// nevertheless cope with a stale pointer by re-querying.
func (r *TrackRegistry) Get(key track.Key) (*TrackEntry, bool) {
        r.mu.RLock()
        defer r.mu.RUnlock()
        e, ok := r.tracks[key]
        return e, ok
}

// MatchNamespace returns every entry whose Track Namespace has prefix as a
// prefix: the tracks a PUBLISH_NAMESPACE for prefix covers (§9.5).
func (r *TrackRegistry) MatchNamespace(prefix wire.TrackNamespace) []*TrackEntry {
        r.mu.RLock()
        defer r.mu.RUnlock()
        var out []*TrackEntry
        for _, e := range r.tracks {
                if e.FullName.Namespace.HasPrefix(prefix) {
                        out = append(out, e)
                }
        }
        return out
}

// ClaimUpstream marks an upstream SUBSCRIBE for key on sess as in flight, so a
// §9.5 late-publisher SUBSCRIBE does not duplicate one. On success the caller
// opens and registers the upstream, then calls release. ok is false when sess
// already has a registered upstream for key or another claim is in flight.
func (r *TrackRegistry) ClaimUpstream(sess *session.Session, key track.Key) (release func(), ok bool) {
        c := upstreamClaim{sess: sess, key: key}
        r.mu.Lock()
        defer r.mu.Unlock()
        if _, busy := r.claims[c]; busy {
                return nil, false
        }
        if e, exists := r.tracks[key]; exists && e.HasUpstreamOn(sess) {
                return nil, false
        }
        if r.claims == nil {
                r.claims = make(map[upstreamClaim]struct{})
        }
        r.claims[c] = struct{}{}
        return func() {
                r.mu.Lock()
                delete(r.claims, c)
                r.mu.Unlock()
        }, true
}

// inflightRequest is one kind of request the relay sends for a track on a
// session.
type inflightRequest struct {
        upstreamClaim

        typ message.Type
}

// BeginRequest marks a request of type typ for key on sess as in flight until
// done; see [TrackRegistry.RequestPending].
func (r *TrackRegistry) BeginRequest(typ message.Type, sess *session.Session, key track.Key) (done func()) {
        c := inflightRequest{typ: typ, sess: sess, key: key}
        r.mu.Lock()
        defer r.mu.Unlock()
        if r.inflight == nil {
                r.inflight = make(map[inflightRequest]int)
        }
        r.inflight[c]++
        return func() {
                r.mu.Lock()
                defer r.mu.Unlock()
                if r.inflight[c]--; r.inflight[c] == 0 {
                        delete(r.inflight, c)
                }
        }
}

// RequestPending reports whether a request of type typ for key on sess is in
// flight.
func (r *TrackRegistry) RequestPending(typ message.Type, sess *session.Session, key track.Key) bool {
        r.mu.RLock()
        defer r.mu.RUnlock()
        return r.inflight[inflightRequest{typ: typ, sess: sess, key: key}] > 0
}

// ReleaseIfUnsubscribed removes the on-demand upstream up from the entry for
// fullName and tears it down if the entry has no downstream left. It covers
// an upstream whose last downstream left during the SUBSCRIBE round trip,
// before [TrackRegistry.RemoveDownstream] could strip it. Reports whether up
// was released.
func (r *TrackRegistry) ReleaseIfUnsubscribed(fullName track.FullTrackName, up *UpstreamSub) bool {
        key := fullName.Key()
        r.mu.Lock()
        entry, ok := r.tracks[key]
        if !ok {
                r.mu.Unlock()
                return false
        }
        entry.mu.Lock()
        before := len(entry.Upstream)
        if len(entry.Downstream) == 0 {
                entry.Upstream = slices.DeleteFunc(entry.Upstream, func(u *UpstreamSub) bool { return u == up })
        }
        released := len(entry.Upstream) < before
        upstreamEmpty := len(entry.Upstream) == 0
        entry.mu.Unlock()
        if released && upstreamEmpty {
                delete(r.tracks, key)
                // Still under r.mu: see [TrackRegistry.unpublishTrackFromDiscovery].
                r.unpublishTrackFromDiscovery(entry)
        }
        r.mu.Unlock()
        if released {
                up.CloseOnDemand()
        }
        return released
}

// GetOrCreateNew is [TrackRegistry.GetOrCreate] that also reports whether this
// call created the entry. Callers that create one speculatively — before the
// request that will populate it is known to succeed — need to know, because
// only the creator may take it back again (see [TrackRegistry.DeleteIfUnused]).
func (r *TrackRegistry) GetOrCreateNew(fullName track.FullTrackName) (entry *TrackEntry, created bool) {
        key := fullName.Key()
        r.mu.Lock()
        defer r.mu.Unlock()
        if e, ok := r.tracks[key]; ok {
                return e, false
        }
        return r.getOrCreateLocked(fullName), true
}

// DeleteIfUnused removes fullName's entry if nothing has happened to it since
// it was created: no upstream, no downstream, no watermark and nothing cached.
// It is the counterpart to a speculative [TrackRegistry.GetOrCreateNew] — the
// PUBLISH path creates the entry before the request is known to succeed, so an
// inbound data stream cannot arrive against a routable Track Alias with no
// entry to route it to (§11.1, §10.11). When the request then fails, the entry
// must not linger: [TrackRegistry.Get] answering for it is read as "track
// known" by the FETCH path, which turns a DOES_NOT_EXIST into an INVALID_RANGE.
//
// The watermark and cache are checked as well as the two slices because a
// DIFFERENT session may have adopted this entry in the meantime and be part
// way through its own §10.11 window — publishing objects into a track whose
// AddUpstream has not run yet. Deleting it there would strand exactly the
// streams this mechanism exists to protect. Callers must additionally only
// call this when GetOrCreateNew reported created, so an entry another session
// created is never a candidate.
//
// Idempotent, and a no-op once anything has registered against the entry.
func (r *TrackRegistry) DeleteIfUnused(fullName track.FullTrackName) {
        key := fullName.Key()
        r.mu.Lock()
        defer r.mu.Unlock()
        entry, ok := r.tracks[key]
        if !ok {
                return
        }
        // r.mu then entry.mu throughout, matching AddUpstream / RemoveSession.
        // GetLargest takes entry.mu itself, so it runs before the block below
        // rather than inside it.
        if _, hasLargest := entry.GetLargest(); hasLargest {
                return
        }
        if entry.Cache.Len() > 0 {
                return
        }
        entry.mu.Lock()
        unused := len(entry.Upstream) == 0 && len(entry.Downstream) == 0
        entry.mu.Unlock()
        if unused {
                delete(r.tracks, key)
        }
}

// GetOrCreate returns the existing entry for fullName, or creates and inserts
// a new one if none exists yet. The fullName argument (rather than just a
// Key) is required so a freshly-created entry can be populated with the
// {namespace, name} tuple §9.6 needs to echo back on outbound replies.
//
// NOTE: GetOrCreate by itself is not sufficient to protect against the
// resurrection race — between this call returning and the caller acquiring
// the entry's own lock, a concurrent Remove may delete the entry from the
// registry map even though the returned pointer remains valid. Callers that
// mutate Upstream/Downstream MUST go through [TrackRegistry.AddUpstream] /
// [TrackRegistry.AddDownstream] (which hold the registry lock for the whole
// add operation) rather than calling GetOrCreate themselves. GetOrCreate is
// exported because read-only callers (tests, metrics) legitimately want a
// "find me an entry, create if missing" primitive.
func (r *TrackRegistry) GetOrCreate(fullName track.FullTrackName) *TrackEntry {
        r.mu.Lock()
        defer r.mu.Unlock()
        return r.getOrCreateLocked(fullName)
}

// getOrCreateLocked is the inner helper used by AddUpstream / AddDownstream.
// The caller must hold r.mu for writing.
func (r *TrackRegistry) getOrCreateLocked(fullName track.FullTrackName) *TrackEntry {
        key := fullName.Key()
        if e, ok := r.tracks[key]; ok {
                return e
        }
        e := &TrackEntry{
                Key:      key,
                FullName: fullName,
                Cache:    cache.NewObjectCache(r.cacheMaxSize, r.resolveCacheTTL(fullName)),
        }
        r.tracks[key] = e
        return e
}

// resolveCacheTTL picks the per-track Object Cache TTL for fullName,
// consulting [TrackRegistry.cacheTTLPolicy] if one was installed. It
// maps the policy's return value onto the [cache.ObjectCache]
// convention (where a non-positive TTL means "no time-based eviction"):
//
//   - a negative duration (the public relay.CacheTTLInfinite sentinel)
//     becomes 0, disabling time-based eviction for the track;
//   - a positive duration is used as-is;
//   - 0, or no policy at all, falls through to the registry-wide
//     default from [WithCacheConfig].
//
// Keeping this translation inside the registry means policy authors only
// ever deal with the public relay.CacheTTLPolicy vocabulary.
func (r *TrackRegistry) resolveCacheTTL(fullName track.FullTrackName) time.Duration {
        if r.cacheTTLPolicy == nil {
                return r.cacheMaxDuration
        }
        switch d := r.cacheTTLPolicy(fullName); {
        case d < 0:
                return 0 // cache.ObjectCache: <=0 means "no TTL filtering"
        case d > 0:
                return d
        default:
                return r.cacheMaxDuration
        }
}

// Len returns the number of tracks currently held. Primarily useful for
// tests and metrics; not part of the relay's hot path.
func (r *TrackRegistry) Len() int {
        r.mu.RLock()
        defer r.mu.RUnlock()
        return len(r.tracks)
}

// AddUpstream appends sub to the entry for fullName (creating the entry if
// necessary) and returns the entry.
//
// Returns (entry, becameNonEmpty). becameNonEmpty is true when this call
// installed the first upstream subscription on the entry, which is the
// signal the registry uses to publish the track to the Discovery Store.
// The boolean also lets tests assert the "first publisher" transition.
//
// The whole add operation runs under the registry write lock. This is
// stricter than strictly necessary, but it eliminates a subtle race where
// a Remove on another goroutine could delete the entry from the map after
// GetOrCreate returns but before the caller locks the entry — leaving the
// caller mutating a [TrackEntry] that no future Get can reach. Add/Remove
// frequency is dwarfed by fanout (which uses [TrackRegistry.Get] +
// [TrackEntry.CopyDownstream] and only takes the registry RLock), so the
// extra serialisation does not affect the hot path.
func (r *TrackRegistry) AddUpstream(
        fullName track.FullTrackName,
        sub *UpstreamSub,
        opts ...AddUpstreamOption,
) (entry *TrackEntry, becameNonEmpty bool) {
        var conf addUpstreamConfig
        for _, opt := range opts {
                opt(&conf)
        }
        r.mu.Lock()
        defer r.mu.Unlock()
        entry = r.getOrCreateLocked(fullName)
        entry.mu.Lock()
        defer entry.mu.Unlock()
        becameNonEmpty = len(entry.Upstream) == 0
        entry.Upstream = append(entry.Upstream, sub)
        entry.upstreamEpoch = upstreamEpochs.Add(1)
        if conf.setProperties && len(entry.Properties) == 0 {
                // Set Properties INSIDE the entry lock so the Discovery
                // publish below sees them. Skip if Properties were already
                // captured by a prior caller — §9.6 expects them to be
                // stable for the lifetime of the track entry, so the first
                // setter wins.
                entry.setPropertiesLocked(conf.properties)
        }
        if becameNonEmpty {
                r.publishTrackToDiscovery(entry)
        }
        r.arrivals.notifyLocked(fullName.Key())
        return entry, becameNonEmpty
}

// AwaitUpstream returns a channel closed once an upstream is next added for
// key, for a SUBSCRIBE held until its track has one (§10.2.6), and the stop
// that ends the wait.
func (r *TrackRegistry) AwaitUpstream(key track.Key) (arrived <-chan struct{}, stop func()) {
        r.mu.Lock()
        defer r.mu.Unlock()
        return r.arrivals.waitLocked(key, nil, &r.mu)
}

// AddUpstreamOption tweaks an [TrackRegistry.AddUpstream] call.
type AddUpstreamOption func(*addUpstreamConfig)

type addUpstreamConfig struct {
        setProperties bool
        properties    []byte
}

// WithProperties attaches Track Properties (§9.6) to the entry
// atomically with the first upstream-sub insertion, so the Discovery
// publish triggered by the same call sees them. Without this, a
// caller that sets Properties after AddUpstream returns sees an
// initial Discovery event with empty Properties followed by no
// update — the §9.6 properties end up missing from the
// cross-relay record. Passing them through AddUpstream avoids the gap.
func WithProperties(props []byte) AddUpstreamOption {
        return func(c *addUpstreamConfig) {
                c.setProperties = true
                c.properties = props
        }
}

// AddDownstream appends sub to the entry for fullName and returns the entry.
// Concurrency rules match [TrackRegistry.AddUpstream].
func (r *TrackRegistry) AddDownstream(fullName track.FullTrackName, sub *DownstreamSub) *TrackEntry {
        r.mu.Lock()
        defer r.mu.Unlock()
        entry := r.getOrCreateLocked(fullName)
        entry.mu.Lock()
        defer entry.mu.Unlock()
        entry.Downstream = append(entry.Downstream, sub)
        entry.downstreamGen++
        return entry
}

// AddDownstreamSnapshotLargest atomically appends sub to the entry's
// Downstream slice AND captures the current LargestObject watermark,
// both under a single entry.mu.Lock acquisition.
//
// Why atomic: handleSubscribe needs a [DownstreamSub.LargestAtSubscribe]
// snapshot that is consistent with the moment the sub becomes eligible
// for live fanout delivery. If the snapshot and append happen in
// separate lock cycles, a publisher write between them can:
//   - Run the fanout's UpdateLargest (under entry.mu) → advances Largest
//   - Cache the object
//   - Not deliver to this sub via live (the fanout's CopyDownstream
//     snapshot pre-dates our append)
//
// resulting in an object whose Location is > our snapshot AND was never
// pushed to us via live — a gap a fill fetch stream can't cover (its
// range ends at our snapshot, which doesn't include the missed object).
//
// Holding entry.mu across both operations serialises with
// [TrackEntry.UpdateLargest] (which also locks entry.mu): either we
// snapshot the pre-update Largest AND appear in any post-update
// CopyDownstream, or we snapshot the post-update Largest. Either way,
// every object the publisher has emitted is either covered by FETCH
// (via the snapshot) or delivered via live (via Downstream inclusion).
// AddDownstreamSnapshotLargest never creates an entry and requires at least
// one upstream to still be registered: ok=false means the track's last
// upstream vanished between the caller's establish check and this call
// (the §9.4 TOCTOU) — registering the downstream anyway would resurrect a
// sourceless entry whose subscriber then hangs with neither objects nor
// PUBLISH_DONE. The caller retries the establish step or rejects.
func (r *TrackRegistry) AddDownstreamSnapshotLargest(
        fullName track.FullTrackName,
        sub *DownstreamSub,
) (entry *TrackEntry, largest message.Location, hasLargest, ok bool) {
        r.mu.Lock()
        defer r.mu.Unlock()
        entry, exists := r.tracks[fullName.Key()]
        if !exists {
                return nil, message.Location{}, false, false
        }
        entry.mu.Lock()
        defer entry.mu.Unlock()
        if len(entry.Upstream) == 0 {
                return nil, message.Location{}, false, false
        }
        entry.Downstream = append(entry.Downstream, sub)
        entry.downstreamGen++
        return entry, entry.LargestObject, entry.HasLargestObject, true
}

// RemoveUpstream removes the upstream subscription with the given ID from
// the entry for fullName. Returns (removed, upstreamEmpty, entryDeleted):
//
//   - removed reports whether an entry with that ID was found and dropped.
//     It is false if the track is unknown or no upstream with subID was
//     present.
//   - upstreamEmpty reports whether the entry's Upstream slice is empty
//     after this call. The registry uses this signal to unpublish the
//     track from the Discovery Store.
//   - entryDeleted reports whether the whole [TrackEntry] was removed from
//     the registry as a consequence (both Upstream and Downstream became
//     empty). The bool is informational — the entry pointer is no longer
//     reachable through [TrackRegistry.Get] after this returns true.
//
// The fullName argument mirrors [TrackRegistry.AddUpstream] for API parity;
// internally we use only its Key. Remove never creates an entry.
//
// Like Add*, the whole remove operation runs under the registry write lock
// so the "decide to delete, then delete" sequence cannot race a concurrent
// Add that resurrects the entry.
func (r *TrackRegistry) RemoveUpstream(
        fullName track.FullTrackName,
        subID uint64,
) (removed, upstreamEmpty, entryDeleted bool) {
        key := fullName.Key()
        r.mu.Lock()
        entry, ok := r.tracks[key]
        if !ok {
                r.mu.Unlock()
                return false, false, false
        }

        entry.mu.Lock()
        before := len(entry.Upstream)
        var gone *UpstreamSub
        entry.Upstream = slices.DeleteFunc(entry.Upstream, func(s *UpstreamSub) bool {
                if s.ID == subID {
                        gone = s
                        return true
                }
                return false
        })
        removed = len(entry.Upstream) < before
        upstreamEmpty = len(entry.Upstream) == 0
        if !removed {
                entry.mu.Unlock()
                r.mu.Unlock()
                return false, upstreamEmpty, false
        }
        // Snapshot the downstream subs while we still hold the entry lock,
        // so we can notify them outside the registry locks: termination takes
        // each subscription's own locks, which must not nest inside r.mu.
        var notifyDownstreams []*DownstreamSub
        if upstreamEmpty && len(entry.Downstream) > 0 {
                notifyDownstreams = append([]*DownstreamSub(nil), entry.Downstream...)
        }
        allEmpty := upstreamEmpty && len(entry.Downstream) == 0
        entry.mu.Unlock()

        if allEmpty {
                delete(r.tracks, key)
                entryDeleted = true
        }
        if upstreamEmpty {
                // Still under r.mu: see [TrackRegistry.unpublishTrackFromDiscovery]
                // for why the unpublish must serialize with AddUpstream's publish.
                r.unpublishTrackFromDiscovery(entry)
        }
        r.mu.Unlock()

        if upstreamEmpty {
                // §10.12: carry the last upstream's reason; see DownstreamDoneCode.
                code := DownstreamDoneCode(gone.publishDone())
                for _, sub := range notifyDownstreams {
                        sub.TerminateWithPublishDone(code, "relay: upstream gone")
                }
        }
        return true, upstreamEmpty, entryDeleted
}

// RemoveDownstream removes the downstream subscription with the given ID
// from the entry for fullName. The return contract mirrors
// [TrackRegistry.RemoveUpstream]: (removed, downstreamEmpty, entryDeleted).
//
// When the removed subscription was the entry's LAST downstream, the
// relay's on-demand upstream subscriptions (§9.4 aggregation) have no
// consumers left: they are stripped from the entry in the same critical
// section — so a concurrent SUBSCRIBE cannot latch onto a dying upstream —
// and torn down via [UpstreamSub.CloseOnDemand] after the locks are
// released. PUBLISH-fed upstreams are untouched (their stream belongs to
// the publisher; future subscribers reuse it).
func (r *TrackRegistry) RemoveDownstream(
        fullName track.FullTrackName,
        subID uint64,
) (removed, downstreamEmpty, entryDeleted bool) {
        key := fullName.Key()
        r.mu.Lock()
        entry, ok := r.tracks[key]
        if !ok {
                r.mu.Unlock()
                return false, false, false
        }

        entry.mu.Lock()
        before := len(entry.Downstream)
        entry.Downstream = slices.DeleteFunc(entry.Downstream, func(s *DownstreamSub) bool {
                return s.ID == subID
        })
        removed = len(entry.Downstream) < before
        downstreamEmpty = len(entry.Downstream) == 0
        if !removed {
                entry.mu.Unlock()
                r.mu.Unlock()
                return false, downstreamEmpty, false
        }

        var stranded []*UpstreamSub
        hadUpstream := len(entry.Upstream) > 0
        if downstreamEmpty {
                stranded = stripOnDemandLocked(entry)
        }
        upstreamEmpty := len(entry.Upstream) == 0
        allEmpty := downstreamEmpty && upstreamEmpty
        entry.mu.Unlock()

        if allEmpty {
                delete(r.tracks, key)
                entryDeleted = true
        }
        if hadUpstream && upstreamEmpty {
                // Still under r.mu: see [TrackRegistry.unpublishTrackFromDiscovery]
                // for why the unpublish must serialize with AddUpstream's publish.
                r.unpublishTrackFromDiscovery(entry)
        }
        r.mu.Unlock()

        // Stream I/O happens outside the registry locks.
        for _, u := range stranded {
                u.CloseOnDemand()
        }
        return true, downstreamEmpty, entryDeleted
}

// stripOnDemandLocked removes every OnDemand upstream from entry and returns
// them for teardown; the caller must hold entry.mu and must CloseOnDemand
// the returned subs only after releasing the registry locks (stream I/O).
//
// Deliberate trade-off: when this empties the entry, the entry — including
// its object cache and LARGEST_OBJECT watermark — is deleted with it, so a
// FETCH arriving after the last subscriber left cold-starts via a fresh
// upstream instead of hitting warm cache. The §9.4 aggregation exists to
// serve live downstreams, not to keep publishers streaming into a void.
func stripOnDemandLocked(entry *TrackEntry) (stranded []*UpstreamSub) {
        entry.Upstream = slices.DeleteFunc(entry.Upstream, func(u *UpstreamSub) bool {
                if u.OnDemand {
                        stranded = append(stranded, u)
                }
                return u.OnDemand
        })
        return stranded
}

// RemoveSession bulk-evicts every UpstreamSub and DownstreamSub owned by
// sess across every track. Used by the session handler's defer in
// [Relay.handleConn] as a belt-and-suspenders measure: per-request handler
// defers already remove individual subscriptions on a clean shutdown, but
// they cannot run if a handler goroutine is wedged on a stale stream or
// raced past Stop. RemoveSession guarantees the registry is consistent
// after a session terminates regardless of why.
//
// Returns the number of upstream and downstream subscriptions removed (in
// that order). Tracks whose subscription slices both become empty are
// deleted from the registry in the same critical section as the slice
// edits.
func (r *TrackRegistry) RemoveSession(sess *session.Session) (upstreamRemoved, downstreamRemoved int) {
        r.mu.Lock()

        // Collect entries whose upstream slice transitions to empty so we
        // can notify their dependent downstream subscribers after releasing
        // the locks — termination takes each subscription's own locks, which
        // must not nest inside r.mu.
        // Their Discovery unpublish, by contrast, happens before r.mu is
        // released: see [TrackRegistry.unpublishTrackFromDiscovery].
        type orphaned struct {
                entry       *TrackEntry
                downstreams []*DownstreamSub
        }
        var (
                orphans  []orphaned
                stranded []*UpstreamSub
        )

        for key, entry := range r.tracks {
                entry.mu.Lock()
                beforeU := len(entry.Upstream)
                entry.Upstream = slices.DeleteFunc(entry.Upstream, func(s *UpstreamSub) bool {
                        return s.Session == sess
                })
                upstreamRemoved += beforeU - len(entry.Upstream)
                hadUpstream := beforeU > 0
                nowEmptyU := len(entry.Upstream) == 0
                if hadUpstream && nowEmptyU {
                        // Snapshot the surviving downstreams BEFORE we strip
                        // the ones owned by sess — a publisher session that
                        // also has downstreams on the same track (a relay
                        // chain configuration) shouldn't notify itself.
                        var notify []*DownstreamSub
                        for _, d := range entry.Downstream {
                                if d.Session != sess {
                                        notify = append(notify, d)
                                }
                        }
                        orphans = append(orphans, orphaned{entry: entry, downstreams: notify})
                }

                beforeD := len(entry.Downstream)
                entry.Downstream = slices.DeleteFunc(entry.Downstream, func(s *DownstreamSub) bool {
                        return s.Session == sess
                })
                removedHereD := beforeD - len(entry.Downstream)
                downstreamRemoved += removedHereD

                // The dying session may have been a track's last downstream: the
                // relay's on-demand upstream subscriptions on OTHER sessions then
                // have no consumers left — strip them (same rule as
                // [TrackRegistry.RemoveDownstream]) and tear them down after the
                // locks drop. Gated on removedHereD so a disconnect never touches
                // tracks this session had no downstream on: a zero-downstream
                // entry may be another handler's in-flight registration (upstream
                // established, downstream not yet added). Upstreams on sess itself
                // were already removed above.
                if removedHereD > 0 && len(entry.Downstream) == 0 && len(entry.Upstream) > 0 {
                        if s := stripOnDemandLocked(entry); len(s) > 0 {
                                stranded = append(stranded, s...)
                                if len(entry.Upstream) == 0 {
                                        orphans = append(orphans, orphaned{entry: entry})
                                }
                        }
                }

                empty := len(entry.Upstream) == 0 && len(entry.Downstream) == 0
                entry.mu.Unlock()
                if empty {
                        delete(r.tracks, key)
                }
        }
        for _, o := range orphans {
                // Still under r.mu: see [TrackRegistry.unpublishTrackFromDiscovery]
                // for why the unpublish must serialize with AddUpstream's publish.
                r.unpublishTrackFromDiscovery(o.entry)
        }
        r.mu.Unlock()

        for _, u := range stranded {
                u.CloseOnDemand()
        }
        for _, o := range orphans {
                for _, sub := range o.downstreams {
                        sub.TerminateWithPublishDone(moqt.PublishDoneTrackEnded,
                                "relay: publisher session gone")
                }
        }
        return upstreamRemoved, downstreamRemoved
}

// publishTrackToDiscovery advertises the entry to the Discovery store
// if one is configured. Called when the first UpstreamSub lands on a
// track. The caller MUST hold entry.mu — Properties is read under it.
// The Discovery call itself runs synchronously on the caller's
// goroutine, with a short context to avoid wedging the hot path on a
// misbehaving backend. Errors are logged at Warn but never propagated.
func (r *TrackRegistry) publishTrackToDiscovery(entry *TrackEntry) {
        if r.discovery == nil {
                return
        }
        info := discovery.TrackInfo{
                Key:        entry.Key,
                FullName:   entry.FullName,
                Properties: entry.Properties,
                RelayAddr:  r.relayAddr,
        }
        ctx, cancel := context.WithTimeout(context.Background(), discoveryCallTimeout)
        defer cancel()
        // ErrWithdrawn is not a failure: the relay is shutting down and the store
        // has deliberately stopped accepting advertisements, so every track still
        // draining would log one of these.
        if err := r.discovery.PublishTrack(ctx, info); err != nil && !errors.Is(err, discovery.ErrWithdrawn) {
                r.log.Warn("discovery: PublishTrack failed", "err", err.Error(), "key", info.Key)
        }
}

// unpublishTrackFromDiscovery is the counterpart called when the last
// UpstreamSub leaves a track.
//
// The caller MUST hold r.mu. Both this and [publishTrackToDiscovery] run
// under the registry lock so the store receives publish/unpublish calls in
// exactly the order the registry's upstream count crossed 0 — a late
// unpublish issued after releasing r.mu could race a concurrent
// AddUpstream's publish and erase the re-published track's record, leaving
// a live track invisible cross-relay until its upstream cycles (nothing
// re-publishes without another 0→1 transition). The Discovery call is
// bounded by [discoveryCallTimeout] — and the interface requires backends
// to honor ctx deadlines — so the lock hold is bounded too. RemoveSession
// pays this once per orphaned track, serially; with a degraded backend
// that is N × the timeout, an accepted worst case for a best-effort
// advertisement fabric.
func (r *TrackRegistry) unpublishTrackFromDiscovery(entry *TrackEntry) {
        if r.discovery == nil {
                return
        }
        ctx, cancel := context.WithTimeout(context.Background(), discoveryCallTimeout)
        defer cancel()
        if err := r.discovery.UnpublishTrack(ctx, entry.Key, r.relayAddr); err != nil {
                r.log.Warn("discovery: UnpublishTrack failed", "err", err.Error(), "key", entry.Key)
        }
}

// Package relaytest holds helpers shared across the relay tests (the
// relay_test and registry_test packages). Keeping them here avoids
// duplicating the same helper across test files that cannot otherwise share
// unexported code.
package relaytest

import (
        "context"
        "errors"
        "strings"
        "testing"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
)

// FormatNamespace renders a Track Namespace as a readable slash-joined string
// for test failure messages. Shared by the relay_test and registry_test
// packages.
func FormatNamespace(ns wire.TrackNamespace) string {
        if len(ns) == 0 {
                return "<root>"
        }
        var out strings.Builder
        for i, f := range ns {
                if i > 0 {
                        out.WriteString("/")
                }
                out.Write(f)
        }
        return out.String()
}

// ReadNextMessage parses one full MoQT control message off stream, failing
// the test if reading takes longer than the deadline allows. A context
// cancellation while blocked in Parse is treated as a clean unblock and
// returns the (possibly nil) partial message rather than failing.
func ReadNextMessage(t *testing.T, stream session.Stream, deadline <-chan time.Time) message.Message {
        t.Helper()
        done := make(chan struct{})
        var (
                msg message.Message
                err error
        )
        go func() {
                defer close(done)
                msg, err = message.Parse(stream)
        }()
        select {
        case <-done:
                if err != nil && !errors.Is(err, context.Canceled) {
                        t.Fatalf("message.Parse: %v", err)
                }
                return msg
        case <-deadline:
                t.Fatal("timeout waiting for next message")
                return nil
        }
}

package relay

import "sync"

// sessionLimiter enforces the per-session resource caps from §13.1
// (subscription amplification) and §13.7.1 (relay state maintenance): a bound
// on the number of concurrently-active subscriptions and on concurrently-active
// namespace-state requests (PUBLISH_NAMESPACE / SUBSCRIBE_NAMESPACE /
// SUBSCRIBE_TRACKS) a single session may hold. A non-positive max disables the
// corresponding limit (the relay's default — these are deployment policies).
//
// The counts track in-flight request handlers: acquire is called at dispatch
// before a handler is spawned, release when it returns (a handler runs for its
// request's whole lifetime); a TRACK_STATUS holds a subscription slot only
// while it is forwarded upstream. An over-limit request is rejected with
// REQUEST_ERROR EXCESSIVE_LOAD before any shared state is mutated.
type sessionLimiter struct {
        mu      sync.Mutex
        subs    int
        ns      int
        maxSubs int
        maxNS   int
}

func (l *sessionLimiter) acquireSub() bool       { return l.acquire(&l.subs, l.maxSubs) }
func (l *sessionLimiter) releaseSub()            { l.release(&l.subs, l.maxSubs) }
func (l *sessionLimiter) acquireNamespace() bool { return l.acquire(&l.ns, l.maxNS) }
func (l *sessionLimiter) releaseNamespace()      { l.release(&l.ns, l.maxNS) }

// acquire reserves a slot in the counter *n bounded by limit. It returns false
// (without incrementing) when the limit is already reached, and true otherwise.
// A non-positive limit means unlimited.
func (l *sessionLimiter) acquire(n *int, limit int) bool {
        if limit <= 0 {
                return true
        }
        l.mu.Lock()
        defer l.mu.Unlock()
        if *n >= limit {
                return false
        }
        *n++
        return true
}

// release returns a slot previously taken by acquire. It is a no-op when the
// limit is disabled, and clamps at zero defensively.
func (l *sessionLimiter) release(n *int, limit int) {
        if limit <= 0 {
                return
        }
        l.mu.Lock()
        if *n > 0 {
                *n--
        }
        l.mu.Unlock()
}

package relay

// Leg says which side of a relay mesh the session carrying an event sits on.
//
// It is deliberately a property of the *session*, not of the peer's role: a
// relay knows for certain who dialled whom, and nothing else about a peer is
// trustworthy. A peer relay that dials *in* is therefore [LegLocal], exactly
// like a browser — the relay has no reliable way to tell them apart, and
// guessing from the MOQT_IMPLEMENTATION string a peer volunteers would be a
// label an operator could not trust.
//
// That asymmetry is not a gap, because a cross-relay hop is observed from both
// ends. The consuming relay reports the hop as [LegUpstream] (it dialled), and
// the producing relay reports the same hop among its [LegLocal] traffic (it was
// dialled). Scrape both instances and the hop is the difference between the
// two: objects the consumer received on its upstream leg, against objects the
// producer forwarded. Objects that leave one and never arrive at the other are
// lost in the middle.
type Leg uint8

const (
        // LegLocal is a session a peer opened to this relay: an ordinary
        // publisher or subscriber, or a peer relay that dialled in.
        LegLocal Leg = iota

        // LegUpstream is a session this relay dialled out to a peer it found
        // through [discovery.DiscoveryStore] — the cross-relay hop, from the
        // consuming side.
        LegUpstream
)

// String returns a stable, lowercase name suitable for use as a metric label
// value. New Leg values may be added; an unknown one renders as "unknown"
// rather than a number, so a label never turns into a cardinality surprise.
func (l Leg) String() string {
        switch l {
        case LegLocal:
                return "local"
        case LegUpstream:
                return "upstream"
        default:
                return "unknown"
        }
}

// TrackRef identifies the track an event happened on. It is passed by value on
// the per-object hot path and holds no pointers, so it costs a copy and no
// allocation.
//
// Name is the track name half of the full track name — for a Media Sync Format
// producer, names like "catalog", "video" and "audio". The *namespace* is
// deliberately absent: it carries the publisher's identity (in a conference,
// one namespace per participant), so a metrics backend keyed on it would grow a
// new time series per participant per call and never retire them. A backend
// that wants per-publisher detail should sample it out-of-band, not from the
// hot path.
//
// Name comes off the wire and is chosen by the publisher, so it is NOT
// inherently bounded either. An implementation that turns it into a label MUST
// fold unrecognised names into a catch-all bucket.
type TrackRef struct {
        Name string
        Leg  Leg
}

// ResetCause explains why the relay tore down a subgroup stream or a whole
// subscription. It is the distinction that matters when a subscriber's picture
// breaks up between keyframes: the relay abandoning a subgroup stream mid-group
// loses the rest of that group's objects, and each cause below implies a
// different fix.
type ResetCause uint8

const (
        // ResetCauseGap is a §11.4.3 reopen: the object to forward was not known
        // to be "the next Object" after the last one written, so the current
        // outbound stream was reset and a fresh one opened. That is correct
        // behaviour, but it usually follows an earlier drop, or Objects arriving
        // from several upstreams, and a subscriber sees the hole.
        ResetCauseGap ResetCause = iota

        // ResetCauseDeliveryTimeout is §8: an object sat unsent past the
        // resolved publisher/subscriber delivery timeout, so
        // [session.OutgoingSubgroupStream] reset this one stream with
        // DELIVERY_TIMEOUT. The subscription survives; the rest of that
        // subgroup does not.
        ResetCauseDeliveryTimeout

        // ResetCauseTooFarBehind is the §3.3.4 TOO_FAR_BEHIND verdict: an
        // object waited in the send queue longer than [Config.MaxFanoutLag],
        // so the subscription was terminated rather than allowed to trail the
        // live edge indefinitely.
        ResetCauseTooFarBehind

        // ResetCauseExcessiveLoad is the optional [Config.MaxDropsBeforeReset]
        // backstop: cumulative drops on one subscription passed the cap and it
        // was terminated with EXCESSIVE_LOAD.
        ResetCauseExcessiveLoad

        // ResetCauseInboundReset is §11.4.3 propagation: the upstream stream
        // feeding this subgroup was reset (or its session went away), so the
        // corresponding downstream stream is reset rather than FIN'd.
        ResetCauseInboundReset

        // ResetCauseWriteError is a transport write failure on the outbound
        // stream — the subscriber's session is in trouble, not the relay's
        // scheduling.
        ResetCauseWriteError

        // ResetCauseProperties is a reopen to carry Object Properties (§2.5):
        // the outbound stream's SUBGROUP_HEADER had PROPERTIES clear (§11.4.2),
        // taken from another upstream of the same Subgroup (§9.3), so a fresh
        // stream with the bit set was opened. Objects already written survive
        // the reset only where RESET_STREAM_AT is in use (see
        // [session.OutgoingSubgroupStream.MarkReliable]); elsewhere the unacked
        // ones are lost, as with [ResetCauseGap].
        ResetCauseProperties
)

// String returns a stable, lowercase name suitable for use as a metric label
// value. Unknown values render as "unknown" rather than a number.
func (c ResetCause) String() string {
        switch c {
        case ResetCauseGap:
                return "gap"
        case ResetCauseDeliveryTimeout:
                return "delivery_timeout"
        case ResetCauseTooFarBehind:
                return "too_far_behind"
        case ResetCauseExcessiveLoad:
                return "excessive_load"
        case ResetCauseInboundReset:
                return "inbound_reset"
        case ResetCauseWriteError:
                return "write_error"
        case ResetCauseProperties:
                return "properties"
        default:
                return "unknown"
        }
}

// Metrics receives lifecycle and hot-path event notifications from a Relay so
// operators can wire relay activity into their own telemetry backend
// (Prometheus, OpenTelemetry, statsd, …) without this package depending on any
// of them. Install one via [Config.Metrics]; the default is [NopMetrics].
//
// All methods are invoked from relay goroutines, concurrently and — for
// ObjectReceived / ObjectForwarded / ObjectDropped — on the per-object fanout
// hot path, while the subgroup's fanout lock is held. An implementation MUST be
// safe for concurrent use and MUST NOT block: do the cheap thing (e.g. an
// atomic increment, or a counter handle looked up once and cached) and
// aggregate elsewhere. Blocking here stalls the inbound read loop for every
// subscriber of the subgroup, not just one.
//
// The interface may grow over time. Embed [NopMetrics] in your implementation
// so unimplemented methods default to a no-op and future additions stay
// backward-compatible.
type Metrics interface {
        // SessionOpened is called when a session completes SETUP and is
        // registered; SessionClosed is called exactly once per SessionOpened when
        // that session's handler tears down. Together they track the live-session
        // gauge, split by [Leg] so the count of live cross-relay hops is visible
        // separately from client sessions.
        SessionOpened(leg Leg)
        SessionClosed(leg Leg)

        // SubscriptionOpened is called when a downstream SUBSCRIBE is accepted and
        // registered for fanout; SubscriptionClosed is called exactly once per
        // SubscriptionOpened when the subscription is removed. Together they track
        // the active-subscription gauge.
        SubscriptionOpened(t TrackRef)
        SubscriptionClosed(t TrackRef)

        // ObjectReceived is called once for each object read off an inbound
        // subgroup stream and won by this contributor — objects discarded as
        // §9.3 duplicates of a redundant upstream are not counted. Compared
        // against ObjectForwarded it separates "the relay never got it" from
        // "the relay got it and shed it".
        ObjectReceived(t TrackRef, subgroup uint64)

        // ObjectForwarded is called once for each object successfully enqueued for
        // delivery to a downstream subscriber. It is counted per subscriber, so a
        // single received object fanned out to N subscribers reports N times —
        // which is why it is not directly comparable to ObjectReceived without
        // dividing by the subscriber count.
        ObjectForwarded(t TrackRef, subgroup uint64)

        // ObjectDropped is called when a downstream subscriber's bounded send
        // queue overflows and the object is dropped (§8 slow-reader pressure).
        // The subgroup is reported because it is how a layered publisher marks
        // what is disposable: shedding an enhancement layer is the design
        // working, and shedding the base layer is the picture breaking.
        ObjectDropped(t TrackRef, subgroup uint64)

        // SubgroupStreamReset is called when one outbound subgroup stream is torn
        // down before its subgroup ended, for any of the [ResetCause] reasons. The
        // subscription itself survives; the remainder of that subgroup does not
        // reach this subscriber.
        SubgroupStreamReset(t TrackRef, subgroup uint64, cause ResetCause)

        // SubscriptionResetSlowReader is called when the relay forcibly resets a
        // subscriber's outbound stream and terminates the whole subscription
        // because it fell too far behind: an object waited longer than
        // [Config.MaxFanoutLag] in the send queue
        // ([ResetCauseTooFarBehind], the primary trigger), or the optional
        // cumulative [Config.MaxDropsBeforeReset] cap was exceeded
        // ([ResetCauseExcessiveLoad]).
        SubscriptionResetSlowReader(t TrackRef, cause ResetCause)

        // FetchServed is called when a FETCH is answered from the relay's object
        // cache, with the number of objects returned (0 when the requested range
        // produced no cached objects).
        FetchServed(t TrackRef, objects int)

        // UpstreamDialFailed is called when the upstream pool could not establish
        // a relay-to-relay session with a peer advertised in Discovery. relayAddr
        // is the peer address that failed, for logging and exemplars — it is
        // operator-controlled but grows with the mesh, so an implementation
        // SHOULD NOT make it a label.
        UpstreamDialFailed(relayAddr string)

        // NamespaceResolved is called after each Discovery FindNamespace lookup
        // on the cross-relay path, with the number of peer relays advertising the
        // namespace (0 when nobody does — the case where a subscriber gets
        // nothing and no error explains why).
        NamespaceResolved(advertisers int)
}

// NopMetrics is the no-op [Metrics] installed when [Config.Metrics] is nil.
// Embed it in a custom implementation to inherit no-op defaults for the methods
// you don't care about — which also keeps your type compiling as the [Metrics]
// interface grows:
//
//        type myMetrics struct {
//                relay.NopMetrics
//                dropped atomic.Int64
//        }
//
//        func (m *myMetrics) ObjectDropped(relay.TrackRef, uint64) { m.dropped.Add(1) }
type NopMetrics struct{}

var _ Metrics = NopMetrics{}

func (NopMetrics) SessionOpened(Leg)                                {}
func (NopMetrics) SessionClosed(Leg)                                {}
func (NopMetrics) SubscriptionOpened(TrackRef)                      {}
func (NopMetrics) SubscriptionClosed(TrackRef)                      {}
func (NopMetrics) ObjectReceived(TrackRef, uint64)                  {}
func (NopMetrics) ObjectForwarded(TrackRef, uint64)                 {}
func (NopMetrics) ObjectDropped(TrackRef, uint64)                   {}
func (NopMetrics) SubgroupStreamReset(TrackRef, uint64, ResetCause) {}
func (NopMetrics) SubscriptionResetSlowReader(TrackRef, ResetCause) {}
func (NopMetrics) FetchServed(TrackRef, int)                        {}
func (NopMetrics) UpstreamDialFailed(string)                        {}
func (NopMetrics) NamespaceResolved(int)                            {}

// This file holds the relay's lifecycle scaffold: the transport-agnostic
// Listener interface, the Relay struct, and Start/Stop. The remaining
// components (Track Registry, Namespace Registry, Subscription Fanout,
// Object Cache, Discovery Store) live in sibling files and plug into this
// scaffold. See doc.go for the package overview and the file-layer map.
package relay

import (
        "context"
        "errors"
        "fmt"
        "log/slog"
        "net"
        "sync"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/relay/discovery"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// Listener yields ready-to-use MOQT transport connections. The caller is
// responsible for TLS, ALPN ("moqt-20"), and — for WebTransport — the HTTP/3
// CONNECT upgrade before returning a Conn. The relay never binds sockets or
// terminates TLS itself.
//
// Implementations of this interface live in the transport adapter packages:
//
//   - quicconn.NewListener wraps a *quic.Listener.
//   - wtconn.NewListener wraps a webtransport.Server mounted on an
//     http.Handler.
//   - sessiontest provides an in-memory pipe listener for tests.
type Listener interface {
        // Accept blocks until the next MOQT-ready Conn is available, or ctx is
        // cancelled, or the listener is closed. The returned Conn has TLS and
        // ALPN already negotiated; the relay only needs to drive the MOQT
        // SETUP handshake on top.
        Accept(ctx context.Context) (session.Conn, error)

        // Addr returns the network address the listener is bound to. May be nil
        // for purely in-process listeners.
        Addr() net.Addr

        // Close stops the listener. After Close returns, Accept must return
        // promptly with an error. Close is safe to call from any goroutine and
        // may be invoked more than once.
        Close() error
}

// discoveryWithdrawTimeout bounds [Relay.Stop]'s Discovery withdrawal. It is the
// first step of shutdown, so a backend that has gone away must not be able to
// delay the GOAWAY broadcast behind it: the advertisements expire on their own
// once the liveness TTL lapses. Deliberately far longer than the registries'
// per-call discoveryCallTimeout — this is one RPC on the shutdown path, not a
// call made under a registry lock. Worst case it adds this much to Stop on top
// of GoawayTimeout, which still leaves room inside a typical 30s orchestrator
// termination grace period.
const discoveryWithdrawTimeout = 5 * time.Second

// Config carries all relay knobs. It bundles transport-agnostic
// scheduling parameters (queue sizes, reset thresholds, cache bounds),
// pluggable hooks (Authorizer, Discovery), the GOAWAY grace period,
// and SETUP-time SessionOptions.
type Config struct {
        // GoawayTimeout is the grace period the relay grants downstream
        // sessions to migrate after Stop sends GOAWAY before forcibly closing
        // them. Zero means "do not send GOAWAY, just close" — useful in tests.
        GoawayTimeout time.Duration

        // SessionOptions are forwarded to session.Server() for every accepted
        // connection. Use this to advertise implementation name, GREASE,
        // MAX_AUTH_TOKEN_CACHE_SIZE, etc. Optional.
        SessionOptions []session.Option

        // KnownMandatoryTrackProperties lists the Mandatory Track Property types
        // (0x4000–0x7FFF) the relay may forward; a track carrying any other is
        // refused with UNSUPPORTED_EXTENSION (§2.5.1). Empty (the default)
        // refuses every Mandatory Track Property. Set it here rather than with
        // session.WithKnownMandatoryTrackProperties in SessionOptions, which
        // overrides this field.
        KnownMandatoryTrackProperties []message.PropertyType

        // Logger is used for relay-level events (accept loop start/stop,
        // session setup failures, GOAWAY broadcast). If nil, slog.Default() is
        // used. Per-session loggers are derived via Logger.With(...).
        Logger *slog.Logger

        // Authorizer gates every incoming request before the relay performs
        // any state mutation. If nil, [AllowAllAuthorizer] is used, which is
        // appropriate for tests and trusted in-process deployments.
        // Production should supply a token- or session-attestation-aware
        // implementation. See [Authorizer] for the full contract.
        Authorizer Authorizer

        // Metrics receives relay lifecycle and hot-path event notifications for
        // telemetry. nil (the default) installs [NopMetrics]. See [Metrics] for
        // the contract; implementations MUST be non-blocking and safe for
        // concurrent use.
        Metrics Metrics

        // MaxFilterRanges is the MAX_FILTER_RANGES (§10.3.1.6) budget this relay
        // advertises in SETUP: the largest total number of Range Filter ranges
        // (§5.1.4) it will accept across every Range Filter parameter on one
        // SUBSCRIBE or FETCH. Over-budget requests are answered INVALID_FILTER.
        //
        // Zero means: use [DefaultMaxFilterRanges]. A negative value advertises 0,
        // which prohibits Range Filters outright — the session default, and the
        // reason this field exists. The filters are implemented and enforced
        // throughout, but [session.WithMaxFilterRanges] defaults to 0, so a relay
        // that never sets it rejects every Range Filter it is sent. That silently
        // disables SUBGROUP_FILTER, which is how a subscriber declines a track's
        // upper temporal layers or fetches only its base layer — a request that
        // looks supported, and is, until the SETUP budget refuses it.
        MaxFilterRanges int

        // MaxCacheSize bounds the per-track Object Cache by object count.
        // Zero means: use [registry.DefaultCacheMaxSize]. The bound is applied
        // independently to every track the relay observes; a noisy track
        // cannot evict a quiet one's entries.
        MaxCacheSize int

        // MaxCacheDuration bounds the per-track Object Cache by object age.
        // Zero means: use [registry.DefaultCacheMaxDuration]. Objects older than this
        // are eligible for time-based eviction on the next Put.
        MaxCacheDuration time.Duration

        // CacheTTLPolicy, when non-nil, may override [MaxCacheDuration] on
        // a per-track basis. See [CacheTTLPolicy] for the contract; the
        // function is invoked once per [registry.TrackEntry] at creation time and
        // never on the fanout hot path. Use this to give well-known tracks
        // (e.g. an MSF catalog track) infinite retention without changing
        // the default for everything else. [pkg/relay] does not own the
        // rule; the binary supplies it.
        CacheTTLPolicy CacheTTLPolicy

        // SendQueueSize is the per-downstream-subscriber bounded channel
        // size used by the fanout writer. Each subscriber's writer goroutine
        // consumes from this queue; the fanout publishes to all queues with
        // a non-blocking send and drops the object on overflow. A larger
        // queue absorbs more transient burst but lets a slow reader keep
        // more memory locked up.
        // Zero means: use the default of 64.
        SendQueueSize int

        // MaxFanoutLag bounds how far behind the live edge a downstream
        // subscriber may fall before the relay resets its outbound subgroup
        // streams and terminates the subscription. The fanout writer measures
        // the time each forwarded object spends queued before it is written; an
        // object that waited longer than MaxFanoutLag means the subscriber has
        // been unable to keep up for that long, so it is dropped. This is a
        // latency window, not a drop count: a subscriber that loses the
        // occasional object but stays current is left alone, while one that
        // steadily falls behind is shed. Zero means: use the default of 2s.
        MaxFanoutLag time.Duration

        // MaxDropsBeforeReset is an OPTIONAL hard cap on the cumulative number of
        // objects dropped to one subscriber's overflowing send queue, after which
        // the relay resets and terminates the subscription. It is a coarse
        // backstop to [MaxFanoutLag] (e.g. to bound memory for a peer that
        // accepts a stream but never reads it); the time window is the primary
        // slow-reader signal. Zero (the default) disables the cap.
        MaxDropsBeforeReset int

        // MaxSubscriptionsPerSession bounds the number of concurrently-active
        // SUBSCRIBE requests a single session may hold (§13.1, subscription
        // amplification), counting a TRACK_STATUS while the relay forwards it
        // upstream. Excess requests are rejected with REQUEST_ERROR
        // EXCESSIVE_LOAD before any state is mutated. Zero (the default) means
        // unlimited — limits are a deployment policy the operator opts into.
        MaxSubscriptionsPerSession int

        // MaxNamespaceRequestsPerSession bounds the number of concurrently-active
        // namespace-state requests (PUBLISH_NAMESPACE, SUBSCRIBE_NAMESPACE,
        // SUBSCRIBE_TRACKS) a single session may hold (§13.7.1, relay state
        // maintenance). Excess requests are rejected with REQUEST_ERROR
        // EXCESSIVE_LOAD. Zero (the default) means unlimited.
        MaxNamespaceRequestsPerSession int

        // MaxRendezvousTimeout caps how long the relay holds a SUBSCRIBE that
        // carries RENDEZVOUS_TIMEOUT for a track with no publisher, waiting for
        // one to appear (§10.2.6: "The relay MAY use a shorter timeout than
        // requested by the subscriber"). What is left of the hold is forwarded
        // on the relay's upstream SUBSCRIBEs to other relays, so the cap bounds
        // theirs too. Zero means: use the default of 30s. A negative value holds
        // none: the SUBSCRIBE is answered DOES_NOT_EXIST at once, as if
        // RENDEZVOUS_TIMEOUT were 0.
        MaxRendezvousTimeout time.Duration

        // Discovery is the cross-instance track + namespace advertisement
        // fabric. nil means "no discovery" — the relay still works as a
        // single-instance setup with no cross-relay routing. Single-process
        // tests typically leave this nil; multi-instance deployments inject
        // a [discovery.MemoryStore] (local-only) or a distributed backend
        // (NATS / Redis).
        Discovery discovery.DiscoveryStore

        // RelayAddr is the address this relay registers itself as in
        // Discovery entries. Empty for single-instance deployments
        // (Discovery still works, RelayAddr just stays empty). NATS / Redis
        // backends use this to route upstream connections to the right
        // peer.
        RelayAddr string

        // Dialer establishes an outbound transport connection to another relay
        // instance, given the RelayAddr that instance advertised in Discovery.
        // It is the outbound counterpart of [Listener]: the relay stays
        // transport-agnostic, so the caller owns TLS, ALPN ("moqt-20"), and — for
        // WebTransport — the HTTP/3 CONNECT upgrade, returning a ready
        // [session.Conn] on which the relay drives the MOQT SETUP handshake as a
        // client.
        //
        // nil (the default) disables cross-relay dialing: the relay serves only
        // from locally-connected publishers and never follows a Discovery
        // [discovery.FindNamespace] result to a remote peer. Set this together
        // with Discovery to enable on-demand cross-relay upstream SUBSCRIBE.
        //
        // The relay pools and reuses one session per RelayAddr; the Dialer is
        // invoked at most once per address while a session to it is live, and
        // again only after that session ends.
        Dialer func(ctx context.Context, relayAddr string) (session.Conn, error)

        // UpstreamFanIn optionally bounds how many remote relays a cross-relay
        // upstream SUBSCRIBE fans into for one namespace. §9.5 requires a relay to
        // subscribe to every publisher that advertised the namespace, so the zero
        // value (the default) does exactly that — full fan-in, no data loss in any
        // topology. Like the other limits on this struct, zero means "no limit," and
        // bounding is a deployment policy the operator opts into.
        //
        // A positive N is that opt-in: the pool ranks the advertising relays by
        // rendezvous (HRW) weight — a deterministic order every relay computes
        // identically — and subscribes only to the top N that dial successfully.
        // Because the ranking is identical fleet-wide, relays converge on the same
        // few upstreams per namespace, collapsing the relay-to-relay stream count
        // from a full O(n²) mesh toward a tree (N is the redundancy width: 1 is a
        // pure tree, 2 keeps one backup upstream). This is a deliberate deviation
        // from §9.5's "subscribe to all," sound only where the advertisers are
        // redundant sources of the same objects — which the relay's fanout already
        // dedups — never where different relays hold distinct objects for the track.
        //
        // Only meaningful when Dialer and Discovery are set; ignored otherwise.
        UpstreamFanIn int
}

// resolved Config defaults; kept as constants so tests can reference them
// without poking at private fields.
const (
        defaultSendQueueSize        = 64
        defaultMaxFanoutLag         = 2 * time.Second
        defaultMaxRendezvousTimeout = 30 * time.Second
)

// DefaultMaxFilterRanges is the MAX_FILTER_RANGES (§10.3.1.6) budget a relay
// advertises when [Config.MaxFilterRanges] is left at zero.
//
// Sixteen ranges across all of a request's Range Filters. The uses this exists
// for are small — a SUBGROUP_FILTER naming one layer or a contiguous band of
// them is one range, an OBJECTID_FILTER picking a group's base-layer ID range
// is another — so sixteen is several such filters at once and still bounds the
// per-object matching work to something a fanout can afford. It is a budget
// against a peer asking for arbitrarily many bands, not a working limit.
const DefaultMaxFilterRanges = 16

// resolveMaxFilterRanges maps [Config.MaxFilterRanges] onto the value
// advertised in SETUP: zero takes the default, negative prohibits Range
// Filters, positive is taken as given.
func resolveMaxFilterRanges(configured int) uint64 {
        switch {
        case configured == 0:
                return DefaultMaxFilterRanges
        case configured < 0:
                return 0
        default:
                return uint64(configured)
        }
}

// Relay is a single MOQT relay instance. It owns one Listener, accepts
// session.Conn values from it, drives the MOQT SETUP handshake, and dispatches
// each established Session to a handler goroutine.
//
// A Relay is created with New and started with Start. Start blocks until the
// context is cancelled or Stop is called. Stop is safe to call concurrently
// with Start and may be invoked at most once meaningfully — subsequent calls
// are no-ops.
type Relay struct {
        listener Listener
        cfg      Config
        log      *slog.Logger

        // tracks and names are the relay-wide registries shared across every
        // session handler.
        tracks *registry.TrackRegistry
        names  *registry.NamespaceRegistry

        // fetch rendezvouses upstream FETCH response streams (dispatched by the
        // upstream session's data loop) with the downstream handler that issued
        // the FETCH. Shared across every session handler.
        fetch *registry.FetchRouter

        // statusRounds shares forwarded TRACK_STATUS rounds across every
        // session handler; see [sessionHandler.forwardTrackStatus].
        statusRounds *trackStatusRounds

        // upstreams dials and pools relay-to-relay sessions for Discovery-driven
        // cross-relay upstream SUBSCRIBE. nil when Config.Dialer is unset (the
        // single-instance case); session handlers treat a nil pool as "no
        // cross-relay routing available".
        upstreams *upstreamPool

        // watchWG tracks the optional Discovery WatchNamespaces consumer
        // goroutine started in Start, so Stop joins it before returning.
        watchWG sync.WaitGroup

        // sessions tracks every Session that has completed SETUP and not yet
        // been torn down. Stop iterates it under sessionsMu to broadcast GOAWAY
        // and to wait for drain. stopCtx, Stop's ctx, is set (under sessionsMu, by
        // beginShutdown) when Stop snapshots the set; addSession reads it under the
        // same lock to decide whether a newly-registered session is a straggler
        // Stop's snapshot missed, and bounds that straggler's drain by it.
        sessionsMu sync.Mutex
        sessions   map[*session.Session]struct{}
        stopCtx    context.Context

        // stopOnce guards Stop so the second caller short-circuits. stopCh is
        // closed by Stop to signal the accept loop to exit and to release any
        // per-session handlers blocked waiting on shutdown.
        stopOnce sync.Once
        stopCh   chan struct{}

        // handlers tracks per-session handler goroutines so Stop can wait for
        // them to finish before returning.
        handlers sync.WaitGroup
}

// New constructs a Relay backed by listener and configured by cfg. listener
// must be non-nil; New panics otherwise, because a relay without a transport
// source is never useful and the misconfiguration would otherwise surface as
// a confusing nil-pointer panic deep inside Start.
func New(listener Listener, cfg Config) *Relay {
        if listener == nil {
                panic("relay.New: listener is required")
        }
        log := cfg.Logger
        if log == nil {
                log = slog.Default()
        }
        if cfg.Authorizer == nil {
                cfg.Authorizer = AllowAllAuthorizer{}
        }
        if cfg.Metrics == nil {
                cfg.Metrics = NopMetrics{}
        }
        if cfg.SendQueueSize <= 0 {
                cfg.SendQueueSize = defaultSendQueueSize
        }
        if cfg.MaxFanoutLag <= 0 {
                cfg.MaxFanoutLag = defaultMaxFanoutLag
        }
        if cfg.MaxRendezvousTimeout == 0 {
                cfg.MaxRendezvousTimeout = defaultMaxRendezvousTimeout
        }
        // Prepended, so it is the SETUP budget unless the caller states one — and
        // stated twice it is advertised twice, which is why [Config.MaxFilterRanges]
        // is the way to change it rather than another WithMaxFilterRanges here.
        knownMandatory := make(map[message.PropertyType]struct{}, len(cfg.KnownMandatoryTrackProperties))
        for _, t := range cfg.KnownMandatoryTrackProperties {
                knownMandatory[t] = struct{}{}
        }
        cfg.SessionOptions = append(
                []session.Option{
                        session.WithMaxFilterRanges(resolveMaxFilterRanges(cfg.MaxFilterRanges)),
                        session.WithKnownMandatoryTrackProperties(knownMandatory),
                },
                cfg.SessionOptions...,
        )
        // MaxDropsBeforeReset is an opt-in hard cap: 0 means "disabled", so no
        // default is applied.
        if cfg.MaxCacheSize <= 0 {
                cfg.MaxCacheSize = registry.DefaultCacheMaxSize
        }
        if cfg.MaxCacheDuration <= 0 {
                cfg.MaxCacheDuration = registry.DefaultCacheMaxDuration
        }
        trackOpts := []registry.TrackRegistryOption{
                registry.WithCacheConfig(cfg.MaxCacheSize, cfg.MaxCacheDuration),
                registry.WithTrackRegistryLogger(log),
        }
        if cfg.CacheTTLPolicy != nil {
                trackOpts = append(trackOpts, registry.WithCacheTTLPolicy(registry.CacheTTLPolicy(cfg.CacheTTLPolicy)))
        }
        var nameOpts []registry.NamespaceRegistryOption
        nameOpts = append(nameOpts, registry.WithNamespaceRegistryLogger(log))
        if cfg.Discovery != nil {
                trackOpts = append(trackOpts, registry.WithTrackDiscovery(cfg.Discovery, cfg.RelayAddr))
                nameOpts = append(nameOpts, registry.WithNamespaceDiscovery(cfg.Discovery, cfg.RelayAddr))
        }

        r := &Relay{
                listener:     listener,
                cfg:          cfg,
                log:          log.With("component", "relay"),
                tracks:       registry.NewTrackRegistry(trackOpts...),
                names:        registry.NewNamespaceRegistry(nameOpts...),
                fetch:        registry.NewFetchRouter(),
                statusRounds: newTrackStatusRounds(),
                sessions:     make(map[*session.Session]struct{}),
                stopCh:       make(chan struct{}),
        }

        // When a Dialer is configured, the relay can follow Discovery
        // FindNamespace results to a remote peer. The pool dials and reuses one
        // session per RelayAddr and runs the relay's normal per-session loops on
        // each dialled session (via serveSession) so its inbound data streams fan
        // out and its FETCH responses route through the fetch router exactly as an
        // accepted session's do.
        if cfg.Dialer != nil {
                // Cross-relay routing keys on RelayAddr: it identifies this instance in
                // Discovery (so peers can dial it) and is the self-exclusion key that
                // keeps the relay from dialing or reflecting its own advertisements. An
                // empty RelayAddr breaks both — every Discovery entry looks unaddressable
                // / "ours", so FindNamespace dials nothing and the namespace watcher
                // reflects nothing. Warn loudly rather than fail silently.
                if cfg.RelayAddr == "" {
                        r.log.Warn("relay: Config.Dialer is set but Config.RelayAddr is empty; " +
                                "cross-relay routing is disabled (Discovery entries are indistinguishable " +
                                "from this relay's own and the relay is unaddressable). Set a unique RelayAddr.")
                }
                r.upstreams = newUpstreamPool(upstreamPoolConfig{
                        dialer:       cfg.Dialer,
                        discovery:    cfg.Discovery,
                        relayAddr:    cfg.RelayAddr,
                        sessionOpts:  cfg.SessionOptions,
                        log:          r.log,
                        metrics:      cfg.Metrics,
                        serveSession: r.serveUpstreamSession,
                        fanIn:        cfg.UpstreamFanIn,
                })
        }

        return r
}

// serveUpstreamSession starts the relay's per-session loops on a dialled
// upstream session in a tracked goroutine and invokes onClose when the session
// ends. It mirrors the accept-path bookkeeping ([Relay.handleConn]): the
// handler goroutine is registered with r.handlers so [Relay.Stop] joins it.
// The Add happens synchronously (before the goroutine) so it cannot race
// Stop's handlers.Wait. Called only by the upstream pool.
func (r *Relay) serveUpstreamSession(sess *session.Session, onClose func()) {
        r.handlers.Go(func() {
                defer onClose()
                r.serveSession(r.upstreams.baseCtx, sess, LegUpstream)
        })
}

// Addr returns the address the underlying Listener is bound to. Convenience
// wrapper; useful for tests that need to dial back into the relay.
func (r *Relay) Addr() net.Addr { return r.listener.Addr() }

// Authorizer returns the authorization hook the relay is currently using.
// Guaranteed non-nil after [New]: a nil [Config.Authorizer] is replaced with
// [AllowAllAuthorizer]. Primarily useful for tests; production code injects
// its policy through Config.
func (r *Relay) Authorizer() Authorizer { return r.cfg.Authorizer }

// Run serves the relay until ctx is cancelled or the Listener fails fatally,
// then performs the full §10.4 graceful shutdown — GOAWAY, [Config.GoawayTimeout]
// grace period, force-close — and returns only once that drain has finished. It
// returns [Relay.Start]'s error. This is the entry point a binary driven by
// [os/signal.NotifyContext] wants; Start and Stop remain available for callers
// that need to drive the phases themselves.
//
// ctx is the shutdown *trigger*, not the lifetime of the sessions. Run
// deliberately does not hand ctx to Start, because Start propagates its ctx down
// to every per-session handler: a signal-cancelled ctx would tear the sessions
// down underneath the drain and their peers would never see the GOAWAY. Sessions
// run under an internal context that outlives ctx and is unwound by Stop.
//
// shutdownTimeout caps the GOAWAY grace period: once it elapses Stop stops
// waiting for peers to migrate and force-closes whatever is left, even if
// [Config.GoawayTimeout] has not elapsed. Zero lets the grace period run to
// Config.GoawayTimeout. It does not bound the join of in-flight handlers that
// follows, so a wedged handler can still delay the return.
func (r *Relay) Run(ctx context.Context, shutdownTimeout time.Duration) error {
        // Trigger the drain on ctx without letting ctx reach the sessions.
        // AfterFunc's stop is deferred so a signal arriving after Run has already
        // returned cannot kick off a second drain.
        stopTrigger := context.AfterFunc(ctx, func() { r.shutdown(shutdownTimeout) })
        defer stopTrigger()

        // WithoutCancel keeps ctx's values (logging scope, tracing) while dropping
        // its cancellation. Nothing cancels this context: Start returns when Stop
        // closes the listener, and the sessions it parents are unwound by Stop's
        // force-close.
        err := r.Start(context.WithoutCancel(ctx))

        // Start returns as soon as Stop closes the listener, so the drain the
        // trigger started is normally still in flight. Join it: returning here would let the
        // caller exit the process mid-GOAWAY. Stop is idempotent and a second call
        // blocks until the first completes, which also covers the path where Start
        // failed on its own and no shutdown has begun yet.
        r.shutdown(shutdownTimeout)
        return err
}

// shutdown runs [Relay.Stop] under an optional timeout, logging a failure. The
// error is not propagated: it reports trouble closing the listener, which says
// nothing useful about why the relay is shutting down.
func (r *Relay) shutdown(timeout time.Duration) {
        ctx := context.Background()
        if timeout > 0 {
                var cancel context.CancelFunc
                ctx, cancel = context.WithTimeout(ctx, timeout)
                defer cancel()
        }
        if err := r.Stop(ctx); err != nil {
                r.log.LogAttrs(ctx, slog.LevelError, "relay stop failed", slog.String("err", err.Error()))
        }
}

// Start runs the relay accept loop until ctx is cancelled, Stop is called, or
// the Listener returns a fatal error. The returned error reports the cause:
//
//   - nil when shutdown was initiated cleanly via ctx cancellation or Stop.
//   - the Listener's Accept error otherwise.
//
// ctx is also the parent context of every accepted session's handler loops, so
// cancelling it terminates live sessions immediately — without GOAWAY. Do not
// wire a signal context straight into Start; use [Relay.Run], which separates
// the shutdown trigger from the sessions' lifetime.
//
// Start is intended to be called exactly once per Relay. Calling it twice
// concurrently is undefined.
func (r *Relay) Start(ctx context.Context) error {
        r.log.LogAttrs(ctx, slog.LevelInfo, "relay accept loop starting",
                slog.Any("addr", r.listener.Addr()))

        // Tie ctx to stopCh so a Stop call from another goroutine unblocks
        // Accept the same way a cancelled context would.
        acceptCtx, cancelAccept := context.WithCancel(ctx)
        defer cancelAccept()
        go func() {
                select {
                case <-r.stopCh:
                        cancelAccept()
                case <-acceptCtx.Done():
                }
        }()

        // Consume Discovery namespace events: forward namespaces advertised by
        // *other* relays to this relay's local SUBSCRIBE_NAMESPACE holders, so a
        // downstream subscriber learns about a namespace served elsewhere and can
        // then SUBSCRIBE (which the on-demand cross-relay path resolves via
        // FindNamespace). acceptCtx is cancelled by Stop (stopCh) or ctx, so the
        // watcher unwinds with the accept loop. Skipped without Discovery.
        if r.cfg.Discovery != nil {
                r.watchWG.Go(func() {
                        r.runNamespaceWatch(acceptCtx)
                })
        }

        for {
                conn, err := r.listener.Accept(acceptCtx)
                if err != nil {
                        // Shutdown paths look like context cancellation or
                        // net.ErrClosed; surface them as a clean nil so callers
                        // can distinguish "I asked to stop" from real listener
                        // failures.
                        if isShutdownErr(err) || acceptCtx.Err() != nil {
                                r.log.LogAttrs(ctx, slog.LevelInfo, "relay accept loop stopped")
                                return nil
                        }
                        r.log.LogAttrs(ctx, slog.LevelError, "relay listener accept failed",
                                slog.String("err", err.Error()))
                        return fmt.Errorf("relay: listener accept: %w", err)
                }

                r.handlers.Add(1)
                go r.handleConn(ctx, conn)
        }
}

// handleConn performs the MOQT SETUP handshake on conn and, on success, runs
// the per-session handler loops. SETUP failures close the underlying conn and
// log the cause; they do not propagate up to Start because one bad client
// must not take the relay down.
//
// This method owns the lifecycle: register the Session, run the
// per-session request / data / datagram loops, and unregister on exit.
func (r *Relay) handleConn(ctx context.Context, conn session.Conn) {
        defer r.handlers.Done()

        sess, err := session.Server(ctx, conn, r.cfg.SessionOptions...)
        if err != nil {
                r.log.LogAttrs(ctx, slog.LevelWarn, "relay SETUP failed",
                        slog.String("err", err.Error()))
                // session.Server already closed conn on failure; nothing more
                // to do here.
                return
        }

        // Start documents cancelling its ctx as terminating live sessions, and a
        // handler ending does not close its session: left open, a relay-scoped
        // reader on one of its request streams, and with it Stop, would wait on
        // it. NO_ERROR (§3.5): no GOAWAY was sent, so none ran out.
        closeSess := func() { _ = sess.Close(moqt.SessionNoError, "relay: stopped") }
        stop := context.AfterFunc(ctx, closeSess)
        r.serveSession(ctx, sess, LegLocal)
        // The handler can end on ctx before AfterFunc has run closeSess, and stop
        // then keeps it from running at all.
        stop()
        if ctx.Err() != nil {
                closeSess()
        }
}

// serveSession runs the per-session lifecycle for a Session that has already
// completed SETUP — whether accepted inbound by [Relay.handleConn] or dialled
// outbound by the [upstreamPool]. It registers the session, watches for Stop /
// GOAWAY, runs the per-session protocol loops, and sweeps the registries on
// exit. It blocks until the session ends.
//
// Both directions share this body so a dialled upstream relay session behaves
// identically to an accepted one: it lands in r.sessions (covered by Stop's
// GOAWAY/drain) and its handler fans out inbound data + routes FETCH responses.
// The only difference is the SETUP role (Server vs Client), handled by the
// caller before calling this — and leg, which records that difference for
// [Metrics] so cross-relay traffic is separable from client traffic.
func (r *Relay) serveSession(ctx context.Context, sess *session.Session, leg Leg) {
        r.addSession(sess, leg)
        defer func() {
                r.removeSession(sess, leg)
                // Belt-and-suspenders: per-request handlers unregister themselves on
                // clean shutdown, but a handler that raced past Stop or wedged could
                // leave dangling refs. Sweep both registries so the post-condition
                // "serveSession returned ⇒ no registry entry references sess" holds.
                r.tracks.RemoveSession(sess)
                r.names.RemoveSession(sess)
        }()

        // Shutdown drain is owned elsewhere: Stop runs the GOAWAY / grace / close
        // lifecycle for every session in the snapshot it takes under sessionsMu,
        // and addSession runs it for any straggler that registered after that
        // snapshot (see addSession). serveSession itself does not watch for Stop.

        handler := newSessionHandler(
                sess, r.log, r.tracks, r.names,
                r.cfg.Authorizer, r.cfg.Metrics, leg, r.fetch, r.statusRounds, r.upstreams,
                r.cfg.Discovery, r.cfg.RelayAddr,
                r.cfg.SendQueueSize, r.cfg.MaxDropsBeforeReset, r.cfg.MaxFanoutLag,
                r.cfg.MaxSubscriptionsPerSession, r.cfg.MaxNamespaceRequestsPerSession,
                max(r.cfg.MaxRendezvousTimeout, 0),
                r.handlers.Go,
        )
        if err := handler.run(ctx); err != nil {
                r.log.LogAttrs(ctx, slog.LevelDebug, "relay session handler returned",
                        slog.String("err", err.Error()))
        }
}

// Stop initiates graceful shutdown.
// Stop is idempotent. Concurrent calls share the same shutdown sequence: only
// the first call performs work, and later calls block until it has finished, so
// a caller can use a second Stop to join a drain another goroutine started.
func (r *Relay) Stop(ctx context.Context) error {
        var firstErr error
        r.stopOnce.Do(func() {
                r.log.LogAttrs(ctx, slog.LevelInfo, "relay stopping")
                close(r.stopCh)

                // 1. Withdraw from Discovery first, before anything else: a peer that
                //    resolves this relay via FindTrack / FindNamespace after step 2 has
                //    closed the listener would dial a dead endpoint. Doing it ahead of
                //    the listener close leaves only the harmless inverse window
                //    (unadvertised but still accepting). Withdraw leaves the store
                //    usable, so the rest of the drain can still resolve *other* relays.
                //    Bounded by discoveryWithdrawTimeout as well as ctx: Stop may be
                //    called with a deadline-free context, and an unreachable backend
                //    must not hold the listener close and every GOAWAY behind it.
                if r.cfg.Discovery != nil {
                        wctx, cancelWithdraw := context.WithTimeout(ctx, discoveryWithdrawTimeout)
                        err := r.cfg.Discovery.Withdraw(wctx, r.cfg.RelayAddr)
                        cancelWithdraw()
                        if err != nil {
                                // Not promoted to firstErr: the advertisements also expire on
                                // their own once the backend's liveness TTL lapses, so a failed
                                // withdrawal delays peer convergence but does not fail shutdown.
                                r.log.LogAttrs(ctx, slog.LevelWarn, "relay discovery withdrawal failed",
                                        slog.String("err", err.Error()))
                        } else {
                                r.log.LogAttrs(ctx, slog.LevelInfo, "relay withdrawn from discovery")
                        }
                }

                // 2. Close the listener; this unblocks the Accept loop. Block new
                //    upstream dials too — starting a cross-relay subscription while
                //    draining is pointless — but leave the established upstream sessions
                //    running: dropping them is "unsubscribing from upstream publishers",
                //    which §3.6 puts after the downstream GOAWAY (step 6). They are in
                //    the snapshot below, so they get GOAWAY'd / force-closed like
                //    accepted ones.
                if err := r.listener.Close(); err != nil && !isShutdownErr(err) {
                        firstErr = fmt.Errorf("relay: listener close: %w", err)
                }
                if r.upstreams != nil {
                        r.upstreams.stopDialing()
                }

                // 3. Mark shutdown in progress and snapshot the session set in one
                //    atomic step, so we can iterate without holding the lock while
                //    doing potentially-blocking session work. The atomicity partitions
                //    sessions cleanly: every session is either in this snapshot (its
                //    drain is owned by steps 4–7 below) or registered later (it observes
                //    stopCtx in addSession and owns its own drain) — never both.
                sessions := r.beginShutdown(ctx)

                // 4. Send GOAWAY to each session if a grace period is set. A
                //    zero timeout means "don't bother with GOAWAY"; close
                //    everything immediately. A relay-to-relay deployment may
                //    want to include a New Session URI here — extend
                //    SessionOptions or Config when that arrives.
                goawaySent := make(map[*session.Session]bool, len(sessions))
                if r.cfg.GoawayTimeout > 0 {
                        for _, sess := range sessions {
                                if err := sess.SendGoaway(r.cfg.GoawayTimeout, ""); err != nil {
                                        // A session that has already sent GOAWAY
                                        // or is closed is fine to skip.
                                        r.log.LogAttrs(ctx, slog.LevelDebug, "relay GOAWAY send skipped",
                                                slog.String("err", err.Error()))
                                        continue
                                }
                                goawaySent[sess] = true
                        }
                }

                // 5. Wait up to GoawayTimeout for sessions to drain. Whichever
                //    finishes first — drain or timeout — wins.
                drained := make(chan struct{})
                go func() {
                        for _, sess := range sessions {
                                <-sess.Done()
                        }
                        close(drained)
                }()

                timedOut := false
                select {
                case <-drained:
                case <-time.After(r.cfg.GoawayTimeout):
                        timedOut = true
                        r.log.LogAttrs(ctx, slog.LevelWarn, "relay GOAWAY drain timed out, force-closing sessions")
                case <-ctx.Done():
                        r.log.LogAttrs(ctx, slog.LevelWarn, "relay Stop ctx cancelled, force-closing sessions")
                }

                // 6. Now that every downstream subscriber has had its GOAWAY and the
                //    grace period is over, unsubscribe from upstream publishers by
                //    cancelling the pool's base context. §3.6: "When the server is a
                //    subscriber, it SHOULD send a GOAWAY message to downstream
                //    subscribers prior to unsubscribing from upstream publishers."
                if r.upstreams != nil {
                        r.upstreams.close()
                }
                // Forwarded TRACK_STATUS rounds are upstream work too: cut short
                // here, joined by handlers.Wait below.
                r.statusRounds.end()

                // 7. Force-close anything still standing, with GOAWAY_TIMEOUT
                //    only where it is true: "the peer took too long to close the
                //    session in response to a GOAWAY" (§3.5). A session sent no
                //    GOAWAY, or cut short by ctx, is closed with NO_ERROR. Closing
                //    an already-closed session is a no-op via Session's internal
                //    closeOnce.
                for _, sess := range sessions {
                        _ = sess.Close(shutdownCloseCode(goawaySent[sess] && timedOut), "relay shutdown")
                }

                // 8. Wait for all handler goroutines to exit. This is
                //    important: returning while handlers are still running
                //    would race with anything the caller does next (e.g.
                //    closing a test's fake transport).
                r.handlers.Wait()

                // 9. Join the Discovery namespace watcher (if started). acceptCtx
                //    was cancelled via stopCh above, so it is already unwinding.
                r.watchWG.Wait()
                r.log.LogAttrs(ctx, slog.LevelInfo, "relay stopped")
        })
        return firstErr
}

func (r *Relay) addSession(s *session.Session, leg Leg) {
        r.sessionsMu.Lock()
        r.sessions[s] = struct{}{}
        stopCtx := r.stopCtx
        r.sessionsMu.Unlock()
        r.cfg.Metrics.SessionOpened(leg)

        // Straggler cover: if shutdown was already in progress when we registered,
        // Stop's snapshot — taken under sessionsMu together with stopCtx (see
        // beginShutdown) — does NOT include this session, so Stop will neither
        // GOAWAY nor close it. Own that lifecycle here. When shutdown began after
        // we registered, stopCtx is nil and Stop's snapshot covers us;
        // exactly one owner either way. The drain runs under r.handlers so Stop's
        // handlers.Wait joins it (safe: this runs inside serveSession, itself a
        // tracked handler, so the WaitGroup counter is already non-zero).
        if stopCtx != nil {
                r.handlers.Go(func() { r.drainStraggler(stopCtx, s) })
        }
}

// drainStraggler runs the GOAWAY grace + force-close lifecycle for a single
// session that registered after Stop snapshotted the live-session set, so
// Stop's bulk drain (Stop steps 4–7) does not cover it. It mirrors that bulk
// drain for one session: GOAWAY, wait for the peer to drain or the grace period
// to elapse, or Stop's ctx to end, then force-close. Spawned by addSession only
// during shutdown.
func (r *Relay) drainStraggler(stopCtx context.Context, s *session.Session) {
        goawayExpired := false
        if r.cfg.GoawayTimeout > 0 {
                sent := s.SendGoaway(r.cfg.GoawayTimeout, "") == nil
                timer := time.NewTimer(r.cfg.GoawayTimeout)
                defer timer.Stop()
                select {
                case <-timer.C:
                        goawayExpired = sent
                case <-s.Done():
                        return // peer drained within the grace period
                case <-stopCtx.Done():
                        // Cut short, as Stop's bulk drain is: the grace period did not
                        // run out, so NO_ERROR (§3.5).
                }
        }
        _ = s.Close(shutdownCloseCode(goawayExpired), "relay shutdown")
}

// shutdownCloseCode is the code a shutdown force-closes a session with:
// GOAWAY_TIMEOUT when the relay sent it a GOAWAY and the grace period ran out
// (§3.5), NO_ERROR otherwise.
func shutdownCloseCode(goawayExpired bool) moqt.SessionErrorCode {
        if goawayExpired {
                return moqt.SessionGoawayTimeout
        }
        return moqt.SessionNoError
}

func (r *Relay) removeSession(s *session.Session, leg Leg) {
        r.sessionsMu.Lock()
        delete(r.sessions, s)
        r.sessionsMu.Unlock()
        r.cfg.Metrics.SessionClosed(leg)
}

// beginShutdown marks the relay as shutting down and returns a snapshot of the
// currently-registered sessions, atomically under sessionsMu. The atomicity is
// what lets addSession partition sessions into exactly two non-overlapping
// groups: those in the returned snapshot (drained by Stop) and those registered
// afterward (which see stopCtx and drain themselves via drainStraggler, bounded
// by ctx, Stop's).
func (r *Relay) beginShutdown(ctx context.Context) []*session.Session {
        r.sessionsMu.Lock()
        defer r.sessionsMu.Unlock()
        r.stopCtx = ctx
        out := make([]*session.Session, 0, len(r.sessions))
        for s := range r.sessions {
                out = append(out, s)
        }
        return out
}

// isShutdownErr reports whether err is one of the "the world is going away"
// signals that should be treated as a clean shutdown rather than a failure:
// net.ErrClosed, context.Canceled, or context.DeadlineExceeded (transports
// surface one of these when the conn/listener is closed under a loop).
func isShutdownErr(err error) bool {
        if err == nil {
                return false
        }
        if errors.Is(err, net.ErrClosed) {
                return true
        }
        if errors.Is(err, context.Canceled) {
                return true
        }
        if errors.Is(err, context.DeadlineExceeded) {
                return true
        }
        return false
}

package relay

import (
        "context"
        "errors"
        "log/slog"
        "time"

        "github.com/floatdrop/moq-go/pkg/relay/discovery"
)

// runNamespaceWatch consumes [discovery.DiscoveryStore.WatchNamespaces] and
// records namespaces advertised by *other* relays in the namespace registry,
// which announces them to local SUBSCRIBE_NAMESPACE holders. It runs as one
// goroutine started in [Relay.Start] when Discovery is configured, until ctx
// is cancelled or the store is closed.
//
// A watch that fails to start is retried with backoff; one whose channel
// closes (the store's signal that this consumer fell behind) is restarted,
// and its snapshot reconciled via [registry.NamespaceRegistry.ReplaceRemote].
func (r *Relay) runNamespaceWatch(ctx context.Context) {
        backoff := namespaceWatchBackoffInitial
        for first := true; ; first = false {
                if ctx.Err() != nil {
                        return // shutting down: the watch closing is not a failure
                }
                ch, err := r.cfg.Discovery.WatchNamespaces(ctx)
                if errors.Is(err, discovery.ErrClosed) {
                        return // "After Close all methods return ErrClosed"
                }
                if err != nil {
                        r.log.LogAttrs(ctx, slog.LevelWarn, "discovery: WatchNamespaces failed",
                                slog.String("err", err.Error()), slog.Duration("retry_in", backoff))
                        select {
                        case <-ctx.Done():
                                return
                        case <-time.After(backoff):
                        }
                        backoff = min(2*backoff, namespaceWatchBackoffCap)
                        continue
                }
                backoff = namespaceWatchBackoffInitial
                r.log.LogAttrs(ctx, slog.LevelDebug, "discovery namespace watch started", slog.Bool("restart", !first))
                if !r.consumeNamespaceWatch(ctx, ch) {
                        return
                }
        }
}

// consumeNamespaceWatch applies events from ch until it closes, reporting
// true, or ctx is cancelled, reporting false. The snapshot is applied as a
// whole at OpSnapshotDone, so a watch that ends before then applies nothing.
func (r *Relay) consumeNamespaceWatch(ctx context.Context, ch <-chan discovery.NamespaceEvent) bool {
        var (
                snapshot []discovery.NamespaceInfo
                synced   bool
        )
        for {
                select {
                case <-ctx.Done():
                        return false
                case ev, ok := <-ch:
                        if !ok {
                                return true
                        }
                        switch {
                        case synced:
                                r.forwardNamespaceEvent(ctx, ev)
                        case ev.Op == discovery.OpSnapshotDone:
                                r.names.ReplaceRemote(snapshot)
                                snapshot, synced = nil, true
                        case ev.Op == discovery.OpPublish && ev.Info.RelayAddr != r.cfg.RelayAddr:
                                // Own-relay advertisements are counted locally.
                                snapshot = append(snapshot, ev.Info)
                        }
                }
        }
}

// The watch restarts after these delays, doubling, when it fails to start.
const (
        namespaceWatchBackoffInitial = 100 * time.Millisecond
        namespaceWatchBackoffCap     = 10 * time.Second
)

// forwardNamespaceEvent records one remote namespace event in the namespace
// registry, counted together with local publishers of the same namespace
// (§10.18: NAMESPACE_DONE is per namespace). Own-relay events are skipped.
func (r *Relay) forwardNamespaceEvent(_ context.Context, ev discovery.NamespaceEvent) {
        if ev.Info.RelayAddr == r.cfg.RelayAddr {
                return // our own advertisement — already counted locally
        }
        switch ev.Op {
        case discovery.OpPublish:
                r.names.RemoteNamespace(ev.Info.Prefix, ev.Info.RelayAddr, true)
        case discovery.OpUnpublish:
                r.names.RemoteNamespace(ev.Info.Prefix, ev.Info.RelayAddr, false)
        case discovery.OpSnapshotDone:
                // Handled by consumeNamespaceWatch.
        }
}

package relay

import (
        "cmp"
        "context"
        "errors"
        "fmt"
        "hash/fnv"
        "log/slog"
        "slices"
        "sync"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/wire"
        "github.com/floatdrop/moq-go/pkg/relay/discovery"
)

// upstreamDialTimeout bounds a single relay-to-relay dial + MOQT SETUP. A hung
// dial must not pin a pool entry forever (other callers wait on it), so the
// dial context is cancelled after this regardless of the pool's lifetime.
const upstreamDialTimeout = 10 * time.Second

// upstreamPool dials and reuses relay-to-relay sessions, keyed by the RelayAddr
// a peer advertised in [discovery.DiscoveryStore]. It is the consume-side
// counterpart of the advertise-side Discovery wiring in the registries: when a
// downstream SUBSCRIBE has no local publisher, the SUBSCRIBE handler asks the
// pool to resolve a remote relay (via FindNamespace) and hand back a live
// session to issue an upstream SUBSCRIBE on.
//
// One session is kept per RelayAddr. Concurrent resolves for the same address
// dial once and share the result (the in-flight entry is published before the
// dial, so later callers block on its ready channel rather than racing a
// second dial). A session is evicted when its per-session handler loop returns,
// so the next resolve re-dials.
type upstreamPool struct {
        dialer       func(ctx context.Context, relayAddr string) (session.Conn, error)
        discovery    discovery.DiscoveryStore
        relayAddr    string
        sessionOpts  []session.Option
        log          *slog.Logger
        metrics      Metrics
        baseCtx      context.Context
        cancelBase   context.CancelFunc
        serveSession func(sess *session.Session, onClose func())

        // fanIn optionally caps how many rendezvous-ranked upstreams
        // resolveUpstreams subscribes to per namespace (Config.UpstreamFanIn).
        // Zero (or negative) means unbounded — fan in to every advertiser, the
        // §9.5 default.
        fanIn int

        mu      sync.Mutex
        entries map[string]*poolEntry
        // noDial is set by stopDialing once shutdown begins: existing upstream
        // sessions keep running, but no new one is established.
        noDial bool
}

// poolEntry is the per-RelayAddr slot. ready is closed once sess/err are set,
// so concurrent callers that found an in-flight entry block on it instead of
// dialing again.
type poolEntry struct {
        ready chan struct{}
        sess  *session.Session
        err   error
}

// upstreamPoolConfig carries the pool's dependencies from [New].
type upstreamPoolConfig struct {
        dialer       func(ctx context.Context, relayAddr string) (session.Conn, error)
        discovery    discovery.DiscoveryStore
        relayAddr    string
        sessionOpts  []session.Option
        log          *slog.Logger
        metrics      Metrics
        serveSession func(sess *session.Session, onClose func())
        // fanIn is Config.UpstreamFanIn verbatim; zero (the default) means
        // unbounded fan-in.
        fanIn int
}

func newUpstreamPool(cfg upstreamPoolConfig) *upstreamPool {
        // The pool's base context spans its whole lifetime: dialled sessions and
        // their handler loops run under it, and close() cancels it from Relay.Stop.
        base, cancel := context.WithCancel(context.Background())
        // [New] already defaults Config.Metrics, but the pool is also constructed
        // directly (tests), and resolveUpstreams calls into this on a path that
        // only runs once a cross-relay subscribe happens — a nil here would be a
        // panic nothing local reproduces.
        if cfg.metrics == nil {
                cfg.metrics = NopMetrics{}
        }
        return &upstreamPool{
                dialer:       cfg.dialer,
                discovery:    cfg.discovery,
                relayAddr:    cfg.relayAddr,
                sessionOpts:  cfg.sessionOpts,
                log:          cfg.log,
                metrics:      cfg.metrics,
                baseCtx:      base,
                cancelBase:   cancel,
                serveSession: cfg.serveSession,
                fanIn:        cfg.fanIn,
                entries:      make(map[string]*poolEntry),
        }
}

// resolveUpstreams finds the remote relays that serve ns and returns a live
// session to each, ranked by rendezvous (HRW) weight. The ranking is a
// deterministic function of (ns, candidate addresses) alone, so every relay
// sharing the same Discovery view computes the same order. On its own that only
// fixes the dial order; with a positive fanIn (see below) it also bounds how
// many upstreams are taken, so relays converge on the same small set and the
// relay-to-relay stream count stays bounded instead of trending toward a full
// O(n²) mesh — a tree rooted at ns's highest-weighted relays.
//
// §9.5 requires subscribing to every matching publisher (the fanout then dedups
// the redundant copies they push), so fanIn == 0 (the default) fans into all of
// them. A positive fanIn (Config.UpstreamFanIn) is an opt-in deviation: it
// bounds the subscription to the top fanIn ranked upstreams — 1 is a pure tree,
// 2 keeps one backup — trading the §9.5 fan-in for a bounded relay mesh. That is
// only sound where the advertisers are redundant sources of the same objects,
// never where different relays hold distinct objects for the track. Candidates
// are dialled in rank order and, when bounded, the first fanIn that connect are
// returned; a dead top-ranked relay still lingering in Discovery during its
// lease TTL falls through transparently to the next-ranked one.
//
// Returns nil when Discovery knows no usable remote (none advertised, only this
// relay itself, or every candidate failed to dial). draining reports whether a
// candidate was skipped for having sent GOAWAY (§10.4), which the caller
// answers as it would a draining publisher. Discovery-lookup and
// per-peer dial failures are logged and treated as "skip that candidate" —
// consistent with the best-effort advertise side: the local registry / a clean
// SUBSCRIBE rejection is the fallback, never a torn-down session.
//
// Loop prevention is minimal: candidates whose RelayAddr equals this relay's
// own (or is empty / unaddressable) are skipped so the relay never subscribes
// to itself. Duplicate RelayAddrs collapse to one session (the pool keys by
// address). Multi-hop cycle detection (A→B→C→A) is out of scope — see the
// package limitations.
func (p *upstreamPool) resolveUpstreams(
        ctx context.Context,
        ns wire.TrackNamespace,
) (out []*session.Session, draining bool) {
        if p == nil || p.discovery == nil {
                return nil, false
        }
        infos, err := p.discovery.FindNamespace(ctx, ns)
        if err != nil {
                // A Discovery lookup failure is transient (etcd RPC timeout, leader
                // election) and collapses to the same nil return as a genuinely empty
                // result — a caller cannot tell "the fabric hiccupped" from "no relay
                // serves this namespace". Log it at Warn so that distinction survives to
                // production, where the default level hides Debug.
                p.log.LogAttrs(ctx, slog.LevelWarn, "upstream pool: FindNamespace failed",
                        slog.String("namespace", fmt.Sprintf("%v", ns)),
                        slog.String("err", err.Error()))
                return nil, false
        }
        p.log.LogAttrs(ctx, slog.LevelInfo, "upstream pool: FindNamespace resolved",
                slog.String("namespace", fmt.Sprintf("%v", ns)),
                slog.Int("advertisers_found", len(infos)))
        p.metrics.NamespaceResolved(len(infos))
        // Rank so the dial order is identical fleet-wide; a positive fanIn then
        // takes the same top-fanIn upstreams everywhere.
        rankByAffinity(ns, infos)

        seen := make(map[string]bool, len(infos))
        for _, info := range infos {
                if info.RelayAddr == "" || info.RelayAddr == p.relayAddr || seen[info.RelayAddr] {
                        continue // self / unaddressable / already dialled this address
                }
                seen[info.RelayAddr] = true
                sess, err := p.get(info.RelayAddr)
                if err != nil {
                        p.log.LogAttrs(ctx, slog.LevelDebug, "upstream pool dial failed",
                                slog.String("relay_addr", info.RelayAddr),
                                slog.String("err", err.Error()))
                        p.metrics.UpstreamDialFailed(info.RelayAddr)
                        continue // fall through to the next-ranked relay
                }
                if goingAway(sess) {
                        // §10.4: a draining relay takes no new requests, so it must not
                        // hold a fan-in slot.
                        draining = true
                        continue
                }
                out = append(out, sess)
                if p.fanIn > 0 && len(out) >= p.fanIn {
                        break // opt-in bound reached; deeper candidates are the fallback pool
                }
        }
        return out, draining
}

// rankByAffinity sorts infos in place by descending rendezvous (HRW) weight for
// ns. The weight hashes (ns, RelayAddr), so the order depends only on the
// namespace and the candidate set — every relay in the fleet derives the same
// order and, taking the top few, converges on the same upstreams. RelayAddr
// breaks weight ties, keeping the order total (and identical everywhere) even
// on the rare hash collision.
func rankByAffinity(ns wire.TrackNamespace, infos []discovery.NamespaceInfo) {
        nsKey := namespaceAffinityKey(ns)
        slices.SortFunc(infos, func(a, b discovery.NamespaceInfo) int {
                if c := cmp.Compare(hrwWeight(nsKey, b.RelayAddr), hrwWeight(nsKey, a.RelayAddr)); c != 0 {
                        return c
                }
                return cmp.Compare(a.RelayAddr, b.RelayAddr)
        })
}

// hrwWeight is the highest-random-weight score for placing ns on the relay at
// addr: FNV-1a over the namespace's canonical bytes followed by the address.
// nsKey is canonical and self-delimiting (a field count then length-prefixed
// fields), so the concatenation is injective in (ns, addr) and needs no
// separator.
func hrwWeight(nsKey []byte, addr string) uint64 {
        // hash.Hash.Write never returns an error (documented on the interface); the
        // blank assignments are just to satisfy the errcheck/gosec linters.
        h := fnv.New64a()
        _, _ = h.Write(nsKey)
        _, _ = h.Write([]byte(addr))
        return h.Sum64()
}

// namespaceAffinityKey is the canonical §2.4.1 wire encoding of ns, used as the
// stable per-namespace seed for hrwWeight. Reusing the wire encoding (the same
// bytes a DiscoveryStore backend keys namespaces by) keeps nested tuples
// unambiguous.
func namespaceAffinityKey(ns wire.TrackNamespace) []byte {
        w := wire.NewWriter(nil)
        w.TrackNamespace(ns)
        return w.Bytes()
}

// get returns a pooled session for relayAddr, dialing one if none is live.
// Concurrent calls for the same address dial once and share the outcome.
func (p *upstreamPool) get(relayAddr string) (*session.Session, error) {
        p.mu.Lock()
        if e, ok := p.entries[relayAddr]; ok {
                p.mu.Unlock()
                <-e.ready
                if e.err != nil {
                        return nil, e.err
                }
                // Reuse only if the session is still live; otherwise drop the stale
                // entry and dial afresh. The eviction goroutine also clears dead
                // entries, but a caller can race ahead of it.
                select {
                case <-e.sess.Done():
                        p.mu.Lock()
                        if p.entries[relayAddr] == e {
                                delete(p.entries, relayAddr)
                        }
                        p.mu.Unlock()
                        return p.get(relayAddr)
                default:
                        return e.sess, nil
                }
        }
        if p.noDial {
                p.mu.Unlock()
                // Callers treat this like any other failed upstream resolution and skip
                // the candidate.
                return nil, errors.New("relay: upstream dialing stopped for shutdown")
        }
        // Publish an in-flight entry before dialing so concurrent callers wait on
        // it rather than starting a second dial.
        e := &poolEntry{ready: make(chan struct{})}
        p.entries[relayAddr] = e
        p.mu.Unlock()

        sess, err := p.dial(relayAddr)
        e.sess, e.err = sess, err
        close(e.ready)
        if err != nil {
                // Failed dial: drop the entry so a later resolve retries.
                p.mu.Lock()
                if p.entries[relayAddr] == e {
                        delete(p.entries, relayAddr)
                }
                p.mu.Unlock()
                return nil, err
        }

        // Run the relay's per-session loops on the dialled session and evict the
        // entry when that handler returns (session ended).
        p.serveSession(sess, func() {
                p.mu.Lock()
                if p.entries[relayAddr] == e {
                        delete(p.entries, relayAddr)
                }
                p.mu.Unlock()
        })
        return sess, nil
}

// dial performs one outbound dial + client-side MOQT SETUP, bounded by
// upstreamDialTimeout. It uses the pool's base context (not a caller's request
// context) so the resulting session outlives the SUBSCRIBE that triggered it.
func (p *upstreamPool) dial(relayAddr string) (*session.Session, error) {
        dialCtx, cancel := context.WithTimeout(p.baseCtx, upstreamDialTimeout)
        defer cancel()
        conn, err := p.dialer(dialCtx, relayAddr)
        if err != nil {
                return nil, err
        }
        sess, err := session.Client(dialCtx, conn, p.sessionOpts...)
        if err != nil {
                return nil, err
        }
        return sess, nil
}

// stopDialing blocks any further upstream dial, leaving sessions already
// established running. [Relay.Stop] calls it as shutdown begins: starting a new
// cross-relay subscription while draining is pointless, but tearing the live ones
// down is "unsubscribing from upstream publishers", which §3.6 says SHOULD happen
// only after the downstream GOAWAY has gone out — so that half is [upstreamPool.close].
func (p *upstreamPool) stopDialing() {
        p.mu.Lock()
        p.noDial = true
        p.mu.Unlock()
}

// close cancels the pool's base context, unwinding any in-flight dial and the
// handler loops of dialled sessions. This is the "unsubscribe from upstream"
// step, so [Relay.Stop] calls it only after the GOAWAY broadcast and drain
// (§3.6).
func (p *upstreamPool) close() {
        p.cancelBase()
}

package relaynet

import (
        "context"
        "crypto/tls"
        "errors"
        "fmt"
        "log/slog"
        "net"
        "net/http"
        "slices"
        "sync"

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

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/session/quicconn"
        "github.com/floatdrop/moq-go/pkg/moqt/session/wtconn"
)

// DualALPNs lists the ALPNs of both MOQT transport mappings, for a listener that
// serves them on one socket — see [Listen]. A TLS config built with these accepts
// a raw-QUIC client offering "moqt-NN" and an HTTP/3 client offering "h3"; each
// connection's negotiated ALPN then says which mapping it is.
var DualALPNs = slices.Concat(MOQTQUICALPNs, WebTransportALPNs)

// dualBacklog bounds the queue of accepted-but-not-yet-Accepted connections. Both
// halves feed it, and the relay drains it promptly; the bound only matters for a
// burst arriving faster than the accept loop consumes.
const dualBacklog = 16

// Listen serves both MOQT transport mappings on a single UDP socket: raw QUIC for
// peers and native clients that dial a moqt URI, and WebTransport (HTTP/3) at
// wtPath for anything dialing the https form of the same URI (§3.1.3, §3.1.4) —
// browsers included. tlsCfg must advertise [DualALPNs].
//
// This is what a relay behind a load balancer wants, and it is why no transport
// flag is needed: the two mappings differ only in ALPN, so one listener can offer
// both and decide per connection. Clients choose by URL scheme, peer relays keep
// dialing raw QUIC, and nothing has to agree deployment-wide.
//
// The returned listener owns the socket; Close releases it along with both halves.
// A connection whose ALPN is not "h3" is treated as raw QUIC: the ALPN set the
// handshake selected from is tlsCfg's, so nothing else can get that far.
//
// CheckOrigin accepts every origin, as [ListenWebTransport] does — see the
// package doc. Serving both mappings means a relay is reachable from a browser by
// default, so a deployment that cares about which pages may open sessions needs
// its own policy here.
//
// opts tune the QUIC config this listener serves on, independently of whatever a
// cross-relay [DialQUIC] uses — see [WithQUICConfig].
func Listen(addr, wtPath string, tlsCfg *tls.Config, logger *slog.Logger, opts ...Option) (*DualListener, error) {
        if logger == nil {
                logger = slog.Default()
        }
        if wtPath == "" {
                // ServeMux panics on an empty pattern; fail at startup with a message
                // instead of taking the process down inside NewListener.
                return nil, fmt.Errorf("relaynet: empty WebTransport path (use %q for the default)", "/moq")
        }
        qcfg := quicConfig(opts)
        // Neither of these is the caller's to switch off here, whatever WithQUICConfig
        // asked for: webtransport.Server.ServeQUICConn checks them one at a time and
        // refuses the connection on the first one missing, so a listener lacking
        // either would silently serve only half of what it advertises. MOQT can use
        // both (§11.3, §11.4.3), but neither is what forces the hand here.
        qcfg.EnableDatagrams = true
        qcfg.EnableStreamResetPartialDelivery = true

        // Not ListenEarly: an early listener yields connections before the handshake
        // completes, which would break this listener's contract that ALPN is already
        // negotiated (the dispatch below reads it) and would have the relay write
        // SETUP as 0.5-RTT data to a peer whose certificate is unverified. 0-RTT would
        // need Config.Allow0RTT, which defaultQUICConfig deliberately leaves unset.
        ql, err := quic.ListenAddr(addr, tlsCfg, qcfg)
        if err != nil {
                return nil, fmt.Errorf("relaynet: listen %s: %w", addr, err)
        }

        mux := http.NewServeMux()
        if wtPath != "/" {
                mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
                        logger.WarnContext(r.Context(), "webtransport: unmatched request",
                                "method", r.Method, "host", r.Host, "path", r.URL.Path,
                                "proto", r.Proto, "upgrade", r.Header.Get(":protocol"))
                        http.NotFound(w, r)
                })
        }
        // This half is what the two re-asserted fields above are for; see there.
        h3 := &http3.Server{TLSConfig: tlsCfg, Handler: mux}
        webtransport.ConfigureHTTP3Server(h3)
        wts := &webtransport.Server{
                H3: h3,
                // The WebTransport sub-protocol, not the TLS ALPN, carries the draft
                // version for this mapping (§3.1) — same identifiers as raw QUIC so the
                // two version signals share one source.
                ApplicationProtocols: MOQTQUICALPNs,
                CheckOrigin:          func(*http.Request) bool { return true },
        }

        ctx, cancel := context.WithCancel(context.Background())
        l := &DualListener{
                ql:     ql,
                wts:    wts,
                conns:  make(chan session.Conn, dualBacklog),
                ctx:    ctx,
                cancel: cancel,
                log:    logger,
        }
        l.wt = wtconn.NewListener(wts, mux, wtPath, ql.Addr(), dualBacklog)

        go l.acceptQUIC()
        go l.pumpWebTransport()
        return l, nil
}

// DualListener is the [Listen] listener: one QUIC socket whose connections are
// split by negotiated ALPN into the raw-QUIC and WebTransport halves, then merged
// into one Accept queue so the relay cannot tell them apart.
type DualListener struct {
        ql  *quic.Listener
        wt  *wtconn.Listener
        wts *webtransport.Server

        conns  chan session.Conn
        ctx    context.Context
        cancel context.CancelFunc
        log    *slog.Logger

        closeOnce sync.Once
}

// acceptQUIC is the demultiplexer: HTTP/3 connections go to the WebTransport
// server, whose upgrade handler feeds the wtconn listener that pumpWebTransport
// drains; everything else is a raw-QUIC MOQT connection and is queued directly.
func (l *DualListener) acceptQUIC() {
        for {
                conn, err := l.ql.Accept(l.ctx)
                if err != nil {
                        return // listener closed
                }
                if conn.ConnectionState().TLS.NegotiatedProtocol == http3.NextProtoH3 {
                        go func() {
                                // Returns when the HTTP/3 connection ends, which is routine.
                                if err := l.wts.ServeQUICConn(conn); err != nil && l.ctx.Err() == nil {
                                        l.log.Debug("relaynet: http/3 connection ended", "err", err.Error())
                                }
                        }()
                        continue
                }
                // The mapping is not recorded on the conn, but it stays recoverable by
                // type — quicconn and wtconn produce distinct implementations — which is
                // what a future §10.3.1.1/§10.3.1.2 check would need (PATH and AUTHORITY
                // MUST NOT be used over WebTransport).
                l.deliver(quicconn.New(conn))
        }
}

// pumpWebTransport moves upgraded WebTransport sessions onto the shared queue, so
// Accept has a single source regardless of transport.
func (l *DualListener) pumpWebTransport() {
        for {
                conn, err := l.wt.Accept(l.ctx)
                if err != nil {
                        return // Close, or the wtconn listener shut down
                }
                l.deliver(conn)
        }
}

// deliver queues conn, or closes it if the listener is shutting down — a conn
// nobody will Accept must not be left believing it has a session.
func (l *DualListener) deliver(conn session.Conn) {
        select {
        case l.conns <- conn:
        case <-l.ctx.Done():
                _ = conn.CloseWithError(uint64(moqt.SessionNoError), "listener closed")
        }
}

// Accept returns the next connection from either transport. It satisfies the
// relay's Listener interface.
func (l *DualListener) Accept(ctx context.Context) (session.Conn, error) {
        select {
        case conn := <-l.conns:
                return conn, nil
        case <-ctx.Done():
                return nil, ctx.Err()
        case <-l.ctx.Done():
                return nil, net.ErrClosed
        }
}

// Addr returns the UDP address both transports are served on.
func (l *DualListener) Addr() net.Addr { return l.ql.Addr() }

// Close stops accepting new connections. Connections already accepted keep
// working, and the UDP socket stays open until the last of them ends — quic-go
// releases it once its transport has no connections left.
//
// That is load-bearing, not incidental: [relay.Relay.Stop] closes the listener as
// an early step and only then broadcasts GOAWAY and waits out the grace period
// (§10.4, §3.6). A Close that dropped the socket would kill every draining
// session with it, and no peer would ever see its GOAWAY.
//
// Close is idempotent and joins the failures of every step.
func (l *DualListener) Close() error {
        var err error
        l.closeOnce.Do(func() {
                l.cancel()
                err = errors.Join(l.wt.Close(), l.wts.Close(), l.ql.Close())
        })
        return err
}

package relaynet

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

// Option customises the QUIC plumbing that [Listen], [DialQUIC] and
// [DialWebTransport] build. Every entry point takes them variadically, so a
// caller can tune each leg of a relay independently — the downstream listener
// and a cross-relay upstream dial do not have to agree.
type Option func(*quic.Config)

// WithQUICConfig returns an [Option] that runs tune over the [quic.Config] the
// entry point is about to use, once the relay's defaults have been populated.
// Passing it more than once applies the hooks in order.
//
// The hook mutates those defaults rather than replacing them, so a caller
// changing one knob keeps tracking every other default this package sets, and
// gains any it adds later. That is what makes it usable for the knob it exists
// for: quic.Config fields that only a patched or forked quic-go defines — a
// pluggable congestion controller, say — can be set here without this package
// ever naming them, and without the caller having to restate the settings MOQT
// needs around them.
//
// Two of those settings are not the caller's to turn off on a [Listen]: the
// dual listener's WebTransport half refuses a connection missing either DATAGRAM
// or stream-reset partial delivery, so Listen re-asserts both after the hooks
// run. Neither dial entry point does. [DialWebTransport] does not need to —
// webtransport-go refuses the same omission up front, before any packet is sent.
// [DialQUIC] has no such guard and none is added here: the config goes straight
// to quic-go, and a caller who disables DATAGRAM on a cross-relay leg owns the
// consequence, which is that objects this relay would have sent as datagrams
// (§11.3) stop crossing that hop with only a Debug line to show for it.
func WithQUICConfig(tune func(*quic.Config)) Option { return Option(tune) }

// quicConfig applies opts to a fresh copy of the relay's default QUIC tuning.
func quicConfig(opts []Option) *quic.Config {
        cfg := defaultQUICConfig()
        for _, opt := range opts {
                opt(cfg)
        }
        return cfg
}

// Package relaynet holds the QUIC + TLS plumbing shared by the relay command
// binaries: self-signed dev certificates, the relay's QUIC tuning, and the
// listener/dialer constructors that bridge quic-go to the transport-agnostic
// [session.Conn] the relay operates on.
//
// It exists so more than one relay binary shares one copy of this setup rather
// than each carrying its own: cmd/relay here, and the out-of-tree relay built
// on a distributed [relay.DiscoveryStore], which imports this package. The helpers here are aimed at local
// development and single-operator deployments — [SelfSignedCert],
// [InsecureClientTLSConfig] and [Listen] are explicitly not for production: the
// first two skip real trust, and the third accepts every browser Origin.
package relaynet

import (
        "context"
        "crypto/ecdsa"
        "crypto/elliptic"
        "crypto/rand"
        "crypto/tls"
        "crypto/x509"
        "crypto/x509/pkix"
        "encoding/pem"
        "fmt"
        "log/slog"
        "math/big"
        "net"
        "time"

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

        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/session/quicconn"
)

// MOQTQUICALPNs lists the raw-QUIC MOQT ALPNs the relay accepts. SETUP
// carries no version field (§10.3); MOQT negotiates the version with ALPN
// (§3.1), so the "moqt-NN" ALPN is itself the draft-version signal — the
// negotiated ALPN fixes the draft. We advertise only "moqt-20", the draft this
// implementation speaks. The older "moqt-19"/"-18"/"-17"/"-16" and "moq-00"
// (used before -15, followed by version negotiation in SETUP — §3.1) are
// deliberately not offered: our -20 wire behavior can't complete a SETUP with
// a peer that selected any of them, so advertising them would only let such a
// peer clear TLS and then fail.
var MOQTQUICALPNs = []string{"moqt-20"}

// defaultQUICConfig returns the QUIC tuning the relay listens and dials with:
// a 30s idle timeout with 5s keep-alives, datagrams enabled (MOQT may deliver
// objects as QUIC datagrams), and RESET_STREAM_AT partial delivery (§11.4.3).
//
// It is the base every entry point starts from; [WithQUICConfig] is how a caller
// adjusts it per leg.
func defaultQUICConfig() *quic.Config {
        return &quic.Config{
                MaxIdleTimeout:                   30 * time.Second,
                KeepAlivePeriod:                  5 * time.Second,
                EnableDatagrams:                  true,
                EnableStreamResetPartialDelivery: true,
        }
}

// TLSConfig returns a server TLS config for the chosen MOQT transport. If
// certFile and keyFile are both non-empty the pair is loaded from disk;
// otherwise an ephemeral self-signed certificate is generated in memory (see
// [SelfSignedCert]). alpns lists the acceptable ALPNs in server-preference
// order.
func TLSConfig(certFile, keyFile string, alpns []string) (*tls.Config, error) {
        var (
                cert tls.Certificate
                err  error
        )
        if certFile != "" && keyFile != "" {
                cert, err = tls.LoadX509KeyPair(certFile, keyFile)
        } else {
                slog.Default().Info("relaynet: no cert/key supplied; generating ephemeral self-signed certificate")
                cert, err = SelfSignedCert()
        }
        if err != nil {
                return nil, fmt.Errorf("tls: %w", err)
        }
        return &tls.Config{
                Certificates: []tls.Certificate{cert},
                NextProtos:   alpns,
        }, nil
}

// InsecureClientTLSConfig returns a client TLS config that offers alpns and
// SKIPS certificate verification. It exists for relay-to-relay dialing in
// development, where peers present self-signed certs; production deployments
// MUST supply a config with a real trust store instead.
func InsecureClientTLSConfig(alpns []string) *tls.Config {
        return &tls.Config{
                InsecureSkipVerify: true, //nolint:gosec // dev-only cross-relay dialing against self-signed peers; documented on the func.
                NextProtos:         alpns,
        }
}

// DialQUIC dials addr over raw QUIC with the relay's default QUIC tuning and
// returns the established connection as a [session.Conn], ready for the relay to
// drive the client-side MOQT SETUP on. It is the shape a relay Dialer expects.
//
// [quicconn.Dial] owns the address handling, including the multi-address,
// RFC 6724-ordered resolution that keeps a dual-stack peer named by hostname
// from being dialed over the wrong family — see its doc comment.
//
// opts tune the QUIC config this leg dials with, leaving the listener and any
// other dial untouched — see [WithQUICConfig].
func DialQUIC(ctx context.Context, addr string, tlsCfg *tls.Config, opts ...Option) (session.Conn, error) {
        return quicconn.Dial(ctx, addr, tlsCfg, quicConfig(opts))
}

// SelfSignedCert generates an ephemeral ECDSA-P256 self-signed certificate for
// localhost, valid for 10 days (within Chrome's ≤14-day tolerance for
// serverCertificateHashes pinning). It is for local development only.
func SelfSignedCert() (tls.Certificate, error) {
        key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
        if err != nil {
                return tls.Certificate{}, err
        }
        template := &x509.Certificate{
                SerialNumber: big.NewInt(1),
                Subject:      pkix.Name{CommonName: "mediamesh-relay"},
                DNSNames:     []string{"localhost"},
                IPAddresses:  []net.IP{net.IPv4(127, 0, 0, 1), net.IPv6loopback},
                // serverAuth + a key-usage that permits TLS server handshakes is
                // mandatory for clients that validate the cert against a trust store.
                // Without it the QUIC/h3 handshake fails before any MOQT logic runs.
                KeyUsage:              x509.KeyUsageDigitalSignature,
                ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
                BasicConstraintsValid: true,
                NotBefore:             time.Now().Add(-time.Minute),
                NotAfter:              time.Now().Add(10 * 24 * time.Hour),
        }
        certDER, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
        if err != nil {
                return tls.Certificate{}, err
        }
        keyDER, err := x509.MarshalECPrivateKey(key)
        if err != nil {
                return tls.Certificate{}, err
        }
        certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
        // SEC1 EC keys use the "EC PRIVATE KEY" PEM type (RFC 5915).
        keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
        return tls.X509KeyPair(certPEM, keyPEM)
}

package relaynet

import (
        "context"
        "crypto/tls"
        "fmt"

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

        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/session/wtconn"
)

// WebTransportALPNs lists the TLS ALPNs of the MOQT-over-WebTransport mapping.
// WebTransport rides HTTP/3, whose ALPN is "h3" — the "moqt-NN" identifiers belong
// to raw QUIC ([MOQTQUICALPNs]). The draft version is instead negotiated as the
// WebTransport sub-protocol (§3.1), which [Listen] offers.
//
// Pass these alone to [TLSConfig] for a listener that serves *only* WebTransport;
// [DualALPNs] serves both mappings.
var WebTransportALPNs = []string{http3.NextProtoH3}

// DialWebTransport dials rawURL — the https URL of a WebTransport endpoint, i.e.
// the §3.1.4 conversion of a moqt URI — and returns the established session as a
// [session.Conn], ready for the caller to drive the client-side MOQT SETUP on.
// It is the WebTransport counterpart of [DialQUIC] and has the shape a relay
// Dialer expects; tlsCfg must advertise [WebTransportALPNs]. opts tune the QUIC
// config this leg dials with — see [WithQUICConfig].
func DialWebTransport(ctx context.Context, rawURL string, tlsCfg *tls.Config, opts ...Option) (session.Conn, error) {
        d := webtransport.Transport{
                TLSClientConfig: tlsCfg,
                // Unlike [Listen], nothing forces DATAGRAM and stream-reset partial
                // delivery back on after the opts run: webtransport-go's Dial rejects a
                // config missing either before it sends a packet, so the mistake surfaces
                // here rather than as a mapping that quietly stops working. The error is
                // a bare errors.New, not a sentinel, so it is reportable but not testable
                // with errors.Is.
                QUICConfig: quicConfig(opts),
                // §3.1.4: "The client includes MOQT protocol identifiers in the
                // WT-Available-Protocols header." That header is how a WebTransport
                // session negotiates the draft version, the way ALPN does it for raw
                // QUIC (§3.1) — without it the upgrade completes with an empty protocol
                // and the version is never agreed. webtransport-go builds the header
                // from this list and rejects a selection it did not offer.
                ApplicationProtocols: MOQTQUICALPNs,
        }
        // The extended-CONNECT response body is the stream the WebTransport session
        // rides on, so it must NOT be closed here: doing so would tear down the very
        // session being returned. Its lifetime belongs to the returned conn.
        //nolint:bodyclose // response body is the session stream; owned by wtSess.
        _, wtSess, err := d.Dial(ctx, rawURL, nil)
        if err != nil {
                return nil, fmt.Errorf("relaynet: dial webtransport %s: %w", rawURL, err)
        }
        return wtconn.New(wtSess), nil
}

package relay

import (
        "context"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "math/rand/v2"
        "slices"
        "sync"
        "sync/atomic"
        "time"

        "github.com/floatdrop/moq-go/pkg/moqt"
        "github.com/floatdrop/moq-go/pkg/moqt/message"
        "github.com/floatdrop/moq-go/pkg/moqt/session"
        "github.com/floatdrop/moq-go/pkg/moqt/track"
        "github.com/floatdrop/moq-go/pkg/relay/cache"
        "github.com/floatdrop/moq-go/pkg/relay/discovery"
        "github.com/floatdrop/moq-go/pkg/relay/internal/registry"
)

// sessionHandler owns the per-session request and data-stream loops, one per
// accepted [session.Session]. The relay's shared state (registries, authorizer)
// is injected by [Relay.handleConn] and referenced read-only.
//
// Concurrency: [sessionHandler.run] drives the request, data-stream, and
// datagram loops on separate goroutines, plus per-request handler goroutines
// spawned by the dispatch loop and tracked via wg for a clean join on teardown.
type sessionHandler struct {
        // nsPrefixes / trackPrefixes are this session's established
        // SUBSCRIBE_NAMESPACE / SUBSCRIBE_TRACKS prefixes, for §10.19 / §10.20
        // PREFIX_OVERLAP. The two types have independent overlap spaces.
        nsPrefixes, trackPrefixes prefixSet

        sess    *session.Session
        log     *slog.Logger
        tracks  *registry.TrackRegistry
        names   *registry.NamespaceRegistry
        auth    Authorizer
        metrics Metrics
        fetch   *registry.FetchRouter
        // statusRounds shares forwarded TRACK_STATUS rounds across the relay's
        // sessions; see [sessionHandler.forwardTrackStatus].
        statusRounds *trackStatusRounds
        // leg records whether this session was dialled by the relay
        // (LegUpstream) or by the peer (LegLocal). Every [Metrics] call this
        // handler makes carries it, so an operator can separate what the
        // cross-relay hop is doing from what clients are doing.
        leg                 Leg
        upstreams           *upstreamPool
        discovery           discovery.DiscoveryStore
        relayAddr           string
        sendQueueSize       int
        maxDropsBeforeReset int
        maxFanoutLag        time.Duration
        // maxRendezvous caps RENDEZVOUS_TIMEOUT (§10.2.6); 0 holds no SUBSCRIBE.
        maxRendezvous time.Duration

        // limiter enforces the §13.1 / §13.7.1 per-session resource caps.
        limiter sessionLimiter

        // wg tracks per-request goroutines spawned by the dispatch loop.
        wg sync.WaitGroup

        // earlyStreams counts the subgroup streams waiting for their Track Alias
        // to be registered; see [sessionHandler.resolveInboundTrack].
        earlyStreams atomic.Int32

        // subscribing holds, per track, this session's SUBSCRIBEs that have not
        // registered their downstream yet; see [sessionHandler.beginSubscribe].
        subscribingMu sync.Mutex
        subscribing   map[track.Key]*inflightSubscribe

        // relayGo runs fn on a RELAY-scoped goroutine (joined by Relay.Stop,
        // not by this handler's run). Used for work whose lifetime must outlive
        // this session — e.g. the reader of an on-demand upstream stream, which
        // serves every downstream subscriber of the track, not just the one on
        // this session (§9.4 aggregation).
        relayGo func(func())
}

// newSessionHandler constructs a handler. Callers provide the shared
// dependencies; the handler does not allocate them itself, which makes it
// trivial to fan-in a test handler with a fake registry or authorizer.
func newSessionHandler(
        sess *session.Session,
        log *slog.Logger,
        tracks *registry.TrackRegistry,
        names *registry.NamespaceRegistry,
        auth Authorizer,
        metrics Metrics,
        leg Leg,
        fetch *registry.FetchRouter,
        statusRounds *trackStatusRounds,
        upstreams *upstreamPool,
        discovery discovery.DiscoveryStore,
        relayAddr string,
        sendQueueSize int,
        maxDropsBeforeReset int,
        maxFanoutLag time.Duration,
        maxSubsPerSession int,
        maxNamespaceReqsPerSession int,
        maxRendezvous time.Duration,
        relayGo func(func()),
) *sessionHandler {
        return &sessionHandler{
                sess:                sess,
                log:                 log.With("moqt.session", fmt.Sprintf("%p", sess)),
                tracks:              tracks,
                names:               names,
                auth:                auth,
                metrics:             metrics,
                leg:                 leg,
                fetch:               fetch,
                statusRounds:        statusRounds,
                upstreams:           upstreams,
                discovery:           discovery,
                relayAddr:           relayAddr,
                sendQueueSize:       sendQueueSize,
                maxDropsBeforeReset: maxDropsBeforeReset,
                maxFanoutLag:        maxFanoutLag,
                maxRendezvous:       maxRendezvous,
                limiter:             sessionLimiter{maxSubs: maxSubsPerSession, maxNS: maxNamespaceReqsPerSession},
                relayGo:             relayGo,
        }
}

// trackRef labels a [Metrics] event with the track it happened on and the leg
// this session sits on.
//
// track.FullTrackName holds the name as []byte, so this conversion allocates.
// Callers MUST hoist it out of any per-object loop — build one TrackRef when a
// stream, subscription or FETCH begins and reuse it — because [Metrics]
// promises implementations a hot path with nothing to spare.
func (h *sessionHandler) trackRef(name track.FullTrackName) TrackRef {
        return TrackRef{Name: string(name.Name), Leg: h.leg}
}

// saveLargestLocation folds a LARGEST_OBJECT parameter the upstream sent into
// the track's watermark. §10.2.17: a relay advertises the largest of the
// values received in SUBSCRIBE_OK, PUBLISH or REQUEST_UPDATE_OK and the
// Objects it received.
//
// Call it on every path carrying the parameter, unconditionally, and not only
// on a Forward State 0→1 transition: subscribeUpstreamOnSession sends
// FORWARD=0 when no downstream forwards, and a track published once would then
// report no Largest Object, so no fill fetch stream could backfill it.
func saveLargestLocation(entry *registry.TrackEntry, ps message.Parameters) {
        if p, ok := ps.Find(message.ParamLargestObject); ok {
                entry.UpdateLargest(message.Location{Group: p.Group, Object: p.Object})
        }
}

// logInboundGoaway records the peer's GOAWAY (§10.4). The relay does not close
// the session: enforcing the Timeout is the sender's job. It only stops
// initiating requests to the peer (see [goingAway]).
//
// Deviation: the relay neither migrates its subscriptions to NewSessionURI nor
// closes the session once none remain (§9.4.1, §3.6); downstream clients
// re-subscribe when the session ends.
func (h *sessionHandler) logInboundGoaway(ctx context.Context) {
        g := h.sess.PeerGoaway()
        //nolint:gosec // G115: g.Timeout is a peer-supplied ms value; an out-of-range value yields a wrong duration, not a memory-safety issue.
        timeout := time.Duration(g.Timeout) * time.Millisecond
        h.log.LogAttrs(ctx, slog.LevelInfo, "relay received inbound GOAWAY",
                slog.Duration("timeout", timeout),
                slog.String("new_session_uri", string(g.NewSessionURI)))
}

// goingAway reports whether sess has a GOAWAY in either direction, so the
// relay initiates no new request on it (§10.4): having received one, an
// endpoint "SHOULD NOT initiate new requests to the peer"; having sent one,
// it "SHOULD avoid initiating requests unless required by migration". Every
// relay-initiated SUBSCRIBE, FETCH and PUBLISH checks this first.
func goingAway(sess *session.Session) bool { return sess.PeerGoaway() != nil || sess.GoawaySent() }

// subIDCounter allocates process-globally unique subscription IDs. It MUST
// be global, not per-handler: a TrackEntry aggregates subscriptions from
// many sessions (§9.4), and the registry removes by ID — two handlers'
// per-handler counters would collide, so one subscriber unsubscribing would
// silently delete another session's subscription from the same track.
var subIDCounter atomic.Uint64

// allocSubID returns a fresh, process-unique subscription ID. Used when
// instantiating registry.UpstreamSub / registry.DownstreamSub from inside the request handlers.
func (h *sessionHandler) allocSubID() uint64 {
        return subIDCounter.Add(1)
}

// run blocks until the session ends, returning nil on a clean close (peer or
// ctx-driven shutdown) or the request/data-loop error otherwise. It spawns the
// protocol loops, joins them and all in-flight per-request goroutines, and does
// not close the session itself except on a protocol violation detected by a loop.
func (h *sessionHandler) run(ctx context.Context) error {
        // Watcher ties runCtx to the parent ctx and the session's Done channel so
        // loops unblock as soon as the session terminates.
        runCtx, cancel := context.WithCancel(ctx)
        defer cancel()
        go func() {
                select {
                case <-h.sess.GoawayReceived():
                        h.logInboundGoaway(ctx)
                case <-h.sess.Done():
                case <-runCtx.Done():
                }
                select {
                case <-h.sess.Done():
                case <-runCtx.Done():
                }
                cancel()
        }()

        var (
                loops   sync.WaitGroup
                reqErr  error
                dataErr error
        )
        // The request and data loops are load-bearing: when either dies, the
        // session is no longer usable, so each cancels runCtx to unwind the rest.
        loops.Go(func() {
                reqErr = h.runRequestLoop(runCtx)
                cancel() // wake sibling loops if the request loop dies first
        })
        loops.Go(func() {
                dataErr = h.runDataLoop(runCtx)
                cancel() // wake sibling loops if the data loop dies first
        })
        // Datagrams are OPTIONAL (§11.3): a transport or peer without DATAGRAM
        // support fails ReceiveDatagram on the first call, which must not take down
        // SUBSCRIBE/PUBLISH handling. So the datagram loop neither cancels its
        // siblings nor promotes its error as a session fault — it just stops.
        loops.Go(func() {
                if err := h.runDatagramLoop(runCtx); err != nil && !isShutdownErr(err) {
                        h.log.LogAttrs(ctx, slog.LevelDebug,
                                "relay datagram loop ended; datagrams unavailable on this session",
                                slog.String("err", err.Error()))
                }
        })

        loops.Wait()
        h.wg.Wait()

        // Promote the first non-shutdown error (request > data) to the caller. All
        // errors are logged at Debug for postmortem.
        if reqErr != nil && !isShutdownErr(reqErr) {
                h.log.LogAttrs(ctx, slog.LevelDebug, "relay request loop ended", slog.String("err", reqErr.Error()))
                return reqErr
        }
        if dataErr != nil && !isShutdownErr(dataErr) {
                h.log.LogAttrs(ctx, slog.LevelDebug, "relay data loop ended", slog.String("err", dataErr.Error()))
                return dataErr
        }
        return nil
}

// runRequestLoop reads requests off the session and dispatches each to the
// appropriate handler. Each handler is responsible for its own bidi stream
// lifecycle — the loop hands off the [*session.Request] and does NOT wait
// for the handler to finish.
//
// The loop terminates when:
//
//   - ctx is cancelled (returns ctx.Err()),
//   - the session emits an unrecoverable error from AcceptRequest,
//   - a non-shutdown read failure occurs.
//
// Per-request failures (auth, rejected requests, a stream ended before its
// request) do NOT terminate the loop. A protocol violation does, and
// AcceptRequest has already closed the session with its code.
func (h *sessionHandler) runRequestLoop(ctx context.Context) error {
        return h.requestMux(ctx).Run(ctx, h.sess)
}

// runDataLoop accepts inbound data streams and routes each by type: subgroup
// streams to [sessionHandler.runFanout], fetch response streams to the fetch
// router (see the inline comments below).
//
// AcceptDataStream skips abandoned streams and closes the session on a fatal
// header itself, so any error it returns ends the loop.
func (h *sessionHandler) runDataLoop(ctx context.Context) error {
        for {
                ds, err := h.sess.AcceptDataStream(ctx)
                if err != nil {
                        if errors.Is(err, session.ErrPaddingStream) {
                                // §11.5.1 padding stream — silently discarded
                                // by AcceptDataStream itself; loop and try again.
                                continue
                        }
                        return err
                }
                switch s := ds.(type) {
                case *session.IncomingSubgroupStream:
                        h.spawn(func() { h.runFanout(ctx, s) })
                case *session.IncomingFetchStream:
                        // Body side of a FETCH the relay issued upstream on this
                        // session. Hand it to the downstream handler waiting on the
                        // matching (session, RequestID) via the fetch router; if none
                        // is registered (no stitch in flight, or a duplicate/late
                        // response), reset it to keep the upstream's flow control free.
                        if !h.fetch.Deliver(h.sess, s.Header.RequestID, s) {
                                h.log.LogAttrs(ctx, slog.LevelDebug, "relay dropped unmatched IncomingFetchStream",
                                        slog.Uint64("request_id", s.Header.RequestID))
                                s.Cancel(moqt.StreamResetInternalError)
                        }
                default:
                        h.log.LogAttrs(ctx, slog.LevelDebug, "relay dropped unknown data stream",
                                slog.String("type", fmt.Sprintf("%T", ds)))
                }
        }
}

// requestMux builds the per-session [session.RequestMux] that routes each inbound
// request to the handler responsible for its First-message type. Each handler is
// expected to:
//
//  1. Authorize the request.
//  2. Reply with either *_OK or REQUEST_ERROR.
//  3. Keep the bidi stream open for as long as the subscription's lifetime
//     warrants (or close it cleanly on rejection).
//  4. Update [registry.TrackRegistry] / [registry.NamespaceRegistry] as appropriate.
//
// Two cross-cutting policies are shared across the per-type handlers:
// verifyRequest applies the §10.2.2 token-verification pre-step, and
// namespaceRequest folds in the §13.7.1 per-session cap for the three
// namespace-state requests (the §13.1 subscription cap is inline on SUBSCRIBE).
//
// Any other first message is a §3.3 PROTOCOL_VIOLATION that
// [session.Session.AcceptRequest] handles before dispatch.
func (h *sessionHandler) requestMux(ctx context.Context) *session.RequestMux {
        mux := session.NewRequestMux()

        mux.HandleType(func(req *session.Request, msg *message.Subscribe) {
                if !h.verifyRequest(ctx, req) {
                        return
                }
                // §13.1: bound concurrent subscriptions per session.
                if !h.limiter.acquireSub() {
                        h.rejectExcessiveLoad(ctx, req, "subscription")
                        return
                }
                h.spawn(func() { defer h.limiter.releaseSub(); h.handleSubscribe(ctx, req, msg) })
        })
        mux.HandleType(func(req *session.Request, msg *message.Publish) {
                if !h.verifyRequest(ctx, req) {
                        return
                }
                h.spawn(func() { h.handlePublish(ctx, req, msg) })
        })
        mux.HandleType(func(req *session.Request, msg *message.Fetch) {
                if !h.verifyRequest(ctx, req) {
                        return
                }
                h.spawn(func() { h.handleFetch(ctx, req, msg) })
        })
        mux.HandleType(func(req *session.Request, msg *message.TrackStatus) {
                if !h.verifyRequest(ctx, req) {
                        return
                }
                h.spawn(func() { h.handleTrackStatus(ctx, req, msg) })
        })
        mux.HandleType(func(req *session.Request, msg *message.PublishNamespace) {
                h.namespaceRequest(ctx, req, func() { h.handlePublishNamespace(ctx, req, msg) })
        })
        mux.HandleType(func(req *session.Request, msg *message.SubscribeNamespace) {
                h.namespaceRequest(ctx, req, func() { h.handleSubscribeNamespace(ctx, req, msg) })
        })
        mux.HandleType(func(req *session.Request, msg *message.SubscribeTracks) {
                h.namespaceRequest(ctx, req, func() { h.handleSubscribeTracks(ctx, req, msg) })
        })

        return mux
}

// verifyRequest runs the per-request dispatch log and the §10.2.2 token
// verification shared by every known request type. It returns false — after
// replying REQUEST_ERROR with the mapped code — when the request's resolved
// AUTHORIZATION_TOKEN is denied; the session stays up (a denial is per-request).
func (h *sessionHandler) verifyRequest(ctx context.Context, req *session.Request) bool {
        h.log.LogAttrs(ctx, slog.LevelDebug, "relay dispatching request",
                slog.String("type", fmt.Sprintf("%T", req.First)))
        if err := h.sess.VerifyRequestTokens(ctx, req); err != nil {
                h.rejectTokenDenied(ctx, req, err)
                return false
        }
        return true
}

// namespaceRequest wraps a namespace-state handler (PUBLISH_NAMESPACE,
// SUBSCRIBE_NAMESPACE, SUBSCRIBE_TRACKS) with the shared token verification and
// the §13.7.1 per-session cap, spawning fn under the limiter when admitted.
func (h *sessionHandler) namespaceRequest(ctx context.Context, req *session.Request, fn func()) {
        if !h.verifyRequest(ctx, req) {
                return
        }
        if !h.limiter.acquireNamespace() {
                h.rejectExcessiveLoad(ctx, req, "namespace request")
                return
        }
        h.spawn(func() { defer h.limiter.releaseNamespace(); fn() })
}

// spawn registers a goroutine with the handler's wg so run() can join it
// during shutdown. Handlers are responsible for handling their own ctx
// cancellation; spawn does not impose a timeout.
func (h *sessionHandler) spawn(fn func()) {
        h.wg.Go(fn)
}

// ---------------------------------------------------------------------------
// Request rejection helpers
// ---------------------------------------------------------------------------

// rejectAuth writes a REQUEST_ERROR with the code derived from the authorizer
// error and FINs the bidi stream. Any write failure is logged but otherwise
// swallowed — the stream is being torn down anyway.
func (h *sessionHandler) rejectAuth(ctx context.Context, req *session.Request, kind string, authErr error) {
        code := CodeForAuthorizerError(authErr)
        reason := ReasonForAuthorizerError(authErr)
        if err := req.RejectError(code, reason); err != nil && !errors.Is(err, context.Canceled) {
                h.log.LogAttrs(ctx, slog.LevelDebug, "relay reject write failed",
                        slog.String("kind", kind), slog.String("err", err.Error()))
        }
}

// excessiveLoadRetry is the least wait an EXCESSIVE_LOAD rejection invites; a
// guess, since the relay cannot predict when a per-session cap frees up.
const excessiveLoadRetry = time.Second

// retryIntervalAfter is a Retry Interval (§10.6.2) inviting a retry after d
// plus up to 50% jitter, encoded as milliseconds plus one.
//
//nolint:unparam // callers pass separate policy constants, equal only today.
func retryIntervalAfter(d time.Duration) uint64 {
        ms := uint64(d / time.Millisecond) //nolint:gosec // G115: callers pass a positive constant duration.
        return ms + rand.Uint64N(ms/2) + 1 //nolint:gosec // G404: retry jitter, not a secret.
}

// rejectExcessiveLoad rejects a request that exceeds a per-session resource cap
// (§13.1 / §13.7.1) with REQUEST_ERROR EXCESSIVE_LOAD and FINs the bidi stream.
// what names the limit category for the log/reason. The reject happens before
// any registry mutation, so no cleanup is needed.
func (h *sessionHandler) rejectExcessiveLoad(ctx context.Context, req *session.Request, what string) {
        h.log.LogAttrs(ctx, slog.LevelDebug, "relay rejecting request: per-session limit reached",
                slog.String("limit", what))
        if err := req.Reject(&session.RequestRejectedError{
                Code:          moqt.RequestExcessiveLoad,
                Reason:        "relay: " + what + " limit reached",
                RetryInterval: retryIntervalAfter(excessiveLoadRetry),
        }); err != nil &&
                !errors.Is(err, context.Canceled) {
                h.log.LogAttrs(ctx, slog.LevelDebug, "relay EXCESSIVE_LOAD reject write failed",
                        slog.String("err", err.Error()))
        }
}

// rejectTokenDenied writes a REQUEST_ERROR for a token-verification denial and
// FINs the bidi stream. The error is always a [*session.TokenDeniedError]
// (VerifyRequestTokens normalises plain verifier errors into one), so its
// RequestErrorCode — e.g. [moqt.RequestExpiredAuthToken] or the default
// [moqt.RequestUnauthorized] — maps straight onto the wire reply. Like
// rejectAuth, a write failure is logged and otherwise swallowed.
func (h *sessionHandler) rejectTokenDenied(ctx context.Context, req *session.Request, denyErr error) {
        code, reason := tokenDenial(denyErr)
        h.log.LogAttrs(ctx, slog.LevelDebug, "relay rejecting request on token verification",
                slog.String("err", denyErr.Error()), slog.Uint64("code", uint64(code)))
        if err := req.RejectError(code, reason); err != nil && !errors.Is(err, context.Canceled) {
                h.log.LogAttrs(ctx, slog.LevelDebug, "relay token-denied reject write failed",
                        slog.String("err", err.Error()))
        }
}

// tokenDenial is the REQUEST_ERROR code and reason for a token-verification
// denial (see [sessionHandler.rejectTokenDenied]).
func tokenDenial(denyErr error) (moqt.RequestErrorCode, string) {
        if denied, ok := errors.AsType[*session.TokenDeniedError](denyErr); ok {
                if denied.Reason != "" {
                        return denied.RequestErrorCode(), denied.Reason
                }
                return denied.RequestErrorCode(), denyErr.Error()
        }
        return moqt.RequestUnauthorized, denyErr.Error()
}

// refuseUpdateTokens is the REQUEST_ERROR refusing a REQUEST_UPDATE whose
// AUTHORIZATION_TOKENs the TokenVerifier denies, or nil: they authorize "the
// operation carrying the parameter" (§10.2.2) as an opener's do.
func (h *sessionHandler) refuseUpdateTokens(ctx context.Context, toks []session.ResolvedToken) *message.RequestError {
        err := h.sess.VerifyTokens(ctx, toks)
        if err == nil {
                return nil
        }
        code, reason := tokenDenial(err)
        return &message.RequestError{ErrorCode: code, ErrorReason: reason}
}

// handleFollowupRequestID validates a peer REQUEST_UPDATE's Request ID —
// §10.1: an update consumes an ID from the sender's space, and the readers
// that parse follow-ups directly bypass AcceptRequest's checking. A
// wrong-parity or duplicate ID is session-fatal (INVALID_REQUEST_ID);
// returns false when the session was closed, in which case the caller's
// read loop should stop.
func (h *sessionHandler) handleFollowupRequestID(ctx context.Context, upd *message.RequestUpdate) bool {
        err := h.sess.CheckPeerRequestID(upd.RequestID)
        if err == nil {
                return true
        }
        h.log.LogAttrs(ctx, slog.LevelDebug, "relay closing session on follow-up Request ID violation",
                slog.String("err", err.Error()))
        _ = h.sess.Close(moqt.SessionInvalidRequestID, err.Error())
        return false
}

// handleRequestUpdateLimit charges one credit on lim for a received
// REQUEST_UPDATE and enforces the per-stream MAX_REQUEST_UPDATES limit
// (§10.3.1.7). Exceeding it is session-fatal (TOO_MANY_REQUEST_UPDATES);
// returns false when the session was closed, in which case the caller's read
// loop should stop. The caller invokes lim.Responded once it has written the
// mandated REQUEST_OK/REQUEST_ERROR.
func (h *sessionHandler) handleRequestUpdateLimit(ctx context.Context, lim *session.RequestUpdateLimiter) bool {
        err := lim.Received()
        if err == nil {
                return true
        }
        h.log.LogAttrs(ctx, slog.LevelDebug, "relay closing session on REQUEST_UPDATE limit",
                slog.String("err", err.Error()))
        _ = h.sess.Close(moqt.SessionTooManyRequestUpdates, err.Error())
        return false
}

// handleFollowupTokens routes a follow-up message's AUTHORIZATION_TOKEN
// parameters through the session token cache — §10.2.2 allows REQUEST_UPDATE
// to REGISTER or DELETE aliases, and the readers that parse follow-ups
// directly bypass AcceptRequest's processing. It returns the resolved tokens,
// and ok false when a token fault closed the session, in which case the
// caller's read loop should stop.
func (h *sessionHandler) handleFollowupTokens(
        ctx context.Context,
        msg message.Message,
) (toks []session.ResolvedToken, ok bool) {
        toks, err := h.sess.ProcessFollowupTokens(msg)
        if err == nil {
                return toks, true
        }
        if tce, ok := errors.AsType[*session.TokenCacheError](err); ok {
                h.log.LogAttrs(ctx, slog.LevelDebug, "relay closing session on follow-up token cache error",
                        slog.String("err", err.Error()),
                        slog.Uint64("code", uint64(tce.Code)))
                _ = h.sess.Close(tce.Code, tce.Error())
                return nil, false
        }
        h.log.LogAttrs(ctx, slog.LevelDebug, "follow-up token processing failed",
                slog.String("err", err.Error()))
        return nil, false
}

// errRequestCancelled is the cancellation cause of a context bound to a
// request its peer cancelled (§3.3.3); see [ctxResetCode].
var errRequestCancelled = errors.New("relay: request cancelled")

// ctxResetCode is the §3.3.4 code for a stream reset because ctx ended:
// CANCELLED for a cancelled request, MALFORMED_TRACK for a fetch stream of a
// malformed track (§2.4.2), SESSION_CLOSED otherwise.
func ctxResetCode(ctx context.Context) moqt.StreamResetCode {
        switch cause := context.Cause(ctx); {
        case errors.Is(cause, errRequestCancelled):
                return moqt.StreamResetCancelled
        case errors.Is(cause, session.ErrMalformedTrack):
                return moqt.StreamResetMalformedTrack
        }
        return moqt.StreamResetSessionClosed
}

// readRequestStream owns all reads on an established request stream: it
// parses follow-up messages off the stream and dispatches each to onMsg
// until the peer ends its side (FIN or reset), onMsg returns false,
// or ctx is cancelled (the read side is then reset with
// StreamResetSessionClosed to unblock the parse). A follow-up that cannot be
// read — any non-EOF error — resets the read side with
// StreamResetInternalError so the peer learns reads stopped; a malformed one
// also closes the session (§10), as does a REQUEST_OK or REQUEST_ERROR, since
// the relay sends no REQUEST_UPDATE on the streams it reads this way (§10.9).
//
// It reports fin when the requester ended its side with a FIN, which is not a
// cancellation (§3.3.2); see [awaitRequestEnd].
func readRequestStream(
        ctx context.Context,
        sess *session.Session,
        stream session.Stream,
        onMsg func(message.Message) bool,
) (fin bool) {
        // done carries the fin result, so nothing else escapes to the heap.
        done := make(chan bool, 1)
        go func() {
                var goaways session.RequestGoaways
                for {
                        m, err := message.Parse(stream)
                        if err != nil {
                                // Covers peer resets too (a STOP_SENDING on an
                                // already-reset stream is a transport no-op), and may run
                                // after the ctx arm's SessionClosed CancelRead — the first
                                // code sent wins on every bundled adapter.
                                eof := errors.Is(err, io.EOF)
                                if !eof {
                                        stream.CancelRead(uint64(moqt.StreamResetInternalError))
                                }
                                // §10, §10.2: a malformed message MUST close the session.
                                if errors.Is(err, message.ErrMalformedMessage) {
                                        _ = sess.Close(moqt.SessionProtocolViolation, err.Error())
                                }
                                done <- eof
                                return
                        }
                        // §10.4: a second GOAWAY, or one with a URI from a client,
                        // closed the session.
                        if g, ok := m.(*message.Goaway); ok && goaways.Received(sess, g) != nil {
                                done <- false
                                return
                        }
                        // §10.9: the relay sends no REQUEST_UPDATE on these streams, so
                        // a REQUEST_OK or REQUEST_ERROR answers nothing: on a forwarded
                        // PUBLISH it is a second response (§5.1).
                        switch m.(type) {
                        case *message.RequestOK, *message.RequestError:
                                _ = sess.Close(moqt.SessionProtocolViolation,
                                        fmt.Sprintf("%s with no REQUEST_UPDATE to answer", m.Type()))
                                done <- false
                                return
                        }
                        if !onMsg(m) {
                                done <- false
                                return
                        }
                }
        }()
        select {
        case fin = <-done:
                return fin
        case <-ctx.Done():
                stream.CancelRead(uint64(moqt.StreamResetSessionClosed))
                <-done
                return false
        }
}

// isPeerStateNotify reports a PUBLISH_STATE_NOTIFY from the requester of a
// request the relay is answering, and closes the session for it (§10.10: "is
// sent only by the publisher").
func (h *sessionHandler) isPeerStateNotify(m message.Message) bool {
        if _, ok := m.(*message.PublishStateNotify); !ok {
                return false
        }
        _ = h.sess.Close(moqt.SessionProtocolViolation,
                "PUBLISH_STATE_NOTIFY from the requester")
        return true
}

// awaitRequestEnd keeps a request whose requester FINned its side alive until
// it really ends (§3.3.2: a FIN "is not a request cancellation"). The send
// Context ends on the requester's STOP_SENDING (§3.3.3) or when the relay ends
// its own side.
func awaitRequestEnd(ctx context.Context, stream session.Stream) {
        select {
        case <-stream.Context().Done():
        case <-ctx.Done():
        }
}

// serveFetchObjects is the response tail of the FETCH handler: stream the
// stitched range, FIN, and park in the §10.9 follow-up loop until the
// requester resets or FINs the request stream. kind tags log lines.
//
// A cancel before the FIN (the requester's STOP_SENDING, §3.3.3) resets both
// streams with CANCELLED. §5.2: the publisher "MUST reset the bidi request
// stream and unidirectional data stream associated with the FETCH". A data
// stream reset because the track must not be forwarded (§2.4.2, §2.5.1) ends
// the request too: the relay cancels it with the same code (§3.3.3).
func (h *sessionHandler) serveFetchObjects(
        ctx context.Context,
        req *session.Request,
        kind string,
        requestID uint64,
        entry *registry.TrackEntry,
        fullName track.FullTrackName,
        start, end message.Location,
        order message.GroupOrder,
        fillTimeout time.Duration,
        rangeFilters *message.RangeFilterSet,
) {
        fetchCtx, cancel := context.WithCancelCause(ctx)
        defer cancel(nil)
        stop := context.AfterFunc(req.Stream.Context(), func() { cancel(errRequestCancelled) })
        // §2.4.2: a relay that detects a malformed track MUST "reset any fetch
        // streams with Status Code MALFORMED_TRACK"; see endMalformedTrack.
        remove := entry.AddFetch(cancel)
        out, refused, code := h.streamFetchRange(fetchCtx, kind, nil, requestID, entry, fullName,
                start, end, order, fillTimeout, rangeFilters)
        stop()
        remove()
        if out == nil {
                // The requester's own signal, not fetchCtx's cause: a write can fail on
                // its STOP_SENDING for the data stream before the cause is set.
                switch {
                case req.Stream.Context().Err() != nil:
                        code = moqt.StreamResetCancelled
                case !refused:
                        return
                }
                // §3.3.3: "RESET_STREAM for a direction they are sending and
                // STOP_SENDING for a direction they are receiving".
                req.Stream.CancelRead(uint64(code))
                req.Stream.CancelWrite(uint64(code))
                return
        }

        h.readFetchUpdates(ctx, req, out)
}

// streamFetchRange opens a unidirectional fetch stream, writes the stitched
// range to it, and FINs. It is the shared body of a FETCH response (§10.13)
// and of a fill fetch stream (§5.1.3), which differ only in what opens them
// and in what happens afterwards — a FETCH parks in the §10.9 follow-up loop,
// a fill is simply done.
//
// It returns the FINed stream, or nil when the stream could not be opened, the
// write failed, or the track must not be forwarded; the stream is then already
// reset. refused reports the last case, a malformed track (§2.4.2) or an
// upstream refusal (§2.5.1), and, when refused, code the data stream was reset
// with.
func (h *sessionHandler) streamFetchRange(
        ctx context.Context,
        kind string,
        sub *registry.DownstreamSub,
        requestID uint64,
        entry *registry.TrackEntry,
        fullName track.FullTrackName,
        start, end message.Location,
        order message.GroupOrder,
        fillTimeout time.Duration,
        rangeFilters *message.RangeFilterSet,
) (_ *session.OutgoingFetchStream, refused bool, code moqt.StreamResetCode) {
        out, err := openFillOrFetchStream(h.sess, sub, requestID)
        if err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "OpenFetchStream failed",
                        slog.String("kind", kind), slog.String("err", err.Error()))
                return nil, false, 0
        }
        if sub != nil {
                // A fill stream's subscription holds its PUBLISH_DONE until the
                // stream closes (§10.12).
                defer sub.StreamClosed()
        }
        // ctx ending resets the stream rather than completing it: CANCELLED when
        // its cause is errRequestCancelled (a fill's cancelled subscription,
        // §5.1.3.1, or a cancelled FETCH, §5.2), MALFORMED_TRACK for a malformed
        // track (§2.4.2; callers register ctx's cancel with
        // [registry.TrackEntry.AddFetch]), else SESSION_CLOSED (§3.3.4).
        cancelOut := func() { out.Cancel(ctxResetCode(ctx)) }
        unwatch := context.AfterFunc(ctx, cancelOut)
        defer unwatch()
        // ctxEnded resets the stream for ctx, before the deferred StreamClosed
        // (§10.12), and reports it.
        ctxEnded := func() (*session.OutgoingFetchStream, bool, moqt.StreamResetCode) {
                cancelOut()
                c := ctxResetCode(ctx)
                return nil, c == moqt.StreamResetMalformedTrack, c
        }

        // Gather cached objects, asking an upstream about what the cache cannot
        // vouch for (§10.13).
        objs, refusal := h.stitchedFetchObjects(ctx, entry, fullName, start, end, order, fillTimeout)
        if ctx.Err() != nil {
                return ctxEnded()
        }
        if refusal != nil {
                // §2.5.1: with FETCH_OK (or SUBSCRIBE_OK) already sent, only a
                // reset is left (an interpretation: no Object was forwarded yet).
                // Unparseable Track Properties (§12.7, §2.4.2) and a malformed upstream
                // Object (§2.4.2) reset with MALFORMED_TRACK.
                code = moqt.StreamResetInternalError
                if errors.Is(refusal, session.ErrMalformedTrackProperties) || errors.Is(refusal, session.ErrMalformedTrack) {
                        code = moqt.StreamResetMalformedTrack
                }
                h.log.LogAttrs(ctx, slog.LevelDebug, "upstream FETCH refused",
                        slog.String("kind", kind), slog.String("err", refusal.Error()))
                out.Cancel(code)
                return nil, true, code
        }

        // §5.1.4: drop objects that fail the request's Range Filters. §11.4.4.2
        // end-of-range markers are not objects and are always kept — they carry no
        // Subgroup ID, Priority or Properties, so matching one against a filter
        // tests zero values and drops it, turning the span into a plain gap that
        // §10.13 reads as authoritative non-existence.
        if rangeFilters != nil {
                objs = slices.DeleteFunc(objs, func(o *cache.CachedObject) bool {
                        return !o.IsRangeMarker() &&
                                !rangeFilters.MatchesObject(o.SubgroupID, o.ObjectID, o.PublisherPriority, o.Properties)
                })
        }

        written, err := streamFetchObjects(out, objs, entry.Cache.Expired)
        h.metrics.FetchServed(h.trackRef(fullName), written)
        if err != nil {
                h.log.LogAttrs(ctx, slog.LevelDebug, "fetch stream write failed",
                        slog.String("kind", kind), slog.String("err", err.Error()))
                if ctx.Err() != nil {
                        return ctxEnded() // ctx's reset failed the write; keep its code
                }
                out.Cancel(moqt.StreamResetInternalError)
                return nil, false, 0
        }
        if !unwatch() {
                return ctxEnded() // ctx ended first
        }
        _ = out.Close()
        return out, false, 0
}