Skip to main content

netget/utils/
quic.rs

1//! Shared owned QUIC transport and authenticated client configuration.
2use crate::protocol::StartupParams;
3use anyhow::{bail, ensure, Context, Result};
4use std::{sync::Arc, time::Duration};
5pub const MAX_STREAMS: usize = 32;
6pub const MAX_CONNECTIONS: usize = 64;
7pub const MAX_BYTES: usize = 1024 * 1024;
8pub const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
9pub const EXCHANGE_TIMEOUT: Duration = Duration::from_secs(30);
10pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
11pub const RAW_ALPN: &[u8] = b"netget-quic";
12/// RFC 9221 DATAGRAM frames buffered per connection before Quinn drops the oldest. A
13/// datagram is unreliable by definition, so a peer that outruns the handler loses
14/// datagrams rather than growing this; 64 KiB holds about fifty full-size ones.
15pub const DATAGRAM_BUFFER_BYTES: usize = 64 * 1024;
16/// Closing on Drop matters when the owning registered task is aborted during removal.
17/// Quinn's driver otherwise retains the endpoint until all connection handles disappear.
18pub struct EndpointGuard(pub quinn::Endpoint);
19impl Drop for EndpointGuard {
20    fn drop(&mut self) {
21        self.0.close(0u32.into(), b"endpoint stopped");
22    }
23}
24pub struct ConnectionGuard(pub quinn::Connection);
25impl Drop for ConnectionGuard {
26    fn drop(&mut self) {
27        self.0.close(0u32.into(), b"connection stopped");
28    }
29}
30
31pub fn parameter(
32    name: &str,
33    description: &str,
34    example: serde_json::Value,
35    default: Option<serde_json::Value>,
36) -> crate::llm::actions::ParameterDefinition {
37    crate::llm::actions::ParameterDefinition {
38        name: name.into(),
39        description: description.into(),
40        type_hint: if example.is_number() {
41            "number"
42        } else {
43            "string"
44        }
45        .into(),
46        required: false,
47        example,
48        default,
49    }
50}
51
52pub fn bounded_parameter(
53    params: Option<&crate::protocol::StartupParams>,
54    name: &str,
55    default: u64,
56    max: u64,
57) -> Result<u64> {
58    let value = params
59        .map(|p| p.get_optional_u64(name))
60        .transpose()?
61        .flatten()
62        .unwrap_or(default);
63    if value == 0 || value > max {
64        bail!("{name} must be between 1 and {max}");
65    }
66    Ok(value)
67}
68
69pub fn client_parameters() -> Vec<crate::llm::actions::ParameterDefinition> {
70    vec![
71        parameter(
72            "server_name",
73            "TLS authentication name; defaults to remote hostname",
74            serde_json::json!("localhost"),
75            None,
76        ),
77        parameter(
78            "ca_cert_path",
79            "Additional PEM CA or self-signed certificate to trust",
80            serde_json::json!("ca.pem"),
81            None,
82        ),
83        parameter(
84            "handshake_timeout_secs",
85            "Resolution and TLS handshake deadline, 1..60 seconds",
86            serde_json::json!(10),
87            Some(serde_json::json!(10)),
88        ),
89        parameter(
90            "exchange_timeout_secs",
91            "Whole stream exchange deadline, 1..300 seconds",
92            serde_json::json!(30),
93            Some(serde_json::json!(30)),
94        ),
95        parameter(
96            "idle_timeout_secs",
97            "Connection idle deadline, 1..3600 seconds",
98            serde_json::json!(300),
99            Some(serde_json::json!(300)),
100        ),
101    ]
102}
103pub fn transport(idle: Duration, bidi: u32, uni: u32) -> Arc<quinn::TransportConfig> {
104    let mut config = quinn::TransportConfig::default();
105    config
106        .max_concurrent_bidi_streams(bidi.into())
107        .max_concurrent_uni_streams(uni.into());
108    config.max_idle_timeout(Some(idle.try_into().expect("bounded idle timeout")));
109    config.stream_receive_window((MAX_BYTES as u32).into());
110    config.receive_window((MAX_BYTES as u32 * MAX_STREAMS as u32).into());
111    // Advertises max_datagram_frame_size (RFC 9221 ยง3), so a peer may send DATAGRAM frames;
112    // both buffers are bounded.
113    config.datagram_receive_buffer_size(Some(DATAGRAM_BUFFER_BYTES));
114    config.datagram_send_buffer_size(DATAGRAM_BUFFER_BYTES);
115    Arc::new(config)
116}
117pub async fn connect(
118    remote: &str,
119    params: Option<&StartupParams>,
120    alpn: &[u8],
121    uni: u32,
122) -> Result<(EndpointGuard, ConnectionGuard)> {
123    let handshake =
124        Duration::from_secs(bounded_parameter(params, "handshake_timeout_secs", 10, 60)?);
125    let idle = Duration::from_secs(bounded_parameter(params, "idle_timeout_secs", 300, 3600)?);
126    let auth = params
127        .map(|p| p.get_optional_string("server_name"))
128        .transpose()?
129        .flatten();
130    let ca = params
131        .map(|p| p.get_optional_string("ca_cert_path"))
132        .transpose()?
133        .flatten();
134    let name = remote
135        .rsplit_once(':')
136        .context("Remote address must be host:port or [IPv6]:port")?
137        .0
138        .trim_matches(['[', ']']);
139    let name = auth.unwrap_or_else(|| name.to_owned());
140    ensure!(
141        !alpn.is_empty() && alpn.len() <= 255,
142        "ALPN must contain 1..255 bytes"
143    );
144    let mut roots = rustls::RootCertStore::empty();
145    roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
146    if let Some(path) = ca {
147        let pem = std::fs::read(path).context("Read QUIC trust certificate")?;
148        let certs = rustls_pemfile::certs(&mut pem.as_slice())
149            .collect::<std::result::Result<Vec<_>, _>>()?;
150        ensure!(!certs.is_empty(), "Trust file contains no certificates");
151        for cert in certs {
152            roots.add(cert)?;
153        }
154    }
155    let _ = rustls::crypto::ring::default_provider().install_default();
156    let mut tls = rustls::ClientConfig::builder()
157        .with_root_certificates(roots)
158        .with_no_client_auth();
159    tls.alpn_protocols = vec![alpn.to_vec()];
160    tls.enable_early_data = false;
161    let mut config = quinn::ClientConfig::new(Arc::new(
162        quinn::crypto::rustls::QuicClientConfig::try_from(tls)?,
163    ));
164    config.transport_config(transport(idle, 0, uni));
165    tokio::time::timeout(handshake, async {
166        let remote = tokio::net::lookup_host(remote)
167            .await?
168            .next()
169            .context("No address resolved")?;
170        let mut endpoint = quinn::Endpoint::client(
171            if remote.is_ipv6() {
172                "[::]:0"
173            } else {
174                "0.0.0.0:0"
175            }
176            .parse()?,
177        )?;
178        endpoint.set_default_client_config(config);
179        let endpoint = EndpointGuard(endpoint);
180        let connection = endpoint
181            .0
182            .connect(remote, &name)?
183            .await
184            .context("QUIC TLS handshake (check ca_cert_path and server_name)")?;
185        Ok((endpoint, ConnectionGuard(connection)))
186    })
187    .await
188    .context("QUIC handshake deadline exceeded")?
189}