sova_quic/
datagram.rs

1use bytes::Bytes;
2use sova_core::extend::{wait_shutdown, BoxFuture as SovaBoxFuture, StateMap};
3use sova_core::{BackgroundService, Error, Result, Shutdown};
4use std::net::SocketAddr;
5use std::path::Path;
6use std::sync::Arc;
7use tokio::task::JoinSet;
8
9use crate::config::{build_client_config, build_server_config};
10use crate::endpoint::{accept_handshake, bind_server, prepare_incoming, PrepareIncoming};
11
12/// Handler invoked for each received QUIC datagram.
13pub type QuicDatagramHandler =
14    Arc<dyn Fn(SocketAddr, Vec<u8>, quinn::Connection) -> SovaBoxFuture<()> + Send + Sync>;
15
16/// QUIC server-side datagram listener (`BackgroundService`).
17pub struct QuicDatagramService {
18    bind_addr: SocketAddr,
19    tls: Arc<sova_core::Tls>,
20    alpn: Vec<Vec<u8>>,
21    handler: QuicDatagramHandler,
22    preshared_token: Option<Vec<u8>>,
23    name: String,
24}
25
26impl QuicDatagramService {
27    /// Build from an existing [`sova_core::Tls`] (share with HTTPS via [`Clone`]).
28    pub fn from_tls(
29        bind_addr: SocketAddr,
30        tls: sova_core::Tls,
31        alpn: Vec<Vec<u8>>,
32        retry: bool,
33        handler: QuicDatagramHandler,
34    ) -> Result<Self> {
35        if alpn.is_empty() {
36            return Err(Error::Internal(
37                "QuicDatagramService: `alpn` must not be empty".into(),
38            ));
39        }
40        if !retry {
41            tracing::debug!(
42                retry,
43                "quic anti-amplification retry=false is ignored; retry is always enabled"
44            );
45        }
46        Ok(Self {
47            bind_addr,
48            tls: Arc::new(tls),
49            alpn,
50            handler,
51            preshared_token: None,
52            name: format!("quic:{bind_addr}"),
53        })
54    }
55
56    pub fn from_pem(
57        bind_addr: SocketAddr,
58        cert_path: impl AsRef<Path>,
59        key_path: impl AsRef<Path>,
60        alpn: Vec<Vec<u8>>,
61        retry: bool,
62        handler: QuicDatagramHandler,
63    ) -> Result<Self> {
64        let tls = sova_core::Tls::from_pem(cert_path, key_path)?;
65        Self::from_tls(bind_addr, tls, alpn, retry, handler)
66    }
67
68    /// Self-signed certificate for local development (requires feature `dev-tls`).
69    #[cfg(feature = "dev-tls")]
70    pub fn self_signed(
71        bind_addr: SocketAddr,
72        hosts: &[&str],
73        alpn: Vec<Vec<u8>>,
74        retry: bool,
75        handler: QuicDatagramHandler,
76    ) -> Result<Self> {
77        let tls = sova_core::Tls::self_signed(hosts)?;
78        Self::from_tls(bind_addr, tls, alpn, retry, handler)
79    }
80
81    /// Optional preshared token inside datagrams: `token + b':' + payload`.
82    ///
83    /// If `token` is `None`, no validation is performed.
84    pub fn with_preshared_token(mut self, token: Option<String>) -> Self {
85        self.preshared_token = token.map(|t| t.into_bytes());
86        self
87    }
88
89    /// Register as a [`BackgroundService`] and attach TLS to the app HTTPS listener.
90    ///
91    /// Prefer this over `app.service(…)` + `bind(…).tls(…)` — one `Tls` handle wires both.
92    pub fn install(self, app: &mut sova_core::App) {
93        app.use_tls((*self.tls).clone());
94        app.service(self);
95    }
96}
97
98impl BackgroundService for QuicDatagramService {
99    fn name(&self) -> &str {
100        &self.name
101    }
102
103    fn run(
104        self: Box<Self>,
105        _state: Arc<StateMap>,
106        shutdown: Shutdown,
107    ) -> SovaBoxFuture<()> {
108        Box::pin(async move {
109            let server_config = match build_server_config(&self.tls, self.alpn.clone(), true) {
110                Ok(v) => v,
111                Err(e) => {
112                    tracing::error!(error = %e, "quic server TLS config build failed");
113                    return;
114                }
115            };
116
117            let endpoint = match bind_server(server_config, self.bind_addr, "quic datagram") {
118                Some(ep) => ep,
119                None => return,
120            };
121
122            tracing::info!(addr = %self.bind_addr, "quic datagram listening");
123
124            let mut recv_tasks: JoinSet<()> = JoinSet::new();
125            loop {
126                tokio::select! {
127                    _ = wait_shutdown(shutdown.clone()) => break,
128                    incoming = endpoint.accept() => {
129                        let Some(incoming) = incoming else { break; };
130
131                        let incoming = match prepare_incoming(incoming) {
132                            PrepareIncoming::RetrySent => continue,
133                            PrepareIncoming::Ready(i) => *i,
134                        };
135
136                        let conn = match accept_handshake(incoming, "quic handshake failed").await {
137                            Some(c) => c,
138                            None => continue,
139                        };
140
141                        let handler = Arc::clone(&self.handler);
142                        let token = self.preshared_token.clone();
143                        let shutdown = shutdown.clone();
144                        recv_tasks.spawn(async move {
145                            recv_loop(conn, handler, token, shutdown).await
146                        });
147                    }
148                }
149            }
150
151            // Stop accepting new connections; existing recv loops exit via shutdown.
152            endpoint.close(0u8.into(), b"shutdown");
153
154            while let Some(res) = recv_tasks.join_next().await {
155                if let Err(e) = res {
156                    tracing::debug!(error = %e, "quic recv task join error");
157                }
158            }
159        })
160    }
161}
162
163async fn recv_loop(
164    conn: quinn::Connection,
165    handler: QuicDatagramHandler,
166    preshared_token: Option<Vec<u8>>,
167    shutdown: Shutdown,
168) {
169    let peer = conn.remote_address();
170    loop {
171        tokio::select! {
172            _ = wait_shutdown(shutdown.clone()) => break,
173            res = conn.read_datagram() => {
174                let Ok(data) = res else {
175                    break;
176                };
177                let data = data.to_vec();
178
179                if let Some(max) = conn.max_datagram_size() {
180                    if data.len() > max {
181                        tracing::error!(
182                            peer = %peer,
183                            payload_len = data.len(),
184                            max_datagram_size = max,
185                            "quic datagram payload exceeds max_datagram_size; dropping"
186                        );
187                        continue;
188                    }
189                } else {
190                    tracing::debug!("quic connection has no max_datagram_size; skipping datagram");
191                    continue;
192                }
193
194                let payload = if let Some(expected) = preshared_token.as_deref() {
195                    if let Some(colon) = data.iter().position(|b| *b == b':') {
196                        if data[..colon] == expected[..] {
197                            data[colon + 1..].to_vec()
198                        } else {
199                            tracing::debug!("quic datagram token mismatch; dropping");
200                            continue;
201                        }
202                    } else {
203                        data
204                    }
205                } else {
206                    data
207                };
208
209                (handler)(peer, payload, conn.clone()).await;
210            }
211        }
212    }
213}
214
215/// QUIC client for sending datagrams and receiving a response.
216pub struct QuicDatagramClient {
217    _endpoint: quinn::Endpoint,
218    conn: quinn::Connection,
219}
220
221impl QuicDatagramClient {
222    /// Connect to a QUIC datagram server.
223    pub async fn connect(
224        bind_addr: SocketAddr,
225        server_addr: SocketAddr,
226        host: &str,
227        alpn: Vec<Vec<u8>>,
228    ) -> Result<Self> {
229        if alpn.is_empty() {
230            return Err(Error::Internal(
231                "QuicDatagramClient: `alpn` must not be empty".into(),
232            ));
233        }
234
235        let client_config = build_client_config(alpn)?;
236
237        let mut endpoint = quinn::Endpoint::client(bind_addr)
238            .map_err(|e| Error::Internal(format!("quic endpoint client bind failed: {e}")))?;
239        endpoint.set_default_client_config(client_config);
240
241        let connecting = endpoint
242            .connect(server_addr, host)
243            .map_err(|e| Error::Internal(format!("quic endpoint connect failed: {e}")))?;
244        let conn = connecting
245            .await
246            .map_err(|e| Error::Internal(format!("quic connect failed: {e}")))?;
247
248        Ok(Self {
249            _endpoint: endpoint,
250            conn,
251        })
252    }
253
254    pub fn max_datagram_size(&self) -> Option<usize> {
255        self.conn.max_datagram_size()
256    }
257
258    /// Returns peer certificate chain (leaf first) as raw DER bytes.
259    pub fn peer_certificates(&self) -> Option<Vec<Vec<u8>>> {
260        let identity = self.conn.peer_identity()?;
261        identity
262            .downcast::<Vec<rustls::pki_types::CertificateDer<'static>>>()
263            .ok()
264            .map(|certs| certs.into_iter().map(|c| c.as_ref().to_vec()).collect())
265    }
266
267    pub async fn send_datagram(&self, data: &[u8]) -> Result<()> {
268        self.conn
269            .send_datagram(Bytes::copy_from_slice(data))
270            .map_err(|e| Error::Internal(format!("send_datagram failed: {e}")))
271    }
272
273    pub async fn recv_datagram(&self) -> Result<Vec<u8>> {
274        let bytes = self
275            .conn
276            .read_datagram()
277            .await
278            .map_err(|e| Error::Internal(format!("recv_datagram failed: {e}")))?;
279        Ok(bytes.to_vec())
280    }
281
282    pub async fn send_and_recv(
283        &self,
284        data: &[u8],
285        timeout: std::time::Duration,
286    ) -> Result<Vec<u8>> {
287        self.send_datagram(data).await?;
288        tokio::time::timeout(timeout, self.recv_datagram())
289            .await
290            .map_err(|_| Error::Internal("timed out waiting for datagram".into()))?
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297    use sova_core::extend::StateMap;
298    use rcgen::generate_simple_self_signed;
299    use rustls::pki_types::CertificateDer;
300    use std::fs;
301    use tempfile::TempDir;
302
303    fn unused_local_addr() -> SocketAddr {
304        let sock = std::net::UdpSocket::bind("127.0.0.1:0").expect("bind udp");
305        sock.local_addr().expect("local addr")
306    }
307
308    fn write_self_signed_pem(tmp: &TempDir, subject: &[&str]) -> (std::path::PathBuf, std::path::PathBuf, Vec<u8>) {
309        let subject_alt_names: Vec<String> = subject.iter().map(|s| (*s).into()).collect();
310        let cert = generate_simple_self_signed(subject_alt_names).expect("rcgen");
311        let cert_pem = cert.cert.pem();
312        let key_pem = cert.key_pair.serialize_pem();
313
314        let cert_path = tmp.path().join("cert.pem");
315        let key_path = tmp.path().join("key.pem");
316        fs::write(&cert_path, cert_pem).expect("write cert");
317        fs::write(&key_path, key_pem).expect("write key");
318
319        let mut rd = std::io::BufReader::new(fs::File::open(&cert_path).expect("open cert pem"));
320        let certs = rustls_pemfile::certs(&mut rd).collect::<std::result::Result<Vec<_>, _>>().expect("parse certs");
321        let leaf: CertificateDer<'static> = certs.into_iter().next().expect("at least 1 cert");
322
323        (cert_path, key_path, leaf.as_ref().to_vec())
324    }
325
326    #[tokio::test]
327    async fn echo_roundtrip() {
328        let bind_addr = unused_local_addr();
329        let client_bind = unused_local_addr();
330        let server_alpn = vec![b"sova_quic-udp".to_vec()];
331
332        let tmp = TempDir::new().unwrap();
333        let (cert_path, key_path, _cert_der) = write_self_signed_pem(&tmp, &["localhost"]);
334
335        let handler: QuicDatagramHandler = Arc::new(|_peer, data, conn| {
336            Box::pin(async move {
337                let _ = conn.send_datagram(Bytes::from(data));
338            })
339        });
340
341        let svc = QuicDatagramService::from_pem(
342            bind_addr,
343            &cert_path,
344            &key_path,
345            server_alpn.clone(),
346            true,
347            handler,
348        )
349        .unwrap();
350
351        let (tx, shutdown) = sova_core::shutdown_channel();
352        let handle = tokio::spawn(Box::new(svc).run(Arc::new(StateMap::new()), shutdown));
353
354        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
355
356        let client = QuicDatagramClient::connect(client_bind, bind_addr, "localhost", server_alpn)
357            .await
358            .unwrap();
359
360        let resp = client
361            .send_and_recv(b"ping", std::time::Duration::from_secs(1))
362            .await
363            .unwrap();
364        assert_eq!(resp, b"ping");
365
366        let _ = tx.send(true);
367        let _ = handle.await;
368    }
369
370    #[tokio::test]
371    async fn handshake_requires_alpn() {
372        let bind_addr = unused_local_addr();
373        let client_bind = unused_local_addr();
374        let server_alpn = vec![b"A".to_vec()];
375        let client_alpn = vec![b"B".to_vec()];
376
377        let tmp = TempDir::new().unwrap();
378        let (cert_path, key_path, _cert_der) = write_self_signed_pem(&tmp, &["localhost"]);
379
380        let handler: QuicDatagramHandler = Arc::new(|_peer, _data, _conn| Box::pin(async {}));
381
382        let svc = QuicDatagramService::from_pem(
383            bind_addr,
384            &cert_path,
385            &key_path,
386            server_alpn.clone(),
387            true,
388            handler,
389        )
390        .unwrap();
391
392        let (tx, shutdown) = sova_core::shutdown_channel();
393        let handle = tokio::spawn(Box::new(svc).run(Arc::new(StateMap::new()), shutdown));
394
395        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
396
397        let res = QuicDatagramClient::connect(client_bind, bind_addr, "localhost", client_alpn).await;
398        assert!(res.is_err());
399
400        let _ = tx.send(true);
401        let _ = handle.await;
402    }
403
404    #[tokio::test]
405    async fn max_datagram_size_accepted() {
406        let bind_addr = unused_local_addr();
407        let client_bind = unused_local_addr();
408        let server_alpn = vec![b"sova_quic-udp".to_vec()];
409
410        let tmp = TempDir::new().unwrap();
411        let (cert_path, key_path, _cert_der) = write_self_signed_pem(&tmp, &["localhost"]);
412
413        let handler: QuicDatagramHandler = Arc::new(|_peer, data, conn| {
414            Box::pin(async move {
415                let _ = conn.send_datagram(Bytes::from(data));
416            })
417        });
418
419        let svc = QuicDatagramService::from_pem(
420            bind_addr,
421            &cert_path,
422            &key_path,
423            server_alpn.clone(),
424            true,
425            handler,
426        )
427        .unwrap();
428
429        let (tx, shutdown) = sova_core::shutdown_channel();
430        let handle = tokio::spawn(Box::new(svc).run(Arc::new(StateMap::new()), shutdown));
431        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
432
433        let client = QuicDatagramClient::connect(client_bind, bind_addr, "localhost", server_alpn)
434            .await
435            .unwrap();
436
437        let max = client.max_datagram_size().expect("max_datagram_size");
438        let len = max.saturating_sub(64).max(1);
439        let payload = vec![0xABu8; len];
440
441        let resp = client
442            .send_and_recv(&payload, std::time::Duration::from_secs(1))
443            .await
444            .unwrap();
445        assert_eq!(resp, payload);
446
447        let _ = tx.send(true);
448        let _ = handle.await;
449    }
450
451    #[tokio::test]
452    async fn shutdown_stops_recv_loop() {
453        let bind_addr = unused_local_addr();
454        let client_bind = unused_local_addr();
455        let server_alpn = vec![b"sova_quic-udp".to_vec()];
456
457        let tmp = TempDir::new().unwrap();
458        let (cert_path, key_path, _cert_der) = write_self_signed_pem(&tmp, &["localhost"]);
459
460        let handler: QuicDatagramHandler = Arc::new(|_peer, _data, _conn| Box::pin(async {}));
461
462        let svc = QuicDatagramService::from_pem(
463            bind_addr,
464            &cert_path,
465            &key_path,
466            server_alpn.clone(),
467            true,
468            handler,
469        )
470        .unwrap();
471
472        let (tx, shutdown) = sova_core::shutdown_channel();
473        let handle = tokio::spawn(Box::new(svc).run(Arc::new(StateMap::new()), shutdown));
474
475        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
476        let _client =
477            QuicDatagramClient::connect(client_bind, bind_addr, "localhost", server_alpn)
478            .await
479            .unwrap();
480
481        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
482        let _ = tx.send(true);
483
484        tokio::time::timeout(std::time::Duration::from_secs(1), handle)
485            .await
486            .expect("service did not stop after shutdown")
487            .expect("join error");
488    }
489
490    #[tokio::test]
491    async fn tls_reload_affects_new_connections() {
492        let bind_addr = unused_local_addr();
493        let client_bind1 = unused_local_addr();
494        let client_bind2 = unused_local_addr();
495        let server_alpn = vec![b"sova_quic-udp".to_vec()];
496
497        let tmp = TempDir::new().unwrap();
498        let (cert_path, key_path, cert1_der) = write_self_signed_pem(&tmp, &["localhost"]);
499        let svc = QuicDatagramService::from_pem(
500            bind_addr,
501            &cert_path,
502            &key_path,
503            server_alpn.clone(),
504            true,
505            Arc::new(|_peer, _data, _conn| Box::pin(async {})),
506        )
507        .unwrap();
508
509        let tls = svc.tls.clone();
510        let (tx, shutdown) = sova_core::shutdown_channel();
511        let handle = tokio::spawn(Box::new(svc).run(Arc::new(StateMap::new()), shutdown));
512
513        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
514        let client1 = QuicDatagramClient::connect(client_bind1, bind_addr, "localhost", server_alpn.clone())
515            .await
516            .unwrap();
517        let peer1 = client1
518            .peer_certificates()
519            .and_then(|mut v| v.pop())
520            .expect("peer cert");
521
522        let (_cert_path2, _key_path2, cert2_der) = write_self_signed_pem(&tmp, &["localhost", "127.0.0.1"]);
523        fs::write(&cert_path, fs::read(_cert_path2).unwrap()).unwrap();
524        fs::write(&key_path, fs::read(_key_path2).unwrap()).unwrap();
525        tls.reload().unwrap();
526
527        let client2 = QuicDatagramClient::connect(client_bind2, bind_addr, "localhost", server_alpn)
528            .await
529            .unwrap();
530        let peer2 = client2
531            .peer_certificates()
532            .and_then(|mut v| v.pop())
533            .expect("peer cert");
534
535        assert_eq!(peer1, cert1_der);
536        assert_eq!(peer2, cert2_der);
537
538        let _ = tx.send(true);
539        let _ = handle.await;
540    }
541}
sova_quic/datagram.rs

sova_quic/
datagram.rs
