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

compio_quic/
endpoint.rs
