simoom_quic/
endpoint.rs

1//! An endpoint: routes datagrams to connections by connection ID, creates server
2//! connections, and answers statelessly (Version Negotiation, Retry, stateless reset).
3//! It also keeps one deadline heap for all its connections.
4
5use crate::config::Config;
6use crate::connection::{Connection, EndpointEvent, Event, Transmit};
7use crate::error::TransportError;
8use crate::packet::{self, PacketType, MIN_INITIAL_SIZE, MIN_RESET_SIZE};
9use crate::params::Cid;
10use crate::token::{self, Validated};
11use std::cmp::Reverse;
12use std::collections::{BinaryHeap, HashMap, VecDeque};
13use std::net::SocketAddr;
14use std::sync::Arc;
15use std::time::Instant;
16
17/// Identifies a connection inside an endpoint.
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
19pub struct ConnHandle(pub usize);
20
21/// What happened to a received datagram.
22#[derive(Debug, PartialEq, Eq)]
23pub enum DatagramEvent {
24    /// Delivered to an existing connection.
25    Connection(ConnHandle),
26    /// A new server connection was created and the datagram delivered to it.
27    NewConnection(ConnHandle),
28    /// A stateless reply was written to the response buffer.
29    Response(Transmit),
30    Dropped,
31}
32
33struct Slot {
34    conn: Connection,
35    gen: u64,
36    deadline: Option<Instant>,
37    dirty: bool,
38    handshaking: bool,
39    cids: Vec<Cid>,
40    tokens: Vec<[u8; 16]>,
41}
42
43pub struct Endpoint {
44    cfg: Arc<Config>,
45    server: bool,
46    slots: Vec<Option<Slot>>,
47    free: Vec<usize>,
48    by_cid: HashMap<Cid, usize>,
49    by_token: HashMap<[u8; 16], usize>,
50    timers: BinaryHeap<Reverse<(Instant, usize, u64)>>,
51    dirty: VecDeque<usize>,
52    next_gen: u64,
53    handshaking: usize,
54    live: usize,
55}
56
57impl Endpoint {
58    /// `server` enables accepting connections.
59    pub fn new(cfg: Arc<Config>, server: bool) -> Self {
60        assert!((4..=20).contains(&cfg.cid_len), "cid_len must be 4..=20");
61        Endpoint {
62            cfg,
63            server,
64            slots: Vec::new(),
65            free: Vec::new(),
66            by_cid: HashMap::new(),
67            by_token: HashMap::new(),
68            timers: BinaryHeap::new(),
69            dirty: VecDeque::new(),
70            next_gen: 0,
71            handshaking: 0,
72            live: 0,
73        }
74    }
75
76    pub fn config(&self) -> &Arc<Config> {
77        &self.cfg
78    }
79
80    /// Live connections.
81    pub fn len(&self) -> usize {
82        self.live
83    }
84    pub fn is_empty(&self) -> bool {
85        self.live == 0
86    }
87
88    fn insert(&mut self, conn: Connection) -> usize {
89        self.next_gen += 1;
90        let slot = Slot { handshaking: conn.is_handshaking(), conn, gen: self.next_gen, deadline: None, dirty: false, cids: Vec::new(), tokens: Vec::new() };
91        if slot.handshaking {
92            self.handshaking += 1;
93        }
94        self.live += 1;
95        if let Some(i) = self.free.pop() {
96            self.slots[i] = Some(slot);
97            i
98        } else {
99            self.slots.push(Some(slot));
100            self.slots.len() - 1
101        }
102    }
103
104    /// Starts a client connection.
105    pub fn connect(
106        &mut self,
107        now: Instant,
108        remote: SocketAddr,
109        server_name: &str,
110        token: &[u8],
111    ) -> Result<ConnHandle, TransportError> {
112        self.cfg.epoch.get_or_init(|| now);
113        let conn = Connection::connect(self.cfg.clone(), now, remote, server_name, token)?;
114        let i = self.insert(conn);
115        self.process(i);
116        Ok(ConnHandle(i))
117    }
118
119    /// Mutable access to a connection; it is scheduled for transmission.
120    pub fn conn_mut(&mut self, h: ConnHandle) -> Option<&mut Connection> {
121        let slot = self.slots.get_mut(h.0)?.as_mut()?;
122        if !slot.dirty {
123            slot.dirty = true;
124            self.dirty.push_back(h.0);
125        }
126        Some(&mut slot.conn)
127    }
128
129    pub fn conn(&self, h: ConnHandle) -> Option<&Connection> {
130        self.slots.get(h.0)?.as_ref().map(|s| &s.conn)
131    }
132
133    /// The next event of a connection. A drained connection is removed once its events
134    /// have been read.
135    pub fn poll_event(&mut self, h: ConnHandle) -> Option<Event> {
136        let slot = self.slots.get_mut(h.0)?.as_mut()?;
137        let e = slot.conn.poll_event();
138        if e.is_none() && slot.conn.is_drained() {
139            self.remove(h.0);
140        }
141        e
142    }
143
144    /// Handles of all live connections.
145    pub fn handles(&self) -> impl Iterator<Item = ConnHandle> + '_ {
146        self.slots.iter().enumerate().filter(|(_, s)| s.is_some()).map(|(i, _)| ConnHandle(i))
147    }
148
149    fn remove(&mut self, i: usize) {
150        if let Some(slot) = self.slots[i].take() {
151            for c in &slot.cids {
152                if self.by_cid.get(c) == Some(&i) {
153                    self.by_cid.remove(c);
154                }
155            }
156            for t in &slot.tokens {
157                if self.by_token.get(t) == Some(&i) {
158                    self.by_token.remove(t);
159                }
160            }
161            if slot.handshaking {
162                self.handshaking -= 1;
163            }
164            self.live -= 1;
165            self.free.push(i);
166        }
167    }
168
169    /// Applies routing changes, updates the deadline and schedules transmission.
170    fn process(&mut self, i: usize) {
171        let Some(slot) = self.slots[i].as_mut() else { return };
172        while let Some(ev) = slot.conn.poll_endpoint_event() {
173            match ev {
174                EndpointEvent::NewCid(c) => {
175                    self.by_cid.insert(c, i);
176                    slot.cids.push(c);
177                }
178                EndpointEvent::RetireCid(c) => {
179                    if self.by_cid.get(&c) == Some(&i) {
180                        self.by_cid.remove(&c);
181                    }
182                    slot.cids.retain(|x| *x != c);
183                }
184                EndpointEvent::ResetToken(t) => {
185                    self.by_token.insert(t, i);
186                    slot.tokens.push(t);
187                }
188            }
189        }
190        // A connection counts as handshaking (for the Retry threshold) until it is
191        // established, including when it fails and closes during the handshake: otherwise
192        // a flood of bad ClientHellos from spoofed addresses would leave closing
193        // connections that never trigger Retry.
194        let hs = slot.handshaking && !slot.conn.is_established();
195        if slot.handshaking && !hs {
196            self.handshaking -= 1;
197        }
198        slot.handshaking = hs;
199        let d = slot.conn.poll_timeout();
200        if d != slot.deadline {
201            slot.deadline = d;
202            if let Some(d) = d {
203                self.timers.push(Reverse((d, i, slot.gen)));
204            }
205        }
206        if !slot.dirty {
207            slot.dirty = true;
208            self.dirty.push_back(i);
209        }
210    }
211
212    /// Processes a datagram from `remote`. Stateless replies are written to `response`.
213    pub fn recv(&mut self, now: Instant, remote: SocketAddr, data: &mut [u8], response: &mut Vec<u8>) -> DatagramEvent {
214        self.cfg.epoch.get_or_init(|| now);
215        response.clear();
216        let Some(h) = packet::parse_header(data, 0, self.cfg.cid_len) else { return DatagramEvent::Dropped };
217        let dcid = Cid::new(&data[h.dcid.clone()]);
218        if let Some(&i) = self.by_cid.get(&dcid) {
219            if let Some(slot) = self.slots[i].as_mut() {
220                slot.conn.recv(now, remote, data);
221                self.process(i);
222                return DatagramEvent::Connection(ConnHandle(i));
223            }
224        }
225        if h.ty == PacketType::OneRtt {
226            // A stateless reset from a peer that lost our connection (RFC 9000 §10.3.1).
227            if data.len() >= MIN_RESET_SIZE {
228                let tail: [u8; 16] = data[data.len() - 16..].try_into().unwrap();
229                if let Some(&i) = self.by_token.get(&tail) {
230                    if let Some(slot) = self.slots[i].as_mut() {
231                        slot.conn.recv(now, remote, data);
232                        self.process(i);
233                        return DatagramEvent::Connection(ConnHandle(i));
234                    }
235                }
236            }
237            return self.stateless_reset(remote, data.len(), &dcid, response);
238        }
239        if !self.server || h.ty == PacketType::VersionNegotiation {
240            return DatagramEvent::Dropped;
241        }
242        let scid = &data[h.scid.clone()];
243        if !self.cfg.versions.contains(&h.version) {
244            if data.len() < MIN_INITIAL_SIZE || scid.len() > 20 || dcid.len() > 20 {
245                return DatagramEvent::Dropped;
246            }
247            let mut r = [0u8; 1];
248            self.cfg.provider.fill_random(&mut r);
249            packet::version_negotiation(response, r[0], &dcid, scid, &self.cfg.versions);
250            return DatagramEvent::Response(Transmit { destination: remote, size: response.len(), segment_size: None });
251        }
252        if h.ty != PacketType::Initial || data.len() < MIN_INITIAL_SIZE || dcid.len() < 8 {
253            return DatagramEvent::Dropped;
254        }
255        if self.live >= self.cfg.max_connections {
256            return DatagramEvent::Dropped;
257        }
258        let scid = Cid::new(scid);
259        let tok = &data[h.token.clone()];
260        let mut odcid = dcid;
261        let mut retry_scid = None;
262        let mut validated = false;
263        if !tok.is_empty() {
264            match token::validate(&self.cfg, tok, remote, now) {
265                Some(Validated::Retry { odcid: o, rscid }) => {
266                    if rscid != dcid {
267                        return DatagramEvent::Dropped;
268                    }
269                    odcid = o;
270                    retry_scid = Some(rscid);
271                    validated = true;
272                }
273                Some(Validated::NewToken) => validated = true,
274                None => {
275                    // An invalid Retry token means a forged or stale Retry: drop. An invalid
276                    // NEW_TOKEN is treated as no token (RFC 9000 §8.1.3).
277                    if tok[0] == 1 {
278                        return DatagramEvent::Dropped;
279                    }
280                }
281            }
282        }
283        if !validated && (self.cfg.retry || self.handshaking >= self.cfg.retry_threshold) {
284            let mut b = [0u8; 21];
285            self.cfg.provider.fill_random(&mut b);
286            let new_scid = Cid::new(&b[1..1 + self.cfg.cid_len]);
287            let t = token::retry_token(&self.cfg, remote, now, &dcid, &new_scid);
288            packet::retry(response, &*self.cfg.provider, h.version, b[0], &dcid, &scid, &new_scid, &t);
289            return DatagramEvent::Response(Transmit { destination: remote, size: response.len(), segment_size: None });
290        }
291        let mut b = [0u8; 20];
292        self.cfg.provider.fill_random(&mut b[..self.cfg.cid_len]);
293        let local = Cid::new(&b[..self.cfg.cid_len]);
294        let conn = match Connection::accept(self.cfg.clone(), now, remote, h.version, local, dcid, scid, odcid, retry_scid, validated) {
295            Ok(c) => c,
296            Err(_) => return DatagramEvent::Dropped,
297        };
298        let i = self.insert(conn);
299        // Route both our CID and the client's chosen DCID (for Initial retransmissions).
300        for c in [local, dcid] {
301            self.by_cid.insert(c, i);
302            self.slots[i].as_mut().unwrap().cids.push(c);
303        }
304        self.slots[i].as_mut().unwrap().conn.recv(now, remote, data);
305        self.process(i);
306        DatagramEvent::NewConnection(ConnHandle(i))
307    }
308
309    /// RFC 9000 §10.3: a reset smaller than the packet that triggered it, at least 21 bytes.
310    fn stateless_reset(&self, remote: SocketAddr, len: usize, dcid: &Cid, response: &mut Vec<u8>) -> DatagramEvent {
311        if len <= MIN_RESET_SIZE {
312            return DatagramEvent::Dropped;
313        }
314        let size = (len - 1).min(MIN_RESET_SIZE + 22);
315        response.resize(size - 16, 0);
316        self.cfg.provider.fill_random(response);
317        response[0] = 0x40 | (response[0] & 0x3f);
318        response.extend_from_slice(&self.cfg.reset_token(dcid));
319        DatagramEvent::Response(Transmit { destination: remote, size, segment_size: None })
320    }
321
322    /// The next datagram batch to send from any connection.
323    pub fn poll_transmit(&mut self, now: Instant, buf: &mut Vec<u8>, max_datagrams: usize) -> Option<(ConnHandle, Transmit)> {
324        while let Some(&i) = self.dirty.front() {
325            let Some(slot) = self.slots[i].as_mut() else {
326                self.dirty.pop_front();
327                continue;
328            };
329            if let Some(t) = slot.conn.poll_transmit(now, buf, max_datagrams) {
330                return Some((ConnHandle(i), t));
331            }
332            slot.dirty = false;
333            self.dirty.pop_front();
334            self.process_timer_only(i);
335        }
336        None
337    }
338
339    fn process_timer_only(&mut self, i: usize) {
340        let Some(slot) = self.slots[i].as_mut() else { return };
341        while let Some(ev) = slot.conn.poll_endpoint_event() {
342            match ev {
343                EndpointEvent::NewCid(c) => {
344                    self.by_cid.insert(c, i);
345                    slot.cids.push(c);
346                }
347                EndpointEvent::RetireCid(c) => {
348                    if self.by_cid.get(&c) == Some(&i) {
349                        self.by_cid.remove(&c);
350                    }
351                    slot.cids.retain(|x| *x != c);
352                }
353                EndpointEvent::ResetToken(t) => {
354                    self.by_token.insert(t, i);
355                    slot.tokens.push(t);
356                }
357            }
358        }
359        let d = slot.conn.poll_timeout();
360        if d != slot.deadline {
361            slot.deadline = d;
362            if let Some(d) = d {
363                self.timers.push(Reverse((d, i, slot.gen)));
364            }
365        }
366    }
367
368    /// The earliest deadline over all connections.
369    pub fn poll_timeout(&mut self) -> Option<Instant> {
370        while let Some(&Reverse((t, i, gen))) = self.timers.peek() {
371            let live = self.slots.get(i).and_then(|s| s.as_ref()).is_some_and(|s| s.gen == gen && s.deadline == Some(t));
372            if live {
373                return Some(t);
374            }
375            self.timers.pop();
376        }
377        None
378    }
379
380    /// Fires due timers; appends the affected connections to `touched`.
381    pub fn handle_timeout(&mut self, now: Instant, touched: &mut Vec<ConnHandle>) {
382        while let Some(&Reverse((t, i, gen))) = self.timers.peek() {
383            if t > now {
384                break;
385            }
386            self.timers.pop();
387            let Some(slot) = self.slots.get_mut(i).and_then(|s| s.as_mut()) else { continue };
388            if slot.gen != gen || slot.deadline != Some(t) {
389                continue;
390            }
391            slot.deadline = None;
392            slot.conn.handle_timeout(now);
393            touched.push(ConnHandle(i));
394            self.process(i);
395        }
396    }
397}
simoom_quic/endpoint.rs

simoom_quic/
endpoint.rs
