From 70872fa632fc431597e46f632d81027a789e8dc4 Mon Sep 17 00:00:00 2001 From: Doloro1978 Date: Mon, 6 Jul 2026 12:31:55 +0100 Subject: [PATCH] add: h265 & AV1 support (with a lot of fixes) --- .cargo/config.toml | 2 + .gitignore | 4 +- CLAUDE.md | 252 +++++++++++--- Cargo.toml | 6 + README.md | 7 +- crates/server/src/codec/av1.rs | 138 ++++++++ crates/server/src/codec/flv_replay_test.rs | 95 ++++++ crates/server/src/codec/h264.rs | 175 ++++++++++ crates/server/src/codec/h265.rs | 142 ++++++++ crates/server/src/codec/mod.rs | 18 + crates/server/src/http.rs | 63 ++-- crates/server/src/main.rs | 51 ++- crates/server/src/media.rs | 147 -------- crates/server/src/rtmp.rs | 369 ++++++++++++--------- crates/server/src/webrtc.rs | 244 +++++++++----- crates/server/src/webrtc_proxy.rs | 39 ++- 16 files changed, 1259 insertions(+), 493 deletions(-) create mode 100644 .cargo/config.toml create mode 100644 crates/server/src/codec/av1.rs create mode 100644 crates/server/src/codec/flv_replay_test.rs create mode 100644 crates/server/src/codec/h264.rs create mode 100644 crates/server/src/codec/h265.rs create mode 100644 crates/server/src/codec/mod.rs delete mode 100644 crates/server/src/media.rs diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..03f4879 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,2 @@ +[build] +rustflags = ["-Ctarget-cpu=native"] diff --git a/.gitignore b/.gitignore index 54321a2..e15da0f 100644 --- a/.gitignore +++ b/.gitignore @@ -11,4 +11,6 @@ devenv.local.yaml # pre-commit .pre-commit-config.yaml /stream.db -/db.sqlite +# /db.sqlite +# /db +/db/db.sqlite diff --git a/CLAUDE.md b/CLAUDE.md index 6da9a2e..6d5dd65 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -31,80 +31,228 @@ crates/ ## Architecture -This is an RTMP-to-WHEP bridge: accepts an RTMP video publish stream and re-streams it to browsers via WebRTC (WHEP signaling protocol). +This is an RTMP-to-WHIP/WHEP bridge: accepts RTMP video+audio publish streams and re-streams to browsers via WebRTC. **Signal flow:** ``` -OBS/encoder → RTMP (port 8123) → H264Parser → async_broadcast channel - ↓ -Browser ← WebRTC/UDP ← str0m Rtc ← WHEP HTTP (port 3000) +OBS/encoder → RTMP (port 1935) → H264Parser / AACParser → async_broadcast channels + ↓ +Browser ← WebRTC/UDP ← str0m Rtc ← WHIP/WHEP HTTP (port 3000) + ↑ + WebrtcProxy (UDP port from RTC_PORT env, default 6969) ``` -**Runtime:** tokio (not smol). The HTTP server runs as a tokio task using hyper. +**Runtime:** tokio. All top-level workers run in a `JoinSet`; if any exits unexpectedly the process aborts all others and exits. -### RTMP ingestion (port 8123) — `crates/server/src/main.rs` +--- -Each incoming TCP connection is handled in a detached tokio task: +### Entry point — `crates/server/src/main.rs` -1. **Handshake** — reads C0+C1 (1537 bytes) via `rml_rtmp::Handshake`, sends S0+S1+S2, reads C2 (1536 bytes). -2. **Session setup** — creates an `rml_rtmp::ServerSession`, writes its initial response bytes. -3. **Event loop** — reads 4096-byte chunks, calls `rtmp_session.handle_input`. Handles: +Initialises everything, then hands each worker to a `tokio::task::JoinSet`: + +```rust +workers.spawn(http.run()); +workers.spawn(proxy.run()); +workers.spawn(webrtc.run()); +workers.spawn(rtmp.run()); +``` + +`tokio::select!` waits for either Ctrl-C or a worker exiting; on either branch `workers.abort_all()` is called. + +**`AppState`** — `Arc>` wrapping a `DashMap` keyed by `stream_key.id`. + +**`StreamSession`** fields: +- `stream_key_id: i32` +- `stream_key_label: String` +- `frame_channel: async_broadcast::Sender>` +- `audio_channel: async_broadcast::Sender>` +- `codec: Option` — `H264 | H265 | AV1` (AV1 unsupported) + +Shared channels between main components: +- `offer_tx/offer_rx: mpsc::channel<(request_id: i32, stream_key_id: i32, sdp_body: String)>(64)` — HTTP → Webrtc +- `answer_tx/answer_rx: async_broadcast<(request_id: i32, Option)>(64)` — Webrtc → HTTP + +DB: `sqlite://./db/db.sqlite?mode=rwc`. Migrations run at startup via `Migrator::up`. `stream_session::Model::clean_unended_streams` is called at startup to repair sessions left open by a previous crash. + +--- + +### RTMP ingestion — `crates/server/src/rtmp.rs` + +`pub struct Rtmp { listener, stream_sessions, db }` + +`async fn run(self)` — accepts TCP connections on port 1935, one `tokio::spawn` per connection: + +1. **Handshake** — `rml_rtmp::Handshake` (C0+C1 → S0+S1+S2 → C2). +2. **Session setup** — `rml_rtmp::ServerSession`. +3. **Event loop** — 4096-byte reads, `session.handle_input`: - `ConnectionRequested` → accepted unconditionally. - - `PublishStreamRequested` → accepted for any stream key; a `StreamSession` with an `async_broadcast::Sender>` is inserted into `AppState`. - - `VideoDataReceived` → HEVC check (bytes 1–4 == `hvc1`; drops connection if true), then passed to `H264Parser::parse`. Parsed frames are broadcast on the channel. - - `PublishStreamFinished` → entry removed from `AppState`. + - `PublishStreamRequested` → looks up stream key in DB; rejects if not found or if a `stream_session` record is already active. On accept: inserts a `StreamSession` into `AppState` and creates a `stream_session` row in DB. + - `VideoDataReceived` → dispatched to `H264Parser::parse` (H.265/AV1 drops the connection). Parsed `VideoFrame`s are broadcast on `frame_channel`. + - `AudioDataReceived` → dispatched to `AACParser::parse_aac`, then `AudioProcesser` transcodes AAC→Opus. `OpusAudioFrame`s broadcast on `audio_channel`. + - `PublishStreamFinished` → removes `StreamSession` from `AppState`, sets `ended_at` on the `stream_session` DB row. -`AppState` is a `Arc>` wrapping a `DashMap`. +--- -### HTTP API (port 3000) — `crates/server/src/http.rs` +### HTTP API — `crates/server/src/http.rs` -Async hyper server running in a tokio task. Routes: +`pub struct HttpServer` fields: `offer_tx`, `accept_rx` (broadcast), `appstate`, `request_count: AtomicI32`, `db`, `config: Arc`. -- `GET /api/catalog` — returns JSON `{ active_streams: [String] }` listing currently publishing stream keys. -- `GET /api/meow` — returns `"meow"` (health check / placeholder). -- WHEP signaling is not yet wired into this HTTP server (see webrtc.rs for the channel plumbing). +`async fn run(self)` — wraps self in `Arc`, builds axum router with CORS (allows `localhost:5173` and `stream.h.doloro.co.uk`), serves on `0.0.0.0:3000`. -SDP offer/answer exchange uses `tokio::sync::mpsc` channels between HttpServer and the Webrtc task. +**Routes:** -### WebRTC negotiation and media loop — `crates/server/src/webrtc.rs` +| Method | Path | Handler | +|--------|------|---------| +| `GET` | `/api/catalog` | `catalog_handler` — returns `{ active_streams: [{ label, id, user }] }` from `AppState` joined with DB user lookup | +| `POST` | `/api/user` | `create_user_handler` — creates user; requires `SIGNUP_CODE` header matching env var | +| `GET/POST` | `/api/stream-key` | `get_all_stream_keys` / `create_stream_key_handler` | +| `POST` | `/api/whip` | `handle_whip_injest` (in `webrtc_ingest.rs`) — WHIP ingest endpoint | +| `POST` | `/api/login` | `login_handler` — verifies password hash, creates `auth_session`, returns session cookie | +| `POST` | `/api/stream/{slug}` | `stream_handler` — WHEP offer: sends SDP offer over `offer_tx`, waits on `accept_rx` for matching `request_id`, returns SDP answer | -The `Webrtc` task receives `(stream_key, sdp_body)` tuples from `offer_rx`: +Auth: session token sent as `session` header or cookie; `auth_session` entity looked up from DB. `FromRequestParts` extractor `AuthSession` handles this for protected routes. -1. Binds a UDP socket to `127.0.0.1:0` — only ICE candidate advertised (host, UDP, loopback). -2. Builds `str0m::Rtc` with H.264 explicitly configured for PTs 102, 104, 106 (profiles `0x42e01f`, `0x4d001f`, `0x64001f`). Default H.264 support is disabled first. -3. Adds a `SendOnly` video media track, calls `changes.accept_offer(offer_sdp)` to produce the SDP answer. -4. Sends the answer back on `accept_tx`, then spawns a per-connection tokio task. +`request_count: AtomicI32` tracks in-flight request IDs (monotonic, `fetch_add(1, Relaxed)`). -**Per-connection loop** (`Webrtc::detach_connection`): +--- -- Drains `rtc.poll_output()` until `Output::Timeout`. Each iteration sends UDP datagrams (`Output::Transmit`) or handles events: - - `Event::MediaAdded` — picks the PT with the highest `profile_level_id`, stores in `video_pt`. - - `Event::Connected` — sets `connected = true`; media sending begins. -- When connected, lazily subscribes to the `async_broadcast` channel for the stream key, then drains up to 8 frames per iteration via `try_recv`, writing each with `writer.write(pt, now, rtp_time, frame.data)`. `rtp_time` is computed as `MediaTime::from_90khz(timestamp_ms * 90)`. -- Waits (capped at 20 ms) with `tokio::select!` for either the str0m deadline or a UDP datagram. Incoming datagrams are fed to `rtc.handle_input(Input::Receive(...))`. +### WebRTC proxy — `crates/server/src/webrtc_proxy.rs` + +`#[derive(Clone)] pub struct WebrtcProxy` — all fields are `Arc`-wrapped: +- `clients_ufrag: Arc>>` — pending ICE ufrag → per-client channel +- `clients_addr: Arc>>` — established addr → per-client channel +- `socket: Arc` — shared UDP socket bound to `0.0.0.0:{RTC_PORT}` (default 6969) +- `public_addr: SocketAddr` — resolved via `PUBLIC_DOMAIN` env var DNS lookup or STUN discovery + +**`async fn run(self)`** — UDP receive loop: +1. Receives datagrams on the shared socket. +2. If source addr is already in `clients_addr`, forwards to that client's channel. +3. Otherwise parses STUN binding request to extract ufrag (`username` attribute, part before `:`), looks up `clients_ufrag`, promotes to `clients_addr`, forwards. + +**`fn add_client(ufrag, …) -> (Arc, Receiver<…>)`** — called by `Webrtc` when setting up a new peer connection. Registers the ufrag and returns the shared socket + a per-client receive channel. + +**`fn public_addr()`** — returns the public address advertised in ICE candidates. + +--- + +### WebRTC negotiation — `crates/server/src/webrtc.rs` + +`pub struct Webrtc { offer_rx, accept_tx, sessions_ref, db, proxy: Arc }` + +**`async fn run(mut self)`** — receives `(request_id, stream_id, sdp_body)` from `offer_rx`: + +1. Looks up `stream_key` in DB; sends `None` answer and continues on error/not-found. +2. Gets `public_addr` from proxy for ICE candidate. +3. Builds `str0m::Rtc` with H.264 PTs 102 (`0x42e01f`), 104 (`0x4d001f`), 106 (`0x64001f`); default H.264 disabled. +4. Calls `add_client(ufrag)` on proxy to register ICE ufrag and get the UDP socket + channel. +5. Accepts SDP offer → produces SDP answer → broadcasts answer on `accept_tx`. +6. Spawns `detach_connection` task. + +**`detach_connection`** per-peer loop: +- Drains `rtc.poll_output()`: sends transmits via the shared proxy UDP socket, handles `Event::MediaAdded` (selects best PT by `profile_level_id`) and `Event::Connected`. +- When connected, subscribes to `frame_channel` from `AppState` for the stream, drains up to 8 frames per tick via `try_recv`, writes with `writer.write(pt, now, MediaTime::from_90khz(ts * 90))`. +- `tokio::select!` (capped 20 ms) on str0m deadline or incoming UDP datagram from proxy channel. + +--- + +### WHIP ingest — `crates/server/src/webrtc_ingest.rs` + +`async fn handle_whip_injest` — axum handler for `POST /api/whip`. Currently a stub; extracts `State>` and the request body. (Implementation in progress.) + +--- + +### Adding a new codec + +To add support for a new ingest codec: + +1. **`rtmp.rs` — `parse_video_codec`**: add a FourCC arm (enhanced RTMP) or legacy codec ID. `StreamCodec` enum lives in `main.rs`. + +2. **New parser struct** (e.g. `crates/server/src/codec/mycodec.rs`): + - Field for each parameter set (`Vec`) + - `parse_sequence_header(&mut self, payload: &[u8])` — parses the decoder config record, caches parameter sets + - `to_annexb(&self, payload: &[u8], is_keyframe: bool) -> Option>` — converts length-prefixed NALUs to Annex-B, prepends parameter sets before keyframes + +3. **`rtmp.rs` — `VideoDataReceived` handler**: branch on `StreamCodec`, slice the payload correctly for each packet type, call `parse_sequence_header` on type `0` and `to_annexb` on types `1`/`3`, broadcast the resulting `VideoFrame`. + +4. **`webrtc.rs` — codec config**: configure the correct PT via `codec_config()` (e.g. `enable_h265`, `add_h264`). Fix PT selection in `Event::MediaAdded` if the new codec uses a different profile field than `profile_level_id`. + +5. **`StreamSession`** (`main.rs`): `codec: StreamCodec` field — set it when inserting into `stream_sessions` in `rtmp.rs` so the WebRTC layer can know what codec the session is using. + +--- ### H.264 parsing — `crates/server/src/media.rs` -`H264Parser` converts raw RTMP `VideoDataReceived` payloads (AVCC) to Annex-B: +`H264Parser` converts RTMP `VideoDataReceived` AVCC payloads → Annex-B `VideoFrame`s: -- **Byte 0**: upper nibble = frame type (1 = keyframe), lower nibble = codec ID (7 = H.264; anything else dropped). -- **Byte 1**: AVC packet type — `0` = sequence header, `1` = NAL unit data. -- **Bytes 5+**: payload. +- Byte 0: frame type (upper nibble, 1=keyframe) + codec ID (lower nibble, 7=H.264). +- Byte 1: AVC packet type — `0`=sequence header, `1`=NAL data. +- Packet type `0`: parses `AVCDecoderConfigurationRecord`, caches SPS+PPS. +- Packet type `1`: converts AVCC length-prefixed NALUs to `00 00 00 01` Annex-B. Prepends SPS+PPS before the first NALU of each keyframe. -Packet type `0` walks `AVCDecoderConfigurationRecord` to cache SPS and PPS byte arrays. +`pub struct VideoFrame { pub data: Bytes, pub is_keyframe: bool, pub timestamp_ms: u32 }` -Packet type `1` converts AVCC (4-byte big-endian length prefix per NALU) to Annex-B (`00 00 00 01` start code). Before the first NALU of every keyframe, prepends SPS+PPS in Annex-B form. +Also defines `pub struct AudioFrame { pub data: Bytes, pub timestamp_ms: u32 }` (distinct from `OpusAudioFrame`). -### Unused stub — `crates/server/src/rtmp.rs` +--- -Early manual RTMP handshake implementation, not used in the current flow. +### Audio — `crates/server/src/audio.rs` + +**`AACParser`** — parses raw RTMP `AudioDataReceived` payloads: +- Byte 0: codec (upper nibble, 10=AAC). +- Byte 1: AAC packet type — `0`=AudioSpecificConfig (codec init), `1`=raw AAC frame. +- Packet type `0`: initialises a **Symphonia** AAC decoder with the config bytes as `extra_data`. +- Packet type `1`: decodes via Symphonia, converts to interleaved f32 PCM, returns `AudioFrame { data, timestamp_ms, sample_rate }`. + +**`AudioProcesser`** — AAC→Opus transcoder: +- `encoder: opus::Encoder` — 48kHz stereo, `LowDelay` application mode. +- `resampler: rubato::FftFixedIn` — resamples 44100 Hz → 48000 Hz when needed. +- `pcm_buf: Vec` — accumulates samples until a full 960-sample (20 ms) Opus frame is ready. +- `samples_emitted: u64` — monotonic 48kHz counter; `timestamp_ms = samples_emitted / 48` (independent of RTMP timestamps). +- `encode(frame) -> Vec`: resamples if not already 48kHz, drains `pcm_buf` in 960-sample chunks, emits one `OpusAudioFrame` per chunk. + +`pub struct OpusAudioFrame { pub data: Bytes, pub timestamp_ms: u32 }` + +--- + +### Password hashing — `crates/server/src/hash.rs` + +- `hash_password(password: &str) -> Result` — Argon2id hash via `Argon2::default()` with a random `OsRng` salt; returns PHC-format string. +- `verify_password(password: &str, hash: &str) -> bool` — parses PHC string and verifies with Argon2. Panics if `hash` is not valid PHC format. + +--- + +### H.265 parsing — `crates/server/src/codec/h265.rs` + +`H265Parser` converts enhanced RTMP HEVC payloads → Annex-B `VideoFrame`s. + +**Enhanced RTMP detection** (`parse_video_codec` in `rtmp.rs`): +- Byte 0 bit 7 (`0x80`) set = ExVideoHeader (enhanced RTMP format) +- Bits 4–6 of byte 0 = frame type (1=keyframe) +- Bits 0–3 of byte 0 = packet type: `0`=SequenceStart, `1`=CodedFrames, `3`=CodedFramesX +- Bytes 1–4 = FourCC: `hvc1`=H.265, `avc1`=H.264, `av01`=AV1 + +**Payload offsets by packet type:** +- Type `0` (SequenceStart): payload at `data[5..]` — `HEVCDecoderConfigurationRecord` +- Type `1` (CodedFrames): payload at `data[8..]` — 3 bytes composition time skipped +- Type `3` (CodedFramesX): payload at `data[5..]` — no composition time + +**`HEVCDecoderConfigurationRecord`** parsing: +- Skip first 22 bytes (profile/level/tier info, not needed for forwarding) +- Byte 22 = `numOfArrays`; each array: 1 byte `nal_unit_type` (lower 6 bits) + 2 byte NALU count + length-prefixed NALUs +- NAL types: VPS=32, SPS=33, PPS=34 — cached as `Vec` on the parser struct + +**HVCC → Annex-B conversion:** +- Same as AVCC: replace 4-byte big-endian length prefix with `00 00 00 01` start code +- Prepend VPS+SPS+PPS (each with start code) before every keyframe + +--- ## Database (SeaORM + SQLite) -DB file: `stream.db`. Migrations: `crates/migration/`. Entities: `crates/entity/`. +DB file: `./db/db.sqlite`. Migrations: `crates/migration/`. Entities: `crates/entity/`. -Run migrations: `Migrator::up(&db, None).await?` — idempotent, tracked in `seaql_migrations`. +Run migrations: `Migrator::up(&db, None).await` — idempotent, tracked in `seaql_migrations`. ### Entities @@ -112,22 +260,30 @@ Run migrations: `Migrator::up(&db, None).await?` — idempotent, tracked in `sea - Fields: `id`, `username`, `hashed_password` - Relations: `has_many` → `stream_key` - `Entity::create(db, username, hashed_password)` — inserts a new user -- `ActiveModel::update_username(db, username)` — updates username -- `ActiveModel::update_password(db, hashed_password)` — updates password hash -- Passwords must be hashed before being passed to these methods +- `ActiveModel::update_username` / `update_password` — mutation helpers +- Passwords must be hashed via `hash::hash_password` (Argon2id) before storing **`stream_key`** (`crates/entity/src/stream_key.rs`) - Fields: `id`, `key_value` (unique), `user_id`, `label`, `is_active`, `is_unlisted`, `created_at` - Relations: `belongs_to` → `users`, `has_many` → `stream_session` +- `Entity::find_by_key(db, key_value)` — lookup by raw stream key string **`stream_session`** (`crates/entity/src/stream_session.rs`) - Fields: `id`, `stream_key_id`, `started_at`, `ended_at` (nullable) - Relations: `belongs_to` → `stream_key` +- `Model::get_active_by_stream_key_id(db, id)` — finds open session (no `ended_at`) +- `Model::clean_unended_streams(db)` — sets `ended_at = now` on all sessions missing it (crash recovery) + +**`auth_session`** (`crates/entity/src/auth_session.rs`) +- Fields: `id`, `user_id`, `token`, `created_at` +- Used for cookie-based auth; token matched against `session` header/cookie ### SeaORM conventions - Query methods go on `Entity` (e.g. `Entity::find_by_x`). -- Mutation helpers that intercept save logic (e.g. setting timestamps, pre-save transforms) go on `ActiveModel`. -- For destructive schema changes in prod, use expand-contract: add new structure → backfill → switch app code → drop old structure in a later migration. +- Mutation helpers that intercept save logic go on `ActiveModel`. +- For destructive schema changes in prod, use expand-contract: add → backfill → switch code → drop old in a later migration. -**Test page:** `index.html` — open in a browser to view the stream via WHEP without any extra tooling. +--- + +**Test page:** `index.html` — open in a browser to play the stream via WHEP without extra tooling. diff --git a/Cargo.toml b/Cargo.toml index 0a31ca9..0dd9592 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,12 @@ inherits = "release" debug = true force-frame-pointers = true +[profile.release] +lto = true +codegen-units = 1 +panic = "abort" +# codegen-backend = "cranelift" + [workspace.metadata.crane] name = "rtmp-to-whip" diff --git a/README.md b/README.md index cc02de1..3b3e0b8 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,10 @@ # Simple Instant Stream - [ ] axum cookie jar support (manual set-cookie headers isnt yum) -- [ ] more codecs - - [ ] h265 - - [ ] av1 +- [X] more codecs + - [X] h265 + - [X] av1 +- [ ] Stream status rtc data channel (i.e, 'you're dropping frames! ,, 'wrong codec!') which is shown in the stream player page - [ ] Whep -> Whip support - [ ] Write customs error types for everything that touches HTTP for [thiserror](https://docs.rs/thiserror/latest/thiserror/) and make it impl axum's IntoReponce - [ ] HTTP diff --git a/crates/server/src/codec/av1.rs b/crates/server/src/codec/av1.rs new file mode 100644 index 0000000..62cbc70 --- /dev/null +++ b/crates/server/src/codec/av1.rs @@ -0,0 +1,138 @@ +use bytes::Bytes; +use tracing::debug; + +use crate::codec::{CodecParser, VideoFrame}; + +pub struct Av1CodecParser { + // Raw OBU bytes from AV1CodecConfigurationRecord configOBUs field + // (sequence header + optional metadata, low-overhead/ISOBMFF format with size fields). + // Prepended to keyframes so str0m's Av1Packetizer sees a Sequence Header OBU + // and sets the N bit in the RTP aggregation header. + config_obus: Option>, + // True until the first coded frame is emitted after a sequence header arrives. + // OBS AV1 may not set FrameType=1 in the RTMP header for keyframes, so we + // bootstrap the decoder by treating the first post-config frame as a keyframe. + first_coded_frame: bool, +} + +impl Av1CodecParser { + pub fn new() -> Self { + Self { + config_obus: None, + first_coded_frame: true, + } + } + + // AV1CodecConfigurationRecord layout (ISOBMFF AV1 spec §2.3.3): + // [0] marker(1) | version(7) — always 0x81 + // [1] seq_profile(3) | seq_level_idx_0(5) + // [2] seq_tier_0(1) | high_bitdepth(1) | twelve_bit(1) | monochrome(1) | chroma_subsampling_x/y(2) | chroma_sample_position(2) + // [3] reserved(3) | initial_presentation_delay_present(1) | initial_presentation_delay_minus_one or reserved(4) + // [4..] configOBUs — sequence header OBU + optional metadata OBUs + fn parse_sequence_header(&mut self, payload: &[u8]) { + if payload.len() < 4 { + return; + } + self.config_obus = Some(payload[4..].to_vec()); + self.first_coded_frame = true; + debug!(bytes = payload.len() - 4, "AV1 sequence header stored"); + } + + // Scan low-overhead-format OBUs to see if any is a SEQUENCE_HEADER (obu_type=1). + // This detects keyframes when OBS inlines the sequence header into the coded frame + // rather than signalling it via FrameType=1 in the RTMP header. + fn obus_contain_sequence_header(data: &[u8]) -> bool { + let mut i = 0; + while i < data.len() { + let header = data[i]; + let obu_type = (header >> 3) & 0x1F; + let has_extension = (header >> 2) & 1 != 0; + let has_size = (header >> 1) & 1 != 0; + i += 1; + if has_extension { + if i >= data.len() { return false; } + i += 1; + } + if obu_type == 1 { + return true; // OBU_SEQUENCE_HEADER + } + if has_size { + // LEB128 decode + let mut size: usize = 0; + let mut shift = 0; + loop { + if i >= data.len() { return false; } + let b = data[i] as usize; + i += 1; + size |= (b & 0x7F) << shift; + shift += 7; + if b & 0x80 == 0 { break; } + if shift > 32 { return false; } + } + i += size; + } else { + // No size field means this OBU spans to end of data; can't advance further. + return false; + } + } + false + } + + fn obus_for_frame(&mut self, payload: &[u8], rtmp_is_keyframe: bool) -> Option<(Vec, bool)> { + if payload.is_empty() { + return None; + } + + // Determine keyframe status: + // 1. RTMP FrameType=1 (reliable when OBS sets it) + // 2. OBU scan: sequence header OBU inline in the coded payload (some encoders do this) + // 3. First coded frame after receiving the config record (bootstrap fallback — OBS AV1 + // may never set FrameType=1, so we start the decoder from the first available frame) + let obu_is_keyframe = Self::obus_contain_sequence_header(payload); + let is_keyframe = rtmp_is_keyframe + || obu_is_keyframe + || (self.first_coded_frame && self.config_obus.is_some()); + + self.first_coded_frame = false; + + if is_keyframe { + if let Some(config) = &self.config_obus { + let mut out = Vec::with_capacity(config.len() + payload.len()); + out.extend_from_slice(config); + out.extend_from_slice(payload); + return Some((out, true)); + } + } + + Some((payload.to_vec(), is_keyframe)) + } +} + +impl CodecParser for Av1CodecParser { + fn parse(&mut self, data: &[u8], timestamp_ms: u32) -> Option { + // AV1 only arrives via enhanced RTMP (bit 7 set, FourCC "av01") + if data.len() < 5 || data[0] & 0x80 == 0 { + return None; + } + + let rtmp_is_keyframe = (data[0] >> 4) & 0x07 == 1; + let packet_type = data[0] & 0x0F; + + match packet_type { + 0 => { + // SequenceStart: data[5..] = AV1CodecConfigurationRecord + self.parse_sequence_header(data.get(5..)?); + None + } + 1 | 3 => { + // CodedFrames / CodedFramesX: data[5..] = OBUs. + // Unlike HEVC, AV1 has no 3-byte composition-time field after the + // FourCC — enhanced RTMP only defines CTS for hvc1 CodedFrames. + let payload = data.get(5..)?; + let (obus, is_keyframe) = self.obus_for_frame(payload, rtmp_is_keyframe)?; + Some(VideoFrame { data: Bytes::from(obus), is_keyframe, timestamp_ms }) + } + _ => None, + } + } +} diff --git a/crates/server/src/codec/flv_replay_test.rs b/crates/server/src/codec/flv_replay_test.rs new file mode 100644 index 0000000..fe7003c --- /dev/null +++ b/crates/server/src/codec/flv_replay_test.rs @@ -0,0 +1,95 @@ +//! Temporary harness: replays video tags from an FLV file through a CodecParser +//! and dumps the resulting elementary stream so it can be validated with ffprobe. +//! Driven by env vars: FLV_IN (input .flv), ES_OUT (output elementary stream). +//! Run: FLV_IN=... ES_OUT=... cargo test -p server flv_replay -- --ignored --nocapture + +use super::{CodecParser, av1::Av1CodecParser, h264::H264CodecParser, h265::H265CodecParser}; + +fn parse_flv(path: &str) -> Vec<(u32, Vec)> { + let buf = std::fs::read(path).unwrap(); + assert_eq!(&buf[0..3], b"FLV"); + let mut tags = Vec::new(); + let mut i = 9 + 4; // header + first prev-tag-size + while i + 11 <= buf.len() { + let tag_type = buf[i]; + let size = u32::from_be_bytes([0, buf[i + 1], buf[i + 2], buf[i + 3]]) as usize; + let ts = u32::from_be_bytes([buf[i + 7], buf[i + 4], buf[i + 5], buf[i + 6]]); + let data_start = i + 11; + if data_start + size > buf.len() { + break; + } + if tag_type == 9 { + tags.push((ts, buf[data_start..data_start + size].to_vec())); + } + i = data_start + size + 4; + } + tags +} + +#[test] +#[ignore] +fn flv_replay() { + let flv_in = std::env::var("FLV_IN").unwrap(); + let es_out = std::env::var("ES_OUT").unwrap(); + + let tags = parse_flv(&flv_in); + assert!(!tags.is_empty(), "no video tags found"); + + // Detect codec from first tag + let first = &tags[0].1; + let mut parser: Box = if first[0] & 0x80 != 0 { + match &first[1..5] { + b"hvc1" => Box::new(H265CodecParser::new()), + b"av01" => Box::new(Av1CodecParser::new()), + b"avc1" => Box::new(H264CodecParser::new()), + other => panic!("unknown fourcc {:?}", other), + } + } else { + Box::new(H264CodecParser::new()) + }; + println!( + "first tag byte0={:#x} fourcc={:?}", + first[0], + String::from_utf8_lossy(&first[1..5.min(first.len())]) + ); + + let mut es = Vec::new(); + let mut n_frames = 0; + let mut n_keyframes = 0; + let mut pkt_type_histogram = std::collections::BTreeMap::new(); + let mut out_ts = Vec::new(); + for (ts, tag) in &tags { + let pt = if tag[0] & 0x80 != 0 { + tag[0] & 0x0F + } else { + tag[1] + }; + *pkt_type_histogram.entry(pt).or_insert(0u32) += 1; + if let Some(frame) = parser.parse(tag, *ts) { + n_frames += 1; + if frame.is_keyframe { + n_keyframes += 1; + } + out_ts.push((*ts, frame.timestamp_ms)); + // Temporal delimiter OBU so ffmpeg's obu demuxer can frame the + // stream (not sent over RTP; only needed for offline validation). + es.extend_from_slice(&[0x12, 0x00]); + es.extend_from_slice(&frame.data); + } + } + println!( + "tags={} frames={} keyframes={} pkt_types={:?}", + tags.len(), + n_frames, + n_keyframes, + pkt_type_histogram + ); + println!("first 12 (dts, pts): {:?}", &out_ts[..12.min(out_ts.len())]); + let mut pts_sorted: Vec = out_ts.iter().map(|(_, p)| *p).collect(); + pts_sorted.sort_unstable(); + pts_sorted.dedup(); + assert_eq!(pts_sorted.len(), out_ts.len(), "duplicate PTS values"); + std::fs::write(&es_out, &es).unwrap(); + assert!(n_frames > 0, "no frames produced"); + assert!(n_keyframes > 0, "no keyframes detected"); +} diff --git a/crates/server/src/codec/h264.rs b/crates/server/src/codec/h264.rs new file mode 100644 index 0000000..ddc8a57 --- /dev/null +++ b/crates/server/src/codec/h264.rs @@ -0,0 +1,175 @@ +use bytes::Bytes; + +use crate::codec::{CodecParser, VideoFrame}; + +pub struct H264CodecParser { + sps: Option>, + pps: Option>, +} + +impl CodecParser for H264CodecParser { + fn parse(&mut self, data: &[u8], timestamp_ms: u32) -> Option { + self.parse_inner(data, timestamp_ms) + } +} + +impl H264CodecParser { + pub fn new() -> Self { + Self { + sps: None, + pps: None, + } + } + + /// Parse an RTMP VideoDataReceived payload. Handles both legacy and enhanced RTMP. + /// Returns None for sequence header packets (no displayable frame). + fn parse_inner(&mut self, bytes: &[u8], timestamp_ms: u32) -> Option { + if bytes.len() < 2 { + return None; + } + + // Bit 7 of byte 0 distinguishes enhanced RTMP from legacy. + let is_keyframe = (bytes[0] >> 4) & 0x07 == 1; + + if bytes[0] & 0x80 != 0 { + // Enhanced RTMP: byte 0 = 0x80 | (frame_type << 4) | packet_type + // bytes 1-4 = FourCC + if bytes.len() < 5 { + return None; + } + let packet_type = bytes[0] & 0x0F; + match packet_type { + 0 => { + // SequenceStart: bytes[5..] = AVCDecoderConfigurationRecord + self.parse_sequence_header(bytes.get(5..)?); + None + } + 1 => { + // CodedFrames: bytes 5-7 = SI24 composition time offset, bytes 8+ = AVCC NALUs. + // RTP timestamps must be presentation time (RFC 6184), so emit + // PTS = DTS + CTS; stamping DTS makes B-frame streams stutter. + if bytes.len() < 8 { + return None; + } + let pts_ms = Self::pts_ms(timestamp_ms, &bytes[5..8]); + let data = self.avcc_to_annexb(bytes.get(8..)?, is_keyframe)?; + Some(VideoFrame { data: Bytes::from(data), is_keyframe, timestamp_ms: pts_ms }) + } + 3 => { + // CodedFramesX: no CTS field, bytes 5+ = AVCC NALUs + let data = self.avcc_to_annexb(bytes.get(5..)?, is_keyframe)?; + Some(VideoFrame { data: Bytes::from(data), is_keyframe, timestamp_ms }) + } + _ => None, + } + } else { + // Legacy RTMP: byte 0 = (frame_type << 4) | codec_id (7 = H.264) + // byte 1 = AVC packet type, bytes 2-4 = CTS, bytes 5+ = AVCC + if bytes[0] & 0x0F != 7 || bytes.len() < 5 { + return None; + } + match bytes[1] { + 0 => { + self.parse_sequence_header(&bytes[5..]); + None + } + 1 => { + // PTS = DTS + CTS — see note on CodedFrames above. + let pts_ms = Self::pts_ms(timestamp_ms, &bytes[2..5]); + let data = self.avcc_to_annexb(&bytes[5..], is_keyframe)?; + Some(VideoFrame { data: Bytes::from(data), is_keyframe, timestamp_ms: pts_ms }) + } + _ => None, + } + } + } + + // Sign-extend a 3-byte big-endian composition time offset and add it to the DTS. + fn pts_ms(timestamp_ms: u32, cts_bytes: &[u8]) -> u32 { + let cts = i32::from_be_bytes([0, cts_bytes[0], cts_bytes[1], cts_bytes[2]]) << 8 >> 8; + (timestamp_ms as i64 + cts as i64).max(0) as u32 + } + + fn parse_sequence_header(&mut self, payload: &[u8]) { + // AVCDecoderConfigurationRecord layout: + // [0] configurationVersion + // [1] AVCProfileIndication + // [2] profile_compatibility + // [3] AVCLevelIndication + // [4] 0xFF (lower 2 bits = lengthSizeMinusOne, always 3 meaning 4-byte lengths) + // [5] 0xE0 | numSPS + // [6..] SPS entries: 2-byte length + bytes + // then: numPPS, PPS entries: 2-byte length + bytes + if payload.len() < 7 { + return; + } + + let mut i = 5; + + let num_sps = (payload[i] & 0x1F) as usize; + i += 1; + + for _ in 0..num_sps { + if i + 2 > payload.len() { + return; + } + let len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize; + i += 2; + if i + len > payload.len() { + return; + } + self.sps = Some(payload[i..i + len].to_vec()); + i += len; + } + + if i >= payload.len() { + return; + } + + let num_pps = payload[i] as usize; + i += 1; + + for _ in 0..num_pps { + if i + 2 > payload.len() { + return; + } + let len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize; + i += 2; + if i + len > payload.len() { + return; + } + self.pps = Some(payload[i..i + len].to_vec()); + i += len; + } + } + + fn avcc_to_annexb(&self, payload: &[u8], is_keyframe: bool) -> Option> { + let mut out = Vec::new(); + + // Prepend SPS+PPS before every keyframe so str0m's packetizer + // can bundle them into a STAP-A alongside the IDR NALU. + if is_keyframe { + if let (Some(sps), Some(pps)) = (&self.sps, &self.pps) { + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(sps); + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(pps); + } + } + + // Convert each length-prefixed NALU to an Annex B start-code NALU. + let mut i = 0; + while i + 4 <= payload.len() { + let nalu_len = u32::from_be_bytes(payload[i..i + 4].try_into().unwrap()) as usize; + i += 4; + if i + nalu_len > payload.len() { + break; + } + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(&payload[i..i + nalu_len]); + i += nalu_len; + } + + if out.is_empty() { None } else { Some(out) } + } +} diff --git a/crates/server/src/codec/h265.rs b/crates/server/src/codec/h265.rs new file mode 100644 index 0000000..a475d20 --- /dev/null +++ b/crates/server/src/codec/h265.rs @@ -0,0 +1,142 @@ +use bytes::Bytes; + +use crate::codec::{CodecParser, VideoFrame}; + +pub struct H265CodecParser { + vps: Option>, + sps: Option>, + pps: Option>, +} + +impl H265CodecParser { + pub fn new() -> Self { + Self { + vps: None, + sps: None, + pps: None, + } + } + + // HEVCDecoderConfigurationRecord layout (ISO 14496-15 §8.3.3): + // [0] configurationVersion (always 1) + // [1..2] general_profile_space(2b) | general_tier_flag(1b) | general_profile_idc(5b) + // [2..6] general_profile_compatibility_flags + // [6..12] general_constraint_indicator_flags + // [12] general_level_idc + // [13..15] min_spatial_segmentation_idc (lower 12 bits) + // [15] parallelismType (lower 2 bits) + // [16] chroma_format_idc (lower 2 bits) + // [17] bit_depth_luma_minus8 (lower 3 bits) + // [18] bit_depth_chroma_minus8 (lower 3 bits) + // [19..21] avgFrameRate + // [21] constantFrameRate(2b) | numTemporalLayers(3b) | temporalIdNested(1b) | lengthSizeMinusOne(2b) + // [22] numOfArrays + // [23..] arrays: [ array_completeness(1b) | reserved(1b) | NAL_unit_type(6b), numNalus(2b), + // [ naluLength(2b), nalu(naluLength) ] ] + fn parse_sequence_header(&mut self, payload: &[u8]) { + if payload.len() < 23 { + return; + } + + let num_arrays = payload[22] as usize; + let mut i = 23; + + for _ in 0..num_arrays { + if i + 3 > payload.len() { + return; + } + let nal_type = payload[i] & 0x3F; + let num_nalus = u16::from_be_bytes([payload[i + 1], payload[i + 2]]) as usize; + i += 3; + + for _ in 0..num_nalus { + if i + 2 > payload.len() { + return; + } + let nalu_len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize; + i += 2; + if i + nalu_len > payload.len() { + return; + } + + match nal_type { + 32 => self.vps = Some(payload[i..i + nalu_len].to_vec()), + 33 => self.sps = Some(payload[i..i + nalu_len].to_vec()), + 34 => self.pps = Some(payload[i..i + nalu_len].to_vec()), + _ => {} + } + + i += nalu_len; + } + } + } + + fn hvcc_to_annexb(&self, payload: &[u8], is_keyframe: bool) -> Option> { + let mut out = Vec::with_capacity(payload.len()); + + if is_keyframe { + if let (Some(vps), Some(sps), Some(pps)) = (&self.vps, &self.sps, &self.pps) { + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(vps); + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(sps); + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(pps); + } + } + + let mut i = 0; + while i + 4 <= payload.len() { + let nalu_len = u32::from_be_bytes(payload[i..i + 4].try_into().unwrap()) as usize; + i += 4; + if i + nalu_len > payload.len() { + break; + } + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(&payload[i..i + nalu_len]); + i += nalu_len; + } + + if out.is_empty() { None } else { Some(out) } + } +} + +impl CodecParser for H265CodecParser { + fn parse(&mut self, data: &[u8], timestamp_ms: u32) -> Option { + // H.265 only comes via enhanced RTMP (bit 7 set, FourCC "hvc1") + if data.len() < 5 || data[0] & 0x80 == 0 { + return None; + } + + let is_keyframe = (data[0] >> 4) & 0x07 == 1; + let packet_type = data[0] & 0x0F; + // bytes 1-4 are FourCC "hvc1" — already validated by rtmp.rs + + match packet_type { + 0 => { + // SequenceStart: bytes[5..] = HEVCDecoderConfigurationRecord + self.parse_sequence_header(data.get(5..)?); + None + } + 1 => { + // CodedFrames: bytes 5-7 = SI24 composition time offset, bytes 8+ = HVCC. + // RTP timestamps must be presentation time (RFC 7798), so emit + // PTS = DTS + CTS. Encoders use this packet type exactly when CTS != 0 + // (B-frames present); stamping DTS instead makes playback stutter. + if data.len() < 8 { + return None; + } + let cts = i32::from_be_bytes([0, data[5], data[6], data[7]]) << 8 >> 8; + let pts_ms = (timestamp_ms as i64 + cts as i64).max(0) as u32; + let annexb = self.hvcc_to_annexb(data.get(8..)?, is_keyframe)?; + Some(VideoFrame { data: Bytes::from(annexb), is_keyframe, timestamp_ms: pts_ms }) + } + 3 => { + // CodedFramesX: no CTS, bytes 5+ = HVCC + let annexb = self.hvcc_to_annexb(data.get(5..)?, is_keyframe)?; + Some(VideoFrame { data: Bytes::from(annexb), is_keyframe, timestamp_ms }) + } + _ => None, + } + } +} diff --git a/crates/server/src/codec/mod.rs b/crates/server/src/codec/mod.rs new file mode 100644 index 0000000..eb6cf9f --- /dev/null +++ b/crates/server/src/codec/mod.rs @@ -0,0 +1,18 @@ +use bytes::Bytes; + +pub mod av1; +pub mod h264; +pub mod h265; + +#[cfg(test)] +mod flv_replay_test; + +pub trait CodecParser: Send { + fn parse(&mut self, data: &[u8], timestamp_ms: u32) -> Option; +} + +pub struct VideoFrame { + pub data: Bytes, + pub is_keyframe: bool, + pub timestamp_ms: u32, +} diff --git a/crates/server/src/http.rs b/crates/server/src/http.rs index cd81b40..30d7d5e 100644 --- a/crates/server/src/http.rs +++ b/crates/server/src/http.rs @@ -1,4 +1,10 @@ -use std::{net::SocketAddr, sync::Arc}; +use std::{ + net::SocketAddr, + sync::{ + Arc, + atomic::{AtomicI32, Ordering}, + }, +}; use axum::{ Json, Router, @@ -41,13 +47,13 @@ pub struct HttpServer { pub offer_tx: Sender<(i32, i32, String)>, pub accept_rx: async_broadcast::InactiveReceiver<(i32, Option)>, pub appstate: Arc>, - pub request_count: Mutex, + pub request_count: AtomicI32, pub db: DatabaseConnection, - pub config: HttpServerConfig, + pub config: Arc, } impl HttpServer { - pub fn start(self) -> Result<(), Box> { + pub async fn run(self) { let state = Arc::new(self); let origins = [ @@ -66,27 +72,29 @@ impl HttpServer { .allow_credentials(true) .allow_origin(origins); - tokio::spawn(async move { - let app = Router::new() - .route("/api/catalog", get(catalog_handler)) - .route("/api/user", post(create_user_handler)) - .route( - "/api/stream-key", - post(create_stream_key_handler).get(get_all_stream_keys), - ) - // .route("/api/admin/server_stats", get(todo!())) - .route("/api/whip", post(handle_whip_injest)) - .route("/api/login", post(login_handler)) - .route("/api/stream/{slug}", post(stream_handler)) - .route("/api/meow", get(meow_handler)) - .layer(cors) - .with_state(state); + // TODO: Add an error type that impls IntoResponse, enum HttpError {}; impl IntoResponse for + // HttpError + // https://docs.rs/thiserror/latest/thiserror/ + // + // Its rather shitty + let app = Router::new() + .route("/api/catalog", get(catalog_handler)) + .route("/api/user", post(create_user_handler)) + .route( + "/api/stream-key", + post(create_stream_key_handler).get(get_all_stream_keys), + ) + // .route("/api/admin/server_stats", get(todo!())) + .route("/api/whip", post(handle_whip_injest)) + .route("/api/login", post(login_handler)) + .route("/api/stream/{slug}", post(stream_handler)) + .route("/api/meow", get(meow_handler)) + .layer(cors) + .with_state(state); - let addr = SocketAddr::from(([0, 0, 0, 0], 3000)); - let listener = TcpListener::bind(addr).await.unwrap(); - axum::serve(listener, app).await.unwrap(); - }); - Ok(()) + let addr = SocketAddr::from(([0, 0, 0, 0], 3000)); + let listener = TcpListener::bind(addr).await.unwrap(); + axum::serve(listener, app).await.unwrap(); } } @@ -263,6 +271,7 @@ async fn login_handler( let auth = auth_session::Entity::create(&state.db, x.id).await.unwrap(); let token = auth.value; let mut meow = Response::new("".to_string()); + // TODO: Replace with axum cookie jar https://docs.rs/axum-extra/latest/axum_extra/extract/cookie/struct.CookieJar.html meow.headers_mut().insert( SET_COOKIE, format!("session={token}; SameSite=Strict; Path=/; Max-Age=2592000") @@ -324,6 +333,7 @@ async fn create_user_handler( }; let token = session.value; let mut headers = HeaderMap::new(); + // TODO: Replace with axum cookie jar https://docs.rs/axum-extra/latest/axum_extra/extract/cookie/struct.CookieJar.html headers.insert( SET_COOKIE, format!("session={token}; HttpOnly; SameSite=Strict; Path=/") @@ -338,10 +348,7 @@ async fn stream_handler( Path(slug): Path, body: String, ) -> impl IntoResponse { - let mut request_id = state.request_count.lock().await; - *request_id += 1; - let request_id_clone = *request_id; - drop(request_id); + let request_id_clone = state.request_count.fetch_add(1, Ordering::Relaxed); let stream_key_id = { let app = state.appstate.lock().await; diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index f57bf9f..95b4910 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -17,21 +17,22 @@ use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::TcpListener, sync::Mutex, + task::JoinSet, time::Instant, }; use tracing_subscriber::EnvFilter; use crate::{ audio::OpusAudioFrame, + codec::VideoFrame, http::{HttpServer, HttpServerConfig}, - media::{H264Parser, VideoFrame}, webrtc_proxy::WebRtcProxyConfig, }; mod audio; +mod codec; mod hash; mod http; -mod media; mod rtmp; mod webrtc; mod webrtc_ingest; @@ -42,11 +43,20 @@ pub struct AppState { pub stream_sessions: Arc>, } +#[derive(Clone, Debug)] +pub enum StreamCodec { + H264, + H265, + AV1, // Unsupported +} + pub struct StreamSession { pub stream_key_id: i32, pub stream_key_label: String, pub frame_channel: async_broadcast::Sender>, pub audio_channel: async_broadcast::Sender>, + // + pub codec: Option, } #[tokio::main] @@ -65,7 +75,7 @@ async fn main() -> Result<(), Box> { let db = Database::connect("sqlite://./db/db.sqlite?mode=rwc") .await - .unwrap(); + .expect("Can't make or read ./db/db.sqlite (create db dir maybe)"); info!("database connected"); Migrator::up(&db, None).await.unwrap(); info!("migrations complete"); @@ -90,17 +100,16 @@ async fn main() -> Result<(), Box> { offer_tx, accept_rx: answer_rx.deactivate(), appstate: appstate.clone(), - request_count: Mutex::new(0), + request_count: std::sync::atomic::AtomicI32::new(0), db: db.clone(), config: HttpServerConfig { signup_code: env::var("SIGNUP_CODE").unwrap_or_else(|_| { warn!("SIGNUP_CODE not set; signup will be disabled"); String::new() }), - }, + } + .into(), }; - http.start()?; - let proxyconfig = WebRtcProxyConfig { proxy_port: env::var("RTC_PORT") .unwrap_or("6969".into()) @@ -108,7 +117,6 @@ async fn main() -> Result<(), Box> { .expect("RTC_PORT needs to be a number (i32)"), }; let proxy = webrtc_proxy::WebrtcProxy::new(proxyconfig).await.unwrap(); - proxy.start().unwrap(); let app = appstate.lock().await; let webrtc = webrtc::Webrtc { @@ -116,19 +124,36 @@ async fn main() -> Result<(), Box> { accept_tx: answer_tx, sessions_ref: app.stream_sessions.clone(), db: db.clone(), - proxy: proxy.into(), + proxy: proxy.clone().into(), }; - webrtc.start()?; let rtmp = rtmp::Rtmp { db: db.clone(), stream_sessions: app.stream_sessions.clone(), - listener: listener, + listener, }; - rtmp.start()?; drop(app); - tokio::signal::ctrl_c().await?; + let mut workers: JoinSet<()> = JoinSet::new(); + workers.spawn(http.run()); + workers.spawn(proxy.run()); + workers.spawn(webrtc.run()); + workers.spawn(rtmp.run()); + + tokio::select! { + _ = tokio::signal::ctrl_c() => { + info!("received ctrl-c, shutting down"); + } + res = workers.join_next() => { + if let Some(Err(e)) = res { + tracing::error!("worker panicked: {:?}", e); + } else { + tracing::error!("a worker exited unexpectedly"); + } + } + } + + workers.abort_all(); Ok(()) } diff --git a/crates/server/src/media.rs b/crates/server/src/media.rs deleted file mode 100644 index 98ef41e..0000000 --- a/crates/server/src/media.rs +++ /dev/null @@ -1,147 +0,0 @@ -// Claude slop... im not skilled amount to do this bullshit - -use bytes::Bytes; - -pub struct VideoFrame { - pub data: Bytes, - pub is_keyframe: bool, - pub timestamp_ms: u32, -} - -pub struct AudioFrame { - pub data: Bytes, - pub timestamp_ms: u32, -} - -pub struct H264Parser { - sps: Option>, - pps: Option>, -} - -impl H264Parser { - pub fn new() -> Self { - Self { - sps: None, - pps: None, - } - } - - /// Parse an RTMP VideoDataReceived payload. Returns None for sequence - /// header packets (which carry SPS/PPS but no displayable frame). - pub fn parse(&mut self, bytes: &[u8], timestamp_ms: u32) -> Option { - if bytes.len() < 5 { - return None; - } - - let frame_type = (bytes[0] >> 4) & 0x0F; - let codec_id = bytes[0] & 0x0F; - - if codec_id != 7 { - return None; // not H.264 - } - - let avc_packet_type = bytes[1]; - // bytes[2..5] are the composition time offset — not needed for sending - let payload = &bytes[5..]; - - match avc_packet_type { - 0 => { - self.parse_sequence_header(payload); - None - } - 1 => { - let is_keyframe = frame_type == 1; - let data = self.avcc_to_annexb(payload, is_keyframe)?; - Some(VideoFrame { - data: Bytes::from(data), - is_keyframe, - timestamp_ms, - }) - } - _ => None, - } - } - - fn parse_sequence_header(&mut self, payload: &[u8]) { - // AVCDecoderConfigurationRecord layout: - // [0] configurationVersion - // [1] AVCProfileIndication - // [2] profile_compatibility - // [3] AVCLevelIndication - // [4] 0xFF (lower 2 bits = lengthSizeMinusOne, always 3 meaning 4-byte lengths) - // [5] 0xE0 | numSPS - // [6..] SPS entries: 2-byte length + bytes - // then: numPPS, PPS entries: 2-byte length + bytes - if payload.len() < 7 { - return; - } - - let mut i = 5; - - let num_sps = (payload[i] & 0x1F) as usize; - i += 1; - - for _ in 0..num_sps { - if i + 2 > payload.len() { - return; - } - let len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize; - i += 2; - if i + len > payload.len() { - return; - } - self.sps = Some(payload[i..i + len].to_vec()); - i += len; - } - - if i >= payload.len() { - return; - } - - let num_pps = payload[i] as usize; - i += 1; - - for _ in 0..num_pps { - if i + 2 > payload.len() { - return; - } - let len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize; - i += 2; - if i + len > payload.len() { - return; - } - self.pps = Some(payload[i..i + len].to_vec()); - i += len; - } - } - - fn avcc_to_annexb(&self, payload: &[u8], is_keyframe: bool) -> Option> { - let mut out = Vec::new(); - - // Prepend SPS+PPS before every keyframe so str0m's packetizer - // can bundle them into a STAP-A alongside the IDR NALU. - if is_keyframe { - if let (Some(sps), Some(pps)) = (&self.sps, &self.pps) { - out.extend_from_slice(&[0, 0, 0, 1]); - out.extend_from_slice(sps); - out.extend_from_slice(&[0, 0, 0, 1]); - out.extend_from_slice(pps); - } - } - - // Convert each length-prefixed NALU to an Annex B start-code NALU. - let mut i = 0; - while i + 4 <= payload.len() { - let nalu_len = u32::from_be_bytes(payload[i..i + 4].try_into().unwrap()) as usize; - i += 4; - if i + nalu_len > payload.len() { - break; - } - out.extend_from_slice(&[0, 0, 0, 1]); - out.extend_from_slice(&payload[i..i + nalu_len]); - i += nalu_len; - } - - if out.is_empty() { None } else { Some(out) } - } -} diff --git a/crates/server/src/rtmp.rs b/crates/server/src/rtmp.rs index e57d172..892c434 100644 --- a/crates/server/src/rtmp.rs +++ b/crates/server/src/rtmp.rs @@ -15,9 +15,14 @@ use tokio::{ use tracing::{debug, info, warn}; use crate::{ - StreamSession, + StreamCodec, StreamSession, audio::{AACParser, AudioProcesser, OpusAudioFrame}, - media::{H264Parser, VideoFrame}, + codec::{ + CodecParser, VideoFrame, + av1::Av1CodecParser, + h264::H264CodecParser, + h265::H265CodecParser, + }, }; pub struct Rtmp { @@ -35,6 +40,30 @@ async fn write_outbound(socket: &mut TcpStream, results: Vec Result> { + if payload.is_empty() { + return Err("empty video payload".into()); + } + + let is_ex = payload[0] & 0x80 != 0; + + if is_ex { + if payload.len() < 5 { + return Err("enhanced RTMP payload too short".into()); + } + match &payload[1..5] { + b"hvc1" => Ok(StreamCodec::H265), + b"avc1" => Ok(StreamCodec::H264), + b"av01" => Ok(StreamCodec::AV1), + _ => Err("unsupported FourCC".into()), + } + } else { + match payload[0] & 0x0F { + 7 => Ok(StreamCodec::H264), + _ => Err("unsupported legacy codec ID".into()), + } + } + } async fn handshake( mut socket: TcpStream, ) -> Result<(ServerSession, TcpStream), Box> { @@ -64,139 +93,129 @@ impl Rtmp { Ok((rtmp_session, socket)) } - pub fn start(self) -> Result<(), Box> { + pub async fn run(self) { let Self { listener, stream_sessions, db, } = self; - tokio::spawn(async move { - loop { - let (socket, peer_addr) = listener.accept().await.unwrap(); - info!(%peer_addr, "RTMP connection accepted"); - let (mut session, mut socket) = Rtmp::handshake(socket).await.unwrap(); - info!(%peer_addr, "RTMP handshake complete"); - let (mut video_tx, mut video_rx) = broadcast::>(32); - let (mut audio_tx, mut audio_rx) = broadcast::>(32); - // video_rx.cycle + loop { + let (socket, peer_addr) = listener.accept().await.unwrap(); + info!(%peer_addr, "RTMP connection accepted"); + let (mut session, mut socket) = Rtmp::handshake(socket).await.unwrap(); + info!(%peer_addr, "RTMP handshake complete"); + let (mut video_tx, mut video_rx) = broadcast::>(32); + let (mut audio_tx, mut audio_rx) = broadcast::>(32); + // video_rx.cycle - video_tx.set_overflow(true); - audio_tx.set_overflow(true); - let mut parser = H264Parser::new(); - let mut aac_parser = AACParser::new(); - let mut audio_proc = AudioProcesser::new(); - let db = db.clone(); - let stream_sessions = stream_sessions.clone(); + video_tx.set_overflow(true); + audio_tx.set_overflow(true); + let mut parser: Option> = None; + let mut aac_parser = AACParser::new(); + let mut audio_proc = AudioProcesser::new(); + let db = db.clone(); + let stream_sessions = stream_sessions.clone(); - tokio::spawn(async move { - loop { - let mut buf = [0u8; 4096]; - let n = socket.read(&mut buf).await.unwrap(); - if n == 0 { - debug!("RTMP connection closed by peer"); - return; - } - let events = session.handle_input(&buf[..n]).unwrap(); - // Blankly using it, so it doesnt drop - video_rx.is_closed(); - audio_rx.is_closed(); + tokio::spawn(async move { + let mut current_stream_key_id: Option = None; + let mut codec_stamped = false; + loop { + let mut buf = [0u8; 4096]; + let n = socket.read(&mut buf).await.unwrap(); + if n == 0 { + debug!("RTMP connection closed by peer"); + return; + } + let events = session.handle_input(&buf[..n]).unwrap(); + // Blankly using it, so it doesnt drop + video_rx.is_closed(); + audio_rx.is_closed(); - for event in events { - match event { - ServerSessionResult::OutboundResponse(p) => { - socket.write_all(&p.bytes).await.unwrap(); + for event in events { + match event { + ServerSessionResult::OutboundResponse(p) => { + socket.write_all(&p.bytes).await.unwrap(); + } + ServerSessionResult::RaisedEvent(e) => match e { + ServerSessionEvent::ConnectionRequested { request_id, .. } => { + debug!("RTMP ConnectionRequested, accepting"); + let reply = session.accept_request(request_id).unwrap(); + write_outbound(&mut socket, reply).await; } - ServerSessionResult::RaisedEvent(e) => match e { - ServerSessionEvent::ConnectionRequested { - request_id, .. - } => { - debug!("RTMP ConnectionRequested, accepting"); - let reply = session.accept_request(request_id).unwrap(); - write_outbound(&mut socket, reply).await; - } - ServerSessionEvent::PublishStreamRequested { - request_id, - stream_key, - .. - } => { - info!(stream_key = %stream_key, "publish stream requested"); - let key = entity::stream_key::Entity::find_by_key( - &db, - &stream_key, - ) - .await; + ServerSessionEvent::PublishStreamRequested { + request_id, + stream_key, + .. + } => { + info!(stream_key = %stream_key, "publish stream requested"); + let key = + entity::stream_key::Entity::find_by_key(&db, &stream_key) + .await; - let key = if let Ok(Some(key)) = key { - key - } else { - warn!(stream_key = %stream_key, "stream key not found, rejecting"); - let reply = session - .reject_request( - request_id, - "", - "Stream key invalid", - ) - .unwrap(); - write_outbound(&mut socket, reply).await; - break; - }; - - let already_live = - stream_session::Model::get_active_by_stream_key_id( - &db, key.id, - ) - .await + let key = if let Ok(Some(key)) = key { + key + } else { + warn!(stream_key = %stream_key, "stream key not found, rejecting"); + let reply = session + .reject_request(request_id, "", "Stream key invalid") .unwrap(); - if already_live.is_some() { - warn!(stream_key_id = key.id, label = %key.label, "stream key already live, rejecting duplicate publish"); - let reply = session - .reject_request( - request_id, - "", - "You're already streaming...", - ) - .unwrap(); - write_outbound(&mut socket, reply).await; - break; - } - - info!(stream_key_id = key.id, label = %key.label, "stream started"); - stream_sessions.insert( - key.id, - StreamSession { - stream_key_id: key.id, - stream_key_label: key.label, - frame_channel: video_tx.clone(), - audio_channel: audio_tx.clone(), - }, - ); - - let reply = session.accept_request(request_id).unwrap(); - stream_session::Model::create_stream_session( - &db, - key.id, - Local::now().into(), - ) - .await - .unwrap(); write_outbound(&mut socket, reply).await; - } - ServerSessionEvent::PublishStreamFinished { - stream_key, - .. - } => { - info!(stream_key = %stream_key, "publish stream finished"); - let key = entity::stream_key::Entity::find_by_key( - &db, - &stream_key, - ) - .await - .unwrap() - .unwrap(); - stream_sessions.remove(&key.id); + break; + }; + + let already_live = stream_session::Model::get_active_by_stream_key_id( &db, key.id, ) + .await + .unwrap(); + if already_live.is_some() { + warn!(stream_key_id = key.id, label = %key.label, "stream key already live, rejecting duplicate publish"); + let reply = session + .reject_request( + request_id, + "", + "You're already streaming...", + ) + .unwrap(); + write_outbound(&mut socket, reply).await; + break; + } + + info!(stream_key_id = key.id, label = %key.label, "stream started"); + current_stream_key_id = Some(key.id); + stream_sessions.insert( + key.id, + StreamSession { + stream_key_id: key.id, + stream_key_label: key.label, + frame_channel: video_tx.clone(), + audio_channel: audio_tx.clone(), + codec: None, + }, + ); + + let reply = session.accept_request(request_id).unwrap(); + stream_session::Model::create_stream_session( + &db, + key.id, + Local::now().into(), + ) + .await + .unwrap(); + write_outbound(&mut socket, reply).await; + } + ServerSessionEvent::PublishStreamFinished { + stream_key, .. + } => { + info!(stream_key = %stream_key, "publish stream finished"); + let key = + entity::stream_key::Entity::find_by_key(&db, &stream_key) + .await + .unwrap() + .unwrap(); + stream_sessions.remove(&key.id); + stream_session::Model::get_active_by_stream_key_id(&db, key.id) .await .unwrap() .unwrap() @@ -204,44 +223,92 @@ impl Rtmp { .finish_stream_session(&db, Local::now().into()) .await .unwrap(); - stream_sessions.get(&key.id).unwrap().frame_channel.close(); - } - ServerSessionEvent::VideoDataReceived { - data, - timestamp, - .. - } => { - if data.len() >= 5 && &data[1..5] == b"hvc1" { - warn!("HEVC/H.265 not supported, closing connection"); - return; + current_stream_key_id = None; + } + // TODO: We can totally replace the broadcast with a + // circular_buff + // Arc>> + ServerSessionEvent::VideoDataReceived { + data, timestamp, .. + } => match Self::parse_video_codec(&data) { + Ok(codec) => { + if !codec_stamped { + if let Some(id) = current_stream_key_id { + if let Some(mut session) = stream_sessions.get_mut(&id) { + session.codec = Some(codec.clone()); + } + } + codec_stamped = true; } - if let Some(frame) = parser.parse(&data, timestamp.value) { - video_tx.broadcast(Arc::new(frame)).await.ok(); + match codec { + StreamCodec::H264 => { + let p = parser.get_or_insert_with(|| { + Box::new(H264CodecParser::new()) + }); + if let Some(frame) = p.parse(&data, timestamp.value) { + video_tx.broadcast(Arc::new(frame)).await.ok(); + } + } + StreamCodec::H265 => { + let p = parser.get_or_insert_with(|| { + Box::new(H265CodecParser::new()) + }); + if let Some(frame) = p.parse(&data, timestamp.value) { + video_tx.broadcast(Arc::new(frame)).await.ok(); + } + } + StreamCodec::AV1 => { + let p = parser.get_or_insert_with(|| { + Box::new(Av1CodecParser::new()) + }); + let pkt_type = data[0] & 0x0F; + match p.parse(&data, timestamp.value) { + Some(frame) => { + debug!( + pkt_type, + is_keyframe = frame.is_keyframe, + ts = frame.timestamp_ms, + bytes = frame.data.len(), + "AV1 frame → broadcast" + ); + video_tx.broadcast(Arc::new(frame)).await.ok(); + } + None => { + debug!(pkt_type, "AV1 packet produced no frame (seq header or unknown type)"); + } + } + } } } - ServerSessionEvent::AudioDataReceived { - data, - timestamp, - .. - } => { - // Consume the non-Send error before any await point. - let opus_frames: Vec<_> = match aac_parser.parse(&data, timestamp.value) { - Err(e) => { warn!("AAC parse error: {}", e); vec![] } - Ok(frame) => frame.map(|f| audio_proc.encode(f)).unwrap_or_default(), - }; - for frame in opus_frames { - audio_tx.broadcast(Arc::new(frame)).await.ok(); - } + Err(_err) => { + warn!(""); } - _ => {} }, + ServerSessionEvent::AudioDataReceived { + data, timestamp, .. + } => { + // Consume the non-Send error before any await point. + let opus_frames: Vec<_> = + match aac_parser.parse(&data, timestamp.value) { + Err(e) => { + warn!("AAC parse error: {}", e); + vec![] + } + Ok(frame) => frame + .map(|f| audio_proc.encode(f)) + .unwrap_or_default(), + }; + for frame in opus_frames { + audio_tx.broadcast(Arc::new(frame)).await.ok(); + } + } _ => {} - } + }, + _ => {} } } - }); - } - }); - Ok(()) + } + }); + } } } diff --git a/crates/server/src/webrtc.rs b/crates/server/src/webrtc.rs index 316de81..7de31e5 100644 --- a/crates/server/src/webrtc.rs +++ b/crates/server/src/webrtc.rs @@ -21,7 +21,9 @@ use str0m::{ net::{Protocol, Receive}, }; -use crate::{StreamSession, audio::OpusAudioFrame, media::VideoFrame, webrtc_proxy::WebrtcProxy}; +use crate::{ + StreamCodec, StreamSession, audio::OpusAudioFrame, codec::VideoFrame, webrtc_proxy::WebrtcProxy, +}; pub struct Webrtc { pub offer_rx: Receiver<(i32, i32, String)>, @@ -34,92 +36,144 @@ pub struct Webrtc { const PER_CLIENT_CONNECTION_BUF: usize = 65535; impl Webrtc { - pub fn start(mut self) -> Result<(), Box> { - tokio::spawn(async move { - while let Some(offer) = self.offer_rx.recv().await { - let (request_id, stream_id, sdp_body) = offer; - info!(request_id, stream_id, "processing offer"); + pub async fn run(mut self) { + while let Some(offer) = self.offer_rx.recv().await { + let (request_id, stream_id, sdp_body) = offer; + info!(request_id, stream_id, "processing offer"); - let stream_key = stream_key::Entity::find_by_id(stream_id) - .one(&self.db) - .await; + let stream_key = stream_key::Entity::find_by_id(stream_id) + .one(&self.db) + .await; - if let Ok(None) = stream_key { - warn!(request_id, stream_id, "stream key not found in DB, rejecting offer"); - self.accept_tx.broadcast((request_id, None)).await.unwrap(); - continue; - } - if let Err(ref e) = stream_key { - warn!(request_id, stream_id, "DB error looking up stream key: {:?}", e); - self.accept_tx.broadcast((request_id, None)).await.unwrap(); - continue; - } - - // let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); - // let local_addr = socket.local_addr().unwrap(); - - let local_addr = self.proxy.public_addr(); - - let mut builder = Rtc::builder(); - { - let cc = builder.codec_config(); - cc.enable_h264(false); - cc.add_h264(102.into(), None, true, 0x42e01f); - cc.add_h264(104.into(), None, true, 0x4d001f); - cc.add_h264(106.into(), None, true, 0x64001f); - } - let mut rtc = builder.build(Instant::now()); - - let candidate = Candidate::host(local_addr, Protocol::Udp).unwrap(); - rtc.add_local_candidate(candidate); - - let offer_sdp = SdpOffer::from_sdp_string(&sdp_body).unwrap(); - let mut changes = rtc.sdp_api(); - let mid = changes.add_media( - MediaKind::Video, - str0m::media::Direction::SendOnly, - Some(stream_id.to_string()), - Some("0".to_string()), - None, + if let Ok(None) = stream_key { + warn!( + request_id, + stream_id, "stream key not found in DB, rejecting offer" ); - changes.add_media( - MediaKind::Audio, - str0m::media::Direction::SendOnly, - Some(stream_id.to_string()), - Some("0".to_string()), - None, - ); - changes.add_channel("meow".into()); - let offer_answer = match changes.accept_offer(offer_sdp) { - Ok(a) => a, - Err(e) => { - error!("accept_offer failed: {:?}", e); - continue; - } - }; - let answer_sdp = offer_answer.to_sdp_string(); - let ufrag = answer_sdp - .lines() - .find(|l| l.starts_with("a=ice-ufrag:")) - .and_then(|l| l.strip_prefix("a=ice-ufrag:")) - .map(|s| s.trim().to_string()); - info!("Serving webrtc ufrag: {:?}", ufrag); - let (socket, rx) = self.proxy.add_client(ufrag.unwrap()); - - debug!(request_id, "sending answer back"); - self.accept_tx - .broadcast((request_id, Some(answer_sdp))) - .await - .unwrap(); - - let sessions_ref = self.sessions_ref.clone(); - let public_addr = self.proxy.public_addr(); - tokio::spawn(async move { - Webrtc::detach_connection(socket, rx, rtc, sessions_ref, stream_id, mid, public_addr).await; - }); + self.accept_tx.broadcast((request_id, None)).await.unwrap(); + continue; } - }); - Ok(()) + if let Err(ref e) = stream_key { + warn!( + request_id, + stream_id, "DB error looking up stream key: {:?}", e + ); + self.accept_tx.broadcast((request_id, None)).await.unwrap(); + continue; + } + + // let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + // let local_addr = socket.local_addr().unwrap(); + + let local_addr = self.proxy.public_addr(); + + let stream_codec = self + .sessions_ref + .get(&stream_id) + .and_then(|s| s.codec.clone()); + info!(request_id, stream_id, codec = ?match &stream_codec { + Some(StreamCodec::H264) => "H264", + Some(StreamCodec::H265) => "H265", + Some(StreamCodec::AV1) => "AV1", + None => "unknown", + }, "configuring RTC codec"); + + let mut builder = Rtc::builder(); + { + let cc = builder.codec_config(); + cc.clear(); + cc.enable_opus(true); + match stream_codec { + Some(StreamCodec::H265) => { + cc.enable_h265(true); + } + Some(StreamCodec::AV1) => { + cc.enable_av1(true); + } + _ => { + cc.add_h264(102.into(), None, true, 0x4d001f); // Main + } + } + } + let mut rtc = builder.build(Instant::now()); + + let candidate = Candidate::host(local_addr, Protocol::Udp).unwrap(); + rtc.add_local_candidate(candidate); + + let offer_sdp = SdpOffer::from_sdp_string(&sdp_body).unwrap(); + let mut changes = rtc.sdp_api(); + let mid = changes.add_media( + MediaKind::Video, + str0m::media::Direction::SendOnly, + Some(stream_id.to_string()), + Some("0".to_string()), + None, + ); + changes.add_media( + MediaKind::Audio, + str0m::media::Direction::SendOnly, + Some(stream_id.to_string()), + Some("0".to_string()), + None, + ); + changes.add_channel("meow".into()); + let offer_answer = match changes.accept_offer(offer_sdp) { + Ok(a) => a, + Err(e) => { + error!("accept_offer failed: {:?}", e); + continue; + } + }; + let answer_sdp = offer_answer.to_sdp_string(); + info!(request_id, "SDP answer:\n{}", answer_sdp); + + // Detect the case where str0m couldn't match any video codec. + // str0m serialises the m-line with an empty PT list, which is invalid SDP + // that browsers (Chrome in particular) reject with a parse error. + // This happens when the browser's offer doesn't include the codec we configured + // (e.g., browser doesn't offer AV1, or fmtp level-idx mismatch). + let video_has_pts = answer_sdp.lines().any(|line| { + line.starts_with("m=video") && line.split_whitespace().nth(3).is_some() + }); + if !video_has_pts { + warn!( + request_id, stream_id, + "no video codec negotiated — browser likely doesn't support {:?}; rejecting offer", + stream_codec + ); + self.accept_tx.broadcast((request_id, None)).await.unwrap(); + continue; + } + + let ufrag = answer_sdp + .lines() + .find(|l| l.starts_with("a=ice-ufrag:")) + .and_then(|l| l.strip_prefix("a=ice-ufrag:")) + .map(|s| s.trim().to_string()); + info!("Serving webrtc ufrag: {:?}", ufrag); + let (socket, rx) = self.proxy.add_client(ufrag.unwrap()); + + debug!(request_id, "sending answer back"); + self.accept_tx + .broadcast((request_id, Some(answer_sdp))) + .await + .unwrap(); + + let sessions_ref = self.sessions_ref.clone(); + let public_addr = self.proxy.public_addr(); + tokio::spawn(async move { + Webrtc::detach_connection( + socket, + rx, + rtc, + sessions_ref, + stream_id, + mid, + public_addr, + ) + .await; + }); + } } async fn detach_connection( @@ -152,16 +206,26 @@ impl Webrtc { } Ok(Output::Event(e)) => match e { Event::MediaAdded(ma) => { + info!(stream_id, kind = ?ma.kind, mid = ?ma.mid, "MediaAdded"); if ma.kind == MediaKind::Video { if let Some(writer) = rtc.writer(ma.mid) { + let all_pts: Vec<_> = writer + .payload_params() + .map(|p| (p.pt(), p.spec().codec)) + .collect(); + info!(stream_id, ?all_pts, "video payload params"); let best = writer.payload_params().max_by_key(|p| { p.spec().format.profile_level_id.unwrap_or(0) }); if let Some(params) = best { - info!("selected PT {:?}", params.pt()); + info!(stream_id, pt = ?params.pt(), codec = ?params.spec().codec, "selected video PT"); video_pt = Some(params.pt()); video_mid = Some(ma.mid); + } else { + warn!(stream_id, "video writer has no payload params — codec not negotiated"); } + } else { + warn!(stream_id, mid = ?ma.mid, "no writer for video mid"); } } if ma.kind == MediaKind::Audio { @@ -170,7 +234,7 @@ impl Webrtc { p.spec().format.profile_level_id.unwrap_or(0) }); if let Some(params) = best { - info!("selected PT {:?}", params.pt()); + info!(stream_id, pt = ?params.pt(), "selected audio PT"); audio_pt = Some(params.pt()); audio_mid = Some(ma.mid); } @@ -237,12 +301,13 @@ impl Webrtc { for _ in 0..8 { match stream.try_recv() { Ok(frame) => { + debug!(stream_id, is_keyframe = frame.is_keyframe, ts = frame.timestamp_ms, bytes = frame.data.len(), "WebRTC received video frame"); if !saw_keyframe { if !frame.is_keyframe { continue; } saw_keyframe = true; - debug!(stream_id, "first keyframe, starting RTP send"); + info!(stream_id, ts = frame.timestamp_ms, "first keyframe — starting RTP send"); } let now = Instant::now(); let rtp_time = @@ -251,7 +316,7 @@ impl Webrtc { (Some(pt), Some(writer)) => { match writer.write(pt, now, rtp_time, frame.data.to_vec()) { Ok(_) => wrote_any = true, - Err(e) => warn!("video RTP write error: {:?}", e), + Err(e) => error!(stream_id, "video RTP write error: {:?}", e), } } _ => warn!( @@ -268,7 +333,10 @@ impl Webrtc { break; } Err(async_broadcast::TryRecvError::Overflowed(_)) => { - // saw_keyframe = false; + // Frames were dropped from the ring buffer; resuming mid-GOP + // would give the decoder frames without their references. + // Wait for the next keyframe. + saw_keyframe = false; continue; } } diff --git a/crates/server/src/webrtc_proxy.rs b/crates/server/src/webrtc_proxy.rs index b067b44..e3791f2 100644 --- a/crates/server/src/webrtc_proxy.rs +++ b/crates/server/src/webrtc_proxy.rs @@ -1,4 +1,9 @@ -use std::{env, error::Error, net::SocketAddr, sync::Arc}; +use std::{ + env, + error::Error, + net::{IpAddr, SocketAddr}, + sync::Arc, +}; use bytes::Bytes; use dashmap::DashMap; @@ -13,6 +18,7 @@ pub struct WebRtcProxyConfig { pub proxy_port: i32, } +#[derive(Clone)] pub struct WebrtcProxy { clients_ufrag: Arc>>, clients_addr: Arc>>, @@ -34,9 +40,17 @@ impl WebrtcProxy { ip } Err(_) => { - let ip = stun_public_ip().await?; - info!(%ip, "discovered public IP via STUN for WebRTC candidates"); - ip + if cfg!(debug_assertions) { + // For testing + info!( + "using ip 0.0.0.0 for WebRTC candidates (because we are in a debug build)" + ); + IpAddr::from([127, 0, 0, 1]) + } else { + let ip = stun_public_ip().await?; + info!(%ip, "discovered public IP via STUN for WebRTC candidates"); + ip + } } }; let public_addr = SocketAddr::new(public_ip, port); @@ -49,13 +63,11 @@ impl WebrtcProxy { public_addr, }) } - pub fn start(&self) -> Result<(), Box> { - // let self_arc = Arc::new(self); - let by_ufrag = self.clients_ufrag.clone(); - let by_addr = self.clients_addr.clone(); - let socket = self.socket.clone(); - - tokio::spawn(async move { + pub async fn run(self) { + let by_ufrag = self.clients_ufrag; + let by_addr = self.clients_addr; + let socket = self.socket; + { let mut buf = vec![0u8; 65535]; loop { let (b, from) = match socket.recv_from(&mut buf).await { @@ -91,7 +103,7 @@ impl WebrtcProxy { }; let Some((_, tx)) = by_ufrag.remove(&ufrag) else { - warn!("STUN packet ({}), isnt registored", ufrag); + // warn!("STUN packet ({}), isnt registored", ufrag); continue; }; @@ -109,8 +121,7 @@ impl WebrtcProxy { } }; } - }); - Ok(()) + } } pub fn add_client(&self, ufrag: String) -> (Arc, Receiver<(Bytes, SocketAddr)>) { debug!("Added client {}", ufrag);