1use 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";
12pub const DATAGRAM_BUFFER_BYTES: usize = 64 * 1024;
16pub 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 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}