netget/server/srt/
mod.rs

1//! SRT listener over srt-tokio (ARQ, TSBPD, optional AES). Rust owns the handshake, stream-ID
2//! parsing, one publisher per resource and the relay to every reader of it; the handler admits
3//! each caller and can push text messages to a reader.
4pub mod actions;
5pub mod streamid;
6
7use crate::llm::action_helper::call_llm;
8use crate::llm::actions::protocol_trait::ActionResult;
9use crate::logging::emit::Log;
10use crate::protocol::{Event, SpawnContext};
11use crate::server::connection::ConnectionId;
12use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
13use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
14use crate::utils::clock::Instant;
15use anyhow::Result;
16use bytes::Bytes;
17use futures::{SinkExt, StreamExt};
18use serde_json::{json, Value};
19use srt_tokio::access::{RejectReason, ServerRejectReason};
20use srt_tokio::{SocketStatistics, SrtListener, SrtSocket};
21use std::collections::HashMap;
22use std::net::SocketAddr;
23use std::sync::{Arc, Mutex};
24use std::time::Duration;
25use tokio::sync::mpsc;
26
27pub const DEFAULT_LATENCY: Duration = Duration::from_millis(120);
28pub const MAX_READERS: usize = 64;
29pub const IDLE_TIMEOUT: Duration = Duration::from_secs(10);
30const READER_QUEUE: usize = 4096;
31
32/// Handler-facing reject reasons and their SRT server reject codes.
33pub const REJECT_CODES: &[(&str, ServerRejectReason)] = &[
34    ("bad_request", ServerRejectReason::BadRequest),
35    ("unauthorized", ServerRejectReason::Unauthorized),
36    ("overload", ServerRejectReason::Overload),
37    ("forbidden", ServerRejectReason::Forbidden),
38    ("not_found", ServerRejectReason::Notfound),
39    ("conflict", ServerRejectReason::Conflict),
40];
41
42#[derive(Default)]
43struct Resource {
44    publisher: Option<ConnectionId>,
45    readers: HashMap<ConnectionId, mpsc::Sender<Bytes>>,
46}
47
48struct Shared {
49    ctx: SpawnContext,
50    resources: Mutex<HashMap<String, Resource>>,
51    idle: Duration,
52}
53
54/// srt-tokio 0.4.4 announces its live payload size (1316) in the handshake's MSS field, where
55/// libsrt expects the MTU, so a libsrt peer limits itself to 1272-byte messages and cannot send
56/// the standard 1316 (seven TS packets). Raising the payload size to the live maximum, 1456, makes
57/// the announced MSS large enough for that; NetGet's own messages are 1316 bytes either way.
58pub fn announce_mss(o: &mut srt_tokio::options::SocketOptions) {
59    o.sender.max_payload_size = srt_tokio::options::PacketSize(1456);
60}
61
62pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
63    let p = ctx.startup_params.as_ref();
64    let latency_ms = p
65        .map(|p| p.get_optional_u64("latency_ms"))
66        .transpose()?
67        .flatten()
68        .unwrap_or(DEFAULT_LATENCY.as_millis() as u64);
69    anyhow::ensure!(
70        (20..=8000).contains(&latency_ms),
71        "latency_ms must be 20..=8000"
72    );
73    let idle = p
74        .map(|p| p.get_optional_u64("idle_timeout_secs"))
75        .transpose()?
76        .flatten()
77        .unwrap_or(IDLE_TIMEOUT.as_secs());
78    anyhow::ensure!(
79        (1..=3600).contains(&idle),
80        "idle_timeout_secs must be 1..=3600"
81    );
82    let passphrase = p
83        .map(|p| p.get_optional_string("passphrase"))
84        .transpose()?
85        .flatten();
86    if let Some(pw) = &passphrase {
87        anyhow::ensure!(
88            (10..=79).contains(&pw.len()),
89            "passphrase must be 10 to 79 characters (SRT)"
90        );
91    }
92    let mut builder = SrtListener::builder()
93        .latency(Duration::from_millis(latency_ms))
94        .set(announce_mss);
95    if let Some(pw) = passphrase {
96        builder = builder.encryption(16, pw);
97    }
98    // Bound here so the real address (port 0 included) is known; srt-tokio does not report it.
99    let udp = tokio::net::UdpSocket::bind(ctx.legacy_listen_addr()).await?;
100    let local = udp.local_addr()?;
101    let (listener, mut incoming) = builder.socket(udp).bind(local).await?;
102    Log::new(Some(&ctx.status_tx)).info(format!(
103        "SRT listener at srt://{local} (latency {latency_ms} ms)"
104    ));
105    let shared = Arc::new(Shared {
106        ctx: ctx.clone(),
107        resources: Mutex::default(),
108        idle: Duration::from_secs(idle),
109    });
110    let server_id = ctx.server_id;
111    let accept = tokio::spawn(async move {
112        // The listener closes when dropped; it lives as long as this task.
113        let _listener = listener;
114        while let Some(request) = incoming.incoming().next().await {
115            let id = ConnectionId::new(shared.ctx.state.get_next_unified_id().await);
116            let now = Instant::now();
117            shared
118                .ctx
119                .state
120                .add_connection_to_server(
121                    server_id,
122                    ConnectionState {
123                        id,
124                        remote_addr: request.remote(),
125                        local_addr: local,
126                        bytes_sent: 0,
127                        bytes_received: 0,
128                        packets_sent: 0,
129                        packets_received: 0,
130                        last_activity: now,
131                        status: ConnectionStatus::Active,
132                        status_changed_at: now,
133                        protocol_info: ProtocolConnectionInfo::empty(),
134                    },
135                )
136                .await;
137            let child = shared.clone();
138            shared
139                .ctx
140                .state
141                .spawn_server_task(server_id, async move {
142                    if let Err(e) = connection(&child, id, request).await {
143                        Log::new(Some(&child.ctx.status_tx))
144                            .debug(format!("SRT connection {id}: {e:#}"));
145                    }
146                    child
147                        .ctx
148                        .state
149                        .remove_peer_handle(server_id, id.as_u32())
150                        .await;
151                    child
152                        .ctx
153                        .state
154                        .update_connection_status(server_id, id, ConnectionStatus::Closed)
155                        .await;
156                    let _ = child.ctx.status_tx.send("__UPDATE_UI__".into());
157                })
158                .await;
159        }
160    });
161    ctx.state.register_server_task(server_id, accept).await;
162    Ok(local)
163}
164
165fn outcome(ctx: &SpawnContext, id: ConnectionId, operation: &str, decision: &str) {
166    let summary = format!("SRT connection {id} operation={operation} decision={decision}");
167    let log = Log::new(Some(&ctx.status_tx));
168    if decision.starts_with("fail_closed_") || decision == "model_silent" {
169        log.error(summary);
170    } else {
171        log.info(summary);
172    }
173}
174
175/// Ask the handler; Ok(None) admits, Ok(Some(reason)) refuses, Err is a failure.
176async fn ask(
177    shared: &Shared,
178    id: ConnectionId,
179    event: Event,
180    operation: &str,
181) -> Result<Option<ServerRejectReason>, ()> {
182    let ctx = &shared.ctx;
183    let result = match call_llm(
184        &ctx.llm_client,
185        &ctx.state,
186        ctx.server_id,
187        Some(id),
188        &event,
189        &actions::SrtProtocol,
190    )
191    .await
192    {
193        Ok(r) => r,
194        Err(_) => {
195            outcome(ctx, id, operation, "fail_closed_llm_error");
196            return Err(());
197        }
198    };
199    let mut answers = Vec::new();
200    let mut pending = result.protocol_results;
201    while let Some(r) = pending.pop() {
202        match r {
203            ActionResult::Custom { data, .. } => answers.push(data),
204            ActionResult::Multiple(items) => pending.extend(items),
205            _ => {}
206        }
207    }
208    match (result.failures.is_empty(), answers.as_slice()) {
209        (true, [a]) if a["type"] == "srt_accept" || a["type"] == "srt_ignore" => {
210            outcome(ctx, id, operation, "model_answer");
211            Ok(None)
212        }
213        (true, [a]) if a["type"] == "srt_reject" => {
214            outcome(ctx, id, operation, "model_reject");
215            let code = REJECT_CODES
216                .iter()
217                .find(|(n, _)| a["reason"] == *n)
218                .map(|(_, c)| *c)
219                .unwrap_or(ServerRejectReason::Forbidden);
220            Ok(Some(code))
221        }
222        (true, []) => {
223            outcome(ctx, id, operation, "model_silent");
224            Err(())
225        }
226        _ => {
227            outcome(ctx, id, operation, "fail_closed_invalid_reply");
228            Err(())
229        }
230    }
231}
232
233fn stats_json(s: &SocketStatistics) -> Value {
234    json!({
235        "rx_packets": s.rx_data, "rx_bytes": s.rx_bytes, "rx_lost_packets": s.rx_loss_data, "rx_retransmitted_packets": s.rx_retransmit_data,
236        "rx_dropped_packets": s.rx_dropped_data, "tx_packets": s.tx_data, "tx_bytes": s.tx_bytes, "tx_retransmitted_packets": s.tx_retransmit_data,
237        "tx_lost_packets": s.tx_loss_data, "rtt_ms": s.tx_average_rtt.as_millis() as u64,
238    })
239}
240
241async fn connection(
242    shared: &Shared,
243    id: ConnectionId,
244    request: srt_tokio::ConnectionRequest,
245) -> Result<()> {
246    let ctx = &shared.ctx;
247    let raw = request
248        .stream_id()
249        .map(|s| s.to_string())
250        .unwrap_or_default();
251    let target = match streamid::parse(&raw) {
252        Ok(t) if !t.resource.is_empty() => t,
253        _ => {
254            outcome(ctx, id, "connect", "protocol_refusal");
255            request
256                .reject(RejectReason::Server(ServerRejectReason::BadRequest))
257                .await?;
258            return Ok(());
259        }
260    };
261    let busy = target.mode == "publish"
262        && shared
263            .resources
264            .lock()
265            .unwrap_or_else(|e| e.into_inner())
266            .get(&target.resource)
267            .is_some_and(|r| r.publisher.is_some());
268    if busy {
269        outcome(ctx, id, "connect", "protocol_refusal");
270        request
271            .reject(RejectReason::Server(ServerRejectReason::Conflict))
272            .await?;
273        return Ok(());
274    }
275    let event = Event::new(
276        &actions::CONNECT_EVENT,
277        json!({"stream_id": raw, "resource": target.resource, "mode": target.mode, "user": target.user, "remote": request.remote().to_string()}),
278    );
279    match ask(shared, id, event, "connect").await {
280        Ok(None) => {}
281        Ok(Some(code)) => {
282            request.reject(RejectReason::Server(code)).await?;
283            return Ok(());
284        }
285        Err(()) => {
286            request
287                .reject(RejectReason::Server(
288                    ServerRejectReason::InternalServerError,
289                ))
290                .await?;
291            return Ok(());
292        }
293    }
294    // Claim the resource now that the caller is admitted (a racing publisher loses here).
295    let (tx, mut relayed) = mpsc::channel::<Bytes>(READER_QUEUE);
296    let refused = {
297        let mut resources = shared.resources.lock().unwrap_or_else(|e| e.into_inner());
298        let r = resources.entry(target.resource.clone()).or_default();
299        if target.mode == "publish" {
300            r.publisher
301                .is_some()
302                .then_some(ServerRejectReason::Conflict)
303                .or_else(|| {
304                    r.publisher = Some(id);
305                    None
306                })
307        } else if r.readers.len() >= MAX_READERS {
308            Some(ServerRejectReason::Overload)
309        } else {
310            r.readers.insert(id, tx);
311            None
312        }
313    };
314    if let Some(code) = refused {
315        request.reject(RejectReason::Server(code)).await?;
316        return Ok(());
317    }
318    let started = Instant::now();
319    let result = match request.accept(None).await {
320        Ok(socket) => session(shared, id, socket, &target, &mut relayed).await,
321        Err(e) => Err((e.into(), None)),
322    };
323    {
324        let mut resources = shared.resources.lock().unwrap_or_else(|e| e.into_inner());
325        if let Some(r) = resources.get_mut(&target.resource) {
326            if r.publisher == Some(id) {
327                r.publisher = None;
328            }
329            r.readers.remove(&id);
330            if r.publisher.is_none() && r.readers.is_empty() {
331                resources.remove(&target.resource);
332            }
333        }
334    }
335    let (reason, stats) = match result {
336        Ok((reason, stats)) => (reason, stats),
337        Err((e, stats)) => (format!("{e:#}"), stats),
338    };
339    let mut data = json!({"resource": target.resource, "mode": target.mode, "duration_ms": started.elapsed().as_millis() as u64, "reason": reason});
340    if let Some(s) = stats {
341        data["statistics"] = stats_json(&s);
342    }
343    let _ = ask(
344        shared,
345        id,
346        Event::new(&actions::CLOSED_EVENT, data),
347        "closed",
348    )
349    .await;
350    Ok(())
351}
352
353type Ended = std::result::Result<
354    (String, Option<SocketStatistics>),
355    (anyhow::Error, Option<SocketStatistics>),
356>;
357
358async fn session(
359    shared: &Shared,
360    id: ConnectionId,
361    mut socket: SrtSocket,
362    target: &streamid::Target,
363    relayed: &mut mpsc::Receiver<Bytes>,
364) -> Ended {
365    let ctx = &shared.ctx;
366    let mut stats_stream = socket.statistics().clone();
367    let mut latest: Option<SocketStatistics> = None;
368    let mut commands = if target.mode == "request" {
369        Some(
370            crate::server::peer_support::register_peer_channel(
371                &ctx.state,
372                ctx.server_id,
373                id.as_u32(),
374            )
375            .await,
376        )
377    } else {
378        None
379    };
380    let publishing = target.mode == "publish";
381    // A fixed deadline from the last data: a timeout around each read would restart whenever
382    // the statistics stream or a relay woke the loop, and never fire.
383    let mut last_data = tokio::time::Instant::now();
384    loop {
385        enum Wake {
386            Data(Option<std::io::Result<(std::time::Instant, Bytes)>>),
387            Relay(Option<Bytes>),
388            Stats(Option<SocketStatistics>),
389            Command(Option<ClientCommand>),
390            Idle,
391        }
392        let wake = {
393            let command = async {
394                match commands.as_mut() {
395                    Some(c) => c.recv().await,
396                    None => std::future::pending().await,
397                }
398            };
399            let relay = async {
400                if publishing {
401                    std::future::pending().await
402                } else {
403                    relayed.recv().await
404                }
405            };
406            tokio::select! {
407                d = socket.next() => Wake::Data(d),
408                _ = tokio::time::sleep_until(last_data + shared.idle), if publishing => Wake::Idle,
409                r = relay => Wake::Relay(r),
410                s = stats_stream.next() => Wake::Stats(s),
411                c = command => Wake::Command(c),
412            }
413        };
414        match wake {
415            Wake::Idle if publishing => {
416                return Ok((
417                    "the publisher sent nothing for the idle timeout".into(),
418                    latest,
419                ))
420            }
421            Wake::Idle => {}
422            Wake::Data(None) => return Ok(("the caller closed the connection".into(), latest)),
423            Wake::Data(Some(Err(e))) => return Err((e.into(), latest)),
424            Wake::Data(Some(Ok((_, data)))) => {
425                last_data = tokio::time::Instant::now();
426                ctx.state
427                    .update_connection_stats(
428                        ctx.server_id,
429                        id,
430                        Some(data.len() as u64),
431                        None,
432                        Some(1),
433                        None,
434                    )
435                    .await;
436                if publishing {
437                    let mut resources = shared.resources.lock().unwrap_or_else(|e| e.into_inner());
438                    if let Some(r) = resources.get_mut(&target.resource) {
439                        // A reader whose queue is full is too slow to keep.
440                        r.readers.retain(|_, tx| tx.try_send(data.clone()).is_ok());
441                    }
442                }
443            }
444            Wake::Relay(None) => {
445                return Ok(("the reader was dropped for falling behind".into(), latest))
446            }
447            Wake::Relay(Some(data)) => {
448                let n = data.len() as u64;
449                if let Err(e) = socket.send((std::time::Instant::now(), data)).await {
450                    return Err((e.into(), latest));
451                }
452                ctx.state
453                    .update_connection_stats(ctx.server_id, id, None, Some(n), None, Some(1))
454                    .await;
455            }
456            Wake::Stats(Some(s)) => latest = Some(s),
457            Wake::Stats(None) => {}
458            Wake::Command(None) => commands = None,
459            Wake::Command(Some(cmd)) => {
460                let a = cmd.action.clone();
461                let outcome = match a["type"].as_str() {
462                    Some("disconnect") => ClientSendOutcome::Disconnected,
463                    Some("srt_send_text") => match actions::check_answer(&a) {
464                        Err(e) => ClientSendOutcome::Rejected {
465                            error: e.to_string(),
466                        },
467                        Ok(()) => {
468                            let text =
469                                Bytes::from(a["text"].as_str().unwrap_or_default().to_owned());
470                            let n = text.len();
471                            match socket.send((std::time::Instant::now(), text)).await {
472                                Ok(()) => ClientSendOutcome::Sent { bytes_sent: n },
473                                Err(e) => ClientSendOutcome::Rejected {
474                                    error: e.to_string(),
475                                },
476                            }
477                        }
478                    },
479                    _ => ClientSendOutcome::Rejected {
480                        error: "an SRT reader accepts srt_send_text or disconnect".into(),
481                    },
482                };
483                ctx.state
484                    .record_access_log(
485                        crate::state::AccessLogOwner::Server(ctx.server_id.as_u32()),
486                        "SRT",
487                        Some(id.as_u32()),
488                        "injected_action",
489                        json!({"type": a["type"]}),
490                        vec![serde_json::to_value(&outcome).unwrap_or(Value::Null)],
491                    )
492                    .await;
493                let done = matches!(outcome, ClientSendOutcome::Disconnected);
494                let _ = cmd.reply_tx.send(Ok(outcome));
495                if done {
496                    let _ = socket.close().await;
497                    return Ok(("the operator disconnected the reader".into(), latest));
498                }
499            }
500        }
501    }
502}
netget/server/srt/mod.rs

netget/server/srt/
mod.rs
