ant-quic 0.27.4

// Copyright 2024 Saorsa Labs Ltd.
//
// This Saorsa Network Software is licensed under the General Public License (GPL), version 3.
// Please see the file LICENSE-GPL, or visit <http://www.gnu.org/licenses/> for the full text.
//
// Full details available at https://saorsalabs.com/licenses

use std::{
    collections::VecDeque,
    fmt,
    future::Future,
    io,
    io::IoSliceMut,
    mem,
    net::{SocketAddr, SocketAddrV6},
    pin::Pin,
    str,
    sync::{Arc, Mutex},
    task::{Context, Poll, Waker},
};

#[cfg(not(wasm_browser))]
use super::runtime::default_runtime;
use super::{
    runtime::{AsyncUdpSocket, Runtime},
    udp_transmit,
};
use crate::Instant;
use crate::{
    ClientConfig, ConnectError, ConnectionError, ConnectionHandle, DatagramEvent, EndpointEvent,
    ServerConfig,
};
use bytes::{Bytes, BytesMut};
use pin_project_lite::pin_project;
use quinn_udp::{BATCH_SIZE, RecvMeta};
use rustc_hash::FxHashMap;
#[cfg(all(not(wasm_browser), feature = "network-discovery"))]
use socket2::{Domain, Protocol, Socket, Type};
use tokio::sync::{Notify, futures::Notified, mpsc};
use tracing::{Instrument, Span, debug, error};

use super::{
    ConnectionEvent, IO_LOOP_BOUND, RECV_TIME_BOUND, connection::Connecting,
    work_limiter::WorkLimiter,
};
use crate::{EndpointConfig, VarInt};

fn is_expected_background_io_error(error: &io::Error) -> bool {
    matches!(
        error.kind(),
        io::ErrorKind::AddrNotAvailable
            | io::ErrorKind::ConnectionRefused
            | io::ErrorKind::ConnectionReset
            | io::ErrorKind::HostUnreachable
            | io::ErrorKind::NetworkUnreachable
            | io::ErrorKind::NotConnected
            | io::ErrorKind::TimedOut
    ) || matches!(error.raw_os_error(), Some(49 | 51 | 65 | 101 | 113))
}

/// A QUIC endpoint.
///
/// An endpoint corresponds to a single UDP socket, may host many connections, and may act as both
/// client and server for different connections.
///
/// May be cloned to obtain another handle to the same endpoint.
#[derive(Debug, Clone)]
pub struct Endpoint {
    pub(crate) inner: EndpointRef,
    pub(crate) default_client_config: Option<ClientConfig>,
    runtime: Arc<dyn Runtime>,
}

impl Endpoint {
    /// Helper to construct an endpoint for use with outgoing connections only
    ///
    /// Note that `addr` is the *local* address to bind to, which should usually be a wildcard
    /// address like `0.0.0.0:0` or `[::]:0`, which allow communication with any reachable IPv4 or
    /// IPv6 address respectively from an OS-assigned port.
    ///
    /// If an IPv6 address is provided, attempts to make the socket dual-stack so as to allow
    /// communication with both IPv4 and IPv6 addresses. As such, calling `Endpoint::client` with
    /// the address `[::]:0` is a reasonable default to maximize the ability to connect to other
    /// address. For example:
    ///
    /// ```
    /// # use std::net::{Ipv6Addr, SocketAddr};
    /// # fn example() -> std::io::Result<()> {
    /// use ant_quic::high_level::Endpoint;
    ///
    /// let addr: SocketAddr = (Ipv6Addr::UNSPECIFIED, 0).into();
    /// let endpoint = Endpoint::client(addr)?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// Some environments may not allow creation of dual-stack sockets, in which case an IPv6
    /// client will only be able to connect to IPv6 servers. An IPv4 client is never dual-stack.
    #[cfg(all(not(wasm_browser), feature = "network-discovery"))]
    pub fn client(addr: SocketAddr) -> io::Result<Self> {
        let socket = Socket::new(Domain::for_address(addr), Type::DGRAM, Some(Protocol::UDP))?;
        if addr.is_ipv6() {
            if let Err(e) = socket.set_only_v6(false) {
                tracing::debug!(%e, "unable to make socket dual-stack");
            }
        }

        // Apply platform-appropriate buffer sizes to avoid WSAEMSGSIZE errors on Windows
        // and ensure reliable QUIC connections, especially with PQC
        use crate::config::buffer_defaults;
        let buffer_size = buffer_defaults::PLATFORM_DEFAULT;
        if let Err(e) = socket.set_send_buffer_size(buffer_size) {
            tracing::debug!(%e, "unable to set send buffer size to {}", buffer_size);
        }
        if let Err(e) = socket.set_recv_buffer_size(buffer_size) {
            tracing::debug!(%e, "unable to set recv buffer size to {}", buffer_size);
        }

        socket.bind(&addr.into())?;
        let runtime =
            default_runtime().ok_or_else(|| io::Error::other("no async runtime found"))?;
        Self::new_with_abstract_socket(
            EndpointConfig::default(),
            None,
            runtime.wrap_udp_socket(socket.into())?,
            runtime,
        )
    }

    /// Helper to construct an endpoint for use with outgoing connections only
    /// (fallback without network-discovery feature)
    #[cfg(all(not(wasm_browser), not(feature = "network-discovery")))]
    pub fn client(addr: SocketAddr) -> io::Result<Self> {
        let socket = std::net::UdpSocket::bind(addr)?;
        let runtime =
            default_runtime().ok_or_else(|| io::Error::other("no async runtime found"))?;
        Self::new_with_abstract_socket(
            EndpointConfig::default(),
            None,
            runtime.wrap_udp_socket(socket)?,
            runtime,
        )
    }

    /// Returns relevant stats from this Endpoint
    pub fn stats(&self) -> EndpointStats {
        self.inner
            .state
            .lock()
            .map(|state| state.stats)
            .unwrap_or_else(|_| {
                error!("Endpoint state mutex poisoned");
                EndpointStats::default()
            })
    }

    /// Helper to construct an endpoint for use with both incoming and outgoing connections
    ///
    /// When binding to an IPv6 address, this creates a dual-stack socket (IPV6_V6ONLY=0)
    /// that can accept both IPv4 and IPv6 connections. IPv4 connections will appear as
    /// IPv4-mapped IPv6 addresses (::ffff:x.x.x.x).
    ///
    /// Platform defaults for dual-stack sockets vary. For example, any socket bound to a wildcard
    /// IPv6 address on Windows will not by default be able to communicate with IPv4
    /// addresses. This method explicitly enables dual-stack for IPv6 sockets.
    #[cfg(all(not(wasm_browser), feature = "network-discovery"))]
    pub fn server(config: ServerConfig, addr: SocketAddr) -> io::Result<Self> {
        let socket = Socket::new(Domain::for_address(addr), Type::DGRAM, Some(Protocol::UDP))?;

        // Enable dual-stack for IPv6 sockets (consistent with client() behavior)
        if addr.is_ipv6() {
            if let Err(e) = socket.set_only_v6(false) {
                tracing::debug!(%e, "unable to make server socket dual-stack");
            }
        }

        socket.set_nonblocking(true)?;

        // Apply platform-appropriate buffer sizes to avoid WSAEMSGSIZE errors on Windows
        // and ensure reliable QUIC connections, especially with PQC
        use crate::config::buffer_defaults;
        let buffer_size = buffer_defaults::PLATFORM_DEFAULT;
        if let Err(e) = socket.set_send_buffer_size(buffer_size) {
            tracing::debug!(%e, "unable to set send buffer size to {}", buffer_size);
        }
        if let Err(e) = socket.set_recv_buffer_size(buffer_size) {
            tracing::debug!(%e, "unable to set recv buffer size to {}", buffer_size);
        }

        socket.bind(&addr.into())?;
        let runtime =
            default_runtime().ok_or_else(|| io::Error::other("no async runtime found"))?;
        Self::new_with_abstract_socket(
            EndpointConfig::default(),
            Some(config),
            runtime.wrap_udp_socket(socket.into())?,
            runtime,
        )
    }

    /// Helper to construct an endpoint for use with both incoming and outgoing connections
    /// (fallback without network-discovery feature)
    #[cfg(all(not(wasm_browser), not(feature = "network-discovery")))]
    pub fn server(config: ServerConfig, addr: SocketAddr) -> io::Result<Self> {
        let socket = std::net::UdpSocket::bind(addr)?;
        let runtime =
            default_runtime().ok_or_else(|| io::Error::other("no async runtime found"))?;
        Self::new_with_abstract_socket(
            EndpointConfig::default(),
            Some(config),
            runtime.wrap_udp_socket(socket)?,
            runtime,
        )
    }

    /// Construct an endpoint with arbitrary configuration and socket
    #[cfg(not(wasm_browser))]
    pub fn new(
        config: EndpointConfig,
        server_config: Option<ServerConfig>,
        socket: std::net::UdpSocket,
        runtime: Arc<dyn Runtime>,
    ) -> io::Result<Self> {
        let socket = runtime.wrap_udp_socket(socket)?;
        Self::new_with_abstract_socket(config, server_config, socket, runtime)
    }

    /// Construct an endpoint with arbitrary configuration and pre-constructed abstract socket
    ///
    /// Useful when `socket` has additional state (e.g. sidechannels) attached for which shared
    /// ownership is needed.
    pub fn new_with_abstract_socket(
        config: EndpointConfig,
        server_config: Option<ServerConfig>,
        socket: Arc<dyn AsyncUdpSocket>,
        runtime: Arc<dyn Runtime>,
    ) -> io::Result<Self> {
        let addr = socket.local_addr()?;
        let allow_mtud = !socket.may_fragment();
        let rc = EndpointRef::new(
            socket,
            crate::endpoint::Endpoint::new(
                Arc::new(config),
                server_config.map(Arc::new),
                allow_mtud,
                None,
            ),
            addr.is_ipv6(),
            runtime.clone(),
        );
        let driver = EndpointDriver(rc.clone());
        runtime.spawn(Box::pin(
            async {
                if let Err(e) = driver.await {
                    if is_expected_background_io_error(&e) {
                        debug!("background I/O path ended: {}", e);
                    } else {
                        tracing::error!("I/O error: {}", e);
                    }
                }
            }
            .instrument(Span::current()),
        ));
        Ok(Self {
            inner: rc,
            default_client_config: None,
            runtime,
        })
    }

    /// Get the next incoming connection attempt from a client
    ///
    /// Yields `Incoming`s, or `None` if the endpoint is [`close`](Self::close)d. `Incoming`
    /// can be `await`ed to obtain the final [`Connection`](crate::Connection), or used to e.g.
    /// filter connection attempts or force address validation, or converted into an intermediate
    /// `Connecting` future which can be used to e.g. send 0.5-RTT data.
    pub fn accept(&self) -> Accept<'_> {
        Accept {
            endpoint: self,
            notify: self.inner.shared.incoming.notified(),
        }
    }

    /// Set the client configuration used by `connect`
    pub fn set_default_client_config(&mut self, config: ClientConfig) {
        self.default_client_config = Some(config);
    }

    /// Connect to a remote endpoint
    ///
    /// `server_name` must be covered by the certificate presented by the server. This prevents a
    /// connection from being intercepted by an attacker with a valid certificate for some other
    /// server.
    ///
    /// May fail immediately due to configuration errors, or in the future if the connection could
    /// not be established.
    pub fn connect(&self, addr: SocketAddr, server_name: &str) -> Result<Connecting, ConnectError> {
        let config = match &self.default_client_config {
            Some(config) => config.clone(),
            None => return Err(ConnectError::NoDefaultClientConfig),
        };

        self.connect_with(config, addr, server_name)
    }

    /// Connect to a remote endpoint using a custom configuration.
    ///
    /// See [`connect()`] for details.
    ///
    /// [`connect()`]: Endpoint::connect
    pub fn connect_with(
        &self,
        config: ClientConfig,
        addr: SocketAddr,
        server_name: &str,
    ) -> Result<Connecting, ConnectError> {
        let mut endpoint = self
            .inner
            .state
            .lock()
            .map_err(|_| ConnectError::EndpointStopping)?;
        if endpoint.driver_lost || endpoint.recv_state.connections.close.is_some() {
            return Err(ConnectError::EndpointStopping);
        }
        if addr.is_ipv6() && !endpoint.ipv6 {
            return Err(ConnectError::InvalidRemoteAddress(addr));
        }
        let addr = if endpoint.ipv6 {
            SocketAddr::V6(ensure_ipv6(addr))
        } else {
            addr
        };

        let (ch, conn) = endpoint
            .inner
            .connect(self.runtime.now(), config, addr, server_name)?;

        let socket = endpoint.socket.clone();
        endpoint.stats.outgoing_handshakes += 1;
        Ok(endpoint
            .recv_state
            .connections
            .insert(ch, conn, socket, self.runtime.clone()))
    }

    /// Switch to a new UDP socket
    ///
    /// See [`Endpoint::rebind_abstract()`] for details.
    #[cfg(not(wasm_browser))]
    pub fn rebind(&self, socket: std::net::UdpSocket) -> io::Result<()> {
        self.rebind_abstract(self.runtime.wrap_udp_socket(socket)?)
    }

    /// Switch to a new UDP socket
    ///
    /// Allows the endpoint's address to be updated live, affecting all active connections. Incoming
    /// connections and connections to servers unreachable from the new address will be lost.
    ///
    /// On error, the old UDP socket is retained.
    pub fn rebind_abstract(&self, socket: Arc<dyn AsyncUdpSocket>) -> io::Result<()> {
        let addr = socket.local_addr()?;
        let mut inner = self
            .inner
            .state
            .lock()
            .map_err(|_| io::Error::other("Endpoint state mutex poisoned"))?;
        inner.prev_socket = Some(mem::replace(&mut inner.socket, socket));
        inner.ipv6 = addr.is_ipv6();

        // Update connection socket references
        for sender in inner.recv_state.connections.senders.values() {
            // Ignoring errors from dropped connections
            let _ = sender.send(ConnectionEvent::Rebind(inner.socket.clone()));
        }

        Ok(())
    }

    /// Replace the server configuration, affecting new incoming connections only
    ///
    /// Useful for e.g. refreshing TLS certificates without disrupting existing connections.
    pub fn set_server_config(&self, server_config: Option<ServerConfig>) {
        if let Ok(mut state) = self.inner.state.lock() {
            state.inner.set_server_config(server_config.map(Arc::new));
        } else {
            error!("Failed to set server config: endpoint state mutex poisoned");
        }
    }

    /// Register a peer ID for an existing connection, enabling PUNCH_ME_NOW relay
    /// routing by peer identity instead of socket address.
    pub fn register_connection_peer_id(
        &self,
        addr: SocketAddr,
        peer_id: crate::nat_traversal_api::PeerId,
    ) {
        if let Ok(mut state) = self.inner.state.lock() {
            let handle = state.inner.connection_handle_for_addr(&addr);
            if let Some(ch) = handle {
                state.inner.set_connection_peer_id(ch, peer_id);
                tracing::info!(
                    "Registered peer ID {} for connection {} at low-level endpoint",
                    hex::encode(&peer_id.0[..8]),
                    addr
                );
            } else {
                tracing::debug!(
                    "No connection handle found for {} — peer ID not registered",
                    addr
                );
            }
        }
    }

    /// Set the channel for forwarding peer address updates to the upper layer.
    pub fn set_peer_address_update_tx(&self, tx: mpsc::UnboundedSender<(SocketAddr, SocketAddr)>) {
        if let Ok(mut state) = self.inner.state.lock() {
            state.peer_address_update_tx = Some(tx);
        }
    }

    /// Get the remote address of a peer's connection by peer ID.
    pub fn peer_connection_addr_by_id(&self, peer_id: &[u8; 32]) -> Option<SocketAddr> {
        let state = self.inner.state.lock().ok()?;
        let pid = crate::nat_traversal_api::PeerId(*peer_id);
        state.inner.peer_connection_addr(&pid)
    }

    /// Get the local `SocketAddr` the underlying socket is bound to
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.inner
            .state
            .lock()
            .map_err(|_| io::Error::other("Endpoint state mutex poisoned"))?
            .socket
            .local_addr()
    }

    /// Get the number of connections that are currently open
    pub fn open_connections(&self) -> usize {
        self.inner
            .state
            .lock()
            .map(|state| state.inner.open_connections())
            .unwrap_or(0)
    }

    /// Close all of this endpoint's connections immediately and cease accepting new connections.
    ///
    /// See [`Connection::close()`] for details.
    ///
    /// [`Connection::close()`]: crate::Connection::close
    pub fn close(&self, error_code: VarInt, reason: &[u8]) {
        let reason = Bytes::copy_from_slice(reason);
        let mut endpoint = match self.inner.state.lock() {
            Ok(endpoint) => endpoint,
            Err(_) => {
                error!("Failed to close endpoint: state mutex poisoned");
                return;
            }
        };
        endpoint.recv_state.connections.close = Some((error_code, reason.clone()));
        for sender in endpoint.recv_state.connections.senders.values() {
            // Ignoring errors from dropped connections
            let _ = sender.send(ConnectionEvent::Close {
                error_code,
                reason: reason.clone(),
            });
        }
        self.inner.shared.incoming.notify_waiters();
    }

    /// Wait for all connections on the endpoint to be cleanly shut down
    ///
    /// Waiting for this condition before exiting ensures that a good-faith effort is made to notify
    /// peers of recent connection closes, whereas exiting immediately could force them to wait out
    /// the idle timeout period.
    ///
    /// Does not proactively close existing connections or cause incoming connections to be
    /// rejected. Consider calling [`close()`] if that is desired.
    ///
    /// [`close()`]: Endpoint::close
    pub async fn wait_idle(&self) {
        loop {
            {
                let endpoint = match self.inner.state.lock() {
                    Ok(endpoint) => endpoint,
                    Err(_) => {
                        error!("Failed to wait for idle: state mutex poisoned");
                        break;
                    }
                };
                if endpoint.recv_state.connections.is_empty() {
                    break;
                }
                // Construct future while lock is held to avoid race
                self.inner.shared.idle.notified()
            }
            .await;
        }
    }
}

/// Statistics on [Endpoint] activity
#[non_exhaustive]
#[derive(Debug, Default, Copy, Clone)]
pub struct EndpointStats {
    /// Cummulative number of Quic handshakes accepted by this [Endpoint]
    pub accepted_handshakes: u64,
    /// Cummulative number of Quic handshakees sent from this [Endpoint]
    pub outgoing_handshakes: u64,
    /// Cummulative number of Quic handshakes refused on this [Endpoint]
    pub refused_handshakes: u64,
    /// Cummulative number of Quic handshakes ignored on this [Endpoint]
    pub ignored_handshakes: u64,
}

/// A future that drives IO on an endpoint
///
/// This task functions as the switch point between the UDP socket object and the
/// `Endpoint` responsible for routing datagrams to their owning `Connection`.
/// In order to do so, it also facilitates the exchange of different types of events
/// flowing between the `Endpoint` and the tasks managing `Connection`s. As such,
/// running this task is necessary to keep the endpoint's connections running.
///
/// `EndpointDriver` futures terminate when all clones of the `Endpoint` have been dropped, or when
/// an I/O error occurs.
#[must_use = "endpoint drivers must be spawned for I/O to occur"]
#[derive(Debug)]
pub(crate) struct EndpointDriver(pub(crate) EndpointRef);

impl Future for EndpointDriver {
    type Output = Result<(), io::Error>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let mut endpoint = match self.0.state.lock() {
            Ok(endpoint) => endpoint,
            Err(_) => {
                return Poll::Ready(Err(io::Error::other("Endpoint state mutex poisoned")));
            }
        };
        if endpoint.driver.is_none() {
            endpoint.driver = Some(cx.waker().clone());
        }

        let now = endpoint.runtime.now();
        let mut keep_going = false;
        keep_going |= endpoint.drive_recv(cx, now)?;
        keep_going |= endpoint.handle_events(cx, &self.0.shared);

        if !endpoint.recv_state.incoming.is_empty() {
            self.0.shared.incoming.notify_waiters();
        }

        if endpoint.ref_count == 0 && endpoint.recv_state.connections.is_empty() {
            Poll::Ready(Ok(()))
        } else {
            drop(endpoint);
            // If there is more work to do schedule the endpoint task again.
            // `wake_by_ref()` is called outside the lock to minimize
            // lock contention on a multithreaded runtime.
            if keep_going {
                cx.waker().wake_by_ref();
            }
            Poll::Pending
        }
    }
}

impl Drop for EndpointDriver {
    fn drop(&mut self) {
        if let Ok(mut endpoint) = self.0.state.lock() {
            endpoint.driver_lost = true;
            self.0.shared.incoming.notify_waiters();
            // Drop all outgoing channels, signaling the termination of the endpoint to the associated
            // connections.
            endpoint.recv_state.connections.senders.clear();
        } else {
            error!("Failed to lock endpoint state in drop - mutex poisoned");
        }
    }
}

#[derive(Debug)]
pub(crate) struct EndpointInner {
    pub(crate) state: Mutex<State>,
    pub(crate) shared: Shared,
}

impl EndpointInner {
    pub(crate) fn accept(
        &self,
        incoming: crate::Incoming,
        server_config: Option<Arc<ServerConfig>>,
    ) -> Result<Connecting, ConnectionError> {
        let mut state = self.state.lock().map_err(|_| {
            ConnectionError::TransportError(crate::transport_error::Error::INTERNAL_ERROR(
                "Endpoint state mutex poisoned".to_string(),
            ))
        })?;
        let mut response_buffer = Vec::new();
        let now = state.runtime.now();
        match state
            .inner
            .accept(incoming, now, &mut response_buffer, server_config)
        {
            Ok((handle, conn)) => {
                state.stats.accepted_handshakes += 1;
                let socket = state.socket.clone();
                let runtime = state.runtime.clone();
                Ok(state
                    .recv_state
                    .connections
                    .insert(handle, conn, socket, runtime))
            }
            Err(error) => {
                if let Some(transmit) = error.response {
                    respond(transmit, &response_buffer, &*state.socket);
                }
                Err(error.cause)
            }
        }
    }

    pub(crate) fn refuse(&self, incoming: crate::Incoming) {
        let mut state = match self.state.lock() {
            Ok(state) => state,
            Err(_) => {
                error!("Failed to refuse connection: endpoint state mutex poisoned");
                return;
            }
        };
        state.stats.refused_handshakes += 1;
        let mut response_buffer = Vec::new();
        let transmit = state.inner.refuse(incoming, &mut response_buffer);
        respond(transmit, &response_buffer, &*state.socket);
    }

    pub(crate) fn retry(
        &self,
        incoming: crate::Incoming,
    ) -> Result<(), crate::endpoint::RetryError> {
        let mut state = match self.state.lock() {
            Ok(state) => state,
            Err(_) => {
                error!("Failed to retry connection: endpoint state mutex poisoned");
                return Err(crate::endpoint::RetryError::incoming(incoming));
            }
        };
        let mut response_buffer = Vec::new();
        let transmit = state.inner.retry(incoming, &mut response_buffer)?;
        respond(transmit, &response_buffer, &*state.socket);
        Ok(())
    }

    pub(crate) fn ignore(&self, incoming: crate::Incoming) {
        if let Ok(mut state) = self.state.lock() {
            state.stats.ignored_handshakes += 1;
            state.inner.ignore(incoming);
        } else {
            error!("Failed to ignore incoming connection: endpoint state mutex poisoned");
        }
    }
}

#[derive(Debug)]
pub(crate) struct State {
    socket: Arc<dyn AsyncUdpSocket>,
    /// During an active migration, abandoned_socket receives traffic
    /// until the first packet arrives on the new socket.
    prev_socket: Option<Arc<dyn AsyncUdpSocket>>,
    inner: crate::endpoint::Endpoint,
    recv_state: RecvState,
    driver: Option<Waker>,
    ipv6: bool,
    events: mpsc::UnboundedReceiver<(ConnectionHandle, EndpointEvent)>,
    /// Number of live handles that can be used to initiate or handle I/O; excludes the driver
    ref_count: usize,
    driver_lost: bool,
    runtime: Arc<dyn Runtime>,
    stats: EndpointStats,
    /// Channel for forwarding peer address updates to the upper layer.
    peer_address_update_tx: Option<mpsc::UnboundedSender<(SocketAddr, SocketAddr)>>,
}

#[derive(Debug)]
pub(crate) struct Shared {
    incoming: Notify,
    idle: Notify,
}

impl State {
    fn drive_recv(&mut self, cx: &mut Context, now: Instant) -> Result<bool, io::Error> {
        let get_time = || self.runtime.now();
        self.recv_state.recv_limiter.start_cycle(get_time);
        if let Some(socket) = &self.prev_socket {
            // We don't care about the `PollProgress` from old sockets.
            let poll_res =
                self.recv_state
                    .poll_socket(cx, &mut self.inner, &**socket, &*self.runtime, now);
            if poll_res.is_err() {
                self.prev_socket = None;
            }
        };
        let poll_res =
            self.recv_state
                .poll_socket(cx, &mut self.inner, &*self.socket, &*self.runtime, now);
        self.recv_state.recv_limiter.finish_cycle(get_time);
        let poll_res = poll_res?;
        if poll_res.received_connection_packet {
            // Traffic has arrived on self.socket, therefore there is no need for the abandoned
            // one anymore. TODO: Account for multiple outgoing connections.
            self.prev_socket = None;
        }
        Ok(poll_res.keep_going)
    }

    fn handle_events(&mut self, cx: &mut Context, shared: &Shared) -> bool {
        let mut did_work = false;

        for _ in 0..IO_LOOP_BOUND {
            let (ch, event) = match self.events.poll_recv(cx) {
                Poll::Ready(Some(x)) => x,
                Poll::Ready(None) => unreachable!("EndpointInner owns one sender"),
                Poll::Pending => {
                    break;
                }
            };

            did_work = true;

            if event.is_drained() {
                self.recv_state.connections.senders.remove(&ch);
                if self.recv_state.connections.is_empty() {
                    shared.idle.notify_waiters();
                }
            }
            let Some(event) = self.inner.handle_event(ch, event) else {
                continue;
            };
            // Ignoring errors from dropped connections that haven't yet been cleaned up
            if let Some(sender) = self.recv_state.connections.senders.get_mut(&ch) {
                let _ = sender.send(ConnectionEvent::Proto(event));
            }
        }

        // Process relay events generated by the endpoint
        // These are PUNCH_ME_NOW frames that need to be forwarded to target connections
        for (ch, event) in self.inner.drain_relay_events() {
            did_work = true;
            if let Some(sender) = self.recv_state.connections.senders.get_mut(&ch) {
                tracing::debug!("Sending relay event to connection {:?}", ch);
                let _ = sender.send(ConnectionEvent::Proto(event));
            } else {
                tracing::warn!(
                    "Cannot send relay event: connection {:?} not found in senders",
                    ch
                );
            }
        }

        // Forward peer address updates from ADD_ADDRESS frames to the
        // NatTraversalEndpoint so it can update the DHT routing table.
        let address_updates: Vec<(SocketAddr, SocketAddr)> =
            self.inner.drain_peer_address_updates().collect();
        for (peer_addr, advertised_addr) in address_updates {
            did_work = true;
            if let Some(ref tx) = self.peer_address_update_tx {
                let _ = tx.send((peer_addr, advertised_addr));
            }
        }

        did_work
    }
}

impl Drop for State {
    fn drop(&mut self) {
        for incoming in self.recv_state.incoming.drain(..) {
            self.inner.ignore(incoming);
        }
    }
}

fn respond(transmit: crate::Transmit, response_buffer: &[u8], socket: &dyn AsyncUdpSocket) {
    // Send if there's kernel buffer space; otherwise, drop it
    //
    // As an endpoint-generated packet, we know this is an
    // immediate, stateless response to an unconnected peer,
    // one of:
    //
    // - A version negotiation response due to an unknown version
    // - A `CLOSE` due to a malformed or unwanted connection attempt
    // - A stateless reset due to an unrecognized connection
    // - A `Retry` packet due to a connection attempt when
    //   `use_retry` is set
    //
    // In each case, a well-behaved peer can be trusted to retry a
    // few times, which is guaranteed to produce the same response
    // from us. Repeated failures might at worst cause a peer's new
    // connection attempt to time out, which is acceptable if we're
    // under such heavy load that there's never room for this code
    // to transmit. This is morally equivalent to the packet getting
    // lost due to congestion further along the link, which
    // similarly relies on peer retries for recovery.
    _ = socket.try_send(&udp_transmit(&transmit, &response_buffer[..transmit.size]));
}

#[inline]
fn proto_ecn(ecn: quinn_udp::EcnCodepoint) -> crate::EcnCodepoint {
    match ecn {
        quinn_udp::EcnCodepoint::Ect0 => crate::EcnCodepoint::Ect0,
        quinn_udp::EcnCodepoint::Ect1 => crate::EcnCodepoint::Ect1,
        quinn_udp::EcnCodepoint::Ce => crate::EcnCodepoint::Ce,
    }
}

#[derive(Debug)]
struct ConnectionSet {
    /// Senders for communicating with the endpoint's connections
    senders: FxHashMap<ConnectionHandle, mpsc::UnboundedSender<ConnectionEvent>>,
    /// Stored to give out clones to new ConnectionInners
    sender: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>,
    /// Set if the endpoint has been manually closed
    close: Option<(VarInt, Bytes)>,
}

impl ConnectionSet {
    fn insert(
        &mut self,
        handle: ConnectionHandle,
        conn: crate::Connection,
        socket: Arc<dyn AsyncUdpSocket>,
        runtime: Arc<dyn Runtime>,
    ) -> Connecting {
        let (send, recv) = mpsc::unbounded_channel();
        if let Some((error_code, ref reason)) = self.close {
            let _ = send.send(ConnectionEvent::Close {
                error_code,
                reason: reason.clone(),
            });
        }
        self.senders.insert(handle, send);
        Connecting::new(handle, conn, self.sender.clone(), recv, socket, runtime)
    }

    fn is_empty(&self) -> bool {
        self.senders.is_empty()
    }
}

fn ensure_ipv6(x: SocketAddr) -> SocketAddrV6 {
    match x {
        SocketAddr::V6(x) => x,
        SocketAddr::V4(x) => SocketAddrV6::new(x.ip().to_ipv6_mapped(), x.port(), 0, 0),
    }
}

pin_project! {
    /// Future produced by [`Endpoint::accept`]
    pub struct Accept<'a> {
        endpoint: &'a Endpoint,
        #[pin]
        notify: Notified<'a>,
    }
}

impl Future for Accept<'_> {
    type Output = Option<super::incoming::Incoming>;
    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();
        let mut endpoint = match this.endpoint.inner.state.lock() {
            Ok(endpoint) => endpoint,
            Err(_) => return Poll::Ready(None),
        };
        if endpoint.driver_lost {
            return Poll::Ready(None);
        }
        if let Some(incoming) = endpoint.recv_state.incoming.pop_front() {
            // Release the mutex lock on endpoint so cloning it doesn't deadlock
            drop(endpoint);
            let incoming = super::incoming::Incoming::new(incoming, this.endpoint.inner.clone());
            return Poll::Ready(Some(incoming));
        }
        if endpoint.recv_state.connections.close.is_some() {
            return Poll::Ready(None);
        }
        loop {
            match this.notify.as_mut().poll(ctx) {
                // `state` lock ensures we didn't race with readiness
                Poll::Pending => return Poll::Pending,
                // Spurious wakeup, get a new future
                Poll::Ready(()) => this
                    .notify
                    .set(this.endpoint.inner.shared.incoming.notified()),
            }
        }
    }
}

#[derive(Debug)]
pub(crate) struct EndpointRef(Arc<EndpointInner>);

impl EndpointRef {
    pub(crate) fn new(
        socket: Arc<dyn AsyncUdpSocket>,
        inner: crate::endpoint::Endpoint,
        ipv6: bool,
        runtime: Arc<dyn Runtime>,
    ) -> Self {
        let (sender, events) = mpsc::unbounded_channel();
        let recv_state = RecvState::new(sender, socket.max_receive_segments(), &inner);
        Self(Arc::new(EndpointInner {
            shared: Shared {
                incoming: Notify::new(),
                idle: Notify::new(),
            },
            state: Mutex::new(State {
                socket,
                prev_socket: None,
                inner,
                ipv6,
                events,
                driver: None,
                ref_count: 0,
                driver_lost: false,
                recv_state,
                runtime,
                stats: EndpointStats::default(),
                peer_address_update_tx: None,
            }),
        }))
    }
}

impl Clone for EndpointRef {
    fn clone(&self) -> Self {
        if let Ok(mut state) = self.0.state.lock() {
            state.ref_count += 1;
        }
        Self(self.0.clone())
    }
}

impl Drop for EndpointRef {
    fn drop(&mut self) {
        if let Ok(mut endpoint) = self.0.state.lock() {
            if let Some(x) = endpoint.ref_count.checked_sub(1) {
                endpoint.ref_count = x;
                if x == 0 {
                    // If the driver is about to be on its own, ensure it can shut down if the last
                    // connection is gone.
                    if let Some(task) = endpoint.driver.take() {
                        task.wake();
                    }
                }
            }
        } else {
            error!("Failed to drop EndpointRef: state mutex poisoned");
        }
    }
}

impl std::ops::Deref for EndpointRef {
    type Target = EndpointInner;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// State directly involved in handling incoming packets
struct RecvState {
    incoming: VecDeque<crate::Incoming>,
    connections: ConnectionSet,
    recv_buf: Box<[u8]>,
    recv_limiter: WorkLimiter,
}

impl RecvState {
    fn new(
        sender: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>,
        max_receive_segments: usize,
        endpoint: &crate::endpoint::Endpoint,
    ) -> Self {
        // Use a receive buffer size large enough to handle any incoming packet.
        // This is especially important for PQC handshakes which can send 4096+ byte datagrams
        // before transport parameters are exchanged. We use the maximum of:
        // - The configured max_udp_payload_size (what we expect to receive)
        // - PQC minimum MTU (4096 bytes) to handle PQC handshakes regardless of config
        // - Capped at 64KB for practical memory usage
        const PQC_MIN_RECV_SIZE: u64 = 4096;
        let configured_size = endpoint.config().get_max_udp_payload_size();
        let effective_size = configured_size.max(PQC_MIN_RECV_SIZE).min(64 * 1024) as usize;

        let recv_buf = vec![0; effective_size * max_receive_segments * BATCH_SIZE];
        Self {
            connections: ConnectionSet {
                senders: FxHashMap::default(),
                sender,
                close: None,
            },
            incoming: VecDeque::new(),
            recv_buf: recv_buf.into(),
            recv_limiter: WorkLimiter::new(RECV_TIME_BOUND),
        }
    }

    fn poll_socket(
        &mut self,
        cx: &mut Context,
        endpoint: &mut crate::endpoint::Endpoint,
        socket: &dyn AsyncUdpSocket,
        runtime: &dyn Runtime,
        now: Instant,
    ) -> Result<PollProgress, io::Error> {
        let mut received_connection_packet = false;
        let mut metas = [RecvMeta::default(); BATCH_SIZE];
        let mut iovs: [IoSliceMut; BATCH_SIZE] = {
            let mut bufs = self
                .recv_buf
                .chunks_mut(self.recv_buf.len() / BATCH_SIZE)
                .map(IoSliceMut::new);

            // expect() safe as self.recv_buf is chunked into BATCH_SIZE items
            // and iovs will be of size BATCH_SIZE, thus from_fn is called
            // exactly BATCH_SIZE times.
            std::array::from_fn(|_| {
                bufs.next().unwrap_or_else(|| {
                    error!("Insufficient buffers for BATCH_SIZE");
                    IoSliceMut::new(&mut [])
                })
            })
        };
        loop {
            match socket.poll_recv(cx, &mut iovs, &mut metas) {
                Poll::Ready(Ok(msgs)) => {
                    self.recv_limiter.record_work(msgs);
                    for (meta, buf) in metas.iter().zip(iovs.iter()).take(msgs) {
                        let mut data: BytesMut = buf[0..meta.len].into();
                        while !data.is_empty() {
                            let buf = data.split_to(meta.stride.min(data.len()));
                            let mut response_buffer = Vec::new();
                            match endpoint.handle(
                                now,
                                meta.addr,
                                meta.dst_ip,
                                meta.ecn.map(proto_ecn),
                                buf,
                                &mut response_buffer,
                            ) {
                                Some(DatagramEvent::NewConnection(incoming)) => {
                                    if self.connections.close.is_none() {
                                        self.incoming.push_back(incoming);
                                    } else {
                                        let transmit =
                                            endpoint.refuse(incoming, &mut response_buffer);
                                        respond(transmit, &response_buffer, socket);
                                    }
                                }
                                Some(DatagramEvent::ConnectionEvent(handle, event)) => {
                                    // Ignoring errors from dropped connections that haven't yet been cleaned up
                                    received_connection_packet = true;
                                    if let Some(sender) = self.connections.senders.get_mut(&handle)
                                    {
                                        let _ = sender.send(ConnectionEvent::Proto(event));
                                    }
                                }
                                Some(DatagramEvent::Response(transmit)) => {
                                    respond(transmit, &response_buffer, socket);
                                }
                                None => {}
                            }
                        }
                    }
                }
                Poll::Pending => {
                    return Ok(PollProgress {
                        received_connection_packet,
                        keep_going: false,
                    });
                }
                // Ignore ECONNRESET as it's undefined in QUIC and may be injected by an
                // attacker
                Poll::Ready(Err(ref e)) if e.kind() == io::ErrorKind::ConnectionReset => {
                    continue;
                }
                Poll::Ready(Err(e)) => {
                    return Err(e);
                }
            }
            if !self.recv_limiter.allow_work(|| runtime.now()) {
                return Ok(PollProgress {
                    received_connection_packet,
                    keep_going: true,
                });
            }
        }
    }
}

impl fmt::Debug for RecvState {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("RecvState")
            .field("incoming", &self.incoming)
            .field("connections", &self.connections)
            // recv_buf too large
            .field("recv_limiter", &self.recv_limiter)
            .finish_non_exhaustive()
    }
}

#[derive(Default)]
struct PollProgress {
    /// Whether a datagram was routed to an existing connection
    received_connection_packet: bool,
    /// Whether datagram handling was interrupted early by the work limiter for fairness
    keep_going: bool,
}