Skip to main content

netget/client/rtp/
mod.rs

1//! RTP endpoint (client role): sends synthesized G.711 to `remote_addr` in paced 20 ms frames,
2//! receives on `listen`, and summarises what arrives per SSRC — one event when a stream
3//! starts, one when it ends — rather than per packet, which at 50 packets a second would be
4//! a model call every 20 ms.
5pub mod actions;
6use crate::client::llm_budget::call_llm_for_client;
7use crate::llm::actions::client_trait::{Client, ClientActionResult};
8use crate::logging::emit::Log;
9use crate::protocol::{ConnectContext, Event};
10use crate::server::rtp::media::{self, AudioCodec};
11use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
12use crate::state::{AccessLogOwner, ClientStatus};
13pub use actions::RtpClientProtocol;
14use anyhow::{Context, Result};
15use serde_json::{json, Value};
16use std::collections::HashMap;
17use std::net::SocketAddr;
18use std::sync::Arc;
19use std::time::Duration;
20use tokio::net::UdpSocket;
21use tokio::sync::mpsc;
22
23/// A handler chain (answer → event → answer …) stops after this many follow-ups.
24pub const MAX_FOLLOWUP_DEPTH: usize = 8;
25/// A stream with nothing for this long has ended.
26pub const STREAM_IDLE: Duration = Duration::from_secs(1);
27/// Streams tracked at once; past it new SSRCs are ignored until one ends.
28pub const MAX_STREAMS: usize = 64;
29/// Decoded samples kept per stream for the tone estimate (ten seconds at 8 kHz).
30pub const MAX_ANALYSED_SAMPLES: usize = 80_000;
31const FRAME: Duration = Duration::from_millis(20);
32
33/// G.711 µ-law to linear PCM (ITU-T G.711).
34pub fn ulaw_to_linear(u: u8) -> i16 {
35    let u = !u;
36    let exponent = (u >> 4) & 0x07;
37    let mantissa = i32::from(u & 0x0F);
38    let magnitude = (((mantissa << 3) + 0x84) << exponent) - 0x84;
39    (if u & 0x80 != 0 { -magnitude } else { magnitude }) as i16
40}
41
42/// G.711 A-law to linear PCM (ITU-T G.711).
43pub fn alaw_to_linear(a: u8) -> i16 {
44    let a = a ^ 0x55;
45    let exponent = (a >> 4) & 0x07;
46    let mantissa = i32::from(a & 0x0F);
47    let magnitude = if exponent == 0 {
48        (mantissa << 4) + 8
49    } else {
50        ((mantissa << 4) + 0x108) << (exponent - 1)
51    };
52    (if a & 0x80 != 0 { magnitude } else { -magnitude }) as i16
53}
54
55/// The dominant frequency of a signal from its zero crossings (with a small hysteresis so
56/// quantisation noise around zero does not count), and its RMS level in dBFS.
57pub fn analyse(samples: &[i16], rate: f64) -> (Option<f64>, f64) {
58    if samples.is_empty() {
59        return (None, f64::NEG_INFINITY);
60    }
61    let rms =
62        (samples.iter().map(|s| f64::from(*s).powi(2)).sum::<f64>() / samples.len() as f64).sqrt();
63    let level = 20.0 * (rms / 32768.0).max(1e-9).log10();
64    let threshold = (rms * 0.1).max(64.0);
65    let (mut crossings, mut state) = (0u64, 0i8);
66    for s in samples {
67        let v = f64::from(*s);
68        let now = if v > threshold {
69            1
70        } else if v < -threshold {
71            -1
72        } else {
73            state
74        };
75        if state != 0 && now != state {
76            crossings += 1;
77        }
78        state = now;
79    }
80    let seconds = samples.len() as f64 / rate;
81    let tone = (crossings > 4).then(|| crossings as f64 / 2.0 / seconds);
82    (tone, level)
83}
84
85struct Stream {
86    payload_type: u8,
87    first_seq: u16,
88    last_seq: u16,
89    first_ts: u32,
90    last_ts: u32,
91    last_frame_len: usize,
92    packets: u64,
93    octets: u64,
94    samples: Vec<i16>,
95    last_seen: crate::utils::clock::Instant,
96}
97
98/// Inbound RTP accounting per SSRC, shared with the RTSP client: a start event for a new
99/// stream, and an end event (loss, length, decoded tone and level) once it goes quiet.
100#[derive(Default)]
101pub struct Tracker {
102    streams: HashMap<u32, Stream>,
103}
104
105impl Tracker {
106    /// Account one RTP datagram; the `rtp_stream_started` data when it opens a stream.
107    pub fn on_rtp(&mut self, data: &[u8], from: SocketAddr) -> Option<Value> {
108        let h = media::parse_rtp(data)?;
109        let payload = &data[data.len() - h.payload_len..];
110        let now = crate::utils::clock::Instant::now();
111        let mut started = None;
112        if !self.streams.contains_key(&h.ssrc) {
113            if self.streams.len() >= MAX_STREAMS {
114                return None;
115            }
116            let codec = match h.payload_type {
117                0 => "pcmu",
118                8 => "pcma",
119                _ => "unknown",
120            };
121            started = Some(
122                json!({"ssrc": h.ssrc, "payload_type": h.payload_type, "codec": codec, "from": from.to_string()}),
123            );
124            self.streams.insert(
125                h.ssrc,
126                Stream {
127                    payload_type: h.payload_type,
128                    first_seq: h.sequence,
129                    last_seq: h.sequence,
130                    first_ts: h.timestamp,
131                    last_ts: h.timestamp,
132                    last_frame_len: 0,
133                    packets: 0,
134                    octets: 0,
135                    samples: Vec::new(),
136                    last_seen: now,
137                },
138            );
139        }
140        if let Some(s) = self.streams.get_mut(&h.ssrc) {
141            s.packets += 1;
142            s.octets += h.payload_len as u64;
143            s.last_seen = now;
144            // Sequence numbers wrap; a packet "after" the last within half the space is newer.
145            if h.sequence.wrapping_sub(s.last_seq) < 0x8000 {
146                s.last_seq = h.sequence;
147                s.last_ts = h.timestamp;
148                s.last_frame_len = h.payload_len;
149            }
150            if s.samples.len() < MAX_ANALYSED_SAMPLES {
151                let decode: Option<fn(u8) -> i16> = match s.payload_type {
152                    0 => Some(ulaw_to_linear),
153                    8 => Some(alaw_to_linear),
154                    _ => None,
155                };
156                if let Some(d) = decode {
157                    s.samples.extend(payload.iter().map(|b| d(*b)));
158                }
159            }
160        }
161        started
162    }
163
164    /// The `rtp_stream_ended` data of every stream quiet for `STREAM_IDLE`, which are dropped.
165    pub fn sweep(&mut self) -> Vec<Value> {
166        let now = crate::utils::clock::Instant::now();
167        let done: Vec<u32> = self
168            .streams
169            .iter()
170            .filter(|(_, s)| now.duration_since(s.last_seen) >= STREAM_IDLE)
171            .map(|(k, _)| *k)
172            .collect();
173        done.into_iter()
174            .filter_map(|ssrc| self.streams.remove(&ssrc).map(|s| ended(ssrc, &s)))
175            .collect()
176    }
177}
178
179fn ended(ssrc: u32, s: &Stream) -> Value {
180    let expected = u64::from(s.last_seq.wrapping_sub(s.first_seq)) + 1;
181    let samples = u64::from(s.last_ts.wrapping_sub(s.first_ts)) + s.last_frame_len as u64;
182    let mut data = json!({"ssrc": ssrc, "packets": s.packets, "lost": expected.saturating_sub(s.packets),
183                          "octets": s.octets, "duration_ms": samples * 1000 / u64::from(media::G711_CLOCK_HZ)});
184    if !s.samples.is_empty() {
185        let (tone, level) = analyse(&s.samples, f64::from(media::G711_CLOCK_HZ));
186        data["tone_hz"] = json!(tone.map(|t| t.round()));
187        data["level_dbfs"] = json!((level * 10.0).round() / 10.0);
188    }
189    data
190}
191
192pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
193    let listen = ctx
194        .startup_params
195        .as_ref()
196        .map(|p| p.get_optional_string("listen"))
197        .transpose()?
198        .flatten()
199        .unwrap_or_else(|| actions::DEFAULT_LISTEN.to_string());
200    let remote: SocketAddr = tokio::net::lookup_host(&ctx.remote_addr)
201        .await
202        .with_context(|| format!("cannot resolve {}", ctx.remote_addr))?
203        .next()
204        .context("remote_addr resolves to nothing")?;
205    let socket = Arc::new(
206        UdpSocket::bind(&listen)
207            .await
208            .with_context(|| format!("cannot bind {listen}"))?,
209    );
210    let local = socket.local_addr()?;
211    ctx.state
212        .update_client_status(ctx.client_id, ClientStatus::Connected)
213        .await;
214    let external =
215        crate::client::command_support::register_command_channel(&ctx.state, ctx.client_id).await;
216    let (internal_tx, internal_rx) = mpsc::channel::<(Value, usize)>(32);
217    let (event_tx, mut event_rx) = mpsc::channel::<(Event, usize)>(64);
218    event_tx.try_send((
219        Event::new(
220            &actions::READY_EVENT,
221            json!({"local_addr": local.to_string(), "remote_addr": remote.to_string()}),
222        ),
223        0,
224    ))?;
225    let events_ctx = ctx.clone();
226    let dispatcher = tokio::spawn(async move {
227        while let Some((event, depth)) = event_rx.recv().await {
228            events_ctx
229                .state
230                .record_access_log(
231                    AccessLogOwner::Client(events_ctx.client_id.as_u32()),
232                    "RTP",
233                    None,
234                    event.id(),
235                    event.data.clone(),
236                    vec![],
237                )
238                .await;
239            let instruction = events_ctx
240                .state
241                .get_instruction_for_client(events_ctx.client_id)
242                .await
243                .unwrap_or_default();
244            let memory = events_ctx
245                .state
246                .get_memory_for_client(events_ctx.client_id)
247                .await
248                .unwrap_or_default();
249            match call_llm_for_client(
250                &events_ctx.llm_client,
251                &events_ctx.state,
252                events_ctx.client_id.to_string(),
253                &instruction,
254                &memory,
255                Some(&event),
256                &RtpClientProtocol,
257                &events_ctx.status_tx,
258            )
259            .await
260            {
261                Ok(result) => {
262                    if let Some(memory) = result.memory_updates {
263                        events_ctx
264                            .state
265                            .set_memory_for_client(events_ctx.client_id, memory)
266                            .await;
267                    }
268                    for action in result.actions {
269                        if internal_tx.send((action, depth + 1)).await.is_err() {
270                            return;
271                        }
272                    }
273                }
274                Err(e) => {
275                    Log::new(Some(&events_ctx.status_tx)).warn(format!("RTP client handler: {e}"))
276                }
277            }
278        }
279    });
280    let dispatcher_abort = dispatcher.abort_handle();
281    ctx.state
282        .register_client_task(ctx.client_id, dispatcher)
283        .await;
284    let session_ctx = ctx.clone();
285    let task = tokio::spawn(async move {
286        session(
287            &session_ctx,
288            socket,
289            remote,
290            external,
291            internal_rx,
292            event_tx,
293        )
294        .await;
295        dispatcher_abort.abort();
296        session_ctx
297            .state
298            .update_client_status(session_ctx.client_id, ClientStatus::Disconnected)
299            .await;
300        session_ctx
301            .state
302            .remove_client_handle(session_ctx.client_id)
303            .await;
304        let _ = session_ctx.status_tx.send("__UPDATE_UI__".into());
305    });
306    ctx.state.register_client_task(ctx.client_id, task).await;
307    Ok(local)
308}
309
310/// Send one stream, paced at its frame rate; report it when done or cut short.
311async fn send_stream(
312    socket: Arc<UdpSocket>,
313    remote: SocketAddr,
314    action: Value,
315    depth: usize,
316    events: mpsc::Sender<(Event, usize)>,
317    cut: tokio::sync::watch::Receiver<u64>,
318    generation: u64,
319) {
320    let codec = action
321        .get("payload_type")
322        .and_then(Value::as_str)
323        .map(AudioCodec::parse)
324        .unwrap_or(Ok(AudioCodec::Pcmu))
325        .unwrap_or(AudioCodec::Pcmu);
326    let Ok(content) = media::parse_audio_content(&action) else {
327        return;
328    };
329    let Ok(bytes) = media::synthesize(
330        codec,
331        &content,
332        action["duration_ms"].as_u64().unwrap_or(1000),
333    ) else {
334        return;
335    };
336    let ssrc = action["ssrc"]
337        .as_u64()
338        .map(|s| s as u32)
339        .unwrap_or_else(rand::random);
340    let mut packetizer = media::RtpPacketizer::new(ssrc, codec.payload_type(), None, None);
341    let packets = packetizer.packetize(&bytes, media::G711_SAMPLES_PER_FRAME);
342    let mut interval = tokio::time::interval(FRAME);
343    let mut sent = 0u64;
344    let mut complete = true;
345    for p in &packets {
346        interval.tick().await;
347        if *cut.borrow() != generation {
348            complete = false;
349            break;
350        }
351        if socket.send_to(p, remote).await.is_err() {
352            complete = false;
353            break;
354        }
355        sent += 1;
356    }
357    let data =
358        json!({"ssrc": ssrc, "packets": sent, "duration_ms": sent * 20, "complete": complete});
359    let _ = events.try_send((Event::new(&actions::SENT_EVENT, data), depth));
360}
361
362async fn session(
363    ctx: &ConnectContext,
364    socket: Arc<UdpSocket>,
365    remote: SocketAddr,
366    mut external: mpsc::Receiver<ClientCommand>,
367    mut internal: mpsc::Receiver<(Value, usize)>,
368    events: mpsc::Sender<(Event, usize)>,
369) {
370    let log = Log::new(Some(&ctx.status_tx));
371    let mut tracker = Tracker::default();
372    let mut buf = vec![0u8; 2048];
373    let mut tick = tokio::time::interval(Duration::from_millis(250));
374    let (cut_tx, cut_rx) = tokio::sync::watch::channel(0u64);
375    let mut generation = 0u64;
376    loop {
377        let (action, depth, mut injected) = tokio::select! {
378            r = socket.recv_from(&mut buf) => {
379                let Ok((n, from)) = r else { return };
380                let data = &buf[..n];
381                if media::is_rtcp(data) {
382                    let ssrc = data.get(4..8).map(|b| u32::from_be_bytes([b[0], b[1], b[2], b[3]])).unwrap_or(0);
383                    let _ = events.try_send((Event::new(&actions::RTCP_EVENT, json!({"packet_type": data[1], "ssrc": ssrc, "from": from.to_string()})), 0));
384                    continue;
385                }
386                if let Some(started) = tracker.on_rtp(data, from) {
387                    let _ = events.try_send((Event::new(&actions::STREAM_STARTED_EVENT, started), 0));
388                }
389                continue;
390            }
391            _ = tick.tick() => {
392                for data in tracker.sweep() {
393                    let _ = events.try_send((Event::new(&actions::STREAM_ENDED_EVENT, data), 0));
394                }
395                continue;
396            }
397            command = external.recv() => match command {
398                Some(c) => (c.action.clone(), 0, Some(c)),
399                None => return,
400            },
401            action = internal.recv() => match action {
402                Some((a, depth)) => (a, depth, None),
403                None => return,
404            },
405        };
406        let reply = |injected: &mut Option<ClientCommand>, outcome: ClientSendOutcome| {
407            if let Some(command) = injected.take() {
408                crate::client::command_support::reply(command, Ok(outcome));
409            }
410        };
411        match RtpClientProtocol.execute_action(action.clone()) {
412            Ok(ClientActionResult::Disconnect) => {
413                let _ = cut_tx.send(u64::MAX);
414                reply(&mut injected, ClientSendOutcome::Disconnected);
415                return;
416            }
417            Ok(_) => {}
418            Err(e) => {
419                log.warn(format!("RTP client action refused: {e}"));
420                reply(
421                    &mut injected,
422                    ClientSendOutcome::Rejected {
423                        error: e.to_string(),
424                    },
425                );
426                continue;
427            }
428        }
429        if depth > MAX_FOLLOWUP_DEPTH {
430            log.warn(format!(
431                "RTP client: handler chain stopped after {MAX_FOLLOWUP_DEPTH} follow-ups"
432            ));
433            continue;
434        }
435        if injected.is_some() {
436            ctx.state
437                .record_access_log(
438                    AccessLogOwner::Client(ctx.client_id.as_u32()),
439                    "RTP",
440                    None,
441                    "injected_action",
442                    action.clone(),
443                    vec![],
444                )
445                .await;
446        }
447        match action["type"].as_str().unwrap_or_default() {
448            actions::SEND_AUDIO => {
449                generation += 1;
450                let _ = cut_tx.send(generation);
451                let task = tokio::spawn(send_stream(
452                    socket.clone(),
453                    remote,
454                    action,
455                    depth,
456                    events.clone(),
457                    cut_rx.clone(),
458                    generation,
459                ));
460                ctx.state.register_client_task(ctx.client_id, task).await;
461                reply(
462                    &mut injected,
463                    ClientSendOutcome::Executed {
464                        detail: "streaming".into(),
465                    },
466                );
467            }
468            _ => {
469                let sr = media::build_rtcp_sender_report(
470                    action["ssrc"]
471                        .as_u64()
472                        .map(|s| s as u32)
473                        .unwrap_or_else(rand::random),
474                    action["rtp_timestamp"].as_u64().unwrap_or(0) as u32,
475                    action["packet_count"].as_u64().unwrap_or(0) as u32,
476                    action["octet_count"].as_u64().unwrap_or(0) as u32,
477                );
478                let outcome = match socket.send_to(&sr, remote).await {
479                    Ok(n) => ClientSendOutcome::Sent { bytes_sent: n },
480                    Err(e) => ClientSendOutcome::Rejected {
481                        error: e.to_string(),
482                    },
483                };
484                reply(&mut injected, outcome);
485            }
486        }
487    }
488}