netget/client/quic/
mod.rs

1//! An authenticated reusable QUIC connection, with independent bounded streams.
2//! The registered owner polls all query and handler futures; no detached tasks
3//! survive client removal and a parked handler never prevents command injection.
4pub mod actions;
5use crate::client::{command_support, llm_budget::call_llm_for_client};
6use crate::llm::actions::{
7    client_trait::{Client, ClientActionResult},
8    protocol_trait::Protocol,
9};
10use crate::logging::emit::Log;
11use crate::protocol::{ConnectContext, Event};
12use crate::state::{
13    client_handles::{ClientCommand, ClientSendOutcome},
14    AccessLogOwner, ClientStatus,
15};
16use crate::utils::quic::*;
17use actions::{
18    QuicClientProtocol, QUIC_CONNECTED_EVENT, QUIC_DATAGRAM_EVENT, QUIC_ERROR_EVENT,
19    QUIC_RESPONSE_EVENT,
20};
21use anyhow::{ensure, Context, Result};
22use futures::{future::BoxFuture, stream::FuturesUnordered, FutureExt, StreamExt};
23use serde_json::{json, Value};
24use std::{net::SocketAddr, time::Duration};
25
26const MAX_FOLLOWUP_DEPTH: u8 = 4;
27/// Marks a handler turn a received DATAGRAM started, in the depth byte it carries.
28const FROM_DATAGRAM: u8 = 0x80;
29type HandlerFuture = BoxFuture<'static, (u8, Result<Vec<Value>>)>;
30type QueryFuture = BoxFuture<'static, (u8, Event)>;
31
32pub struct QuicClient;
33impl QuicClient {
34    pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
35        let params = ctx.startup_params.as_ref();
36        let exchange =
37            Duration::from_secs(bounded_parameter(params, "exchange_timeout_secs", 30, 300)?);
38        let alpn = params
39            .map(|p| p.get_optional_string("alpn"))
40            .transpose()?
41            .flatten()
42            .unwrap_or_else(|| "netget-quic".into());
43        let (endpoint, connection) =
44            crate::utils::quic::connect(&ctx.remote_addr, params, alpn.as_bytes(), 0).await?;
45        let local_addr = endpoint.0.local_addr()?;
46        ctx.state
47            .update_client_status(ctx.client_id, ClientStatus::Connected)
48            .await;
49        let command_rx = command_support::register_command_channel(&ctx.state, ctx.client_id).await;
50        let state = ctx.state.clone();
51        let id = ctx.client_id;
52        let task = tokio::spawn(Self::run(ctx, endpoint, connection, command_rx, exchange));
53        state.register_client_task(id, task).await;
54        Ok(local_addr)
55    }
56
57    async fn run(
58        ctx: ConnectContext,
59        _endpoint: EndpointGuard,
60        connection: ConnectionGuard,
61        mut commands: tokio::sync::mpsc::Receiver<ClientCommand>,
62        exchange: Duration,
63    ) {
64        let mut handlers: FuturesUnordered<HandlerFuture> = FuturesUnordered::new();
65        let mut queries: FuturesUnordered<QueryFuture> = FuturesUnordered::new();
66        handlers.push(handle_event(
67            ctx.clone(),
68            Event::new(
69                &QUIC_CONNECTED_EVENT,
70                json!({
71                    "remote_addr": ctx.remote_addr,
72                    "max_datagram_size": connection.0.max_datagram_size(),
73                }),
74            ),
75            0,
76        ));
77        let mut disconnect = false;
78        // Consecutive datagrams sent in answer to a received datagram. An echoing peer and a
79        // model that answers every datagram would otherwise trade them for as long as the LLM
80        // budget lasts; past MAX_FOLLOWUP_DEPTH such answers are dropped until the operator or
81        // a non-datagram event sends something.
82        let mut datagram_replies: u8 = 0;
83        while !disconnect {
84            tokio::select! {
85                command = commands.recv() => {
86                    let Some(command) = command else { break; };
87                    datagram_replies = 0;
88                    disconnect = enqueue(&ctx,&connection.0,exchange,&mut queries,command.action.clone(),Some(command),0).await;
89                }
90                Some((depth,event)) = queries.next(), if !queries.is_empty() => {
91                    // The final response still reaches the handler; only subsequent actions
92                    // are stopped at the depth boundary. A returned disconnect is honoured.
93                    if handlers.len() < MAX_STREAMS {
94                        handlers.push(handle_event(ctx.clone(),event,depth));
95                    } else { Log::new(Some(&ctx.status_tx)).warn("QUIC response handler limit reached"); }
96                }
97                Some((tagged,result)) = handlers.next(), if !handlers.is_empty() => {
98                    let from_datagram = tagged & FROM_DATAGRAM != 0;
99                    let depth = tagged & !FROM_DATAGRAM;
100                    match result {
101                        Ok(actions) if actions.len() <= MAX_STREAMS => {
102                            for action in actions {
103                                if action["type"] == "send_quic_datagram" {
104                                    if !from_datagram {
105                                        datagram_replies = 0;
106                                    } else if datagram_replies >= MAX_FOLLOWUP_DEPTH {
107                                        Log::new(Some(&ctx.status_tx)).warn("QUIC datagram reply limit reached");
108                                        continue;
109                                    } else {
110                                        datagram_replies += 1;
111                                    }
112                                }
113                                if depth >= MAX_FOLLOWUP_DEPTH && action["type"] != "disconnect" {
114                                    if action["type"] == "send_quic_data" || action["type"] == "send_quic_datagram" { Log::new(Some(&ctx.status_tx)).warn("QUIC follow-up limit reached"); }
115                                    continue;
116                                }
117                                if enqueue(&ctx,&connection.0,exchange,&mut queries,action,None,depth).await { disconnect = true; break; }
118                            }
119                        }
120                        Ok(_) => Log::new(Some(&ctx.status_tx)).warn("QUIC handler returned more than 32 actions"),
121                        Err(e) => Log::new(Some(&ctx.status_tx)).warn(format!("QUIC event handler failed: {e}")),
122                    }
123                }
124                // An RFC 9221 DATAGRAM: one more handler turn. While the handler set is full,
125                // datagrams wait in Quinn's bounded buffer and the oldest are dropped there.
126                datagram = connection.0.read_datagram(), if handlers.len() < MAX_STREAMS => {
127                    let Ok(datagram) = datagram else { break; };
128                    let (data, encoding) = crate::server::quic::actions::encode_quic_payload(&datagram);
129                    handlers.push(handle_event(
130                        ctx.clone(),
131                        Event::new(
132                            &QUIC_DATAGRAM_EVENT,
133                            json!({"data":data,"encoding":encoding,"bytes":datagram.len()}),
134                        ),
135                        FROM_DATAGRAM,
136                    ));
137                }
138                _ = connection.0.closed() => break,
139            }
140        }
141        drop(queries);
142        drop(handlers);
143        connection.0.close(0u32.into(), b"client disconnected");
144        ctx.state.remove_client_handle(ctx.client_id).await;
145        ctx.state
146            .update_client_status(ctx.client_id, ClientStatus::Disconnected)
147            .await;
148        let _ = ctx.status_tx.send("__UPDATE_UI__".into());
149    }
150}
151
152pub fn payload(data: &Value) -> Result<Vec<u8>> {
153    let text = data["data"].as_str().context("Missing data")?;
154    let bytes = crate::server::quic::actions::decode_quic_payload(text, data["encoding"].as_str())?;
155    ensure!(bytes.len() <= MAX_BYTES, "QUIC payload exceeds 1 MiB");
156    Ok(bytes)
157}
158
159fn handle_event(ctx: ConnectContext, event: Event, depth: u8) -> HandlerFuture {
160    async move {
161        let result = async {
162            let instruction = ctx
163                .state
164                .get_instruction_for_client(ctx.client_id)
165                .await
166                .unwrap_or_default();
167            let memory = ctx
168                .state
169                .get_memory_for_client(ctx.client_id)
170                .await
171                .unwrap_or_default();
172            let result = call_llm_for_client(
173                &ctx.llm_client,
174                &ctx.state,
175                ctx.client_id.to_string(),
176                &instruction,
177                &memory,
178                Some(&event),
179                &QuicClientProtocol::new(),
180                &ctx.status_tx,
181            )
182            .await?;
183            if let Some(memory) = result.memory_updates {
184                ctx.state.set_memory_for_client(ctx.client_id, memory).await;
185            }
186            Ok(result.actions)
187        }
188        .await;
189        (depth, result)
190    }
191    .boxed()
192}
193
194async fn enqueue(
195    ctx: &ConnectContext,
196    connection: &quinn::Connection,
197    deadline: Duration,
198    queries: &mut FuturesUnordered<QueryFuture>,
199    action: Value,
200    command: Option<ClientCommand>,
201    depth: u8,
202) -> bool {
203    let outcome = match QuicClientProtocol::new().execute_action(action.clone()) {
204        Ok(ClientActionResult::Disconnect) => Ok(ClientSendOutcome::Disconnected),
205        Ok(ClientActionResult::WaitForMore | ClientActionResult::NoAction) => {
206            Ok(ClientSendOutcome::Executed {
207                detail: "Waiting for more queries".into(),
208            })
209        }
210        Ok(ClientActionResult::Custom { name, data })
211            if name == crate::server::quic::actions::DATAGRAM_RESULT =>
212        {
213            match data.as_str().and_then(|h| hex::decode(h).ok()) {
214                Some(bytes) => {
215                    let n = bytes.len();
216                    connection
217                        .send_datagram(bytes::Bytes::from(bytes))
218                        .map(|()| ClientSendOutcome::Executed {
219                            detail: format!("QUIC datagram of {n} bytes sent"),
220                        })
221                        .map_err(|e| {
222                            anyhow::anyhow!(
223                                "QUIC datagram of {n} bytes not sent (peer limit {:?}): {e}",
224                                connection.max_datagram_size()
225                            )
226                        })
227                }
228                None => Err(anyhow::anyhow!("QUIC datagram payload was not hex")),
229            }
230        }
231        Ok(ClientActionResult::Custom { name, data }) if name == "quic_exchange" => {
232            if queries.len() >= MAX_STREAMS {
233                Err(anyhow::anyhow!("QUIC client busy: 32 active transactions"))
234            } else {
235                queries.push(query_future(
236                    ctx.clone(),
237                    connection.clone(),
238                    deadline,
239                    data,
240                    action,
241                    command,
242                    depth + 1,
243                ));
244                return false;
245            }
246        }
247        Ok(_) => Err(anyhow::anyhow!("Unsupported QUIC client action")),
248        Err(e) => Err(e),
249    };
250    let disconnect = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
251    record_action(ctx, &action, &outcome).await;
252    if let Some(command) = command {
253        command_support::reply(command, outcome);
254    }
255    disconnect
256}
257
258async fn record_action(ctx: &ConnectContext, action: &Value, outcome: &Result<ClientSendOutcome>) {
259    let result = match outcome {
260        Ok(v) => serde_json::to_value(v).unwrap_or(Value::Null),
261        Err(e) => json!({"error":e.to_string()}),
262    };
263    ctx.state
264        .record_access_log(
265            AccessLogOwner::Client(ctx.client_id.as_u32()),
266            QuicClientProtocol::new().protocol_name(),
267            None,
268            "client_action",
269            action.clone(),
270            vec![result],
271        )
272        .await;
273}
274
275fn query_future(
276    ctx: ConnectContext,
277    connection: quinn::Connection,
278    deadline: Duration,
279    data: Value,
280    action: Value,
281    command: Option<ClientCommand>,
282    depth: u8,
283) -> QueryFuture {
284    async move {
285        let result = async {
286            let query = payload(&data)?;
287            exchange(&connection, &query, deadline).await
288        }
289        .await;
290        let (outcome, event) = match result {
291            Ok((bytes, stream)) => {
292                let (data, encoding) = crate::server::quic::actions::encode_quic_payload(&bytes);
293                let event = Event::new(
294                    &QUIC_RESPONSE_EVENT,
295                    json!({"stream_id":stream,"data":data,"encoding":encoding,"bytes":bytes.len()}),
296                );
297                (
298                    Ok(ClientSendOutcome::Executed {
299                        detail: format!("QUIC stream {stream}: {} response bytes", bytes.len()),
300                    }),
301                    event,
302                )
303            }
304            Err(e) => {
305                Log::new(Some(&ctx.status_tx)).warn(format!("QUIC stream failed: {e:#}"));
306                let event = Event::new(&QUIC_ERROR_EVENT, json!({"error":format!("{e:#}")}));
307                (Err(e), event)
308            }
309        };
310        record_action(&ctx, &action, &outcome).await;
311        if let Some(command) = command {
312            command_support::reply(command, outcome);
313        }
314        (depth, event)
315    }
316    .boxed()
317}
318
319/// One raw bidirectional exchange: write payload and FIN, then read until peer FIN.
320pub async fn exchange(
321    connection: &quinn::Connection,
322    payload: &[u8],
323    deadline: Duration,
324) -> Result<(Vec<u8>, u64)> {
325    ensure!(payload.len() <= MAX_BYTES, "QUIC payload exceeds 1 MiB");
326    let end = tokio::time::Instant::now() + deadline;
327    let (send, recv) = tokio::time::timeout_at(end, connection.open_bi())
328        .await
329        .context("QUIC stream-open deadline")??;
330    let mut streams = CancelStreams {
331        send,
332        recv,
333        done: false,
334    };
335    let result = tokio::time::timeout_at(end, async {
336        streams.send.write_all(payload).await?;
337        streams.send.finish()?;
338        let response = streams.recv.read_to_end(MAX_BYTES).await?;
339        Ok((response, u64::from(streams.recv.id())))
340    })
341    .await
342    .context("QUIC stream deadline exceeded")?;
343    if result.is_ok() {
344        streams.done = true;
345    }
346    result
347}
348struct CancelStreams {
349    send: quinn::SendStream,
350    recv: quinn::RecvStream,
351    done: bool,
352}
353impl Drop for CancelStreams {
354    fn drop(&mut self) {
355        if !self.done {
356            let _ = self.recv.stop(quinn::VarInt::from_u32(3));
357            let _ = self.send.reset(quinn::VarInt::from_u32(3));
358        }
359    }
360}
netget/client/quic/mod.rs

netget/client/quic/
mod.rs
