netget/client/rtmp/
mod.rs

1//! RTMP client over the server's handshake, chunk and AMF0 code.
2pub mod actions;
3use crate::client::llm_budget::call_llm_for_client;
4use crate::llm::actions::client_trait::{Client, ClientActionResult};
5use crate::logging::emit::Log;
6use crate::protocol::{ConnectContext, Event};
7use crate::server::rtmp::chunk::{self, Message};
8use crate::server::rtmp::{amf0, AbortOnDrop};
9use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
10use crate::state::{AccessLogOwner, ClientStatus};
11use crate::utils::clock::Instant;
12pub use actions::RtmpClientProtocol;
13use anyhow::{bail, ensure, Context, Result};
14use serde_json::{json, Map, Value};
15use std::net::SocketAddr;
16use std::time::Duration;
17use tokio::io::{AsyncWriteExt, WriteHalf};
18use tokio::net::TcpStream;
19use tokio::sync::mpsc;
20
21pub const DEFAULT_APP: &str = "live";
22const TIMEOUT: Duration = Duration::from_secs(15);
23const MAX_FLV: u64 = 64 * 1024 * 1024;
24const OUT_CHUNK_SIZE: usize = 4096;
25const MAX_DATA_MESSAGES: usize = 20;
26
27struct Conn {
28    /// Where `rtmp_publish` may read FLV files from; see `client::media_root`.
29    media_root: crate::client::media_root::MediaRoot,
30    w: WriteHalf<TcpStream>,
31    msgs: mpsc::Receiver<Result<(Message, u64)>>,
32    writer: chunk::Writer,
33    window: u32,
34    acked: u64,
35    next_tx: f64,
36    _reader: AbortOnDrop,
37}
38
39impl Conn {
40    async fn send(&mut self, csid: u32, m: &Message) -> Result<()> {
41        let bytes = self.writer.encode(csid, m)?;
42        self.w.write_all(&bytes).await?;
43        Ok(())
44    }
45
46    async fn command(&mut self, stream_id: u32, name: &str, args: &[Value]) -> Result<f64> {
47        self.next_tx += 1.0;
48        let tx = if matches!(name, "play" | "publish" | "deleteStream" | "closeStream") {
49            0.0
50        } else {
51            self.next_tx
52        };
53        let mut values = vec![json!(name), json!(tx)];
54        values.extend_from_slice(args);
55        self.send(
56            3,
57            &Message {
58                type_id: chunk::COMMAND_AMF0,
59                stream_id,
60                timestamp: 0,
61                payload: amf0::encode(&values)?,
62            },
63        )
64        .await?;
65        Ok(tx)
66    }
67
68    /// The next message from the server, answering protocol control on the way; None on close.
69    async fn next(&mut self, deadline: Duration) -> Result<Option<Message>> {
70        loop {
71            let m = match tokio::time::timeout(deadline, self.msgs.recv()).await {
72                Err(_) => bail!("the server sent nothing for {deadline:?}"),
73                Ok(None) => return Ok(None),
74                Ok(Some(m)) => m?,
75            };
76            let (m, total) = m;
77            if total - self.acked >= self.window as u64 {
78                self.acked = total;
79                self.send(
80                    2,
81                    &chunk::control(chunk::ACK, (total as u32).to_be_bytes().to_vec()),
82                )
83                .await?;
84            }
85            match m.type_id {
86                chunk::WINDOW_ACK_SIZE => {
87                    self.window = u32::from_be_bytes(
88                        m.payload
89                            .get(..4)
90                            .context("short Window Acknowledgement Size")?
91                            .try_into()?,
92                    )
93                    .max(4096);
94                }
95                chunk::USER_CONTROL
96                    if m.payload.get(..2) == Some(&[0, 6]) && m.payload.len() >= 6 =>
97                {
98                    let t = u32::from_be_bytes(m.payload[2..6].try_into()?);
99                    self.send(2, &chunk::user_control(7, t)).await?;
100                }
101                chunk::SET_CHUNK_SIZE
102                | chunk::ACK
103                | chunk::SET_PEER_BANDWIDTH
104                | chunk::ABORT
105                | chunk::USER_CONTROL => {}
106                _ => return Ok(Some(m)),
107            }
108        }
109    }
110
111    /// Wait for `_result` / `_error` to transaction `tx`, keeping onStatus codes seen on the way.
112    async fn result(&mut self, tx: f64, codes: &mut Vec<String>) -> Result<Vec<Value>> {
113        // One deadline for the whole wait: a server still streaming media would otherwise keep
114        // a per-message timeout from ever firing.
115        let deadline = Instant::now() + TIMEOUT;
116        loop {
117            let left = deadline.saturating_duration_since(Instant::now());
118            ensure!(
119                !left.is_zero(),
120                "no answer to the command within {TIMEOUT:?}"
121            );
122            let m = self
123                .next(left)
124                .await?
125                .context("the server closed the connection")?;
126            if m.type_id != chunk::COMMAND_AMF0 {
127                continue;
128            }
129            let v = amf0::decode_all(&m.payload)?;
130            match v.first().and_then(Value::as_str) {
131                Some("_result" | "_error") if v.get(1).and_then(Value::as_f64) == Some(tx) => {
132                    return Ok(v)
133                }
134                Some("onStatus") => {
135                    codes.extend(v.get(3).and_then(|i| i["code"].as_str()).map(str::to_owned))
136                }
137                _ => {}
138            }
139        }
140    }
141}
142
143pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
144    let app = ctx
145        .startup_params
146        .as_ref()
147        .map(|p| p.get_optional_string("app"))
148        .transpose()?
149        .flatten()
150        .unwrap_or_else(|| DEFAULT_APP.to_owned());
151    ensure!(
152        !app.is_empty() && app.len() <= 256 && !crate::utils::sanitize::has_controls(&app),
153        "app is a name"
154    );
155    let media_root = crate::client::media_root::MediaRoot::new(
156        ctx.startup_params
157            .as_ref()
158            .map(|p| p.get_optional_string(crate::client::media_root::MEDIA_ROOT_PARAM))
159            .transpose()?
160            .flatten()
161            .as_deref(),
162    )?;
163    let mut stream = tokio::time::timeout(TIMEOUT, TcpStream::connect(&ctx.remote_addr))
164        .await
165        .context("RTMP connect timed out")??;
166    let local = stream.local_addr()?;
167    tokio::time::timeout(TIMEOUT, chunk::connect(&mut stream))
168        .await
169        .context("RTMP handshake timed out")??;
170    let (mut r, w) = tokio::io::split(stream);
171    let (msg_tx, msgs) = mpsc::channel::<Result<(Message, u64)>>(256);
172    let reader = AbortOnDrop(tokio::spawn(async move {
173        let mut reader = chunk::Reader::default();
174        loop {
175            let m = reader.read(&mut r).await;
176            let failed = m.is_err();
177            if let Ok(m) = &m {
178                if m.type_id == chunk::SET_CHUNK_SIZE {
179                    if let Err(e) = reader.set_chunk_size(&m.payload) {
180                        let _ = msg_tx.send(Err(e)).await;
181                        return;
182                    }
183                }
184            }
185            if msg_tx.send(m.map(|m| (m, reader.bytes))).await.is_err() || failed {
186                return;
187            }
188        }
189    }));
190    let mut conn = Conn {
191        media_root,
192        w,
193        msgs,
194        writer: chunk::Writer::default(),
195        window: 2_500_000,
196        acked: 0,
197        next_tx: 0.0,
198        _reader: reader,
199    };
200    conn.send(
201        2,
202        &chunk::control(
203            chunk::SET_CHUNK_SIZE,
204            (OUT_CHUNK_SIZE as u32).to_be_bytes().to_vec(),
205        ),
206    )
207    .await?;
208    conn.writer.chunk_size = OUT_CHUNK_SIZE;
209    let tc_url = format!("rtmp://{}/{app}", ctx.remote_addr);
210    let tx = conn.command(0, "connect", &[json!({"app": app, "type": "nonprivate", "flashVer": "FMLE/3.0 (compatible; NetGet)", "tcUrl": tc_url})]).await?;
211    let reply = conn.result(tx, &mut vec![]).await?;
212    if reply[0] == "_error" {
213        bail!(
214            "connect refused: {}",
215            reply
216                .get(3)
217                .and_then(|i| i["description"].as_str().or(i["code"].as_str()))
218                .unwrap_or("no reason given")
219        );
220    }
221    ctx.state
222        .update_client_status(ctx.client_id, ClientStatus::Connected)
223        .await;
224    let external =
225        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
226    let (internal_tx, internal_rx) = mpsc::channel::<Value>(16);
227    let (event_tx, mut event_rx) = mpsc::channel::<Event>(16);
228    event_tx.try_send(Event::new(
229        &actions::CONNECTED_EVENT,
230        json!({"app": app, "server": reply.get(2).cloned().unwrap_or(Value::Null)}),
231    ))?;
232    let events_ctx = ctx.clone();
233    let dispatcher = tokio::spawn(async move {
234        let protocol = RtmpClientProtocol;
235        while let Some(event) = event_rx.recv().await {
236            let instruction = events_ctx
237                .state
238                .get_instruction_for_client(events_ctx.client_id)
239                .await
240                .unwrap_or_default();
241            let memory = events_ctx
242                .state
243                .get_memory_for_client(events_ctx.client_id)
244                .await
245                .unwrap_or_default();
246            match call_llm_for_client(
247                &events_ctx.llm_client,
248                &events_ctx.state,
249                events_ctx.client_id.to_string(),
250                &instruction,
251                &memory,
252                Some(&event),
253                &protocol,
254                &events_ctx.status_tx,
255            )
256            .await
257            {
258                Ok(result) => {
259                    if let Some(memory) = result.memory_updates {
260                        events_ctx
261                            .state
262                            .set_memory_for_client(events_ctx.client_id, memory)
263                            .await;
264                    }
265                    for action in result.actions {
266                        if internal_tx.send(action).await.is_err() {
267                            return;
268                        }
269                    }
270                }
271                Err(e) => {
272                    Log::new(Some(&events_ctx.status_tx)).warn(format!("RTMP client handler: {e}"))
273                }
274            }
275        }
276    });
277    ctx.state
278        .register_client_task(ctx.client_id, dispatcher)
279        .await;
280    let session_ctx = ctx.clone();
281    let task = tokio::spawn(async move {
282        if let Err(e) = run(&session_ctx, &mut conn, external, internal_rx, &event_tx).await {
283            Log::new(Some(&session_ctx.status_tx)).warn(format!("RTMP client ended: {e:#}"));
284        }
285        session_ctx
286            .state
287            .update_client_status(session_ctx.client_id, ClientStatus::Disconnected)
288            .await;
289        session_ctx
290            .state
291            .remove_client_handle(session_ctx.client_id)
292            .await;
293        let _ = session_ctx.status_tx.send("__UPDATE_UI__".into());
294    });
295    ctx.state.register_client_task(ctx.client_id, task).await;
296    Ok(local)
297}
298
299async fn run(
300    ctx: &ConnectContext,
301    conn: &mut Conn,
302    mut external: mpsc::Receiver<ClientCommand>,
303    mut internal: mpsc::Receiver<Value>,
304    events: &mpsc::Sender<Event>,
305) -> Result<()> {
306    loop {
307        let (action, command) = tokio::select! {
308            c = external.recv() => match c { Some(c) => (c.action.clone(), Some(c)), None => return Ok(()) },
309            a = internal.recv() => match a { Some(a) => (a, None), None => return Ok(()) },
310            m = conn.msgs.recv() => match m {
311                None | Some(Err(_)) => return Ok(()),
312                // Between operations only control traffic is expected; nothing to report.
313                Some(Ok(_)) => continue,
314            },
315        };
316        let outcome = match RtmpClientProtocol.execute_action(action.clone()) {
317            Err(e) => Ok(ClientSendOutcome::Rejected {
318                error: e.to_string(),
319            }),
320            Ok(ClientActionResult::Disconnect) => {
321                if let Some(c) = command {
322                    crate::client::command_support::reply(c, Ok(ClientSendOutcome::Disconnected));
323                }
324                return Ok(());
325            }
326            Ok(_) => {
327                let report = if action["type"] == "rtmp_play" {
328                    play(conn, &action).await
329                } else {
330                    publish(conn, &action).await
331                };
332                match report {
333                    Ok(r) => {
334                        events
335                            .send(Event::new(&actions::REPORT_EVENT, r))
336                            .await
337                            .ok();
338                        Ok(ClientSendOutcome::Sent { bytes_sent: 0 })
339                    }
340                    Err(e) => Err(e),
341                }
342            }
343        };
344        if let Some(c) = command {
345            let logged = outcome
346                .as_ref()
347                .map(|o| serde_json::to_value(o).unwrap_or(Value::Null))
348                .unwrap_or_else(|e| json!({"error": e.to_string()}));
349            ctx.state
350                .record_access_log(
351                    AccessLogOwner::Client(ctx.client_id.as_u32()),
352                    "RTMP",
353                    None,
354                    "injected_action",
355                    json!({"type": action["type"], "stream": action["stream"]}),
356                    vec![logged],
357                )
358                .await;
359            crate::client::command_support::reply(c, outcome);
360        } else if let Err(e) = outcome {
361            Log::new(Some(&ctx.status_tx)).warn(format!("RTMP action failed: {e:#}"));
362        }
363    }
364}
365
366fn video_codec(b: &[u8]) -> Option<String> {
367    let b0 = *b.first()?;
368    Some(if b0 & 0x80 != 0 {
369        match b.get(1..5)? {
370            b"avc1" => "avc".into(),
371            b"hvc1" => "hevc".into(),
372            b"av01" => "av1".into(),
373            b"vp09" => "vp9".into(),
374            other => String::from_utf8_lossy(other).into_owned(),
375        }
376    } else {
377        match b0 & 0x0f {
378            7 => "avc".into(),
379            12 => "hevc".into(),
380            n => format!("flv-{n}"),
381        }
382    })
383}
384
385fn audio_codec(b: &[u8]) -> Option<String> {
386    Some(match b.first()? >> 4 {
387        10 => "aac".into(),
388        2 => "mp3".into(),
389        n => format!("flv-{n}"),
390    })
391}
392
393#[derive(Default)]
394struct Tally {
395    codes: Vec<String>,
396    video: u64,
397    audio: u64,
398    keyframes: u64,
399    video_codec: Option<String>,
400    audio_codec: Option<String>,
401    metadata: Option<Value>,
402    first: Option<u32>,
403    last: Option<u32>,
404    data: Vec<Value>,
405}
406
407impl Tally {
408    fn media(&mut self, type_id: u8, ts: u32, payload: &[u8]) {
409        self.first.get_or_insert(ts);
410        self.last = Some(ts);
411        if type_id == chunk::VIDEO {
412            self.video += 1;
413            if payload.first().is_some_and(|b| (b >> 4) & 0x07 == 1) {
414                self.keyframes += 1;
415            }
416            if self.video_codec.is_none() {
417                self.video_codec = video_codec(payload);
418            }
419        } else {
420            self.audio += 1;
421            if self.audio_codec.is_none() {
422                self.audio_codec = audio_codec(payload);
423            }
424        }
425    }
426    fn report(self, operation: &str, stream: &str) -> Value {
427        json!({"operation": operation, "stream": stream, "status_codes": self.codes, "video_messages": self.video, "audio_messages": self.audio, "keyframes": self.keyframes,
428               "video_codec": self.video_codec, "audio_codec": self.audio_codec, "metadata": self.metadata, "first_timestamp": self.first, "last_timestamp": self.last, "data_messages": self.data})
429    }
430}
431
432async fn create_stream(conn: &mut Conn, codes: &mut Vec<String>) -> Result<u32> {
433    let tx = conn.command(0, "createStream", &[Value::Null]).await?;
434    let r = conn.result(tx, codes).await?;
435    ensure!(r[0] == "_result", "createStream refused");
436    r.get(3)
437        .and_then(Value::as_f64)
438        .map(|f| f as u32)
439        .context("createStream gave no stream ID")
440}
441
442fn is_error(code: &str) -> bool {
443    code.contains("NotFound")
444        || code.contains("Failed")
445        || code.contains("Denied")
446        || code.contains("BadName")
447        || code.contains("Rejected")
448}
449
450async fn play(conn: &mut Conn, a: &Value) -> Result<Value> {
451    let stream = a["stream"].as_str().unwrap_or_default().to_owned();
452    let seconds = a["seconds"].as_f64().unwrap_or(5.0);
453    let mut t = Tally::default();
454    let sid = create_stream(conn, &mut t.codes).await?;
455    conn.command(sid, "play", &[Value::Null, json!(stream), json!(-2000.0)])
456        .await?;
457    let mut buffer = 3u16.to_be_bytes().to_vec();
458    buffer.extend(sid.to_be_bytes());
459    buffer.extend(3000u32.to_be_bytes());
460    conn.send(2, &chunk::control(chunk::USER_CONTROL, buffer))
461        .await?;
462    let until = Instant::now() + Duration::from_secs_f64(seconds);
463    loop {
464        let left = until.saturating_duration_since(Instant::now());
465        if left.is_zero() {
466            break;
467        }
468        let m = match tokio::time::timeout(left, conn.next(Duration::from_secs(3600))).await {
469            Err(_) => break,
470            Ok(m) => match m? {
471                Some(m) => m,
472                None => break,
473            },
474        };
475        match m.type_id {
476            chunk::AUDIO | chunk::VIDEO => t.media(m.type_id, m.timestamp, &m.payload),
477            chunk::COMMAND_AMF0 => {
478                let v = amf0::decode_all(&m.payload)?;
479                if v.first().and_then(Value::as_str) == Some("onStatus") {
480                    if let Some(code) = v.get(3).and_then(|i| i["code"].as_str()) {
481                        t.codes.push(code.to_owned());
482                        if is_error(code) {
483                            break;
484                        }
485                    }
486                }
487            }
488            chunk::DATA_AMF0 => {
489                let v = amf0::decode_all(&m.payload)?;
490                match v.first().and_then(Value::as_str) {
491                    Some("onMetaData") => t.metadata = v.get(1).cloned(),
492                    Some("|RtmpSampleAccess") => {}
493                    Some(h) if t.data.len() < MAX_DATA_MESSAGES => {
494                        t.data.push(json!({"handler": h, "data": v.get(1)}))
495                    }
496                    _ => {}
497                }
498            }
499            _ => {}
500        }
501    }
502    conn.command(0, "deleteStream", &[Value::Null, json!(sid)])
503        .await?;
504    Ok(t.report("play", &stream))
505}
506
507/// One FLV tag: (type, timestamp, data).
508fn flv_tags(file: &[u8]) -> Result<Vec<(u8, u32, &[u8])>> {
509    ensure!(
510        file.len() >= 13 && &file[..3] == b"FLV" && file[3] == 1,
511        "not an FLV version 1 file"
512    );
513    let header = u32::from_be_bytes(file[5..9].try_into()?) as usize;
514    let mut i = header.checked_add(4).context("bad FLV header size")?;
515    let mut out = Vec::new();
516    while i + 11 <= file.len() {
517        let kind = file[i] & 0x1f;
518        let size = u32::from_be_bytes([0, file[i + 1], file[i + 2], file[i + 3]]) as usize;
519        let ts = u32::from_be_bytes([file[i + 7], file[i + 4], file[i + 5], file[i + 6]]);
520        let data = file
521            .get(i + 11..i + 11 + size)
522            .context("FLV tag runs past the end of the file")?;
523        ensure!(size <= chunk::MAX_MESSAGE, "FLV tag over the message bound");
524        out.push((kind, ts, data));
525        i += 11 + size + 4;
526    }
527    Ok(out)
528}
529
530async fn publish(conn: &mut Conn, a: &Value) -> Result<Value> {
531    let stream = a["stream"].as_str().unwrap_or_default().to_owned();
532    let path = a["flv_file"].as_str().unwrap_or_default();
533    let realtime = a["realtime"].as_bool().unwrap_or(true);
534    // The peer this file goes to is the peer whose responses the model reads: confined
535    // to `media_root`, or it is a file-exfiltration primitive. See `client::media_root`.
536    let path = conn
537        .media_root
538        .resolve(path, "the rtmp_publish 'flv_file'")?;
539    let meta = tokio::fs::metadata(&path)
540        .await
541        .with_context(|| format!("reading {}", path.display()))?;
542    ensure!(meta.len() <= MAX_FLV, "{} is over 64 MiB", path.display());
543    let file = tokio::fs::read(&path).await?;
544    let tags = flv_tags(&file)?;
545    let mut t = Tally::default();
546    conn.command(0, "releaseStream", &[Value::Null, json!(stream)])
547        .await?;
548    conn.command(0, "FCPublish", &[Value::Null, json!(stream)])
549        .await?;
550    let sid = create_stream(conn, &mut t.codes).await?;
551    conn.command(sid, "publish", &[Value::Null, json!(stream), json!("live")])
552        .await?;
553    loop {
554        let m = conn
555            .next(TIMEOUT)
556            .await?
557            .context("the server closed the connection")?;
558        if m.type_id != chunk::COMMAND_AMF0 {
559            continue;
560        }
561        let v = amf0::decode_all(&m.payload)?;
562        if v.first().and_then(Value::as_str) != Some("onStatus") {
563            continue;
564        }
565        let code = v
566            .get(3)
567            .and_then(|i| i["code"].as_str())
568            .unwrap_or_default()
569            .to_owned();
570        t.codes.push(code.clone());
571        if is_error(&code) {
572            return Ok(t.report("publish", &stream));
573        }
574        if code == "NetStream.Publish.Start" {
575            break;
576        }
577    }
578    let started = Instant::now();
579    let first = tags
580        .iter()
581        .find(|(k, _, _)| *k == 8 || *k == 9)
582        .map(|(_, ts, _)| *ts)
583        .unwrap_or(0);
584    for (kind, ts, data) in tags {
585        let (type_id, csid, payload) = match kind {
586            8 => (chunk::AUDIO, 4, data.to_vec()),
587            9 => (chunk::VIDEO, 6, data.to_vec()),
588            18 => {
589                let v = amf0::decode_all(data)?;
590                if v.first().and_then(Value::as_str) != Some("onMetaData") {
591                    continue;
592                }
593                let meta: Map<String, Value> = v
594                    .get(1)
595                    .and_then(Value::as_object)
596                    .cloned()
597                    .unwrap_or_default();
598                t.metadata = Some(Value::Object(meta.clone()));
599                let mut p = amf0::encode(&[json!("@setDataFrame"), json!("onMetaData")])?;
600                p.extend(amf0::encode_ecma(&meta)?);
601                (chunk::DATA_AMF0, 5, p)
602            }
603            _ => continue,
604        };
605        if realtime && type_id != chunk::DATA_AMF0 {
606            let due = Duration::from_millis(ts.saturating_sub(first) as u64);
607            let elapsed = started.elapsed();
608            if due > elapsed {
609                tokio::time::sleep(due - elapsed).await;
610            }
611        }
612        if type_id != chunk::DATA_AMF0 {
613            t.media(type_id, ts, &payload);
614        }
615        conn.send(
616            csid,
617            &Message {
618                type_id,
619                stream_id: sid,
620                timestamp: ts,
621                payload,
622            },
623        )
624        .await?;
625        // Keep acknowledging while publishing.
626        while let Ok(Ok((m, total))) = conn.msgs.try_recv() {
627            if total - conn.acked >= conn.window as u64 {
628                conn.acked = total;
629                conn.send(
630                    2,
631                    &chunk::control(chunk::ACK, (total as u32).to_be_bytes().to_vec()),
632                )
633                .await?;
634            }
635            if m.type_id == chunk::COMMAND_AMF0 {
636                if let Some(code) = amf0::decode_all(&m.payload)?
637                    .get(3)
638                    .and_then(|i| i["code"].as_str())
639                {
640                    t.codes.push(code.to_owned());
641                }
642            }
643        }
644    }
645    conn.command(0, "FCUnpublish", &[Value::Null, json!(stream)])
646        .await?;
647    conn.command(0, "deleteStream", &[Value::Null, json!(sid)])
648        .await?;
649    Ok(t.report("publish", &stream))
650}
netget/client/rtmp/mod.rs

netget/client/rtmp/
mod.rs
