stun_proto/

agent.rs

// Copyright (C) 2020 Matthew Waters <matthew@centricular.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//
// SPDX-License-Identifier: MIT OR Apache-2.0

//! # STUN agent
//!
//! A STUN Agent that follows the procedures of [RFC5389] and [RFC8489] and is implemented with the
//! sans-IO pattern. This agent does no IO processing and operates solely on inputs it is
//! provided.
//!
//! [RFC8489]: https://tools.ietf.org/html/rfc8489
//! [RFC5389]: https://tools.ietf.org/html/rfc5389

use core::net::SocketAddr;
use core::sync::atomic::{AtomicUsize, Ordering};

use alloc::collections::{BTreeMap, BTreeSet};
use alloc::vec;
use alloc::vec::Vec;
use core::time::Duration;

use crate::stats::StunAgentStats;
use crate::Instant;

use stun_types::attribute::*;
use stun_types::data::Data;
use stun_types::message::*;

use stun_types::TransportType;

use tracing::{debug, trace};

static STUN_AGENT_COUNT: AtomicUsize = AtomicUsize::new(0);

/// Implementation of a STUN agent.
#[derive(Debug)]
pub struct StunAgent {
    id: usize,
    transport: TransportType,
    local_addr: SocketAddr,
    remote_addr: Option<SocketAddr>,
    validated_peers: BTreeSet<SocketAddr>,
    outstanding_requests: BTreeMap<TransactionId, StunRequestState>,
    request_timeouts: Vec<Duration>,
    last_retransmit_timeout: Duration,
    stats: Option<StunAgentStats>,
}

/// Builder struct for a [`StunAgent`].
#[derive(Debug)]
pub struct StunAgentBuilder {
    transport: TransportType,
    local_addr: SocketAddr,
    remote_addr: Option<SocketAddr>,
    rto: RequestRto,
    stats: bool,
}

impl StunAgentBuilder {
    /// Set the remote address the [`StunAgent`] will be configured to only send data to.
    pub fn remote_addr(mut self, addr: SocketAddr) -> Self {
        self.remote_addr = Some(addr);
        self
    }

    /// Configure the default timeouts and retransmissions for each STUN request.
    ///
    /// - `initial` - the initial time between consecutive transmissions. If 0, or 1, then only a
    ///   single request will be performed.
    /// - `max` - the maximum amount of time between consecutive retransmits.
    /// - `retransmits` - the total number of transmissions of the request.
    /// - `final_retransmit_timeout` - the amount of time after the final transmission to wait
    ///   for a response before considering the request as having timed out.
    ///
    /// As specified in RFC 8489, `initial_rto` should be >= 500ms (unless specific information is
    /// available on the RTT, `max` is `Duration::MAX`, `retransmits` has a default value of 7,
    /// and `last_retransmit_timeout` should be `16 * initial_rto`.
    ///
    /// STUN transactions over TCP will only send a single request and have a timeout of the sum of
    /// the timeouts of a UDP transaction.
    pub fn request_retransmits(
        mut self,
        initial: Duration,
        max: Duration,
        retransmits: u32,
        final_retransmit_timeout: Duration,
    ) -> Self {
        self.rto.initial = initial;
        self.rto.max = max;
        self.rto.retransmits = retransmits;
        self.rto.last_retransmit = final_retransmit_timeout;
        self
    }

    /// Enable statistics tracking on the built [`StunAgent`].
    ///
    /// When enabled, the agent tracks round-trip times, bytes sent and received,
    /// transaction counts, and timeout/cancellation counts. Access the collected
    /// statistics via [`StunAgent::stats`] or [`StunAgent::stats_mut`].
    pub fn stats(mut self, stats: bool) -> Self {
        self.stats = stats;
        self
    }

    /// Build the [`StunAgent`].
    pub fn build(self) -> StunAgent {
        let id = STUN_AGENT_COUNT.fetch_add(1, Ordering::SeqCst);
        let (request_timeouts, last_retransmit_timeout) =
            self.rto.calculate_timeouts(self.transport);
        let stats = if self.stats {
            Some(StunAgentStats::default())
        } else {
            None
        };
        StunAgent {
            id,
            transport: self.transport,
            local_addr: self.local_addr,
            remote_addr: self.remote_addr,
            validated_peers: Default::default(),
            outstanding_requests: Default::default(),
            request_timeouts,
            last_retransmit_timeout,
            stats,
        }
    }
}

impl StunAgent {
    /// Create a new [`StunAgentBuilder`].
    pub fn builder(transport: TransportType, local_addr: SocketAddr) -> StunAgentBuilder {
        StunAgentBuilder {
            transport,
            local_addr,
            remote_addr: None,
            rto: Default::default(),
            stats: false,
        }
    }

    /// The [`TransportType`] of this [`StunAgent`].
    pub fn transport(&self) -> TransportType {
        self.transport
    }

    /// The local address of this [`StunAgent`].
    pub fn local_addr(&self) -> SocketAddr {
        self.local_addr
    }

    /// The remote address of this [`StunAgent`].
    pub fn remote_addr(&self) -> Option<SocketAddr> {
        self.remote_addr
    }

    /// Returns the statistics for this agent if stats collection is enabled.
    pub fn stats(&self) -> Option<&StunAgentStats> {
        self.stats.as_ref()
    }

    /// Returns the mutable statistics for this agent if stats collection is enabled.
    pub fn stats_mut(&mut self) -> Option<&mut StunAgentStats> {
        self.stats.as_mut()
    }

    /// Perform any operations needed to be able to send data to a peer.
    pub fn send_data<T: AsRef<[u8]>>(&self, bytes: T, to: SocketAddr) -> Transmit<T> {
        send_data(self.transport, bytes, self.local_addr, to)
    }

    /// Perform any operations needed to be able to send a [`Message`] to a peer.
    ///
    /// The returned [`Transmit`] must be sent to the respective peer after this call.
    ///
    /// # Panics
    ///
    /// - If the STUN Message is a request. Use [`send_request()`](StunAgent::send_request) instead.
    #[tracing::instrument(name = "stun_agent_send",
        skip(self, msg),
        fields(
            transport = %self.transport,
            from = %self.local_addr,
            transaction_id,
        )
    )]
    pub fn send<T: AsRef<[u8]>>(
        &mut self,
        msg: T,
        to: SocketAddr,
        now: Instant,
    ) -> Result<Transmit<T>, StunError> {
        let data = msg.as_ref();
        let hdr = MessageHeader::from_bytes(data)?;
        tracing::Span::current().record(
            "transaction_id",
            tracing::field::display(hdr.transaction_id()),
        );
        let cls = hdr.get_type().class();
        assert_ne!(cls, MessageClass::Request);
        trace!("Sending {} to {to}", hdr.get_type());
        if let Some(s) = &mut self.stats {
            if cls == MessageClass::Indication {
                s.record_indication_sent(data.len() as u64);
            } else {
                s.record_response_sent(data.len() as u64);
            }
        }
        Ok(Transmit::new(msg, self.transport, self.local_addr, to))
    }

    /// Perform any operations needed to be able to send a request [`Message`] to a peer.
    ///
    /// The returned [`Transmit`] must be sent to the respective peer after this call.
    ///
    /// # Panics
    ///
    /// - If the STUN Message is not a request. Use [`send()`](StunAgent::send) instead.
    #[tracing::instrument(name = "stun_agent_send_request",
        skip(self, msg),
        fields(
            transport = %self.transport,
            from = %self.local_addr,
            transaction_id,
        )
    )]
    pub fn send_request<'a, T: AsRef<[u8]>>(
        &'a mut self,
        msg: T,
        to: SocketAddr,
        now: Instant,
    ) -> Result<Transmit<Data<'a>>, StunError> {
        let data = msg.as_ref();
        let hdr = MessageHeader::from_bytes(data)?;
        assert!(hdr.get_type().has_class(MessageClass::Request));
        let transaction_id = hdr.transaction_id();
        tracing::Span::current().record("transaction_id", tracing::field::display(transaction_id));
        let state = match self.outstanding_requests.entry(transaction_id) {
            alloc::collections::btree_map::Entry::Vacant(entry) => {
                let integrity_algorithm = MessageAttributesIter::new(data)
                    .filter_map(|(_offset, attr)| match attr.get_type() {
                        MessageIntegrity::TYPE => Some(IntegrityAlgorithm::Sha1),
                        MessageIntegritySha256::TYPE => Some(IntegrityAlgorithm::Sha256),
                        _ => None,
                    })
                    .last();
                trace!("Adding request to {to} with integrity algorithm: {integrity_algorithm:?}");
                if let Some(s) = &mut self.stats {
                    s.record_request_sent(data.len() as u64);
                }
                entry.insert(StunRequestState::new(
                    msg,
                    self.transport,
                    self.local_addr,
                    to,
                    transaction_id,
                    integrity_algorithm,
                    self.request_timeouts.clone(),
                    self.last_retransmit_timeout,
                ))
            }
            alloc::collections::btree_map::Entry::Occupied(_entry) => {
                return Err(StunError::AlreadyInProgress);
            }
        };
        let Some(transmit) = state.poll_transmit(now) else {
            unreachable!();
        };
        Ok(Transmit::new(
            Data::from(transmit.data),
            transmit.transport,
            transmit.from,
            transmit.to,
        ))
    }

    /// Returns whether this agent has received or sent a STUN message with this peer. Failure may
    /// be the result of an attacker and the caller must drop any non-STUN data received before this
    /// functions returns `true`.
    ///
    /// If non-STUN data is received over a TCP connection from an unvalidated peer, the caller
    /// must immediately close the TCP connection.
    pub fn is_validated_peer(&self, remote_addr: SocketAddr) -> bool {
        self.validated_peers.contains(&remote_addr)
    }

    /// Indicate to the STUN agent that STUN messages have been sent/received to/from a peer.
    #[tracing::instrument(
        name = "stun_validated_peer"
        skip(self),
        fields(stun_id = self.id)
    )]
    pub fn validated_peer(&mut self, addr: SocketAddr) {
        if !self.validated_peers.contains(&addr) {
            debug!("validated peer {:?}", addr);
            self.validated_peers.insert(addr);
        }
    }

    /// Provide data received on a socket from a peer for handling by the [`StunAgent`] after it
    /// has successfully passed authentication.
    ///
    /// For responses, this will cause the associated request to be removed from the agent if it
    /// exists.
    ///
    /// The return value indicates whether the message passes internal checks and should be acted
    /// upon.
    ///
    /// This function is provided for backwards compatibility and new code should use
    /// [`StunAgent::handle_stun_message_with_time`] in order to compute round trip statistics.
    #[tracing::instrument(
        name = "stun_handle_message"
        skip(self, msg, from),
        fields(
            transaction_id = %msg.transaction_id(),
        )
    )]
    #[deprecated = "Use handle_stun_message_with_time() to be able to retrieve round trip statistics"]
    // FIXME: 3.0: remove
    pub fn handle_stun_message(&mut self, msg: &Message<'_>, from: SocketAddr) -> bool {
        self.handle_stun_message_internal(msg, from, None)
    }

    /// Provide data received on a socket from a peer for handling by the [`StunAgent`] after it
    /// has successfully passed authentication.
    ///
    /// For responses, this will cause the associated request to be removed from the agent if it
    /// exists. If statistics are enabled, then `time` is used to compute round trip statistics.
    ///
    /// The return value indicates whether the message passes internal checks and should be acted
    /// upon.
    pub fn handle_stun_message_with_time(
        &mut self,
        msg: &Message<'_>,
        from: SocketAddr,
        now: Instant,
    ) -> bool {
        self.handle_stun_message_internal(msg, from, Some(now))
    }

    fn handle_stun_message_internal(
        &mut self,
        msg: &Message<'_>,
        from: SocketAddr,
        now: Option<Instant>,
    ) -> bool {
        let outstanding = if msg.is_response() {
            self.take_outstanding_request(&msg.transaction_id())
        } else {
            None
        };
        if msg.is_response() && outstanding.is_none() {
            trace!("original request disappeared");
            return false;
        }
        if let Some(s) = &mut self.stats {
            let msg_len = msg.as_bytes().len() as u64;
            match msg.class() {
                MessageClass::Request => s.record_request_received(msg_len),
                MessageClass::Indication => s.record_indication_received(msg_len),
                MessageClass::Error | MessageClass::Success => {
                    if let Some(state) = &outstanding {
                        if let Some(rtt) = now
                            .zip(
                                state
                                    .last_send_time
                                    // Only requests that did not require a retransmit count towards rtt.
                                    .filter(|_| state.timeout_i == 0)
                                    .zip(state.first_send_time),
                            )
                            .map(|(now, (last, _first))| now - last)
                        {
                            s.record_rtt(rtt);
                        }
                    }
                    s.record_response_received(msg_len);
                }
            }
        }
        self.validated_peer(from);
        true
    }

    #[tracing::instrument(
        skip(self, transaction_id),
        fields(transaction_id = %transaction_id)
    )]
    fn take_outstanding_request(
        &mut self,
        transaction_id: &TransactionId,
    ) -> Option<StunRequestState> {
        if let Some(request) = self.outstanding_requests.remove(transaction_id) {
            trace!("removing request");
            Some(request)
        } else {
            trace!("no outstanding request");
            None
        }
    }

    /// Retrieve a reference to an outstanding STUN request. Outstanding requests are kept until
    /// either:
    /// - [`handle_stun_message()`](StunAgent::handle_stun_message) is called, or
    /// - [`poll()`](StunAgent::poll) returns [`StunAgentPollRet::TransactionCancelled`] or
    ///   [`StunAgentPollRet::TransactionTimedOut`] for the request.
    pub fn request_transaction(&self, transaction_id: TransactionId) -> Option<StunRequest<'_>> {
        self.request_state(transaction_id)
            .map(|request| StunRequest {
                agent: self,
                peer_address: request.to,
                request_integrity: request.request_integrity,
            })
    }

    /// Retrieve a mutable reference to an outstanding STUN request. Outstanding requests are kept
    /// until either:
    /// - [`handle_stun_message()`](StunAgent::handle_stun_message) is called, or
    /// - [`poll()`](StunAgent::poll) returns [`StunAgentPollRet::TransactionCancelled`] or
    ///   [`StunAgentPollRet::TransactionTimedOut`] for the request.
    pub fn mut_request_transaction(
        &mut self,
        transaction_id: TransactionId,
    ) -> Option<StunRequestMut<'_>> {
        if let Some(request) = self.mut_request_state(transaction_id) {
            let peer_address = request.to;
            let request_integrity = request.request_integrity;
            Some(StunRequestMut {
                agent: self,
                transaction_id,
                peer_address,
                request_integrity,
            })
        } else {
            None
        }
    }

    fn mut_request_state(
        &mut self,
        transaction_id: TransactionId,
    ) -> Option<&mut StunRequestState> {
        self.outstanding_requests.get_mut(&transaction_id)
    }

    fn request_state(&self, transaction_id: TransactionId) -> Option<&StunRequestState> {
        self.outstanding_requests.get(&transaction_id)
    }

    /// Poll the agent for making further progress on any outstanding requests. The returned value
    /// indicates the current state and anything the caller needs to perform.
    ///
    /// Upon expiry of the timer from [`StunAgentPollRet::WaitUntil`],
    /// [`poll_transmit()`](StunAgent::poll_transmit) must be called.
    #[tracing::instrument(
        name = "stun_agent_poll"
        level = "debug",
        skip(self),
    )]
    pub fn poll(&mut self, now: Instant) -> StunAgentPollRet {
        let mut lowest_wait = now + Duration::from_secs(3600);
        let mut timeout = None;
        let mut cancelled = None;
        for (transaction_id, request) in self.outstanding_requests.iter_mut() {
            debug_assert_eq!(transaction_id, &request.transaction_id);
            match request.poll(now) {
                StunRequestPollRet::Cancelled => {
                    cancelled = Some(*transaction_id);
                    break;
                }
                StunRequestPollRet::WaitUntil(wait_until) => {
                    if wait_until < lowest_wait {
                        lowest_wait = wait_until;
                    }
                }
                StunRequestPollRet::TimedOut => {
                    timeout = Some(*transaction_id);
                    break;
                }
            }
        }
        if let Some(transaction) = timeout {
            if let Some(_state) = self.outstanding_requests.remove(&transaction) {
                if let Some(s) = &mut self.stats {
                    s.record_timeout();
                }
                return StunAgentPollRet::TransactionTimedOut(transaction);
            }
        }
        if let Some(transaction) = cancelled {
            if let Some(_state) = self.outstanding_requests.remove(&transaction) {
                if let Some(s) = &mut self.stats {
                    s.record_cancelled();
                }
                return StunAgentPollRet::TransactionCancelled(transaction);
            }
        }
        StunAgentPollRet::WaitUntil(lowest_wait)
    }

    /// Poll for any transmissions that may need to be performed.
    #[tracing::instrument(
        name = "stun_agent_poll_transmit"
        level = "debug",
        skip(self),
    )]
    pub fn poll_transmit(&mut self, now: Instant) -> Option<Transmit<&[u8]>> {
        let transmit = self
            .outstanding_requests
            .values_mut()
            .filter_map(|request| request.poll_transmit(now))
            .next();
        if let Some(t) = &transmit {
            if let Some(s) = &mut self.stats {
                s.record_retransmit_bytes(t.data.as_ref().len() as u64);
            }
        }
        transmit
    }
}

/// Return value for [`StunAgent::poll`].
#[derive(Debug)]
pub enum StunAgentPollRet {
    /// An oustanding transaction timed out and has been removed from the agent.
    TransactionTimedOut(TransactionId),
    /// An oustanding transaction was cancelled and has been removed from the agent.
    TransactionCancelled(TransactionId),
    /// Wait until the specified time has passed.
    WaitUntil(Instant),
}

fn send_data<T: AsRef<[u8]>>(
    transport: TransportType,
    bytes: T,
    from: SocketAddr,
    to: SocketAddr,
) -> Transmit<T> {
    Transmit::new(bytes, transport, from, to)
}

/// A piece of data that needs to, or has been transmitted.
#[derive(Debug)]
pub struct Transmit<T: AsRef<[u8]>> {
    /// The data blob.
    pub data: T,
    /// The transport for the transmission.
    pub transport: TransportType,
    /// The source address of the transmission.
    pub from: SocketAddr,
    /// The destination address of the transmission.
    pub to: SocketAddr,
}

impl<T: AsRef<[u8]>> core::fmt::Display for Transmit<T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(
            f,
            "Transmit({}: {} -> {} of {} bytes)",
            self.transport,
            self.from,
            self.to,
            self.data.as_ref().len()
        )
    }
}

impl<T: AsRef<[u8]>> Transmit<T> {
    /// Construct a new [`Transmit`] with the specifid data and 5-tuple.
    pub fn new(data: T, transport: TransportType, from: SocketAddr, to: SocketAddr) -> Self {
        Self {
            data,
            transport,
            from,
            to,
        }
    }

    /// Reinterpret the data of a [`Transmit`] into a different type.
    ///
    /// # Examples
    ///
    /// ```
    /// # use stun_proto::agent::Transmit;
    /// # use stun_proto::types::TransportType;
    /// # use core::net::SocketAddr;
    /// let local_addr = "10.0.0.1:1000".parse().unwrap();
    /// let remote_addr = "10.0.0.2:2000".parse().unwrap();
    /// let slice = [42; 8];
    /// let transmit = Transmit::new(slice.clone(), TransportType::Udp, local_addr, remote_addr);
    /// // change the data type of the `Transmit` into a `Vec<u8>`.
    /// let transmit = transmit.reinterpret_data(|data| data.to_vec());
    /// # assert_eq!(transmit.transport, TransportType::Udp);
    /// # assert_eq!(transmit.from, local_addr);
    /// # assert_eq!(transmit.to, remote_addr);
    /// assert_eq!(transmit.data, slice.as_slice());
    /// ```
    pub fn reinterpret_data<O: AsRef<[u8]>, F: FnOnce(T) -> O>(self, f: F) -> Transmit<O> {
        Transmit {
            data: f(self.data),
            transport: self.transport,
            from: self.from,
            to: self.to,
        }
    }
}

impl Transmit<Data<'_>> {
    /// Construct a new owned [`Transmit`] from a provided [`Transmit`].
    pub fn into_owned<'b>(self) -> Transmit<Data<'b>> {
        self.reinterpret_data(|data| data.into_owned())
    }
}

/// Return value for [`StunRequest::poll`].
#[derive(Debug)]
enum StunRequestPollRet {
    /// Wait until the specified time has passed.
    WaitUntil(Instant),
    /// The request has been cancelled and will not make further progress.
    Cancelled,
    /// The request timed out.
    TimedOut,
}

#[derive(Debug)]
struct RequestRto {
    initial: Duration,
    max: Duration,
    retransmits: u32,
    last_retransmit: Duration,
}

impl Default for RequestRto {
    fn default() -> Self {
        Self {
            initial: Duration::from_millis(500),
            max: Duration::MAX,
            retransmits: 7,
            last_retransmit: Duration::from_millis(8),
        }
    }
}

impl RequestRto {
    fn calculate_timeouts(&self, transport: TransportType) -> (Vec<Duration>, Duration) {
        match transport {
            TransportType::Udp => {
                let timeouts = (0..self.retransmits.max(1) - 1)
                    .map(|i| (self.initial * 2u32.pow(i)).min(self.max))
                    .collect::<Vec<_>>();
                (timeouts, self.last_retransmit)
            }
            TransportType::Tcp => {
                let timeouts = vec![];
                let last_retransmit_timeout = self.last_retransmit
                    + (0..self.retransmits.max(1) - 1).fold(Duration::ZERO, |acc, i| {
                        acc + (self.initial * 2u32.pow(i)).min(self.max)
                    });
                (timeouts, last_retransmit_timeout)
            }
        }
    }
}

#[derive(Debug)]
struct StunRequestState {
    transaction_id: TransactionId,
    request_integrity: Option<IntegrityAlgorithm>,
    bytes: Vec<u8>,
    transport: TransportType,
    from: SocketAddr,
    to: SocketAddr,
    timeouts: Vec<Duration>,
    last_retransmit_timeout: Duration,
    recv_cancelled: bool,
    send_cancelled: bool,
    timeout_i: usize,
    first_send_time: Option<Instant>,
    last_send_time: Option<Instant>,
}

impl StunRequestState {
    #[allow(clippy::too_many_arguments)]
    fn new<T: AsRef<[u8]>>(
        request: T,
        transport: TransportType,
        from: SocketAddr,
        to: SocketAddr,
        transaction_id: TransactionId,
        integrity_algorithm: Option<IntegrityAlgorithm>,
        timeouts: Vec<Duration>,
        last_retransmit_timeout: Duration,
    ) -> Self {
        let data = request.as_ref();
        Self {
            transaction_id,
            bytes: data.to_vec(),
            transport,
            from,
            to,
            request_integrity: integrity_algorithm,
            timeouts,
            timeout_i: 0,
            last_retransmit_timeout,
            recv_cancelled: false,
            send_cancelled: false,
            first_send_time: None,
            last_send_time: None,
        }
    }

    #[tracing::instrument(skip(self, now), level = "trace")]
    fn next_send_time(&self, now: Instant) -> Option<Instant> {
        let Some(last_send) = self.last_send_time else {
            trace!("not sent yet -> send immediately");
            return Some(now);
        };
        if self.timeout_i >= self.timeouts.len() {
            let next_send = last_send + self.last_retransmit_timeout;
            trace!("final retransmission, final timeout ends at {next_send:?}");
            if next_send > now {
                return Some(next_send);
            }
            return None;
        }
        let next_send = last_send + self.timeouts[self.timeout_i];
        Some(next_send)
    }

    #[tracing::instrument(
        name = "stun_request_poll"
        level = "debug",
        ret,
        skip(self, now),
        fields(transaction_id = %self.transaction_id),
    )]
    fn poll(&mut self, now: Instant) -> StunRequestPollRet {
        if self.recv_cancelled {
            return StunRequestPollRet::Cancelled;
        }
        // TODO: account for TCP connect in timeout
        let Some(next_send) = self.next_send_time(now) else {
            return StunRequestPollRet::TimedOut;
        };
        if next_send >= now {
            if self.send_cancelled && self.timeout_i >= self.timeouts.len() {
                // this cancellation may need a different value
                return StunRequestPollRet::Cancelled;
            }
            return StunRequestPollRet::WaitUntil(next_send);
        }
        StunRequestPollRet::WaitUntil(now)
    }

    #[tracing::instrument(
        name = "stun_request_poll_transmit",
        skip(self, now),
        fields(transaction_id = %self.transaction_id)
    )]
    fn poll_transmit(&mut self, now: Instant) -> Option<Transmit<&[u8]>> {
        if self.recv_cancelled {
            return None;
        };
        let next_send = self.next_send_time(now)?;

        if next_send > now {
            return None;
        }
        if self.last_send_time.is_some() {
            self.timeout_i += 1;
        }
        if self.first_send_time.is_none() {
            self.first_send_time = Some(now);
        }
        self.last_send_time = Some(now);
        if self.send_cancelled {
            return None;
        };
        trace!(
            "sending {} bytes over {:?} from {:?} to {:?}",
            self.bytes.len(),
            self.transport,
            self.from,
            self.to
        );
        Some(send_data(
            self.transport,
            self.bytes.as_slice(),
            self.from,
            self.to,
        ))
    }
}

/// A STUN Request.
#[derive(Debug, Clone)]
pub struct StunRequest<'a> {
    agent: &'a StunAgent,
    peer_address: SocketAddr,
    request_integrity: Option<IntegrityAlgorithm>,
}

impl StunRequest<'_> {
    /// The remote address the request is sent to.
    pub fn peer_address(&self) -> SocketAddr {
        self.peer_address
    }

    /// The integrity algorithm present on the request.
    pub fn integrity(&self) -> Option<IntegrityAlgorithm> {
        self.request_integrity
    }

    /// The [`StunAgent`] this request is being sent with.
    pub fn agent(&self) -> &StunAgent {
        self.agent
    }
}

/// A STUN Request.
#[derive(Debug)]
pub struct StunRequestMut<'a> {
    agent: &'a mut StunAgent,
    transaction_id: TransactionId,
    peer_address: SocketAddr,
    request_integrity: Option<IntegrityAlgorithm>,
}

impl StunRequestMut<'_> {
    /// The remote address the request is sent to.
    pub fn peer_address(&self) -> SocketAddr {
        self.peer_address
    }

    /// The integrity algorithm present on the request.
    pub fn integrity(&self) -> Option<IntegrityAlgorithm> {
        self.request_integrity
    }

    /// Do not retransmit further. This will still allow for a reply to occur within the configured
    /// timeouts, but will never send a retransmission. If no response is received, this will cause
    /// [`StunAgent::poll()`] to return [`StunAgentPollRet::TransactionCancelled`] for this request.
    pub fn cancel_retransmissions(&mut self) {
        if let Some(state) = self.agent.mut_request_state(self.transaction_id) {
            state.send_cancelled = true;
        }
    }

    /// Do not wait for any kind of response. This will cause [`StunAgent::poll()`] to return
    /// [`StunAgentPollRet::TransactionCancelled`] for this request.
    pub fn cancel(&mut self) {
        if let Some(state) = self.agent.mut_request_state(self.transaction_id) {
            state.send_cancelled = true;
            state.recv_cancelled = true;
        }
    }

    /// The [`StunAgent`] this request is being sent with.
    pub fn agent(&self) -> &StunAgent {
        self.agent
    }

    /// The mutable [`StunAgent`] this request is being sent with.
    pub fn mut_agent(&mut self) -> &mut StunAgent {
        self.agent
    }

    /// Configure the timeouts and retransmissions for the STUN request.
    ///
    /// This the same as calling `[configure_timeout_with_max`] with a `max` of `Duration::MAX`.
    pub fn configure_timeout(
        &mut self,
        initial_rto: Duration,
        retransmits: u32,
        last_retransmit_timeout: Duration,
    ) {
        self.configure_timeout_with_max(
            initial_rto,
            retransmits,
            last_retransmit_timeout,
            Duration::MAX,
        );
    }

    /// Configure the timeouts and retransmissions for the STUN request.
    ///
    /// - `initial` - the initial time between consecutive transmissions. If 0, or 1, then only a
    ///   single request will be performed.
    /// - `max` - the maximum amount of time between consecutive retransmits.
    /// - `retransmits` - the total number of transmissions of the request.
    /// - `final_retransmit_timeout` - the amount of time after the final transmission to wait
    ///   for a response before considering the request as having timed out.
    ///
    /// As specified in RFC 8489, `initial_rto` should be >= 500ms (unless specific information is
    /// available on the RTT, `max` is `Duration::MAX`, `retransmits` has a default value of 7,
    /// and `last_retransmit_timeout` should be `16 * initial_rto`.
    ///
    /// STUN transactions over TCP will only send a single request and have a timeout of the sum of
    /// the timeouts of a UDP transaction.
    pub fn configure_timeout_with_max(
        &mut self,
        initial_rto: Duration,
        retransmits: u32,
        last_retransmit_timeout: Duration,
        max_rto: Duration,
    ) {
        if let Some(state) = self.agent.mut_request_state(self.transaction_id) {
            let (timeouts, final_wait) = RequestRto {
                initial: initial_rto,
                max: max_rto,
                retransmits,
                last_retransmit: last_retransmit_timeout,
            }
            .calculate_timeouts(state.transport);
            state.timeouts = timeouts;
            state.last_retransmit_timeout = final_wait;
        }
    }
}

/// STUN errors.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum StunError {
    /// The operation is already in progress.
    #[error("The operation is already in progress")]
    AlreadyInProgress,
    /// A resource was not found.
    #[error("A required resource could not be found")]
    ResourceNotFound,
    /// An operation timed out without a response.
    #[error("An operation timed out")]
    TimedOut,
    /// Unexpected data was received or an operation is not allowed at this time.
    #[error("Unexpected data was received")]
    ProtocolViolation,
    /// An operation was cancelled.
    #[error("Operation was aborted")]
    Aborted,
    /// A parsing error. The contained error contains more details.
    #[error("{}", .0)]
    ParseError(StunParseError),
    /// A writing error. The contained error contains more details.
    #[error("{}", .0)]
    WriteError(StunWriteError),
}

impl From<StunParseError> for StunError {
    fn from(e: StunParseError) -> Self {
        StunError::ParseError(e)
    }
}

impl From<StunWriteError> for StunError {
    fn from(e: StunWriteError) -> Self {
        StunError::WriteError(e)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use alloc::string::String;
    use tracing::error;

    use crate::auth::ShortTermAuth;

    use super::*;

    #[test]
    fn agent_getters_setters() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();
        let agent = StunAgent::builder(TransportType::Udp, local_addr)
            .remote_addr(remote_addr)
            .build();

        assert_eq!(agent.transport(), TransportType::Udp);
        assert_eq!(agent.local_addr(), local_addr);
        assert_eq!(agent.remote_addr(), Some(remote_addr));
    }

    #[test]
    fn request() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .remote_addr(remote_addr)
            .build();
        let now = Instant::ZERO;

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let transmit = agent
            .send_request(msg.finish(), remote_addr, now)
            .unwrap()
            .into_owned();
        let request = agent.request_transaction(transaction_id).unwrap();
        assert!(request.integrity().is_none());
        assert_eq!(transmit.transport, TransportType::Udp);
        assert_eq!(transmit.from, local_addr);
        assert_eq!(transmit.to, remote_addr);
        let request = Message::from_bytes(&transmit.data).unwrap();
        let response = Message::builder_error(&request, MessageWriteVec::new());
        let resp_data = response.finish();
        let response = Message::from_bytes(&resp_data).unwrap();
        assert!(agent.handle_stun_message_with_time(&response, remote_addr, now));
        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());

        let ret = agent.poll(now);
        assert!(matches!(ret, StunAgentPollRet::WaitUntil(_)));
    }

    #[test]
    fn indication_with_invalid_response() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let now = Instant::ZERO;

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .remote_addr(remote_addr)
            .build();
        let transaction_id = TransactionId::generate();
        let msg = Message::builder(
            MessageType::from_class_method(MessageClass::Indication, BINDING),
            transaction_id,
            MessageWriteVec::new(),
        );
        let transmit = agent
            .send(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();
        assert_eq!(transmit.transport, TransportType::Udp);
        assert_eq!(transmit.from, local_addr);
        assert_eq!(transmit.to, remote_addr);
        let _indication = Message::from_bytes(&transmit.data).unwrap();
        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());
        // you should definitely never do this ;). Indications should never get replies.
        let response = Message::builder(
            MessageType::from_class_method(MessageClass::Error, BINDING),
            transaction_id,
            MessageWriteVec::new(),
        );
        let resp_data = response.finish();
        let response = Message::from_bytes(&resp_data).unwrap();
        // response without a request is dropped.
        assert!(!agent.handle_stun_message_with_time(&response, remote_addr, now))
    }

    #[test]
    fn request_with_credentials() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();
        let now = Instant::ZERO;

        let mut auth = ShortTermAuth::new();
        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();
        let credentials = ShortTermCredentials::new(String::from("local_password"));
        auth.set_credentials(credentials.clone(), IntegrityAlgorithm::Sha1);

        // unvalidated peer data should be dropped
        assert!(!agent.is_validated_peer(remote_addr));

        let mut msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        msg.add_message_integrity(&credentials.clone().into(), IntegrityAlgorithm::Sha1)
            .unwrap();
        error!("send");
        let transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();
        error!("sent");

        let request = Message::from_bytes(&transmit.data).unwrap();

        error!("generate response");
        let mut response = Message::builder_success(&request, MessageWriteVec::new());
        let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
        response.add_attribute(&xor_addr).unwrap();
        response
            .add_message_integrity(&credentials.into(), IntegrityAlgorithm::Sha1)
            .unwrap();
        error!("{response:?}");

        let data = response.finish();
        error!("{data:?}");
        let response = Message::from_bytes(&data).unwrap();
        error!("{response}");
        assert_eq!(
            auth.validate_incoming_message(&response).unwrap(),
            Some(IntegrityAlgorithm::Sha1)
        );
        let request = agent
            .request_transaction(response.transaction_id())
            .unwrap();
        assert_eq!(request.integrity(), Some(IntegrityAlgorithm::Sha1));
        assert!(agent.handle_stun_message_with_time(&response, remote_addr, now));

        assert_eq!(response.transaction_id(), transaction_id);
        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());
        assert!(agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn request_unanswered() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .remote_addr(remote_addr)
            .build();
        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();
        let mut now = Instant::ZERO;
        loop {
            let _ = agent.poll_transmit(now);
            match agent.poll(now) {
                StunAgentPollRet::WaitUntil(new_now) => {
                    now = new_now;
                }
                StunAgentPollRet::TransactionTimedOut(_) => break,
                _ => unreachable!(),
            }
        }
        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());

        // unvalidated peer data should be dropped
        assert!(!agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn request_custom_timeout() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .remote_addr(remote_addr)
            .build();
        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let mut now = Instant::ZERO;
        agent.send_request(msg.finish(), remote_addr, now).unwrap();
        let mut transaction = agent.mut_request_transaction(transaction_id).unwrap();
        transaction.configure_timeout_with_max(
            Duration::from_secs(1),
            4,
            Duration::from_secs(10),
            Duration::from_secs(2),
        );
        let StunAgentPollRet::WaitUntil(wait) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(wait - now, Duration::from_secs(1));
        now = wait;
        // a poll with the same instant should not busy loop
        let StunAgentPollRet::WaitUntil(wait) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(wait, now);
        let Some(_) = agent.poll_transmit(now) else {
            unreachable!();
        };
        let StunAgentPollRet::WaitUntil(wait) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(wait - now, Duration::from_secs(2));
        now = wait;
        let Some(_) = agent.poll_transmit(now) else {
            unreachable!();
        };
        let StunAgentPollRet::WaitUntil(wait) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(wait - now, Duration::from_secs(2));
        now = wait;
        let Some(_) = agent.poll_transmit(now) else {
            unreachable!();
        };
        let StunAgentPollRet::WaitUntil(wait) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(wait - now, Duration::from_secs(10));
        now = wait;
        let StunAgentPollRet::TransactionTimedOut(timed_out) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(timed_out, transaction_id);

        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());

        // unvalidated peer data should be dropped
        assert!(!agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn request_no_retransmit() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .remote_addr(remote_addr)
            .build();
        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let mut now = Instant::ZERO;
        agent.send_request(msg.finish(), remote_addr, now).unwrap();
        let mut transaction = agent.mut_request_transaction(transaction_id).unwrap();
        transaction.configure_timeout(Duration::from_secs(1), 0, Duration::from_secs(10));
        let StunAgentPollRet::WaitUntil(wait) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(wait - now, Duration::from_secs(10));
        now = wait;
        let StunAgentPollRet::TransactionTimedOut(timed_out) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(timed_out, transaction_id);

        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());

        // unvalidated peer data should be dropped
        assert!(!agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn request_tcp_custom_timeout() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let mut agent = StunAgent::builder(TransportType::Tcp, local_addr)
            .remote_addr(remote_addr)
            .request_retransmits(
                Duration::from_secs(1),
                Duration::from_secs(2),
                4,
                Duration::from_secs(3),
            )
            .build();
        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let mut now = Instant::ZERO;
        agent.send_request(msg.finish(), remote_addr, now).unwrap();
        let StunAgentPollRet::WaitUntil(wait) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(wait - now, Duration::from_secs(1 + 2 + 2 + 3));
        now = wait;
        let StunAgentPollRet::TransactionTimedOut(timed_out) = agent.poll(now) else {
            unreachable!();
        };
        assert_eq!(timed_out, transaction_id);

        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());

        // unvalidated peer data should be dropped
        assert!(!agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn request_without_credentials() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();
        let now = Instant::ZERO;

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();

        // unvalidated peer data should be dropped
        assert!(!agent.is_validated_peer(remote_addr));

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();

        let request = Message::from_bytes(&transmit.data).unwrap();

        let mut response = Message::builder_success(&request, MessageWriteVec::new());
        let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
        response.add_attribute(&xor_addr).unwrap();

        let data = response.finish();
        let to = transmit.to;
        trace!("data: {data:?}");
        let response = Message::from_bytes(&data).unwrap();
        let request = agent
            .request_transaction(response.transaction_id())
            .unwrap();
        assert_eq!(request.integrity(), None);
        assert!(agent.handle_stun_message_with_time(&response, to, now));
        assert_eq!(response.transaction_id(), transaction_id);
        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());
        assert!(agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn response_with_incorrect_credentials() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();
        let now = Instant::ZERO;

        let mut auth = ShortTermAuth::new();
        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();
        let credentials = ShortTermCredentials::new(String::from("local_password"));
        let wrong_credentials = ShortTermCredentials::new(String::from("wrong_password"));
        auth.set_credentials(credentials.clone(), IntegrityAlgorithm::Sha1);

        let mut msg = Message::builder_request(BINDING, MessageWriteVec::new());
        msg.add_message_integrity(&credentials.clone().into(), IntegrityAlgorithm::Sha1)
            .unwrap();
        let transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();
        let data = transmit.data;

        let request = Message::from_bytes(&data).unwrap();

        let mut response = Message::builder_success(&request, MessageWriteVec::new());
        let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
        response.add_attribute(&xor_addr).unwrap();
        // wrong credentials, should be `remote_credentials`
        response
            .add_message_integrity(&wrong_credentials.into(), IntegrityAlgorithm::Sha1)
            .unwrap();

        let data = response.finish();
        let response = Message::from_bytes(&data).unwrap();
        // reply is ignored as it does not have credentials
        let request = agent
            .request_transaction(response.transaction_id())
            .unwrap();
        assert_eq!(request.integrity(), Some(IntegrityAlgorithm::Sha1));
        assert!(matches!(
            auth.validate_incoming_message(&response),
            Err(ValidateError::IntegrityFailed)
        ));

        // unvalidated peer data should be dropped
        assert!(!agent.is_validated_peer(remote_addr));

        // however signifying success will cause peer validation to succeed
        assert!(agent.handle_stun_message_with_time(&response, remote_addr, now));
        assert!(!agent.handle_stun_message_with_time(&response, remote_addr, now));
        assert!(agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn duplicate_response_ignored() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();
        let now = Instant::ZERO;

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();
        assert!(!agent.is_validated_peer(remote_addr));

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();
        let data = transmit.data;

        let request = Message::from_bytes(&data).unwrap();

        let mut response = Message::builder_success(&request, MessageWriteVec::new());
        let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
        response.add_attribute(&xor_addr).unwrap();

        let data = response.finish();
        let to = transmit.to;
        let response = Message::from_bytes(&data).unwrap();
        assert!(agent.handle_stun_message_with_time(&response, to, now));

        let response = Message::from_bytes(&data).unwrap();
        assert!(!agent.handle_stun_message_with_time(&response, to, now));
    }

    #[test]
    fn request_cancel() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let _transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();

        let mut request = agent.mut_request_transaction(transaction_id).unwrap();
        assert_eq!(request.integrity(), None);
        assert_eq!(request.agent().local_addr(), local_addr);
        assert_eq!(request.mut_agent().local_addr(), local_addr);
        assert_eq!(request.peer_address(), remote_addr);
        request.cancel();

        let ret = agent.poll(Instant::ZERO);
        let StunAgentPollRet::TransactionCancelled(_request) = ret else {
            unreachable!();
        };
        assert_eq!(transaction_id, transaction_id);
        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());
        assert!(!agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn request_cancel_send() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let _transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();

        let mut request = agent.mut_request_transaction(transaction_id).unwrap();
        assert_eq!(request.integrity(), None);
        assert_eq!(request.agent().local_addr(), local_addr);
        assert_eq!(request.mut_agent().local_addr(), local_addr);
        assert_eq!(request.peer_address(), remote_addr);
        request.cancel_retransmissions();

        let mut now = Instant::ZERO;
        let start = now;
        loop {
            match agent.poll(now) {
                StunAgentPollRet::WaitUntil(new_now) => {
                    assert_ne!(new_now, now);
                    now = new_now;
                }
                StunAgentPollRet::TransactionCancelled(_) => break,
                _ => unreachable!(),
            }
            let _ = agent.poll_transmit(now);
        }
        assert!(now - start > Duration::from_secs(20));
        assert!(agent.request_transaction(transaction_id).is_none());
        assert!(agent.mut_request_transaction(transaction_id).is_none());
        assert!(!agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn request_duplicate() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();
        let now = Instant::ZERO;

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let msg = msg.finish();
        let transmit = agent
            .send_request(msg.clone(), remote_addr, Instant::ZERO)
            .unwrap();
        let to = transmit.to;
        let request = Message::from_bytes(&transmit.data).unwrap();

        let mut response = Message::builder_success(&request, MessageWriteVec::new());
        let xor_addr = XorMappedAddress::new(transmit.from, transaction_id);
        response.add_attribute(&xor_addr).unwrap();

        assert!(matches!(
            agent.send_request(msg, remote_addr, Instant::ZERO),
            Err(StunError::AlreadyInProgress)
        ));

        // the original transaction should still exist
        let request = agent.request_transaction(transaction_id).unwrap();
        assert_eq!(request.peer_address(), remote_addr);

        let data = response.finish();
        let response = Message::from_bytes(&data).unwrap();
        assert!(agent.handle_stun_message_with_time(&response, to, now));

        assert!(agent.is_validated_peer(to));
    }

    #[test]
    fn incoming_request() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();
        let now = Instant::ZERO;

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let data = msg.finish();
        let stun = Message::from_bytes(&data).unwrap();
        error!("{stun:?}");
        assert!(agent.handle_stun_message_with_time(&stun, remote_addr, now));
        agent.validated_peer(remote_addr);
        assert!(agent.is_validated_peer(remote_addr));
    }

    #[test]
    fn tcp_request() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let mut agent = StunAgent::builder(TransportType::Tcp, local_addr)
            .remote_addr(remote_addr)
            .build();

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();
        assert_eq!(transmit.transport, TransportType::Tcp);
        assert_eq!(transmit.from, local_addr);
        assert_eq!(transmit.to, remote_addr);

        let request = Message::from_bytes(&transmit.data).unwrap();
        assert_eq!(request.transaction_id(), transaction_id);
    }

    #[test]
    fn transmit_into_owned() {
        let data = [0x10, 0x20];
        let transport = TransportType::Udp;
        let from = "127.0.0.1:1000".parse().unwrap();
        let to = "127.0.0.1:2000".parse().unwrap();
        let transmit = Transmit::new(Data::from(data.as_ref()), TransportType::Udp, from, to);
        let owned = transmit.into_owned();
        assert_eq!(owned.data.as_ref(), data.as_ref());
        assert_eq!(owned.transport, transport);
        assert_eq!(owned.from, from);
        assert_eq!(owned.to, to);
        error!("{owned}");
    }

    #[test]
    fn transmit_display() {
        let data = [0x10, 0x20];
        let from = "127.0.0.1:1000".parse().unwrap();
        let to = "127.0.0.1:2000".parse().unwrap();
        assert_eq!(
            alloc::format!(
                "{}",
                Transmit::new(Data::from(data.as_ref()), TransportType::Udp, from, to)
            ),
            String::from("Transmit(UDP: 127.0.0.1:1000 -> 127.0.0.1:2000 of 2 bytes)")
        );
    }

    #[test]
    fn request_retransmits() {
        let _log = crate::tests::test_init_log();
        let rto = RequestRto {
            initial: Duration::from_millis(1),
            max: Duration::MAX,
            retransmits: 0,
            last_retransmit: Duration::from_secs(1),
        };
        let (timeouts, last_transmit_timeout) = rto.calculate_timeouts(TransportType::Udp);
        assert_eq!(timeouts, vec![]);
        assert_eq!(last_transmit_timeout, Duration::from_secs(1));
        let (timeouts, last_transmit_timeout) = rto.calculate_timeouts(TransportType::Tcp);
        assert_eq!(timeouts, vec![]);
        assert_eq!(last_transmit_timeout, Duration::from_secs(1));

        let rto = RequestRto {
            initial: Duration::from_millis(1),
            max: Duration::MAX,
            retransmits: 1,
            last_retransmit: Duration::from_secs(1),
        };
        let (timeouts, last_transmit_timeout) = rto.calculate_timeouts(TransportType::Udp);
        assert_eq!(timeouts, vec![]);
        assert_eq!(last_transmit_timeout, Duration::from_secs(1));
        let (timeouts, last_transmit_timeout) = rto.calculate_timeouts(TransportType::Tcp);
        assert_eq!(timeouts, vec![]);
        assert_eq!(last_transmit_timeout, Duration::from_secs(1));

        let rto = RequestRto {
            initial: Duration::from_millis(1),
            max: Duration::MAX,
            retransmits: 2,
            last_retransmit: Duration::from_secs(1),
        };
        let (timeouts, last_transmit_timeout) = rto.calculate_timeouts(TransportType::Udp);
        assert_eq!(timeouts, vec![Duration::from_millis(1)]);
        assert_eq!(last_transmit_timeout, Duration::from_secs(1));
        let (timeouts, last_transmit_timeout) = rto.calculate_timeouts(TransportType::Tcp);
        assert_eq!(timeouts, vec![]);
        assert_eq!(
            last_transmit_timeout,
            Duration::from_secs(1) + Duration::from_millis(1)
        );
    }

    #[test]
    fn stats_send_receive_request() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let now = Instant::ZERO;

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .stats(true)
            .build();

        assert!(agent.stats().is_some());
        let stats = agent.stats().unwrap();
        assert_eq!(stats.requests_sent(), 0);
        assert_eq!(stats.responses_received(), 0);
        assert_eq!(stats.bytes_sent(), 0);
        assert_eq!(stats.rtt_count(), 0);

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transmit = agent.send_request(msg.finish(), remote_addr, now).unwrap();
        let request_data = transmit.data.as_ref().to_vec();

        let stats = agent.stats().unwrap();
        assert_eq!(stats.requests_sent(), 1);
        assert_eq!(stats.bytes_sent(), request_data.len() as u64);

        // Build and handle response with RTT tracking
        let request = Message::from_bytes(&request_data).unwrap();
        let mut response = Message::builder_success(&request, MessageWriteVec::new());
        let xor_addr =
            XorMappedAddress::new("10.0.0.1:12345".parse().unwrap(), request.transaction_id());
        response.add_attribute(&xor_addr).unwrap();
        let response_data = response.finish();
        let response = Message::from_bytes(&response_data).unwrap();

        let receive_time = Duration::from_millis(10);
        assert!(agent.poll_transmit(now + receive_time).is_none());
        assert!(agent.handle_stun_message_with_time(&response, remote_addr, now + receive_time));

        let stats = agent.stats().unwrap();
        assert_eq!(stats.responses_received(), 1);
        assert_eq!(stats.bytes_received(), response_data.len() as u64);
        assert_eq!(stats.rtt_count(), 1);
        assert_eq!(stats.rtt_min(), Some(receive_time));
        assert_eq!(stats.rtt_max(), Some(receive_time));
        assert_eq!(stats.rtt_average(), Some(receive_time));

        assert!(agent.handle_stun_message_with_time(&request, remote_addr, now));
        let stats = agent.stats().unwrap();
        assert_eq!(stats.requests_sent(), 1);
        assert_eq!(
            stats.bytes_received(),
            (request_data.len() + response_data.len()) as u64
        );

        agent.send(&response, remote_addr, now).unwrap();
        let stats = agent.stats().unwrap();
        assert_eq!(stats.responses_sent(), 1);
        assert_eq!(
            stats.bytes_sent(),
            (request_data.len() + response_data.len()) as u64
        );
    }

    #[test]
    fn stats_disabled_by_default() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let agent = StunAgent::builder(TransportType::Udp, local_addr).build();
        assert!(agent.stats().is_none());
    }

    #[test]
    fn stats_timeout_and_cancel() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .stats(true)
            .build();

        // Send a request that will timeout
        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();

        let mut now = Instant::ZERO;
        loop {
            let _ = agent.poll_transmit(now);
            match agent.poll(now) {
                StunAgentPollRet::WaitUntil(new_now) => now = new_now,
                StunAgentPollRet::TransactionTimedOut(_) => break,
                _ => unreachable!(),
            }
        }

        let stats = agent.stats().unwrap();
        assert_eq!(stats.transactions_timed_out(), 1);

        // Send another request that we'll cancel
        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let cancel_id = msg.transaction_id();
        agent.send_request(msg.finish(), remote_addr, now).unwrap();

        let mut req = agent.mut_request_transaction(cancel_id).unwrap();
        req.cancel();

        let ret = agent.poll(now);
        assert!(matches!(ret, StunAgentPollRet::TransactionCancelled(_)));

        let stats = agent.stats().unwrap();
        assert_eq!(stats.transactions_cancelled(), 1);
    }

    #[test]
    fn stats_rtt_only_no_retransmit() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .request_retransmits(
                Duration::from_millis(1),
                Duration::MAX,
                2,
                Duration::from_millis(1),
            )
            .stats(true)
            .build();

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();
        let request_data = transmit.data.as_ref().to_vec();
        let from_addr = transmit.from;

        // Retransmit once by polling
        let retransmit_time = Instant::ZERO + Duration::from_millis(1);
        let _retransmit = agent.poll_transmit(retransmit_time).unwrap();

        // Now build and handle response
        let request = Message::from_bytes(&request_data).unwrap();
        let mut response = Message::builder_success(&request, MessageWriteVec::new());
        let xor_addr = XorMappedAddress::new(from_addr, request.transaction_id());
        response.add_attribute(&xor_addr).unwrap();
        let response_data = response.finish();
        let response = Message::from_bytes(&response_data).unwrap();

        let receive_time = retransmit_time + Duration::from_millis(10);
        assert!(agent.handle_stun_message_with_time(&response, remote_addr, receive_time));

        let stats = agent.stats().unwrap();
        assert_eq!(stats.responses_received(), 1);
        assert_eq!(stats.rtt_count(), 0);
    }

    #[test]
    fn stats_error_response() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let now = Instant::ZERO;

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .stats(true)
            .build();

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();
        let request_data = transmit.data.as_ref().to_vec();

        let request = Message::from_bytes(&request_data).unwrap();
        let mut error_response = Message::builder_error(&request, MessageWriteVec::new());
        let error_code = ErrorCode::builder(ErrorCode::BAD_REQUEST).build().unwrap();
        error_response.add_attribute(&error_code).unwrap();
        let error_data = error_response.finish();
        let error_msg = Message::from_bytes(&error_data).unwrap();

        assert!(agent.handle_stun_message_with_time(&error_msg, remote_addr, now));

        let stats = agent.stats().unwrap();
        assert_eq!(stats.responses_received(), 1);
        assert_eq!(stats.bytes_received(), error_data.len() as u64);
    }

    #[test]
    fn stats_send_receive_indication() {
        let _log = crate::tests::test_init_log();
        let local_addr = "127.0.0.1:2000".parse().unwrap();
        let remote_addr = "127.0.0.1:1000".parse().unwrap();
        let mut agent = StunAgent::builder(TransportType::Udp, local_addr)
            .stats(true)
            .build();
        let now = Instant::ZERO;

        // Send a request that will timeout
        let msg = Message::builder_indication(BINDING, MessageWriteVec::new());
        let msg_data = msg.finish();
        agent
            .send(msg_data.clone(), remote_addr, Instant::ZERO)
            .unwrap();

        let stats = agent.stats().unwrap();
        assert_eq!(stats.indications_sent(), 1);
        assert_eq!(stats.bytes_sent(), msg_data.len() as u64);

        let msg = Message::from_bytes(&msg_data).unwrap();
        assert!(agent.handle_stun_message_with_time(&msg, remote_addr, now));

        let stats = agent.stats().unwrap();
        assert_eq!(stats.indications_received(), 1);
        assert_eq!(stats.bytes_received(), msg_data.len() as u64);
    }

    #[test]
    fn request_removal_from_mut_agent() {
        let _log = crate::tests::test_init_log();
        let local_addr = "10.0.0.1:12345".parse().unwrap();
        let remote_addr = "10.0.0.2:3478".parse().unwrap();
        let now = Instant::ZERO;

        let mut agent = StunAgent::builder(TransportType::Udp, local_addr).build();

        let msg = Message::builder_request(BINDING, MessageWriteVec::new());
        let transaction_id = msg.transaction_id();
        let transmit = agent
            .send_request(msg.finish(), remote_addr, Instant::ZERO)
            .unwrap();

        let request = Message::from_bytes(&transmit.data).unwrap();

        let mut response = Message::builder_success(&request, MessageWriteVec::new());
        let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
        response.add_attribute(&xor_addr).unwrap();

        let data = response.finish();
        let to = transmit.to;
        trace!("data: {data:?}");
        let response = Message::from_bytes(&data).unwrap();
        assert_eq!(response.transaction_id(), transaction_id);

        // handling the success response while holding the outstanding request should not panic and
        // still result in valid retrieval of configuration.
        let mut request = agent
            .mut_request_transaction(response.transaction_id())
            .unwrap();
        assert_eq!(request.integrity(), None);
        assert_eq!(request.peer_address(), to);
        assert!(request
            .mut_agent()
            .handle_stun_message_with_time(&response, to, now));

        assert_eq!(request.integrity(), None);
        assert_eq!(request.peer_address(), to);
        // attempting to modify the request is ignored.
        request.cancel_retransmissions();
        request.cancel();
        request.configure_timeout(Duration::from_secs(1), 3, Duration::from_secs(2));

        // the request is no longer stored by the agent.
        assert!(request
            .agent()
            .request_transaction(transaction_id)
            .is_none());
        assert!(request
            .mut_agent()
            .mut_request_transaction(transaction_id)
            .is_none());
    }
}