tokio_quiche/quic/connection/
map.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 super::Incoming;
28use super::InitialQuicConnection;
29use crate::metrics::Metrics;
30
31use datagram_socket::DatagramSocketSend;
32use quiche::ConnectionId;
33use quiche::MAX_CONN_ID_LEN;
34use std::collections::BTreeMap;
35use std::collections::HashMap;
36use tokio::sync::mpsc;
37
38const U64_SZ: usize = std::mem::size_of::<u64>();
39const MAX_CONN_ID_QUADS: usize = MAX_CONN_ID_LEN.div_ceil(U64_SZ);
40const CONN_ID_USABLE_LEN: usize = min_usize(
41    // Last byte in CidOwned::Optimized stores CID length
42    MAX_CONN_ID_QUADS * U64_SZ - 1,
43    // CID length must fit in 1 byte
44    min_usize(MAX_CONN_ID_LEN, u8::MAX as _),
45);
46
47const fn min_usize(v1: usize, v2: usize) -> usize {
48    if v1 < v2 {
49        v1
50    } else {
51        v2
52    }
53}
54
55/// A non unique connection identifier, multiple Cids can map to the same
56/// conenction.
57#[derive(PartialEq, Eq, PartialOrd, Ord)]
58enum CidOwned {
59    /// The QUIC connections IDs theoretically have unbounded length, so for the
60    /// generic case a boxed slice is used to store the ID.
61    Generic(Box<[u8]>),
62    /// For QUIC version 1 (the one that actually exists) the maximal ID size is
63    /// `20`, which should correspond to the `MAX_CONN_ID_LEN` value. For
64    /// this common case, we store the ID in a u64 array for faster
65    /// comparison (and therefore BTreeMap lookups).
66    Optimized([u64; MAX_CONN_ID_QUADS]),
67}
68
69impl From<&ConnectionId<'_>> for CidOwned {
70    #[inline(always)]
71    fn from(value: &ConnectionId<'_>) -> Self {
72        if value.len() > CONN_ID_USABLE_LEN {
73            return CidOwned::Generic(value.as_ref().into());
74        }
75
76        let mut cid = [0; MAX_CONN_ID_QUADS];
77
78        value
79            .chunks(U64_SZ)
80            .map(|c| match c.try_into() {
81                Ok(v) => u64::from_le_bytes(v),
82                Err(_) => {
83                    let mut remainder = [0u8; U64_SZ];
84                    remainder[..c.len()].copy_from_slice(c);
85                    u64::from_le_bytes(remainder)
86                },
87            })
88            .enumerate()
89            .for_each(|(i, v)| cid[i] = v);
90
91        // In order to differentiate cids with zeroes as opposed to shorter cids,
92        // append the cid length.
93        *cid.last_mut().unwrap() |= (value.len() as u64) << 56;
94
95        CidOwned::Optimized(cid)
96    }
97}
98
99/// A unique identifier quiche assigns to a connection.
100type QuicheId = u64;
101
102/// A map for QUIC connections.
103///
104/// Due to the fact that QUIC connections can be identified by multiple QUIC
105/// connection IDs, we have to be able to map multiple IDs to the same
106/// connection.
107#[derive(Default)]
108pub(crate) struct ConnectionMap {
109    quic_id_map: BTreeMap<CidOwned, (QuicheId, mpsc::Sender<Incoming>)>,
110    conn_map: HashMap<QuicheId, mpsc::Sender<Incoming>>,
111}
112
113impl ConnectionMap {
114    pub(crate) fn insert<Tx, M>(
115        &mut self, cid: ConnectionId<'_>, conn: &InitialQuicConnection<Tx, M>,
116    ) where
117        Tx: DatagramSocketSend + Send + 'static,
118        M: Metrics,
119    {
120        let id = conn.id;
121        let ev_sender = conn.incoming_ev_sender.clone();
122
123        self.conn_map.insert(id, ev_sender.clone());
124        self.quic_id_map.insert((&cid).into(), (id, ev_sender));
125    }
126
127    pub(crate) fn remove(&mut self, cid: &ConnectionId<'_>) {
128        if let Some((id, _)) = self.quic_id_map.remove(&cid.into()) {
129            self.conn_map.remove(&id);
130        }
131    }
132
133    pub(crate) fn map_cid<Tx, M>(
134        &mut self, cid: ConnectionId<'_>, conn: &InitialQuicConnection<Tx, M>,
135    ) where
136        Tx: DatagramSocketSend + Send + 'static,
137        M: Metrics,
138    {
139        let id = conn.id;
140
141        if let Some(ev_sender) = self.conn_map.get(&id) {
142            self.quic_id_map
143                .insert((&cid).into(), (id, ev_sender.clone()));
144        }
145    }
146
147    pub(crate) fn unmap_cid(&mut self, cid: &ConnectionId<'_>) {
148        self.quic_id_map.remove(&cid.into());
149    }
150
151    pub(crate) fn get(
152        &self, id: &ConnectionId,
153    ) -> Option<&mpsc::Sender<Incoming>> {
154        if id.len() == MAX_CONN_ID_LEN {
155            // Although both branches run the same code, the one here will
156            // generate an optimized version for the length we are
157            // using, as opposed to temporary cids sent by clients.
158            self.quic_id_map.get(&id.into()).map(|(_id, sender)| sender)
159        } else {
160            self.quic_id_map.get(&id.into()).map(|(_id, sender)| sender)
161        }
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168    use quiche::ConnectionId;
169
170    #[test]
171    fn cid_storage() {
172        let max_v1_cid = ConnectionId::from_ref(&[0xfa; MAX_CONN_ID_LEN]);
173        let optimized = CidOwned::from(&max_v1_cid);
174        assert!(
175            matches!(optimized, CidOwned::Optimized(_)),
176            "QUIC v1 CID is not stored inline"
177        );
178
179        let oversize_cid = ConnectionId::from_ref(&[0x1b; MAX_CONN_ID_LEN + 20]);
180        let boxed = CidOwned::from(&oversize_cid);
181        assert!(
182            matches!(boxed, CidOwned::Generic(_)),
183            "Oversized CID is not boxed"
184        );
185    }
186}
tokio_quiche/quic/connection/map.rs

tokio_quiche/quic/connection/
map.rs
