447 lines
16 KiB
Rust
447 lines
16 KiB
Rust
//! An Enhanced RTMP publisher of our own, for the gateway's tests: the handshake, the
|
||
//! commands, `onMetaData`, and FLV tag bodies, legacy and enhanced, built from Enhanced
|
||
//! RTMP v2's tables; and the synthetic frames and configuration records it sends.
|
||
|
||
#![allow(dead_code)]
|
||
|
||
use super::{Auth, World, STEP};
|
||
use joule_moq_ingest::rtmp::amf0::{self, Value};
|
||
use joule_moq_ingest::rtmp::chunk::{self, types, Decoder, Message};
|
||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||
use tokio::time::timeout;
|
||
|
||
/// A stream key from the gateway's `/rtmp` resource, under `auth`'s grant.
|
||
pub async fn stream_key(w: &World, auth: &Auth) -> String {
|
||
let r = auth.send("POST", &format!("{}/rtmp", w.base()), &[], b"").await;
|
||
assert_eq!(r.status, 201, "POST /rtmp: {}", r.text());
|
||
let v: serde_json::Value = serde_json::from_slice(&r.body).unwrap();
|
||
v["stream_key"].as_str().unwrap().to_string()
|
||
}
|
||
|
||
/// One status the gateway sent: its code, level and description.
|
||
#[derive(Clone, Debug)]
|
||
pub struct Status {
|
||
pub code: String,
|
||
pub level: String,
|
||
pub description: String,
|
||
}
|
||
|
||
pub struct Client<S> {
|
||
io: S,
|
||
dec: Decoder,
|
||
out_chunk: usize,
|
||
pub received: Vec<Message>,
|
||
}
|
||
|
||
impl<S: AsyncRead + AsyncWrite + Unpin> Client<S> {
|
||
/// The simple handshake over `io`.
|
||
pub async fn start(mut io: S) -> Client<S> {
|
||
let mut c0c1 = vec![3u8];
|
||
c0c1.extend_from_slice(&[0; 8]);
|
||
c0c1.extend((0..chunk::HANDSHAKE_LEN - 8).map(|i| (i * 13) as u8));
|
||
io.write_all(&c0c1).await.unwrap();
|
||
let mut s = vec![0u8; 1 + 2 * chunk::HANDSHAKE_LEN];
|
||
timeout(STEP, io.read_exact(&mut s)).await.expect("S0 S1 S2 in time").unwrap();
|
||
assert_eq!(s[0], 3);
|
||
assert_eq!(&s[1 + chunk::HANDSHAKE_LEN + 8..], &c0c1[9..], "S2 echoes C1");
|
||
io.write_all(&s[1..1 + chunk::HANDSHAKE_LEN]).await.unwrap();
|
||
Client { io, dec: Decoder::new(1 << 24, 1 << 24), out_chunk: 128, received: Vec::new() }
|
||
}
|
||
|
||
pub async fn send(&mut self, m: Message) {
|
||
let b = chunk::encode(&m, self.out_chunk);
|
||
self.io.write_all(&b).await.unwrap();
|
||
if m.type_id == types::SET_CHUNK_SIZE {
|
||
self.out_chunk = u32::from_be_bytes(m.payload[..4].try_into().unwrap()) as usize;
|
||
}
|
||
}
|
||
|
||
pub async fn command(&mut self, msid: u32, values: &[Value]) {
|
||
self.send(Message { csid: 3, type_id: types::COMMAND_AMF0, msid, ts: 0, payload: amf0::encode(values) }).await;
|
||
}
|
||
|
||
/// The next message the gateway sends, bounded.
|
||
pub async fn next(&mut self) -> Option<Message> {
|
||
let mut buf = vec![0u8; 8192];
|
||
loop {
|
||
if let Some(m) = self.dec.next_message().expect("the gateway's chunk stream reads") {
|
||
if m.type_id == types::SET_CHUNK_SIZE {
|
||
self.dec.set_chunk_size(u32::from_be_bytes(m.payload[..4].try_into().unwrap())).unwrap();
|
||
}
|
||
self.received.push(m.clone());
|
||
return Some(m);
|
||
}
|
||
match timeout(STEP, self.io.read(&mut buf)).await {
|
||
Ok(Ok(0)) | Ok(Err(_)) | Err(_) => return None,
|
||
Ok(Ok(n)) => self.dec.feed(&buf[..n]),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// The values of the next command message, skipping the others.
|
||
pub async fn next_command(&mut self) -> Option<Vec<Value>> {
|
||
while let Some(m) = self.next().await {
|
||
if m.type_id == types::COMMAND_AMF0 {
|
||
return Some(amf0::read_all(&m.payload).expect("a command reads"));
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
/// The next `onStatus`.
|
||
pub async fn status(&mut self) -> Option<Status> {
|
||
while let Some(v) = self.next_command().await {
|
||
if v.first().and_then(Value::as_str) == Some("onStatus") {
|
||
let info = v.get(3)?;
|
||
let s = |k| info.get(k).and_then(Value::as_str).unwrap_or("").to_string();
|
||
return Some(Status { code: s("code"), level: s("level"), description: s("description") });
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
/// `connect` with `capsEx` and the FourCCs it can send, then the `_result`, and a
|
||
/// chunk size of 4096 from here.
|
||
pub async fn connect(&mut self, caps_ex: u32) -> Vec<Value> {
|
||
let obj = Value::obj(&[
|
||
("app", Value::str("live")),
|
||
("type", Value::str("nonprivate")),
|
||
("tcUrl", Value::str("rtmp://127.0.0.1/live")),
|
||
("capsEx", Value::Number(caps_ex as f64)),
|
||
("fourCcList", Value::StrictArray(["av01", "vp09", "vp08", "hvc1", "avc1", "Opus", "fLaC", "mp4a"].iter().map(|f| Value::str(f)).collect())),
|
||
]);
|
||
self.command(0, &[Value::str("connect"), Value::Number(1.0), obj]).await;
|
||
let r = loop {
|
||
let v = self.next_command().await.expect("the connect _result");
|
||
if v.first().and_then(Value::as_str) == Some("_result") {
|
||
break v;
|
||
}
|
||
};
|
||
self.send(chunk::control(types::SET_CHUNK_SIZE, 4096u32.to_be_bytes().to_vec())).await;
|
||
r
|
||
}
|
||
|
||
/// `createStream` and `publish` with `key`; the gateway's status.
|
||
pub async fn publish(&mut self, key: &str) -> Status {
|
||
self.command(0, &[Value::str("releaseStream"), Value::Number(2.0), Value::Null, Value::str(key)]).await;
|
||
self.command(0, &[Value::str("FCPublish"), Value::Number(3.0), Value::Null, Value::str(key)]).await;
|
||
self.command(0, &[Value::str("createStream"), Value::Number(4.0), Value::Null]).await;
|
||
let id = loop {
|
||
let v = self.next_command().await.expect("the createStream _result");
|
||
if v.first().and_then(Value::as_str) == Some("_result") && v.get(1).and_then(Value::as_f64) == Some(4.0) {
|
||
break v.get(3).and_then(Value::as_f64).unwrap() as u32;
|
||
}
|
||
};
|
||
self.command(id, &[Value::str("publish"), Value::Number(5.0), Value::Null, Value::str(key), Value::str("live")]).await;
|
||
self.status().await.expect("a status for publish")
|
||
}
|
||
|
||
/// `@setDataFrame` `onMetaData` with `pairs`.
|
||
pub async fn metadata(&mut self, pairs: Vec<(String, Value)>) {
|
||
let p = amf0::encode(&[Value::str("@setDataFrame"), Value::str("onMetaData"), Value::EcmaArray(pairs)]);
|
||
self.send(Message { csid: 4, type_id: types::DATA_AMF0, msid: 1, ts: 0, payload: p }).await;
|
||
}
|
||
|
||
pub async fn video(&mut self, ts: u32, body: Vec<u8>) {
|
||
self.send(Message { csid: 6, type_id: types::VIDEO, msid: 1, ts, payload: body }).await;
|
||
}
|
||
|
||
pub async fn audio(&mut self, ts: u32, body: Vec<u8>) {
|
||
self.send(Message { csid: 4, type_id: types::AUDIO, msid: 1, ts, payload: body }).await;
|
||
}
|
||
|
||
pub async fn delete_stream(&mut self) {
|
||
self.command(0, &[Value::str("FCUnpublish"), Value::Number(6.0), Value::Null, Value::str("")]).await;
|
||
self.command(0, &[Value::str("deleteStream"), Value::Number(7.0), Value::Null, Value::Number(1.0)]).await;
|
||
}
|
||
|
||
pub async fn shutdown(&mut self) {
|
||
let _ = self.io.shutdown().await;
|
||
}
|
||
|
||
/// Read until the gateway closes the connection, bounded.
|
||
pub async fn drain(&mut self) {
|
||
while self.next().await.is_some() {}
|
||
}
|
||
}
|
||
|
||
// ── FLV tag bodies (Enhanced RTMP v2, Enhanced Video and Enhanced Audio) ──────────
|
||
|
||
pub fn legacy_avc(packet_type: u8, key: bool, cts: i32, body: &[u8]) -> Vec<u8> {
|
||
let mut o = vec![if key { 0x17 } else { 0x27 }, packet_type];
|
||
o.extend_from_slice(&cts.to_be_bytes()[1..]);
|
||
o.extend_from_slice(body);
|
||
o
|
||
}
|
||
|
||
pub fn legacy_aac(packet_type: u8, body: &[u8]) -> Vec<u8> {
|
||
let mut o = vec![0xaf, packet_type];
|
||
o.extend_from_slice(body);
|
||
o
|
||
}
|
||
|
||
/// An enhanced video body: `isExVideoHeader`, the frame type, the packet type, the
|
||
/// FourCC, then `body` (which, for CodedFrames of avc1, hvc1 and vvc1, begins with its
|
||
/// SI24 composition time offset).
|
||
pub fn ex_video(key: bool, packet_type: u8, fourcc: &[u8; 4], body: &[u8]) -> Vec<u8> {
|
||
let mut o = vec![0x80 | if key { 0x10 } else { 0x20 } | packet_type];
|
||
o.extend_from_slice(fourcc);
|
||
o.extend_from_slice(body);
|
||
o
|
||
}
|
||
|
||
/// The same, after a ModEx `TimestampOffsetNano` of `nano`.
|
||
pub fn ex_video_nano(key: bool, packet_type: u8, fourcc: &[u8; 4], nano: u32, body: &[u8]) -> Vec<u8> {
|
||
let mut o = vec![0x80 | if key { 0x10 } else { 0x20 } | 7, 2];
|
||
o.extend_from_slice(&nano.to_be_bytes()[1..]);
|
||
o.push(packet_type);
|
||
o.extend_from_slice(fourcc);
|
||
o.extend_from_slice(body);
|
||
o
|
||
}
|
||
|
||
/// A multitrack video body of type ManyTracksManyCodecs: each track its FourCC, its id,
|
||
/// its size and its body.
|
||
pub fn ex_video_tracks(key: bool, packet_type: u8, tracks: &[(&[u8; 4], u8, Vec<u8>)]) -> Vec<u8> {
|
||
let mut o = vec![0x80 | if key { 0x10 } else { 0x20 } | 6, 2 << 4 | packet_type];
|
||
for (f, id, b) in tracks {
|
||
o.extend_from_slice(*f);
|
||
o.push(*id);
|
||
o.extend_from_slice(&(b.len() as u32).to_be_bytes()[1..]);
|
||
o.extend_from_slice(b);
|
||
}
|
||
o
|
||
}
|
||
|
||
pub fn ex_audio(packet_type: u8, fourcc: &[u8; 4], body: &[u8]) -> Vec<u8> {
|
||
let mut o = vec![0x90 | packet_type];
|
||
o.extend_from_slice(fourcc);
|
||
o.extend_from_slice(body);
|
||
o
|
||
}
|
||
|
||
/// A multitrack audio body of type ManyTracksManyCodecs.
|
||
pub fn ex_audio_tracks(packet_type: u8, tracks: &[(&[u8; 4], u8, Vec<u8>)]) -> Vec<u8> {
|
||
let mut o = vec![0x90 | 5, 2 << 4 | packet_type];
|
||
for (f, id, b) in tracks {
|
||
o.extend_from_slice(*f);
|
||
o.push(*id);
|
||
o.extend_from_slice(&(b.len() as u32).to_be_bytes()[1..]);
|
||
o.extend_from_slice(b);
|
||
}
|
||
o
|
||
}
|
||
|
||
// ── Synthetic configurations and frames ────────────────────────────────────────────
|
||
|
||
/// Bytes none of which is zero, so no start code appears inside a NAL unit.
|
||
pub fn body(seed: u64, n: usize) -> Vec<u8> {
|
||
(0..n).map(|j| ((seed.wrapping_mul(31).wrapping_add(j as u64 * 7)) % 255 + 1) as u8).collect()
|
||
}
|
||
|
||
/// NAL units, each after a four-byte length.
|
||
pub fn lengths(nals: &[Vec<u8>]) -> Vec<u8> {
|
||
let mut o = Vec::new();
|
||
for n in nals {
|
||
o.extend_from_slice(&(n.len() as u32).to_be_bytes());
|
||
o.extend_from_slice(n);
|
||
}
|
||
o
|
||
}
|
||
|
||
/// NAL units as an Annex B access unit, each after a four-byte start code.
|
||
pub fn annexb(nals: &[Vec<u8>]) -> Vec<u8> {
|
||
let mut o = Vec::new();
|
||
for n in nals {
|
||
o.extend_from_slice(&[0, 0, 0, 1]);
|
||
o.extend_from_slice(n);
|
||
}
|
||
o
|
||
}
|
||
|
||
pub const AVC_SPS: [u8; 6] = [0x67, 0x64, 0x00, 0x1f, 0xac, 0xd9];
|
||
pub const AVC_PPS: [u8; 4] = [0x68, 0xeb, 0xe3, 0xcb];
|
||
|
||
/// An AVCDecoderConfigurationRecord: High profile, level 3.1, four-byte lengths.
|
||
pub fn avc_record() -> Vec<u8> {
|
||
let mut o = vec![1, 0x64, 0x00, 0x1f, 0xff, 0xe1];
|
||
o.extend_from_slice(&(AVC_SPS.len() as u16).to_be_bytes());
|
||
o.extend_from_slice(&AVC_SPS);
|
||
o.push(1);
|
||
o.extend_from_slice(&(AVC_PPS.len() as u16).to_be_bytes());
|
||
o.extend_from_slice(&AVC_PPS);
|
||
o
|
||
}
|
||
|
||
/// AVC access unit `i`'s NAL units: an IDR slice when `key`, a non-IDR one otherwise.
|
||
pub fn avc_nals(i: u64, key: bool) -> Vec<Vec<u8>> {
|
||
let mut n = vec![if key { 0x65 } else { 0x41 }];
|
||
n.extend(body(i, if key { 2000 } else { 400 + (i % 5) as usize * 30 }));
|
||
vec![n]
|
||
}
|
||
|
||
pub fn hevc_sets() -> Vec<Vec<u8>> {
|
||
vec![vec![0x40, 0x01, 0x0c, 0x01], vec![0x42, 0x01, 0x01, 0x01, 0x60], vec![0x44, 0x01, 0xc1, 0x72]]
|
||
}
|
||
|
||
/// An HEVCDecoderConfigurationRecord: Main profile, level 3.1 (93), four-byte lengths, its
|
||
/// VPS, SPS and PPS.
|
||
pub fn hevc_record() -> Vec<u8> {
|
||
let mut o = vec![1, 0x01, 0x60, 0, 0, 0, 0x90, 0, 0, 0, 0, 0, 93, 0xf0, 0, 0xfc, 0xfd, 0xf8, 0xf8, 0, 0, 0x0f, 3];
|
||
for (ty, s) in [32u8, 33, 34].iter().zip(hevc_sets()) {
|
||
o.push(0x80 | ty);
|
||
o.extend_from_slice(&1u16.to_be_bytes());
|
||
o.extend_from_slice(&(s.len() as u16).to_be_bytes());
|
||
o.extend_from_slice(&s);
|
||
}
|
||
o
|
||
}
|
||
|
||
/// HEVC access unit `i`: an IDR_W_RADL slice (type 19) when `key`, TRAIL_R (1) otherwise.
|
||
pub fn hevc_nals(i: u64, key: bool) -> Vec<Vec<u8>> {
|
||
let mut n = if key { vec![0x26, 0x01] } else { vec![0x02, 0x01] };
|
||
n.extend(body(i + 50, if key { 1800 } else { 350 }));
|
||
vec![n]
|
||
}
|
||
|
||
struct Bw(Vec<u8>, usize);
|
||
impl Bw {
|
||
fn put(&mut self, v: u64, n: usize) {
|
||
for i in (0..n).rev() {
|
||
if self.1.is_multiple_of(8) {
|
||
self.0.push(0);
|
||
}
|
||
if (v >> i) & 1 == 1 {
|
||
*self.0.last_mut().unwrap() |= 0x80 >> (self.1 % 8);
|
||
}
|
||
self.1 += 1;
|
||
}
|
||
}
|
||
}
|
||
|
||
fn obu(ty: u8, payload: &[u8]) -> Vec<u8> {
|
||
let mut o = vec![(ty << 3) | 0x02];
|
||
let mut n = payload.len();
|
||
loop {
|
||
let b = (n & 0x7f) as u8;
|
||
n >>= 7;
|
||
if n == 0 {
|
||
o.push(b);
|
||
break;
|
||
}
|
||
o.push(b | 0x80);
|
||
}
|
||
o.extend_from_slice(payload);
|
||
o
|
||
}
|
||
|
||
/// An AV1 sequence header OBU: profile 0, one operating point at level 8, the size.
|
||
pub fn av1_sequence_header(w: u32, h: u32) -> Vec<u8> {
|
||
let mut b = Bw(Vec::new(), 0);
|
||
b.put(0, 3);
|
||
b.put(0, 1);
|
||
b.put(0, 1);
|
||
b.put(0, 1);
|
||
b.put(0, 1);
|
||
b.put(0, 5);
|
||
b.put(0, 12);
|
||
b.put(8, 5);
|
||
b.put(0, 1);
|
||
b.put(15, 4);
|
||
b.put(15, 4);
|
||
b.put(w as u64 - 1, 16);
|
||
b.put(h as u64 - 1, 16);
|
||
obu(1, &b.0)
|
||
}
|
||
|
||
/// An AV1CodecConfigurationRecord with the sequence header as its configOBUs.
|
||
pub fn av1_record(w: u32, h: u32) -> Vec<u8> {
|
||
let mut o = vec![0x81, 0x08, 0x0c, 0x00];
|
||
o.extend(av1_sequence_header(w, h));
|
||
o
|
||
}
|
||
|
||
/// AV1 temporal unit `i`: a frame OBU (with a temporal delimiter before it, which the
|
||
/// gateway drops and its framing restores).
|
||
pub fn av1_tu(i: u64, key: bool) -> Vec<u8> {
|
||
let mut o = vec![0x12, 0x00];
|
||
o.extend(obu(6, &body(i + 90, if key { 1500 } else { 300 })));
|
||
o
|
||
}
|
||
|
||
/// VP9 frame `i`: an uncompressed header (keyframe with its sync code, colour config and
|
||
/// size; or an interframe) and bytes.
|
||
pub fn vp9_frame(i: u64, key: bool, w: u32, h: u32) -> Vec<u8> {
|
||
let mut b = Bw(Vec::new(), 0);
|
||
b.put(2, 2);
|
||
b.put(0, 2);
|
||
b.put(0, 1);
|
||
b.put(if key { 0 } else { 1 }, 1);
|
||
b.put(1, 1);
|
||
b.put(0, 1);
|
||
if key {
|
||
b.put(0x49_83_42, 24);
|
||
b.put(1, 3);
|
||
b.put(0, 1);
|
||
b.put(w as u64 - 1, 16);
|
||
b.put(h as u64 - 1, 16);
|
||
}
|
||
let mut f = b.0;
|
||
f.extend(body(i + 300, if key { 1200 } else { 250 }));
|
||
f
|
||
}
|
||
|
||
/// A VPCodecConfigurationRecord: profile 0, level 3.1, 8 bits.
|
||
pub fn vp9_record() -> Vec<u8> {
|
||
vec![0, 31, 0x80, 1, 1, 1, 0, 0]
|
||
}
|
||
|
||
/// An AudioSpecificConfig: AAC-LC, 48 kHz, stereo.
|
||
pub const AAC_ASC: [u8; 2] = [0x11, 0x90];
|
||
|
||
pub fn aac_raw(i: u64) -> Vec<u8> {
|
||
body(i + 1000, 120 + (i % 3) as usize)
|
||
}
|
||
|
||
/// The ADTS frame the gateway builds for a raw frame under `AAC_ASC`: profile 1, index 3,
|
||
/// two channels.
|
||
pub fn adts(raw: &[u8]) -> Vec<u8> {
|
||
let len = raw.len() + 7;
|
||
let mut o = vec![0xff, 0xf1, (1 << 6) | (3 << 2), (2 << 6) | ((len >> 11) & 3) as u8, ((len >> 3) & 0xff) as u8, (((len & 7) << 5) as u8) | 0x1f, 0xfc];
|
||
o.extend_from_slice(raw);
|
||
o
|
||
}
|
||
|
||
pub fn opus_packet(i: u64) -> Vec<u8> {
|
||
let mut f = vec![0xfc];
|
||
f.extend(body(i + 7_000, 80 + (i % 7) as usize));
|
||
f
|
||
}
|
||
|
||
/// An Opus ID header (RFC 7845 §5.1): two channels, mapping family 0.
|
||
pub fn opus_head() -> Vec<u8> {
|
||
let mut o = b"OpusHead".to_vec();
|
||
o.extend_from_slice(&[1, 2, 0x38, 0x01, 0x80, 0xbb, 0, 0, 0, 0, 0]);
|
||
o
|
||
}
|
||
|
||
/// `fLaC` and a STREAMINFO block (not marked last): 48 kHz, two channels, 16 bits.
|
||
pub fn flac_header() -> Vec<u8> {
|
||
let mut o = b"fLaC".to_vec();
|
||
o.extend_from_slice(&[0x00, 0, 0, 34]);
|
||
let mut si = vec![0x10, 0x00, 0x10, 0x00, 0, 0, 0, 0, 0, 0];
|
||
// sample rate 48 000 (20 bits), channels − 1 = 1 (3 bits), bits − 1 = 15 (5 bits),
|
||
// total samples 0 (36 bits).
|
||
let rate = 48_000u64;
|
||
let v: u64 = rate << 44 | 1 << 41 | 15 << 36;
|
||
si.extend_from_slice(&v.to_be_bytes());
|
||
si.extend_from_slice(&[0; 16]);
|
||
o.extend(si);
|
||
o
|
||
}
|
||
|
||
pub fn flac_frame(i: u64) -> Vec<u8> {
|
||
let mut f = vec![0xff, 0xf8, 0x69, 0x18, i as u8];
|
||
f.extend(body(i + 5_000, 200));
|
||
f
|
||
}
|