Skip to main content

tokio_quiche/quic/router/
connector.rs

1// Copyright (C) 2025, Cloudflare, Inc.
2// All rights reserved.
3//
4// Redistribution and use in source and binary forms, with or without
5// modification, are permitted provided that the following conditions are
6// met:
7//
8//     * Redistributions of source code must retain the above copyright notice,
9//       this list of conditions and the following disclaimer.
10//
11//     * Redistributions in binary form must reproduce the above copyright
12//       notice, this list of conditions and the following disclaimer in the
13//       documentation and/or other materials provided with the distribution.
14//
15// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
16// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
17// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
19// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
23// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
24// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
25// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26
27use std::io;
28use std::mem;
29use std::sync::Arc;
30use std::task::Context;
31use std::task::Poll;
32use std::time::Instant;
33
34use datagram_socket::DatagramSocketSend;
35use datagram_socket::DatagramSocketSendExt;
36use datagram_socket::MaybeConnectedSocket;
37use datagram_socket::MAX_DATAGRAM_SIZE;
38use foundations::telemetry::log;
39use quiche::ConnectionId;
40use quiche::Header;
41use tokio_util::time::delay_queue::Key;
42use tokio_util::time::DelayQueue;
43
44use crate::quic::router::InitialPacketHandler;
45use crate::quic::router::NewConnection;
46use crate::quic::Incoming;
47use crate::quic::QuicheConnection;
48
49/// A [`ClientConnector`] manages client-initiated [`quiche::Connection`]s. When
50/// a connection is established, this struct returns the connection to the
51/// [`InboundPacketRouter`](super::InboundPacketRouter) for further processing.
52pub(crate) struct ClientConnector<Tx> {
53    socket_tx: MaybeConnectedSocket<Arc<Tx>>,
54    connection: ConnectionState,
55    timeout_queue: DelayQueue<ConnectionId<'static>>,
56}
57
58/// State the connecting connection is in.
59enum ConnectionState {
60    /// Connection hasn't had any initials sent for it
61    Queued(Box<QuicheConnection>),
62    /// It's currently in a QUIC handshake
63    Pending(PendingConnection),
64    /// It's been returned to the
65    /// [`InboundPacketRouter`](super::InboundPacketRouter).
66    Returned,
67}
68
69impl ConnectionState {
70    fn take_if_queued(&mut self) -> Option<Box<QuicheConnection>> {
71        match mem::replace(self, Self::Returned) {
72            Self::Queued(conn) => Some(conn),
73            state => {
74                *self = state;
75                None
76            },
77        }
78    }
79
80    fn take_if_pending_and_id_matches(
81        &mut self, scid: &ConnectionId<'static>,
82    ) -> Option<PendingConnection> {
83        match mem::replace(self, Self::Returned) {
84            Self::Pending(pending) if *scid == pending.conn.source_id() =>
85                Some(pending),
86            state => {
87                *self = state;
88                None
89            },
90        }
91    }
92}
93
94/// A [`PendingConnection`] holds an internal [`quiche::Connection`] and an
95/// optional timeout [`Key`].
96struct PendingConnection {
97    conn: Box<QuicheConnection>,
98    timeout_key: Option<Key>,
99    handshake_start_time: Instant,
100}
101
102impl<Tx> ClientConnector<Tx>
103where
104    Tx: DatagramSocketSend + Send + 'static,
105{
106    pub(crate) fn new(
107        socket_tx: Arc<Tx>, connection: Box<QuicheConnection>,
108    ) -> Self {
109        Self {
110            socket_tx: MaybeConnectedSocket::new(socket_tx),
111            connection: ConnectionState::Queued(connection),
112            timeout_queue: Default::default(),
113        }
114    }
115
116    /// Sets the connection to it's pending state. Await [`Incoming`] packets.
117    ///
118    /// This sends any pending packets and arms the connection's timeout timer.
119    fn set_connection_to_pending(
120        &mut self, mut conn: Box<QuicheConnection>,
121    ) -> io::Result<()> {
122        simple_conn_send(&self.socket_tx, &mut conn)?;
123
124        let timeout_key = conn.timeout_instant().map(|instant| {
125            self.timeout_queue
126                .insert_at(conn.source_id().into_owned(), instant.into())
127        });
128
129        self.connection = ConnectionState::Pending(PendingConnection {
130            conn,
131            timeout_key,
132            handshake_start_time: Instant::now(),
133        });
134
135        Ok(())
136    }
137
138    /// Handles an incoming packet (or packets) designated for this pending
139    /// connection.
140    ///
141    /// If the connection is pending, we return it
142    fn on_incoming(
143        &mut self, mut incoming: Incoming, hdr: Header<'static>,
144    ) -> io::Result<Option<NewConnection>> {
145        let Some(PendingConnection {
146            mut conn,
147            timeout_key,
148            handshake_start_time,
149        }) = self.connection.take_if_pending_and_id_matches(&hdr.dcid)
150        else {
151            log::debug!("Received Initial packet for unknown connection ID"; "scid" => ?hdr.dcid);
152            return Ok(None);
153        };
154
155        let recv_info = quiche::RecvInfo {
156            from: incoming.peer_addr,
157            to: incoming.local_addr,
158        };
159
160        if let Some(gro) = incoming.gro {
161            for dgram in incoming.buf.chunks_mut(gro as usize) {
162                // Log error here if recv fails
163                let _ = conn.recv(dgram, recv_info);
164            }
165        } else {
166            // Log error here if recv fails
167            let _ = conn.recv(&mut incoming.buf, recv_info);
168        }
169
170        // disarm the timer since we're either going to immediately rearm it or
171        // return an established connection.
172        if let Some(key) = timeout_key {
173            self.timeout_queue.remove(&key);
174        }
175
176        let scid = conn.source_id();
177        if conn.is_established() {
178            log::debug!("QUIC connection established"; "scid" => ?scid);
179
180            Ok(Some(NewConnection {
181                conn,
182                pending_cid: None,
183                initial_pkt: None,
184                cid_generator: None,
185                handshake_start_time,
186            }))
187        } else if conn.is_closed() {
188            let scid = conn.source_id();
189            log::error!("QUIC connection closed on_incoming"; "scid" => ?scid);
190
191            Err(io::Error::new(
192                io::ErrorKind::TimedOut,
193                format!("connection {scid:?} timed out"),
194            ))
195        } else {
196            self.set_connection_to_pending(conn).map(|()| None)
197        }
198    }
199
200    /// [`ClientConnector::on_timeout`] runs [`quiche::Connection::on_timeout`]
201    /// for a pending connection. If the connection is closed, this sends an
202    /// error upstream.
203    fn on_timeout(&mut self, scid: ConnectionId<'static>) -> io::Result<()> {
204        log::debug!("connection timedout"; "scid" => ?scid);
205
206        let Some(mut pending) =
207            self.connection.take_if_pending_and_id_matches(&scid)
208        else {
209            log::debug!("timedout connection missing from pending map"; "scid" => ?scid);
210            return Ok(());
211        };
212
213        pending.conn.on_timeout();
214
215        if pending.conn.is_closed() {
216            log::error!("pending connection closed on_timeout"; "scid" => ?scid);
217
218            return Err(io::Error::new(
219                io::ErrorKind::TimedOut,
220                format!("connection {scid:?} timed out"),
221            ));
222        }
223
224        self.set_connection_to_pending(pending.conn)
225    }
226
227    /// [`ClientConnector::update`] handles expired pending connections and
228    /// checks starts the inner connection if not started yet.
229    fn update(&mut self, cx: &mut Context) -> io::Result<()> {
230        while let Poll::Ready(Some(expired)) = self.timeout_queue.poll_expired(cx)
231        {
232            let scid = expired.into_inner();
233            self.on_timeout(scid)?;
234        }
235
236        if let Some(conn) = self.connection.take_if_queued() {
237            self.set_connection_to_pending(conn)?;
238        }
239
240        Ok(())
241    }
242}
243
244impl<Tx> InitialPacketHandler for ClientConnector<Tx>
245where
246    Tx: DatagramSocketSend + Send + 'static,
247{
248    fn update(&mut self, ctx: &mut Context<'_>) -> io::Result<()> {
249        ClientConnector::update(self, ctx)
250    }
251
252    fn handle_initials(
253        &mut self, incoming: Incoming, hdr: Header<'static>,
254        _: &mut quiche::Config,
255    ) -> io::Result<Option<NewConnection>> {
256        self.on_incoming(incoming, hdr)
257    }
258}
259
260/// Repeatedly send packets until quiche reports that it's done.
261///
262/// This does not have to be efficent, since once a connection is established
263/// the [`crate::quic::io::worker::IoWorker`] will take over sending and
264/// receiving.
265fn simple_conn_send<Tx: DatagramSocketSend + Send + Sync + 'static>(
266    socket_tx: &MaybeConnectedSocket<Arc<Tx>>, conn: &mut QuicheConnection,
267) -> io::Result<()> {
268    let scid = conn.source_id().into_owned();
269    log::debug!("sending client Initials to peer"; "scid" => ?scid);
270
271    loop {
272        let scid = scid.clone();
273        let mut buf = [0; MAX_DATAGRAM_SIZE];
274        let send_res = conn.send(&mut buf);
275
276        let socket_clone = socket_tx.clone();
277        match send_res {
278            Ok((n, send_info)) => {
279                tokio::spawn({
280                    let buf = buf[0..n].to_vec();
281                    async move {
282                        socket_clone.send_to(&buf, send_info.to).await.inspect_err(|error| {
283                        log::error!("error sending client Initial packets to peer"; "scid" => ?scid, "peer_addr" => send_info.to, "error" => error.to_string());
284                    })
285                    }
286                });
287            },
288            Err(quiche::Error::Done) => break Ok(()),
289            Err(error) => {
290                log::error!("error writing packets to quiche's internal buffer"; "scid" => ?scid, "error" => error.to_string());
291                break Err(std::io::Error::other(error));
292            },
293        }
294    }
295}