Skip to main content

karyon_net/transports/
quic.rs

1use std::{
2    io,
3    net::SocketAddr,
4    pin::Pin,
5    sync::Arc,
6    task::{Context, Poll},
7    time::Duration,
8};
9
10use quinn::Incoming;
11use rustls_pki_types::{CertificateDer, PrivateKeyDer};
12
13use crate::async_rustls::rustls;
14
15use karyon_core::async_runtime::io::{AsyncRead, AsyncWrite};
16
17#[cfg(feature = "tokio")]
18use tokio::io::{AsyncWrite as TokioAsyncWrite, ReadBuf};
19
20#[cfg(feature = "smol")]
21use futures_util::io::{AsyncRead as FutAsyncRead, AsyncWrite as FutAsyncWrite};
22
23use crate::{stream::ByteStream, stream_mux::StreamMux, Bytes, Endpoint, Error, Result};
24
25/// Default read chunk size for QUIC streams.
26const DEFAULT_READ_CHUNK_SIZE: usize = 1024 * 1024; // 1MB
27
28/// A QUIC send stream. Re-exported from quinn for use by higher layers.
29pub type QuicSendStream = quinn::SendStream;
30
31/// A QUIC receive stream. Re-exported from quinn for use by higher layers.
32pub type QuicRecvStream = quinn::RecvStream;
33
34/// QUIC configuration.
35#[derive(Clone)]
36pub struct QuicConfig {
37    /// Maximum concurrent bidirectional streams.
38    pub max_bi_streams: u64,
39    /// Maximum concurrent unidirectional streams.
40    pub max_uni_streams: u64,
41    /// Keep-alive interval. None to disable.
42    pub keep_alive_interval: Option<Duration>,
43    /// Idle timeout.
44    pub idle_timeout: Option<Duration>,
45    /// Enable datagrams.
46    pub enable_datagrams: bool,
47    /// Read chunk size for stream reads (bytes).
48    pub read_chunk_size: usize,
49}
50
51impl Default for QuicConfig {
52    fn default() -> Self {
53        Self {
54            max_bi_streams: 100,
55            max_uni_streams: 100,
56            keep_alive_interval: Some(Duration::from_secs(5)),
57            idle_timeout: Some(Duration::from_secs(30)),
58            enable_datagrams: false,
59            read_chunk_size: DEFAULT_READ_CHUNK_SIZE,
60        }
61    }
62}
63
64/// Server-side QUIC configuration. Build from either a cert chain +
65/// private key (`new`) or a pre-built rustls config (`from_rustls`).
66#[derive(Clone)]
67pub struct ServerQuicConfig {
68    source: ServerSource,
69    config: QuicConfig,
70}
71
72#[derive(Clone)]
73enum ServerSource {
74    Certs {
75        cert_chain: Vec<CertificateDer<'static>>,
76        private_key: Arc<PrivateKeyDer<'static>>,
77    },
78    Rustls(Box<rustls::ServerConfig>),
79}
80
81impl ServerQuicConfig {
82    /// Create a config from a cert chain + private key.
83    pub fn new(
84        cert_chain: Vec<CertificateDer<'static>>,
85        private_key: Arc<PrivateKeyDer<'static>>,
86    ) -> Self {
87        Self {
88            source: ServerSource::Certs {
89                cert_chain,
90                private_key,
91            },
92            config: QuicConfig::default(),
93        }
94    }
95
96    /// Create a config from a pre-built rustls `ServerConfig` (for
97    /// custom verifiers, client-auth, etc).
98    pub fn from_rustls(rustls_config: rustls::ServerConfig) -> Self {
99        Self {
100            source: ServerSource::Rustls(Box::new(rustls_config)),
101            config: QuicConfig::default(),
102        }
103    }
104
105    /// Override the QUIC transport parameters.
106    pub fn with_config(mut self, config: QuicConfig) -> Self {
107        self.config = config;
108        self
109    }
110
111    pub(crate) fn build(self) -> Result<quinn::ServerConfig> {
112        let mut server_config = match self.source {
113            ServerSource::Certs {
114                cert_chain,
115                private_key,
116            } => quinn::ServerConfig::with_single_cert(cert_chain, private_key.clone_key())
117                .map_err(|e| Error::TlsConfig(e.to_string()))?,
118            ServerSource::Rustls(rustls_config) => {
119                let quic_config = quinn::crypto::rustls::QuicServerConfig::try_from(*rustls_config)
120                    .map_err(|e| Error::QuicConfigError(e.to_string()))?;
121                quinn::ServerConfig::with_crypto(Arc::new(quic_config))
122            }
123        };
124        server_config.transport_config(Arc::new(build_transport_config(&self.config)));
125        Ok(server_config)
126    }
127}
128
129/// Client-side QUIC configuration. Build from either a root cert list
130/// (`new`) or a pre-built rustls config (`from_rustls`).
131#[derive(Clone)]
132pub struct ClientQuicConfig {
133    source: ClientSource,
134    server_name: String,
135    config: QuicConfig,
136}
137
138#[derive(Clone)]
139enum ClientSource {
140    Roots(Vec<CertificateDer<'static>>),
141    Rustls(Box<rustls::ClientConfig>),
142}
143
144impl ClientQuicConfig {
145    /// Create a config from a list of trusted root certs + server name.
146    pub fn new(root_certs: Vec<CertificateDer<'static>>, server_name: impl Into<String>) -> Self {
147        Self {
148            source: ClientSource::Roots(root_certs),
149            server_name: server_name.into(),
150            config: QuicConfig::default(),
151        }
152    }
153
154    /// Create a config from a pre-built rustls `ClientConfig` + server name
155    /// (for custom verifiers, client-auth certs, etc).
156    pub fn from_rustls(
157        rustls_config: rustls::ClientConfig,
158        server_name: impl Into<String>,
159    ) -> Self {
160        Self {
161            source: ClientSource::Rustls(Box::new(rustls_config)),
162            server_name: server_name.into(),
163            config: QuicConfig::default(),
164        }
165    }
166
167    /// Override the QUIC transport parameters.
168    pub fn with_config(mut self, config: QuicConfig) -> Self {
169        self.config = config;
170        self
171    }
172
173    pub fn server_name(&self) -> &str {
174        &self.server_name
175    }
176
177    pub(crate) fn build(self) -> Result<quinn::ClientConfig> {
178        let mut client_config = match self.source {
179            ClientSource::Roots(root_certs) => {
180                let mut root_store = rustls::RootCertStore::empty();
181                for cert in root_certs {
182                    root_store
183                        .add(cert)
184                        .map_err(|e| Error::TlsConfig(e.to_string()))?;
185                }
186                quinn::ClientConfig::with_root_certificates(Arc::new(root_store))
187                    .map_err(|e| Error::TlsConfig(e.to_string()))?
188            }
189            ClientSource::Rustls(rustls_config) => {
190                let quic_config = quinn::crypto::rustls::QuicClientConfig::try_from(*rustls_config)
191                    .map_err(|e| Error::QuicConfigError(e.to_string()))?;
192                quinn::ClientConfig::new(Arc::new(quic_config))
193            }
194        };
195        client_config.transport_config(Arc::new(build_transport_config(&self.config)));
196        Ok(client_config)
197    }
198}
199
200fn build_transport_config(config: &QuicConfig) -> quinn::TransportConfig {
201    let mut transport = quinn::TransportConfig::default();
202    transport.max_concurrent_bidi_streams(
203        quinn::VarInt::from_u64(config.max_bi_streams).unwrap_or(quinn::VarInt::from_u32(100u32)),
204    );
205    transport.max_concurrent_uni_streams(
206        quinn::VarInt::from_u64(config.max_uni_streams).unwrap_or(quinn::VarInt::from_u32(100u32)),
207    );
208    if let Some(interval) = config.keep_alive_interval {
209        transport.keep_alive_interval(Some(interval));
210    }
211    if let Some(timeout) = config.idle_timeout {
212        if let Ok(idle) = quinn::IdleTimeout::try_from(timeout) {
213            transport.max_idle_timeout(Some(idle));
214        }
215    }
216    // Disable datagrams explicitly when the caller opts out.
217    if !config.enable_datagrams {
218        transport.datagram_receive_buffer_size(None);
219    }
220    transport
221}
222
223/// A QUIC endpoint that can listen for and initiate connections.
224pub struct QuicEndpoint {
225    inner: quinn::Endpoint,
226    local_endpoint: Endpoint,
227}
228
229impl QuicEndpoint {
230    /// Bind to a local address and start listening with the given server config.
231    pub async fn listen(endpoint: &Endpoint, config: ServerQuicConfig) -> Result<Self> {
232        let addr = SocketAddr::try_from(endpoint.clone())?;
233        let server_config = config.build()?;
234        let inner = quinn::Endpoint::server(server_config, addr)?;
235        let local_addr = inner.local_addr()?;
236        let local_endpoint = Endpoint::new_quic_addr(local_addr);
237        Ok(Self {
238            inner,
239            local_endpoint,
240        })
241    }
242
243    /// Connect to a remote QUIC endpoint.
244    pub async fn dial(endpoint: &Endpoint, config: ClientQuicConfig) -> Result<QuicConn> {
245        let addr = SocketAddr::try_from(endpoint.clone())?;
246        let server_name = config.server_name.clone();
247        let client_config = config.build()?;
248        // Bind the client socket in the same address family as the target.
249        let bind_addr: SocketAddr = if addr.is_ipv6() {
250            "[::]:0".parse().unwrap()
251        } else {
252            "0.0.0.0:0".parse().unwrap()
253        };
254        let mut quinn_endpoint = quinn::Endpoint::client(bind_addr)?;
255        quinn_endpoint.set_default_client_config(client_config);
256
257        let connection = quinn_endpoint.connect(addr, &server_name)?.await?;
258        let peer_endpoint = Endpoint::new_quic_addr(connection.remote_address());
259        let local_addr = quinn_endpoint.local_addr()?;
260        let local_endpoint = Endpoint::new_quic_addr(local_addr);
261
262        Ok(QuicConn {
263            inner: connection,
264            peer_endpoint,
265            local_endpoint,
266        })
267    }
268
269    /// Accept an incoming QUIC connection without running its handshake.
270    ///
271    /// Servers should use this and drive [`QuicIncoming::handshake`] in a
272    /// separate task, so a slow peer cannot stall the accept loop.
273    pub async fn accept_incoming(&self) -> Result<QuicIncoming> {
274        let inner = self.inner.accept().await.ok_or(Error::ConnectionClosed)?;
275        Ok(QuicIncoming {
276            inner,
277            local_endpoint: self.local_endpoint.clone(),
278        })
279    }
280
281    /// Accept an incoming QUIC connection and run its handshake.
282    pub async fn accept(&self) -> Result<QuicConn> {
283        self.accept_incoming().await?.handshake().await
284    }
285
286    /// Returns the local endpoint.
287    pub fn local_endpoint(&self) -> Result<Endpoint> {
288        Ok(self.local_endpoint.clone())
289    }
290
291    /// Close the endpoint.
292    pub fn close(&self, code: u32, reason: &[u8]) {
293        self.inner.close(quinn::VarInt::from_u32(code), reason);
294    }
295}
296
297/// An accepted QUIC connection whose handshake has not run yet.
298pub struct QuicIncoming {
299    inner: Incoming,
300    local_endpoint: Endpoint,
301}
302
303impl QuicIncoming {
304    /// Runs the QUIC/TLS handshake.
305    pub async fn handshake(self) -> Result<QuicConn> {
306        let connection = self.inner.await?;
307        let peer_endpoint = Endpoint::new_quic_addr(connection.remote_address());
308        Ok(QuicConn {
309            inner: connection,
310            peer_endpoint,
311            local_endpoint: self.local_endpoint,
312        })
313    }
314
315    /// Remote address of the peer. Known before the handshake runs.
316    pub fn peer_endpoint(&self) -> Endpoint {
317        Endpoint::new_quic_addr(self.inner.remote_address())
318    }
319}
320
321/// A QUIC connection. Manages streams and datagrams.
322/// This is NOT a single read/write channel — it is a stream factory.
323pub struct QuicConn {
324    inner: quinn::Connection,
325    peer_endpoint: Endpoint,
326    local_endpoint: Endpoint,
327}
328
329impl QuicConn {
330    /// Open a new bidirectional stream.
331    pub async fn open_bi(&self) -> Result<(QuicSendStream, QuicRecvStream)> {
332        let (send, recv) = self.inner.open_bi().await?;
333        Ok((send, recv))
334    }
335
336    /// Open a new unidirectional (send-only) stream.
337    pub async fn open_uni(&self) -> Result<QuicSendStream> {
338        let send = self.inner.open_uni().await?;
339        Ok(send)
340    }
341
342    /// Accept a bidirectional stream opened by the peer.
343    pub async fn accept_bi(&self) -> Result<(QuicSendStream, QuicRecvStream)> {
344        let (send, recv) = self.inner.accept_bi().await?;
345        Ok((send, recv))
346    }
347
348    /// Accept a unidirectional (receive-only) stream from the peer.
349    pub async fn accept_uni(&self) -> Result<QuicRecvStream> {
350        let recv = self.inner.accept_uni().await?;
351        Ok(recv)
352    }
353
354    /// Send an unreliable datagram over the connection. Zero-copy —
355    /// ownership of the `Bytes` allocation is passed to quinn.
356    pub fn send_datagram(&self, data: Bytes) -> Result<()> {
357        self.inner
358            .send_datagram(data.into_inner())
359            .map_err(|e| Error::QuicConfigError(e.to_string()))?;
360        Ok(())
361    }
362
363    /// Receive an unreliable datagram. Zero-copy — wraps the allocation
364    /// returned by quinn.
365    pub async fn recv_datagram(&self) -> Result<Bytes> {
366        let data = self.inner.read_datagram().await?;
367        Ok(Bytes::from_inner(data))
368    }
369
370    /// Maximum datagram size the peer supports, or None if unsupported.
371    pub fn max_datagram_size(&self) -> Option<usize> {
372        self.inner.max_datagram_size()
373    }
374
375    /// Remote peer's address.
376    pub fn peer_endpoint(&self) -> Result<Endpoint> {
377        Ok(self.peer_endpoint.clone())
378    }
379
380    /// Local address.
381    pub fn local_endpoint(&self) -> Result<Endpoint> {
382        Ok(self.local_endpoint.clone())
383    }
384
385    /// Current round-trip time estimate.
386    pub fn rtt(&self) -> Duration {
387        self.inner.rtt()
388    }
389
390    /// Close the connection gracefully.
391    pub fn close(&self, code: u32, reason: &[u8]) {
392        self.inner.close(quinn::VarInt::from_u32(code), reason);
393    }
394
395    /// Wait for the connection to be closed (by us or the peer).
396    pub async fn closed(&self) -> quinn::ConnectionError {
397        self.inner.closed().await
398    }
399
400    /// Returns a reference to the inner quinn connection.
401    pub fn inner(&self) -> &quinn::Connection {
402        &self.inner
403    }
404
405    /// Peer certificate chain from the QUIC TLS handshake. Quinn returns
406    /// a type-erased `Box<dyn Any>`; for rustls-based QUIC (which is what
407    /// karyon uses) it downcasts to `Vec<CertificateDer<'static>>`.
408    pub fn peer_certificates(&self) -> Option<Vec<CertificateDer<'static>>> {
409        let any = self.inner.peer_identity()?;
410        any.downcast::<Vec<CertificateDer<'static>>>()
411            .ok()
412            .map(|b| *b)
413    }
414}
415
416// -- StreamMux impl --
417
418impl StreamMux for QuicConn {
419    async fn open_stream(&self) -> Result<Box<dyn ByteStream>> {
420        let (send, recv) = self.open_bi().await?;
421        Ok(Box::new(QuicBiStream {
422            send,
423            recv,
424            peer_endpoint: self.peer_endpoint.clone(),
425            local_endpoint: self.local_endpoint.clone(),
426        }))
427    }
428
429    async fn accept_stream(&self) -> Result<Box<dyn ByteStream>> {
430        let (send, recv) = self.accept_bi().await?;
431        Ok(Box::new(QuicBiStream {
432            send,
433            recv,
434            peer_endpoint: self.peer_endpoint.clone(),
435            local_endpoint: self.local_endpoint.clone(),
436        }))
437    }
438
439    fn peer_endpoint(&self) -> Option<Endpoint> {
440        Some(self.peer_endpoint.clone())
441    }
442
443    fn local_endpoint(&self) -> Option<Endpoint> {
444        Some(self.local_endpoint.clone())
445    }
446}
447
448/// Single QUIC bidirectional stream as a ByteStream.
449pub struct QuicBiStream {
450    send: QuicSendStream,
451    recv: QuicRecvStream,
452    peer_endpoint: Endpoint,
453    local_endpoint: Endpoint,
454}
455
456impl ByteStream for QuicBiStream {
457    fn peer_endpoint(&self) -> Option<Endpoint> {
458        Some(self.peer_endpoint.clone())
459    }
460    fn local_endpoint(&self) -> Option<Endpoint> {
461        Some(self.local_endpoint.clone())
462    }
463}
464
465// Quinn streams use tokio IO traits natively.
466// For smol builds, delegate manually via poll methods.
467
468#[cfg(feature = "tokio")]
469impl AsyncRead for QuicBiStream {
470    fn poll_read(
471        mut self: Pin<&mut Self>,
472        cx: &mut Context<'_>,
473        buf: &mut ReadBuf<'_>,
474    ) -> Poll<io::Result<()>> {
475        Pin::new(&mut self.recv).poll_read(cx, buf)
476    }
477}
478
479#[cfg(feature = "smol")]
480impl AsyncRead for QuicBiStream {
481    fn poll_read(
482        mut self: Pin<&mut Self>,
483        cx: &mut Context<'_>,
484        buf: &mut [u8],
485    ) -> Poll<io::Result<usize>> {
486        FutAsyncRead::poll_read(Pin::new(&mut self.recv), cx, buf)
487    }
488}
489
490#[cfg(feature = "tokio")]
491impl AsyncWrite for QuicBiStream {
492    fn poll_write(
493        mut self: Pin<&mut Self>,
494        cx: &mut Context<'_>,
495        buf: &[u8],
496    ) -> Poll<io::Result<usize>> {
497        TokioAsyncWrite::poll_write(Pin::new(&mut self.send), cx, buf)
498    }
499
500    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
501        TokioAsyncWrite::poll_flush(Pin::new(&mut self.send), cx)
502    }
503
504    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
505        TokioAsyncWrite::poll_shutdown(Pin::new(&mut self.send), cx)
506    }
507}
508
509#[cfg(feature = "smol")]
510impl AsyncWrite for QuicBiStream {
511    fn poll_write(
512        mut self: Pin<&mut Self>,
513        cx: &mut Context<'_>,
514        buf: &[u8],
515    ) -> Poll<io::Result<usize>> {
516        FutAsyncWrite::poll_write(Pin::new(&mut self.send), cx, buf)
517    }
518
519    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
520        FutAsyncWrite::poll_flush(Pin::new(&mut self.send), cx)
521    }
522
523    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
524        FutAsyncWrite::poll_close(Pin::new(&mut self.send), cx)
525    }
526}