netget/server/rtmp/
chunk.rs

1//! RTMP handshake and chunk stream (Adobe RTMP 1.0 §5.2–5.4): the simple version-3 handshake,
2//! chunk basic and message headers in all four formats with extended timestamps, message
3//! reassembly per chunk stream, and a chunk writer.
4use anyhow::{bail, ensure, Context, Result};
5use std::collections::HashMap;
6use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
7
8pub const HANDSHAKE_SIZE: usize = 1536;
9/// Largest message accepted (a keyframe of a high-bitrate stream fits comfortably).
10pub const MAX_MESSAGE: usize = 8 * 1024 * 1024;
11/// Largest chunk size either side may set.
12pub const MAX_CHUNK_SIZE: usize = 1024 * 1024;
13const MAX_CHUNK_STREAMS: usize = 64;
14
15pub const SET_CHUNK_SIZE: u8 = 1;
16pub const ABORT: u8 = 2;
17pub const ACK: u8 = 3;
18pub const USER_CONTROL: u8 = 4;
19pub const WINDOW_ACK_SIZE: u8 = 5;
20pub const SET_PEER_BANDWIDTH: u8 = 6;
21pub const AUDIO: u8 = 8;
22pub const VIDEO: u8 = 9;
23pub const DATA_AMF3: u8 = 15;
24pub const COMMAND_AMF3: u8 = 17;
25pub const DATA_AMF0: u8 = 18;
26pub const COMMAND_AMF0: u8 = 20;
27
28/// A reassembled RTMP message.
29#[derive(Debug, Clone)]
30pub struct Message {
31    pub type_id: u8,
32    pub stream_id: u32,
33    pub timestamp: u32,
34    pub payload: Vec<u8>,
35}
36
37fn random_block() -> Vec<u8> {
38    use ring::rand::SecureRandom;
39    let mut b = vec![0u8; HANDSHAKE_SIZE];
40    let _ = ring::rand::SystemRandom::new().fill(&mut b[8..]);
41    b[..8].fill(0);
42    b
43}
44
45/// Server side: C0+C1 in, S0+S1+S2 out, C2 in.
46pub async fn accept<S: AsyncRead + AsyncWrite + Unpin>(s: &mut S) -> Result<()> {
47    let mut c0c1 = vec![0u8; 1 + HANDSHAKE_SIZE];
48    s.read_exact(&mut c0c1).await?;
49    ensure!(c0c1[0] == 3, "RTMP version {} is not 3", c0c1[0]);
50    let mut out = vec![3u8];
51    out.extend(random_block());
52    out.extend_from_slice(&c0c1[1..]);
53    s.write_all(&out).await?;
54    let mut c2 = vec![0u8; HANDSHAKE_SIZE];
55    s.read_exact(&mut c2).await?;
56    Ok(())
57}
58
59/// Client side: C0+C1 out, S0+S1+S2 in, C2 (an echo of S1) out.
60pub async fn connect<S: AsyncRead + AsyncWrite + Unpin>(s: &mut S) -> Result<()> {
61    let mut out = vec![3u8];
62    out.extend(random_block());
63    s.write_all(&out).await?;
64    let mut s0s1s2 = vec![0u8; 1 + 2 * HANDSHAKE_SIZE];
65    s.read_exact(&mut s0s1s2).await?;
66    ensure!(
67        s0s1s2[0] == 3,
68        "the server answered RTMP version {}",
69        s0s1s2[0]
70    );
71    s.write_all(&s0s1s2[1..1 + HANDSHAKE_SIZE]).await?;
72    Ok(())
73}
74
75#[derive(Default, Clone)]
76struct InStream {
77    timestamp: u32,
78    delta: u32,
79    length: usize,
80    type_id: u8,
81    stream_id: u32,
82    extended: bool,
83    buf: Vec<u8>,
84}
85
86/// Reassembles messages from chunks.
87pub struct Reader {
88    pub chunk_size: usize,
89    streams: HashMap<u32, InStream>,
90    /// Bytes read so far, for acknowledgements.
91    pub bytes: u64,
92}
93
94impl Default for Reader {
95    fn default() -> Self {
96        Self {
97            chunk_size: 128,
98            streams: HashMap::new(),
99            bytes: 0,
100        }
101    }
102}
103
104impl Reader {
105    async fn u8<R: AsyncRead + Unpin>(&mut self, r: &mut R) -> Result<u8> {
106        self.bytes += 1;
107        Ok(r.read_u8().await?)
108    }
109    async fn be(&mut self, r: &mut (impl AsyncRead + Unpin), n: usize) -> Result<u32> {
110        let mut b = [0u8; 4];
111        r.read_exact(&mut b[4 - n..]).await?;
112        self.bytes += n as u64;
113        Ok(u32::from_be_bytes(b))
114    }
115
116    pub fn abort(&mut self, csid: u32) {
117        if let Some(s) = self.streams.get_mut(&csid) {
118            s.buf.clear();
119        }
120    }
121
122    /// The next complete message.
123    pub async fn read<R: AsyncRead + Unpin>(&mut self, r: &mut R) -> Result<Message> {
124        loop {
125            let b0 = self.u8(r).await?;
126            let fmt = b0 >> 6;
127            let csid = match b0 & 0x3f {
128                0 => self.u8(r).await? as u32 + 64,
129                1 => {
130                    let lo = self.u8(r).await? as u32;
131                    let hi = self.u8(r).await? as u32;
132                    hi * 256 + lo + 64
133                }
134                n => n as u32,
135            };
136            if !self.streams.contains_key(&csid) {
137                ensure!(
138                    self.streams.len() < MAX_CHUNK_STREAMS,
139                    "more than {MAX_CHUNK_STREAMS} chunk streams"
140                );
141                ensure!(
142                    fmt == 0,
143                    "chunk stream {csid} starts without a type-0 header"
144                );
145            }
146            let mut st = self.streams.remove(&csid).unwrap_or_default();
147            let starting = st.buf.is_empty();
148            match fmt {
149                0 => {
150                    let ts = self.be(r, 3).await?;
151                    st.length = self.be(r, 3).await? as usize;
152                    st.type_id = self.u8(r).await?;
153                    let mut sid = [0u8; 4];
154                    r.read_exact(&mut sid).await?;
155                    self.bytes += 4;
156                    st.stream_id = u32::from_le_bytes(sid);
157                    st.extended = ts == 0xFF_FFFF;
158                    st.timestamp = if st.extended {
159                        self.be(r, 4).await?
160                    } else {
161                        ts
162                    };
163                    st.delta = 0;
164                    ensure!(
165                        starting,
166                        "type-0 header in the middle of a message on chunk stream {csid}"
167                    );
168                }
169                1 | 2 => {
170                    let d = self.be(r, 3).await?;
171                    if fmt == 1 {
172                        st.length = self.be(r, 3).await? as usize;
173                        st.type_id = self.u8(r).await?;
174                    }
175                    st.extended = d == 0xFF_FFFF;
176                    st.delta = if st.extended { self.be(r, 4).await? } else { d };
177                    ensure!(
178                        starting,
179                        "type-{fmt} header in the middle of a message on chunk stream {csid}"
180                    );
181                    st.timestamp = st.timestamp.wrapping_add(st.delta);
182                }
183                _ => {
184                    if st.extended {
185                        self.be(r, 4).await?;
186                    }
187                    if starting {
188                        st.timestamp = st.timestamp.wrapping_add(st.delta);
189                    }
190                }
191            }
192            ensure!(
193                st.length <= MAX_MESSAGE,
194                "message of {} bytes exceeds {MAX_MESSAGE}",
195                st.length
196            );
197            let n = (st.length - st.buf.len()).min(self.chunk_size);
198            let start = st.buf.len();
199            st.buf.resize(start + n, 0);
200            r.read_exact(&mut st.buf[start..]).await?;
201            self.bytes += n as u64;
202            if st.buf.len() >= st.length {
203                let payload = std::mem::take(&mut st.buf);
204                let msg = Message {
205                    type_id: st.type_id,
206                    stream_id: st.stream_id,
207                    timestamp: st.timestamp,
208                    payload,
209                };
210                self.streams.insert(csid, st);
211                return Ok(msg);
212            }
213            self.streams.insert(csid, st);
214        }
215    }
216
217    /// Apply a peer's Set Chunk Size.
218    pub fn set_chunk_size(&mut self, payload: &[u8]) -> Result<()> {
219        let v = u32::from_be_bytes(
220            payload
221                .get(..4)
222                .context("short Set Chunk Size")?
223                .try_into()?,
224        ) & 0x7FFF_FFFF;
225        ensure!(
226            (1..=MAX_CHUNK_SIZE as u32).contains(&v),
227            "chunk size {v} out of bounds"
228        );
229        self.chunk_size = v as usize;
230        Ok(())
231    }
232}
233
234/// Splits messages into chunks.
235pub struct Writer {
236    pub chunk_size: usize,
237}
238
239impl Default for Writer {
240    fn default() -> Self {
241        Self { chunk_size: 128 }
242    }
243}
244
245fn basic_header(fmt: u8, csid: u32, out: &mut Vec<u8>) {
246    match csid {
247        2..=63 => out.push((fmt << 6) | csid as u8),
248        64..=319 => out.extend([fmt << 6, (csid - 64) as u8]),
249        _ => out.extend([
250            (fmt << 6) | 1,
251            ((csid - 64) & 0xff) as u8,
252            ((csid - 64) >> 8) as u8,
253        ]),
254    }
255}
256
257impl Writer {
258    pub fn encode(&self, csid: u32, m: &Message) -> Result<Vec<u8>> {
259        if m.payload.len() > 0xFF_FFFF {
260            bail!("message too long for RTMP");
261        }
262        let extended = m.timestamp >= 0xFF_FFFF;
263        let mut out =
264            Vec::with_capacity(m.payload.len() + 16 + m.payload.len() / self.chunk_size.max(1) * 5);
265        basic_header(0, csid, &mut out);
266        let ts = if extended { 0xFF_FFFF } else { m.timestamp };
267        out.extend(&ts.to_be_bytes()[1..]);
268        out.extend(&(m.payload.len() as u32).to_be_bytes()[1..]);
269        out.push(m.type_id);
270        out.extend(m.stream_id.to_le_bytes());
271        if extended {
272            out.extend(m.timestamp.to_be_bytes());
273        }
274        for (i, chunk) in m.payload.chunks(self.chunk_size.max(1)).enumerate() {
275            if i > 0 {
276                basic_header(3, csid, &mut out);
277                if extended {
278                    out.extend(m.timestamp.to_be_bytes());
279                }
280            }
281            out.extend(chunk);
282        }
283        Ok(out)
284    }
285}
286
287pub fn control(type_id: u8, payload: Vec<u8>) -> Message {
288    Message {
289        type_id,
290        stream_id: 0,
291        timestamp: 0,
292        payload,
293    }
294}
295
296/// User Control event (type 4): StreamBegin = 0, StreamEOF = 1, PingRequest = 6, PingResponse = 7.
297pub fn user_control(event: u16, data: u32) -> Message {
298    let mut p = event.to_be_bytes().to_vec();
299    p.extend(data.to_be_bytes());
300    control(USER_CONTROL, p)
301}
netget/server/rtmp/chunk.rs

netget/server/rtmp/
chunk.rs
