compio_quic/
endpoint.rs

1use std::{
2    collections::VecDeque,
3    fmt::Debug,
4    future::{Future, poll_fn},
5    io,
6    mem::ManuallyDrop,
7    net::{SocketAddr, SocketAddrV6},
8    ops::Deref,
9    pin::pin,
10    ptr,
11    sync::Arc,
12    task::{Context, Poll, Waker},
13    time::Instant,
14};
15
16use compio_buf::{BufResult, bytes::Bytes};
17use compio_log::{Instrument, error};
18#[cfg(rustls)]
19use compio_net::ToSocketAddrsAsync;
20use compio_net::UdpSocket;
21use compio_runtime::{JoinHandle, SpawnMeta};
22use flume::{Receiver, Sender, unbounded};
23use futures_util::{FutureExt, StreamExt, future, select, task::AtomicWaker};
24use quinn_proto::{
25    ClientConfig, ConnectError, ConnectionError, ConnectionHandle, DatagramEvent, EndpointConfig,
26    EndpointEvent, ServerConfig, Transmit, VarInt,
27};
28use rustc_hash::FxHashMap as HashMap;
29
30use crate::{
31    Connecting, ConnectionEvent, Incoming, RecvMeta, Socket,
32    sync::{mutex_blocking::Mutex, shared::Shared},
33};
34
35#[derive(Debug)]
36struct EndpointState {
37    endpoint: quinn_proto::Endpoint,
38    worker: Option<JoinHandle<()>>,
39    connections: HashMap<ConnectionHandle, Sender<ConnectionEvent>>,
40    close: Option<(VarInt, Bytes)>,
41    exit_on_idle: bool,
42    incoming: VecDeque<quinn_proto::Incoming>,
43    incoming_wakers: VecDeque<Waker>,
44    stats: EndpointStats,
45}
46
47/// Statistics on [Endpoint] activity
48#[non_exhaustive]
49#[derive(Debug, Default, Copy, Clone)]
50pub struct EndpointStats {
51    /// Cumulative number of Quic handshakes accepted by this [Endpoint]
52    pub accepted_handshakes: u64,
53    /// Cumulative number of Quic handshakes sent from this [Endpoint]
54    pub outgoing_handshakes: u64,
55    /// Cumulative number of Quic handshakes refused on this [Endpoint]
56    pub refused_handshakes: u64,
57    /// Cumulative number of Quic handshakes ignored on this [Endpoint]
58    pub ignored_handshakes: u64,
59}
60
61impl EndpointState {
62    fn handle_data(&mut self, meta: RecvMeta, buf: &[u8], respond_fn: impl Fn(Vec<u8>, Transmit)) {
63        // Empty UDP datagrams are not valid QUIC packets. Calling
64        // `chunks(0)` would panic with "chunk size must be non-zero".
65        if meta.len == 0 {
66            return;
67        }
68        let now = Instant::now();
69        for data in buf[..meta.len]
70            .chunks(meta.stride.min(meta.len))
71            .map(Into::into)
72        {
73            let mut resp_buf = Vec::new();
74            match self.endpoint.handle(
75                now,
76                meta.remote,
77                meta.local_ip,
78                meta.ecn,
79                data,
80                &mut resp_buf,
81            ) {
82                Some(DatagramEvent::NewConnection(incoming)) => {
83                    if self.close.is_none() {
84                        self.incoming.push_back(incoming);
85                    } else {
86                        let transmit = self.endpoint.refuse(incoming, &mut resp_buf);
87                        respond_fn(resp_buf, transmit);
88                    }
89                }
90                Some(DatagramEvent::ConnectionEvent(ch, event)) => {
91                    let _ = self
92                        .connections
93                        .get(&ch)
94                        .unwrap()
95                        .send(ConnectionEvent::Proto(event));
96                }
97                Some(DatagramEvent::Response(transmit)) => respond_fn(resp_buf, transmit),
98                None => {}
99            }
100        }
101    }
102
103    fn handle_event(&mut self, ch: ConnectionHandle, event: EndpointEvent) {
104        if event.is_drained() {
105            self.connections.remove(&ch);
106        }
107        if let Some(event) = self.endpoint.handle_event(ch, event) {
108            let _ = self
109                .connections
110                .get(&ch)
111                .unwrap()
112                .send(ConnectionEvent::Proto(event));
113        }
114    }
115
116    fn is_idle(&self) -> bool {
117        self.connections.is_empty()
118    }
119
120    fn poll_incoming(&mut self, cx: &mut Context) -> Poll<Option<quinn_proto::Incoming>> {
121        if self.close.is_none() {
122            if let Some(incoming) = self.incoming.pop_front() {
123                Poll::Ready(Some(incoming))
124            } else {
125                self.incoming_wakers.push_back(cx.waker().clone());
126                Poll::Pending
127            }
128        } else {
129            Poll::Ready(None)
130        }
131    }
132
133    #[track_caller]
134    fn new_connection(
135        &mut self,
136        handle: ConnectionHandle,
137        conn: quinn_proto::Connection,
138        socket: Socket,
139        events_tx: Sender<(ConnectionHandle, EndpointEvent)>,
140    ) -> Connecting {
141        let (tx, rx) = unbounded();
142        if let Some((error_code, reason)) = &self.close {
143            tx.send(ConnectionEvent::Close(*error_code, reason.clone()))
144                .unwrap();
145        }
146        self.connections.insert(handle, tx);
147        Connecting::new(handle, conn, socket, events_tx, rx)
148    }
149}
150
151impl Drop for EndpointState {
152    fn drop(&mut self) {
153        for incoming in self.incoming.drain(..) {
154            self.endpoint.ignore(incoming);
155        }
156    }
157}
158
159type ChannelPair<T> = (Sender<T>, Receiver<T>);
160
161#[derive(Debug)]
162pub(crate) struct EndpointInner {
163    state: Mutex<EndpointState>,
164    socket: Socket,
165    ipv6: bool,
166    events: ChannelPair<(ConnectionHandle, EndpointEvent)>,
167    done: AtomicWaker,
168}
169
170impl EndpointInner {
171    fn new(
172        socket: UdpSocket,
173        config: EndpointConfig,
174        server_config: Option<ServerConfig>,
175    ) -> io::Result<Self> {
176        let socket = Socket::new(socket)?;
177        let ipv6 = socket.local_addr()?.is_ipv6();
178        let allow_mtud = !socket.may_fragment();
179
180        Ok(Self {
181            state: Mutex::new(EndpointState {
182                endpoint: quinn_proto::Endpoint::new(
183                    Arc::new(config),
184                    server_config.map(Arc::new),
185                    allow_mtud,
186                    None,
187                ),
188                worker: None,
189                connections: HashMap::default(),
190                close: None,
191                exit_on_idle: false,
192                incoming: VecDeque::new(),
193                incoming_wakers: VecDeque::new(),
194                stats: EndpointStats::default(),
195            }),
196            socket,
197            ipv6,
198            events: unbounded(),
199            done: AtomicWaker::new(),
200        })
201    }
202
203    #[track_caller]
204    fn connect(
205        &self,
206        remote: SocketAddr,
207        server_name: &str,
208        config: ClientConfig,
209    ) -> Result<Connecting, ConnectError> {
210        let mut state = self.state.lock();
211
212        if state.worker.is_none() {
213            return Err(ConnectError::EndpointStopping);
214        }
215        if remote.is_ipv6() && !self.ipv6 {
216            return Err(ConnectError::InvalidRemoteAddress(remote));
217        }
218        let remote = if self.ipv6 {
219            SocketAddr::V6(match remote {
220                SocketAddr::V4(addr) => {
221                    SocketAddrV6::new(addr.ip().to_ipv6_mapped(), addr.port(), 0, 0)
222                }
223                SocketAddr::V6(addr) => addr,
224            })
225        } else {
226            remote
227        };
228
229        let (handle, conn) = state
230            .endpoint
231            .connect(Instant::now(), config, remote, server_name)?;
232        state.stats.outgoing_handshakes += 1;
233
234        Ok(state.new_connection(handle, conn, self.socket.clone(), self.events.0.clone()))
235    }
236
237    fn respond(&self, buf: Vec<u8>, transmit: Transmit) {
238        let socket = self.socket.clone();
239        // Name the task: the caller of this is compio rather than user code, so
240        // its location alone does not say what the task is.
241        let meta = SpawnMeta::capture().named("quic::respond");
242        compio_runtime::spawn_at(
243            async move {
244                socket.send(buf, &transmit).await;
245            },
246            meta,
247        )
248        .detach();
249    }
250
251    #[track_caller]
252    pub(crate) fn accept(
253        &self,
254        incoming: quinn_proto::Incoming,
255        server_config: Option<ServerConfig>,
256    ) -> Result<Connecting, ConnectionError> {
257        let mut state = self.state.lock();
258        let mut resp_buf = Vec::new();
259        let now = Instant::now();
260        match state
261            .endpoint
262            .accept(incoming, now, &mut resp_buf, server_config.map(Arc::new))
263        {
264            Ok((handle, conn)) => {
265                state.stats.accepted_handshakes += 1;
266                Ok(state.new_connection(handle, conn, self.socket.clone(), self.events.0.clone()))
267            }
268            Err(err) => {
269                if let Some(transmit) = err.response {
270                    self.respond(resp_buf, transmit);
271                }
272                Err(err.cause)
273            }
274        }
275    }
276
277    pub(crate) fn refuse(&self, incoming: quinn_proto::Incoming) {
278        let mut state = self.state.lock();
279        state.stats.refused_handshakes += 1;
280        let mut resp_buf = Vec::new();
281        let transmit = state.endpoint.refuse(incoming, &mut resp_buf);
282        self.respond(resp_buf, transmit);
283    }
284
285    #[allow(clippy::result_large_err)]
286    pub(crate) fn retry(
287        &self,
288        incoming: quinn_proto::Incoming,
289    ) -> Result<(), quinn_proto::RetryError> {
290        let mut state = self.state.lock();
291        let mut resp_buf = Vec::new();
292        let transmit = state.endpoint.retry(incoming, &mut resp_buf)?;
293        self.respond(resp_buf, transmit);
294        Ok(())
295    }
296
297    pub(crate) fn ignore(&self, incoming: quinn_proto::Incoming) {
298        let mut state = self.state.lock();
299        state.stats.ignored_handshakes += 1;
300        state.endpoint.ignore(incoming);
301    }
302
303    async fn run(&self) -> io::Result<()> {
304        let respond_fn = |buf: Vec<u8>, transmit: Transmit| self.respond(buf, transmit);
305
306        let mut recv_fut = pin!(
307            self.socket
308                .recv(Vec::with_capacity(
309                    self.state
310                        .lock()
311                        .endpoint
312                        .config()
313                        .get_max_udp_payload_size()
314                        .min(64 * 1024) as usize
315                        * self.socket.max_gro_segments(),
316                ))
317                .fuse()
318        );
319
320        let mut event_stream = self.events.1.stream().ready_chunks(100);
321
322        loop {
323            let mut state = select! {
324                BufResult(res, recv_buf) = recv_fut => {
325                    let mut state = self.state.lock();
326                    match res {
327                        Ok(meta) => state.handle_data(meta, &recv_buf, respond_fn),
328                        Err(e) if e.kind() == io::ErrorKind::ConnectionReset => {}
329                        #[cfg(windows)]
330                        Err(e) if e.raw_os_error() == Some(windows_sys::Win32::Foundation::ERROR_PORT_UNREACHABLE as _) => {}
331                        Err(e) => break Err(e),
332                    }
333                    recv_fut.set(self.socket.recv(recv_buf).fuse());
334                    state
335                },
336                events = event_stream.select_next_some() => {
337                    let mut state = self.state.lock();
338                    for (ch, event) in events {
339                        state.handle_event(ch, event);
340                    }
341                    state
342                },
343            };
344
345            if state.exit_on_idle && state.is_idle() {
346                break Ok(());
347            }
348            if !state.incoming.is_empty() {
349                let n = state.incoming.len().min(state.incoming_wakers.len());
350                state.incoming_wakers.drain(..n).for_each(Waker::wake);
351            }
352        }
353    }
354}
355
356#[derive(Debug, Clone)]
357pub(crate) struct EndpointRef(Shared<EndpointInner>);
358
359impl EndpointRef {
360    fn into_inner(self) -> Shared<EndpointInner> {
361        let this = ManuallyDrop::new(self);
362        // SAFETY: `this` is not dropped here, and we're consuming Self
363        unsafe { ptr::read(&this.0) }
364    }
365
366    async fn shutdown(self) -> io::Result<()> {
367        let (worker, idle) = {
368            let mut state = self.0.state.lock();
369            let idle = state.is_idle();
370            if !idle {
371                state.exit_on_idle = true;
372            }
373            (state.worker.take(), idle)
374        };
375        if let Some(worker) = worker {
376            if idle {
377                worker.cancel().await;
378            } else {
379                _ = worker.await;
380            }
381        }
382
383        let mut this = Some(self.into_inner());
384        let inner = poll_fn(move |cx| {
385            let s = match Shared::try_unwrap(this.take().unwrap()) {
386                Ok(inner) => return Poll::Ready(inner),
387                Err(s) => s,
388            };
389
390            s.done.register(cx.waker());
391
392            match Shared::try_unwrap(s) {
393                Ok(inner) => Poll::Ready(inner),
394                Err(s) => {
395                    this.replace(s);
396                    Poll::Pending
397                }
398            }
399        })
400        .await;
401
402        inner.socket.close().await
403    }
404}
405
406impl Drop for EndpointRef {
407    fn drop(&mut self) {
408        if Shared::strong_count(&self.0) == 2 {
409            // There are actually two cases:
410            // 1. User is trying to shutdown the socket.
411            self.0.done.wake();
412            // 2. User dropped the endpoint but the worker is still running.
413            self.0.state.lock().exit_on_idle = true;
414        }
415    }
416}
417
418impl Deref for EndpointRef {
419    type Target = EndpointInner;
420
421    fn deref(&self) -> &Self::Target {
422        &self.0
423    }
424}
425
426/// The metadata of an endpoint's worker task, captured on behalf of whoever
427/// creates the endpoint.
428///
429/// `#[track_caller]` carries that caller through the constructors, which are
430/// plain `fn`s returning a future for the purpose: it does not propagate
431/// through an `async fn`.
432#[track_caller]
433pub(crate) fn worker_meta() -> SpawnMeta {
434    SpawnMeta::capture().named("quic::endpoint")
435}
436
437/// A QUIC endpoint.
438#[derive(Debug, Clone)]
439pub struct Endpoint {
440    inner: EndpointRef,
441    /// The client configuration used by `connect`
442    pub default_client_config: Option<ClientConfig>,
443}
444
445impl Endpoint {
446    /// Create a QUIC endpoint.
447    #[track_caller]
448    pub fn new(
449        socket: UdpSocket,
450        config: EndpointConfig,
451        server_config: Option<ServerConfig>,
452        default_client_config: Option<ClientConfig>,
453    ) -> io::Result<Self> {
454        Self::new_at(
455            socket,
456            config,
457            server_config,
458            default_client_config,
459            worker_meta(),
460        )
461    }
462
463    pub(crate) fn new_at(
464        socket: UdpSocket,
465        config: EndpointConfig,
466        server_config: Option<ServerConfig>,
467        default_client_config: Option<ClientConfig>,
468        meta: SpawnMeta,
469    ) -> io::Result<Self> {
470        let inner = EndpointRef(Shared::new(EndpointInner::new(
471            socket,
472            config,
473            server_config,
474        )?));
475        // See the note in `respond` on why this task is named. Its metadata is
476        // taken rather than captured here, since the constructors that reach us
477        // return a future, and whoever called one is long gone by the time it
478        // is polled.
479        let worker = compio_runtime::spawn_at(
480            {
481                let inner = inner.clone();
482                async move {
483                    if let Err(e) = inner.run().await {
484                        error!("I/O error: {:?}", e);
485                    }
486                }
487                .in_current_span()
488            },
489            meta,
490        );
491        inner.state.lock().worker = Some(worker);
492        Ok(Self {
493            inner,
494            default_client_config,
495        })
496    }
497
498    /// Helper to construct an endpoint for use with outgoing connections only.
499    ///
500    /// Note that `addr` is the *local* address to bind to, which should usually
501    /// be a wildcard address like `0.0.0.0:0` or `[::]:0`, which allow
502    /// communication with any reachable IPv4 or IPv6 address respectively
503    /// from an OS-assigned port.
504    ///
505    /// If an IPv6 address is provided, the socket may dual-stack depending on
506    /// the platform, so as to allow communication with both IPv4 and IPv6
507    /// addresses. As such, calling this method with the address `[::]:0` is a
508    /// reasonable default to maximize the ability to connect to other
509    /// address.
510    ///
511    /// IPv4 client is never dual-stack.
512    #[cfg(rustls)]
513    #[track_caller]
514    pub fn client(addr: impl ToSocketAddrsAsync) -> impl Future<Output = io::Result<Endpoint>> {
515        Self::client_at(addr, worker_meta())
516    }
517
518    #[cfg(rustls)]
519    pub(crate) async fn client_at(
520        addr: impl ToSocketAddrsAsync,
521        meta: SpawnMeta,
522    ) -> io::Result<Endpoint> {
523        // TODO: try to enable dual-stack on all platforms, notably Windows
524        let socket = UdpSocket::bind(addr).await?;
525        Self::new_at(socket, EndpointConfig::default(), None, None, meta)
526    }
527
528    /// Helper to construct an endpoint for use with both incoming and outgoing
529    /// connections
530    ///
531    /// Platform defaults for dual-stack sockets vary. For example, any socket
532    /// bound to a wildcard IPv6 address on Windows will not by default be
533    /// able to communicate with IPv4 addresses. Portable applications
534    /// should bind an address that matches the family they wish to
535    /// communicate within.
536    #[cfg(rustls)]
537    #[track_caller]
538    pub fn server(
539        addr: impl ToSocketAddrsAsync,
540        config: ServerConfig,
541    ) -> impl Future<Output = io::Result<Self>> {
542        Self::server_at(addr, config, worker_meta())
543    }
544
545    #[cfg(rustls)]
546    pub(crate) async fn server_at(
547        addr: impl ToSocketAddrsAsync,
548        config: ServerConfig,
549        meta: SpawnMeta,
550    ) -> io::Result<Self> {
551        let socket = UdpSocket::bind(addr).await?;
552        Self::new_at(socket, EndpointConfig::default(), Some(config), None, meta)
553    }
554
555    /// Returns relevant stats from this Endpoint
556    pub fn stats(&self) -> EndpointStats {
557        self.inner.state.lock().stats
558    }
559
560    /// Connect to a remote endpoint.
561    #[track_caller]
562    pub fn connect(
563        &self,
564        remote: SocketAddr,
565        server_name: &str,
566        config: Option<ClientConfig>,
567    ) -> Result<Connecting, ConnectError> {
568        let config = config
569            .or_else(|| self.default_client_config.clone())
570            .ok_or(ConnectError::NoDefaultClientConfig)?;
571
572        self.inner.connect(remote, server_name, config)
573    }
574
575    /// Wait for the next incoming connection attempt from a client.
576    ///
577    /// Yields [`Incoming`]s, or `None` if the endpoint is
578    /// [`close`](Self::close)d. [`Incoming`] can be `await`ed to obtain the
579    /// final [`Connection`](crate::Connection), or used to e.g. filter
580    /// connection attempts or force address validation, or converted into an
581    /// intermediate `Connecting` future which can be used to e.g. send 0.5-RTT
582    /// data.
583    pub async fn wait_incoming(&self) -> Option<Incoming> {
584        future::poll_fn(|cx| self.inner.state.lock().poll_incoming(cx))
585            .await
586            .map(|incoming| Incoming::new(incoming, self.inner.clone()))
587    }
588
589    /// Replace the server configuration, affecting new incoming connections
590    /// only.
591    ///
592    /// Useful for e.g. refreshing TLS certificates without disrupting existing
593    /// connections.
594    pub fn set_server_config(&self, server_config: Option<ServerConfig>) {
595        self.inner
596            .state
597            .lock()
598            .endpoint
599            .set_server_config(server_config.map(Arc::new))
600    }
601
602    /// Get the local `SocketAddr` the underlying socket is bound to.
603    pub fn local_addr(&self) -> io::Result<SocketAddr> {
604        self.inner.socket.local_addr()
605    }
606
607    /// Get the number of connections that are currently open.
608    pub fn open_connections(&self) -> usize {
609        self.inner.state.lock().endpoint.open_connections()
610    }
611
612    /// Close all of this endpoint's connections immediately and cease accepting
613    /// new connections.
614    ///
615    /// See [`Connection::close()`] for details.
616    ///
617    /// [`Connection::close()`]: crate::Connection::close
618    pub fn close(&self, error_code: VarInt, reason: &[u8]) {
619        let reason = Bytes::copy_from_slice(reason);
620        let mut state = self.inner.state.lock();
621        if state.close.is_some() {
622            return;
623        }
624        state.close = Some((error_code, reason.clone()));
625        for conn in state.connections.values() {
626            let _ = conn.send(ConnectionEvent::Close(error_code, reason.clone()));
627        }
628        state.incoming_wakers.drain(..).for_each(Waker::wake);
629    }
630
631    /// Gracefully shutdown the endpoint.
632    ///
633    /// Wait for all connections on the endpoint to be cleanly shut down and
634    /// close the underlying socket. This will wait for all clones of the
635    /// endpoint, all connections and all streams to be dropped before
636    /// closing the socket.
637    ///
638    /// Waiting for this condition before exiting ensures that a good-faith
639    /// effort is made to notify peers of recent connection closes, whereas
640    /// exiting immediately could force them to wait out the idle timeout
641    /// period.
642    ///
643    /// Does not proactively close existing connections. Consider calling
644    /// [`close()`] if that is desired.
645    ///
646    /// [`close()`]: Endpoint::close
647    pub async fn shutdown(self) -> io::Result<()> {
648        self.inner.shutdown().await
649    }
650}
compio_quic/endpoint.rs

compio_quic/
endpoint.rs
