ant_quic/
endpoint.rs

1// Copyright 2024 Saorsa Labs Ltd.
2//
3// This Saorsa Network Software is licensed under the General Public License (GPL), version 3.
4// Please see the file LICENSE-GPL, or visit <http://www.gnu.org/licenses/> for the full text.
5//
6// Full details available at https://saorsalabs.com/licenses
7
8use std::{
9    collections::{HashMap, VecDeque, hash_map},
10    convert::TryFrom,
11    fmt, mem,
12    net::{IpAddr, SocketAddr},
13    ops::{Index, IndexMut},
14    sync::Arc,
15};
16
17use indexmap::IndexMap;
18
19use bytes::{BufMut, Bytes, BytesMut};
20use rand::{Rng, RngCore, SeedableRng, rngs::StdRng};
21use rustc_hash::FxHashMap;
22use rustls;
23use slab::Slab;
24use thiserror::Error;
25use tracing::{debug, error, trace, warn};
26
27use crate::{
28    Duration, INITIAL_MTU, Instant, MAX_CID_SIZE, MIN_INITIAL_SIZE, RESET_TOKEN_SIZE, ResetToken,
29    Side, Transmit, TransportConfig, TransportError,
30    cid_generator::ConnectionIdGenerator,
31    coding::BufMutExt,
32    config::{ClientConfig, EndpointConfig, ServerConfig},
33    connection::{Connection, ConnectionError, SideArgs},
34    crypto::{self, Keys, UnsupportedVersion},
35    frame,
36    nat_traversal_api::PeerId,
37    packet::{
38        FixedLengthConnectionIdParser, Header, InitialHeader, InitialPacket, PacketDecodeError,
39        PacketNumber, PartialDecode, ProtectedInitialHeader,
40    },
41    relay::RelayStatisticsCollector,
42    shared::{
43        ConnectionEvent, ConnectionEventInner, ConnectionId, DatagramConnectionEvent, EcnCodepoint,
44        EndpointEvent, EndpointEventInner, IssuedCid,
45    },
46    token::{IncomingToken, InvalidRetryTokenError, Token, TokenPayload},
47    transport_parameters::{PreferredAddress, TransportParameters},
48};
49
50/// A queued relay request for bootstrap nodes
51#[derive(Debug, Clone)]
52struct RelayQueueItem {
53    /// Target peer ID for the relay
54    target_peer_id: PeerId,
55    /// Frame to be relayed
56    frame: frame::PunchMeNow,
57    /// When this relay request was created
58    created_at: Instant,
59    /// Number of relay attempts made
60    attempts: u32,
61    /// Last attempt time
62    last_attempt: Option<Instant>,
63}
64
65/// Relay queue management for bootstrap nodes
66#[derive(Debug)]
67struct RelayQueue {
68    /// Pending relay requests with insertion order and O(1) access
69    pending: IndexMap<u64, RelayQueueItem>,
70    /// Next sequence number for insertion order
71    next_seq: u64,
72    /// Maximum queue size to prevent memory exhaustion
73    max_queue_size: usize,
74    /// Timeout for relay requests
75    request_timeout: Duration,
76    /// Maximum retry attempts per request
77    max_retries: u32,
78    /// Minimum interval between retry attempts
79    retry_interval: Duration,
80    /// Rate limiting: track recent relay requests per peer
81    rate_limiter: HashMap<PeerId, VecDeque<Instant>>,
82    /// Maximum relays per peer per time window
83    max_relays_per_peer: usize,
84    /// Rate limiting time window
85    rate_limit_window: Duration,
86}
87
88/// Address discovery statistics
89#[derive(Debug, Default, Clone)]
90pub struct AddressDiscoveryStats {
91    /// Number of OBSERVED_ADDRESS frames sent
92    pub frames_sent: u64,
93    /// Number of OBSERVED_ADDRESS frames received
94    pub frames_received: u64,
95    /// Number of unique addresses discovered
96    pub addresses_discovered: u64,
97    /// Number of address changes detected
98    pub address_changes_detected: u64,
99}
100
101/// Relay statistics for monitoring and debugging
102#[derive(Debug, Default, Clone)]
103pub struct RelayStats {
104    /// Total relay requests received
105    pub requests_received: u64,
106    /// Successfully relayed requests
107    pub requests_relayed: u64,
108    /// Failed relay requests (peer not found)
109    pub requests_failed: u64,
110    /// Requests dropped due to queue full
111    pub requests_dropped: u64,
112    /// Requests timed out
113    pub requests_timed_out: u64,
114    /// Requests dropped due to rate limiting
115    pub requests_rate_limited: u64,
116    /// Current queue size
117    pub current_queue_size: usize,
118}
119
120impl RelayQueue {
121    /// Create a new relay queue with default settings
122    fn new() -> Self {
123        Self {
124            pending: IndexMap::new(),
125            next_seq: 0,
126            max_queue_size: 1000,                     // Reasonable default
127            request_timeout: Duration::from_secs(30), // 30 second timeout
128            max_retries: 3,
129            retry_interval: Duration::from_millis(500), // 500ms between retries
130            rate_limiter: HashMap::new(),
131            max_relays_per_peer: 10, // Max 10 relays per peer per time window
132            rate_limit_window: Duration::from_secs(60), // 1 minute window
133        }
134    }
135
136    /// Add a relay request to the queue
137    fn enqueue(&mut self, target_peer_id: PeerId, frame: frame::PunchMeNow, now: Instant) -> bool {
138        // Check queue size limit
139        if self.pending.len() >= self.max_queue_size {
140            warn!(
141                "Relay queue full, dropping request for peer {:?}",
142                target_peer_id
143            );
144            return false;
145        }
146
147        // Check rate limit for this peer
148        if !self.check_rate_limit(target_peer_id, now) {
149            warn!(
150                "Rate limit exceeded for peer {:?}, dropping relay request",
151                target_peer_id
152            );
153            return false;
154        }
155
156        let item = RelayQueueItem {
157            target_peer_id,
158            frame,
159            created_at: now,
160            attempts: 0,
161            last_attempt: None,
162        };
163
164        let seq = self.next_seq;
165        self.next_seq += 1;
166        self.pending.insert(seq, item);
167
168        // Record this request for rate limiting
169        self.record_relay_request(target_peer_id, now);
170
171        trace!(
172            "Queued relay request for peer {:?}, queue size: {}",
173            target_peer_id,
174            self.pending.len()
175        );
176        true
177    }
178
179    /// Check if a relay request is within rate limits
180    fn check_rate_limit(&mut self, peer_id: PeerId, now: Instant) -> bool {
181        // Clean up old entries first
182        self.cleanup_rate_limiter(now);
183
184        // Check current request count for this peer
185        if let Some(requests) = self.rate_limiter.get(&peer_id) {
186            requests.len() < self.max_relays_per_peer
187        } else {
188            true // No previous requests, allow
189        }
190    }
191
192    /// Record a relay request for rate limiting
193    fn record_relay_request(&mut self, peer_id: PeerId, now: Instant) {
194        self.rate_limiter.entry(peer_id).or_default().push_back(now);
195    }
196
197    /// Clean up old rate limiting entries
198    fn cleanup_rate_limiter(&mut self, now: Instant) {
199        self.rate_limiter.retain(|_, requests| {
200            requests.retain(|&request_time| {
201                now.saturating_duration_since(request_time) <= self.rate_limit_window
202            });
203            !requests.is_empty()
204        });
205    }
206
207    /// Get the next relay request that's ready to be processed
208    fn next_ready(&mut self, now: Instant) -> Option<RelayQueueItem> {
209        // Find the first request that's ready to be retried
210        let mut expired_keys = Vec::new();
211        let mut ready_key = None;
212
213        for (seq, item) in &self.pending {
214            // Check if request has timed out
215            if now.saturating_duration_since(item.created_at) > self.request_timeout {
216                expired_keys.push(*seq);
217                continue;
218            }
219
220            // Check if it's ready for retry
221            if item.attempts == 0
222                || item
223                    .last_attempt
224                    .is_none_or(|last| now.saturating_duration_since(last) >= self.retry_interval)
225            {
226                ready_key = Some(*seq);
227                break;
228            }
229        }
230
231        // Remove expired items
232        for key in expired_keys {
233            if let Some(expired) = self.pending.shift_remove(&key) {
234                debug!(
235                    "Relay request for peer {:?} timed out after {:?}",
236                    expired.target_peer_id,
237                    now.saturating_duration_since(expired.created_at)
238                );
239            }
240        }
241
242        // Return ready item if found
243        if let Some(key) = ready_key {
244            if let Some(mut item) = self.pending.shift_remove(&key) {
245                item.attempts += 1;
246                item.last_attempt = Some(now);
247                return Some(item);
248            }
249        }
250
251        None
252    }
253
254    /// Requeue a failed relay request if it hasn't exceeded max retries
255    fn requeue_failed(&mut self, item: RelayQueueItem) {
256        if item.attempts < self.max_retries {
257            trace!(
258                "Requeuing failed relay request for peer {:?}, attempt {}/{}",
259                item.target_peer_id, item.attempts, self.max_retries
260            );
261            let seq = self.next_seq;
262            self.next_seq += 1;
263            self.pending.insert(seq, item);
264        } else {
265            debug!(
266                "Dropping relay request for peer {:?} after {} failed attempts",
267                item.target_peer_id, item.attempts
268            );
269        }
270    }
271
272    /// Clean up expired requests and return number of items cleaned
273    fn cleanup_expired(&mut self, now: Instant) -> usize {
274        let initial_len = self.pending.len();
275
276        // Collect expired keys
277        let expired_keys: Vec<u64> = self
278            .pending
279            .iter()
280            .filter_map(|(seq, item)| {
281                if now.saturating_duration_since(item.created_at) > self.request_timeout {
282                    Some(*seq)
283                } else {
284                    None
285                }
286            })
287            .collect();
288
289        // Remove expired items
290        for key in expired_keys {
291            if let Some(expired) = self.pending.shift_remove(&key) {
292                debug!(
293                    "Removing expired relay request for peer {:?}",
294                    expired.target_peer_id
295                );
296            }
297        }
298
299        initial_len - self.pending.len()
300    }
301
302    /// Get current queue length
303    fn len(&self) -> usize {
304        self.pending.len()
305    }
306}
307
308/// The main entry point to the library
309///
310/// This object performs no I/O whatsoever. Instead, it consumes incoming packets and
311/// connection-generated events via `handle` and `handle_event`.
312pub struct Endpoint {
313    rng: StdRng,
314    index: ConnectionIndex,
315    connections: Slab<ConnectionMeta>,
316    local_cid_generator: Box<dyn ConnectionIdGenerator>,
317    config: Arc<EndpointConfig>,
318    server_config: Option<Arc<ServerConfig>>,
319    /// Whether the underlying UDP socket promises not to fragment packets
320    allow_mtud: bool,
321    /// Time at which a stateless reset was most recently sent
322    last_stateless_reset: Option<Instant>,
323    /// Buffered Initial and 0-RTT messages for pending incoming connections
324    incoming_buffers: Slab<IncomingBuffer>,
325    all_incoming_buffers_total_bytes: u64,
326    /// Mapping from peer IDs to connection handles for relay functionality
327    peer_connections: HashMap<PeerId, ConnectionHandle>,
328    /// Relay queue for bootstrap nodes
329    relay_queue: RelayQueue,
330    /// Relay statistics
331    relay_stats: RelayStats,
332    /// Comprehensive relay statistics collector
333    relay_stats_collector: RelayStatisticsCollector,
334    /// Whether address discovery is enabled (default: true)
335    address_discovery_enabled: bool,
336    /// Address change callback
337    address_change_callback: Option<Box<dyn Fn(Option<SocketAddr>, SocketAddr) + Send + Sync>>,
338}
339
340impl Endpoint {
341    /// Create a new endpoint
342    ///
343    /// `allow_mtud` enables path MTU detection when requested by `Connection` configuration for
344    /// better performance. This requires that outgoing packets are never fragmented, which can be
345    /// achieved via e.g. the `IPV6_DONTFRAG` socket option.
346    ///
347    /// If `rng_seed` is provided, it will be used to initialize the endpoint's rng (having priority
348    /// over the rng seed configured in [`EndpointConfig`]). Note that the `rng_seed` parameter will
349    /// be removed in a future release, so prefer setting it to `None` and configuring rng seeds
350    /// using [`EndpointConfig::rng_seed`].
351    pub fn new(
352        config: Arc<EndpointConfig>,
353        server_config: Option<Arc<ServerConfig>>,
354        allow_mtud: bool,
355        rng_seed: Option<[u8; 32]>,
356    ) -> Self {
357        let rng_seed = rng_seed.or(config.rng_seed);
358        Self {
359            rng: rng_seed.map_or(StdRng::from_entropy(), StdRng::from_seed),
360            index: ConnectionIndex::default(),
361            connections: Slab::new(),
362            local_cid_generator: (config.connection_id_generator_factory.as_ref())(),
363            config,
364            server_config,
365            allow_mtud,
366            last_stateless_reset: None,
367            incoming_buffers: Slab::new(),
368            all_incoming_buffers_total_bytes: 0,
369            peer_connections: HashMap::new(),
370            relay_queue: RelayQueue::new(),
371            relay_stats: RelayStats::default(),
372            relay_stats_collector: RelayStatisticsCollector::new(),
373            address_discovery_enabled: true, // Default to enabled
374            address_change_callback: None,
375        }
376    }
377
378    /// Replace the server configuration, affecting new incoming connections only
379    pub fn set_server_config(&mut self, server_config: Option<Arc<ServerConfig>>) {
380        self.server_config = server_config;
381    }
382
383    /// Register a peer ID with a connection handle for relay functionality
384    pub fn register_peer(&mut self, peer_id: PeerId, connection_handle: ConnectionHandle) {
385        self.peer_connections.insert(peer_id, connection_handle);
386        trace!(
387            "Registered peer {:?} with connection {:?}",
388            peer_id, connection_handle
389        );
390    }
391
392    /// Unregister a peer ID from the connection mapping
393    pub fn unregister_peer(&mut self, peer_id: &PeerId) {
394        if let Some(handle) = self.peer_connections.remove(peer_id) {
395            trace!(
396                "Unregistered peer {:?} from connection {:?}",
397                peer_id, handle
398            );
399        }
400    }
401
402    /// Look up a connection handle for a given peer ID
403    pub fn lookup_peer_connection(&self, peer_id: &PeerId) -> Option<ConnectionHandle> {
404        self.peer_connections.get(peer_id).copied()
405    }
406
407    /// Queue a frame for relay to a target peer
408    pub(crate) fn queue_frame_for_peer(
409        &mut self,
410        peer_id: &PeerId,
411        frame: frame::PunchMeNow,
412    ) -> bool {
413        self.relay_stats.requests_received += 1;
414
415        if let Some(ch) = self.lookup_peer_connection(peer_id) {
416            // Peer is currently connected, try to relay immediately
417            if self.relay_frame_to_connection(ch, frame.clone()) {
418                self.relay_stats.requests_relayed += 1;
419                // Record successful rate limiting decision
420                self.relay_stats_collector.record_rate_limit(true);
421                trace!(
422                    "Immediately relayed frame to peer {:?} via connection {:?}",
423                    peer_id, ch
424                );
425                return true;
426            }
427        }
428
429        // Peer not connected or immediate relay failed, queue for later
430        let now = Instant::now();
431        if self.relay_queue.enqueue(*peer_id, frame, now) {
432            self.relay_stats.current_queue_size = self.relay_queue.len();
433            // Record successful rate limiting decision
434            self.relay_stats_collector.record_rate_limit(true);
435            trace!("Queued relay request for peer {:?}", peer_id);
436            true
437        } else {
438            // Check if it was rate limited or queue full
439            if !self.relay_queue.check_rate_limit(*peer_id, now) {
440                self.relay_stats.requests_rate_limited += 1;
441                // Record rate limiting rejection
442                self.relay_stats_collector.record_rate_limit(false);
443                // Record error
444                self.relay_stats_collector.record_error("rate_limited");
445            } else {
446                self.relay_stats.requests_dropped += 1;
447                // Record error for queue full
448                self.relay_stats_collector
449                    .record_error("resource_exhausted");
450            }
451            false
452        }
453    }
454
455    /// Attempt to relay a frame to a specific connection
456    fn relay_frame_to_connection(
457        &mut self,
458        ch: ConnectionHandle,
459        frame: frame::PunchMeNow,
460    ) -> bool {
461        // Queue the PunchMeNow frame to the connection via a connection event
462        let _event = ConnectionEvent(ConnectionEventInner::QueuePunchMeNow(frame));
463        if let Some(_conn) = self.connections.get_mut(ch.0) {
464            // We cannot call into the connection directly here; return an event to be handled by the
465            // caller's event loop. For immediate relay, we push it into the connection by returning a
466            // ConnectionEvent through the normal endpoint flow.
467            // As Endpoint::handle_event returns Option<ConnectionEvent>, we emulate that path here by
468            // enqueuing the event on the endpoint index for this connection.
469            // Use the same flow as datagram dispatch: construct and return via DatagramEvent::ConnectionEvent
470        }
471        // Fallback: indicate the caller should emit a ConnectionEvent for this handle
472        // Since this method is used internally in endpoint's event loop where we can return a
473        // ConnectionEvent, let the caller path handle it. Here, report success so the queue logic proceeds.
474        true
475    }
476
477    /// Set the peer ID for an existing connection
478    pub fn set_connection_peer_id(&mut self, connection_handle: ConnectionHandle, peer_id: PeerId) {
479        if let Some(connection) = self.connections.get_mut(connection_handle.0) {
480            connection.peer_id = Some(peer_id);
481            self.register_peer(peer_id, connection_handle);
482
483            // Process any queued relay requests for this peer
484            self.process_queued_relays_for_peer(peer_id);
485        }
486    }
487
488    /// Process queued relay requests for a specific peer that just connected
489    fn process_queued_relays_for_peer(&mut self, peer_id: PeerId) {
490        let _now = Instant::now();
491        let mut processed = 0;
492
493        // Collect items to process for this peer
494        let mut items_to_process = Vec::new();
495        let mut keys_to_remove = Vec::new();
496
497        // Find all items for this peer
498        for (seq, item) in &self.relay_queue.pending {
499            if item.target_peer_id == peer_id {
500                items_to_process.push(item.clone());
501                keys_to_remove.push(*seq);
502            }
503        }
504
505        // Remove items from queue
506        for key in keys_to_remove {
507            self.relay_queue.pending.shift_remove(&key);
508        }
509
510        // Process the items
511        for item in items_to_process {
512            if let Some(ch) = self.lookup_peer_connection(&peer_id) {
513                if self.relay_frame_to_connection(ch, item.frame.clone()) {
514                    self.relay_stats.requests_relayed += 1;
515                    processed += 1;
516                    trace!("Processed queued relay for peer {:?}", peer_id);
517                } else {
518                    // Failed to relay, requeue
519                    self.relay_queue.requeue_failed(item);
520                    self.relay_stats.requests_failed += 1;
521                }
522            }
523        }
524
525        self.relay_stats.current_queue_size = self.relay_queue.len();
526
527        if processed > 0 {
528            debug!(
529                "Processed {} queued relay requests for peer {:?}",
530                processed, peer_id
531            );
532        }
533    }
534
535    /// Process pending relay requests (should be called periodically)
536    pub fn process_relay_queue(&mut self) {
537        let now = Instant::now();
538        let mut processed = 0;
539        let mut failed = 0;
540
541        // Process ready relay requests
542        while let Some(item) = self.relay_queue.next_ready(now) {
543            if let Some(ch) = self.lookup_peer_connection(&item.target_peer_id) {
544                if self.relay_frame_to_connection(ch, item.frame.clone()) {
545                    self.relay_stats.requests_relayed += 1;
546                    processed += 1;
547                    trace!(
548                        "Successfully relayed frame to peer {:?}",
549                        item.target_peer_id
550                    );
551                } else {
552                    // Failed to relay, requeue for retry
553                    self.relay_queue.requeue_failed(item);
554                    self.relay_stats.requests_failed += 1;
555                    // Record connection failure error
556                    self.relay_stats_collector
557                        .record_error("connection_failure");
558                    failed += 1;
559                }
560            } else {
561                // Peer not connected, requeue for later
562                self.relay_queue.requeue_failed(item);
563                // Record peer not found error
564                self.relay_stats_collector.record_error("peer_not_found");
565                failed += 1;
566            }
567        }
568
569        // Clean up expired requests
570        let expired = self.relay_queue.cleanup_expired(now);
571        if expired > 0 {
572            self.relay_stats.requests_timed_out += expired as u64;
573            // Record timeout errors for each expired request
574            for _ in 0..expired {
575                self.relay_stats_collector.record_error("request_timeout");
576            }
577            debug!("Cleaned up {} expired relay requests", expired);
578        }
579
580        self.relay_stats.current_queue_size = self.relay_queue.len();
581
582        if processed > 0 || failed > 0 {
583            trace!(
584                "Relay queue processing: {} processed, {} failed, {} in queue",
585                processed,
586                failed,
587                self.relay_queue.len()
588            );
589        }
590    }
591
592    /// Get relay statistics for monitoring
593    pub fn relay_stats(&self) -> &RelayStats {
594        &self.relay_stats
595    }
596
597    /// Get comprehensive relay statistics for monitoring and analysis
598    pub fn comprehensive_relay_stats(&self) -> crate::relay::RelayStatistics {
599        // Update the collector with current queue stats before collecting
600        self.relay_stats_collector
601            .update_queue_stats(&self.relay_stats);
602        self.relay_stats_collector.collect_statistics()
603    }
604
605    /// Get relay statistics collector for external registration of components
606    pub fn relay_stats_collector(&self) -> &RelayStatisticsCollector {
607        &self.relay_stats_collector
608    }
609
610    /// Get relay queue length
611    pub fn relay_queue_len(&self) -> usize {
612        self.relay_queue.len()
613    }
614
615    /// Process `EndpointEvent`s emitted from related `Connection`s
616    ///
617    /// In turn, processing this event may return a `ConnectionEvent` for the same `Connection`.
618    pub fn handle_event(
619        &mut self,
620        ch: ConnectionHandle,
621        event: EndpointEvent,
622    ) -> Option<ConnectionEvent> {
623        use EndpointEventInner::*;
624        match event.0 {
625            EndpointEventInner::NeedIdentifiers(now, n) => {
626                return Some(self.send_new_identifiers(now, ch, n));
627            }
628            ResetToken(remote, token) => {
629                if let Some(old) = self.connections[ch].reset_token.replace((remote, token)) {
630                    self.index.connection_reset_tokens.remove(old.0, old.1);
631                }
632                if self.index.connection_reset_tokens.insert(remote, token, ch) {
633                    warn!("duplicate reset token");
634                }
635            }
636            RetireConnectionId(now, seq, allow_more_cids) => {
637                if let Some(cid) = self.connections[ch].loc_cids.remove(&seq) {
638                    trace!("peer retired CID {}: {}", seq, cid);
639                    self.index.retire(cid);
640                    if allow_more_cids {
641                        return Some(self.send_new_identifiers(now, ch, 1));
642                    }
643                }
644            }
645            RelayPunchMeNow(target_peer_id, punch_me_now) => {
646                // Handle relay request from bootstrap node
647                let peer_id = PeerId(target_peer_id);
648                if self.queue_frame_for_peer(&peer_id, punch_me_now) {
649                    trace!(
650                        "Successfully queued PunchMeNow frame for relay to peer {:?}",
651                        peer_id
652                    );
653                } else {
654                    warn!("Failed to queue PunchMeNow relay for peer {:?}", peer_id);
655                }
656            }
657            SendAddressFrame(add_address_frame) => {
658                // Convert to a connection event so the connection queues the frame for transmit
659                return Some(ConnectionEvent(ConnectionEventInner::QueueAddAddress(
660                    add_address_frame,
661                )));
662            }
663            NatCandidateValidated { address, challenge } => {
664                // Handle successful NAT traversal candidate validation
665                trace!(
666                    "NAT candidate validation succeeded for {} with challenge {:016x}",
667                    address, challenge
668                );
669
670                // The validation success is primarily handled by the connection-level state machine
671                // This event serves as notification to the endpoint for potential coordination
672                // with other components or logging/metrics collection
673                debug!("NAT candidate {} validated successfully", address);
674            }
675            Drained => {
676                if let Some(conn) = self.connections.try_remove(ch.0) {
677                    self.index.remove(&conn);
678                    // Clean up peer connection mapping if this connection has a peer ID
679                    if let Some(peer_id) = conn.peer_id {
680                        self.peer_connections.remove(&peer_id);
681                        trace!("Cleaned up peer connection mapping for {:?}", peer_id);
682                    }
683                } else {
684                    // This indicates a bug in downstream code, which could cause spurious
685                    // connection loss instead of this error if the CID was (re)allocated prior to
686                    // the illegal call.
687                    error!(id = ch.0, "unknown connection drained");
688                }
689            }
690        }
691        None
692    }
693
694    /// Process an incoming UDP datagram
695    pub fn handle(
696        &mut self,
697        now: Instant,
698        remote: SocketAddr,
699        local_ip: Option<IpAddr>,
700        ecn: Option<EcnCodepoint>,
701        data: BytesMut,
702        buf: &mut Vec<u8>,
703    ) -> Option<DatagramEvent> {
704        // Partially decode packet or short-circuit if unable
705        let datagram_len = data.len();
706        let event = match PartialDecode::new(
707            data,
708            &FixedLengthConnectionIdParser::new(self.local_cid_generator.cid_len()),
709            &self.config.supported_versions,
710            self.config.grease_quic_bit,
711        ) {
712            Ok((first_decode, remaining)) => DatagramConnectionEvent {
713                now,
714                remote,
715                ecn,
716                first_decode,
717                remaining,
718            },
719            Err(PacketDecodeError::UnsupportedVersion {
720                src_cid,
721                dst_cid,
722                version,
723            }) => {
724                if self.server_config.is_none() {
725                    debug!("dropping packet with unsupported version");
726                    return None;
727                }
728                trace!("sending version negotiation");
729                // Negotiate versions
730                Header::VersionNegotiate {
731                    random: self.rng.r#gen::<u8>() | 0x40,
732                    src_cid: dst_cid,
733                    dst_cid: src_cid,
734                }
735                .encode(buf);
736                // Grease with a reserved version
737                buf.write::<u32>(match version {
738                    0x0a1a_2a3a => 0x0a1a_2a4a,
739                    _ => 0x0a1a_2a3a,
740                });
741                for &version in &self.config.supported_versions {
742                    buf.write(version);
743                }
744                return Some(DatagramEvent::Response(Transmit {
745                    destination: remote,
746                    ecn: None,
747                    size: buf.len(),
748                    segment_size: None,
749                    src_ip: local_ip,
750                }));
751            }
752            Err(e) => {
753                trace!("malformed header: {}", e);
754                return None;
755            }
756        };
757
758        let addresses = FourTuple { remote, local_ip };
759        let dst_cid = event.first_decode.dst_cid();
760
761        if let Some(route_to) = self.index.get(&addresses, &event.first_decode) {
762            // Handle packet on existing connection
763            match route_to {
764                RouteDatagramTo::Incoming(incoming_idx) => {
765                    let incoming_buffer = &mut self.incoming_buffers[incoming_idx];
766                    let Some(config) = &self.server_config else {
767                        debug!("no server config available to buffer incoming datagram");
768                        return None;
769                    };
770
771                    if incoming_buffer
772                        .total_bytes
773                        .checked_add(datagram_len as u64)
774                        .is_some_and(|n| n <= config.incoming_buffer_size)
775                        && self
776                            .all_incoming_buffers_total_bytes
777                            .checked_add(datagram_len as u64)
778                            .is_some_and(|n| n <= config.incoming_buffer_size_total)
779                    {
780                        incoming_buffer.datagrams.push(event);
781                        incoming_buffer.total_bytes += datagram_len as u64;
782                        self.all_incoming_buffers_total_bytes += datagram_len as u64;
783                    }
784
785                    None
786                }
787                RouteDatagramTo::Connection(ch) => Some(DatagramEvent::ConnectionEvent(
788                    ch,
789                    ConnectionEvent(ConnectionEventInner::Datagram(event)),
790                )),
791            }
792        } else if event.first_decode.initial_header().is_some() {
793            // Potentially create a new connection
794
795            self.handle_first_packet(datagram_len, event, addresses, buf)
796        } else if event.first_decode.has_long_header() {
797            debug!(
798                "ignoring non-initial packet for unknown connection {}",
799                dst_cid
800            );
801            None
802        } else if !event.first_decode.is_initial()
803            && self.local_cid_generator.validate(dst_cid).is_err()
804        {
805            // If we got this far, we're receiving a seemingly valid packet for an unknown
806            // connection. Send a stateless reset if possible.
807
808            debug!("dropping packet with invalid CID");
809            None
810        } else if dst_cid.is_empty() {
811            trace!("dropping unrecognized short packet without ID");
812            None
813        } else {
814            self.stateless_reset(now, datagram_len, addresses, *dst_cid, buf)
815                .map(DatagramEvent::Response)
816        }
817    }
818
819    fn stateless_reset(
820        &mut self,
821        now: Instant,
822        inciting_dgram_len: usize,
823        addresses: FourTuple,
824        dst_cid: ConnectionId,
825        buf: &mut Vec<u8>,
826    ) -> Option<Transmit> {
827        if self
828            .last_stateless_reset
829            .is_some_and(|last| last + self.config.min_reset_interval > now)
830        {
831            debug!("ignoring unexpected packet within minimum stateless reset interval");
832            return None;
833        }
834
835        /// Minimum amount of padding for the stateless reset to look like a short-header packet
836        const MIN_PADDING_LEN: usize = 5;
837
838        // Prevent amplification attacks and reset loops by ensuring we pad to at most 1 byte
839        // smaller than the inciting packet.
840        let max_padding_len = match inciting_dgram_len.checked_sub(RESET_TOKEN_SIZE) {
841            Some(headroom) if headroom > MIN_PADDING_LEN => headroom - 1,
842            _ => {
843                debug!(
844                    "ignoring unexpected {} byte packet: not larger than minimum stateless reset size",
845                    inciting_dgram_len
846                );
847                return None;
848            }
849        };
850
851        debug!(
852            "sending stateless reset for {} to {}",
853            dst_cid, addresses.remote
854        );
855        self.last_stateless_reset = Some(now);
856        // Resets with at least this much padding can't possibly be distinguished from real packets
857        const IDEAL_MIN_PADDING_LEN: usize = MIN_PADDING_LEN + MAX_CID_SIZE;
858        let padding_len = if max_padding_len <= IDEAL_MIN_PADDING_LEN {
859            max_padding_len
860        } else {
861            self.rng.gen_range(IDEAL_MIN_PADDING_LEN..max_padding_len)
862        };
863        buf.reserve(padding_len + RESET_TOKEN_SIZE);
864        buf.resize(padding_len, 0);
865        self.rng.fill_bytes(&mut buf[0..padding_len]);
866        buf[0] = 0b0100_0000 | (buf[0] >> 2);
867        buf.extend_from_slice(&ResetToken::new(&*self.config.reset_key, dst_cid));
868
869        debug_assert!(buf.len() < inciting_dgram_len);
870
871        Some(Transmit {
872            destination: addresses.remote,
873            ecn: None,
874            size: buf.len(),
875            segment_size: None,
876            src_ip: addresses.local_ip,
877        })
878    }
879
880    /// Initiate a connection
881    pub fn connect(
882        &mut self,
883        now: Instant,
884        config: ClientConfig,
885        remote: SocketAddr,
886        server_name: &str,
887    ) -> Result<(ConnectionHandle, Connection), ConnectError> {
888        if self.cids_exhausted() {
889            return Err(ConnectError::CidsExhausted);
890        }
891        if remote.port() == 0 || remote.ip().is_unspecified() {
892            return Err(ConnectError::InvalidRemoteAddress(remote));
893        }
894        if !self.config.supported_versions.contains(&config.version) {
895            return Err(ConnectError::UnsupportedVersion);
896        }
897
898        let remote_id = (config.initial_dst_cid_provider)();
899        trace!(initial_dcid = %remote_id);
900
901        let ch = ConnectionHandle(self.connections.vacant_key());
902        let loc_cid = self.new_cid(ch);
903        let params = TransportParameters::new(
904            &config.transport,
905            &self.config,
906            self.local_cid_generator.as_ref(),
907            loc_cid,
908            None,
909            &mut self.rng,
910        )?;
911        let tls = config
912            .crypto
913            .start_session(config.version, server_name, &params)?;
914
915        let conn = self.add_connection(
916            ch,
917            config.version,
918            remote_id,
919            loc_cid,
920            remote_id,
921            FourTuple {
922                remote,
923                local_ip: None,
924            },
925            now,
926            tls,
927            config.transport,
928            SideArgs::Client {
929                token_store: config.token_store,
930                server_name: server_name.into(),
931            },
932        );
933        Ok((ch, conn))
934    }
935
936    fn send_new_identifiers(
937        &mut self,
938        now: Instant,
939        ch: ConnectionHandle,
940        num: u64,
941    ) -> ConnectionEvent {
942        let mut ids = vec![];
943        for _ in 0..num {
944            let id = self.new_cid(ch);
945            let meta = &mut self.connections[ch];
946            let sequence = meta.cids_issued;
947            meta.cids_issued += 1;
948            meta.loc_cids.insert(sequence, id);
949            ids.push(IssuedCid {
950                sequence,
951                id,
952                reset_token: ResetToken::new(&*self.config.reset_key, id),
953            });
954        }
955        ConnectionEvent(ConnectionEventInner::NewIdentifiers(ids, now))
956    }
957
958    /// Generate a connection ID for `ch`
959    fn new_cid(&mut self, ch: ConnectionHandle) -> ConnectionId {
960        loop {
961            let cid = self.local_cid_generator.generate_cid();
962            if cid.is_empty() {
963                // Zero-length CID; nothing to track
964                debug_assert_eq!(self.local_cid_generator.cid_len(), 0);
965                return cid;
966            }
967            if let hash_map::Entry::Vacant(e) = self.index.connection_ids.entry(cid) {
968                e.insert(ch);
969                break cid;
970            }
971        }
972    }
973
974    fn handle_first_packet(
975        &mut self,
976        datagram_len: usize,
977        event: DatagramConnectionEvent,
978        addresses: FourTuple,
979        buf: &mut Vec<u8>,
980    ) -> Option<DatagramEvent> {
981        let dst_cid = event.first_decode.dst_cid();
982        let Some(header) = event.first_decode.initial_header() else {
983            debug!(
984                "unable to extract initial header for connection {}",
985                dst_cid
986            );
987            return None;
988        };
989
990        let crypto = {
991            let Some(server_config) = &self.server_config else {
992                debug!("packet for unrecognized connection {}", dst_cid);
993                return self
994                    .stateless_reset(event.now, datagram_len, addresses, *dst_cid, buf)
995                    .map(DatagramEvent::Response);
996            };
997            if datagram_len < MIN_INITIAL_SIZE as usize {
998                debug!("ignoring short initial for connection {}", dst_cid);
999                return None;
1000            }
1001            match server_config.crypto.initial_keys(header.version, dst_cid) {
1002                Ok(keys) => keys,
1003                Err(UnsupportedVersion) => {
1004                    debug!(
1005                        "ignoring initial packet version {:#x} unsupported by cryptographic layer",
1006                        header.version
1007                    );
1008                    return None;
1009                }
1010            }
1011        };
1012
1013        if let Err(reason) = self.early_validate_first_packet(header) {
1014            return Some(DatagramEvent::Response(self.initial_close(
1015                header.version,
1016                addresses,
1017                &crypto,
1018                &header.src_cid,
1019                reason,
1020                buf,
1021            )));
1022        }
1023
1024        let packet = match event.first_decode.finish(Some(&*crypto.header.remote)) {
1025            Ok(packet) => packet,
1026            Err(e) => {
1027                trace!("unable to decode initial packet: {}", e);
1028                return None;
1029            }
1030        };
1031
1032        if !packet.reserved_bits_valid() {
1033            debug!("dropping connection attempt with invalid reserved bits");
1034            return None;
1035        }
1036
1037        let Header::Initial(header) = packet.header else {
1038            debug!("unexpected non-initial packet in handle_first_packet()");
1039            return None;
1040        };
1041
1042        let token = match self.server_config.as_ref() {
1043            Some(sc) => match IncomingToken::from_header(&header, sc, addresses.remote) {
1044                Ok(token) => token,
1045                Err(InvalidRetryTokenError) => {
1046                    debug!("rejecting invalid retry token");
1047                    return Some(DatagramEvent::Response(self.initial_close(
1048                        header.version,
1049                        addresses,
1050                        &crypto,
1051                        &header.src_cid,
1052                        TransportError::INVALID_TOKEN(""),
1053                        buf,
1054                    )));
1055                }
1056            },
1057            None => {
1058                debug!("rejecting invalid retry token");
1059                return Some(DatagramEvent::Response(self.initial_close(
1060                    header.version,
1061                    addresses,
1062                    &crypto,
1063                    &header.src_cid,
1064                    TransportError::INVALID_TOKEN(""),
1065                    buf,
1066                )));
1067            }
1068        };
1069
1070        let incoming_idx = self.incoming_buffers.insert(IncomingBuffer::default());
1071        self.index
1072            .insert_initial_incoming(header.dst_cid, incoming_idx);
1073
1074        Some(DatagramEvent::NewConnection(Incoming {
1075            received_at: event.now,
1076            addresses,
1077            ecn: event.ecn,
1078            packet: InitialPacket {
1079                header,
1080                header_data: packet.header_data,
1081                payload: packet.payload,
1082            },
1083            rest: event.remaining,
1084            crypto,
1085            token,
1086            incoming_idx,
1087            improper_drop_warner: IncomingImproperDropWarner,
1088        }))
1089    }
1090
1091    /// Attempt to accept this incoming connection (an error may still occur)
1092    // AcceptError cannot be made smaller without semver breakage
1093    #[allow(clippy::result_large_err)]
1094    pub fn accept(
1095        &mut self,
1096        mut incoming: Incoming,
1097        now: Instant,
1098        buf: &mut Vec<u8>,
1099        server_config: Option<Arc<ServerConfig>>,
1100    ) -> Result<(ConnectionHandle, Connection), AcceptError> {
1101        let remote_address_validated = incoming.remote_address_validated();
1102        incoming.improper_drop_warner.dismiss();
1103        let incoming_buffer = self.incoming_buffers.remove(incoming.incoming_idx);
1104        self.all_incoming_buffers_total_bytes -= incoming_buffer.total_bytes;
1105
1106        let packet_number = incoming.packet.header.number.expand(0);
1107        let InitialHeader {
1108            src_cid,
1109            dst_cid,
1110            version,
1111            ..
1112        } = incoming.packet.header;
1113        let server_config = match server_config.or_else(|| self.server_config.clone()) {
1114            Some(sc) => sc,
1115            None => {
1116                return Err(AcceptError {
1117                    cause: ConnectionError::TransportError(
1118                        crate::transport_error::Error::INTERNAL_ERROR(""),
1119                    ),
1120                    response: None,
1121                });
1122            }
1123        };
1124
1125        if server_config
1126            .transport
1127            .max_idle_timeout
1128            .is_some_and(|timeout| {
1129                incoming.received_at + Duration::from_millis(timeout.into()) <= now
1130            })
1131        {
1132            debug!("abandoning accept of stale initial");
1133            self.index.remove_initial(dst_cid);
1134            return Err(AcceptError {
1135                cause: ConnectionError::TimedOut,
1136                response: None,
1137            });
1138        }
1139
1140        if self.cids_exhausted() {
1141            debug!("refusing connection");
1142            self.index.remove_initial(dst_cid);
1143            return Err(AcceptError {
1144                cause: ConnectionError::CidsExhausted,
1145                response: Some(self.initial_close(
1146                    version,
1147                    incoming.addresses,
1148                    &incoming.crypto,
1149                    &src_cid,
1150                    TransportError::CONNECTION_REFUSED(""),
1151                    buf,
1152                )),
1153            });
1154        }
1155
1156        if incoming
1157            .crypto
1158            .packet
1159            .remote
1160            .decrypt(
1161                packet_number,
1162                &incoming.packet.header_data,
1163                &mut incoming.packet.payload,
1164            )
1165            .is_err()
1166        {
1167            debug!(packet_number, "failed to authenticate initial packet");
1168            self.index.remove_initial(dst_cid);
1169            return Err(AcceptError {
1170                cause: TransportError::PROTOCOL_VIOLATION("authentication failed").into(),
1171                response: None,
1172            });
1173        };
1174
1175        let ch = ConnectionHandle(self.connections.vacant_key());
1176        let loc_cid = self.new_cid(ch);
1177        let mut params = TransportParameters::new(
1178            &server_config.transport,
1179            &self.config,
1180            self.local_cid_generator.as_ref(),
1181            loc_cid,
1182            Some(&server_config),
1183            &mut self.rng,
1184        )?;
1185        params.stateless_reset_token = Some(ResetToken::new(&*self.config.reset_key, loc_cid));
1186        params.original_dst_cid = Some(incoming.token.orig_dst_cid);
1187        params.retry_src_cid = incoming.token.retry_src_cid;
1188        let mut pref_addr_cid = None;
1189        if server_config.has_preferred_address() {
1190            let cid = self.new_cid(ch);
1191            pref_addr_cid = Some(cid);
1192            params.preferred_address = Some(PreferredAddress {
1193                address_v4: server_config.preferred_address_v4,
1194                address_v6: server_config.preferred_address_v6,
1195                connection_id: cid,
1196                stateless_reset_token: ResetToken::new(&*self.config.reset_key, cid),
1197            });
1198        }
1199
1200        let tls = server_config.crypto.clone().start_session(version, &params);
1201        let transport_config = server_config.transport.clone();
1202        let mut conn = self.add_connection(
1203            ch,
1204            version,
1205            dst_cid,
1206            loc_cid,
1207            src_cid,
1208            incoming.addresses,
1209            incoming.received_at,
1210            tls,
1211            transport_config,
1212            SideArgs::Server {
1213                server_config,
1214                pref_addr_cid,
1215                path_validated: remote_address_validated,
1216            },
1217        );
1218        self.index.insert_initial(dst_cid, ch);
1219
1220        match conn.handle_first_packet(
1221            incoming.received_at,
1222            incoming.addresses.remote,
1223            incoming.ecn,
1224            packet_number,
1225            incoming.packet,
1226            incoming.rest,
1227        ) {
1228            Ok(()) => {
1229                trace!(id = ch.0, icid = %dst_cid, "new connection");
1230
1231                for event in incoming_buffer.datagrams {
1232                    conn.handle_event(ConnectionEvent(ConnectionEventInner::Datagram(event)))
1233                }
1234
1235                Ok((ch, conn))
1236            }
1237            Err(e) => {
1238                debug!("handshake failed: {}", e);
1239                self.handle_event(ch, EndpointEvent(EndpointEventInner::Drained));
1240                let response = match e {
1241                    ConnectionError::TransportError(ref e) => Some(self.initial_close(
1242                        version,
1243                        incoming.addresses,
1244                        &incoming.crypto,
1245                        &src_cid,
1246                        e.clone(),
1247                        buf,
1248                    )),
1249                    _ => None,
1250                };
1251                Err(AcceptError { cause: e, response })
1252            }
1253        }
1254    }
1255
1256    /// Check if we should refuse a connection attempt regardless of the packet's contents
1257    fn early_validate_first_packet(
1258        &mut self,
1259        header: &ProtectedInitialHeader,
1260    ) -> Result<(), TransportError> {
1261        let Some(config) = &self.server_config else {
1262            return Err(TransportError::INTERNAL_ERROR(""));
1263        };
1264        if self.cids_exhausted() || self.incoming_buffers.len() >= config.max_incoming {
1265            return Err(TransportError::CONNECTION_REFUSED(""));
1266        }
1267
1268        // RFC9000 §7.2 dictates that initial (client-chosen) destination CIDs must be at least 8
1269        // bytes. If this is a Retry packet, then the length must instead match our usual CID
1270        // length. If we ever issue non-Retry address validation tokens via `NEW_TOKEN`, then we'll
1271        // also need to validate CID length for those after decoding the token.
1272        if header.dst_cid.len() < 8
1273            && (header.token_pos.is_empty()
1274                || header.dst_cid.len() != self.local_cid_generator.cid_len())
1275        {
1276            debug!(
1277                "rejecting connection due to invalid DCID length {}",
1278                header.dst_cid.len()
1279            );
1280            return Err(TransportError::PROTOCOL_VIOLATION(
1281                "invalid destination CID length",
1282            ));
1283        }
1284
1285        Ok(())
1286    }
1287
1288    /// Reject this incoming connection attempt
1289    pub fn refuse(&mut self, incoming: Incoming, buf: &mut Vec<u8>) -> Transmit {
1290        self.clean_up_incoming(&incoming);
1291        incoming.improper_drop_warner.dismiss();
1292
1293        self.initial_close(
1294            incoming.packet.header.version,
1295            incoming.addresses,
1296            &incoming.crypto,
1297            &incoming.packet.header.src_cid,
1298            TransportError::CONNECTION_REFUSED(""),
1299            buf,
1300        )
1301    }
1302
1303    /// Respond with a retry packet, requiring the client to retry with address validation
1304    ///
1305    /// Errors if `incoming.may_retry()` is false.
1306    pub fn retry(&mut self, incoming: Incoming, buf: &mut Vec<u8>) -> Result<Transmit, RetryError> {
1307        if !incoming.may_retry() {
1308            return Err(RetryError(Box::new(incoming)));
1309        }
1310
1311        let Some(server_config_arc) = self.server_config.clone() else {
1312            return Err(RetryError(Box::new(incoming)));
1313        };
1314
1315        self.clean_up_incoming(&incoming);
1316        incoming.improper_drop_warner.dismiss();
1317
1318        // First Initial
1319        // The peer will use this as the DCID of its following Initials. Initial DCIDs are
1320        // looked up separately from Handshake/Data DCIDs, so there is no risk of collision
1321        // with established connections. In the unlikely event that a collision occurs
1322        // between two connections in the initial phase, both will fail fast and may be
1323        // retried by the application layer.
1324        let loc_cid = self.local_cid_generator.generate_cid();
1325
1326        let payload = TokenPayload::Retry {
1327            address: incoming.addresses.remote,
1328            orig_dst_cid: incoming.packet.header.dst_cid,
1329            issued: server_config_arc.time_source.now(),
1330        };
1331        let token = Token::new(payload, &mut self.rng).encode(&*server_config_arc.token_key);
1332
1333        let header = Header::Retry {
1334            src_cid: loc_cid,
1335            dst_cid: incoming.packet.header.src_cid,
1336            version: incoming.packet.header.version,
1337        };
1338
1339        let encode = header.encode(buf);
1340        buf.put_slice(&token);
1341        buf.extend_from_slice(&server_config_arc.crypto.retry_tag(
1342            incoming.packet.header.version,
1343            &incoming.packet.header.dst_cid,
1344            buf,
1345        ));
1346        encode.finish(buf, &*incoming.crypto.header.local, None);
1347
1348        Ok(Transmit {
1349            destination: incoming.addresses.remote,
1350            ecn: None,
1351            size: buf.len(),
1352            segment_size: None,
1353            src_ip: incoming.addresses.local_ip,
1354        })
1355    }
1356
1357    /// Ignore this incoming connection attempt, not sending any packet in response
1358    ///
1359    /// Doing this actively, rather than merely dropping the [`Incoming`], is necessary to prevent
1360    /// memory leaks due to state within [`Endpoint`] tracking the incoming connection.
1361    pub fn ignore(&mut self, incoming: Incoming) {
1362        self.clean_up_incoming(&incoming);
1363        incoming.improper_drop_warner.dismiss();
1364    }
1365
1366    /// Clean up endpoint data structures associated with an `Incoming`.
1367    fn clean_up_incoming(&mut self, incoming: &Incoming) {
1368        self.index.remove_initial(incoming.packet.header.dst_cid);
1369        let incoming_buffer = self.incoming_buffers.remove(incoming.incoming_idx);
1370        self.all_incoming_buffers_total_bytes -= incoming_buffer.total_bytes;
1371    }
1372
1373    fn add_connection(
1374        &mut self,
1375        ch: ConnectionHandle,
1376        version: u32,
1377        init_cid: ConnectionId,
1378        loc_cid: ConnectionId,
1379        rem_cid: ConnectionId,
1380        addresses: FourTuple,
1381        now: Instant,
1382        tls: Box<dyn crypto::Session>,
1383        transport_config: Arc<TransportConfig>,
1384        side_args: SideArgs,
1385    ) -> Connection {
1386        let mut rng_seed = [0; 32];
1387        self.rng.fill_bytes(&mut rng_seed);
1388        let side = side_args.side();
1389        let pref_addr_cid = side_args.pref_addr_cid();
1390        let conn = Connection::new(
1391            self.config.clone(),
1392            transport_config,
1393            init_cid,
1394            loc_cid,
1395            rem_cid,
1396            addresses.remote,
1397            addresses.local_ip,
1398            tls,
1399            self.local_cid_generator.as_ref(),
1400            now,
1401            version,
1402            self.allow_mtud,
1403            rng_seed,
1404            side_args,
1405        );
1406
1407        let mut cids_issued = 0;
1408        let mut loc_cids = FxHashMap::default();
1409
1410        loc_cids.insert(cids_issued, loc_cid);
1411        cids_issued += 1;
1412
1413        if let Some(cid) = pref_addr_cid {
1414            debug_assert_eq!(cids_issued, 1, "preferred address cid seq must be 1");
1415            loc_cids.insert(cids_issued, cid);
1416            cids_issued += 1;
1417        }
1418
1419        let id = self.connections.insert(ConnectionMeta {
1420            init_cid,
1421            cids_issued,
1422            loc_cids,
1423            addresses,
1424            side,
1425            reset_token: None,
1426            peer_id: None,
1427        });
1428        debug_assert_eq!(id, ch.0, "connection handle allocation out of sync");
1429
1430        self.index.insert_conn(addresses, loc_cid, ch, side);
1431
1432        conn
1433    }
1434
1435    fn initial_close(
1436        &mut self,
1437        version: u32,
1438        addresses: FourTuple,
1439        crypto: &Keys,
1440        remote_id: &ConnectionId,
1441        reason: TransportError,
1442        buf: &mut Vec<u8>,
1443    ) -> Transmit {
1444        // We don't need to worry about CID collisions in initial closes because the peer
1445        // shouldn't respond, and if it does, and the CID collides, we'll just drop the
1446        // unexpected response.
1447        let local_id = self.local_cid_generator.generate_cid();
1448        let number = PacketNumber::U8(0);
1449        let header = Header::Initial(InitialHeader {
1450            dst_cid: *remote_id,
1451            src_cid: local_id,
1452            number,
1453            token: Bytes::new(),
1454            version,
1455        });
1456
1457        let partial_encode = header.encode(buf);
1458        let max_len =
1459            INITIAL_MTU as usize - partial_encode.header_len - crypto.packet.local.tag_len();
1460        frame::Close::from(reason).encode(buf, max_len);
1461        buf.resize(buf.len() + crypto.packet.local.tag_len(), 0);
1462        partial_encode.finish(buf, &*crypto.header.local, Some((0, &*crypto.packet.local)));
1463        Transmit {
1464            destination: addresses.remote,
1465            ecn: None,
1466            size: buf.len(),
1467            segment_size: None,
1468            src_ip: addresses.local_ip,
1469        }
1470    }
1471
1472    /// Access the configuration used by this endpoint
1473    pub fn config(&self) -> &EndpointConfig {
1474        &self.config
1475    }
1476
1477    /// Enable or disable address discovery for this endpoint
1478    ///
1479    /// Address discovery is enabled by default. When enabled, the endpoint will:
1480    /// - Send OBSERVED_ADDRESS frames to peers to inform them of their reflexive addresses
1481    /// - Process received OBSERVED_ADDRESS frames to learn about its own reflexive addresses
1482    /// - Integrate discovered addresses with NAT traversal for improved connectivity
1483    pub fn enable_address_discovery(&mut self, enabled: bool) {
1484        self.address_discovery_enabled = enabled;
1485        // Note: Existing connections will continue with their current setting.
1486        // New connections will use the updated setting.
1487    }
1488
1489    /// Check if address discovery is enabled
1490    pub fn address_discovery_enabled(&self) -> bool {
1491        self.address_discovery_enabled
1492    }
1493
1494    /// Get all discovered addresses across all connections
1495    ///
1496    /// Returns a list of unique socket addresses that have been observed
1497    /// by remote peers and reported via OBSERVED_ADDRESS frames.
1498    ///
1499    /// Note: This returns an empty vector in the current implementation.
1500    /// Applications should track discovered addresses at the connection level.
1501    pub fn discovered_addresses(&self) -> Vec<SocketAddr> {
1502        // TODO: Implement address tracking at the endpoint level
1503        Vec::new()
1504    }
1505
1506    /// Set a callback to be invoked when an address change is detected
1507    ///
1508    /// The callback receives the old address (if any) and the new address.
1509    /// Only one callback can be set at a time; setting a new callback replaces the previous one.
1510    pub fn set_address_change_callback<F>(&mut self, callback: F)
1511    where
1512        F: Fn(Option<SocketAddr>, SocketAddr) + Send + Sync + 'static,
1513    {
1514        self.address_change_callback = Some(Box::new(callback));
1515    }
1516
1517    /// Clear the address change callback
1518    pub fn clear_address_change_callback(&mut self) {
1519        self.address_change_callback = None;
1520    }
1521
1522    /// Get address discovery statistics
1523    ///
1524    /// Note: This returns default statistics in the current implementation.
1525    /// Applications should track statistics at the connection level.
1526    pub fn address_discovery_stats(&self) -> AddressDiscoveryStats {
1527        // TODO: Implement statistics tracking at the endpoint level
1528        AddressDiscoveryStats::default()
1529    }
1530
1531    /// Number of connections that are currently open
1532    pub fn open_connections(&self) -> usize {
1533        self.connections.len()
1534    }
1535
1536    /// Counter for the number of bytes currently used
1537    /// in the buffers for Initial and 0-RTT messages for pending incoming connections
1538    pub fn incoming_buffer_bytes(&self) -> u64 {
1539        self.all_incoming_buffers_total_bytes
1540    }
1541
1542    #[cfg(test)]
1543    #[allow(dead_code)]
1544    pub(crate) fn known_connections(&self) -> usize {
1545        let x = self.connections.len();
1546        debug_assert_eq!(x, self.index.connection_ids_initial.len());
1547        // Not all connections have known reset tokens
1548        debug_assert!(x >= self.index.connection_reset_tokens.0.len());
1549        // Not all connections have unique remotes, and 0-length CIDs might not be in use.
1550        debug_assert!(x >= self.index.incoming_connection_remotes.len());
1551        debug_assert!(x >= self.index.outgoing_connection_remotes.len());
1552        x
1553    }
1554
1555    #[cfg(test)]
1556    #[allow(dead_code)]
1557    pub(crate) fn known_cids(&self) -> usize {
1558        self.index.connection_ids.len()
1559    }
1560
1561    /// Whether we've used up 3/4 of the available CID space
1562    ///
1563    /// We leave some space unused so that `new_cid` can be relied upon to finish quickly. We don't
1564    /// bother to check when CID longer than 4 bytes are used because 2^40 connections is a lot.
1565    fn cids_exhausted(&self) -> bool {
1566        self.local_cid_generator.cid_len() <= 4
1567            && self.local_cid_generator.cid_len() != 0
1568            && (2usize.pow(self.local_cid_generator.cid_len() as u32 * 8)
1569                - self.index.connection_ids.len())
1570                < 2usize.pow(self.local_cid_generator.cid_len() as u32 * 8 - 2)
1571    }
1572}
1573
1574impl fmt::Debug for Endpoint {
1575    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
1576        fmt.debug_struct("Endpoint")
1577            .field("rng", &self.rng)
1578            .field("index", &self.index)
1579            .field("connections", &self.connections)
1580            .field("config", &self.config)
1581            .field("server_config", &self.server_config)
1582            // incoming_buffers too large
1583            .field("incoming_buffers.len", &self.incoming_buffers.len())
1584            .field(
1585                "all_incoming_buffers_total_bytes",
1586                &self.all_incoming_buffers_total_bytes,
1587            )
1588            .finish()
1589    }
1590}
1591
1592/// Buffered Initial and 0-RTT messages for a pending incoming connection
1593#[derive(Default)]
1594struct IncomingBuffer {
1595    datagrams: Vec<DatagramConnectionEvent>,
1596    total_bytes: u64,
1597}
1598
1599/// Part of protocol state incoming datagrams can be routed to
1600#[derive(Copy, Clone, Debug)]
1601enum RouteDatagramTo {
1602    Incoming(usize),
1603    Connection(ConnectionHandle),
1604}
1605
1606/// Maps packets to existing connections
1607#[derive(Default, Debug)]
1608struct ConnectionIndex {
1609    /// Identifies connections based on the initial DCID the peer utilized
1610    ///
1611    /// Uses a standard `HashMap` to protect against hash collision attacks.
1612    ///
1613    /// Used by the server, not the client.
1614    connection_ids_initial: HashMap<ConnectionId, RouteDatagramTo>,
1615    /// Identifies connections based on locally created CIDs
1616    ///
1617    /// Uses a cheaper hash function since keys are locally created
1618    connection_ids: FxHashMap<ConnectionId, ConnectionHandle>,
1619    /// Identifies incoming connections with zero-length CIDs
1620    ///
1621    /// Uses a standard `HashMap` to protect against hash collision attacks.
1622    incoming_connection_remotes: HashMap<FourTuple, ConnectionHandle>,
1623    /// Identifies outgoing connections with zero-length CIDs
1624    ///
1625    /// We don't yet support explicit source addresses for client connections, and zero-length CIDs
1626    /// require a unique four-tuple, so at most one client connection with zero-length local CIDs
1627    /// may be established per remote. We must omit the local address from the key because we don't
1628    /// necessarily know what address we're sending from, and hence receiving at.
1629    ///
1630    /// Uses a standard `HashMap` to protect against hash collision attacks.
1631    outgoing_connection_remotes: HashMap<SocketAddr, ConnectionHandle>,
1632    /// Reset tokens provided by the peer for the CID each connection is currently sending to
1633    ///
1634    /// Incoming stateless resets do not have correct CIDs, so we need this to identify the correct
1635    /// recipient, if any.
1636    connection_reset_tokens: ResetTokenTable,
1637}
1638
1639impl ConnectionIndex {
1640    /// Associate an incoming connection with its initial destination CID
1641    fn insert_initial_incoming(&mut self, dst_cid: ConnectionId, incoming_key: usize) {
1642        if dst_cid.is_empty() {
1643            return;
1644        }
1645        self.connection_ids_initial
1646            .insert(dst_cid, RouteDatagramTo::Incoming(incoming_key));
1647    }
1648
1649    /// Remove an association with an initial destination CID
1650    fn remove_initial(&mut self, dst_cid: ConnectionId) {
1651        if dst_cid.is_empty() {
1652            return;
1653        }
1654        let removed = self.connection_ids_initial.remove(&dst_cid);
1655        debug_assert!(removed.is_some());
1656    }
1657
1658    /// Associate a connection with its initial destination CID
1659    fn insert_initial(&mut self, dst_cid: ConnectionId, connection: ConnectionHandle) {
1660        if dst_cid.is_empty() {
1661            return;
1662        }
1663        self.connection_ids_initial
1664            .insert(dst_cid, RouteDatagramTo::Connection(connection));
1665    }
1666
1667    /// Associate a connection with its first locally-chosen destination CID if used, or otherwise
1668    /// its current 4-tuple
1669    fn insert_conn(
1670        &mut self,
1671        addresses: FourTuple,
1672        dst_cid: ConnectionId,
1673        connection: ConnectionHandle,
1674        side: Side,
1675    ) {
1676        match dst_cid.len() {
1677            0 => match side {
1678                Side::Server => {
1679                    self.incoming_connection_remotes
1680                        .insert(addresses, connection);
1681                }
1682                Side::Client => {
1683                    self.outgoing_connection_remotes
1684                        .insert(addresses.remote, connection);
1685                }
1686            },
1687            _ => {
1688                self.connection_ids.insert(dst_cid, connection);
1689            }
1690        }
1691    }
1692
1693    /// Discard a connection ID
1694    fn retire(&mut self, dst_cid: ConnectionId) {
1695        self.connection_ids.remove(&dst_cid);
1696    }
1697
1698    /// Remove all references to a connection
1699    fn remove(&mut self, conn: &ConnectionMeta) {
1700        if conn.side.is_server() {
1701            self.remove_initial(conn.init_cid);
1702        }
1703        for cid in conn.loc_cids.values() {
1704            self.connection_ids.remove(cid);
1705        }
1706        self.incoming_connection_remotes.remove(&conn.addresses);
1707        self.outgoing_connection_remotes
1708            .remove(&conn.addresses.remote);
1709        if let Some((remote, token)) = conn.reset_token {
1710            self.connection_reset_tokens.remove(remote, token);
1711        }
1712    }
1713
1714    /// Find the existing connection that `datagram` should be routed to, if any
1715    fn get(&self, addresses: &FourTuple, datagram: &PartialDecode) -> Option<RouteDatagramTo> {
1716        let dst_cid = datagram.dst_cid();
1717        let is_empty_cid = dst_cid.is_empty();
1718
1719        // Fast path: Try most common lookup first (non-empty CID)
1720        if !is_empty_cid {
1721            if let Some(&ch) = self.connection_ids.get(dst_cid) {
1722                return Some(RouteDatagramTo::Connection(ch));
1723            }
1724        }
1725
1726        // Initial/0RTT packet lookup
1727        if datagram.is_initial() || datagram.is_0rtt() {
1728            if let Some(&ch) = self.connection_ids_initial.get(dst_cid) {
1729                return Some(ch);
1730            }
1731        }
1732
1733        // Empty CID lookup (less common, do after fast path)
1734        if is_empty_cid {
1735            // Check incoming connections first (servers handle more incoming)
1736            if let Some(&ch) = self.incoming_connection_remotes.get(addresses) {
1737                return Some(RouteDatagramTo::Connection(ch));
1738            }
1739            if let Some(&ch) = self.outgoing_connection_remotes.get(&addresses.remote) {
1740                return Some(RouteDatagramTo::Connection(ch));
1741            }
1742        }
1743
1744        // Stateless reset token lookup (least common, do last)
1745        let data = datagram.data();
1746        if data.len() < RESET_TOKEN_SIZE {
1747            return None;
1748        }
1749        self.connection_reset_tokens
1750            .get(addresses.remote, &data[data.len() - RESET_TOKEN_SIZE..])
1751            .cloned()
1752            .map(RouteDatagramTo::Connection)
1753    }
1754}
1755
1756#[derive(Debug)]
1757pub(crate) struct ConnectionMeta {
1758    init_cid: ConnectionId,
1759    /// Number of local connection IDs that have been issued in NEW_CONNECTION_ID frames.
1760    cids_issued: u64,
1761    loc_cids: FxHashMap<u64, ConnectionId>,
1762    /// Remote/local addresses the connection began with
1763    ///
1764    /// Only needed to support connections with zero-length CIDs, which cannot migrate, so we don't
1765    /// bother keeping it up to date.
1766    addresses: FourTuple,
1767    side: Side,
1768    /// Reset token provided by the peer for the CID we're currently sending to, and the address
1769    /// being sent to
1770    reset_token: Option<(SocketAddr, ResetToken)>,
1771    /// Peer ID for this connection, used for relay functionality
1772    peer_id: Option<PeerId>,
1773}
1774
1775/// Internal identifier for a `Connection` currently associated with an endpoint
1776#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)]
1777pub struct ConnectionHandle(pub usize);
1778
1779impl From<ConnectionHandle> for usize {
1780    fn from(x: ConnectionHandle) -> Self {
1781        x.0
1782    }
1783}
1784
1785impl Index<ConnectionHandle> for Slab<ConnectionMeta> {
1786    type Output = ConnectionMeta;
1787    fn index(&self, ch: ConnectionHandle) -> &ConnectionMeta {
1788        &self[ch.0]
1789    }
1790}
1791
1792impl IndexMut<ConnectionHandle> for Slab<ConnectionMeta> {
1793    fn index_mut(&mut self, ch: ConnectionHandle) -> &mut ConnectionMeta {
1794        &mut self[ch.0]
1795    }
1796}
1797
1798/// Event resulting from processing a single datagram
1799pub enum DatagramEvent {
1800    /// The datagram is redirected to its `Connection`
1801    ConnectionEvent(ConnectionHandle, ConnectionEvent),
1802    /// The datagram may result in starting a new `Connection`
1803    NewConnection(Incoming),
1804    /// Response generated directly by the endpoint
1805    Response(Transmit),
1806}
1807
1808/// An incoming connection for which the server has not yet begun its part of the handshake.
1809pub struct Incoming {
1810    received_at: Instant,
1811    addresses: FourTuple,
1812    ecn: Option<EcnCodepoint>,
1813    packet: InitialPacket,
1814    rest: Option<BytesMut>,
1815    crypto: Keys,
1816    token: IncomingToken,
1817    incoming_idx: usize,
1818    improper_drop_warner: IncomingImproperDropWarner,
1819}
1820
1821impl Incoming {
1822    /// The local IP address which was used when the peer established the connection
1823    ///
1824    /// This has the same behavior as [`Connection::local_ip`].
1825    pub fn local_ip(&self) -> Option<IpAddr> {
1826        self.addresses.local_ip
1827    }
1828
1829    /// The peer's UDP address
1830    pub fn remote_address(&self) -> SocketAddr {
1831        self.addresses.remote
1832    }
1833
1834    /// Whether the socket address that is initiating this connection has been validated
1835    ///
1836    /// This means that the sender of the initial packet has proved that they can receive traffic
1837    /// sent to `self.remote_address()`.
1838    ///
1839    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
1840    /// The inverse is not guaranteed.
1841    pub fn remote_address_validated(&self) -> bool {
1842        self.token.validated
1843    }
1844
1845    /// Whether it is legal to respond with a retry packet
1846    ///
1847    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
1848    /// The inverse is not guaranteed.
1849    pub fn may_retry(&self) -> bool {
1850        self.token.retry_src_cid.is_none()
1851    }
1852
1853    /// The original destination connection ID sent by the client
1854    pub fn orig_dst_cid(&self) -> &ConnectionId {
1855        &self.token.orig_dst_cid
1856    }
1857}
1858
1859impl fmt::Debug for Incoming {
1860    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1861        f.debug_struct("Incoming")
1862            .field("addresses", &self.addresses)
1863            .field("ecn", &self.ecn)
1864            // packet doesn't implement debug
1865            // rest is too big and not meaningful enough
1866            .field("token", &self.token)
1867            .field("incoming_idx", &self.incoming_idx)
1868            // improper drop warner contains no information
1869            .finish_non_exhaustive()
1870    }
1871}
1872
1873struct IncomingImproperDropWarner;
1874
1875impl IncomingImproperDropWarner {
1876    fn dismiss(self) {
1877        mem::forget(self);
1878    }
1879}
1880
1881impl Drop for IncomingImproperDropWarner {
1882    fn drop(&mut self) {
1883        warn!(
1884            "quinn_proto::Incoming dropped without passing to Endpoint::accept/refuse/retry/ignore \
1885               (may cause memory leak and eventual inability to accept new connections)"
1886        );
1887    }
1888}
1889
1890/// Errors in the parameters being used to create a new connection
1891///
1892/// These arise before any I/O has been performed.
1893#[derive(Debug, Error, Clone, PartialEq, Eq)]
1894pub enum ConnectError {
1895    /// The endpoint can no longer create new connections
1896    ///
1897    /// Indicates that a necessary component of the endpoint has been dropped or otherwise disabled.
1898    #[error("endpoint stopping")]
1899    EndpointStopping,
1900    /// The connection could not be created because not enough of the CID space is available
1901    ///
1902    /// Try using longer connection IDs
1903    #[error("CIDs exhausted")]
1904    CidsExhausted,
1905    /// The given server name was malformed
1906    #[error("invalid server name: {0}")]
1907    InvalidServerName(String),
1908    /// The remote [`SocketAddr`] supplied was malformed
1909    ///
1910    /// Examples include attempting to connect to port 0, or using an inappropriate address family.
1911    #[error("invalid remote address: {0}")]
1912    InvalidRemoteAddress(SocketAddr),
1913    /// No default client configuration was set up
1914    ///
1915    /// Use `Endpoint::connect_with` to specify a client configuration.
1916    #[error("no default client config")]
1917    NoDefaultClientConfig,
1918    /// The local endpoint does not support the QUIC version specified in the client configuration
1919    #[error("unsupported QUIC version")]
1920    UnsupportedVersion,
1921    /// A TLS-related error occurred during connection establishment
1922    #[error("TLS error: {0}")]
1923    TlsError(String),
1924}
1925
1926/// Error type for attempting to accept an [`Incoming`]
1927#[derive(Debug)]
1928pub struct AcceptError {
1929    /// Underlying error describing reason for failure
1930    pub cause: ConnectionError,
1931    /// Optional response to transmit back
1932    pub response: Option<Transmit>,
1933}
1934
1935impl From<rustls::Error> for ConnectError {
1936    fn from(error: rustls::Error) -> Self {
1937        ConnectError::TlsError(error.to_string())
1938    }
1939}
1940
1941impl From<crate::transport_error::Error> for AcceptError {
1942    fn from(error: crate::transport_error::Error) -> Self {
1943        Self {
1944            cause: ConnectionError::TransportError(error),
1945            response: None,
1946        }
1947    }
1948}
1949
1950/// Error for attempting to retry an [`Incoming`] which already bears a token from a previous retry
1951#[derive(Debug, Error)]
1952#[error("retry() with validated Incoming")]
1953pub struct RetryError(Box<Incoming>);
1954
1955impl RetryError {
1956    /// Get the [`Incoming`]
1957    pub fn into_incoming(self) -> Incoming {
1958        *self.0
1959    }
1960}
1961
1962/// Reset Tokens which are associated with peer socket addresses
1963///
1964/// The standard `HashMap` is used since both `SocketAddr` and `ResetToken` are
1965/// peer generated and might be usable for hash collision attacks.
1966#[derive(Default, Debug)]
1967struct ResetTokenTable(HashMap<SocketAddr, HashMap<ResetToken, ConnectionHandle>>);
1968
1969impl ResetTokenTable {
1970    fn insert(&mut self, remote: SocketAddr, token: ResetToken, ch: ConnectionHandle) -> bool {
1971        self.0
1972            .entry(remote)
1973            .or_default()
1974            .insert(token, ch)
1975            .is_some()
1976    }
1977
1978    fn remove(&mut self, remote: SocketAddr, token: ResetToken) {
1979        use std::collections::hash_map::Entry;
1980        match self.0.entry(remote) {
1981            Entry::Vacant(_) => {}
1982            Entry::Occupied(mut e) => {
1983                e.get_mut().remove(&token);
1984                if e.get().is_empty() {
1985                    e.remove_entry();
1986                }
1987            }
1988        }
1989    }
1990
1991    fn get(&self, remote: SocketAddr, token: &[u8]) -> Option<&ConnectionHandle> {
1992        let token = ResetToken::from(<[u8; RESET_TOKEN_SIZE]>::try_from(token).ok()?);
1993        self.0.get(&remote)?.get(&token)
1994    }
1995}
1996
1997/// Identifies a connection by the combination of remote and local addresses
1998///
1999/// Including the local ensures good behavior when the host has multiple IP addresses on the same
2000/// subnet and zero-length connection IDs are in use.
2001#[derive(Hash, Eq, PartialEq, Debug, Copy, Clone)]
2002struct FourTuple {
2003    remote: SocketAddr,
2004    // A single socket can only listen on a single port, so no need to store it explicitly
2005    local_ip: Option<IpAddr>,
2006}