Skip to main content

moq_noq_proto/
endpoint.rs

1use std::{
2    collections::{HashMap, hash_map},
3    convert::TryFrom,
4    fmt, mem,
5    net::{IpAddr, SocketAddr},
6    ops::{Index, IndexMut},
7    sync::Arc,
8};
9
10use bytes::{Buf, BufMut, Bytes, BytesMut};
11use rand::{Rng, RngExt, SeedableRng, rngs::StdRng};
12use rustc_hash::FxHashMap;
13use slab::Slab;
14use thiserror::Error;
15use tracing::{debug, error, trace, warn};
16
17use crate::{
18    FourTuple, INITIAL_MTU, Instant, MAX_CID_SIZE, MIN_INITIAL_SIZE, PathId, RESET_TOKEN_SIZE,
19    ResetToken, Side, Transmit, TransportConfig, TransportError,
20    cid_generator::ConnectionIdGenerator,
21    coding::{BufMutExt, Decodable, Encodable, UnexpectedEnd},
22    config::{ClientConfig, EndpointConfig, ServerConfig},
23    connection::{Connection, ConnectionError, SideArgs},
24    crypto::{self, Keys, UnsupportedVersion},
25    frame,
26    packet::{
27        FixedLengthConnectionIdParser, Header, InitialHeader, InitialPacket, PacketDecodeError,
28        PacketNumber, PartialDecode, ProtectedInitialHeader,
29    },
30    shared::{
31        ConnectionEvent, ConnectionEventInner, ConnectionId, DatagramConnectionEvent, EcnCodepoint,
32        EndpointEvent, EndpointEventInner, IssuedCid,
33    },
34    token::{IncomingToken, InvalidRetryTokenError, Token, TokenPayload},
35    transport_parameters::{PreferredAddress, TransportParameters},
36};
37
38/// The main entry point to the library
39///
40/// This object performs no I/O whatsoever. Instead, it consumes incoming packets and
41/// connection-generated events via `handle` and `handle_event`.
42pub struct Endpoint {
43    rng: StdRng,
44    index: ConnectionIndex,
45    connections: Slab<ConnectionMeta>,
46    local_cid_generator: Box<dyn ConnectionIdGenerator>,
47    config: Arc<EndpointConfig>,
48    server_config: Option<Arc<ServerConfig>>,
49    /// Whether the underlying UDP socket promises not to fragment packets
50    allow_mtud: bool,
51    /// Time at which a stateless reset was most recently sent
52    last_stateless_reset: Option<Instant>,
53    /// Buffered Initial and 0-RTT messages for pending incoming connections
54    incoming_buffers: Slab<IncomingBuffer>,
55    all_incoming_buffers_total_bytes: u64,
56}
57
58impl Endpoint {
59    /// Create a new endpoint
60    ///
61    /// `allow_mtud` enables path MTU detection when requested by `Connection` configuration for
62    /// better performance. This requires that outgoing packets are never fragmented, which can be
63    /// achieved via e.g. the `IPV6_DONTFRAG` socket option.
64    pub fn new(
65        config: Arc<EndpointConfig>,
66        server_config: Option<Arc<ServerConfig>>,
67        allow_mtud: bool,
68    ) -> Self {
69        Self {
70            rng: config
71                .rng_seed
72                .map_or_else(|| StdRng::from_rng(&mut rand::rng()), StdRng::from_seed),
73            index: ConnectionIndex::default(),
74            connections: Slab::new(),
75            local_cid_generator: (config.connection_id_generator_factory.as_ref())(),
76            config,
77            server_config,
78            allow_mtud,
79            last_stateless_reset: None,
80            incoming_buffers: Slab::new(),
81            all_incoming_buffers_total_bytes: 0,
82        }
83    }
84
85    /// Replace the server configuration, affecting new incoming connections only
86    pub fn set_server_config(&mut self, server_config: Option<Arc<ServerConfig>>) {
87        self.server_config = server_config;
88    }
89
90    /// Process `EndpointEvent`s emitted from related `Connection`s
91    ///
92    /// In turn, processing this event may return a `ConnectionEvent` for the same `Connection`.
93    pub fn handle_event(
94        &mut self,
95        ch: ConnectionHandle,
96        event: EndpointEvent,
97    ) -> Option<ConnectionEvent> {
98        use EndpointEventInner::*;
99        match event.0 {
100            NeedIdentifiers(path_id, now, n) => {
101                return Some(self.send_new_identifiers(path_id, now, ch, n));
102            }
103            ResetToken(path_id, remote, token) => {
104                if let Some(old) = self.connections[ch]
105                    .reset_token
106                    .insert(path_id, (remote, token))
107                {
108                    self.index.connection_reset_tokens.remove(old.0, old.1);
109                }
110                if self.index.connection_reset_tokens.insert(remote, token, ch) {
111                    warn!("duplicate reset token");
112                }
113            }
114            RetireResetToken(path_id) => {
115                if let Some(old) = self.connections[ch].reset_token.remove(&path_id) {
116                    self.index.connection_reset_tokens.remove(old.0, old.1);
117                }
118            }
119            RetireConnectionId(now, path_id, seq, allow_more_cids) => {
120                if let Some(cid) = self.connections[ch]
121                    .local_cids
122                    .get_mut(&path_id)
123                    .and_then(|pcid| pcid.cids.remove(&seq))
124                {
125                    trace!(%path_id, "local CID retired {}: {}", seq, cid);
126                    self.index.retire(cid);
127                    if allow_more_cids {
128                        return Some(self.send_new_identifiers(path_id, now, ch, 1));
129                    }
130                }
131            }
132            Draining => {
133                // Nothing to do.
134            }
135            Drained => {
136                if let Some(conn) = self.connections.try_remove(ch.0) {
137                    self.index.remove(&conn);
138                } else {
139                    // This indicates a bug in downstream code, which could cause spurious
140                    // connection loss instead of this error if the CID was (re)allocated prior to
141                    // the illegal call.
142                    error!(id = ch.0, "unknown connection drained");
143                }
144            }
145        }
146        None
147    }
148
149    /// Process an incoming UDP datagram
150    pub fn handle(
151        &mut self,
152        now: Instant,
153        network_path: FourTuple,
154        ecn: Option<EcnCodepoint>,
155        data: BytesMut,
156        buf: &mut Vec<u8>,
157    ) -> Option<DatagramEvent> {
158        // Partially decode packet or short-circuit if unable
159        let datagram_len = data.len();
160        let mut event = match PartialDecode::new(
161            data,
162            &FixedLengthConnectionIdParser::new(self.local_cid_generator.cid_len()),
163            &self.config.supported_versions,
164            self.config.grease_quic_bit,
165        ) {
166            Ok((first_decode, remaining)) => DatagramConnectionEvent {
167                now,
168                network_path,
169                path_id: PathId::ZERO, // Corrected later for existing paths
170                ecn,
171                first_decode,
172                remaining,
173            },
174            Err(PacketDecodeError::UnsupportedVersion {
175                src_cid,
176                dst_cid,
177                version,
178            }) => {
179                if self.server_config.is_none() {
180                    debug!("dropping packet with unsupported version");
181                    return None;
182                }
183                trace!("sending version negotiation");
184                // Negotiate versions
185                Header::VersionNegotiate {
186                    random: self.rng.random::<u8>() | 0x40,
187                    src_cid: dst_cid,
188                    dst_cid: src_cid,
189                }
190                .encode(buf);
191                // Grease with a reserved version
192                buf.write::<u32>(match version {
193                    0x0a1a_2a3a => 0x0a1a_2a4a,
194                    _ => 0x0a1a_2a3a,
195                });
196                for &version in &self.config.supported_versions {
197                    buf.write(version);
198                }
199                return Some(DatagramEvent::Response(Transmit {
200                    destination: network_path.remote,
201                    ecn: None,
202                    size: buf.len(),
203                    segment_size: None,
204                    src_ip: network_path.local_ip,
205                }));
206            }
207            Err(e) => {
208                trace!("malformed header: {}", e);
209                return None;
210            }
211        };
212
213        let dst_cid = event.first_decode.dst_cid();
214
215        if let Some(route_to) = self.index.get(&network_path, &event.first_decode) {
216            event.path_id = match route_to {
217                RouteDatagramTo::Incoming(_) => PathId::ZERO,
218                RouteDatagramTo::Connection(_, path_id) => path_id,
219            };
220            match route_to {
221                RouteDatagramTo::Incoming(incoming_idx) => {
222                    let incoming_buffer = &mut self.incoming_buffers[incoming_idx];
223                    let config = &self.server_config.as_ref().unwrap();
224
225                    if incoming_buffer
226                        .total_bytes
227                        .checked_add(datagram_len as u64)
228                        .is_some_and(|n| n <= config.incoming_buffer_size)
229                        && self
230                            .all_incoming_buffers_total_bytes
231                            .checked_add(datagram_len as u64)
232                            .is_some_and(|n| n <= config.incoming_buffer_size_total)
233                    {
234                        incoming_buffer.datagrams.push(event);
235                        incoming_buffer.total_bytes += datagram_len as u64;
236                        self.all_incoming_buffers_total_bytes += datagram_len as u64;
237                    }
238
239                    None
240                }
241                RouteDatagramTo::Connection(ch, _path_id) => Some(DatagramEvent::ConnectionEvent(
242                    ch,
243                    ConnectionEvent(ConnectionEventInner::Datagram(event)),
244                )),
245            }
246        } else if event.first_decode.initial_header().is_some() {
247            // Potentially create a new connection
248
249            self.handle_first_packet(datagram_len, event, network_path, buf)
250        } else if event.first_decode.has_long_header() {
251            debug!(
252                "ignoring non-initial packet for unknown connection {}",
253                dst_cid
254            );
255            None
256        } else if !event.first_decode.is_initial()
257            && self.local_cid_generator.validate(dst_cid).is_err()
258        {
259            debug!("dropping packet with invalid CID");
260            None
261        } else if dst_cid.is_empty() {
262            trace!("dropping unrecognized short packet without ID");
263            None
264        } else {
265            // If we got this far, we're receiving a seemingly valid packet for an unknown
266            // connection. Send a stateless reset if possible.
267            self.stateless_reset(now, datagram_len, network_path, dst_cid, buf)
268                .map(DatagramEvent::Response)
269        }
270    }
271
272    /// Builds a stateless reset packet to respond with
273    fn stateless_reset(
274        &mut self,
275        now: Instant,
276        inciting_dgram_len: usize,
277        network_path: FourTuple,
278        dst_cid: ConnectionId,
279        buf: &mut Vec<u8>,
280    ) -> Option<Transmit> {
281        if self
282            .last_stateless_reset
283            .is_some_and(|last| last + self.config.min_reset_interval > now)
284        {
285            debug!("ignoring unexpected packet within minimum stateless reset interval");
286            return None;
287        }
288
289        /// Minimum amount of padding for the stateless reset to look like a short-header packet
290        const MIN_PADDING_LEN: usize = 5;
291
292        // Prevent amplification attacks and reset loops by ensuring we pad to at most 1 byte
293        // smaller than the inciting packet.
294        let max_padding_len = match inciting_dgram_len.checked_sub(RESET_TOKEN_SIZE) {
295            Some(headroom) if headroom > MIN_PADDING_LEN => headroom - 1,
296            _ => {
297                debug!(
298                    "ignoring unexpected {} byte packet: not larger than minimum stateless reset size",
299                    inciting_dgram_len
300                );
301                return None;
302            }
303        };
304
305        debug!(%dst_cid, %network_path.remote, "sending stateless reset");
306        self.last_stateless_reset = Some(now);
307        // Resets with at least this much padding can't possibly be distinguished from real packets
308        const IDEAL_MIN_PADDING_LEN: usize = MIN_PADDING_LEN + MAX_CID_SIZE;
309        let padding_len = if max_padding_len <= IDEAL_MIN_PADDING_LEN {
310            max_padding_len
311        } else {
312            self.rng
313                .random_range(IDEAL_MIN_PADDING_LEN..max_padding_len)
314        };
315        buf.reserve(padding_len + RESET_TOKEN_SIZE);
316        buf.resize(padding_len, 0);
317        self.rng.fill_bytes(&mut buf[0..padding_len]);
318        buf[0] = 0b0100_0000 | (buf[0] >> 2);
319        buf.extend_from_slice(&ResetToken::new(&*self.config.reset_key, dst_cid));
320
321        debug_assert!(buf.len() < inciting_dgram_len);
322
323        Some(Transmit {
324            destination: network_path.remote,
325            ecn: None,
326            size: buf.len(),
327            segment_size: None,
328            src_ip: network_path.local_ip,
329        })
330    }
331
332    /// Initiate a connection
333    pub fn connect(
334        &mut self,
335        now: Instant,
336        config: ClientConfig,
337        remote: SocketAddr,
338        server_name: &str,
339    ) -> Result<(ConnectionHandle, Connection), ConnectError> {
340        if self.cids_exhausted() {
341            return Err(ConnectError::CidsExhausted);
342        }
343        if remote.port() == 0 || remote.ip().is_unspecified() {
344            return Err(ConnectError::InvalidRemoteAddress(remote));
345        }
346        if !self.config.supported_versions.contains(&config.version) {
347            return Err(ConnectError::UnsupportedVersion);
348        }
349
350        let remote_id = (config.initial_dst_cid_provider)();
351        trace!(initial_dcid = %remote_id);
352
353        let ch = ConnectionHandle(self.connections.vacant_key());
354        let local_cid = self.new_cid(ch, PathId::ZERO);
355        let params = TransportParameters::new(
356            &config.transport,
357            &self.config,
358            self.local_cid_generator.as_ref(),
359            local_cid,
360            None,
361            &mut self.rng,
362        );
363        let tls = config
364            .crypto
365            .start_session(config.version, server_name, &params)?;
366
367        let conn = self.add_connection(
368            ch,
369            config.version,
370            remote_id,
371            local_cid,
372            remote_id,
373            FourTuple {
374                remote,
375                local_ip: None,
376            },
377            now,
378            tls,
379            config.transport,
380            SideArgs::Client {
381                token_store: config.token_store,
382                server_name: server_name.into(),
383            },
384            &params,
385        );
386        Ok((ch, conn))
387    }
388
389    /// Generates new CIDs and creates message to send to the connection state
390    fn send_new_identifiers(
391        &mut self,
392        path_id: PathId,
393        now: Instant,
394        ch: ConnectionHandle,
395        num: u64,
396    ) -> ConnectionEvent {
397        let mut ids = vec![];
398        for _ in 0..num {
399            let id = self.new_cid(ch, path_id);
400            let cid_meta = self.connections[ch].local_cids.entry(path_id).or_default();
401            let sequence = cid_meta.issued;
402            cid_meta.issued += 1;
403            cid_meta.cids.insert(sequence, id);
404            ids.push(IssuedCid {
405                path_id,
406                sequence,
407                id,
408                reset_token: ResetToken::new(&*self.config.reset_key, id),
409            });
410        }
411        ConnectionEvent(ConnectionEventInner::NewIdentifiers(
412            ids,
413            now,
414            self.local_cid_generator.cid_len(),
415            self.local_cid_generator.cid_lifetime(),
416        ))
417    }
418
419    /// Generate a connection ID for `ch`
420    fn new_cid(&mut self, ch: ConnectionHandle, path_id: PathId) -> ConnectionId {
421        loop {
422            let cid = self.local_cid_generator.generate_cid();
423            if cid.is_empty() {
424                // Zero-length CID; nothing to track
425                debug_assert_eq!(self.local_cid_generator.cid_len(), 0);
426                return cid;
427            }
428            if let hash_map::Entry::Vacant(e) = self.index.connection_ids.entry(cid) {
429                e.insert((ch, path_id));
430                break cid;
431            }
432        }
433    }
434
435    fn handle_first_packet(
436        &mut self,
437        datagram_len: usize,
438        event: DatagramConnectionEvent,
439        network_path: FourTuple,
440        buf: &mut Vec<u8>,
441    ) -> Option<DatagramEvent> {
442        let dst_cid = event.first_decode.dst_cid();
443        let header = event.first_decode.initial_header().unwrap();
444
445        let Some(server_config) = &self.server_config else {
446            debug!("packet for unrecognized connection {}", dst_cid);
447            return self
448                .stateless_reset(event.now, datagram_len, network_path, dst_cid, buf)
449                .map(DatagramEvent::Response);
450        };
451
452        if datagram_len < MIN_INITIAL_SIZE as usize {
453            debug!("ignoring short initial for connection {}", dst_cid);
454            return None;
455        }
456
457        // Saturation only happens under heavy load, where deriving initial keys per Initial just to
458        // reply with CONNECTION_REFUSED would starve packet processing for existing connections.
459        if self.cids_exhausted() || self.incoming_buffers.len() >= server_config.max_incoming {
460            debug!(
461                "ignoring initial for connection {} due to saturation",
462                dst_cid
463            );
464            return None;
465        }
466
467        let crypto = match server_config.crypto.initial_keys(header.version, dst_cid) {
468            Ok(keys) => keys,
469            Err(UnsupportedVersion) => {
470                // This probably indicates that the user set supported_versions incorrectly in
471                // `EndpointConfig`.
472                debug!(
473                    "ignoring initial packet version {:#x} unsupported by cryptographic layer",
474                    header.version
475                );
476                return None;
477            }
478        };
479
480        if let Err(reason) = self.early_validate_first_packet(header) {
481            return Some(DatagramEvent::Response(self.initial_close(
482                header.version,
483                network_path,
484                &crypto,
485                header.src_cid,
486                reason,
487                buf,
488            )));
489        }
490
491        let packet = match event.first_decode.finish(Some(&*crypto.header.remote)) {
492            Ok(packet) => packet,
493            Err(e) => {
494                trace!("unable to decode initial packet: {}", e);
495                return None;
496            }
497        };
498
499        if !packet.reserved_bits_valid() {
500            debug!("dropping connection attempt with invalid reserved bits");
501            return None;
502        }
503
504        let Header::Initial(header) = packet.header else {
505            panic!("non-initial packet in handle_first_packet()");
506        };
507
508        let server_config = self.server_config.as_ref().unwrap().clone();
509
510        let token = match IncomingToken::from_header(&header, &server_config, network_path.remote) {
511            Ok(token) => token,
512            Err(InvalidRetryTokenError) => {
513                debug!("rejecting invalid retry token");
514                return Some(DatagramEvent::Response(self.initial_close(
515                    header.version,
516                    network_path,
517                    &crypto,
518                    header.src_cid,
519                    TransportError::INVALID_TOKEN(""),
520                    buf,
521                )));
522            }
523        };
524
525        let incoming_idx = self.incoming_buffers.insert(IncomingBuffer::default());
526        self.index
527            .insert_initial_incoming(header.dst_cid, incoming_idx);
528
529        Some(DatagramEvent::NewConnection(Incoming {
530            received_at: event.now,
531            network_path,
532            ecn: event.ecn,
533            packet: InitialPacket {
534                header,
535                header_data: packet.header_data,
536                payload: packet.payload,
537            },
538            rest: event.remaining,
539            crypto,
540            token,
541            incoming_idx,
542            improper_drop_warner: IncomingImproperDropWarner,
543        }))
544    }
545
546    /// Attempt to accept this incoming connection (an error may still occur)
547    // box err to avoid clippy::result_large_err
548    pub fn accept(
549        &mut self,
550        mut incoming: Incoming,
551        now: Instant,
552        buf: &mut Vec<u8>,
553        server_config: Option<Arc<ServerConfig>>,
554    ) -> Result<(ConnectionHandle, Connection), Box<AcceptError>> {
555        let remote_address_validated = incoming.remote_address_validated();
556        incoming.improper_drop_warner.dismiss();
557        let incoming_buffer = self.incoming_buffers.remove(incoming.incoming_idx);
558        self.all_incoming_buffers_total_bytes -= incoming_buffer.total_bytes;
559
560        let packet_number = incoming.packet.header.number.expand(0);
561        let InitialHeader {
562            src_cid,
563            dst_cid,
564            version,
565            ..
566        } = incoming.packet.header;
567        let server_config =
568            server_config.unwrap_or_else(|| self.server_config.as_ref().unwrap().clone());
569
570        // The handshake starts when its Initial arrives, not when it is accepted.
571        if incoming
572            .received_at
573            .checked_add(server_config.transport.handshake_idle_timeout)
574            .is_some_and(|deadline| deadline <= now)
575        {
576            debug!("abandoning accept of stale initial");
577            self.index.remove_initial(dst_cid);
578            return Err(Box::new(AcceptError {
579                cause: ConnectionError::TimedOut,
580                response: None,
581            }));
582        }
583
584        if self.cids_exhausted() {
585            debug!("refusing connection");
586            self.index.remove_initial(dst_cid);
587            return Err(Box::new(AcceptError {
588                cause: ConnectionError::CidsExhausted,
589                response: Some(self.initial_close(
590                    version,
591                    incoming.network_path,
592                    &incoming.crypto,
593                    src_cid,
594                    TransportError::CONNECTION_REFUSED(""),
595                    buf,
596                )),
597            }));
598        }
599
600        if incoming
601            .crypto
602            .packet
603            .remote
604            .decrypt(
605                PathId::ZERO,
606                packet_number,
607                &incoming.packet.header_data,
608                &mut incoming.packet.payload,
609            )
610            .is_err()
611        {
612            debug!(packet_number, "failed to authenticate initial packet");
613            self.index.remove_initial(dst_cid);
614            return Err(Box::new(AcceptError {
615                cause: TransportError::PROTOCOL_VIOLATION("authentication failed").into(),
616                response: None,
617            }));
618        };
619
620        let ch = ConnectionHandle(self.connections.vacant_key());
621        let local_cid = self.new_cid(ch, PathId::ZERO);
622        let mut params = TransportParameters::new(
623            &server_config.transport,
624            &self.config,
625            self.local_cid_generator.as_ref(),
626            local_cid,
627            Some(&server_config),
628            &mut self.rng,
629        );
630        params.stateless_reset_token = Some(ResetToken::new(&*self.config.reset_key, local_cid));
631        params.original_dst_cid = Some(incoming.token.orig_dst_cid);
632        params.retry_src_cid = incoming.token.retry_src_cid;
633        let mut pref_addr_cid = None;
634        if server_config.has_preferred_address() {
635            let cid = self.new_cid(ch, PathId::ZERO);
636            pref_addr_cid = Some(cid);
637            params.preferred_address = Some(PreferredAddress {
638                address_v4: server_config.preferred_address_v4,
639                address_v6: server_config.preferred_address_v6,
640                connection_id: cid,
641                stateless_reset_token: ResetToken::new(&*self.config.reset_key, cid),
642            });
643        }
644
645        let tls = server_config.crypto.start_session(version, &params);
646        let transport_config = server_config.transport.clone();
647        let mut conn = self.add_connection(
648            ch,
649            version,
650            dst_cid,
651            local_cid,
652            src_cid,
653            incoming.network_path,
654            incoming.received_at,
655            tls,
656            transport_config,
657            SideArgs::Server {
658                server_config,
659                pref_addr_cid,
660                path_validated: remote_address_validated,
661            },
662            &params,
663        );
664        self.index.insert_initial(dst_cid, ch);
665
666        match conn.handle_first_packet(
667            incoming.received_at,
668            incoming.network_path,
669            incoming.ecn,
670            packet_number,
671            incoming.packet,
672            incoming.rest,
673        ) {
674            Ok(()) => {
675                trace!(
676                    id = ch.0,
677                    icid = %dst_cid,
678                    network_path = %incoming.network_path,
679                    "new connection",
680                );
681
682                for event in incoming_buffer.datagrams {
683                    conn.handle_event(ConnectionEvent(ConnectionEventInner::Datagram(event)))
684                }
685
686                Ok((ch, conn))
687            }
688            Err(e) => {
689                debug!("handshake failed: {}", e);
690                self.handle_event(ch, EndpointEvent(EndpointEventInner::Drained));
691                let response = match e {
692                    ConnectionError::TransportError(ref e) => Some(self.initial_close(
693                        version,
694                        incoming.network_path,
695                        &incoming.crypto,
696                        src_cid,
697                        e.clone(),
698                        buf,
699                    )),
700                    _ => None,
701                };
702                Err(Box::new(AcceptError { cause: e, response }))
703            }
704        }
705    }
706
707    /// Check if we should refuse a connection attempt regardless of the packet's contents
708    fn early_validate_first_packet(
709        &mut self,
710        header: &ProtectedInitialHeader,
711    ) -> Result<(), TransportError> {
712        // RFC9000 §7.2 dictates that initial (client-chosen) destination CIDs must be at least 8
713        // bytes. If this is a Retry packet, then the length must instead match our usual CID
714        // length. If we ever issue non-Retry address validation tokens via `NEW_TOKEN`, then we'll
715        // also need to validate CID length for those after decoding the token.
716        if header.dst_cid.len() < 8
717            && (header.token_pos.is_empty()
718                || header.dst_cid.len() != self.local_cid_generator.cid_len())
719        {
720            debug!(
721                "rejecting connection due to invalid DCID length {}",
722                header.dst_cid.len()
723            );
724            return Err(TransportError::PROTOCOL_VIOLATION(
725                "invalid destination CID length",
726            ));
727        }
728
729        Ok(())
730    }
731
732    /// Reject this incoming connection attempt
733    pub fn refuse(&mut self, incoming: Incoming, buf: &mut Vec<u8>) -> Transmit {
734        self.clean_up_incoming(&incoming);
735        incoming.improper_drop_warner.dismiss();
736
737        trace!(?incoming.network_path, "refusing incoming");
738
739        self.initial_close(
740            incoming.packet.header.version,
741            incoming.network_path,
742            &incoming.crypto,
743            incoming.packet.header.src_cid,
744            TransportError::CONNECTION_REFUSED(""),
745            buf,
746        )
747    }
748
749    /// Respond with a retry packet, requiring the client to retry with address validation
750    ///
751    /// Errors if `incoming.may_retry()` is false.
752    pub fn retry(&mut self, incoming: Incoming, buf: &mut Vec<u8>) -> Result<Transmit, RetryError> {
753        if !incoming.may_retry() {
754            trace!(
755                ?incoming.network_path,
756                "not responding retry on incoming due to missing src CID"
757            );
758            return Err(RetryError(Box::new(incoming)));
759        }
760
761        trace!(?incoming.network_path, "responding retry on incoming");
762
763        self.clean_up_incoming(&incoming);
764        incoming.improper_drop_warner.dismiss();
765
766        let server_config = self.server_config.as_ref().unwrap();
767
768        // First Initial
769        // The peer will use this as the DCID of its following Initials. Initial DCIDs are
770        // looked up separately from Handshake/Data DCIDs, so there is no risk of collision
771        // with established connections. In the unlikely event that a collision occurs
772        // between two connections in the initial phase, both will fail fast and may be
773        // retried by the application layer.
774        let local_cid = self.local_cid_generator.generate_cid();
775
776        let payload = TokenPayload::Retry {
777            address: incoming.network_path.remote,
778            orig_dst_cid: incoming.packet.header.dst_cid,
779            issued: server_config.time_source.now(),
780        };
781        let token = Token::new(payload, &mut self.rng).encode(&*server_config.token_key);
782
783        let header = Header::Retry {
784            src_cid: local_cid,
785            dst_cid: incoming.packet.header.src_cid,
786            version: incoming.packet.header.version,
787        };
788
789        let encode = header.encode(buf);
790        buf.put_slice(&token);
791        buf.extend_from_slice(&server_config.crypto.retry_tag(
792            incoming.packet.header.version,
793            incoming.packet.header.dst_cid,
794            buf,
795        ));
796        encode.finish(buf, &*incoming.crypto.header.local, None);
797
798        Ok(Transmit {
799            destination: incoming.network_path.remote,
800            ecn: None,
801            size: buf.len(),
802            segment_size: None,
803            src_ip: incoming.network_path.local_ip,
804        })
805    }
806
807    /// Ignore this incoming connection attempt, not sending any packet in response
808    ///
809    /// Doing this actively, rather than merely dropping the [`Incoming`], is necessary to prevent
810    /// memory leaks due to state within [`Endpoint`] tracking the incoming connection.
811    pub fn ignore(&mut self, incoming: Incoming) {
812        self.clean_up_incoming(&incoming);
813        incoming.improper_drop_warner.dismiss();
814
815        trace!(?incoming.network_path, "ignoring incoming");
816    }
817
818    /// Clean up endpoint data structures associated with an `Incoming`.
819    fn clean_up_incoming(&mut self, incoming: &Incoming) {
820        self.index.remove_initial(incoming.packet.header.dst_cid);
821        let incoming_buffer = self.incoming_buffers.remove(incoming.incoming_idx);
822        self.all_incoming_buffers_total_bytes -= incoming_buffer.total_bytes;
823    }
824
825    fn add_connection(
826        &mut self,
827        ch: ConnectionHandle,
828        version: u32,
829        init_cid: ConnectionId,
830        local_cid: ConnectionId,
831        remote_cid: ConnectionId,
832        network_path: FourTuple,
833        now: Instant,
834        tls: Box<dyn crypto::Session>,
835        transport_config: Arc<TransportConfig>,
836        side_args: SideArgs,
837        // Only used for qlog.
838        params: &TransportParameters,
839    ) -> Connection {
840        let mut rng_seed = [0; 32];
841        self.rng.fill_bytes(&mut rng_seed);
842        let side = side_args.side();
843        let pref_addr_cid = side_args.pref_addr_cid();
844
845        let qlog =
846            transport_config.create_qlog_sink(side_args.side(), network_path.remote, init_cid, now);
847
848        qlog.emit_connection_started(
849            now,
850            local_cid,
851            remote_cid,
852            network_path.remote,
853            network_path.local_ip,
854            params,
855        );
856
857        let conn = Connection::new(
858            self.config.clone(),
859            transport_config,
860            init_cid,
861            local_cid,
862            remote_cid,
863            network_path,
864            tls,
865            self.local_cid_generator.as_ref(),
866            now,
867            version,
868            self.allow_mtud,
869            rng_seed,
870            side_args,
871            qlog,
872        );
873
874        let mut path_cids = PathLocalCids::default();
875        path_cids.cids.insert(path_cids.issued, local_cid);
876        path_cids.issued += 1;
877
878        if let Some(cid) = pref_addr_cid {
879            debug_assert_eq!(path_cids.issued, 1, "preferred address cid seq must be 1");
880            path_cids.cids.insert(path_cids.issued, cid);
881            path_cids.issued += 1;
882        }
883
884        let id = self.connections.insert(ConnectionMeta {
885            init_cid,
886            local_cids: FxHashMap::from_iter([(PathId::ZERO, path_cids)]),
887            network_path,
888            side,
889            reset_token: Default::default(),
890        });
891        debug_assert_eq!(id, ch.0, "connection handle allocation out of sync");
892
893        self.index.insert_conn(network_path, local_cid, ch, side);
894
895        conn
896    }
897
898    fn initial_close(
899        &mut self,
900        version: u32,
901        network_path: FourTuple,
902        crypto: &Keys,
903        remote_id: ConnectionId,
904        reason: TransportError,
905        buf: &mut Vec<u8>,
906    ) -> Transmit {
907        // We don't need to worry about CID collisions in initial closes because the peer
908        // shouldn't respond, and if it does, and the CID collides, we'll just drop the
909        // unexpected response.
910        let local_id = self.local_cid_generator.generate_cid();
911        let number = PacketNumber::U8(0);
912        let header = Header::Initial(InitialHeader {
913            dst_cid: remote_id,
914            src_cid: local_id,
915            number,
916            token: Bytes::new(),
917            version,
918        });
919
920        let partial_encode = header.encode(buf);
921        let max_len =
922            INITIAL_MTU as usize - partial_encode.header_len - crypto.packet.local.tag_len();
923        frame::Close::from(reason).encoder(max_len).encode(buf);
924        buf.resize(buf.len() + crypto.packet.local.tag_len(), 0);
925        partial_encode.finish(
926            buf,
927            &*crypto.header.local,
928            Some((0, Default::default(), &*crypto.packet.local)),
929        );
930        Transmit {
931            destination: network_path.remote,
932            ecn: None,
933            size: buf.len(),
934            segment_size: None,
935            src_ip: network_path.local_ip,
936        }
937    }
938
939    /// Access the configuration used by this endpoint
940    pub fn config(&self) -> &EndpointConfig {
941        &self.config
942    }
943
944    /// Number of connections that are currently open
945    pub fn open_connections(&self) -> usize {
946        self.connections.len()
947    }
948
949    /// Counter for the number of bytes currently used
950    /// in the buffers for Initial and 0-RTT messages for pending incoming connections
951    pub fn incoming_buffer_bytes(&self) -> u64 {
952        self.all_incoming_buffers_total_bytes
953    }
954
955    #[cfg(test)]
956    pub(crate) fn known_connections(&self) -> usize {
957        let x = self.connections.len();
958        debug_assert_eq!(x, self.index.connection_ids_initial.len());
959        // Not all connections have known reset tokens
960        debug_assert!(x >= self.index.connection_reset_tokens.0.len());
961        // Not all connections have unique remotes, and 0-length CIDs might not be in use.
962        debug_assert!(x >= self.index.incoming_connection_remotes.len());
963        debug_assert!(x >= self.index.outgoing_connection_remotes.len());
964        x
965    }
966
967    #[cfg(test)]
968    pub(crate) fn known_cids(&self) -> usize {
969        self.index.connection_ids.len()
970    }
971
972    /// Whether we've used up 3/4 of the available CID space
973    ///
974    /// We leave some space unused so that `new_cid` can be relied upon to finish quickly. We don't
975    /// bother to check when CID longer than 4 bytes are used because 2^40 connections is a lot.
976    fn cids_exhausted(&self) -> bool {
977        let cid_len = self.local_cid_generator.cid_len();
978        if cid_len == 0 || cid_len > 4 {
979            return false;
980        }
981
982        // Keep this architecture-independent: on 32-bit targets, 2usize.pow(32) overflows.
983        let bits = (cid_len * 8) as u32;
984        let space = 1u64 << bits;
985        let reserve = 1u64 << (bits - 2);
986        let len = self.index.connection_ids.len() as u64;
987
988        len > (space - reserve)
989    }
990}
991
992impl fmt::Debug for Endpoint {
993    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
994        fmt.debug_struct("Endpoint")
995            .field("rng", &self.rng)
996            .field("index", &self.index)
997            .field("connections", &self.connections)
998            .field("config", &self.config)
999            .field("server_config", &self.server_config)
1000            // incoming_buffers too large
1001            .field("incoming_buffers.len", &self.incoming_buffers.len())
1002            .field(
1003                "all_incoming_buffers_total_bytes",
1004                &self.all_incoming_buffers_total_bytes,
1005            )
1006            .finish()
1007    }
1008}
1009
1010/// Buffered Initial and 0-RTT messages for a pending incoming connection
1011#[derive(Default)]
1012struct IncomingBuffer {
1013    datagrams: Vec<DatagramConnectionEvent>,
1014    total_bytes: u64,
1015}
1016
1017/// Part of protocol state incoming datagrams can be routed to
1018#[derive(Copy, Clone, Debug)]
1019enum RouteDatagramTo {
1020    Incoming(usize),
1021    Connection(ConnectionHandle, PathId),
1022}
1023
1024/// Maps packets to existing connections
1025#[derive(Default, Debug)]
1026struct ConnectionIndex {
1027    /// Identifies connections based on the initial DCID the peer utilized
1028    ///
1029    /// Uses a standard `HashMap` to protect against hash collision attacks.
1030    ///
1031    /// Used by the server, not the client.
1032    connection_ids_initial: HashMap<ConnectionId, RouteDatagramTo>,
1033    /// Identifies connections based on locally created CIDs
1034    ///
1035    /// Uses a cheaper hash function since keys are locally created
1036    connection_ids: FxHashMap<ConnectionId, (ConnectionHandle, PathId)>,
1037    /// Identifies incoming connections with zero-length CIDs
1038    ///
1039    /// Uses a standard `HashMap` to protect against hash collision attacks.
1040    incoming_connection_remotes: HashMap<FourTuple, ConnectionHandle>,
1041    /// Identifies outgoing connections with zero-length CIDs
1042    ///
1043    /// We don't yet support explicit source addresses for client connections, and zero-length CIDs
1044    /// require a unique 4-tuple, so at most one client connection with zero-length local CIDs
1045    /// may be established per remote. We must omit the local address from the key because we don't
1046    /// necessarily know what address we're sending from, and hence receiving at.
1047    ///
1048    /// Uses a standard `HashMap` to protect against hash collision attacks.
1049    // TODO(matheus23): It's possible this could be changed now that we track the full 4-tuple on
1050    // the client side, too.
1051    outgoing_connection_remotes: HashMap<SocketAddr, ConnectionHandle>,
1052    /// Reset tokens provided by the peer for the CID each connection is currently sending to
1053    ///
1054    /// Incoming stateless resets do not have correct CIDs, so we need this to identify the correct
1055    /// recipient, if any.
1056    connection_reset_tokens: ResetTokenTable,
1057}
1058
1059impl ConnectionIndex {
1060    /// Associate an incoming connection with its initial destination CID
1061    fn insert_initial_incoming(&mut self, dst_cid: ConnectionId, incoming_key: usize) {
1062        if dst_cid.is_empty() {
1063            return;
1064        }
1065        self.connection_ids_initial
1066            .insert(dst_cid, RouteDatagramTo::Incoming(incoming_key));
1067    }
1068
1069    /// Remove an association with an initial destination CID
1070    fn remove_initial(&mut self, dst_cid: ConnectionId) {
1071        if dst_cid.is_empty() {
1072            return;
1073        }
1074        let removed = self.connection_ids_initial.remove(&dst_cid);
1075        debug_assert!(removed.is_some());
1076    }
1077
1078    /// Associate a connection with its initial destination CID
1079    fn insert_initial(&mut self, dst_cid: ConnectionId, connection: ConnectionHandle) {
1080        if dst_cid.is_empty() {
1081            return;
1082        }
1083        self.connection_ids_initial.insert(
1084            dst_cid,
1085            RouteDatagramTo::Connection(connection, PathId::ZERO),
1086        );
1087    }
1088
1089    /// Associate a connection with its first locally-chosen destination CID if used, or otherwise
1090    /// its current 4-tuple
1091    fn insert_conn(
1092        &mut self,
1093        network_path: FourTuple,
1094        dst_cid: ConnectionId,
1095        connection: ConnectionHandle,
1096        side: Side,
1097    ) {
1098        match dst_cid.len() {
1099            0 => match side {
1100                Side::Server => {
1101                    self.incoming_connection_remotes
1102                        .insert(network_path, connection);
1103                }
1104                Side::Client => {
1105                    self.outgoing_connection_remotes
1106                        .insert(network_path.remote, connection);
1107                }
1108            },
1109            _ => {
1110                self.connection_ids
1111                    .insert(dst_cid, (connection, PathId::ZERO));
1112            }
1113        }
1114    }
1115
1116    /// Discard a connection ID
1117    fn retire(&mut self, dst_cid: ConnectionId) {
1118        self.connection_ids.remove(&dst_cid);
1119    }
1120
1121    /// Remove all references to a connection
1122    fn remove(&mut self, conn: &ConnectionMeta) {
1123        if conn.side.is_server() {
1124            self.remove_initial(conn.init_cid);
1125        }
1126        for cid in conn
1127            .local_cids
1128            .values()
1129            .flat_map(|pcids| pcids.cids.values())
1130        {
1131            self.connection_ids.remove(cid);
1132        }
1133        self.incoming_connection_remotes.remove(&conn.network_path);
1134        self.outgoing_connection_remotes
1135            .remove(&conn.network_path.remote);
1136        for (remote, token) in conn.reset_token.values() {
1137            self.connection_reset_tokens.remove(*remote, *token);
1138        }
1139    }
1140
1141    /// Find the existing connection that `datagram` should be routed to, if any
1142    fn get(&self, network_path: &FourTuple, datagram: &PartialDecode) -> Option<RouteDatagramTo> {
1143        if !datagram.dst_cid().is_empty()
1144            && let Some(&(ch, path_id)) = self.connection_ids.get(&datagram.dst_cid())
1145        {
1146            return Some(RouteDatagramTo::Connection(ch, path_id));
1147        }
1148        if (datagram.is_initial() || datagram.is_0rtt())
1149            && let Some(&ch) = self.connection_ids_initial.get(&datagram.dst_cid())
1150        {
1151            return Some(ch);
1152        }
1153        if datagram.dst_cid().is_empty() {
1154            if let Some(&ch) = self.incoming_connection_remotes.get(network_path) {
1155                // Never multipath because QUIC-MULTIPATH 1.1 mandates the use of non-zero
1156                // length CIDs.  So this is always PathId::ZERO.
1157                return Some(RouteDatagramTo::Connection(ch, PathId::ZERO));
1158            }
1159            if let Some(&ch) = self.outgoing_connection_remotes.get(&network_path.remote) {
1160                // Like above, QUIC-MULTIPATH 1.1 mandates the use of non-zero length CIDs.
1161                return Some(RouteDatagramTo::Connection(ch, PathId::ZERO));
1162            }
1163        }
1164        let data = datagram.data();
1165        if data.len() < RESET_TOKEN_SIZE {
1166            return None;
1167        }
1168        // For stateless resets the PathId is meaningless since it closes the entire
1169        // connection regardless of path.  So use PathId::ZERO.
1170        self.connection_reset_tokens
1171            .get(network_path.remote, &data[data.len() - RESET_TOKEN_SIZE..])
1172            .cloned()
1173            .map(|ch| RouteDatagramTo::Connection(ch, PathId::ZERO))
1174    }
1175}
1176
1177#[derive(Debug)]
1178pub(crate) struct ConnectionMeta {
1179    init_cid: ConnectionId,
1180    /// Locally issues CIDs for each path
1181    local_cids: FxHashMap<PathId, PathLocalCids>,
1182    /// Remote/local addresses the connection began with
1183    ///
1184    /// Only needed to support connections with zero-length CIDs, which cannot migrate, so we don't
1185    /// bother keeping it up to date.
1186    network_path: FourTuple,
1187    side: Side,
1188    /// Reset tokens provided by the peer for CIDs we're currently sending to
1189    ///
1190    /// Since each reset token is for a CID, it is also for a fixed remote address which is
1191    /// also stored. This allows us to look up which reset tokens we might expect from a
1192    /// given remote address, see [`ResetTokenTable`].
1193    ///
1194    /// Each path has its own active CID. We use the [`PathId`] as a unique index, allowing
1195    /// us to retire the reset token when a path is abandoned.
1196    // TODO(matheus23): Should be migrated to make reset tokens per 4-tuple instead of per remote
1197    // addr
1198    reset_token: FxHashMap<PathId, (SocketAddr, ResetToken)>,
1199}
1200
1201/// Local connection IDs for a single path
1202#[derive(Debug, Default)]
1203struct PathLocalCids {
1204    /// Number of connection IDs that have been issued in (PATH_)NEW_CONNECTION_ID frames
1205    ///
1206    /// Another way of saying this is that this is the next sequence number to be issued.
1207    issued: u64,
1208    /// Issues CIDs indexed by their sequence number.
1209    cids: FxHashMap<u64, ConnectionId>,
1210}
1211
1212/// Internal identifier for a `Connection` currently associated with an endpoint
1213#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)]
1214pub struct ConnectionHandle(pub usize);
1215
1216impl From<ConnectionHandle> for usize {
1217    fn from(x: ConnectionHandle) -> Self {
1218        x.0
1219    }
1220}
1221
1222impl Index<ConnectionHandle> for Slab<ConnectionMeta> {
1223    type Output = ConnectionMeta;
1224    fn index(&self, ch: ConnectionHandle) -> &ConnectionMeta {
1225        &self[ch.0]
1226    }
1227}
1228
1229impl IndexMut<ConnectionHandle> for Slab<ConnectionMeta> {
1230    fn index_mut(&mut self, ch: ConnectionHandle) -> &mut ConnectionMeta {
1231        &mut self[ch.0]
1232    }
1233}
1234
1235/// Event resulting from processing a single datagram
1236pub enum DatagramEvent {
1237    /// The datagram is redirected to its `Connection`
1238    ConnectionEvent(ConnectionHandle, ConnectionEvent),
1239    /// The datagram may result in starting a new `Connection`
1240    NewConnection(Incoming),
1241    /// Response generated directly by the endpoint
1242    Response(Transmit),
1243}
1244
1245/// An incoming connection for which the server has not yet begun its part of the handshake.
1246#[derive(derive_more::Debug)]
1247pub struct Incoming {
1248    #[debug(skip)]
1249    received_at: Instant,
1250    network_path: FourTuple,
1251    ecn: Option<EcnCodepoint>,
1252    #[debug(skip)]
1253    packet: InitialPacket,
1254    #[debug(skip)]
1255    rest: Option<BytesMut>,
1256    #[debug(skip)]
1257    crypto: Keys,
1258    token: IncomingToken,
1259    incoming_idx: usize,
1260    #[debug(skip)]
1261    improper_drop_warner: IncomingImproperDropWarner,
1262}
1263
1264impl Incoming {
1265    /// The local IP address which was used when the peer established the connection
1266    pub fn local_ip(&self) -> Option<IpAddr> {
1267        self.network_path.local_ip
1268    }
1269
1270    /// The peer's UDP address
1271    pub fn remote_address(&self) -> SocketAddr {
1272        self.network_path.remote
1273    }
1274
1275    /// Whether the socket address that is initiating this connection has been validated
1276    ///
1277    /// This means that the sender of the initial packet has proved that they can receive traffic
1278    /// sent to `self.remote_address()`.
1279    ///
1280    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
1281    /// The inverse is not guaranteed.
1282    pub fn remote_address_validated(&self) -> bool {
1283        self.token.validated
1284    }
1285
1286    /// Whether it is legal to respond with a retry packet
1287    ///
1288    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
1289    /// The inverse is not guaranteed.
1290    pub fn may_retry(&self) -> bool {
1291        self.token.retry_src_cid.is_none()
1292    }
1293
1294    /// The original destination connection ID sent by the client
1295    pub fn orig_dst_cid(&self) -> ConnectionId {
1296        self.token.orig_dst_cid
1297    }
1298
1299    /// Decrypt the Initial packet payload
1300    ///
1301    /// This clones and decrypts the packet payload (~1200 bytes).
1302    /// Can be used to extract information from the TLS ClientHello without completing the
1303    /// handshake.
1304    pub fn decrypt(&self) -> Option<DecryptedInitial> {
1305        let packet_number = self.packet.header.number.expand(0);
1306        let mut payload = self.packet.payload.clone();
1307        self.crypto
1308            .packet
1309            .remote
1310            .decrypt(
1311                PathId::ZERO,
1312                packet_number,
1313                &self.packet.header_data,
1314                &mut payload,
1315            )
1316            .ok()?;
1317        Some(DecryptedInitial(payload.freeze()))
1318    }
1319}
1320
1321/// Decrypted payload of a QUIC Initial packet
1322///
1323/// Obtained via [`Incoming::decrypt`]. Can be used to extract information from
1324/// the TLS ClientHello without completing the handshake.
1325pub struct DecryptedInitial(Bytes);
1326
1327impl DecryptedInitial {
1328    /// Best-effort extraction of the ALPN protocols from the TLS ClientHello
1329    ///
1330    /// Parses the CRYPTO frames to extract the ALPN extension. This is intended
1331    /// for routing and filtering; it is not guaranteed to succeed if the
1332    /// ClientHello spans multiple packets. Returns `None` if parsing fails.
1333    pub fn alpns(&self) -> Option<IncomingAlpns> {
1334        let frames = frame::Iter::new(self.0.clone()).ok()?;
1335        let mut first = None;
1336        let mut rest = Vec::new();
1337        for frame in frames {
1338            match frame {
1339                Ok(frame::Frame::Crypto(crypto)) => match first {
1340                    None => first = Some(crypto),
1341                    Some(_) => rest.push(crypto),
1342                },
1343                Err(_) => return None,
1344                _ => {}
1345            }
1346        }
1347        let first = first?;
1348
1349        // Fast path: single CRYPTO frame at offset 0 (no extra allocation)
1350        if rest.is_empty() && first.offset == 0 {
1351            let data = find_alpn_data(&first.data).ok()?;
1352            return Some(IncomingAlpns { data, pos: 0 });
1353        }
1354
1355        // Slow path: reassemble multiple CRYPTO frames
1356        rest.push(first);
1357        let source = assemble_crypto_frames(&mut rest)?;
1358        let data = find_alpn_data(&source).ok()?;
1359        Some(IncomingAlpns { data, pos: 0 })
1360    }
1361}
1362
1363/// TLS handshake type for ClientHello messages
1364/// <https://www.rfc-editor.org/rfc/rfc8446#section-4.1.2>
1365const TLS_HANDSHAKE_TYPE_CLIENT_HELLO: u8 = 0x01;
1366/// TLS extension type for Application-Layer Protocol Negotiation
1367/// <https://www.rfc-editor.org/rfc/rfc7301#section-3.1>
1368const TLS_EXTENSION_TYPE_ALPN: u16 = 0x0010;
1369/// Size of the fixed-length fields in a ClientHello (client_version + random)
1370/// <https://www.rfc-editor.org/rfc/rfc8446#section-4.1.2>
1371const TLS_CLIENT_HELLO_FIXED_LEN: usize = 2 + 32;
1372
1373/// Iterator over ALPN protocol names from a TLS ClientHello
1374///
1375/// Yields protocol names as [`Bytes`] slices. On the common fast path (single
1376/// CRYPTO frame), the only allocation is the payload clone for decryption.
1377pub struct IncomingAlpns {
1378    data: Bytes,
1379    pos: usize,
1380}
1381
1382impl Iterator for IncomingAlpns {
1383    type Item = Result<Bytes, UnexpectedEnd>;
1384
1385    fn next(&mut self) -> Option<Self::Item> {
1386        if self.pos >= self.data.len() {
1387            return None;
1388        }
1389        let len = self.data[self.pos] as usize;
1390        self.pos += 1;
1391        if self.pos + len > self.data.len() {
1392            return Some(Err(UnexpectedEnd));
1393        }
1394        let proto = self.data.slice(self.pos..self.pos + len);
1395        self.pos += len;
1396        Some(Ok(proto))
1397    }
1398}
1399
1400/// Sort CRYPTO frames by offset and concatenate into a contiguous `Bytes`
1401///
1402/// Returns `None` if there are gaps in the stream.
1403fn assemble_crypto_frames(frames: &mut [frame::Crypto]) -> Option<Bytes> {
1404    frames.sort_by_key(|f| f.offset);
1405    let capacity = frames.iter().map(|f| f.data.len()).sum();
1406    let mut buf = Vec::with_capacity(capacity);
1407    for f in frames.iter() {
1408        let start = f.offset as usize;
1409        if start > buf.len() {
1410            return None;
1411        }
1412        let end = start + f.data.len();
1413        if end > buf.len() {
1414            buf.extend_from_slice(&f.data[buf.len() - start..]);
1415        }
1416    }
1417    Some(Bytes::from(buf))
1418}
1419
1420/// Locate the raw ALPN protocol list data within a TLS ClientHello message
1421///
1422/// Parses the ClientHello in `source` and returns a [`Bytes`] containing the
1423/// u8-length-prefixed protocol names (after the outer ProtocolNameList u16
1424/// length prefix). The returned `Bytes` is a zero-copy slice of `source`.
1425fn find_alpn_data(source: &Bytes) -> Result<Bytes, UnexpectedEnd> {
1426    let mut r = &**source;
1427
1428    if u8::decode(&mut r)? != TLS_HANDSHAKE_TYPE_CLIENT_HELLO {
1429        return Err(UnexpectedEnd);
1430    }
1431
1432    // Handshake message length (u24), scopes the remainder
1433    let len = decode_u24(&mut r)?;
1434    let mut body = take(&mut r, len)?;
1435
1436    // Client version + random
1437    skip(&mut body, TLS_CLIENT_HELLO_FIXED_LEN)?;
1438
1439    // Session ID, cipher suites, compression methods
1440    skip_u8_prefixed(&mut body)?;
1441    skip_u16_prefixed(&mut body)?;
1442    skip_u8_prefixed(&mut body)?;
1443
1444    // Extensions
1445    let mut exts = take_u16_prefixed(&mut body)?;
1446    while exts.has_remaining() {
1447        let ext_type = u16::decode(&mut exts)?;
1448        let ext_data = take_u16_prefixed(&mut exts)?;
1449        if ext_type == TLS_EXTENSION_TYPE_ALPN {
1450            let list = take_u16_prefixed(&mut &*ext_data)?;
1451            return Ok(source.slice_ref(list));
1452        }
1453    }
1454    Err(UnexpectedEnd)
1455}
1456
1457/// Decode a big-endian u24 as usize
1458fn decode_u24(r: &mut &[u8]) -> Result<usize, UnexpectedEnd> {
1459    let a = u8::decode(r)?;
1460    let b = u8::decode(r)?;
1461    let c = u8::decode(r)?;
1462    Ok(u32::from_be_bytes([0, a, b, c]) as usize)
1463}
1464
1465/// Take `len` bytes from the front and return them as a sub-slice
1466fn take<'a>(r: &mut &'a [u8], len: usize) -> Result<&'a [u8], UnexpectedEnd> {
1467    if r.remaining() < len {
1468        return Err(UnexpectedEnd);
1469    }
1470    let data = &r[..len];
1471    r.advance(len);
1472    Ok(data)
1473}
1474
1475/// Read a u16 length prefix and return the sub-slice it covers
1476fn take_u16_prefixed<'a>(r: &mut &'a [u8]) -> Result<&'a [u8], UnexpectedEnd> {
1477    let len = u16::decode(r)? as usize;
1478    take(r, len)
1479}
1480
1481/// Advance past `n` bytes
1482fn skip(r: &mut &[u8], len: usize) -> Result<(), UnexpectedEnd> {
1483    take(r, len)?;
1484    Ok(())
1485}
1486
1487/// Skip a u8-length-prefixed field
1488fn skip_u8_prefixed(r: &mut &[u8]) -> Result<(), UnexpectedEnd> {
1489    let len = u8::decode(r)? as usize;
1490    skip(r, len)
1491}
1492
1493/// Skip a u16-length-prefixed field
1494fn skip_u16_prefixed(r: &mut &[u8]) -> Result<(), UnexpectedEnd> {
1495    let len = u16::decode(r)? as usize;
1496    skip(r, len)
1497}
1498
1499struct IncomingImproperDropWarner;
1500
1501impl IncomingImproperDropWarner {
1502    fn dismiss(self) {
1503        mem::forget(self);
1504    }
1505}
1506
1507impl Drop for IncomingImproperDropWarner {
1508    fn drop(&mut self) {
1509        warn!(
1510            "moq_noq_proto::Incoming dropped without passing to Endpoint::accept/refuse/retry/ignore \
1511               (may cause memory leak and eventual inability to accept new connections)"
1512        );
1513    }
1514}
1515
1516/// Errors in the parameters being used to create a new connection
1517///
1518/// These arise before any I/O has been performed.
1519#[derive(Debug, Error, Clone, PartialEq, Eq)]
1520pub enum ConnectError {
1521    /// The endpoint can no longer create new connections
1522    ///
1523    /// Indicates that a necessary component of the endpoint has been dropped or otherwise disabled.
1524    #[error("endpoint stopping")]
1525    EndpointStopping,
1526    /// The connection could not be created because not enough of the CID space is available
1527    ///
1528    /// Try using longer connection IDs
1529    #[error("CIDs exhausted")]
1530    CidsExhausted,
1531    /// The given server name was malformed
1532    #[error("invalid server name: {0}")]
1533    InvalidServerName(String),
1534    /// The remote [`SocketAddr`] supplied was malformed
1535    ///
1536    /// Examples include attempting to connect to port 0, or using an inappropriate address family.
1537    #[error("invalid remote address: {0}")]
1538    InvalidRemoteAddress(SocketAddr),
1539    /// No default client configuration was set up
1540    ///
1541    /// Use `Endpoint::connect_with` to specify a client configuration.
1542    #[error("no default client config")]
1543    NoDefaultClientConfig,
1544    /// The local endpoint does not support the QUIC version specified in the client configuration
1545    #[error("unsupported QUIC version")]
1546    UnsupportedVersion,
1547}
1548
1549/// Error type for attempting to accept an [`Incoming`]
1550#[derive(Debug)]
1551pub struct AcceptError {
1552    /// Underlying error describing reason for failure
1553    pub cause: ConnectionError,
1554    /// Optional response to transmit back
1555    pub response: Option<Transmit>,
1556}
1557
1558/// Error for attempting to retry an [`Incoming`] which already bears a token from a previous retry
1559#[derive(Debug, Error)]
1560#[error("retry() with validated Incoming")]
1561pub struct RetryError(Box<Incoming>);
1562
1563impl RetryError {
1564    /// Get the [`Incoming`]
1565    pub fn into_incoming(self) -> Incoming {
1566        *self.0
1567    }
1568}
1569
1570/// Reset Tokens which are associated with peer socket addresses
1571///
1572/// The standard `HashMap` is used since both `SocketAddr` and `ResetToken` are
1573/// peer generated and might be usable for hash collision attacks.
1574#[derive(Default, Debug)]
1575struct ResetTokenTable(HashMap<SocketAddr, HashMap<ResetToken, ConnectionHandle>>);
1576
1577impl ResetTokenTable {
1578    fn insert(&mut self, remote: SocketAddr, token: ResetToken, ch: ConnectionHandle) -> bool {
1579        self.0
1580            .entry(remote)
1581            .or_default()
1582            .insert(token, ch)
1583            .is_some()
1584    }
1585
1586    fn remove(&mut self, remote: SocketAddr, token: ResetToken) {
1587        use std::collections::hash_map::Entry;
1588        match self.0.entry(remote) {
1589            Entry::Vacant(_) => {}
1590            Entry::Occupied(mut e) => {
1591                e.get_mut().remove(&token);
1592                if e.get().is_empty() {
1593                    e.remove_entry();
1594                }
1595            }
1596        }
1597    }
1598
1599    fn get(&self, remote: SocketAddr, token: &[u8]) -> Option<&ConnectionHandle> {
1600        let token = ResetToken::from(<[u8; RESET_TOKEN_SIZE]>::try_from(token).ok()?);
1601        self.0.get(&remote)?.get(&token)
1602    }
1603}
1604
1605#[cfg(test)]
1606mod tests {
1607    use super::*;
1608
1609    #[test]
1610    fn assemble_contiguous() {
1611        let data = b"hello world";
1612        let mut frames = vec![
1613            frame::Crypto {
1614                offset: 0,
1615                data: Bytes::from_static(&data[..5]),
1616            },
1617            frame::Crypto {
1618                offset: 5,
1619                data: Bytes::from_static(&data[5..]),
1620            },
1621        ];
1622        assert_eq!(&assemble_crypto_frames(&mut frames).unwrap()[..], &data[..]);
1623    }
1624
1625    #[test]
1626    fn assemble_out_of_order() {
1627        let data = b"hello world";
1628        let mut frames = vec![
1629            frame::Crypto {
1630                offset: 5,
1631                data: Bytes::from_static(&data[5..]),
1632            },
1633            frame::Crypto {
1634                offset: 0,
1635                data: Bytes::from_static(&data[..5]),
1636            },
1637        ];
1638        assert_eq!(&assemble_crypto_frames(&mut frames).unwrap()[..], &data[..]);
1639    }
1640
1641    #[test]
1642    fn assemble_with_overlap() {
1643        let data = b"hello world";
1644        let mut frames = vec![
1645            frame::Crypto {
1646                offset: 0,
1647                data: Bytes::from_static(&data[..7]),
1648            },
1649            frame::Crypto {
1650                offset: 5,
1651                data: Bytes::from_static(&data[5..]),
1652            },
1653        ];
1654        assert_eq!(&assemble_crypto_frames(&mut frames).unwrap()[..], &data[..]);
1655    }
1656
1657    #[test]
1658    fn assemble_with_gap() {
1659        let mut frames = vec![frame::Crypto {
1660            offset: 10,
1661            data: Bytes::from_static(b"world"),
1662        }];
1663        assert!(assemble_crypto_frames(&mut frames).is_none());
1664    }
1665}