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        conn.set_expected_peer_id_from_token(incoming.token.peer_id);
1219        self.index.insert_initial(dst_cid, ch);
1220
1221        match conn.handle_first_packet(
1222            incoming.received_at,
1223            incoming.addresses.remote,
1224            incoming.ecn,
1225            packet_number,
1226            incoming.packet,
1227            incoming.rest,
1228        ) {
1229            Ok(()) => {
1230                trace!(id = ch.0, icid = %dst_cid, "new connection");
1231
1232                for event in incoming_buffer.datagrams {
1233                    conn.handle_event(ConnectionEvent(ConnectionEventInner::Datagram(event)))
1234                }
1235
1236                Ok((ch, conn))
1237            }
1238            Err(e) => {
1239                debug!("handshake failed: {}", e);
1240                self.handle_event(ch, EndpointEvent(EndpointEventInner::Drained));
1241                let response = match e {
1242                    ConnectionError::TransportError(ref e) => Some(self.initial_close(
1243                        version,
1244                        incoming.addresses,
1245                        &incoming.crypto,
1246                        &src_cid,
1247                        e.clone(),
1248                        buf,
1249                    )),
1250                    _ => None,
1251                };
1252                Err(AcceptError { cause: e, response })
1253            }
1254        }
1255    }
1256
1257    /// Check if we should refuse a connection attempt regardless of the packet's contents
1258    fn early_validate_first_packet(
1259        &mut self,
1260        header: &ProtectedInitialHeader,
1261    ) -> Result<(), TransportError> {
1262        let Some(config) = &self.server_config else {
1263            return Err(TransportError::INTERNAL_ERROR(""));
1264        };
1265        if self.cids_exhausted() || self.incoming_buffers.len() >= config.max_incoming {
1266            return Err(TransportError::CONNECTION_REFUSED(""));
1267        }
1268
1269        // RFC9000 §7.2 dictates that initial (client-chosen) destination CIDs must be at least 8
1270        // bytes. If this is a Retry packet, then the length must instead match our usual CID
1271        // length. If we ever issue non-Retry address validation tokens via `NEW_TOKEN`, then we'll
1272        // also need to validate CID length for those after decoding the token.
1273        if header.dst_cid.len() < 8
1274            && (header.token_pos.is_empty()
1275                || header.dst_cid.len() != self.local_cid_generator.cid_len())
1276        {
1277            debug!(
1278                "rejecting connection due to invalid DCID length {}",
1279                header.dst_cid.len()
1280            );
1281            return Err(TransportError::PROTOCOL_VIOLATION(
1282                "invalid destination CID length",
1283            ));
1284        }
1285
1286        Ok(())
1287    }
1288
1289    /// Reject this incoming connection attempt
1290    pub fn refuse(&mut self, incoming: Incoming, buf: &mut Vec<u8>) -> Transmit {
1291        self.clean_up_incoming(&incoming);
1292        incoming.improper_drop_warner.dismiss();
1293
1294        self.initial_close(
1295            incoming.packet.header.version,
1296            incoming.addresses,
1297            &incoming.crypto,
1298            &incoming.packet.header.src_cid,
1299            TransportError::CONNECTION_REFUSED(""),
1300            buf,
1301        )
1302    }
1303
1304    /// Respond with a retry packet, requiring the client to retry with address validation
1305    ///
1306    /// Errors if `incoming.may_retry()` is false.
1307    pub fn retry(&mut self, incoming: Incoming, buf: &mut Vec<u8>) -> Result<Transmit, RetryError> {
1308        if !incoming.may_retry() {
1309            return Err(RetryError(Box::new(incoming)));
1310        }
1311
1312        let Some(server_config_arc) = self.server_config.clone() else {
1313            return Err(RetryError(Box::new(incoming)));
1314        };
1315
1316        self.clean_up_incoming(&incoming);
1317        incoming.improper_drop_warner.dismiss();
1318
1319        // First Initial
1320        // The peer will use this as the DCID of its following Initials. Initial DCIDs are
1321        // looked up separately from Handshake/Data DCIDs, so there is no risk of collision
1322        // with established connections. In the unlikely event that a collision occurs
1323        // between two connections in the initial phase, both will fail fast and may be
1324        // retried by the application layer.
1325        let loc_cid = self.local_cid_generator.generate_cid();
1326
1327        let payload = TokenPayload::Retry {
1328            address: incoming.addresses.remote,
1329            orig_dst_cid: incoming.packet.header.dst_cid,
1330            issued: server_config_arc.time_source.now(),
1331        };
1332        let token = Token::new(payload, &mut self.rng).encode(&*server_config_arc.token_key);
1333
1334        let header = Header::Retry {
1335            src_cid: loc_cid,
1336            dst_cid: incoming.packet.header.src_cid,
1337            version: incoming.packet.header.version,
1338        };
1339
1340        let encode = header.encode(buf);
1341        buf.put_slice(&token);
1342        buf.extend_from_slice(&server_config_arc.crypto.retry_tag(
1343            incoming.packet.header.version,
1344            &incoming.packet.header.dst_cid,
1345            buf,
1346        ));
1347        encode.finish(buf, &*incoming.crypto.header.local, None);
1348
1349        Ok(Transmit {
1350            destination: incoming.addresses.remote,
1351            ecn: None,
1352            size: buf.len(),
1353            segment_size: None,
1354            src_ip: incoming.addresses.local_ip,
1355        })
1356    }
1357
1358    /// Ignore this incoming connection attempt, not sending any packet in response
1359    ///
1360    /// Doing this actively, rather than merely dropping the [`Incoming`], is necessary to prevent
1361    /// memory leaks due to state within [`Endpoint`] tracking the incoming connection.
1362    pub fn ignore(&mut self, incoming: Incoming) {
1363        self.clean_up_incoming(&incoming);
1364        incoming.improper_drop_warner.dismiss();
1365    }
1366
1367    /// Clean up endpoint data structures associated with an `Incoming`.
1368    fn clean_up_incoming(&mut self, incoming: &Incoming) {
1369        self.index.remove_initial(incoming.packet.header.dst_cid);
1370        let incoming_buffer = self.incoming_buffers.remove(incoming.incoming_idx);
1371        self.all_incoming_buffers_total_bytes -= incoming_buffer.total_bytes;
1372    }
1373
1374    fn add_connection(
1375        &mut self,
1376        ch: ConnectionHandle,
1377        version: u32,
1378        init_cid: ConnectionId,
1379        loc_cid: ConnectionId,
1380        rem_cid: ConnectionId,
1381        addresses: FourTuple,
1382        now: Instant,
1383        tls: Box<dyn crypto::Session>,
1384        transport_config: Arc<TransportConfig>,
1385        side_args: SideArgs,
1386    ) -> Connection {
1387        let mut rng_seed = [0; 32];
1388        self.rng.fill_bytes(&mut rng_seed);
1389        let side = side_args.side();
1390        let pref_addr_cid = side_args.pref_addr_cid();
1391        let conn = Connection::new(
1392            self.config.clone(),
1393            transport_config,
1394            init_cid,
1395            loc_cid,
1396            rem_cid,
1397            addresses.remote,
1398            addresses.local_ip,
1399            tls,
1400            self.local_cid_generator.as_ref(),
1401            now,
1402            version,
1403            self.allow_mtud,
1404            rng_seed,
1405            side_args,
1406        );
1407
1408        let mut cids_issued = 0;
1409        let mut loc_cids = FxHashMap::default();
1410
1411        loc_cids.insert(cids_issued, loc_cid);
1412        cids_issued += 1;
1413
1414        if let Some(cid) = pref_addr_cid {
1415            debug_assert_eq!(cids_issued, 1, "preferred address cid seq must be 1");
1416            loc_cids.insert(cids_issued, cid);
1417            cids_issued += 1;
1418        }
1419
1420        let id = self.connections.insert(ConnectionMeta {
1421            init_cid,
1422            cids_issued,
1423            loc_cids,
1424            addresses,
1425            side,
1426            reset_token: None,
1427            peer_id: None,
1428        });
1429        debug_assert_eq!(id, ch.0, "connection handle allocation out of sync");
1430
1431        self.index.insert_conn(addresses, loc_cid, ch, side);
1432
1433        conn
1434    }
1435
1436    fn initial_close(
1437        &mut self,
1438        version: u32,
1439        addresses: FourTuple,
1440        crypto: &Keys,
1441        remote_id: &ConnectionId,
1442        reason: TransportError,
1443        buf: &mut Vec<u8>,
1444    ) -> Transmit {
1445        // We don't need to worry about CID collisions in initial closes because the peer
1446        // shouldn't respond, and if it does, and the CID collides, we'll just drop the
1447        // unexpected response.
1448        let local_id = self.local_cid_generator.generate_cid();
1449        let number = PacketNumber::U8(0);
1450        let header = Header::Initial(InitialHeader {
1451            dst_cid: *remote_id,
1452            src_cid: local_id,
1453            number,
1454            token: Bytes::new(),
1455            version,
1456        });
1457
1458        let partial_encode = header.encode(buf);
1459        let max_len =
1460            INITIAL_MTU as usize - partial_encode.header_len - crypto.packet.local.tag_len();
1461        frame::Close::from(reason).encode(buf, max_len);
1462        buf.resize(buf.len() + crypto.packet.local.tag_len(), 0);
1463        partial_encode.finish(buf, &*crypto.header.local, Some((0, &*crypto.packet.local)));
1464        Transmit {
1465            destination: addresses.remote,
1466            ecn: None,
1467            size: buf.len(),
1468            segment_size: None,
1469            src_ip: addresses.local_ip,
1470        }
1471    }
1472
1473    /// Access the configuration used by this endpoint
1474    pub fn config(&self) -> &EndpointConfig {
1475        &self.config
1476    }
1477
1478    /// Enable or disable address discovery for this endpoint
1479    ///
1480    /// Address discovery is enabled by default. When enabled, the endpoint will:
1481    /// - Send OBSERVED_ADDRESS frames to peers to inform them of their reflexive addresses
1482    /// - Process received OBSERVED_ADDRESS frames to learn about its own reflexive addresses
1483    /// - Integrate discovered addresses with NAT traversal for improved connectivity
1484    pub fn enable_address_discovery(&mut self, enabled: bool) {
1485        self.address_discovery_enabled = enabled;
1486        // Note: Existing connections will continue with their current setting.
1487        // New connections will use the updated setting.
1488    }
1489
1490    /// Check if address discovery is enabled
1491    pub fn address_discovery_enabled(&self) -> bool {
1492        self.address_discovery_enabled
1493    }
1494
1495    /// Get all discovered addresses across all connections
1496    ///
1497    /// Returns a list of unique socket addresses that have been observed
1498    /// by remote peers and reported via OBSERVED_ADDRESS frames.
1499    ///
1500    /// Note: This returns an empty vector in the current implementation.
1501    /// Applications should track discovered addresses at the connection level.
1502    pub fn discovered_addresses(&self) -> Vec<SocketAddr> {
1503        // TODO: Implement address tracking at the endpoint level
1504        Vec::new()
1505    }
1506
1507    /// Set a callback to be invoked when an address change is detected
1508    ///
1509    /// The callback receives the old address (if any) and the new address.
1510    /// Only one callback can be set at a time; setting a new callback replaces the previous one.
1511    pub fn set_address_change_callback<F>(&mut self, callback: F)
1512    where
1513        F: Fn(Option<SocketAddr>, SocketAddr) + Send + Sync + 'static,
1514    {
1515        self.address_change_callback = Some(Box::new(callback));
1516    }
1517
1518    /// Clear the address change callback
1519    pub fn clear_address_change_callback(&mut self) {
1520        self.address_change_callback = None;
1521    }
1522
1523    /// Get address discovery statistics
1524    ///
1525    /// Note: This returns default statistics in the current implementation.
1526    /// Applications should track statistics at the connection level.
1527    pub fn address_discovery_stats(&self) -> AddressDiscoveryStats {
1528        // TODO: Implement statistics tracking at the endpoint level
1529        AddressDiscoveryStats::default()
1530    }
1531
1532    /// Number of connections that are currently open
1533    pub fn open_connections(&self) -> usize {
1534        self.connections.len()
1535    }
1536
1537    /// Counter for the number of bytes currently used
1538    /// in the buffers for Initial and 0-RTT messages for pending incoming connections
1539    pub fn incoming_buffer_bytes(&self) -> u64 {
1540        self.all_incoming_buffers_total_bytes
1541    }
1542
1543    #[cfg(test)]
1544    #[allow(dead_code)]
1545    pub(crate) fn known_connections(&self) -> usize {
1546        let x = self.connections.len();
1547        debug_assert_eq!(x, self.index.connection_ids_initial.len());
1548        // Not all connections have known reset tokens
1549        debug_assert!(x >= self.index.connection_reset_tokens.0.len());
1550        // Not all connections have unique remotes, and 0-length CIDs might not be in use.
1551        debug_assert!(x >= self.index.incoming_connection_remotes.len());
1552        debug_assert!(x >= self.index.outgoing_connection_remotes.len());
1553        x
1554    }
1555
1556    #[cfg(test)]
1557    #[allow(dead_code)]
1558    pub(crate) fn known_cids(&self) -> usize {
1559        self.index.connection_ids.len()
1560    }
1561
1562    /// Whether we've used up 3/4 of the available CID space
1563    ///
1564    /// We leave some space unused so that `new_cid` can be relied upon to finish quickly. We don't
1565    /// bother to check when CID longer than 4 bytes are used because 2^40 connections is a lot.
1566    fn cids_exhausted(&self) -> bool {
1567        self.local_cid_generator.cid_len() <= 4
1568            && self.local_cid_generator.cid_len() != 0
1569            && (2usize.pow(self.local_cid_generator.cid_len() as u32 * 8)
1570                - self.index.connection_ids.len())
1571                < 2usize.pow(self.local_cid_generator.cid_len() as u32 * 8 - 2)
1572    }
1573}
1574
1575impl fmt::Debug for Endpoint {
1576    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
1577        fmt.debug_struct("Endpoint")
1578            .field("rng", &self.rng)
1579            .field("index", &self.index)
1580            .field("connections", &self.connections)
1581            .field("config", &self.config)
1582            .field("server_config", &self.server_config)
1583            // incoming_buffers too large
1584            .field("incoming_buffers.len", &self.incoming_buffers.len())
1585            .field(
1586                "all_incoming_buffers_total_bytes",
1587                &self.all_incoming_buffers_total_bytes,
1588            )
1589            .finish()
1590    }
1591}
1592
1593/// Buffered Initial and 0-RTT messages for a pending incoming connection
1594#[derive(Default)]
1595struct IncomingBuffer {
1596    datagrams: Vec<DatagramConnectionEvent>,
1597    total_bytes: u64,
1598}
1599
1600/// Part of protocol state incoming datagrams can be routed to
1601#[derive(Copy, Clone, Debug)]
1602enum RouteDatagramTo {
1603    Incoming(usize),
1604    Connection(ConnectionHandle),
1605}
1606
1607/// Maps packets to existing connections
1608#[derive(Default, Debug)]
1609struct ConnectionIndex {
1610    /// Identifies connections based on the initial DCID the peer utilized
1611    ///
1612    /// Uses a standard `HashMap` to protect against hash collision attacks.
1613    ///
1614    /// Used by the server, not the client.
1615    connection_ids_initial: HashMap<ConnectionId, RouteDatagramTo>,
1616    /// Identifies connections based on locally created CIDs
1617    ///
1618    /// Uses a cheaper hash function since keys are locally created
1619    connection_ids: FxHashMap<ConnectionId, ConnectionHandle>,
1620    /// Identifies incoming connections with zero-length CIDs
1621    ///
1622    /// Uses a standard `HashMap` to protect against hash collision attacks.
1623    incoming_connection_remotes: HashMap<FourTuple, ConnectionHandle>,
1624    /// Identifies outgoing connections with zero-length CIDs
1625    ///
1626    /// We don't yet support explicit source addresses for client connections, and zero-length CIDs
1627    /// require a unique four-tuple, so at most one client connection with zero-length local CIDs
1628    /// may be established per remote. We must omit the local address from the key because we don't
1629    /// necessarily know what address we're sending from, and hence receiving at.
1630    ///
1631    /// Uses a standard `HashMap` to protect against hash collision attacks.
1632    outgoing_connection_remotes: HashMap<SocketAddr, ConnectionHandle>,
1633    /// Reset tokens provided by the peer for the CID each connection is currently sending to
1634    ///
1635    /// Incoming stateless resets do not have correct CIDs, so we need this to identify the correct
1636    /// recipient, if any.
1637    connection_reset_tokens: ResetTokenTable,
1638}
1639
1640impl ConnectionIndex {
1641    /// Associate an incoming connection with its initial destination CID
1642    fn insert_initial_incoming(&mut self, dst_cid: ConnectionId, incoming_key: usize) {
1643        if dst_cid.is_empty() {
1644            return;
1645        }
1646        self.connection_ids_initial
1647            .insert(dst_cid, RouteDatagramTo::Incoming(incoming_key));
1648    }
1649
1650    /// Remove an association with an initial destination CID
1651    fn remove_initial(&mut self, dst_cid: ConnectionId) {
1652        if dst_cid.is_empty() {
1653            return;
1654        }
1655        let removed = self.connection_ids_initial.remove(&dst_cid);
1656        debug_assert!(removed.is_some());
1657    }
1658
1659    /// Associate a connection with its initial destination CID
1660    fn insert_initial(&mut self, dst_cid: ConnectionId, connection: ConnectionHandle) {
1661        if dst_cid.is_empty() {
1662            return;
1663        }
1664        self.connection_ids_initial
1665            .insert(dst_cid, RouteDatagramTo::Connection(connection));
1666    }
1667
1668    /// Associate a connection with its first locally-chosen destination CID if used, or otherwise
1669    /// its current 4-tuple
1670    fn insert_conn(
1671        &mut self,
1672        addresses: FourTuple,
1673        dst_cid: ConnectionId,
1674        connection: ConnectionHandle,
1675        side: Side,
1676    ) {
1677        match dst_cid.len() {
1678            0 => match side {
1679                Side::Server => {
1680                    self.incoming_connection_remotes
1681                        .insert(addresses, connection);
1682                }
1683                Side::Client => {
1684                    self.outgoing_connection_remotes
1685                        .insert(addresses.remote, connection);
1686                }
1687            },
1688            _ => {
1689                self.connection_ids.insert(dst_cid, connection);
1690            }
1691        }
1692    }
1693
1694    /// Discard a connection ID
1695    fn retire(&mut self, dst_cid: ConnectionId) {
1696        self.connection_ids.remove(&dst_cid);
1697    }
1698
1699    /// Remove all references to a connection
1700    fn remove(&mut self, conn: &ConnectionMeta) {
1701        if conn.side.is_server() {
1702            self.remove_initial(conn.init_cid);
1703        }
1704        for cid in conn.loc_cids.values() {
1705            self.connection_ids.remove(cid);
1706        }
1707        self.incoming_connection_remotes.remove(&conn.addresses);
1708        self.outgoing_connection_remotes
1709            .remove(&conn.addresses.remote);
1710        if let Some((remote, token)) = conn.reset_token {
1711            self.connection_reset_tokens.remove(remote, token);
1712        }
1713    }
1714
1715    /// Find the existing connection that `datagram` should be routed to, if any
1716    fn get(&self, addresses: &FourTuple, datagram: &PartialDecode) -> Option<RouteDatagramTo> {
1717        let dst_cid = datagram.dst_cid();
1718        let is_empty_cid = dst_cid.is_empty();
1719
1720        // Fast path: Try most common lookup first (non-empty CID)
1721        if !is_empty_cid {
1722            if let Some(&ch) = self.connection_ids.get(dst_cid) {
1723                return Some(RouteDatagramTo::Connection(ch));
1724            }
1725        }
1726
1727        // Initial/0RTT packet lookup
1728        if datagram.is_initial() || datagram.is_0rtt() {
1729            if let Some(&ch) = self.connection_ids_initial.get(dst_cid) {
1730                return Some(ch);
1731            }
1732        }
1733
1734        // Empty CID lookup (less common, do after fast path)
1735        if is_empty_cid {
1736            // Check incoming connections first (servers handle more incoming)
1737            if let Some(&ch) = self.incoming_connection_remotes.get(addresses) {
1738                return Some(RouteDatagramTo::Connection(ch));
1739            }
1740            if let Some(&ch) = self.outgoing_connection_remotes.get(&addresses.remote) {
1741                return Some(RouteDatagramTo::Connection(ch));
1742            }
1743        }
1744
1745        // Stateless reset token lookup (least common, do last)
1746        let data = datagram.data();
1747        if data.len() < RESET_TOKEN_SIZE {
1748            return None;
1749        }
1750        self.connection_reset_tokens
1751            .get(addresses.remote, &data[data.len() - RESET_TOKEN_SIZE..])
1752            .cloned()
1753            .map(RouteDatagramTo::Connection)
1754    }
1755}
1756
1757#[derive(Debug)]
1758pub(crate) struct ConnectionMeta {
1759    init_cid: ConnectionId,
1760    /// Number of local connection IDs that have been issued in NEW_CONNECTION_ID frames.
1761    cids_issued: u64,
1762    loc_cids: FxHashMap<u64, ConnectionId>,
1763    /// Remote/local addresses the connection began with
1764    ///
1765    /// Only needed to support connections with zero-length CIDs, which cannot migrate, so we don't
1766    /// bother keeping it up to date.
1767    addresses: FourTuple,
1768    side: Side,
1769    /// Reset token provided by the peer for the CID we're currently sending to, and the address
1770    /// being sent to
1771    reset_token: Option<(SocketAddr, ResetToken)>,
1772    /// Peer ID for this connection, used for relay functionality
1773    peer_id: Option<PeerId>,
1774}
1775
1776/// Internal identifier for a `Connection` currently associated with an endpoint
1777#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)]
1778pub struct ConnectionHandle(pub usize);
1779
1780impl From<ConnectionHandle> for usize {
1781    fn from(x: ConnectionHandle) -> Self {
1782        x.0
1783    }
1784}
1785
1786impl Index<ConnectionHandle> for Slab<ConnectionMeta> {
1787    type Output = ConnectionMeta;
1788    fn index(&self, ch: ConnectionHandle) -> &ConnectionMeta {
1789        &self[ch.0]
1790    }
1791}
1792
1793impl IndexMut<ConnectionHandle> for Slab<ConnectionMeta> {
1794    fn index_mut(&mut self, ch: ConnectionHandle) -> &mut ConnectionMeta {
1795        &mut self[ch.0]
1796    }
1797}
1798
1799/// Event resulting from processing a single datagram
1800pub enum DatagramEvent {
1801    /// The datagram is redirected to its `Connection`
1802    ConnectionEvent(ConnectionHandle, ConnectionEvent),
1803    /// The datagram may result in starting a new `Connection`
1804    NewConnection(Incoming),
1805    /// Response generated directly by the endpoint
1806    Response(Transmit),
1807}
1808
1809/// An incoming connection for which the server has not yet begun its part of the handshake.
1810pub struct Incoming {
1811    received_at: Instant,
1812    addresses: FourTuple,
1813    ecn: Option<EcnCodepoint>,
1814    packet: InitialPacket,
1815    rest: Option<BytesMut>,
1816    crypto: Keys,
1817    token: IncomingToken,
1818    incoming_idx: usize,
1819    improper_drop_warner: IncomingImproperDropWarner,
1820}
1821
1822impl Incoming {
1823    /// The local IP address which was used when the peer established the connection
1824    ///
1825    /// This has the same behavior as [`Connection::local_ip`].
1826    pub fn local_ip(&self) -> Option<IpAddr> {
1827        self.addresses.local_ip
1828    }
1829
1830    /// The peer's UDP address
1831    pub fn remote_address(&self) -> SocketAddr {
1832        self.addresses.remote
1833    }
1834
1835    /// Whether the socket address that is initiating this connection has been validated
1836    ///
1837    /// This means that the sender of the initial packet has proved that they can receive traffic
1838    /// sent to `self.remote_address()`.
1839    ///
1840    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
1841    /// The inverse is not guaranteed.
1842    pub fn remote_address_validated(&self) -> bool {
1843        self.token.validated
1844    }
1845
1846    /// Whether it is legal to respond with a retry packet
1847    ///
1848    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
1849    /// The inverse is not guaranteed.
1850    pub fn may_retry(&self) -> bool {
1851        self.token.retry_src_cid.is_none()
1852    }
1853
1854    /// The original destination connection ID sent by the client
1855    pub fn orig_dst_cid(&self) -> &ConnectionId {
1856        &self.token.orig_dst_cid
1857    }
1858}
1859
1860impl fmt::Debug for Incoming {
1861    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1862        f.debug_struct("Incoming")
1863            .field("addresses", &self.addresses)
1864            .field("ecn", &self.ecn)
1865            // packet doesn't implement debug
1866            // rest is too big and not meaningful enough
1867            .field("token", &self.token)
1868            .field("incoming_idx", &self.incoming_idx)
1869            // improper drop warner contains no information
1870            .finish_non_exhaustive()
1871    }
1872}
1873
1874struct IncomingImproperDropWarner;
1875
1876impl IncomingImproperDropWarner {
1877    fn dismiss(self) {
1878        mem::forget(self);
1879    }
1880}
1881
1882impl Drop for IncomingImproperDropWarner {
1883    fn drop(&mut self) {
1884        warn!(
1885            "quinn_proto::Incoming dropped without passing to Endpoint::accept/refuse/retry/ignore \
1886               (may cause memory leak and eventual inability to accept new connections)"
1887        );
1888    }
1889}
1890
1891/// Errors in the parameters being used to create a new connection
1892///
1893/// These arise before any I/O has been performed.
1894#[derive(Debug, Error, Clone, PartialEq, Eq)]
1895pub enum ConnectError {
1896    /// The endpoint can no longer create new connections
1897    ///
1898    /// Indicates that a necessary component of the endpoint has been dropped or otherwise disabled.
1899    #[error("endpoint stopping")]
1900    EndpointStopping,
1901    /// The connection could not be created because not enough of the CID space is available
1902    ///
1903    /// Try using longer connection IDs
1904    #[error("CIDs exhausted")]
1905    CidsExhausted,
1906    /// The given server name was malformed
1907    #[error("invalid server name: {0}")]
1908    InvalidServerName(String),
1909    /// The remote [`SocketAddr`] supplied was malformed
1910    ///
1911    /// Examples include attempting to connect to port 0, or using an inappropriate address family.
1912    #[error("invalid remote address: {0}")]
1913    InvalidRemoteAddress(SocketAddr),
1914    /// No default client configuration was set up
1915    ///
1916    /// Use `Endpoint::connect_with` to specify a client configuration.
1917    #[error("no default client config")]
1918    NoDefaultClientConfig,
1919    /// The local endpoint does not support the QUIC version specified in the client configuration
1920    #[error("unsupported QUIC version")]
1921    UnsupportedVersion,
1922    /// A TLS-related error occurred during connection establishment
1923    #[error("TLS error: {0}")]
1924    TlsError(String),
1925}
1926
1927/// Error type for attempting to accept an [`Incoming`]
1928#[derive(Debug)]
1929pub struct AcceptError {
1930    /// Underlying error describing reason for failure
1931    pub cause: ConnectionError,
1932    /// Optional response to transmit back
1933    pub response: Option<Transmit>,
1934}
1935
1936impl From<rustls::Error> for ConnectError {
1937    fn from(error: rustls::Error) -> Self {
1938        ConnectError::TlsError(error.to_string())
1939    }
1940}
1941
1942impl From<crate::transport_error::Error> for AcceptError {
1943    fn from(error: crate::transport_error::Error) -> Self {
1944        Self {
1945            cause: ConnectionError::TransportError(error),
1946            response: None,
1947        }
1948    }
1949}
1950
1951/// Error for attempting to retry an [`Incoming`] which already bears a token from a previous retry
1952#[derive(Debug, Error)]
1953#[error("retry() with validated Incoming")]
1954pub struct RetryError(Box<Incoming>);
1955
1956impl RetryError {
1957    /// Get the [`Incoming`]
1958    pub fn into_incoming(self) -> Incoming {
1959        *self.0
1960    }
1961}
1962
1963/// Reset Tokens which are associated with peer socket addresses
1964///
1965/// The standard `HashMap` is used since both `SocketAddr` and `ResetToken` are
1966/// peer generated and might be usable for hash collision attacks.
1967#[derive(Default, Debug)]
1968struct ResetTokenTable(HashMap<SocketAddr, HashMap<ResetToken, ConnectionHandle>>);
1969
1970impl ResetTokenTable {
1971    fn insert(&mut self, remote: SocketAddr, token: ResetToken, ch: ConnectionHandle) -> bool {
1972        self.0
1973            .entry(remote)
1974            .or_default()
1975            .insert(token, ch)
1976            .is_some()
1977    }
1978
1979    fn remove(&mut self, remote: SocketAddr, token: ResetToken) {
1980        use std::collections::hash_map::Entry;
1981        match self.0.entry(remote) {
1982            Entry::Vacant(_) => {}
1983            Entry::Occupied(mut e) => {
1984                e.get_mut().remove(&token);
1985                if e.get().is_empty() {
1986                    e.remove_entry();
1987                }
1988            }
1989        }
1990    }
1991
1992    fn get(&self, remote: SocketAddr, token: &[u8]) -> Option<&ConnectionHandle> {
1993        let token = ResetToken::from(<[u8; RESET_TOKEN_SIZE]>::try_from(token).ok()?);
1994        self.0.get(&remote)?.get(&token)
1995    }
1996}
1997
1998/// Identifies a connection by the combination of remote and local addresses
1999///
2000/// Including the local ensures good behavior when the host has multiple IP addresses on the same
2001/// subnet and zero-length connection IDs are in use.
2002#[derive(Hash, Eq, PartialEq, Debug, Copy, Clone)]
2003struct FourTuple {
2004    remote: SocketAddr,
2005    // A single socket can only listen on a single port, so no need to store it explicitly
2006    local_ip: Option<IpAddr>,
2007}