compio_quic/
connection.rs

1use std::{
2    collections::VecDeque,
3    fmt::Debug,
4    net::{IpAddr, SocketAddr},
5    pin::{Pin, pin},
6    task::{Context, Poll, Waker},
7    time::{Duration, Instant},
8};
9
10use compio_buf::bytes::Bytes;
11use compio_log::Instrument;
12use compio_runtime::{JoinHandle, SpawnMeta};
13use flume::{Receiver, Sender};
14use futures_util::{
15    FutureExt, StreamExt,
16    future::{self, Fuse, FusedFuture, LocalBoxFuture},
17    select, stream,
18};
19#[cfg(rustls)]
20use quinn_proto::crypto::rustls::HandshakeData;
21use quinn_proto::{
22    ConnectionHandle, ConnectionStats, Dir, EndpointEvent, Side, StreamEvent, StreamId, VarInt,
23    congestion::Controller,
24};
25use rustc_hash::FxHashMap as HashMap;
26use thiserror::Error;
27
28use crate::{
29    RecvStream, SendStream, Socket,
30    sync::{
31        mutex_blocking::{Mutex, MutexGuard},
32        shared::Shared,
33    },
34};
35
36#[derive(Debug)]
37pub(crate) enum ConnectionEvent {
38    Close(VarInt, Bytes),
39    Proto(quinn_proto::ConnectionEvent),
40}
41
42#[derive(Debug)]
43pub(crate) struct ConnectionState {
44    pub(crate) conn: quinn_proto::Connection,
45    pub(crate) error: Option<ConnectionError>,
46    connected: bool,
47    worker: Option<JoinHandle<()>>,
48    poller: Option<Waker>,
49    on_connected: Option<Waker>,
50    on_handshake_data: Option<Waker>,
51    datagram_received: VecDeque<Waker>,
52    datagrams_unblocked: VecDeque<Waker>,
53    stream_opened: [VecDeque<Waker>; 2],
54    stream_available: [VecDeque<Waker>; 2],
55    pub(crate) writable: HashMap<StreamId, Waker>,
56    pub(crate) readable: HashMap<StreamId, Waker>,
57    pub(crate) stopped: HashMap<StreamId, Waker>,
58}
59
60impl ConnectionState {
61    fn terminate(&mut self, reason: ConnectionError) {
62        self.error = Some(reason);
63        self.connected = false;
64
65        if let Some(waker) = self.on_handshake_data.take() {
66            waker.wake()
67        }
68        if let Some(waker) = self.on_connected.take() {
69            waker.wake()
70        }
71        self.datagram_received.drain(..).for_each(Waker::wake);
72        self.datagrams_unblocked.drain(..).for_each(Waker::wake);
73        for e in &mut self.stream_opened {
74            e.drain(..).for_each(Waker::wake);
75        }
76        for e in &mut self.stream_available {
77            e.drain(..).for_each(Waker::wake);
78        }
79        wake_all_streams(&mut self.writable);
80        wake_all_streams(&mut self.readable);
81        wake_all_streams(&mut self.stopped);
82    }
83
84    fn close(&mut self, error_code: VarInt, reason: Bytes) {
85        self.conn.close(Instant::now(), error_code, reason);
86        self.terminate(ConnectionError::LocallyClosed);
87        self.wake();
88    }
89
90    pub(crate) fn wake(&mut self) {
91        if let Some(waker) = self.poller.take() {
92            waker.wake()
93        }
94    }
95
96    #[cfg(rustls)]
97    fn handshake_data(&self) -> Option<Box<HandshakeData>> {
98        self.conn
99            .crypto_session()
100            .handshake_data()
101            .map(|data| data.downcast::<HandshakeData>().unwrap())
102    }
103
104    pub(crate) fn check_0rtt(&self) -> bool {
105        self.conn.side().is_server() || self.conn.is_handshaking() || self.conn.accepted_0rtt()
106    }
107}
108
109fn wake_stream(stream: StreamId, wakers: &mut HashMap<StreamId, Waker>) {
110    if let Some(waker) = wakers.remove(&stream) {
111        waker.wake();
112    }
113}
114
115fn wake_all_streams(wakers: &mut HashMap<StreamId, Waker>) {
116    wakers.drain().for_each(|(_, waker)| waker.wake())
117}
118
119#[derive(Debug)]
120pub(crate) struct ConnectionInner {
121    state: Mutex<ConnectionState>,
122    handle: ConnectionHandle,
123    socket: Socket,
124    events_tx: Sender<(ConnectionHandle, EndpointEvent)>,
125    events_rx: Receiver<ConnectionEvent>,
126}
127
128fn implicit_close(this: &Shared<ConnectionInner>) {
129    if Shared::strong_count(this) == 2 {
130        this.state().close(0u32.into(), Bytes::new())
131    }
132}
133
134impl ConnectionInner {
135    fn new(
136        handle: ConnectionHandle,
137        conn: quinn_proto::Connection,
138        socket: Socket,
139        events_tx: Sender<(ConnectionHandle, EndpointEvent)>,
140        events_rx: Receiver<ConnectionEvent>,
141    ) -> Self {
142        Self {
143            state: Mutex::new(ConnectionState {
144                conn,
145                connected: false,
146                error: None,
147                worker: None,
148                poller: None,
149                on_connected: None,
150                on_handshake_data: None,
151                datagram_received: VecDeque::new(),
152                datagrams_unblocked: VecDeque::new(),
153                stream_opened: [VecDeque::new(), VecDeque::new()],
154                stream_available: [VecDeque::new(), VecDeque::new()],
155                writable: HashMap::default(),
156                readable: HashMap::default(),
157                stopped: HashMap::default(),
158            }),
159            handle,
160            socket,
161            events_tx,
162            events_rx,
163        }
164    }
165
166    #[inline]
167    pub(crate) fn state(&self) -> MutexGuard<'_, ConnectionState> {
168        self.state.lock()
169    }
170
171    #[inline]
172    pub(crate) fn try_state(&self) -> Result<MutexGuard<'_, ConnectionState>, ConnectionError> {
173        let state = self.state();
174        if let Some(error) = &state.error {
175            Err(error.clone())
176        } else {
177            Ok(state)
178        }
179    }
180
181    async fn run(&self) {
182        let mut poller = stream::poll_fn(|cx| {
183            let mut state = self.state();
184            let ready = state.poller.is_none();
185            match &state.poller {
186                Some(waker) if waker.will_wake(cx.waker()) => {}
187                _ => state.poller = Some(cx.waker().clone()),
188            };
189            if ready {
190                Poll::Ready(Some(()))
191            } else {
192                Poll::Pending
193            }
194        })
195        .fuse();
196
197        let mut timer = Timer::new();
198        let mut event_stream = self.events_rx.stream().ready_chunks(100);
199        let mut send_buf = Some(Vec::with_capacity(self.state().conn.current_mtu() as usize));
200        let mut transmit_fut = pin!(Fuse::terminated());
201
202        loop {
203            let mut state = select! {
204                _ = poller.select_next_some() => self.state(),
205                _ = timer => {
206                    timer.reset(None);
207                    let mut state = self.state();
208                    state.conn.handle_timeout(Instant::now());
209                    state
210                }
211                events = event_stream.select_next_some() => {
212                    let mut state = self.state();
213                    for event in events {
214                        match event {
215                            ConnectionEvent::Close(error_code, reason) => state.close(error_code, reason),
216                            ConnectionEvent::Proto(event) => state.conn.handle_event(event),
217                        }
218                    }
219                    state
220                },
221                buf = transmit_fut => {
222                    // The following line is required to avoid "type annotations needed" error
223                    let mut buf: Vec<_> = buf;
224                    buf.clear();
225                    send_buf = Some(buf);
226                    self.state()
227                },
228            };
229
230            if let Some(mut buf) = send_buf.take() {
231                if let Some(transmit) = state.conn.poll_transmit(
232                    Instant::now(),
233                    self.socket.max_gso_segments(),
234                    &mut buf,
235                ) {
236                    transmit_fut.set(async move { self.socket.send(buf, &transmit).await }.fuse())
237                } else {
238                    send_buf = Some(buf);
239                }
240            }
241
242            timer.reset(state.conn.poll_timeout());
243
244            while let Some(event) = state.conn.poll_endpoint_events() {
245                let _ = self.events_tx.send((self.handle, event));
246            }
247
248            while let Some(event) = state.conn.poll() {
249                use quinn_proto::Event::*;
250                match event {
251                    HandshakeDataReady => {
252                        if let Some(waker) = state.on_handshake_data.take() {
253                            waker.wake()
254                        }
255                    }
256                    Connected => {
257                        state.connected = true;
258                        if let Some(waker) = state.on_connected.take() {
259                            waker.wake()
260                        }
261                        if state.conn.side().is_client() && !state.conn.accepted_0rtt() {
262                            // Wake up rejected 0-RTT streams so they can fail
263                            // immediately with
264                            // `ZeroRttRejected` errors.
265                            wake_all_streams(&mut state.writable);
266                            wake_all_streams(&mut state.readable);
267                            wake_all_streams(&mut state.stopped);
268                        }
269                    }
270                    ConnectionLost { reason } => state.terminate(reason.into()),
271                    Stream(StreamEvent::Readable { id }) => wake_stream(id, &mut state.readable),
272                    Stream(StreamEvent::Writable { id }) => wake_stream(id, &mut state.writable),
273                    Stream(StreamEvent::Finished { id }) => wake_stream(id, &mut state.stopped),
274                    Stream(StreamEvent::Stopped { id, .. }) => {
275                        wake_stream(id, &mut state.stopped);
276                        wake_stream(id, &mut state.writable);
277                    }
278                    Stream(StreamEvent::Available { dir }) => state.stream_available[dir as usize]
279                        .drain(..)
280                        .for_each(Waker::wake),
281                    Stream(StreamEvent::Opened { dir }) => state.stream_opened[dir as usize]
282                        .drain(..)
283                        .for_each(Waker::wake),
284                    DatagramReceived => state.datagram_received.drain(..).for_each(Waker::wake),
285                    DatagramsUnblocked => state.datagrams_unblocked.drain(..).for_each(Waker::wake),
286                }
287            }
288
289            if state.conn.is_drained() {
290                break;
291            }
292        }
293
294        // Break the reference cycle.
295        if let Some(worker) = self.state().worker.take() {
296            worker.detach();
297        }
298    }
299}
300
301macro_rules! conn_fn {
302    () => {
303        /// The side of the connection (client or server)
304        pub fn side(&self) -> Side {
305            self.0.state().conn.side()
306        }
307
308        /// The local IP address which was used when the peer established
309        /// the connection.
310        ///
311        /// This can be different from the address the endpoint is bound to, in case
312        /// the endpoint is bound to a wildcard address like `0.0.0.0` or `::`.
313        ///
314        /// This will return `None` for clients, or when the platform does not
315        /// expose this information.
316        pub fn local_ip(&self) -> Option<IpAddr> {
317            self.0.state().conn.local_ip()
318        }
319
320        /// The peer's UDP address.
321        ///
322        /// Will panic if called after `poll` has returned `Ready`.
323        pub fn remote_address(&self) -> SocketAddr {
324            self.0.state().conn.remote_address()
325        }
326
327        /// Current best estimate of this connection's latency (round-trip-time).
328        pub fn rtt(&self) -> Duration {
329            self.0.state().conn.rtt()
330        }
331
332        /// Connection statistics.
333        pub fn stats(&self) -> ConnectionStats {
334            self.0.state().conn.stats()
335        }
336
337        /// Current state of the congestion control algorithm. (For debugging
338        /// purposes)
339        pub fn congestion_state(&self) -> Box<dyn Controller> {
340            self.0.state().conn.congestion_state().clone_box()
341        }
342
343        /// Cryptographic identity of the peer.
344        pub fn peer_identity(
345            &self,
346        ) -> Option<Box<Vec<rustls::pki_types::CertificateDer<'static>>>> {
347            self.0
348                .state()
349                .conn
350                .crypto_session()
351                .peer_identity()
352                .map(|v| v.downcast().unwrap())
353        }
354
355        /// A stable identifier for this connection
356        ///
357        /// Peer addresses and connection IDs can change, but this value will remain
358        /// fixed for the lifetime of the connection.
359        pub fn stable_id(&self) -> usize {
360            Shared::as_ptr(&self.0) as usize
361        }
362
363        /// Derive keying material from this connection's TLS session secrets.
364        ///
365        /// When both peers call this method with the same `label` and `context`
366        /// arguments and `output` buffers of equal length, they will get the
367        /// same sequence of bytes in `output`. These bytes are cryptographically
368        /// strong and pseudorandom, and are suitable for use as keying material.
369        ///
370        /// This function fails if called with an empty `output` or called prior to
371        /// the handshake completing.
372        ///
373        /// See [RFC5705](https://tools.ietf.org/html/rfc5705) for more information.
374        pub fn export_keying_material(
375            &self,
376            output: &mut [u8],
377            label: &[u8],
378            context: &[u8],
379        ) -> Result<(), quinn_proto::crypto::ExportKeyingMaterialError> {
380            self.0
381                .state()
382                .conn
383                .crypto_session()
384                .export_keying_material(output, label, context)
385        }
386    };
387}
388
389/// In-progress connection attempt future
390#[derive(Debug)]
391#[must_use = "futures/streams/sinks do nothing unless you `.await` or poll them"]
392pub struct Connecting(Shared<ConnectionInner>);
393
394impl Connecting {
395    conn_fn!();
396
397    #[track_caller]
398    pub(crate) fn new(
399        handle: ConnectionHandle,
400        conn: quinn_proto::Connection,
401        socket: Socket,
402        events_tx: Sender<(ConnectionHandle, EndpointEvent)>,
403        events_rx: Receiver<ConnectionEvent>,
404    ) -> Self {
405        let inner = Shared::new(ConnectionInner::new(
406            handle, conn, socket, events_tx, events_rx,
407        ));
408        // Name the task: the caller of this is compio rather than user code, so
409        // its location alone does not say what the task is.
410        // Unlike the endpoint's worker, this one can be attributed to the user:
411        // every path here is a plain `fn`, so the caller reaches us and the
412        // console tells a connection they made apart from one they accepted.
413        let meta = SpawnMeta::capture().named("quic::connection");
414        let worker = compio_runtime::spawn_at(
415            {
416                let inner = inner.clone();
417                async move { inner.run().await }.in_current_span()
418            },
419            meta,
420        );
421        inner.state().worker = Some(worker);
422        Self(inner)
423    }
424
425    /// Parameters negotiated during the handshake.
426    #[cfg(rustls)]
427    pub async fn handshake_data(&mut self) -> Result<Box<HandshakeData>, ConnectionError> {
428        future::poll_fn(|cx| {
429            let mut state = self.0.try_state()?;
430            if let Some(data) = state.handshake_data() {
431                return Poll::Ready(Ok(data));
432            }
433
434            match &state.on_handshake_data {
435                Some(waker) if waker.will_wake(cx.waker()) => {}
436                _ => state.on_handshake_data = Some(cx.waker().clone()),
437            }
438
439            Poll::Pending
440        })
441        .await
442    }
443
444    /// Convert into a 0-RTT or 0.5-RTT connection at the cost of weakened
445    /// security.
446    ///
447    /// Returns `Ok` immediately if the local endpoint is able to attempt
448    /// sending 0/0.5-RTT data. If so, the returned [`Connection`] can be used
449    /// to send application data without waiting for the rest of the handshake
450    /// to complete, at the cost of weakened cryptographic security guarantees.
451    /// The [`Connection::accepted_0rtt`] method resolves when the handshake
452    /// does complete, at which point subsequently opened streams and written
453    /// data will have full cryptographic protection.
454    ///
455    /// ## Outgoing
456    ///
457    /// For outgoing connections, the initial attempt to convert to a
458    /// [`Connection`] which sends 0-RTT data will proceed if the
459    /// [`crypto::ClientConfig`][crate::crypto::ClientConfig] attempts to resume
460    /// a previous TLS session. However, **the remote endpoint may not actually
461    /// _accept_ the 0-RTT data**--yet still accept the connection attempt in
462    /// general. This possibility is conveyed through the
463    /// [`Connection::accepted_0rtt`] method--when the handshake completes, it
464    /// resolves to true if the 0-RTT data was accepted and false if it was
465    /// rejected. If it was rejected, the existence of streams opened and other
466    /// application data sent prior to the handshake completing will not be
467    /// conveyed to the remote application, and local operations on them will
468    /// return `ZeroRttRejected` errors.
469    ///
470    /// A server may reject 0-RTT data at its discretion, but accepting 0-RTT
471    /// data requires the relevant resumption state to be stored in the server,
472    /// which servers may limit or lose for various reasons including not
473    /// persisting resumption state across server restarts.
474    ///
475    /// ## Incoming
476    ///
477    /// For incoming connections, conversion to 0.5-RTT will always fully
478    /// succeed. `into_0rtt` will always return `Ok` and
479    /// [`Connection::accepted_0rtt`] will always resolve to true.
480    ///
481    /// ## Security
482    ///
483    /// On outgoing connections, this enables transmission of 0-RTT data, which
484    /// is vulnerable to replay attacks, and should therefore never invoke
485    /// non-idempotent operations.
486    ///
487    /// On incoming connections, this enables transmission of 0.5-RTT data,
488    /// which may be sent before TLS client authentication has occurred, and
489    /// should therefore not be used to send data for which client
490    /// authentication is being used.
491    pub fn into_0rtt(self) -> Result<Connection, Self> {
492        let is_ok = {
493            let state = self.0.state();
494            state.conn.has_0rtt() || state.conn.side().is_server()
495        };
496        if is_ok {
497            Ok(Connection(self.0.clone()))
498        } else {
499            Err(self)
500        }
501    }
502}
503
504impl Future for Connecting {
505    type Output = Result<Connection, ConnectionError>;
506
507    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
508        let mut state = self.0.try_state()?;
509
510        if state.connected {
511            return Poll::Ready(Ok(Connection(self.0.clone())));
512        }
513
514        match &state.on_connected {
515            Some(waker) if waker.will_wake(cx.waker()) => {}
516            _ => state.on_connected = Some(cx.waker().clone()),
517        }
518
519        Poll::Pending
520    }
521}
522
523impl Drop for Connecting {
524    fn drop(&mut self) {
525        implicit_close(&self.0)
526    }
527}
528
529/// A QUIC connection.
530#[derive(Debug, Clone)]
531pub struct Connection(Shared<ConnectionInner>);
532
533impl Connection {
534    conn_fn!();
535
536    /// Update traffic keys spontaneously
537    ///
538    /// This primarily exists for testing purposes.
539    pub fn force_key_update(&self) {
540        self.0.state().conn.force_key_update()
541    }
542
543    /// Parameters negotiated during the handshake.
544    #[cfg(rustls)]
545    pub fn handshake_data(&mut self) -> Result<Box<HandshakeData>, ConnectionError> {
546        Ok(self.0.try_state()?.handshake_data().unwrap())
547    }
548
549    /// Compute the maximum size of datagrams that may be passed to
550    /// [`send_datagram()`](Self::send_datagram).
551    ///
552    /// Returns `None` if datagrams are unsupported by the peer or disabled
553    /// locally.
554    ///
555    /// This may change over the lifetime of a connection according to variation
556    /// in the path MTU estimate. The peer can also enforce an arbitrarily small
557    /// fixed limit, but if the peer's limit is large this is guaranteed to be a
558    /// little over a kilobyte at minimum.
559    ///
560    /// Not necessarily the maximum size of received datagrams.
561    pub fn max_datagram_size(&self) -> Option<usize> {
562        self.0.state().conn.datagrams().max_size()
563    }
564
565    /// Bytes available in the outgoing datagram buffer.
566    ///
567    /// When greater than zero, calling [`send_datagram()`](Self::send_datagram)
568    /// with a datagram of at most this size is guaranteed not to cause older
569    /// datagrams to be dropped.
570    pub fn datagram_send_buffer_space(&self) -> usize {
571        self.0.state().conn.datagrams().send_buffer_space()
572    }
573
574    /// Modify the number of remotely initiated unidirectional streams that may
575    /// be concurrently open.
576    ///
577    /// No streams may be opened by the peer unless fewer than `count` are
578    /// already open. Large `count`s increase both minimum and worst-case
579    /// memory consumption.
580    pub fn set_max_concurrent_uni_streams(&self, count: VarInt) {
581        let mut state = self.0.state();
582        state.conn.set_max_concurrent_streams(Dir::Uni, count);
583        // May need to send MAX_STREAMS to make progress
584        state.wake();
585    }
586
587    /// See [`quinn_proto::TransportConfig::send_window()`]
588    pub fn set_send_window(&self, send_window: u64) {
589        let mut state = self.0.state();
590        state.conn.set_send_window(send_window);
591        state.wake();
592    }
593
594    /// See [`quinn_proto::TransportConfig::receive_window()`]
595    pub fn set_receive_window(&self, receive_window: VarInt) {
596        let mut state = self.0.state();
597        state.conn.set_receive_window(receive_window);
598        state.wake();
599    }
600
601    /// Modify the number of remotely initiated bidirectional streams that may
602    /// be concurrently open.
603    ///
604    /// No streams may be opened by the peer unless fewer than `count` are
605    /// already open. Large `count`s increase both minimum and worst-case
606    /// memory consumption.
607    pub fn set_max_concurrent_bi_streams(&self, count: VarInt) {
608        let mut state = self.0.state();
609        state.conn.set_max_concurrent_streams(Dir::Bi, count);
610        // May need to send MAX_STREAMS to make progress
611        state.wake();
612    }
613
614    /// Close the connection immediately.
615    ///
616    /// Pending operations will fail immediately with
617    /// [`ConnectionError::LocallyClosed`]. No more data is sent to the peer
618    /// and the peer may drop buffered data upon receiving
619    /// the CONNECTION_CLOSE frame.
620    ///
621    /// `error_code` and `reason` are not interpreted, and are provided directly
622    /// to the peer.
623    ///
624    /// `reason` will be truncated to fit in a single packet with overhead; to
625    /// improve odds that it is preserved in full, it should be kept under
626    /// 1KiB.
627    ///
628    /// # Gracefully closing a connection
629    ///
630    /// Only the peer last receiving application data can be certain that all
631    /// data is delivered. The only reliable action it can then take is to
632    /// close the connection, potentially with a custom error code. The
633    /// delivery of the final CONNECTION_CLOSE frame is very likely if both
634    /// endpoints stay online long enough, and [`Endpoint::shutdown()`] can
635    /// be used to provide sufficient time. Otherwise, the remote peer will
636    /// time out the connection, provided that the idle timeout is not
637    /// disabled.
638    ///
639    /// The sending side can not guarantee all stream data is delivered to the
640    /// remote application. It only knows the data is delivered to the QUIC
641    /// stack of the remote endpoint. Once the local side sends a
642    /// CONNECTION_CLOSE frame in response to calling [`close()`] the remote
643    /// endpoint may drop any data it received but is as yet undelivered to
644    /// the application, including data that was acknowledged as received to
645    /// the local endpoint.
646    ///
647    /// [`ConnectionError::LocallyClosed`]: ConnectionError::LocallyClosed
648    /// [`Endpoint::shutdown()`]: crate::Endpoint::shutdown
649    /// [`close()`]: Connection::close
650    pub fn close(&self, error_code: VarInt, reason: &[u8]) {
651        self.0
652            .state()
653            .close(error_code, Bytes::copy_from_slice(reason));
654    }
655
656    /// Wait for the connection to be closed for any reason.
657    pub async fn closed(&self) -> ConnectionError {
658        let worker = self.0.state().worker.take();
659        if let Some(worker) = worker {
660            let _ = worker.await;
661        }
662
663        self.0.try_state().unwrap_err()
664    }
665
666    /// If the connection is closed, the reason why.
667    ///
668    /// Returns `None` if the connection is still open.
669    pub fn close_reason(&self) -> Option<ConnectionError> {
670        self.0.try_state().err()
671    }
672
673    fn poll_recv_datagram(&self, cx: &mut Context) -> Poll<Result<Bytes, ConnectionError>> {
674        let mut state = self.0.try_state()?;
675        if let Some(bytes) = state.conn.datagrams().recv() {
676            return Poll::Ready(Ok(bytes));
677        }
678        state.datagram_received.push_back(cx.waker().clone());
679        Poll::Pending
680    }
681
682    /// Try to receive an application datagram. Returns None if no datagram is
683    /// available.
684    pub fn try_recv_datagram(&self) -> Result<Option<Bytes>, ConnectionError> {
685        let mut state = self.0.try_state()?;
686        Ok(state.conn.datagrams().recv())
687    }
688
689    /// Receive an application datagram.
690    pub async fn recv_datagram(&self) -> Result<Bytes, ConnectionError> {
691        future::poll_fn(|cx| self.poll_recv_datagram(cx)).await
692    }
693
694    fn try_send_datagram(
695        &self,
696        cx: Option<&mut Context>,
697        data: Bytes,
698    ) -> Result<(), Result<SendDatagramError, Bytes>> {
699        use quinn_proto::SendDatagramError::*;
700        let mut state = self.0.try_state().map_err(|e| Ok(e.into()))?;
701        state
702            .conn
703            .datagrams()
704            .send(data, cx.is_none())
705            .map_err(|err| match err {
706                UnsupportedByPeer => Ok(SendDatagramError::UnsupportedByPeer),
707                Disabled => Ok(SendDatagramError::Disabled),
708                TooLarge => Ok(SendDatagramError::TooLarge),
709                Blocked(data) => {
710                    state
711                        .datagrams_unblocked
712                        .push_back(cx.unwrap().waker().clone());
713                    Err(data)
714                }
715            })?;
716        state.wake();
717        Ok(())
718    }
719
720    /// Transmit `data` as an unreliable, unordered application datagram.
721    ///
722    /// Application datagrams are a low-level primitive. They may be lost or
723    /// delivered out of order, and `data` must both fit inside a single
724    /// QUIC packet and be smaller than the maximum dictated by the peer.
725    pub fn send_datagram(&self, data: Bytes) -> Result<(), SendDatagramError> {
726        self.try_send_datagram(None, data).map_err(Result::unwrap)
727    }
728
729    /// Transmit `data` as an unreliable, unordered application datagram.
730    ///
731    /// Unlike [`send_datagram()`], this method will wait for buffer space
732    /// during congestion conditions, which effectively prioritizes old
733    /// datagrams over new datagrams.
734    ///
735    /// See [`send_datagram()`] for details.
736    ///
737    /// [`send_datagram()`]: Connection::send_datagram
738    pub async fn send_datagram_wait(&self, data: Bytes) -> Result<(), SendDatagramError> {
739        let mut data = Some(data);
740        future::poll_fn(
741            |cx| match self.try_send_datagram(Some(cx), data.take().unwrap()) {
742                Ok(()) => Poll::Ready(Ok(())),
743                Err(Ok(e)) => Poll::Ready(Err(e)),
744                Err(Err(b)) => {
745                    data.replace(b);
746                    Poll::Pending
747                }
748            },
749        )
750        .await
751    }
752
753    fn poll_open_stream(
754        &self,
755        cx: Option<&mut Context>,
756        dir: Dir,
757    ) -> Poll<Result<(StreamId, bool), ConnectionError>> {
758        let mut state = self.0.try_state()?;
759        if let Some(stream) = state.conn.streams().open(dir) {
760            Poll::Ready(Ok((
761                stream,
762                state.conn.side().is_client() && state.conn.is_handshaking(),
763            )))
764        } else {
765            if let Some(cx) = cx {
766                state.stream_available[dir as usize].push_back(cx.waker().clone());
767            }
768            Poll::Pending
769        }
770    }
771
772    /// Initiate a new outgoing unidirectional stream.
773    ///
774    /// Streams are cheap and instantaneous to open. As a consequence, the peer
775    /// won't be notified that a stream has been opened until the stream is
776    /// actually used.
777    pub fn open_uni(&self) -> Result<SendStream, OpenStreamError> {
778        if let Poll::Ready((stream, is_0rtt)) = self.poll_open_stream(None, Dir::Uni)? {
779            Ok(SendStream::new(self.0.clone(), stream, is_0rtt))
780        } else {
781            Err(OpenStreamError::StreamsExhausted)
782        }
783    }
784
785    /// Initiate a new outgoing unidirectional stream.
786    ///
787    /// Unlike [`open_uni()`], this method will wait for the connection to allow
788    /// a new stream to be opened.
789    ///
790    /// See [`open_uni()`] for details.
791    ///
792    /// [`open_uni()`]: crate::Connection::open_uni
793    pub async fn open_uni_wait(&self) -> Result<SendStream, ConnectionError> {
794        let (stream, is_0rtt) =
795            future::poll_fn(|cx| self.poll_open_stream(Some(cx), Dir::Uni)).await?;
796        Ok(SendStream::new(self.0.clone(), stream, is_0rtt))
797    }
798
799    /// Initiate a new outgoing bidirectional stream.
800    ///
801    /// Streams are cheap and instantaneous to open. As a consequence, the peer
802    /// won't be notified that a stream has been opened until the stream is
803    /// actually used.
804    pub fn open_bi(&self) -> Result<(SendStream, RecvStream), OpenStreamError> {
805        if let Poll::Ready((stream, is_0rtt)) = self.poll_open_stream(None, Dir::Bi)? {
806            Ok((
807                SendStream::new(self.0.clone(), stream, is_0rtt),
808                RecvStream::new(self.0.clone(), stream, is_0rtt),
809            ))
810        } else {
811            Err(OpenStreamError::StreamsExhausted)
812        }
813    }
814
815    /// Initiate a new outgoing bidirectional stream.
816    ///
817    /// Unlike [`open_bi()`], this method will wait for the connection to allow
818    /// a new stream to be opened.
819    ///
820    /// See [`open_bi()`] for details.
821    ///
822    /// [`open_bi()`]: crate::Connection::open_bi
823    pub async fn open_bi_wait(&self) -> Result<(SendStream, RecvStream), ConnectionError> {
824        let (stream, is_0rtt) =
825            future::poll_fn(|cx| self.poll_open_stream(Some(cx), Dir::Bi)).await?;
826        Ok((
827            SendStream::new(self.0.clone(), stream, is_0rtt),
828            RecvStream::new(self.0.clone(), stream, is_0rtt),
829        ))
830    }
831
832    fn poll_accept_stream(
833        &self,
834        cx: &mut Context,
835        dir: Dir,
836    ) -> Poll<Result<(StreamId, bool), ConnectionError>> {
837        let mut state = self.0.try_state()?;
838        if let Some(stream) = state.conn.streams().accept(dir) {
839            state.wake();
840            Poll::Ready(Ok((stream, state.conn.is_handshaking())))
841        } else {
842            state.stream_opened[dir as usize].push_back(cx.waker().clone());
843            Poll::Pending
844        }
845    }
846
847    /// Accept the next incoming uni-directional stream
848    pub async fn accept_uni(&self) -> Result<RecvStream, ConnectionError> {
849        let (stream, is_0rtt) = future::poll_fn(|cx| self.poll_accept_stream(cx, Dir::Uni)).await?;
850        Ok(RecvStream::new(self.0.clone(), stream, is_0rtt))
851    }
852
853    /// Accept the next incoming bidirectional stream
854    ///
855    /// **Important Note**: The `Connection` that calls [`open_bi()`] must write
856    /// to its [`SendStream`] before the other `Connection` is able to
857    /// `accept_bi()`. Calling [`open_bi()`] then waiting on the [`RecvStream`]
858    /// without writing anything to [`SendStream`] will never succeed.
859    ///
860    /// [`accept_bi()`]: crate::Connection::accept_bi
861    /// [`open_bi()`]: crate::Connection::open_bi
862    /// [`SendStream`]: crate::SendStream
863    /// [`RecvStream`]: crate::RecvStream
864    pub async fn accept_bi(&self) -> Result<(SendStream, RecvStream), ConnectionError> {
865        let (stream, is_0rtt) = future::poll_fn(|cx| self.poll_accept_stream(cx, Dir::Bi)).await?;
866        Ok((
867            SendStream::new(self.0.clone(), stream, is_0rtt),
868            RecvStream::new(self.0.clone(), stream, is_0rtt),
869        ))
870    }
871
872    /// Wait for the connection to be fully established.
873    ///
874    /// For clients, the resulting value indicates if 0-RTT was accepted. For
875    /// servers, the resulting value is meaningless.
876    pub async fn accepted_0rtt(&self) -> Result<bool, ConnectionError> {
877        future::poll_fn(|cx| {
878            let mut state = self.0.try_state()?;
879
880            if state.connected {
881                return Poll::Ready(Ok(state.conn.accepted_0rtt()));
882            }
883
884            match &state.on_connected {
885                Some(waker) if waker.will_wake(cx.waker()) => {}
886                _ => state.on_connected = Some(cx.waker().clone()),
887            }
888
889            Poll::Pending
890        })
891        .await
892    }
893}
894
895impl PartialEq for Connection {
896    fn eq(&self, other: &Self) -> bool {
897        Shared::ptr_eq(&self.0, &other.0)
898    }
899}
900
901impl Eq for Connection {}
902
903impl Drop for Connection {
904    fn drop(&mut self) {
905        implicit_close(&self.0)
906    }
907}
908
909struct Timer {
910    deadline: Option<Instant>,
911    fut: Fuse<LocalBoxFuture<'static, ()>>,
912}
913
914impl Timer {
915    fn new() -> Self {
916        Self {
917            deadline: None,
918            fut: Fuse::terminated(),
919        }
920    }
921
922    fn reset(&mut self, deadline: Option<Instant>) {
923        if let Some(deadline) = deadline {
924            if self.deadline.is_none() || self.deadline != Some(deadline) {
925                self.fut = compio_runtime::time::sleep_until(deadline)
926                    .boxed_local()
927                    .fuse();
928            }
929        } else {
930            self.fut = Fuse::terminated();
931        }
932        self.deadline = deadline;
933    }
934}
935
936impl Future for Timer {
937    type Output = ();
938
939    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
940        self.fut.poll_unpin(cx)
941    }
942}
943
944impl FusedFuture for Timer {
945    fn is_terminated(&self) -> bool {
946        self.fut.is_terminated()
947    }
948}
949
950/// Reasons why a connection might be lost
951#[derive(Debug, Error, Clone, PartialEq, Eq)]
952pub enum ConnectionError {
953    /// The peer doesn't implement any supported version
954    #[error("peer doesn't implement any supported version")]
955    VersionMismatch,
956    /// The peer violated the QUIC specification as understood by this
957    /// implementation
958    #[error(transparent)]
959    TransportError(#[from] quinn_proto::TransportError),
960    /// The peer's QUIC stack aborted the connection automatically
961    #[error("aborted by peer: {0}")]
962    ConnectionClosed(quinn_proto::ConnectionClose),
963    /// The peer closed the connection
964    #[error("closed by peer: {0}")]
965    ApplicationClosed(quinn_proto::ApplicationClose),
966    /// The peer is unable to continue processing this connection, usually due
967    /// to having restarted
968    #[error("reset by peer")]
969    Reset,
970    /// Communication with the peer has lapsed for longer than the negotiated
971    /// idle timeout
972    ///
973    /// If neither side is sending keep-alives, a connection will time out after
974    /// a long enough idle period even if the peer is still reachable. See
975    /// also [`TransportConfig::max_idle_timeout()`](quinn_proto::TransportConfig::max_idle_timeout())
976    /// and [`TransportConfig::keep_alive_interval()`](quinn_proto::TransportConfig::keep_alive_interval()).
977    #[error("timed out")]
978    TimedOut,
979    /// The local application closed the connection
980    #[error("closed")]
981    LocallyClosed,
982    /// The connection could not be created because not enough of the CID space
983    /// is available
984    ///
985    /// Try using longer connection IDs.
986    #[error("CIDs exhausted")]
987    CidsExhausted,
988}
989
990impl From<quinn_proto::ConnectionError> for ConnectionError {
991    fn from(value: quinn_proto::ConnectionError) -> Self {
992        use quinn_proto::ConnectionError::*;
993
994        match value {
995            VersionMismatch => ConnectionError::VersionMismatch,
996            TransportError(e) => ConnectionError::TransportError(e),
997            ConnectionClosed(e) => ConnectionError::ConnectionClosed(e),
998            ApplicationClosed(e) => ConnectionError::ApplicationClosed(e),
999            Reset => ConnectionError::Reset,
1000            TimedOut => ConnectionError::TimedOut,
1001            LocallyClosed => ConnectionError::LocallyClosed,
1002            CidsExhausted => ConnectionError::CidsExhausted,
1003        }
1004    }
1005}
1006
1007/// Errors that can arise when sending a datagram
1008#[derive(Debug, Error, Clone, Eq, PartialEq)]
1009pub enum SendDatagramError {
1010    /// The peer does not support receiving datagram frames
1011    #[error("datagrams not supported by peer")]
1012    UnsupportedByPeer,
1013    /// Datagram support is disabled locally
1014    #[error("datagram support disabled")]
1015    Disabled,
1016    /// The datagram is larger than the connection can currently accommodate
1017    ///
1018    /// Indicates that the path MTU minus overhead or the limit advertised by
1019    /// the peer has been exceeded.
1020    #[error("datagram too large")]
1021    TooLarge,
1022    /// The connection was lost
1023    #[error("connection lost")]
1024    ConnectionLost(#[from] ConnectionError),
1025}
1026
1027/// Errors that can arise when trying to open a stream
1028#[derive(Debug, Error, Clone, Eq, PartialEq)]
1029pub enum OpenStreamError {
1030    /// The connection was lost
1031    #[error("connection lost")]
1032    ConnectionLost(#[from] ConnectionError),
1033    /// The streams in the given direction are currently exhausted
1034    #[error("streams exhausted")]
1035    StreamsExhausted,
1036}
1037
1038#[cfg(feature = "h3")]
1039pub(crate) mod h3_impl {
1040    use std::sync::Arc;
1041
1042    use compio_buf::bytes::Buf;
1043    use futures_util::ready;
1044    use h3::{
1045        error::Code,
1046        quic::{self, ConnectionErrorIncoming, StreamErrorIncoming, WriteBuf},
1047    };
1048    use h3_datagram::{
1049        datagram::EncodedDatagram,
1050        quic_traits::{
1051            DatagramConnectionExt, RecvDatagram, SendDatagram, SendDatagramErrorIncoming,
1052        },
1053    };
1054
1055    use super::*;
1056    use crate::send_stream::h3_impl::SendStream;
1057
1058    impl From<ConnectionError> for ConnectionErrorIncoming {
1059        fn from(e: ConnectionError) -> Self {
1060            use ConnectionError::*;
1061            match e {
1062                ApplicationClosed(e) => Self::ApplicationClose {
1063                    error_code: e.error_code.into_inner(),
1064                },
1065                TimedOut => Self::Timeout,
1066
1067                e => Self::Undefined(Arc::new(e)),
1068            }
1069        }
1070    }
1071
1072    impl From<ConnectionError> for StreamErrorIncoming {
1073        fn from(e: ConnectionError) -> Self {
1074            Self::ConnectionErrorIncoming {
1075                connection_error: e.into(),
1076            }
1077        }
1078    }
1079
1080    impl From<SendDatagramError> for SendDatagramErrorIncoming {
1081        fn from(e: SendDatagramError) -> Self {
1082            use SendDatagramError::*;
1083            match e {
1084                UnsupportedByPeer | Disabled => Self::NotAvailable,
1085                TooLarge => Self::TooLarge,
1086                ConnectionLost(e) => Self::ConnectionError(e.into()),
1087            }
1088        }
1089    }
1090
1091    impl<B> SendDatagram<B> for Connection
1092    where
1093        B: Buf,
1094    {
1095        fn send_datagram<T: Into<EncodedDatagram<B>>>(
1096            &mut self,
1097            data: T,
1098        ) -> Result<(), SendDatagramErrorIncoming> {
1099            let mut buf: EncodedDatagram<B> = data.into();
1100            let buf = buf.copy_to_bytes(buf.remaining());
1101            Ok(Connection::send_datagram(self, buf)?)
1102        }
1103    }
1104
1105    impl RecvDatagram for Connection {
1106        type Buffer = Bytes;
1107
1108        fn poll_incoming_datagram(
1109            &mut self,
1110            cx: &mut core::task::Context<'_>,
1111        ) -> Poll<Result<Self::Buffer, ConnectionErrorIncoming>> {
1112            Poll::Ready(Ok(ready!(self.poll_recv_datagram(cx))?))
1113        }
1114    }
1115
1116    impl<B: Buf> DatagramConnectionExt<B> for Connection {
1117        type RecvDatagramHandler = Self;
1118        type SendDatagramHandler = Self;
1119
1120        fn send_datagram_handler(&self) -> Self::SendDatagramHandler {
1121            self.clone()
1122        }
1123
1124        fn recv_datagram_handler(&self) -> Self::RecvDatagramHandler {
1125            self.clone()
1126        }
1127    }
1128
1129    /// Bidirectional stream.
1130    pub struct BidiStream<B> {
1131        send: SendStream<B>,
1132        recv: RecvStream,
1133    }
1134
1135    impl<B> BidiStream<B> {
1136        pub(crate) fn new(conn: Shared<ConnectionInner>, stream: StreamId, is_0rtt: bool) -> Self {
1137            Self {
1138                send: SendStream::new(conn.clone(), stream, is_0rtt),
1139                recv: RecvStream::new(conn, stream, is_0rtt),
1140            }
1141        }
1142    }
1143
1144    impl<B> quic::BidiStream<B> for BidiStream<B>
1145    where
1146        B: Buf,
1147    {
1148        type RecvStream = RecvStream;
1149        type SendStream = SendStream<B>;
1150
1151        fn split(self) -> (Self::SendStream, Self::RecvStream) {
1152            (self.send, self.recv)
1153        }
1154    }
1155
1156    impl<B> quic::RecvStream for BidiStream<B>
1157    where
1158        B: Buf,
1159    {
1160        type Buf = Bytes;
1161
1162        fn poll_data(
1163            &mut self,
1164            cx: &mut Context<'_>,
1165        ) -> Poll<Result<Option<Self::Buf>, StreamErrorIncoming>> {
1166            self.recv.poll_data(cx)
1167        }
1168
1169        fn stop_sending(&mut self, error_code: u64) {
1170            self.recv.stop_sending(error_code)
1171        }
1172
1173        fn recv_id(&self) -> quic::StreamId {
1174            self.recv.recv_id()
1175        }
1176    }
1177
1178    impl<B> quic::SendStream<B> for BidiStream<B>
1179    where
1180        B: Buf,
1181    {
1182        fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), StreamErrorIncoming>> {
1183            self.send.poll_ready(cx)
1184        }
1185
1186        fn send_data<T: Into<WriteBuf<B>>>(&mut self, data: T) -> Result<(), StreamErrorIncoming> {
1187            self.send.send_data(data)
1188        }
1189
1190        fn poll_finish(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), StreamErrorIncoming>> {
1191            self.send.poll_finish(cx)
1192        }
1193
1194        fn reset(&mut self, reset_code: u64) {
1195            self.send.reset(reset_code)
1196        }
1197
1198        fn send_id(&self) -> quic::StreamId {
1199            self.send.send_id()
1200        }
1201    }
1202
1203    impl<B> quic::SendStreamUnframed<B> for BidiStream<B>
1204    where
1205        B: Buf,
1206    {
1207        fn poll_send<D: Buf>(
1208            &mut self,
1209            cx: &mut Context<'_>,
1210            buf: &mut D,
1211        ) -> Poll<Result<usize, StreamErrorIncoming>> {
1212            self.send.poll_send(cx, buf)
1213        }
1214    }
1215
1216    /// Stream opener.
1217    #[derive(Clone)]
1218    pub struct OpenStreams(Connection);
1219
1220    impl<B> quic::OpenStreams<B> for OpenStreams
1221    where
1222        B: Buf,
1223    {
1224        type BidiStream = BidiStream<B>;
1225        type SendStream = SendStream<B>;
1226
1227        fn poll_open_bidi(
1228            &mut self,
1229            cx: &mut Context<'_>,
1230        ) -> Poll<Result<Self::BidiStream, StreamErrorIncoming>> {
1231            let (stream, is_0rtt) = ready!(self.0.poll_open_stream(Some(cx), Dir::Bi))?;
1232            Poll::Ready(Ok(BidiStream::new(self.0.0.clone(), stream, is_0rtt)))
1233        }
1234
1235        fn poll_open_send(
1236            &mut self,
1237            cx: &mut Context<'_>,
1238        ) -> Poll<Result<Self::SendStream, StreamErrorIncoming>> {
1239            let (stream, is_0rtt) = ready!(self.0.poll_open_stream(Some(cx), Dir::Uni))?;
1240            Poll::Ready(Ok(SendStream::new(self.0.0.clone(), stream, is_0rtt)))
1241        }
1242
1243        fn close(&mut self, code: Code, reason: &[u8]) {
1244            self.0
1245                .close(code.value().try_into().expect("invalid code"), reason)
1246        }
1247    }
1248
1249    impl<B> quic::OpenStreams<B> for Connection
1250    where
1251        B: Buf,
1252    {
1253        type BidiStream = BidiStream<B>;
1254        type SendStream = SendStream<B>;
1255
1256        fn poll_open_bidi(
1257            &mut self,
1258            cx: &mut Context<'_>,
1259        ) -> Poll<Result<Self::BidiStream, StreamErrorIncoming>> {
1260            let (stream, is_0rtt) = ready!(self.poll_open_stream(Some(cx), Dir::Bi))?;
1261            Poll::Ready(Ok(BidiStream::new(self.0.clone(), stream, is_0rtt)))
1262        }
1263
1264        fn poll_open_send(
1265            &mut self,
1266            cx: &mut Context<'_>,
1267        ) -> Poll<Result<Self::SendStream, StreamErrorIncoming>> {
1268            let (stream, is_0rtt) = ready!(self.poll_open_stream(Some(cx), Dir::Uni))?;
1269            Poll::Ready(Ok(SendStream::new(self.0.clone(), stream, is_0rtt)))
1270        }
1271
1272        fn close(&mut self, code: Code, reason: &[u8]) {
1273            Connection::close(self, code.value().try_into().expect("invalid code"), reason)
1274        }
1275    }
1276
1277    impl<B> quic::Connection<B> for Connection
1278    where
1279        B: Buf,
1280    {
1281        type OpenStreams = OpenStreams;
1282        type RecvStream = RecvStream;
1283
1284        fn poll_accept_recv(
1285            &mut self,
1286            cx: &mut std::task::Context<'_>,
1287        ) -> Poll<Result<Self::RecvStream, ConnectionErrorIncoming>> {
1288            let (stream, is_0rtt) = ready!(self.poll_accept_stream(cx, Dir::Uni))?;
1289            Poll::Ready(Ok(RecvStream::new(self.0.clone(), stream, is_0rtt)))
1290        }
1291
1292        fn poll_accept_bidi(
1293            &mut self,
1294            cx: &mut std::task::Context<'_>,
1295        ) -> Poll<Result<Self::BidiStream, ConnectionErrorIncoming>> {
1296            let (stream, is_0rtt) = ready!(self.poll_accept_stream(cx, Dir::Bi))?;
1297            Poll::Ready(Ok(BidiStream::new(self.0.clone(), stream, is_0rtt)))
1298        }
1299
1300        fn opener(&self) -> Self::OpenStreams {
1301            OpenStreams(self.clone())
1302        }
1303    }
1304}
compio_quic/connection.rs

compio_quic/
connection.rs
