netget/server/webtransport/
session.rs

1//! One established WebTransport session, shared by the server and the client: incoming streams
2//! are read to their end and datagrams taken as they come, each raises one event, and the
3//! handler's answer is carried out here. Handler turns are serialised per session.
4use super::actions::{validate, DATAGRAM_EVENT, STREAM_EVENT, STREAM_REPLY_EVENT};
5use crate::client::command_support;
6use crate::logging::emit::Log;
7use crate::protocol::Event;
8use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
9use anyhow::{bail, ensure, Context, Result};
10use futures::future::BoxFuture;
11use futures::{stream::FuturesUnordered, StreamExt};
12use serde_json::{json, Value};
13use std::sync::Arc;
14use std::time::Duration;
15use tokio::sync::mpsc;
16use wtransport::{Connection, RecvStream, SendStream, VarInt};
17
18/// One stream's whole content, either direction.
19pub const MAX_STREAM_BYTES: usize = 1024 * 1024;
20/// From a stream's first byte to its end.
21pub const STREAM_TIMEOUT: Duration = Duration::from_secs(30);
22/// Incoming streams, datagrams and injected actions being handled at once per session.
23pub const MAX_IN_FLIGHT: usize = 32;
24/// webtransport_open_bi → webtransport_stream_reply → webtransport_open_bi … stops here.
25pub const MAX_FOLLOWUP_DEPTH: usize = 4;
26/// Stream error codes NetGet resets with.
27pub const CODE_REFUSED: u32 = 0x1;
28pub const CODE_FAILED: u32 = 0x2;
29
30pub type AskFn = Arc<dyn Fn(Event) -> BoxFuture<'static, Result<Vec<Value>>> + Send + Sync>;
31pub type TrafficFn = Arc<dyn Fn(u64, u64) -> BoxFuture<'static, ()> + Send + Sync>;
32
33pub struct Session {
34    pub conn: Connection,
35    ask: AskFn,
36    traffic: TrafficFn,
37    status_tx: mpsc::UnboundedSender<String>,
38    label: &'static str,
39    turn: tokio::sync::Mutex<()>,
40}
41
42/// The bytes an action carries in `data`, decoded per `encoding`.
43pub fn payload(v: &Value) -> Result<Vec<u8>> {
44    let data = v["data"].as_str().context("data is text")?;
45    let bytes = match v["encoding"].as_str().unwrap_or("utf8") {
46        "utf8" => data.as_bytes().to_vec(),
47        "hex" => hex::decode(data).context("data is not hex")?,
48        other => bail!("encoding {other:?} is utf8 or hex"),
49    };
50    ensure!(
51        bytes.len() <= MAX_STREAM_BYTES,
52        "data exceeds {MAX_STREAM_BYTES} bytes"
53    );
54    Ok(bytes)
55}
56
57/// Bytes that arrived, as `(data, encoding)`: the text when it is UTF-8, hex otherwise.
58pub fn describe(bytes: &[u8]) -> (String, &'static str) {
59    match std::str::from_utf8(bytes) {
60        Ok(text) => (text.to_owned(), "utf8"),
61        Err(_) => (hex::encode(bytes), "hex"),
62    }
63}
64
65/// Read a stream to its end within the bounds; a stream past them is stopped.
66pub async fn read_all(mut recv: RecvStream) -> Result<Vec<u8>> {
67    let read = async {
68        let mut data = Vec::new();
69        let mut buf = vec![0u8; 16 * 1024];
70        while let Some(n) = recv.read(&mut buf).await? {
71            if n > MAX_STREAM_BYTES - data.len() {
72                return Ok(None);
73            }
74            data.extend_from_slice(&buf[..n]);
75        }
76        Ok::<_, anyhow::Error>(Some(data))
77    };
78    match tokio::time::timeout(STREAM_TIMEOUT, read).await {
79        Ok(Ok(Some(data))) => Ok(data),
80        Ok(Ok(None)) => {
81            recv.stop(VarInt::from_u32(CODE_REFUSED));
82            bail!("the stream exceeds {MAX_STREAM_BYTES} bytes")
83        }
84        Ok(Err(e)) => Err(e),
85        Err(_) => {
86            recv.stop(VarInt::from_u32(CODE_REFUSED));
87            bail!("the stream did not finish within {STREAM_TIMEOUT:?}")
88        }
89    }
90}
91
92async fn write_and_finish(send: &mut SendStream, bytes: &[u8]) -> Result<()> {
93    tokio::time::timeout(STREAM_TIMEOUT, async {
94        send.write_all(bytes).await?;
95        send.finish().await?;
96        Ok::<_, anyhow::Error>(())
97    })
98    .await
99    .context("the peer did not take the data in time")?
100}
101
102impl Session {
103    pub fn new(
104        conn: Connection,
105        ask: AskFn,
106        traffic: TrafficFn,
107        status_tx: mpsc::UnboundedSender<String>,
108        label: &'static str,
109    ) -> Arc<Self> {
110        Arc::new(Self {
111            conn,
112            ask,
113            traffic,
114            status_tx,
115            label,
116            turn: tokio::sync::Mutex::new(()),
117        })
118    }
119
120    fn log(&self) -> Log<'_> {
121        Log::new(Some(&self.status_tx))
122    }
123
124    /// One handler turn; turns on a session never overlap.
125    pub async fn ask(&self, event: Event) -> Result<Vec<Value>> {
126        let _turn = self.turn.lock().await;
127        (self.ask)(event).await
128    }
129
130    /// Serve the session until it closes; returns why it closed.
131    pub async fn run(
132        self: Arc<Self>,
133        initial: Vec<Value>,
134        first: Option<Event>,
135        mut commands: mpsc::Receiver<ClientCommand>,
136    ) -> String {
137        let mut work: FuturesUnordered<BoxFuture<'static, ()>> = FuturesUnordered::new();
138        if !initial.is_empty() {
139            work.push(self.clone().execute(initial, None, 0));
140        }
141        if let Some(event) = first {
142            let s = self.clone();
143            work.push(Box::pin(async move {
144                match s.ask(event).await {
145                    Ok(actions) => s.clone().execute(actions, None, 0).await,
146                    Err(e) => s.log().error(format!(
147                        "{} decision=fail_closed_llm_error; nothing sent: {e}",
148                        s.label
149                    )),
150                }
151            }));
152        }
153        let mut commands_open = true;
154        loop {
155            let room = work.len() < MAX_IN_FLIGHT;
156            tokio::select! {
157                r = self.conn.accept_bi(), if room => match r {
158                    Ok((send, recv)) => work.push(Box::pin(self.clone().incoming(Some(send), recv))),
159                    Err(e) => return e.to_string(),
160                },
161                r = self.conn.accept_uni(), if room => match r {
162                    Ok(recv) => work.push(Box::pin(self.clone().incoming(None, recv))),
163                    Err(e) => return e.to_string(),
164                },
165                r = self.conn.receive_datagram(), if room => match r {
166                    Ok(d) => work.push(Box::pin(self.clone().datagram(d.payload().to_vec()))),
167                    Err(e) => return e.to_string(),
168                },
169                c = commands.recv(), if commands_open && room => match c {
170                    Some(c) => {
171                        let checked = match c.action["type"].as_str() {
172                            Some("webtransport_send_datagram" | "webtransport_open_uni" | "webtransport_open_bi" | "webtransport_close") => validate(&c.action),
173                            _ => Err(anyhow::anyhow!("a live session takes webtransport_send_datagram, webtransport_open_uni, webtransport_open_bi or webtransport_close")),
174                        };
175                        match checked {
176                            Ok(()) => {
177                                let bytes_sent = payload(&c.action).map(|p| p.len()).unwrap_or(0);
178                                work.push(self.clone().execute(vec![c.action.clone()], None, 0));
179                                command_support::reply(c, Ok(ClientSendOutcome::Sent { bytes_sent }));
180                            }
181                            Err(e) => command_support::reply(c, Ok(ClientSendOutcome::Rejected { error: e.to_string() })),
182                        }
183                    }
184                    None => commands_open = false,
185                },
186                _ = work.next(), if !work.is_empty() => {}
187                e = self.conn.closed() => return e.to_string(),
188            }
189        }
190    }
191
192    async fn incoming(self: Arc<Self>, send: Option<SendStream>, recv: RecvStream) {
193        let stream_id = recv.id().into_u64();
194        let data = match read_all(recv).await {
195            Ok(data) => data,
196            Err(e) => {
197                self.log().warn(format!(
198                    "{} decision=protocol_refusal stream={stream_id}: {e}",
199                    self.label
200                ));
201                if let Some(mut send) = send {
202                    let _ = send.reset(VarInt::from_u32(CODE_REFUSED));
203                }
204                return;
205            }
206        };
207        (self.traffic)(data.len() as u64, 0).await;
208        let (text, encoding) = describe(&data);
209        let direction = if send.is_some() {
210            "bidirectional"
211        } else {
212            "unidirectional"
213        };
214        self.log().debug(format!(
215            "{} {direction} stream {stream_id}: {} bytes",
216            self.label,
217            data.len()
218        ));
219        let event = Event::new(
220            &STREAM_EVENT,
221            json!({"stream_id": stream_id, "direction": direction, "data": text, "encoding": encoding}),
222        );
223        match self.ask(event).await {
224            Ok(actions) => self.clone().execute(actions, send, 0).await,
225            Err(e) => {
226                self.log().error(format!(
227                    "{} decision=fail_closed_llm_error stream={stream_id}; resetting it: {e}",
228                    self.label
229                ));
230                if let Some(mut send) = send {
231                    let _ = send.reset(VarInt::from_u32(CODE_FAILED));
232                }
233            }
234        }
235    }
236
237    async fn datagram(self: Arc<Self>, data: Vec<u8>) {
238        (self.traffic)(data.len() as u64, 0).await;
239        let (text, encoding) = describe(&data);
240        let event = Event::new(&DATAGRAM_EVENT, json!({"data": text, "encoding": encoding}));
241        match self.ask(event).await {
242            Ok(actions) => self.clone().execute(actions, None, 0).await,
243            Err(e) => self.log().error(format!(
244                "{} decision=fail_closed_llm_error datagram; nothing sent: {e}",
245                self.label
246            )),
247        }
248    }
249
250    /// Carry out the handler's actions in order. `reply` is the stream a bidirectional
251    /// event arrived on; it is finished empty when nothing answers on it.
252    pub fn execute(
253        self: Arc<Self>,
254        actions: Vec<Value>,
255        mut reply: Option<SendStream>,
256        depth: usize,
257    ) -> BoxFuture<'static, ()> {
258        Box::pin(async move {
259            for a in actions {
260                let kind = a["type"].as_str().unwrap_or_default().to_owned();
261                if !kind.starts_with("webtransport_") {
262                    continue;
263                }
264                if let Err(e) = validate(&a) {
265                    self.log().warn(format!(
266                        "{} decision=fail_closed_invalid_reply {kind}: {e}",
267                        self.label
268                    ));
269                    continue;
270                }
271                if let Err(e) = self.clone().act(&kind, &a, &mut reply, depth).await {
272                    self.log()
273                        .warn(format!("{} {kind} failed: {e:#}", self.label));
274                }
275                if kind == "webtransport_close" {
276                    return;
277                }
278            }
279            if let Some(mut send) = reply {
280                self.log().info(format!(
281                    "{} decision=model_silent; finishing the stream empty",
282                    self.label
283                ));
284                let _ = send.finish().await;
285            }
286        })
287    }
288
289    async fn act(
290        self: Arc<Self>,
291        kind: &str,
292        a: &Value,
293        reply: &mut Option<SendStream>,
294        depth: usize,
295    ) -> Result<()> {
296        match kind {
297            "webtransport_reply" => {
298                let mut send = reply
299                    .take()
300                    .context("webtransport_reply answers a bidirectional stream, and there is none to answer")?;
301                let bytes = payload(a)?;
302                self.log().info(format!(
303                    "{} decision=model_answer; replying {} bytes",
304                    self.label,
305                    bytes.len()
306                ));
307                write_and_finish(&mut send, &bytes).await?;
308                (self.traffic)(0, bytes.len() as u64).await;
309            }
310            "webtransport_send_datagram" => {
311                let bytes = payload(a)?;
312                self.conn.send_datagram(&bytes)?;
313                (self.traffic)(0, bytes.len() as u64).await;
314            }
315            "webtransport_open_uni" => {
316                let bytes = payload(a)?;
317                let mut send = self.conn.open_uni().await?.await?;
318                write_and_finish(&mut send, &bytes).await?;
319                (self.traffic)(0, bytes.len() as u64).await;
320            }
321            "webtransport_open_bi" => {
322                let bytes = payload(a)?;
323                let (mut send, recv) = self.conn.open_bi().await?.await?;
324                let stream_id = send.id().into_u64();
325                write_and_finish(&mut send, &bytes).await?;
326                (self.traffic)(0, bytes.len() as u64).await;
327                let answer = read_all(recv).await?;
328                (self.traffic)(answer.len() as u64, 0).await;
329                if depth + 1 >= MAX_FOLLOWUP_DEPTH {
330                    self.log().warn(format!(
331                        "{} follow-up depth {MAX_FOLLOWUP_DEPTH} reached; the reply on stream {stream_id} is not raised",
332                        self.label
333                    ));
334                    return Ok(());
335                }
336                let (text, encoding) = describe(&answer);
337                let event = Event::new(
338                    &STREAM_REPLY_EVENT,
339                    json!({"stream_id": stream_id, "request": a["data"], "data": text, "encoding": encoding}),
340                );
341                let actions = self.ask(event).await?;
342                self.clone().execute(actions, None, depth + 1).await;
343            }
344            "webtransport_close" => {
345                let code = a["code"].as_u64().unwrap_or(0) as u32;
346                let reason = a["reason"].as_str().unwrap_or_default();
347                self.log().info(format!(
348                    "{} closing the session: code={code} reason={reason:?}",
349                    self.label
350                ));
351                self.conn.close(VarInt::from_u32(code), reason.as_bytes());
352            }
353            // Admission answers are taken before the session exists.
354            "webtransport_accept" | "webtransport_reject" => {}
355            other => bail!("unknown action {other}"),
356        }
357        Ok(())
358    }
359}
netget/server/webtransport/session.rs

netget/server/webtransport/
session.rs
