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
25const DEFAULT_READ_CHUNK_SIZE: usize = 1024 * 1024; pub type QuicSendStream = quinn::SendStream;
30
31pub type QuicRecvStream = quinn::RecvStream;
33
34#[derive(Clone)]
36pub struct QuicConfig {
37 pub max_bi_streams: u64,
39 pub max_uni_streams: u64,
41 pub keep_alive_interval: Option<Duration>,
43 pub idle_timeout: Option<Duration>,
45 pub enable_datagrams: bool,
47 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#[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 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 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 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#[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 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 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 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 if !config.enable_datagrams {
218 transport.datagram_receive_buffer_size(None);
219 }
220 transport
221}
222
223pub struct QuicEndpoint {
225 inner: quinn::Endpoint,
226 local_endpoint: Endpoint,
227}
228
229impl QuicEndpoint {
230 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 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 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 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 pub async fn accept(&self) -> Result<QuicConn> {
283 self.accept_incoming().await?.handshake().await
284 }
285
286 pub fn local_endpoint(&self) -> Result<Endpoint> {
288 Ok(self.local_endpoint.clone())
289 }
290
291 pub fn close(&self, code: u32, reason: &[u8]) {
293 self.inner.close(quinn::VarInt::from_u32(code), reason);
294 }
295}
296
297pub struct QuicIncoming {
299 inner: Incoming,
300 local_endpoint: Endpoint,
301}
302
303impl QuicIncoming {
304 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 pub fn peer_endpoint(&self) -> Endpoint {
317 Endpoint::new_quic_addr(self.inner.remote_address())
318 }
319}
320
321pub struct QuicConn {
324 inner: quinn::Connection,
325 peer_endpoint: Endpoint,
326 local_endpoint: Endpoint,
327}
328
329impl QuicConn {
330 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 pub async fn open_uni(&self) -> Result<QuicSendStream> {
338 let send = self.inner.open_uni().await?;
339 Ok(send)
340 }
341
342 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 pub async fn accept_uni(&self) -> Result<QuicRecvStream> {
350 let recv = self.inner.accept_uni().await?;
351 Ok(recv)
352 }
353
354 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 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 pub fn max_datagram_size(&self) -> Option<usize> {
372 self.inner.max_datagram_size()
373 }
374
375 pub fn peer_endpoint(&self) -> Result<Endpoint> {
377 Ok(self.peer_endpoint.clone())
378 }
379
380 pub fn local_endpoint(&self) -> Result<Endpoint> {
382 Ok(self.local_endpoint.clone())
383 }
384
385 pub fn rtt(&self) -> Duration {
387 self.inner.rtt()
388 }
389
390 pub fn close(&self, code: u32, reason: &[u8]) {
392 self.inner.close(quinn::VarInt::from_u32(code), reason);
393 }
394
395 pub async fn closed(&self) -> quinn::ConnectionError {
397 self.inner.closed().await
398 }
399
400 pub fn inner(&self) -> &quinn::Connection {
402 &self.inner
403 }
404
405 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
416impl 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
448pub 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#[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}