tokio-quiche: Enable active connection migration for clients of tq's QUIC server

tokio-quiche/examples/async_http3_server/main.rs

@@ -37,7 +37,6 @@ use tokio::net::UdpSocket;
 use tokio_quiche::http3::settings::Http3Settings;
 use tokio_quiche::listen;
 use tokio_quiche::metrics::DefaultMetrics;
-use tokio_quiche::quic::SimpleConnectionIdGenerator;
 use tokio_quiche::settings::CertificateKind::{
     self,
 };
@@ -68,7 +67,6 @@ async fn main() {
             },
             Hooks::default(),
         ),
-        SimpleConnectionIdGenerator,
         DefaultMetrics,
     )
     .expect("should be able to create a listener from a UDP socket");

tokio-quiche/src/lib.rs

@@ -47,12 +47,8 @@
 //!
 //! # async fn example() -> tokio_quiche::QuicResult<()> {
 //! let socket = tokio::net::UdpSocket::bind("0.0.0.0:443").await?;
-//! let mut listeners = listen(
-//!     [socket],
-//!     ConnectionParams::default(),
-//!     SimpleConnectionIdGenerator,
-//!     DefaultMetrics,
-//! )?;
+//! let mut listeners =
+//!     listen([socket], ConnectionParams::default(), DefaultMetrics)?;
 //! let mut accept_stream = &mut listeners[0];
 //!
 //! while let Some(conn) = accept_stream.next().await {
@@ -154,7 +150,7 @@ pub type QuicConnectionStream<M> =
 /// previously-yielded connections are closed.
 pub fn listen_with_capabilities<M>(
     sockets: impl IntoIterator<Item = QuicListener>, params: ConnectionParams,
-    cid_generator: impl ConnectionIdGenerator<'static> + Clone, metrics: M,
+    metrics: M,
 ) -> io::Result<Vec<QuicConnectionStream<M>>>
 where
     M: Metrics,
@@ -165,14 +161,7 @@ where
 
     sockets
         .into_iter()
-        .map(|s| {
-            crate::quic::start_listener(
-                s,
-                &params,
-                cid_generator.clone(),
-                metrics.clone(),
-            )
-        })
+        .map(|s| crate::quic::start_listener(s, &params, metrics.clone()))
         .collect()
 }
 
@@ -181,8 +170,7 @@ where
 /// Each socket is converted into a [`QuicListener`] with defaulted socket
 /// parameters. The listeners are then passed to [`listen_with_capabilities`].
 pub fn listen<S, M>(
-    sockets: impl IntoIterator<Item = S>, params: ConnectionParams,
-    cid_generator: impl ConnectionIdGenerator<'static> + Clone, metrics: M,
+    sockets: impl IntoIterator<Item = S>, params: ConnectionParams, metrics: M,
 ) -> io::Result<Vec<QuicConnectionStream<M>>>
 where
     S: TryInto<QuicListener, Error = io::Error>,
@@ -199,7 +187,7 @@ where
         })
         .collect::<io::Result<_>>()?;
 
-    listen_with_capabilities(quic_sockets, params, cid_generator, metrics)
+    listen_with_capabilities(quic_sockets, params, metrics)
 }
 
 static GLOBAL_LOGGER_ONCE: Once = Once::new();

tokio-quiche/src/quic/connection/id.rs

@@ -25,6 +25,7 @@
 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 use quiche::ConnectionId;
+use std::sync::Arc;
 
 use crate::QuicResult;
 
@@ -37,24 +38,16 @@ use crate::QuicResult;
 ///
 /// Clients currently can't configure a [`ConnectionIdGenerator`] and always use
 /// the [`SimpleConnectionIdGenerator`].
-pub trait ConnectionIdGenerator<'a>: Send + 'static {
+pub trait ConnectionIdGenerator<'a>: Send + Sync + 'static {
     /// Creates a new [`ConnectionId`] according to the generator's logic.
-    ///
-    /// The `socket_cookie` is propagated unchanged from the
-    /// [`QuicListener`](crate::socket::QuicListener) that received the
-    /// connection.
-    fn new_connection_id(&self, socket_cookie: u64) -> ConnectionId<'a>;
+    fn new_connection_id(&self) -> ConnectionId<'a>;
 
     /// Verifies whether `cid` was generated by this [`ConnectionIdGenerator`].
-    ///
-    /// `socket_cookie` matches the value that would have been passed to
-    /// `new_connection_id` calls. The method should return an [`Err`] variant
-    /// if the connection ID can't be verified.
-    fn verify_connection_id(
-        &self, socket_cookie: u64, cid: &ConnectionId,
-    ) -> QuicResult<()>;
+    fn verify_connection_id(&self, cid: &ConnectionId) -> QuicResult<()>;
 }
 
+pub type SharedConnectionIdGenerator = Arc<dyn ConnectionIdGenerator<'static>>;
+
 /// A [`ConnectionIdGenerator`] which creates random 20-byte connection IDs.
 ///
 /// Random bytes are pulled directly from the operating system to create an ID.
@@ -63,7 +56,7 @@ pub trait ConnectionIdGenerator<'a>: Send + 'static {
 pub struct SimpleConnectionIdGenerator;
 
 impl ConnectionIdGenerator<'static> for SimpleConnectionIdGenerator {
-    fn new_connection_id(&self, _socket_cookie: u64) -> ConnectionId<'static> {
+    fn new_connection_id(&self) -> ConnectionId<'static> {
         let mut buf = vec![0; 20];
         boring::rand::rand_bytes(&mut buf).unwrap();
 
@@ -72,9 +65,7 @@ impl ConnectionIdGenerator<'static> for SimpleConnectionIdGenerator {
 
     /// Performs no verification, because this generator can create
     /// any valid connection ID.
-    fn verify_connection_id(
-        &self, _socket_cookie: u64, _cid: &ConnectionId<'_>,
-    ) -> QuicResult<()> {
+    fn verify_connection_id(&self, _cid: &ConnectionId<'_>) -> QuicResult<()> {
         Ok(())
     }
 }

tokio-quiche/src/quic/connection/map.rs

@@ -32,7 +32,6 @@ use datagram_socket::DatagramSocketSend;
 use quiche::ConnectionId;
 use quiche::MAX_CONN_ID_LEN;
 use std::collections::BTreeMap;
-use std::collections::HashMap;
 use tokio::sync::mpsc;
 
 const U64_SZ: usize = std::mem::size_of::<u64>();
@@ -96,51 +95,32 @@ impl From<&ConnectionId<'_>> for CidOwned {
     }
 }
 
-/// A unique identifier quiche assigns to a connection.
-type QuicheId = u64;
-
 /// A map for QUIC connections.
 ///
 /// Due to the fact that QUIC connections can be identified by multiple QUIC
 /// connection IDs, we have to be able to map multiple IDs to the same
 /// connection.
 #[derive(Default)]
 pub(crate) struct ConnectionMap {
-    quic_id_map: BTreeMap<CidOwned, (QuicheId, mpsc::Sender<Incoming>)>,
-    conn_map: HashMap<QuicheId, mpsc::Sender<Incoming>>,
+    quic_id_map: BTreeMap<CidOwned, mpsc::Sender<Incoming>>,
 }
 
 impl ConnectionMap {
     pub(crate) fn insert<Tx, M>(
-        &mut self, cid: ConnectionId<'_>, conn: &InitialQuicConnection<Tx, M>,
+        &mut self, cid: &ConnectionId<'_>, conn: &InitialQuicConnection<Tx, M>,
     ) where
         Tx: DatagramSocketSend + Send + 'static,
         M: Metrics,
     {
-        let id = conn.id;
         let ev_sender = conn.incoming_ev_sender.clone();
-
-        self.conn_map.insert(id, ev_sender.clone());
-        self.quic_id_map.insert((&cid).into(), (id, ev_sender));
-    }
-
-    pub(crate) fn remove(&mut self, cid: &ConnectionId<'_>) {
-        if let Some((id, _)) = self.quic_id_map.remove(&cid.into()) {
-            self.conn_map.remove(&id);
-        }
+        self.quic_id_map.insert(cid.into(), ev_sender);
     }
 
-    pub(crate) fn map_cid<Tx, M>(
-        &mut self, cid: ConnectionId<'_>, conn: &InitialQuicConnection<Tx, M>,
-    ) where
-        Tx: DatagramSocketSend + Send + 'static,
-        M: Metrics,
-    {
-        let id = conn.id;
-
-        if let Some(ev_sender) = self.conn_map.get(&id) {
-            self.quic_id_map
-                .insert((&cid).into(), (id, ev_sender.clone()));
+    pub(crate) fn map_cid(
+        &mut self, existing_cid: &ConnectionId<'_>, new_cid: &ConnectionId<'_>,
+    ) {
+        if let Some(ev_sender) = self.quic_id_map.get(&existing_cid.into()) {
+            self.quic_id_map.insert(new_cid.into(), ev_sender.clone());
         }
     }
 
@@ -155,9 +135,9 @@ impl ConnectionMap {
             // Although both branches run the same code, the one here will
             // generate an optimized version for the length we are
             // using, as opposed to temporary cids sent by clients.
-            self.quic_id_map.get(&id.into()).map(|(_id, sender)| sender)
+            self.quic_id_map.get(&id.into())
         } else {
-            self.quic_id_map.get(&id.into()).map(|(_id, sender)| sender)
+            self.quic_id_map.get(&id.into())
         }
     }
 }

tokio-quiche/src/quic/connection/mod.rs

@@ -30,6 +30,7 @@ mod map;
 
 pub use self::error::HandshakeError;
 pub use self::id::ConnectionIdGenerator;
+pub use self::id::SharedConnectionIdGenerator;
 pub use self::id::SimpleConnectionIdGenerator;
 pub(crate) use self::map::ConnectionMap;
 
@@ -231,9 +232,6 @@ where
     Tx: DatagramSocketSend + Send + 'static + ?Sized,
     M: Metrics,
 {
-    /// An internal ID, to uniquely identify the connection across multiple QUIC
-    /// connection IDs.
-    pub(crate) id: u64,
     params: QuicConnectionParams<Tx, M>,
     pub(crate) audit_log_stats: Arc<QuicAuditStats>,
     stats: QuicConnectionStatsShared,
@@ -257,7 +255,6 @@ where
         )));
 
         Self {
-            id: Self::generate_id(),
             params,
             audit_log_stats,
             stats,
@@ -341,6 +338,7 @@ where
             audit_log_stats: self.audit_log_stats,
             write_state: WriteState::default(),
             conn_map_cmd_tx: self.params.conn_map_cmd_tx,
+            cid_generator: self.params.cid_generator,
             #[cfg(feature = "perf-quic-listener-metrics")]
             init_rx_time: self.params.init_rx_time,
             metrics: self.params.metrics.clone(),
@@ -457,14 +455,6 @@ where
 
         conn
     }
-
-    fn generate_id() -> u64 {
-        let mut buf = [0; 8];
-
-        boring::rand::rand_bytes(&mut buf).unwrap();
-
-        u64::from_ne_bytes(buf)
-    }
 }
 
 pub(crate) struct QuicConnectionParams<Tx, M>
@@ -477,6 +467,7 @@ where
     pub shutdown_tx: mpsc::Sender<()>,
     pub conn_map_cmd_tx: mpsc::UnboundedSender<ConnectionMapCommand>, /* channel that signals connection map changes */
     pub scid: ConnectionId<'static>,
+    pub cid_generator: Option<SharedConnectionIdGenerator>,
     pub metrics: M,
     #[cfg(feature = "perf-quic-listener-metrics")]
     pub init_rx_time: Option<SystemTime>,