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

tokio_quiche/quic/connection/
map.rs
