noq/

connection.rs

use std::{
    any::Any,
    fmt,
    future::Future,
    io,
    net::{IpAddr, SocketAddr},
    num::NonZeroUsize,
    pin::Pin,
    sync::{
        Arc, Weak,
        atomic::{AtomicUsize, Ordering},
    },
    task::{Context, Poll, Waker, ready},
};

use bytes::Bytes;
use pin_project_lite::pin_project;
use rustc_hash::FxHashMap;
use thiserror::Error;
use tokio::sync::{Notify, futures::Notified, mpsc, oneshot, watch};
use tracing::{Instrument, Span, debug_span};

use crate::{
    ConnectionEvent, Duration, Instant, Path, VarInt,
    endpoint::ensure_ipv6,
    mutex::{Mutex, MutexGuard},
    path::{OpenPath, PathRef, PathRefOwner},
    recv_stream::RecvStream,
    runtime::{AsyncTimer, Runtime, UdpSender},
    send_stream::SendStream,
    udp_transmit,
};
use proto::{
    ConnectionError, ConnectionHandle, ConnectionStats, Dir, EndpointEvent, FourTuple, PathError,
    PathEvent, PathId, PathStats, PathStatus, Side, StreamEvent, StreamId, TransportError,
    TransportErrorCode, congestion::Controller, n0_nat_traversal,
};

/// In-progress connection attempt future
#[derive(Debug)]
pub struct Connecting {
    conn: Option<ConnectionRef>,
    connected: oneshot::Receiver<bool>,
    handshake_data_ready: Option<oneshot::Receiver<()>>,
}

impl Connecting {
    pub(crate) fn new(
        handle: ConnectionHandle,
        conn: proto::Connection,
        endpoint_events: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>,
        conn_events: mpsc::UnboundedReceiver<ConnectionEvent>,
        sender: Pin<Box<dyn UdpSender>>,
        runtime: Arc<dyn Runtime>,
    ) -> Self {
        let (on_handshake_data_send, on_handshake_data_recv) = oneshot::channel();
        let (on_connected_send, on_connected_recv) = oneshot::channel();

        let conn = ConnectionRef(Arc::new(Arc::new(ConnectionInner {
            state: Mutex::new(State::new(
                conn,
                handle,
                endpoint_events,
                conn_events,
                on_handshake_data_send,
                on_connected_send,
                sender,
                runtime.clone(),
            )),
            shared: Shared::default(),
        })));

        let driver = ConnectionDriver(conn.clone());
        runtime.spawn(Box::pin(
            async {
                if let Err(e) = driver.await {
                    tracing::error!("I/O error: {e}");
                }
            }
            .instrument(Span::current()),
        ));

        Self {
            conn: Some(conn),
            connected: on_connected_recv,
            handshake_data_ready: Some(on_handshake_data_recv),
        }
    }

    /// Convert into a 0-RTT or 0.5-RTT connection at the cost of weakened security
    ///
    /// Returns `Ok` immediately if the local endpoint is able to attempt sending 0/0.5-RTT data.
    /// If so, the returned [`Connection`] can be used to send application data without waiting for
    /// the rest of the handshake to complete, at the cost of weakened cryptographic security
    /// guarantees. The returned [`ZeroRttAccepted`] future resolves when the handshake does
    /// complete, at which point subsequently opened streams and written data will have full
    /// cryptographic protection.
    ///
    /// ## Outgoing
    ///
    /// For outgoing connections, the initial attempt to convert to a [`Connection`] which sends
    /// 0-RTT data will proceed if the [`crypto::ClientConfig`][crate::crypto::ClientConfig]
    /// attempts to resume a previous TLS session. However, **the remote endpoint may not actually
    /// _accept_ the 0-RTT data**--yet still accept the connection attempt in general. This
    /// possibility is conveyed through the [`ZeroRttAccepted`] future--when the handshake
    /// completes, it resolves to true if the 0-RTT data was accepted and false if it was rejected.
    /// If it was rejected, the existence of streams opened and other application data sent prior
    /// to the handshake completing will not be conveyed to the remote application, and local
    /// operations on them will return `ZeroRttRejected` errors.
    ///
    /// A server may reject 0-RTT data at its discretion, but accepting 0-RTT data requires the
    /// relevant resumption state to be stored in the server, which servers may limit or lose for
    /// various reasons including not persisting resumption state across server restarts.
    ///
    /// If manually providing a [`crypto::ClientConfig`][crate::crypto::ClientConfig], check your
    /// implementation's docs for 0-RTT pitfalls.
    ///
    /// ## Incoming
    ///
    /// For incoming connections, conversion to 0.5-RTT will always fully succeed. `into_0rtt` will
    /// always return `Ok` and the [`ZeroRttAccepted`] will always resolve to true.
    ///
    /// If manually providing a [`crypto::ServerConfig`][crate::crypto::ServerConfig], check your
    /// implementation's docs for 0-RTT pitfalls.
    ///
    /// ## Security
    ///
    /// On outgoing connections, this enables transmission of 0-RTT data, which is vulnerable to
    /// replay attacks, and should therefore never invoke non-idempotent operations.
    ///
    /// On incoming connections, this enables transmission of 0.5-RTT data, which may be sent
    /// before TLS client authentication has occurred, and should therefore not be used to send
    /// data for which client authentication is being used.
    pub fn into_0rtt(mut self) -> Result<(Connection, ZeroRttAccepted), Self> {
        // This lock borrows `self` and would normally be dropped at the end of this scope, so we'll
        // have to release it explicitly before returning `self` by value.
        let conn = (self.conn.as_mut().unwrap()).lock_without_waking("into_0rtt");

        let is_ok = conn.inner.has_0rtt() || conn.inner.side().is_server();
        drop(conn);

        if is_ok {
            let conn = self.conn.take().unwrap();
            Ok((Connection(conn), ZeroRttAccepted(self.connected)))
        } else {
            Err(self)
        }
    }

    /// Parameters negotiated during the handshake
    ///
    /// The dynamic type returned is determined by the configured
    /// [`Session`](proto::crypto::Session). For the default `rustls` session, the return value can
    /// be [`downcast`](Box::downcast) to a
    /// [`crypto::rustls::HandshakeData`](crate::crypto::rustls::HandshakeData).
    pub async fn handshake_data(&mut self) -> Result<Box<dyn Any>, ConnectionError> {
        // Taking &mut self allows us to use a single oneshot channel rather than dealing with
        // potentially many tasks waiting on the same event. It's a bit of a hack, but keeps things
        // simple.
        if let Some(x) = self.handshake_data_ready.take() {
            let _ = x.await;
        }
        let conn = self.conn.as_ref().unwrap();
        let inner = conn.lock_without_waking("handshake");
        inner
            .inner
            .crypto_session()
            .handshake_data()
            .ok_or_else(|| {
                inner
                    .error
                    .clone()
                    .expect("spurious handshake data ready notification")
            })
    }

    /// The local IP address which was used when the peer established
    /// the connection
    ///
    /// This can be different from the address the endpoint is bound to, in case
    /// the endpoint is bound to a wildcard address like `0.0.0.0` or `::`.
    ///
    /// This will return `None` for clients, or when the platform does not expose this
    /// information. See [`noq_udp::RecvMeta::dst_ip`](udp::RecvMeta::dst_ip) for a list of
    /// supported platforms when using [`noq_udp`](udp) for I/O, which is the default.
    ///
    /// Will panic if called after `poll` has returned `Ready`.
    pub fn local_ip(&self) -> Option<IpAddr> {
        let conn = self.conn.as_ref().expect("used after yielding Ready");
        let inner = conn.lock_without_waking("local_ip");

        inner
            .inner
            .network_path(PathId::ZERO)
            .expect("PathId::ZERO is the only path during the handshake")
            .local_ip()
    }

    /// The peer's UDP addresses
    ///
    /// Will panic if called after `poll` has returned `Ready`.
    pub fn remote_address(&self) -> SocketAddr {
        let conn_ref: &ConnectionRef = self.conn.as_ref().expect("used after yielding Ready");
        conn_ref
            .lock_without_waking("remote_address")
            .inner
            .network_path(PathId::ZERO)
            .expect("PathId::ZERO is the only path during the handshake")
            .remote()
    }
}

impl Future for Connecting {
    type Output = Result<Connection, ConnectionError>;
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut self.connected).poll(cx).map(|_| {
            let conn = self.conn.take().unwrap();
            let inner = conn.lock_without_waking("connecting");
            if inner.connected {
                drop(inner);
                Ok(Connection(conn))
            } else {
                Err(inner
                    .error
                    .clone()
                    .expect("connected signaled without connection success or error"))
            }
        })
    }
}

/// Future that completes when a connection is fully established
///
/// For clients, the resulting value indicates if 0-RTT was accepted. For servers, the resulting
/// value is meaningless.
pub struct ZeroRttAccepted(oneshot::Receiver<bool>);

impl Future for ZeroRttAccepted {
    type Output = bool;
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut self.0).poll(cx).map(|x| x.unwrap_or(false))
    }
}

/// A future that drives protocol logic for a connection
///
/// This future handles the protocol logic for a single connection, routing events from the
/// `Connection` API object to the `Endpoint` task and the related stream-related interfaces.
/// It also keeps track of outstanding timeouts for the `Connection`.
///
/// If the connection encounters an error condition, this future will yield an error. It will
/// terminate (yielding `Ok(())`) if the connection was closed without error. Unlike other
/// connection-related futures, this waits for the draining period to complete to ensure that
/// packets still in flight from the peer are handled gracefully.
#[must_use = "connection drivers must be spawned for their connections to function"]
#[derive(Debug)]
struct ConnectionDriver(ConnectionRef);

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

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let conn = &mut *self.0.lock_without_waking("poll");

        let span = debug_span!("drive", id = conn.handle.0);
        let _guard = span.enter();

        if let Err(e) = conn.process_conn_events(&self.0.shared, cx) {
            conn.terminate(e, &self.0.shared);
            return Poll::Ready(Ok(()));
        }
        let mut keep_going = conn.drive_transmit(cx)?;
        // If a timer expires, there might be more to transmit. When we transmit something, we
        // might need to reset a timer. Hence, we must loop until neither happens.
        keep_going |= conn.drive_timer(cx);
        conn.forward_endpoint_events();
        conn.forward_app_events(&self.0.shared);

        if !conn.inner.is_drained() {
            if keep_going {
                // If the connection hasn't processed all tasks, schedule it again
                cx.waker().wake_by_ref();
            } else {
                conn.driver = Some(cx.waker().clone());
            }
            return Poll::Pending;
        }
        if conn.error.is_none() {
            unreachable!("drained connections always have an error");
        }
        Poll::Ready(Ok(()))
    }
}

/// A QUIC connection.
///
/// If all references to a connection (including every clone of the `Connection` handle, streams of
/// incoming streams, and the various stream types) have been dropped, then the connection will be
/// automatically closed with an `error_code` of 0 and an empty `reason`. You can also close the
/// connection explicitly by calling [`Connection::close()`].
///
/// Closing the connection immediately abandons efforts to deliver data to the peer.  Upon
/// receiving CONNECTION_CLOSE the peer *may* drop any stream data not yet delivered to the
/// application. [`Connection::close()`] describes in more detail how to gracefully close a
/// connection without losing application data.
///
/// May be cloned to obtain another handle to the same connection.
///
/// [`Connection::close()`]: Connection::close
#[derive(Debug, Clone)]
pub struct Connection(ConnectionRef);

impl Connection {
    /// Returns a weak reference to the inner connection struct.
    pub fn weak_handle(&self) -> WeakConnectionHandle {
        self.0.weak_handle()
    }

    /// Initiate a new outgoing unidirectional stream.
    ///
    /// Streams are cheap and instantaneous to open unless blocked by flow control. As a
    /// consequence, the peer won't be notified that a stream has been opened until the stream is
    /// actually used.
    pub fn open_uni(&self) -> OpenUni<'_> {
        OpenUni {
            conn: &self.0,
            notify: self.0.shared.stream_budget_available[Dir::Uni as usize].notified(),
        }
    }

    /// Initiate a new outgoing bidirectional stream.
    ///
    /// Streams are cheap and instantaneous to open unless blocked by flow control. As a
    /// consequence, the peer won't be notified that a stream has been opened until the stream is
    /// actually used. Calling [`open_bi()`] then waiting on the [`RecvStream`] without writing
    /// anything to [`SendStream`] will never succeed.
    ///
    /// [`open_bi()`]: crate::Connection::open_bi
    /// [`SendStream`]: crate::SendStream
    /// [`RecvStream`]: crate::RecvStream
    pub fn open_bi(&self) -> OpenBi<'_> {
        OpenBi {
            conn: &self.0,
            notify: self.0.shared.stream_budget_available[Dir::Bi as usize].notified(),
        }
    }

    /// Accept the next incoming uni-directional stream
    pub fn accept_uni(&self) -> AcceptUni<'_> {
        AcceptUni {
            conn: &self.0,
            notify: self.0.shared.stream_incoming[Dir::Uni as usize].notified(),
        }
    }

    /// Accept the next incoming bidirectional stream
    ///
    /// **Important Note**: The `Connection` that calls [`open_bi()`] must write to its [`SendStream`]
    /// before the other `Connection` is able to `accept_bi()`. Calling [`open_bi()`] then
    /// waiting on the [`RecvStream`] without writing anything to [`SendStream`] will never succeed.
    ///
    /// [`accept_bi()`]: crate::Connection::accept_bi
    /// [`open_bi()`]: crate::Connection::open_bi
    /// [`SendStream`]: crate::SendStream
    /// [`RecvStream`]: crate::RecvStream
    pub fn accept_bi(&self) -> AcceptBi<'_> {
        AcceptBi {
            conn: &self.0,
            notify: self.0.shared.stream_incoming[Dir::Bi as usize].notified(),
        }
    }

    /// Receive an application datagram
    pub fn read_datagram(&self) -> ReadDatagram<'_> {
        ReadDatagram {
            conn: &self.0,
            notify: self.0.shared.datagram_received.notified(),
        }
    }

    /// Opens a new path if no path exists yet for `network_path`.
    ///
    /// If `network_path` has no local IP set, then this will open a new path
    /// if no path exists for this remote address, independent of the existing
    /// path's local IP. If a local IP is set, it will match against the full
    /// four-tuple of existing paths.
    ///
    /// Otherwise behaves exactly as [`open_path`].
    ///
    /// [`open_path`]: Self::open_path
    pub fn open_path_ensure(
        &self,
        network_path: impl Into<FourTuple>,
        initial_status: PathStatus,
    ) -> OpenPath {
        let network_path = network_path.into();
        let mut state = self.0.lock_and_wake("open_path");

        let network_path = match normalize_network_path(network_path, &state.inner) {
            Ok(network_path) => network_path,
            Err(err) => return OpenPath::rejected(err),
        };

        let now = state.runtime.now();
        let open_res = state
            .inner
            .open_path_ensure(network_path, initial_status, now);
        match open_res {
            Ok((path_id, existed)) if existed => {
                let recv = state.open_path.get(&path_id).map(|tx| tx.subscribe());
                drop(state);
                match recv {
                    Some(recv) => OpenPath::new(path_id, recv, self.0.clone()),
                    None => OpenPath::ready(path_id, self.0.clone()),
                }
            }
            Ok((path_id, _)) => {
                let (tx, rx) = watch::channel(Ok(()));
                state.open_path.insert(path_id, tx);
                drop(state);
                OpenPath::new(path_id, rx, self.0.clone())
            }
            Err(err) => OpenPath::rejected(err),
        }
    }

    /// Opens an additional path if the multipath extension is negotiated.
    ///
    /// This function takes a [`FourTuple`], which contains the remote address and an optional
    /// local IP. If the local IP is set, the path will be opened with this source address,
    /// and the endpoint must support sending from that IP address. You can also pass a
    /// [`SocketAddr`] to only set the remote address and leave the local IP interface unspecified.
    ///
    /// The returned future completes once the path is either fully opened and ready to
    /// carry application data, or if there was an error.
    ///
    /// Dropping the returned future does not cancel the opening of the path, the
    /// [`PathEvent::Established`] event will still be emitted from [`Self::path_events`] if
    /// the path opens.  The [`PathId`] for the events can be extracted from
    /// [`OpenPath::path_id`].
    ///
    /// Failure to open a path can either occur immediately, before polling the returned
    /// future, or at a later time.  If the failure is immediate [`OpenPath::path_id`] will
    /// return `None` and the future will be ready immediately.  If the failure happens
    /// later, a [`PathEvent`] will be emitted.
    pub fn open_path(
        &self,
        network_path: impl Into<FourTuple>,
        initial_status: PathStatus,
    ) -> OpenPath {
        let network_path = network_path.into();
        let mut state = self.0.lock_and_wake("open_path");

        let network_path = match normalize_network_path(network_path, &state.inner) {
            Ok(network_path) => network_path,
            Err(err) => return OpenPath::rejected(err),
        };

        let (on_open_path_send, on_open_path_recv) = watch::channel(Ok(()));
        let now = state.runtime.now();
        let open_res = state.inner.open_path(network_path, initial_status, now);
        match open_res {
            Ok(path_id) => {
                state.open_path.insert(path_id, on_open_path_send);
                drop(state);
                OpenPath::new(path_id, on_open_path_recv, self.0.clone())
            }
            Err(err) => OpenPath::rejected(err),
        }
    }

    /// Returns the [`Path`] structure of an open path
    pub fn path(&self, id: PathId) -> Option<Path> {
        Path::new(&self.0, id)
    }

    /// A stream of [`PathEvent`]s for all paths in this connection.
    ///
    /// The stream will yield a [`PathEvent`] whenever there is a change in the state of any path in this connection.
    /// The events need to be processed immediately, since there isn't an unbounded buffer for them.
    ///
    /// If processing of events lags behind too much, you will get an error of type [`crate::Lagged`] indicating
    /// how many events were lost. The stream continues after a lag, delivering the oldest retained message next.
    pub fn path_events(&self) -> crate::PathEvents {
        crate::PathEvents::new(
            self.0
                .lock_without_waking("path_events")
                .path_events
                .subscribe(),
        )
    }

    /// A stream of NAT traversal updates for this connection.
    ///
    /// The events need to be processed immediately, since there isn't an unbounded buffer for them.
    ///
    /// If processing of events lags behind too much, you will get an error of type [`crate::Lagged`] indicating
    /// how many events were lost. The stream continues after a lag, delivering the oldest retained message next.
    pub fn nat_traversal_updates(&self) -> crate::NatTraversalUpdates {
        crate::NatTraversalUpdates::new(
            self.0
                .lock_without_waking("nat_traversal_updates")
                .nat_traversal_updates
                .subscribe(),
        )
    }

    /// Wait for the connection to be closed for any reason
    ///
    /// Despite the return type's name, closed connections are often not an error condition at the
    /// application layer. Cases that might be routine include [`ConnectionError::LocallyClosed`]
    /// and [`ConnectionError::ApplicationClosed`].
    pub async fn closed(&self) -> ConnectionError {
        {
            let conn = self.0.lock_without_waking("closed");
            if let Some(error) = conn.error.as_ref() {
                return error.clone();
            }
            // Construct the future while the lock is held to ensure we can't miss a wakeup if
            // the `Notify` is signaled immediately after we release the lock. `await` it after
            // the lock guard is out of scope.
            self.0.shared.closed.notified()
        }
        .await;
        self.0
            .lock_without_waking("closed")
            .error
            .as_ref()
            .expect("closed without an error")
            .clone()
    }

    /// Wait for the connection to be closed without keeping a strong reference to the connection
    ///
    /// Returns a future that resolves, once the connection is closed, to a [`Closed`] struct
    /// describing the close reason and final connection and per-path statistics.
    ///
    /// Calling [`Self::closed`] keeps the connection alive until it is either closed locally via [`Connection::close`]
    /// or closed by the remote peer. This function instead does not keep the connection itself alive,
    /// so if all *other* clones of the connection are dropped, the connection will be closed implicitly even
    /// if there are futures returned from this function still being awaited.
    pub fn on_closed(&self) -> OnClosed {
        let (tx, rx) = oneshot::channel();
        let mut state = self.0.lock_without_waking("on_closed");
        if let Some(reason) = state.error.clone() {
            // Connection already closed, send immediately
            let _ = tx.send(Closed::new(&mut state, reason));
        } else {
            state.on_closed.push(tx);
        }
        drop(state);
        OnClosed {
            conn: self.weak_handle(),
            rx,
        }
    }

    /// Whether the connection is closed, and why.
    ///
    /// The close_reason is always set to `Some(ConnectionError)` when a socket is
    /// closed; whether it was closed manually by calling [`Connection::close()`] or due to
    /// an internal error (such as an idle timeout or the peer closing the
    /// connection).
    ///
    /// Note: when the connection is closed, `connection.close_reason().is_some()` will always be true.
    pub fn close_reason(&self) -> Option<ConnectionError> {
        self.0.lock_without_waking("close_reason").error.clone()
    }

    /// Close the connection immediately.
    ///
    /// Pending operations will fail immediately with [`ConnectionError::LocallyClosed`]. No
    /// more data is sent to the peer and the peer may drop buffered data upon receiving
    /// the CONNECTION_CLOSE frame.
    ///
    /// `error_code` and `reason` are not interpreted, and are provided directly to the peer.
    ///
    /// `reason` will be truncated to fit in a single packet with overhead; to improve odds that it
    /// is preserved in full, it should be kept under 1KiB.
    ///
    /// # Gracefully closing a connection
    ///
    /// Only the peer last receiving application data can be certain that all data is
    /// delivered. The only reliable action it can then take is to close the connection,
    /// potentially with a custom error code. The delivery of the final CONNECTION_CLOSE
    /// frame is very likely if both endpoints stay online long enough, and
    /// [`Endpoint::wait_idle()`] can be used to provide sufficient time. Otherwise, the
    /// remote peer will time out the connection, provided that the idle timeout is not
    /// disabled.
    ///
    /// The sending side can not guarantee all stream data is delivered to the remote
    /// application. It only knows the data is delivered to the QUIC stack of the remote
    /// endpoint. Once the local side sends a CONNECTION_CLOSE frame in response to calling
    /// [`close()`] the remote endpoint may drop any data it received but is as yet
    /// undelivered to the application, including data that was acknowledged as received to
    /// the local endpoint.
    ///
    /// [`ConnectionError::LocallyClosed`]: crate::ConnectionError::LocallyClosed
    /// [`Endpoint::wait_idle()`]: crate::Endpoint::wait_idle
    /// [`close()`]: Connection::close
    pub fn close(&self, error_code: VarInt, reason: &[u8]) {
        let conn = &mut *self.0.lock_without_waking("close"); // conn.close self-wakes
        conn.close(error_code, Bytes::copy_from_slice(reason), &self.0.shared);
    }

    /// Wait for the handshake to be confirmed.
    ///
    /// As a server, who must be authenticated by clients,
    /// this happens when the handshake completes
    /// upon receiving a TLS Finished message from the client.
    /// In return, the server send a HANDSHAKE_DONE frame.
    ///
    /// As a client, this happens when receiving a HANDSHAKE_DONE frame.
    /// At this point, the server has either accepted our authentication,
    /// or, if client authentication is not required, accepted our lack of authentication.
    pub async fn handshake_confirmed(&self) -> Result<(), ConnectionError> {
        {
            let conn = self.0.lock_without_waking("handshake_confirmed");
            if let Some(error) = conn.error.as_ref() {
                return Err(error.clone());
            }
            if conn.handshake_confirmed {
                return Ok(());
            }
            // Construct the future while the lock is held to ensure we can't miss a wakeup if
            // the `Notify` is signaled immediately after we release the lock. `await` it after
            // the lock guard is out of scope.
            self.0.shared.handshake_confirmed.notified()
        }
        .await;
        if let Some(error) = self
            .0
            .lock_without_waking("handshake_confirmed")
            .error
            .as_ref()
        {
            Err(error.clone())
        } else {
            Ok(())
        }
    }

    /// Transmit `data` as an unreliable, unordered application datagram
    ///
    /// Application datagrams are a low-level primitive. They may be lost or delivered out of order,
    /// and `data` must both fit inside a single QUIC packet and be smaller than the maximum
    /// dictated by the peer.
    ///
    /// Previously queued datagrams which are still unsent may be discarded to make space for this
    /// datagram, in order of oldest to newest.
    pub fn send_datagram(&self, data: Bytes) -> Result<(), SendDatagramError> {
        let conn = &mut *self.0.lock_and_wake("send_datagram");
        if let Some(ref x) = conn.error {
            return Err(SendDatagramError::ConnectionLost(x.clone()));
        }
        use proto::SendDatagramError::*;
        match conn.inner.datagrams().send(data, true) {
            Ok(()) => Ok(()),
            Err(e) => Err(match e {
                Blocked(..) => unreachable!(),
                UnsupportedByPeer => SendDatagramError::UnsupportedByPeer,
                Disabled => SendDatagramError::Disabled,
                TooLarge => SendDatagramError::TooLarge,
            }),
        }
    }

    /// Transmit `data` as an unreliable, unordered application datagram
    ///
    /// Unlike [`send_datagram()`], this method will wait for buffer space during congestion
    /// conditions, which effectively prioritizes old datagrams over new datagrams.
    ///
    /// See [`send_datagram()`] for details.
    ///
    /// [`send_datagram()`]: Connection::send_datagram
    pub fn send_datagram_wait(&self, data: Bytes) -> SendDatagram<'_> {
        SendDatagram {
            conn: &self.0,
            data: Some(data),
            notify: self.0.shared.datagrams_unblocked.notified(),
        }
    }

    /// Compute the maximum size of datagrams that may be passed to [`send_datagram()`].
    ///
    /// Returns `None` if datagrams are unsupported by the peer or disabled locally.
    ///
    /// This may change over the lifetime of a connection according to variation in the path MTU
    /// estimate. The peer can also enforce an arbitrarily small fixed limit, but if the peer's
    /// limit is large this is guaranteed to be a little over a kilobyte at minimum.
    ///
    /// Not necessarily the maximum size of received datagrams.
    ///
    /// [`send_datagram()`]: Connection::send_datagram
    pub fn max_datagram_size(&self) -> Option<usize> {
        self.0
            .lock_without_waking("max_datagram_size")
            .inner
            .datagrams()
            .max_size()
    }

    /// Bytes available in the outgoing datagram buffer
    ///
    /// When greater than zero, calling [`send_datagram()`](Self::send_datagram) with a datagram of
    /// at most this size is guaranteed not to cause older datagrams to be dropped.
    pub fn datagram_send_buffer_space(&self) -> usize {
        self.0
            .lock_without_waking("datagram_send_buffer_space")
            .inner
            .datagrams()
            .send_buffer_space()
    }

    /// The side of the connection (client or server)
    pub fn side(&self) -> Side {
        self.0.lock_without_waking("side").inner.side()
    }

    /// Current best estimate of this connection's latency (round-trip-time)
    pub fn rtt(&self, path_id: PathId) -> Option<Duration> {
        self.0.lock_without_waking("rtt").inner.rtt(path_id)
    }

    /// Returns connection statistics
    pub fn stats(&self) -> ConnectionStats {
        self.0.lock_without_waking("stats").inner.stats()
    }

    /// Returns path statistics
    pub fn path_stats(&self, path_id: PathId) -> Option<PathStats> {
        self.0.lock_without_waking("path_stats").path_stats(path_id)
    }

    /// Current state of the congestion control algorithm, for debugging purposes
    pub fn congestion_state(&self, path_id: PathId) -> Option<Box<dyn Controller>> {
        self.0
            .lock_without_waking("congestion_state")
            .inner
            .congestion_state(path_id)
            .map(|c| c.clone_box())
    }

    /// Succeeds when an incoming connection is proven not to be a replay attack.
    ///
    /// Only interesting for `Connection`s obtained from [`Connecting::into_0rtt`]. On 1-RTT
    /// connections, always completes immediately. Contrast
    /// [`handshake_confirmed`](Self::handshake_confirmed), which waits longer on clients e.g. to
    /// confirm client authentication.
    ///
    /// For incoming connections, reads from [`RecvStream`]s are guaranteed not to arise from replay
    /// attacks after this succeeds, even for streams accepted or read during 0-RTT. For outgoing
    /// connections, streams opened after this succeeds will never be discarded by the server due to
    /// 0-RTT rejection.
    pub async fn authenticated(&self) -> Result<(), ConnectionError> {
        let notified = {
            let conn = self.0.state.lock("connected");
            if let Some(e) = &conn.error {
                return Err(e.clone());
            }
            if conn.connected {
                return Ok(());
            }
            self.0.shared.connected.notified()
        };
        notified.await;
        let conn = self.0.state.lock("connected");
        conn.error.clone().map_or(Ok(()), Err)
    }

    /// Parameters negotiated during the handshake
    ///
    /// Guaranteed to return `Some` on fully established connections or after
    /// [`Connecting::handshake_data()`] succeeds. See that method's documentations for details on
    /// the returned value.
    ///
    /// [`Connection::handshake_data()`]: crate::Connecting::handshake_data
    pub fn handshake_data(&self) -> Option<Box<dyn Any>> {
        self.0
            .lock_without_waking("handshake_data")
            .inner
            .crypto_session()
            .handshake_data()
    }

    /// Cryptographic identity of the peer
    ///
    /// The dynamic type returned is determined by the configured
    /// [`Session`](proto::crypto::Session). For the default `rustls` session, the return value can
    /// be [`downcast`](Box::downcast) to a <code>Vec<[rustls::pki_types::CertificateDer]></code>
    pub fn peer_identity(&self) -> Option<Box<dyn Any>> {
        self.0
            .lock_without_waking("peer_identity")
            .inner
            .crypto_session()
            .peer_identity()
    }

    /// A stable identifier for this connection
    ///
    /// Peer addresses and connection IDs can change, but this value will remain
    /// fixed for the lifetime of the connection.
    pub fn stable_id(&self) -> usize {
        self.0.stable_id()
    }

    /// Update traffic keys spontaneously
    ///
    /// This primarily exists for testing purposes.
    pub fn force_key_update(&self) {
        self.0
            .lock_and_wake("force_key_update")
            .inner
            .force_key_update()
    }

    /// Derive keying material from this connection's TLS session secrets.
    ///
    /// When both peers call this method with the same `label` and `context`
    /// arguments and `output` buffers of equal length, they will get the
    /// same sequence of bytes in `output`. These bytes are cryptographically
    /// strong and pseudorandom, and are suitable for use as keying material.
    ///
    /// See [RFC5705](https://tools.ietf.org/html/rfc5705) for more information.
    pub fn export_keying_material(
        &self,
        output: &mut [u8],
        label: &[u8],
        context: &[u8],
    ) -> Result<(), proto::crypto::ExportKeyingMaterialError> {
        self.0
            .lock_without_waking("export_keying_material")
            .inner
            .crypto_session()
            .export_keying_material(output, label, context)
    }

    /// Modify the number of remotely initiated unidirectional streams that may be concurrently open
    ///
    /// No streams may be opened by the peer unless fewer than `count` are already open. Large
    /// `count`s increase both minimum and worst-case memory consumption.
    pub fn set_max_concurrent_uni_streams(&self, count: VarInt) {
        let mut conn = self.0.lock_and_wake("set_max_concurrent_uni_streams");
        conn.inner.set_max_concurrent_streams(Dir::Uni, count);
    }

    /// See [`proto::TransportConfig::send_window()`]
    pub fn set_send_window(&self, send_window: u64) {
        let mut conn = self.0.lock_and_wake("set_send_window");
        conn.inner.set_send_window(send_window);
    }

    /// See [`proto::TransportConfig::receive_window()`]
    pub fn set_receive_window(&self, receive_window: VarInt) {
        let mut conn = self.0.lock_and_wake("set_receive_window");
        conn.inner.set_receive_window(receive_window);
    }

    /// Modify the number of remotely initiated bidirectional streams that may be concurrently open
    ///
    /// No streams may be opened by the peer unless fewer than `count` are already open. Large
    /// `count`s increase both minimum and worst-case memory consumption.
    pub fn set_max_concurrent_bi_streams(&self, count: VarInt) {
        let mut conn = self.0.lock_and_wake("set_max_concurrent_bi_streams");
        conn.inner.set_max_concurrent_streams(Dir::Bi, count);
    }

    /// Track changes on our external address as reported by the peer.
    pub fn observed_external_addr(&self) -> crate::ObservedExternalAddr {
        let conn = self.0.lock_without_waking("external_addr");
        crate::ObservedExternalAddr::new(conn.observed_external_addr.subscribe())
    }

    /// Is multipath enabled?
    // TODO(flub): not a useful API, once we do real things with multipath we can remove
    // this again.
    pub fn is_multipath_enabled(&self) -> bool {
        let conn = self.0.lock_without_waking("is_multipath_enabled");
        conn.inner.is_multipath_negotiated()
    }

    /// Registers one address at which this endpoint might be reachable
    ///
    /// When the NAT traversal extension is negotiated, servers send these addresses to clients in
    /// `ADD_ADDRESS` frames. This allows clients to obtain server address candidates to initiate
    /// NAT traversal attempts. Clients provide their own reachable addresses in `REACH_OUT` frames
    /// when [`Self::initiate_nat_traversal_round`] is called.
    pub fn add_nat_traversal_address(
        &self,
        address: SocketAddr,
    ) -> Result<(), n0_nat_traversal::Error> {
        let mut conn = self.0.lock_and_wake("add_nat_traversal_addresses");
        conn.inner.add_nat_traversal_address(address)
    }

    /// Removes one or more addresses from the set of addresses at which this endpoint is reachable
    ///
    /// When the NAT traversal extension is negotiated, servers send address removals to
    /// clients in `REMOVE_ADDRESS` frames. This allows clients to stop using outdated
    /// server address candidates that are no longer valid for NAT traversal.
    ///
    /// For clients, removed addresses will no longer be advertised in `REACH_OUT` frames.
    ///
    /// Addresses not present in the set will be silently ignored.
    pub fn remove_nat_traversal_address(
        &self,
        address: SocketAddr,
    ) -> Result<(), n0_nat_traversal::Error> {
        let mut conn = self.0.lock_and_wake("remove_nat_traversal_addresses");
        conn.inner.remove_nat_traversal_address(address)
    }

    /// Get the current local nat traversal addresses
    pub fn get_local_nat_traversal_addresses(
        &self,
    ) -> Result<Vec<SocketAddr>, n0_nat_traversal::Error> {
        let conn = self
            .0
            .lock_without_waking("get_local_nat_traversal_addresses");
        conn.inner.get_local_nat_traversal_addresses()
    }

    /// Get the currently advertised nat traversal addresses by the server
    pub fn get_remote_nat_traversal_addresses(
        &self,
    ) -> Result<Vec<SocketAddr>, n0_nat_traversal::Error> {
        let conn = self
            .0
            .lock_without_waking("get_remote_nat_traversal_addresses");
        conn.inner.get_remote_nat_traversal_addresses()
    }

    /// Initiates a new nat traversal round
    ///
    /// A nat traversal round involves advertising the client's local addresses in `REACH_OUT`
    /// frames, and initiating probing of the known remote addresses. When a new round is
    /// initiated, the previous one is cancelled, and paths that have not been opened are closed.
    ///
    /// Returns the server addresses that are now being probed.
    pub fn initiate_nat_traversal_round(&self) -> Result<Vec<SocketAddr>, n0_nat_traversal::Error> {
        let mut conn = self.0.lock_and_wake("initiate_nat_traversal_round");
        let now = conn.runtime.now();
        conn.inner.initiate_nat_traversal_round(now)
    }
}

/// Normalizes a [`FourTuple`] against the connection's address family.
///
/// If the connection already uses IPv6 paths, the remote is canonicalised via
/// [`ensure_ipv6`]. If it uses IPv4 and the requested remote is IPv6, this returns
/// [`PathError::InvalidRemoteAddress`].
fn normalize_network_path(
    network_path: FourTuple,
    conn: &proto::Connection,
) -> Result<FourTuple, PathError> {
    // If endpoint::State::ipv6 is true we want to keep all our IP addresses as IPv6.
    // If not, we do not support IPv6.  We can not access endpoint::State from here
    // however, but either all our paths use an IPv6 address, or all our paths use an
    // IPv4 address.  So we can use that information.
    let ipv6 = conn
        .paths()
        .iter()
        .filter_map(|id| {
            conn.network_path(*id)
                .map(|addrs| addrs.remote().is_ipv6())
                .ok()
        })
        .next()
        .unwrap_or_default();
    let remote = network_path.remote();
    if remote.is_ipv6() && !ipv6 {
        Err(PathError::InvalidRemoteAddress(remote))
    } else if ipv6 {
        let remote = SocketAddr::V6(ensure_ipv6(remote));
        Ok(FourTuple::new(remote, network_path.local_ip()))
    } else {
        Ok(network_path)
    }
}

pin_project! {
    /// Future produced by [`Connection::open_uni`]
    pub struct OpenUni<'a> {
        conn: &'a ConnectionRef,
        #[pin]
        notify: Notified<'a>,
    }
}

impl Future for OpenUni<'_> {
    type Output = Result<SendStream, ConnectionError>;
    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let (conn, id, is_0rtt) = ready!(poll_open(ctx, this.conn, this.notify, Dir::Uni))?;
        Poll::Ready(Ok(SendStream::new(conn, id, is_0rtt)))
    }
}

pin_project! {
    /// Future produced by [`Connection::open_bi`]
    pub struct OpenBi<'a> {
        conn: &'a ConnectionRef,
        #[pin]
        notify: Notified<'a>,
    }
}

impl Future for OpenBi<'_> {
    type Output = Result<(SendStream, RecvStream), ConnectionError>;
    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let (conn, id, is_0rtt) = ready!(poll_open(ctx, this.conn, this.notify, Dir::Bi))?;

        Poll::Ready(Ok((
            SendStream::new(conn.clone(), id, is_0rtt),
            RecvStream::new(conn, id, is_0rtt),
        )))
    }
}

fn poll_open<'a>(
    ctx: &mut Context<'_>,
    conn: &'a ConnectionRef,
    mut notify: Pin<&mut Notified<'a>>,
    dir: Dir,
) -> Poll<Result<(ConnectionRef, StreamId, bool), ConnectionError>> {
    let mut state = conn.lock_without_waking("poll_open");
    if let Some(ref e) = state.error {
        return Poll::Ready(Err(e.clone()));
    } else if let Some(id) = state.inner.streams().open(dir) {
        let is_0rtt = state.inner.side().is_client() && state.inner.is_handshaking();
        drop(state); // Release the lock so clone can take it
        return Poll::Ready(Ok((conn.clone(), id, is_0rtt)));
    }
    loop {
        match 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(()) => {
                notify.set(conn.shared.stream_budget_available[dir as usize].notified())
            }
        }
    }
}

pin_project! {
    /// Future produced by [`Connection::accept_uni`]
    pub struct AcceptUni<'a> {
        conn: &'a ConnectionRef,
        #[pin]
        notify: Notified<'a>,
    }
}

impl Future for AcceptUni<'_> {
    type Output = Result<RecvStream, ConnectionError>;

    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let (conn, id, is_0rtt) = ready!(poll_accept(ctx, this.conn, this.notify, Dir::Uni))?;
        Poll::Ready(Ok(RecvStream::new(conn, id, is_0rtt)))
    }
}

pin_project! {
    /// Future produced by [`Connection::accept_bi`]
    pub struct AcceptBi<'a> {
        conn: &'a ConnectionRef,
        #[pin]
        notify: Notified<'a>,
    }
}

impl Future for AcceptBi<'_> {
    type Output = Result<(SendStream, RecvStream), ConnectionError>;

    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let (conn, id, is_0rtt) = ready!(poll_accept(ctx, this.conn, this.notify, Dir::Bi))?;
        Poll::Ready(Ok((
            SendStream::new(conn.clone(), id, is_0rtt),
            RecvStream::new(conn, id, is_0rtt),
        )))
    }
}

fn poll_accept<'a>(
    ctx: &mut Context<'_>,
    conn: &'a ConnectionRef,
    mut notify: Pin<&mut Notified<'a>>,
    dir: Dir,
) -> Poll<Result<(ConnectionRef, StreamId, bool), ConnectionError>> {
    let mut state = conn.lock_and_wake("poll_accept");
    // Check for incoming streams before checking `state.error` so that already-received streams,
    // which are necessarily finite, can be drained from a closed connection.
    if let Some(id) = state.inner.streams().accept(dir) {
        let is_0rtt = state.inner.is_handshaking();
        drop(state); // Release the lock (wake on drop) so clone can take it
        return Poll::Ready(Ok((conn.clone(), id, is_0rtt)));
    } else if let Some(ref e) = state.error {
        return Poll::Ready(Err(e.clone()));
    }
    loop {
        match 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(()) => notify.set(conn.shared.stream_incoming[dir as usize].notified()),
        }
    }
}

pin_project! {
    /// Future produced by [`Connection::read_datagram`]
    pub struct ReadDatagram<'a> {
        conn: &'a ConnectionRef,
        #[pin]
        notify: Notified<'a>,
    }
}

impl Future for ReadDatagram<'_> {
    type Output = Result<Bytes, ConnectionError>;
    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();
        let mut state = this.conn.lock_without_waking("ReadDatagram::poll");
        // Check for buffered datagrams before checking `state.error` so that already-received
        // datagrams, which are necessarily finite, can be drained from a closed connection.
        if let Some(x) = state.inner.datagrams().recv() {
            return Poll::Ready(Ok(x));
        } else if let Some(ref e) = state.error {
            return Poll::Ready(Err(e.clone()));
        }
        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.conn.shared.datagram_received.notified()),
            }
        }
    }
}

pin_project! {
    /// Future produced by [`Connection::send_datagram_wait`]
    pub struct SendDatagram<'a> {
        conn: &'a ConnectionRef,
        data: Option<Bytes>,
        #[pin]
        notify: Notified<'a>,
    }
}

impl Future for SendDatagram<'_> {
    type Output = Result<(), SendDatagramError>;
    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();
        let mut state = this.conn.lock_and_wake("SendDatagram::poll");
        if let Some(ref e) = state.error {
            return Poll::Ready(Err(SendDatagramError::ConnectionLost(e.clone())));
        }
        use proto::SendDatagramError::*;
        match state
            .inner
            .datagrams()
            .send(this.data.take().unwrap(), false)
        {
            Ok(()) => Poll::Ready(Ok(())),
            Err(e) => Poll::Ready(Err(match e {
                Blocked(data) => {
                    this.data.replace(data);
                    loop {
                        match this.notify.as_mut().poll(ctx) {
                            Poll::Pending => return Poll::Pending,
                            // Spurious wakeup, get a new future
                            Poll::Ready(()) => this
                                .notify
                                .set(this.conn.shared.datagrams_unblocked.notified()),
                        }
                    }
                }
                UnsupportedByPeer => SendDatagramError::UnsupportedByPeer,
                Disabled => SendDatagramError::Disabled,
                TooLarge => SendDatagramError::TooLarge,
            })),
        }
    }
}

/// State of a [`Connection`] at the moment it was closed.
///
/// Returned by the [`OnClosed`] future from [`Connection::on_closed`].
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct Closed {
    /// The reason the connection was closed.
    pub reason: ConnectionError,
    /// Aggregate connection statistics at the moment of close.
    pub stats: ConnectionStats,
    /// Per-path statistics for every path the connection knew about at close time.
    ///
    /// This includes paths that haven't been discarded at close time, plus any
    /// already-discarded paths whose final stats had been retained because a [`Path`]
    /// or [`WeakPathHandle`] handle was kept alive.
    ///
    /// [`WeakPathHandle`]: crate::WeakPathHandle
    pub path_stats: Vec<(PathId, PathStats)>,
}

impl Closed {
    /// Snapshot the current connection state into a [`Closed`] value.
    ///
    /// Must only be called once `state.error` has been set.
    pub(crate) fn new(state: &mut State, reason: ConnectionError) -> Self {
        let stats = state.inner.stats();

        let non_discarded_paths = state.inner.paths();
        let mut path_stats =
            Vec::with_capacity(non_discarded_paths.len() + state.final_path_stats.len());

        // Non-discarded paths are tracked by proto::Connection.
        path_stats.extend(
            non_discarded_paths
                .into_iter()
                .filter_map(|id| state.inner.path_stats(id).map(|stats| (id, stats))),
        );
        // Already-discarded paths whose final stats we kept around.
        path_stats.extend(
            state
                .final_path_stats
                .iter()
                .map(|(id, stats)| (*id, *stats)),
        );
        Self {
            reason,
            stats,
            path_stats,
        }
    }
}

/// Future returned by [`Connection::on_closed`]
///
/// Resolves to [`Closed`].
pub struct OnClosed {
    rx: oneshot::Receiver<Closed>,
    conn: WeakConnectionHandle,
}

impl Drop for OnClosed {
    fn drop(&mut self) {
        if self.rx.is_terminated() {
            return;
        };
        if let Some(conn) = self.conn.upgrade() {
            self.rx.close();
            conn.0
                .lock_without_waking("OnClosed::drop")
                .on_closed
                .retain(|tx| !tx.is_closed());
        }
    }
}

impl Future for OnClosed {
    type Output = Closed;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();
        // The `expect` is safe because `State::drop` ensures that all senders are triggered
        // before being dropped.
        Pin::new(&mut this.rx)
            .poll(cx)
            .map(|x| x.expect("on_close sender is never dropped before sending"))
    }
}

#[derive(Debug)]
#[allow(clippy::redundant_allocation)]
pub(crate) struct ConnectionRef(Arc<Arc<ConnectionInner>>);

impl ConnectionRef {
    #[allow(clippy::redundant_allocation)]
    fn from_arc(inner: Arc<Arc<ConnectionInner>>) -> Self {
        inner.shared.ref_count.fetch_add(1, Ordering::Relaxed);
        Self(inner)
    }

    pub(crate) fn stable_id(&self) -> usize {
        &*self.0 as *const _ as usize
    }

    pub(crate) fn weak_handle(&self) -> WeakConnectionHandle {
        WeakConnectionHandle(Arc::downgrade(&self.0))
    }
}

impl Clone for ConnectionRef {
    fn clone(&self) -> Self {
        Self::from_arc(Arc::clone(&self.0))
    }
}

impl Drop for ConnectionRef {
    fn drop(&mut self) {
        if self.shared.ref_count.fetch_sub(1, Ordering::Relaxed) > 1 {
            return;
        }

        let conn = &mut *self.lock_without_waking("drop");

        if !conn.inner.is_closed() {
            // If the driver is alive, it's just it and us, so we'd better shut it down. If it's
            // not, we can't do any harm. If there were any streams being opened, then either
            // the connection will be closed for an unrelated reason or a fresh reference will
            // be constructed for the newly opened stream.
            conn.implicit_close(&self.shared);
        }
    }
}

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

#[derive(Debug)]
pub(crate) struct ConnectionInner {
    /// Kept private intentionally, use [`Self::lock_and_wake`].
    state: Mutex<State>,
    pub(crate) shared: Shared,
}

impl ConnectionInner {
    /// Lock the state and return a guard that wakes the connection driver on drop.
    ///
    /// Use this for operations that may queue frames. The wake ensures the driver sends queued frames.
    /// If that's not needed, use [`Self::lock_without_waking`].
    pub(crate) fn lock_and_wake(&self, purpose: &'static str) -> WakeGuard<'_> {
        WakeGuard {
            guard: self.state.lock(purpose),
            wake: true,
        }
    }

    /// Lock the state and return a guard that unlocks once dropped.
    ///
    /// Use this for operations that don't require any action from the connection driver.
    /// Otherwise, use [`Self::lock_and_wake`] instead.
    pub(crate) fn lock_without_waking(&self, purpose: &'static str) -> WakeGuard<'_> {
        WakeGuard {
            guard: self.state.lock(purpose),
            wake: false,
        }
    }
}

/// [`MutexGuard`] wrapper that calls [`State::wake`] on drop.
#[derive(derive_more::Deref, derive_more::DerefMut)]
pub(crate) struct WakeGuard<'a> {
    #[deref]
    #[deref_mut]
    guard: MutexGuard<'a, State>,
    wake: bool,
}

impl WakeGuard<'_> {
    pub(crate) fn skip_waking(&mut self) {
        self.wake = false;
    }
}

impl Drop for WakeGuard<'_> {
    fn drop(&mut self) {
        if self.wake {
            self.guard.wake();
        }
    }
}

/// A handle to some connection internals, use with care.
///
/// This contains a weak reference to the connection so will not itself keep the connection
/// alive.
#[derive(Debug, Clone)]
pub struct WeakConnectionHandle(Weak<Arc<ConnectionInner>>);

impl WeakConnectionHandle {
    /// Returns `true` if the [`Connection`] associated with this handle is still alive.
    pub fn is_alive(&self) -> bool {
        self.0.upgrade().is_some()
    }

    /// Upgrade the handle to a full `Connection`
    pub fn upgrade(&self) -> Option<Connection> {
        self.upgrade_to_ref().map(Connection)
    }

    pub(crate) fn upgrade_to_ref(&self) -> Option<ConnectionRef> {
        self.0.upgrade().map(ConnectionRef::from_arc)
    }

    /// Returns `true` if the two [`WeakConnectionHandle`] point at the same connection.
    pub fn is_same_connection(&self, other: &Self) -> bool {
        self.0.ptr_eq(&other.0)
    }
}

#[derive(Debug, Default)]
pub(crate) struct Shared {
    handshake_confirmed: Notify,
    /// Notified when new streams may be locally initiated due to an increase in stream ID flow
    /// control budget
    stream_budget_available: [Notify; 2],
    /// Notified when the peer has initiated a new stream
    stream_incoming: [Notify; 2],
    datagram_received: Notify,
    datagrams_unblocked: Notify,
    closed: Notify,
    connected: Notify,
    /// Number of live handles that can be used to initiate or handle I/O; excludes the driver
    ref_count: AtomicUsize,
}

pub(crate) struct State {
    pub(crate) inner: proto::Connection,
    driver: Option<Waker>,
    handle: ConnectionHandle,
    on_handshake_data: Option<oneshot::Sender<()>>,
    on_connected: Option<oneshot::Sender<bool>>,
    connected: bool,
    handshake_confirmed: bool,
    timer: Option<Pin<Box<dyn AsyncTimer>>>,
    timer_deadline: Option<Instant>,
    conn_events: mpsc::UnboundedReceiver<ConnectionEvent>,
    endpoint_events: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>,
    pub(crate) blocked_writers: FxHashMap<StreamId, Waker>,
    pub(crate) blocked_readers: FxHashMap<StreamId, Waker>,
    pub(crate) stopped: FxHashMap<StreamId, Arc<Notify>>,
    /// Always set to Some before the connection becomes drained
    pub(crate) error: Option<ConnectionError>,
    /// Tracks paths being opened
    open_path: FxHashMap<PathId, watch::Sender<Result<(), PathError>>>,
    /// Tracks reference counts for paths.
    ///
    /// I.e. how many [`Path`] and [`WeakPathHandle`] structs are alive for a path.
    /// Each entry's [`PathRefOwner`] holds an [`AtomicUsize`] so that cloning or
    /// dropping a [`PathRef`] (held by [`Path`] or [`WeakPathHandle`]) does not need
    /// to lock this state.
    ///
    /// [`WeakPathHandle`]: crate::path::WeakPathHandle
    pub(crate) path_refs: FxHashMap<PathId, PathRefOwner>,
    /// Final path stats for discarded paths.
    ///
    /// We only insert entries if the discarded path has a non-zero reference count in [`Self::path_refs`].
    /// When the last reference to a path is dropped its entry is removed from both maps.
    pub(crate) final_path_stats: FxHashMap<PathId, PathStats>,
    pub(crate) path_events: tokio::sync::broadcast::Sender<PathEvent>,
    sender: Pin<Box<dyn UdpSender>>,
    pub(crate) runtime: Arc<dyn Runtime>,
    send_buffer: Vec<u8>,
    /// We buffer a transmit when the underlying I/O would block
    buffered_transmit: Option<proto::Transmit>,
    /// Our last external address reported by the peer. When multipath is enabled, this will be the
    /// last report across all paths.
    pub(crate) observed_external_addr: watch::Sender<Option<SocketAddr>>,
    pub(crate) nat_traversal_updates: tokio::sync::broadcast::Sender<n0_nat_traversal::Event>,
    on_closed: Vec<oneshot::Sender<Closed>>,
}

impl State {
    #[allow(clippy::too_many_arguments)]
    fn new(
        inner: proto::Connection,
        handle: ConnectionHandle,
        endpoint_events: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>,
        conn_events: mpsc::UnboundedReceiver<ConnectionEvent>,
        on_handshake_data: oneshot::Sender<()>,
        on_connected: oneshot::Sender<bool>,
        sender: Pin<Box<dyn UdpSender>>,
        runtime: Arc<dyn Runtime>,
    ) -> Self {
        Self {
            inner,
            driver: None,
            handle,
            on_handshake_data: Some(on_handshake_data),
            on_connected: Some(on_connected),
            connected: false,
            handshake_confirmed: false,
            timer: None,
            timer_deadline: None,
            conn_events,
            endpoint_events,
            blocked_writers: FxHashMap::default(),
            blocked_readers: FxHashMap::default(),
            stopped: FxHashMap::default(),
            open_path: FxHashMap::default(),
            error: None,
            sender,
            runtime,
            send_buffer: Vec::new(),
            buffered_transmit: None,
            path_events: tokio::sync::broadcast::channel(32).0,
            observed_external_addr: watch::Sender::new(None),
            nat_traversal_updates: tokio::sync::broadcast::channel(32).0,
            on_closed: Vec::new(),
            final_path_stats: Default::default(),
            path_refs: Default::default(),
        }
    }

    fn drive_transmit(&mut self, cx: &mut Context<'_>) -> io::Result<bool> {
        let now = self.runtime.now();
        let mut transmits = 0;

        let max_datagrams = self
            .sender
            .max_transmit_segments()
            .min(MAX_TRANSMIT_SEGMENTS);

        loop {
            // Retry the last transmit, or get a new one.
            let t = match self.buffered_transmit.take() {
                Some(t) => t,
                None => {
                    self.send_buffer.clear();
                    match self
                        .inner
                        .poll_transmit(now, max_datagrams, &mut self.send_buffer)
                    {
                        Some(t) => {
                            transmits += match t.segment_size {
                                None => 1,
                                Some(s) => t.size.div_ceil(s), // round up
                            };
                            t
                        }
                        None => break,
                    }
                }
            };

            let len = t.size;
            match self
                .sender
                .as_mut()
                .poll_send(&udp_transmit(&t, &self.send_buffer[..len]), cx)
            {
                Poll::Pending => {
                    self.buffered_transmit = Some(t);
                    return Ok(false);
                }
                Poll::Ready(Err(e)) => return Err(e),
                Poll::Ready(Ok(())) => {}
            }

            if transmits >= MAX_TRANSMIT_DATAGRAMS {
                // TODO: What isn't ideal here yet is that if we don't poll all
                // datagrams that could be sent we don't go into the `app_limited`
                // state and CWND continues to grow until we get here the next time.
                // See https://github.com/quinn-rs/quinn/issues/1126
                return Ok(true);
            }
        }

        Ok(false)
    }

    fn forward_endpoint_events(&mut self) {
        while let Some(event) = self.inner.poll_endpoint_events() {
            // If the endpoint driver is gone, noop.
            let _ = self.endpoint_events.send((self.handle, event));
        }
    }

    /// If this returns `Err`, the endpoint is dead, so the driver should exit immediately.
    fn process_conn_events(
        &mut self,
        shared: &Shared,
        cx: &mut Context<'_>,
    ) -> Result<(), ConnectionError> {
        loop {
            match self.conn_events.poll_recv(cx) {
                Poll::Ready(Some(ConnectionEvent::Rebind(sender))) => {
                    self.sender = sender;
                    self.inner.handle_network_change(None, self.runtime.now());
                }
                Poll::Ready(Some(ConnectionEvent::LocalAddressChanged(hint))) => {
                    self.inner
                        .handle_network_change(hint.as_deref().map(|x| x as _), self.runtime.now());
                }
                Poll::Ready(Some(ConnectionEvent::Proto(event))) => {
                    self.inner.handle_event(event);
                }
                Poll::Ready(Some(ConnectionEvent::Close { reason, error_code })) => {
                    self.close(error_code, reason, shared);
                }
                Poll::Ready(None) => {
                    return Err(ConnectionError::TransportError(TransportError::new(
                        TransportErrorCode::INTERNAL_ERROR,
                        "endpoint driver future was dropped".to_string(),
                    )));
                }
                Poll::Pending => {
                    return Ok(());
                }
            }
        }
    }

    fn forward_app_events(&mut self, shared: &Shared) {
        while let Some(event) = self.inner.poll() {
            use proto::Event::*;
            match event {
                HandshakeDataReady => {
                    if let Some(x) = self.on_handshake_data.take() {
                        let _ = x.send(());
                    }
                }
                Connected => {
                    self.connected = true;
                    shared.connected.notify_waiters();
                    if let Some(x) = self.on_connected.take() {
                        // We don't care if the on-connected future was dropped
                        let _ = x.send(self.inner.accepted_0rtt());
                    }
                    if self.inner.side().is_client() && !self.inner.accepted_0rtt() {
                        // Wake up rejected 0-RTT streams so they can fail immediately with
                        // `ZeroRttRejected` errors.
                        wake_all(&mut self.blocked_writers);
                        wake_all(&mut self.blocked_readers);
                        wake_all_notify(&mut self.stopped);
                    }
                }
                HandshakeConfirmed => {
                    self.handshake_confirmed = true;
                    shared.handshake_confirmed.notify_waiters();
                }
                ConnectionLost { reason } => {
                    self.terminate(reason, shared);
                }
                Stream(StreamEvent::Writable { id }) => wake_stream(id, &mut self.blocked_writers),
                Stream(StreamEvent::Opened { dir: Dir::Uni }) => {
                    shared.stream_incoming[Dir::Uni as usize].notify_waiters();
                }
                Stream(StreamEvent::Opened { dir: Dir::Bi }) => {
                    shared.stream_incoming[Dir::Bi as usize].notify_waiters();
                }
                DatagramReceived => {
                    shared.datagram_received.notify_waiters();
                }
                DatagramsUnblocked => {
                    shared.datagrams_unblocked.notify_waiters();
                }
                Stream(StreamEvent::Readable { id }) => wake_stream(id, &mut self.blocked_readers),
                Stream(StreamEvent::Available { dir }) => {
                    // Might mean any number of streams are ready, so we wake up everyone
                    shared.stream_budget_available[dir as usize].notify_waiters();
                }
                Stream(StreamEvent::Finished { id }) => wake_stream_notify(id, &mut self.stopped),
                Stream(StreamEvent::Stopped { id, .. }) => {
                    wake_stream_notify(id, &mut self.stopped);
                    wake_stream(id, &mut self.blocked_writers);
                }
                Path(ref evt @ PathEvent::ObservedAddr { addr: observed, .. }) => {
                    self.path_events.send(evt.clone()).ok();
                    self.observed_external_addr.send_if_modified(|addr| {
                        let old = addr.replace(observed);
                        old != *addr
                    });
                }
                Path(ref evt @ PathEvent::Established { id, .. }) => {
                    self.path_events.send(evt.clone()).ok();
                    if let Some(sender) = self.open_path.remove(&id) {
                        sender.send_modify(|value| *value = Ok(()));
                    }
                }
                Path(
                    ref evt @ PathEvent::Discarded {
                        id, ref path_stats, ..
                    },
                ) => {
                    if self.path_refs.contains_key(&id) {
                        self.final_path_stats.insert(id, *path_stats.clone());
                    }
                    self.path_events.send(evt.clone()).ok();
                }
                Path(ref evt @ PathEvent::Abandoned { id, .. }) => {
                    if let Some(sender) = self.open_path.remove(&id) {
                        // We don't care for the reason why this path was closed here, because semantically
                        // all close reasons for a path that has not yet been opened equals to `ValidationFailed`.
                        // With the noq API, there is no way to application-close a not-yet-opened path, so
                        // `ApplicationClosed` cannot occur. And all other variants will only occur for paths
                        // that have already been opened.
                        // The previous iteration of this code had another event `PathEvent::LocallyClosed` which
                        // contained a `PathError`, but that was only ever set to `ValidationFailed`.
                        let error = PathError::ValidationFailed;
                        sender.send_modify(|value| *value = Err(error));
                    }
                    // this will happen also for already opened paths
                    self.path_events.send(evt.clone()).ok();
                }
                Path(evt @ PathEvent::RemoteStatus { .. }) => {
                    self.path_events.send(evt).ok();
                }
                NatTraversal(update) => {
                    self.nat_traversal_updates.send(update).ok();
                }
                _ => {
                    // PathEvent is #[non_exhaustive].
                    // It's possible that noq is built against a newer noq-proto version.
                    // In that case, we need to ignore path events we can't handle yet.
                    // But for tests, we expect noq and noq-proto to be in sync, so we
                    // should panic in case we don't actually handle new cases.
                    #[cfg(test)]
                    panic!("Unhandled PathEvent variant: {event:?}");
                }
            }
        }
    }

    fn drive_timer(&mut self, cx: &mut Context<'_>) -> bool {
        let Some(deadline) = self.inner.poll_timeout() else {
            self.timer_deadline = None;
            return false;
        };

        // Use the clock rather than the async timer to detect expiry: Sleep::poll
        // respects Tokio's cooperative budget and can return Pending for elapsed
        // deadlines.
        let now = self.runtime.now();
        if now >= deadline {
            self.inner.handle_timeout(now);
            self.timer_deadline = None;
            return true;
        }

        match &mut self.timer {
            // Avoid resetting the timer when the deadline is unchanged.
            Some(delay) if self.timer_deadline != Some(deadline) => {
                delay.as_mut().reset(deadline);
            }
            None => {
                self.timer = Some(self.runtime.new_timer(deadline));
            }
            _ => {}
        }
        self.timer_deadline = Some(deadline);

        let delay = self
            .timer
            .as_mut()
            .expect("timer must exist in this state")
            .as_mut();
        if delay.poll(cx).is_pending() {
            return false;
        }

        // The deadline elapsed in the window between the clock check and poll.
        self.inner.handle_timeout(self.runtime.now());
        self.timer_deadline = None;
        true
    }

    /// Wake up a blocked `Driver` task to process I/O
    pub(crate) fn wake(&mut self) {
        if let Some(x) = self.driver.take() {
            x.wake();
        }
    }

    /// Used to wake up all blocked futures when the connection becomes closed for any reason
    fn terminate(&mut self, reason: ConnectionError, shared: &Shared) {
        self.error = Some(reason.clone());
        if let Some(x) = self.on_handshake_data.take() {
            let _ = x.send(());
        }
        wake_all(&mut self.blocked_writers);
        wake_all(&mut self.blocked_readers);
        shared.stream_budget_available[Dir::Uni as usize].notify_waiters();
        shared.stream_budget_available[Dir::Bi as usize].notify_waiters();
        shared.stream_incoming[Dir::Uni as usize].notify_waiters();
        shared.stream_incoming[Dir::Bi as usize].notify_waiters();
        shared.datagram_received.notify_waiters();
        shared.datagrams_unblocked.notify_waiters();
        if let Some(x) = self.on_connected.take() {
            let _ = x.send(false);
        }
        shared.handshake_confirmed.notify_waiters();
        wake_all_notify(&mut self.stopped);
        shared.closed.notify_waiters();
        // Send to the registered on_closed futures.
        if !self.on_closed.is_empty() {
            let closed = Closed::new(self, reason);
            for tx in self.on_closed.drain(..) {
                tx.send(closed.clone()).ok();
            }
        }
        shared.connected.notify_waiters();
    }

    fn close(&mut self, error_code: VarInt, reason: Bytes, shared: &Shared) {
        self.inner.close(self.runtime.now(), error_code, reason);
        self.terminate(ConnectionError::LocallyClosed, shared);
        self.wake();
    }

    /// Close for a reason other than the application's explicit request
    pub(crate) fn implicit_close(&mut self, shared: &Shared) {
        self.close(0u32.into(), Bytes::new(), shared);
    }

    pub(crate) fn check_0rtt(&self) -> Result<(), ()> {
        if self.inner.is_handshaking()
            || self.inner.accepted_0rtt()
            || self.inner.side().is_server()
        {
            Ok(())
        } else {
            Err(())
        }
    }

    /// Returns [`PathStats`] for a path, if available.
    ///
    /// This gets the stats from [`proto::Connection`]. If that returns `None`
    /// it gets them from `Self::final_path_stats` instead.
    pub(crate) fn path_stats(&mut self, path_id: PathId) -> Option<PathStats> {
        self.inner
            .path_stats(path_id)
            .or_else(|| self.final_path_stats.get(&path_id).copied())
    }

    /// Acquire a new [`PathRef`] for a path id, bumping its reference counter by 1.
    ///
    /// The returned [`PathRef`] is intended to be stored on a [`Path`] or [`WeakPathHandle`].
    /// Its reference count is automatically increased when cloned. When its owner is dropped,
    /// [`PathRef::on_drop`] must be called to decrement the refcount.
    ///
    /// [`WeakPathHandle`]: crate::path::WeakPathHandle
    pub(crate) fn acquire_path_ref(&mut self, path_id: PathId) -> PathRef {
        self.path_refs.entry(path_id).or_default().acquire(path_id)
    }
}

impl Drop for State {
    fn drop(&mut self) {
        if !self.inner.is_drained() {
            // Ensure the endpoint can tidy up
            let _ = self
                .endpoint_events
                .send((self.handle, EndpointEvent::drained()));
        }

        if !self.on_closed.is_empty()
            && let Some(reason) = self.error.clone()
        {
            // Ensure that all on_closed oneshot senders are triggered before dropping.
            let closed = Closed::new(self, reason);
            for tx in self.on_closed.drain(..) {
                tx.send(closed.clone()).ok();
            }
        }
    }
}

impl fmt::Debug for State {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("State").field("inner", &self.inner).finish()
    }
}

fn wake_stream(stream_id: StreamId, wakers: &mut FxHashMap<StreamId, Waker>) {
    if let Some(waker) = wakers.remove(&stream_id) {
        waker.wake();
    }
}

fn wake_all(wakers: &mut FxHashMap<StreamId, Waker>) {
    wakers.drain().for_each(|(_, waker)| waker.wake())
}

fn wake_stream_notify(stream_id: StreamId, wakers: &mut FxHashMap<StreamId, Arc<Notify>>) {
    if let Some(notify) = wakers.remove(&stream_id) {
        notify.notify_waiters()
    }
}

fn wake_all_notify(wakers: &mut FxHashMap<StreamId, Arc<Notify>>) {
    wakers
        .drain()
        .for_each(|(_, notify)| notify.notify_waiters())
}

/// Errors that can arise when sending a datagram
#[derive(Debug, Error, Clone, Eq, PartialEq)]
pub enum SendDatagramError {
    /// The peer does not support receiving datagram frames
    #[error("datagrams not supported by peer")]
    UnsupportedByPeer,
    /// Datagram support is disabled locally
    #[error("datagram support disabled")]
    Disabled,
    /// The datagram is larger than the connection can currently accommodate
    ///
    /// Indicates that the path MTU minus overhead or the limit advertised by the peer has been
    /// exceeded.
    #[error("datagram too large")]
    TooLarge,
    /// The connection was lost
    #[error("connection lost")]
    ConnectionLost(#[from] ConnectionError),
}

/// The maximum amount of datagrams which will be produced in a single `drive_transmit` call
///
/// This limits the amount of CPU resources consumed by datagram generation,
/// and allows other tasks (like receiving ACKs) to run in between.
const MAX_TRANSMIT_DATAGRAMS: usize = 20;

/// The maximum amount of datagrams that are sent in a single transmit
///
/// This can be lower than the maximum platform capabilities, to avoid excessive
/// memory allocations when calling `poll_transmit()`. Benchmarks have shown
/// that numbers around 10 are a good compromise.
const MAX_TRANSMIT_SEGMENTS: NonZeroUsize = NonZeroUsize::new(10).expect("known");