zlayer-overlay 0.12.6

//! Minimal STUN client per RFC 5389
//!
//! Implements just enough of the STUN Binding protocol to discover our
//! server-reflexive (public) address and detect NAT behavior (endpoint-independent
//! vs symmetric). No external STUN crate required.

use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
use std::str::FromStr;

use rand::Rng;
use socket2::{Domain, Protocol, Socket, Type};
use thiserror::Error;
use tokio::net::UdpSocket;

use super::config::StunServerConfig;
use crate::dns::{build_forward_resolver, resolve_upstreams, DnsConfig, RESOLV_CONF_PATH};

// ---- STUN protocol constants (RFC 5389) ------------------------------------

const MAGIC_COOKIE: u32 = 0x2112_A442;
const BINDING_REQUEST: u16 = 0x0001;
const BINDING_RESPONSE: u16 = 0x0101;
const ATTR_MAPPED_ADDRESS: u16 = 0x0001;
const ATTR_XOR_MAPPED_ADDRESS: u16 = 0x0020;
const STUN_HEADER_SIZE: usize = 20;

// ---- Error type ------------------------------------------------------------

/// Errors specific to STUN operations.
#[derive(Debug, Error)]
pub enum StunError {
    /// Response too short to contain a valid STUN header.
    #[error("STUN response too short ({0} bytes, need at least {STUN_HEADER_SIZE})")]
    TooShort(usize),

    /// The response message type is not a Binding Response (0x0101).
    #[error("Expected Binding Response (0x0101), got 0x{0:04X}")]
    NotBindingResponse(u16),

    /// The magic cookie field does not match `0x2112A442`.
    #[error("Bad magic cookie in STUN response")]
    BadMagicCookie,

    /// The 12-byte transaction ID does not match what we sent.
    #[error("Transaction ID mismatch")]
    TransactionMismatch,

    /// Neither XOR-MAPPED-ADDRESS nor MAPPED-ADDRESS was found.
    #[error("No mapped address attribute in STUN response")]
    NoMappedAddress,

    /// An I/O error from the UDP socket.
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    /// Timed out waiting for a STUN response.
    #[error("STUN query timed out")]
    Timeout,

    /// DNS resolution of a STUN server hostname failed.
    #[error("DNS resolution failed for '{0}'")]
    DnsResolution(String),

    /// No STUN servers configured.
    #[error("No STUN servers configured")]
    NoServers,
}

// ---- Public types ----------------------------------------------------------

/// A reflexive (public) address discovered via a STUN server.
#[derive(Debug, Clone)]
pub struct ReflexiveAddress {
    /// The public address:port the STUN server observed.
    pub address: SocketAddr,
    /// Human-readable label of the STUN server that reported this.
    pub server: String,
}

/// Observed NAT mapping behavior.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NatBehavior {
    /// All STUN servers saw the same reflexive port (or we only queried one).
    /// Hole punching is likely to succeed.
    EndpointIndependent,
    /// Different STUN servers saw different reflexive ports.
    /// Hole punching will probably fail; a relay is needed.
    Symmetric,
}

// ---- STUN client -----------------------------------------------------------

/// A lightweight STUN client that discovers reflexive addresses.
pub struct StunClient {
    servers: Vec<StunServerConfig>,
}

impl StunClient {
    /// Create a new STUN client with the given server list.
    #[must_use]
    pub fn new(servers: Vec<StunServerConfig>) -> Self {
        Self { servers }
    }

    /// Query a single STUN server for our reflexive address.
    ///
    /// Creates an ephemeral UDP socket, sends a Binding Request, and waits up
    /// to 3 seconds for a Binding Response.
    ///
    /// # Interface pinning
    ///
    /// When `egress_iface` is a non-empty interface name, the UDP socket is
    /// pinned to that physical NIC via [`crate::egress::bind_to_device`] before
    /// any packet is sent. This stops the STUN probe from egressing through a
    /// mesh VPN tunnel (netbird `wt0`, Tailscale `utun`, …) that has captured
    /// the host's default route — which would otherwise make the discovered
    /// reflexive address the *tunnel's* mapping rather than the real NIC's.
    /// Pinning is Unix-only (the `AsRawFd` bound is unsatisfiable on Windows)
    /// and degrades to a warning on `PermissionDenied` / `InterfaceNotFound` /
    /// any other error rather than failing the probe (per `egress.rs`'s
    /// documented contract).
    ///
    /// # Errors
    ///
    /// Returns [`StunError::Io`] on socket failures, [`StunError::Timeout`] if
    /// no response arrives within 3 seconds, or a parse error if the response
    /// is malformed.
    pub async fn query_server(
        &self,
        server_addr: SocketAddr,
        server_name: &str,
        egress_iface: &str,
    ) -> Result<ReflexiveAddress, StunError> {
        // Create UDP socket via socket2 so we can bind to port 0 (ephemeral)
        let domain = if server_addr.is_ipv4() {
            Domain::IPV4
        } else {
            Domain::IPV6
        };
        let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?;
        socket.set_nonblocking(true)?;

        // Pin to the physical egress NIC before binding/sending so the probe
        // never leaks through a mesh-VPN default route. Unix-only: the
        // `AsRawFd` bound `bind_to_device` requires is unsatisfiable on Windows.
        #[cfg(unix)]
        {
            if !egress_iface.is_empty() {
                match crate::egress::bind_to_device(&socket, egress_iface) {
                    Ok(()) => {}
                    Err(crate::error::OverlayError::PermissionDenied(msg)) => {
                        tracing::warn!(
                            interface = %egress_iface,
                            error = %msg,
                            "STUN socket could not be pinned to physical NIC (insufficient \
                             privilege); continuing unpinned"
                        );
                    }
                    Err(e) => {
                        tracing::warn!(
                            interface = %egress_iface,
                            error = %e,
                            "STUN socket could not be pinned to physical NIC; \
                             continuing unpinned"
                        );
                    }
                }
            }
        }
        #[cfg(not(unix))]
        {
            let _ = egress_iface;
        }

        let bind_addr: SocketAddr = if server_addr.is_ipv4() {
            SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0)
        } else {
            SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 0)
        };
        socket.bind(&bind_addr.into())?;

        // Convert socket2::Socket -> std::net::UdpSocket -> tokio::net::UdpSocket
        let std_socket: std::net::UdpSocket = socket.into();
        let udp = UdpSocket::from_std(std_socket)?;

        // Build Binding Request
        let mut txn_id = [0u8; 12];
        rand::rng().fill(&mut txn_id);
        let request = build_binding_request(&txn_id);

        udp.send_to(&request, server_addr).await?;

        // Wait for response with 3-second timeout
        let mut buf = [0u8; 576]; // RFC 5389 recommended minimum
        let n = tokio::time::timeout(std::time::Duration::from_secs(3), udp.recv(&mut buf))
            .await
            .map_err(|_| StunError::Timeout)?
            .map_err(StunError::Io)?;

        let reflexive = parse_binding_response(&buf[..n], &txn_id)?;

        tracing::debug!(
            server = %server_name,
            reflexive = %reflexive,
            "STUN query succeeded"
        );

        Ok(ReflexiveAddress {
            address: reflexive,
            server: server_name.to_string(),
        })
    }

    /// Query all configured STUN servers in parallel.
    ///
    /// Returns all successfully discovered reflexive addresses and the inferred
    /// NAT behavior. If all servers fail, returns the last error.
    ///
    /// # Surviving a netbird-style host takeover
    ///
    /// This function is hardened against a mesh VPN (netbird, Tailscale, …) that
    /// has captured the host: both the default route *and* the system DNS
    /// resolver (a `~.` systemd-resolved catch-all). Two defenses run per pass:
    ///
    /// 1. **Egress pinning.** [`crate::egress::detect_physical_egress`] resolves
    ///    the real NIC once up front; its name is threaded into every
    ///    [`Self::query_server`] call so each STUN probe egresses through the
    ///    physical interface rather than the mesh tunnel. If detection fails
    ///    (empty interface name or error), pinning is skipped with a warning.
    /// 2. **Resolver bypass.** STUN server hostnames are resolved through the
    ///    *same* upstream resolvers the overlay DNS forwarder uses
    ///    ([`build_forward_resolver`] over auto-detected, stub-filtered
    ///    upstreams) instead of the poisoned host resolver. Literal
    ///    `ip:port` addresses short-circuit DNS entirely. Only if the upstream
    ///    resolver errors do we fall back to [`tokio::net::lookup_host`] (the
    ///    host resolver may still work, and a broken fallback is no worse than
    ///    the pre-hardening behaviour). The resolver is built once per pass.
    ///
    /// # Errors
    ///
    /// Returns [`StunError::NoServers`] if no servers are configured, or the
    /// last per-server error if every query fails.
    pub async fn discover(&self) -> Result<(Vec<ReflexiveAddress>, NatBehavior), StunError> {
        if self.servers.is_empty() {
            return Err(StunError::NoServers);
        }

        // Detect the physical egress NIC once for this discovery pass. An empty
        // interface name (or a detection error) means "pin nothing" — the
        // probes then fall back to the OS default route, which is no worse than
        // the pre-hardening behaviour.
        let egress_iface = match crate::egress::detect_physical_egress().await {
            Ok(egress) => egress.interface,
            Err(e) => {
                tracing::warn!(
                    error = %e,
                    "could not detect physical egress NIC; STUN probes will not be \
                     pinned to a physical interface"
                );
                String::new()
            }
        };

        // Build the resolver-bypass once per pass: forward STUN hostname lookups
        // through the overlay's own upstreams (stub/loopback filtered) rather
        // than the host resolver a mesh VPN may have hijacked.
        let bypass_resolver = build_bypass_resolver();

        let mut set = tokio::task::JoinSet::new();

        for server_cfg in &self.servers {
            let address_str = server_cfg.address.clone();
            let label = server_cfg
                .label
                .clone()
                .unwrap_or_else(|| server_cfg.address.clone());

            // Resolve hostname -> SocketAddr without trusting the host resolver.
            let Some(resolved) = resolve_stun_address(&address_str, bypass_resolver.as_ref()).await
            else {
                continue;
            };

            // We need a fresh client reference per task; clone the necessary data
            let label_clone = label.clone();
            let iface_clone = egress_iface.clone();
            // Build a minimal single-server client for the spawned task
            let servers_clone = vec![StunServerConfig {
                address: address_str,
                label: Some(label_clone.clone()),
            }];

            set.spawn(async move {
                let client = StunClient::new(servers_clone);
                client
                    .query_server(resolved, &label_clone, &iface_clone)
                    .await
            });
        }

        let mut results: Vec<ReflexiveAddress> = Vec::new();
        let mut last_error: Option<StunError> = None;

        while let Some(join_result) = set.join_next().await {
            match join_result {
                Ok(Ok(reflexive)) => results.push(reflexive),
                Ok(Err(e)) => {
                    tracing::debug!(error = %e, "STUN query failed");
                    last_error = Some(e);
                }
                Err(e) => {
                    tracing::debug!(error = %e, "STUN task panicked");
                }
            }
        }

        if results.is_empty() {
            return Err(last_error.unwrap_or(StunError::NoServers));
        }

        let behavior = detect_nat_behavior(&results);

        tracing::debug!(
            count = results.len(),
            behavior = ?behavior,
            "STUN discovery complete"
        );

        Ok((results, behavior))
    }
}

// ---- Protocol helpers (free functions) -------------------------------------

/// Build a 20-byte STUN Binding Request.
///
/// Layout (RFC 5389 Section 6):
///   - 2 bytes: message type (0x0001 = Binding Request)
///   - 2 bytes: message length (0 = no attributes)
///   - 4 bytes: magic cookie (0x2112A442)
///   - 12 bytes: transaction ID
fn build_binding_request(transaction_id: &[u8; 12]) -> [u8; 20] {
    let mut msg = [0u8; STUN_HEADER_SIZE];
    // Message type: Binding Request
    msg[0..2].copy_from_slice(&BINDING_REQUEST.to_be_bytes());
    // Message length: 0 (no attributes in the request)
    msg[2..4].copy_from_slice(&0u16.to_be_bytes());
    // Magic cookie
    msg[4..8].copy_from_slice(&MAGIC_COOKIE.to_be_bytes());
    // Transaction ID
    msg[8..20].copy_from_slice(transaction_id);
    msg
}

/// Parse a STUN Binding Response and extract the reflexive address.
///
/// Validates the header (message type, magic cookie, transaction ID) then
/// searches for XOR-MAPPED-ADDRESS (preferred) or MAPPED-ADDRESS (fallback).
fn parse_binding_response(
    data: &[u8],
    expected_txn_id: &[u8; 12],
) -> Result<SocketAddr, StunError> {
    if data.len() < STUN_HEADER_SIZE {
        return Err(StunError::TooShort(data.len()));
    }

    // Message type
    let msg_type = u16::from_be_bytes([data[0], data[1]]);
    if msg_type != BINDING_RESPONSE {
        return Err(StunError::NotBindingResponse(msg_type));
    }

    // Attributes length (bytes after the 20-byte header)
    let attr_len = u16::from_be_bytes([data[2], data[3]]) as usize;

    // Magic cookie
    let cookie = u32::from_be_bytes([data[4], data[5], data[6], data[7]]);
    if cookie != MAGIC_COOKIE {
        return Err(StunError::BadMagicCookie);
    }

    // Transaction ID
    if data[8..20] != *expected_txn_id {
        return Err(StunError::TransactionMismatch);
    }

    // Walk attributes
    let attr_end = STUN_HEADER_SIZE + attr_len.min(data.len() - STUN_HEADER_SIZE);
    let mut offset = STUN_HEADER_SIZE;
    let mut xor_addr: Option<SocketAddr> = None;
    let mut mapped_addr: Option<SocketAddr> = None;

    while offset + 4 <= attr_end {
        let attr_type = u16::from_be_bytes([data[offset], data[offset + 1]]);
        let value_len = u16::from_be_bytes([data[offset + 2], data[offset + 3]]) as usize;
        let value_start = offset + 4;
        let value_end = value_start + value_len;

        if value_end > data.len() {
            break;
        }

        let value = &data[value_start..value_end];

        match attr_type {
            ATTR_XOR_MAPPED_ADDRESS => {
                xor_addr = parse_xor_mapped_address(value, expected_txn_id);
            }
            ATTR_MAPPED_ADDRESS => {
                mapped_addr = parse_mapped_address(value);
            }
            _ => {}
        }

        // Attributes are padded to 4-byte boundaries
        let padded = (value_len + 3) & !3;
        offset = value_start + padded;
    }

    xor_addr.or(mapped_addr).ok_or(StunError::NoMappedAddress)
}

/// Decode an XOR-MAPPED-ADDRESS attribute value (RFC 5389 Section 15.2).
///
/// The port is XOR-decoded with the high 16 bits of the magic cookie (0x2112).
/// IPv4 addresses are XOR-decoded with the full magic cookie (0x2112A442).
fn parse_xor_mapped_address(value: &[u8], txn_id: &[u8; 12]) -> Option<SocketAddr> {
    if value.len() < 8 {
        return None;
    }

    // Byte 0 is reserved (0x00), byte 1 is address family
    let family = value[1];
    let xored_port = u16::from_be_bytes([value[2], value[3]]);
    let port = xored_port ^ (MAGIC_COOKIE >> 16) as u16;

    match family {
        0x01 => {
            // IPv4: 4 bytes XOR-decoded with magic cookie
            if value.len() < 8 {
                return None;
            }
            let xored_ip = u32::from_be_bytes([value[4], value[5], value[6], value[7]]);
            let ip = Ipv4Addr::from(xored_ip ^ MAGIC_COOKIE);
            Some(SocketAddr::new(ip.into(), port))
        }
        0x02 => {
            // IPv6: 16 bytes XOR-decoded with magic cookie + transaction ID
            if value.len() < 20 {
                return None;
            }
            let mut ip_bytes = [0u8; 16];
            ip_bytes.copy_from_slice(&value[4..20]);
            let cookie_bytes = MAGIC_COOKIE.to_be_bytes();
            for (i, b) in ip_bytes.iter_mut().enumerate() {
                if i < 4 {
                    *b ^= cookie_bytes[i];
                } else {
                    *b ^= txn_id[i - 4];
                }
            }
            let ip = Ipv6Addr::from(ip_bytes);
            Some(SocketAddr::new(ip.into(), port))
        }
        _ => None,
    }
}

/// Decode a MAPPED-ADDRESS attribute value (RFC 5389 Section 15.1).
///
/// Legacy format without XOR obfuscation. Used as fallback for older servers.
fn parse_mapped_address(value: &[u8]) -> Option<SocketAddr> {
    if value.len() < 8 {
        return None;
    }

    let family = value[1];
    let port = u16::from_be_bytes([value[2], value[3]]);

    match family {
        0x01 => {
            // IPv4
            let ip = Ipv4Addr::new(value[4], value[5], value[6], value[7]);
            Some(SocketAddr::new(ip.into(), port))
        }
        0x02 => {
            // IPv6
            if value.len() < 20 {
                return None;
            }
            let mut ip_bytes = [0u8; 16];
            ip_bytes.copy_from_slice(&value[4..20]);
            let ip = Ipv6Addr::from(ip_bytes);
            Some(SocketAddr::new(ip.into(), port))
        }
        _ => None,
    }
}

/// Determine NAT behavior from multiple reflexive address results.
///
/// If all servers report the same reflexive port, the NAT is likely
/// endpoint-independent (good for hole punching). Different ports
/// indicate symmetric NAT (relay needed).
fn detect_nat_behavior(results: &[ReflexiveAddress]) -> NatBehavior {
    if results.len() <= 1 {
        return NatBehavior::EndpointIndependent;
    }

    let first_port = results[0].address.port();
    let all_same = results.iter().all(|r| r.address.port() == first_port);

    if all_same {
        NatBehavior::EndpointIndependent
    } else {
        NatBehavior::Symmetric
    }
}

/// Classification of a STUN server `address` string for resolution purposes.
#[derive(Debug, Clone, PartialEq, Eq)]
enum StunAddress {
    /// The address already parses as a `SocketAddr` (literal `ip:port`,
    /// including bracketed IPv6 like `[::1]:3478`); no DNS is needed.
    Literal(SocketAddr),
    /// A `host:port` pair whose host part must be resolved via DNS.
    HostPort { host: String, port: u16 },
    /// The address has no usable `host:port` shape (no port, empty host, …).
    Invalid,
}

/// Classify a STUN server address string into the resolution path it needs.
///
/// A literal `SocketAddr` (v4 `1.2.3.4:3478`, bracketed v6 `[2001:db8::1]:3478`)
/// short-circuits DNS. A `host:port` whose host part is a DNS name needs
/// resolution. Anything else is `Invalid`.
///
/// The host part is split off the rightmost colon. A residual colon in the host
/// means the input was an *unbracketed* IPv6 literal (e.g. `2001:db8::1`): such
/// a string is rejected as `Invalid` rather than being misparsed as
/// `host="2001:db8:" port=1`, because an unbracketed v6 literal is not a valid
/// `host:port` string (a DNS hostname / IPv4 never contains a colon).
fn classify_stun_address(address: &str) -> StunAddress {
    if let Ok(sock) = SocketAddr::from_str(address) {
        return StunAddress::Literal(sock);
    }

    // Split host:port on the *last* colon so ordinary hostnames work.
    let Some((host, port_str)) = address.rsplit_once(':') else {
        return StunAddress::Invalid;
    };
    let host = host.trim_matches(|c| c == '[' || c == ']');
    // A non-empty host containing a colon is an unbracketed IPv6 leftover, not a
    // resolvable DNS name — reject it. (Bracketed v6 literals already matched the
    // `SocketAddr` parse above.)
    if host.is_empty() || host.contains(':') {
        return StunAddress::Invalid;
    }
    match port_str.parse::<u16>() {
        Ok(port) => StunAddress::HostPort {
            host: host.to_string(),
            port,
        },
        Err(_) => StunAddress::Invalid,
    }
}

/// Build the resolver-bypass used to look up STUN server hostnames without
/// trusting the host resolver.
///
/// Auto-detects upstreams the same way [`crate::dns::DnsServer::new`] does — a
/// [`DnsConfig`] with `upstreams: None` parsed against the production
/// [`RESOLV_CONF_PATH`], with loopback/systemd-resolved-stub entries filtered
/// out and a public fallback. Returns `None` if no resolver can be built (the
/// caller then leans on the [`tokio::net::lookup_host`] fallback).
fn build_bypass_resolver() -> Option<hickory_server::resolver::TokioAsyncResolver> {
    // A zone/bind_addr is required to construct DnsConfig but is irrelevant to
    // upstream detection (only `upstreams: None` + the resolv.conf path matter).
    let config = DnsConfig::new("overlay.local.", Ipv4Addr::UNSPECIFIED.into());
    let upstreams = resolve_upstreams(&config, RESOLV_CONF_PATH);
    match build_forward_resolver(&upstreams) {
        Ok(resolver) => Some(resolver),
        Err(e) => {
            tracing::warn!(
                error = %e,
                "could not build STUN resolver-bypass; falling back to host resolver"
            );
            None
        }
    }
}

/// Resolve a STUN server `address` to a `SocketAddr` without trusting the host
/// resolver where avoidable.
///
/// Resolution order:
/// 1. Literal `ip:port` → returned directly, no DNS.
/// 2. `host:port` → resolved through `bypass` (the overlay upstream resolver).
/// 3. Only if the bypass lookup errors (or no bypass resolver exists) →
///    [`tokio::net::lookup_host`], with a warning.
///
/// Returns `None` (with a warning) when the address is unparseable or every
/// resolution path yields no address.
async fn resolve_stun_address(
    address: &str,
    bypass: Option<&hickory_server::resolver::TokioAsyncResolver>,
) -> Option<SocketAddr> {
    let (host, port) = match classify_stun_address(address) {
        StunAddress::Literal(sock) => return Some(sock),
        StunAddress::HostPort { host, port } => (host, port),
        StunAddress::Invalid => {
            tracing::warn!(server = %address, "STUN server address is not a valid host:port");
            return None;
        }
    };

    // Preferred path: resolve through the overlay's own upstreams.
    if let Some(resolver) = bypass {
        match resolver.lookup_ip(host.as_str()).await {
            Ok(lookup) => {
                if let Some(ip) = lookup.iter().next() {
                    return Some(SocketAddr::new(ip, port));
                }
                tracing::warn!(
                    server = %address,
                    "STUN hostname resolved to no addresses via overlay upstreams; \
                     trying host resolver"
                );
            }
            Err(e) => {
                tracing::warn!(
                    server = %address,
                    error = %e,
                    "STUN hostname lookup via overlay upstreams failed; \
                     falling back to host resolver"
                );
            }
        }
    }

    // Fallback: the host resolver. It may still work, and a broken fallback is
    // no worse than the pre-hardening behaviour.
    match tokio::net::lookup_host((host.clone(), port)).await {
        Ok(mut addrs) => {
            if let Some(addr) = addrs.next() {
                Some(addr)
            } else {
                tracing::warn!(server = %address, "DNS returned no addresses");
                None
            }
        }
        Err(e) => {
            tracing::warn!(server = %address, error = %e, "DNS resolution failed");
            None
        }
    }
}

// ---- Tests -----------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};

    #[test]
    fn test_build_binding_request_header() {
        let txn_id = [1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
        let msg = build_binding_request(&txn_id);

        assert_eq!(
            msg.len(),
            20,
            "STUN Binding Request must be exactly 20 bytes"
        );

        // Message type: 0x0001
        assert_eq!(msg[0], 0x00);
        assert_eq!(msg[1], 0x01);

        // Message length: 0x0000
        assert_eq!(msg[2], 0x00);
        assert_eq!(msg[3], 0x00);

        // Magic cookie: 0x2112A442
        assert_eq!(msg[4], 0x21);
        assert_eq!(msg[5], 0x12);
        assert_eq!(msg[6], 0xA4);
        assert_eq!(msg[7], 0x42);

        // Transaction ID
        assert_eq!(&msg[8..20], &txn_id);
    }

    #[test]
    fn test_parse_binding_response_rfc5769_ipv4() {
        // Construct a synthetic Binding Response with XOR-MAPPED-ADDRESS
        // Reflexive address: 192.0.2.1:32853 (0xC000_0201, port 0x8055)
        let txn_id: [u8; 12] = [
            0xB7, 0xE7, 0xA7, 0x01, 0xBC, 0x34, 0xD6, 0x86, 0xFA, 0x87, 0xDF, 0xAE,
        ];

        let ip = Ipv4Addr::new(192, 0, 2, 1);
        let port: u16 = 32853;

        // XOR the port and IP
        let xored_port = port ^ (MAGIC_COOKIE >> 16) as u16;
        let xored_ip = u32::from_be_bytes(ip.octets()) ^ MAGIC_COOKIE;

        // Build XOR-MAPPED-ADDRESS attribute value
        let mut attr_value = [0u8; 8];
        attr_value[0] = 0x00; // reserved
        attr_value[1] = 0x01; // IPv4
        attr_value[2..4].copy_from_slice(&xored_port.to_be_bytes());
        attr_value[4..8].copy_from_slice(&xored_ip.to_be_bytes());

        // Build complete response
        let attr_header_len: u16 = 4 + 8; // type(2) + length(2) + value(8)
        let mut response = Vec::with_capacity(STUN_HEADER_SIZE + attr_header_len as usize);

        // Header
        response.extend_from_slice(&BINDING_RESPONSE.to_be_bytes());
        response.extend_from_slice(&attr_header_len.to_be_bytes());
        response.extend_from_slice(&MAGIC_COOKIE.to_be_bytes());
        response.extend_from_slice(&txn_id);

        // XOR-MAPPED-ADDRESS attribute
        response.extend_from_slice(&ATTR_XOR_MAPPED_ADDRESS.to_be_bytes());
        response.extend_from_slice(&8u16.to_be_bytes()); // value length
        response.extend_from_slice(&attr_value);

        let result = parse_binding_response(&response, &txn_id).unwrap();
        assert_eq!(result, SocketAddr::new(ip.into(), port));
    }

    #[test]
    fn test_parse_xor_mapped_address_ipv4() {
        // Test reflexive address: 198.51.100.50:12345
        let ip = Ipv4Addr::new(198, 51, 100, 50);
        let port: u16 = 12345;
        let txn_id = [0u8; 12];

        let xored_port = port ^ (MAGIC_COOKIE >> 16) as u16;
        let xored_ip = u32::from_be_bytes(ip.octets()) ^ MAGIC_COOKIE;

        let mut value = [0u8; 8];
        value[0] = 0x00;
        value[1] = 0x01; // IPv4
        value[2..4].copy_from_slice(&xored_port.to_be_bytes());
        value[4..8].copy_from_slice(&xored_ip.to_be_bytes());

        let result = parse_xor_mapped_address(&value, &txn_id).unwrap();
        assert_eq!(result.ip(), ip);
        assert_eq!(result.port(), port);
    }

    #[test]
    fn test_parse_mapped_address_ipv4() {
        let ip = Ipv4Addr::new(10, 0, 0, 1);
        let port: u16 = 51820;

        let mut value = [0u8; 8];
        value[0] = 0x00;
        value[1] = 0x01; // IPv4
        value[2..4].copy_from_slice(&port.to_be_bytes());
        value[4] = 10;
        value[5] = 0;
        value[6] = 0;
        value[7] = 1;

        let result = parse_mapped_address(&value).unwrap();
        assert_eq!(result, SocketAddr::new(ip.into(), port));
    }

    #[test]
    fn test_parse_binding_response_too_short() {
        let data = [0u8; 10];
        let txn_id = [0u8; 12];
        let err = parse_binding_response(&data, &txn_id).unwrap_err();
        assert!(matches!(err, StunError::TooShort(10)));
    }

    #[test]
    fn test_parse_binding_response_wrong_type() {
        let txn_id = [0u8; 12];
        let mut data = [0u8; 20];
        // Set message type to Binding Request instead of Response
        data[0..2].copy_from_slice(&BINDING_REQUEST.to_be_bytes());
        data[4..8].copy_from_slice(&MAGIC_COOKIE.to_be_bytes());
        data[8..20].copy_from_slice(&txn_id);

        let err = parse_binding_response(&data, &txn_id).unwrap_err();
        assert!(matches!(err, StunError::NotBindingResponse(0x0001)));
    }

    #[test]
    fn test_parse_binding_response_bad_cookie() {
        let txn_id = [0u8; 12];
        let mut data = [0u8; 20];
        data[0..2].copy_from_slice(&BINDING_RESPONSE.to_be_bytes());
        data[4..8].copy_from_slice(&0xDEAD_BEEFu32.to_be_bytes());
        data[8..20].copy_from_slice(&txn_id);

        let err = parse_binding_response(&data, &txn_id).unwrap_err();
        assert!(matches!(err, StunError::BadMagicCookie));
    }

    #[test]
    fn test_parse_binding_response_txn_mismatch() {
        let txn_id = [1u8; 12];
        let wrong_txn = [2u8; 12];
        let mut data = [0u8; 20];
        data[0..2].copy_from_slice(&BINDING_RESPONSE.to_be_bytes());
        data[4..8].copy_from_slice(&MAGIC_COOKIE.to_be_bytes());
        data[8..20].copy_from_slice(&wrong_txn);

        let err = parse_binding_response(&data, &txn_id).unwrap_err();
        assert!(matches!(err, StunError::TransactionMismatch));
    }

    #[test]
    fn test_parse_binding_response_no_mapped_address() {
        let txn_id = [0u8; 12];
        let mut data = [0u8; 20];
        data[0..2].copy_from_slice(&BINDING_RESPONSE.to_be_bytes());
        // attr_len = 0
        data[4..8].copy_from_slice(&MAGIC_COOKIE.to_be_bytes());
        data[8..20].copy_from_slice(&txn_id);

        let err = parse_binding_response(&data, &txn_id).unwrap_err();
        assert!(matches!(err, StunError::NoMappedAddress));
    }

    #[test]
    fn test_parse_binding_response_mapped_address_fallback() {
        // Test that MAPPED-ADDRESS (0x0001) is used when XOR-MAPPED-ADDRESS
        // is absent (legacy STUN server).
        let txn_id = [0u8; 12];
        let ip = Ipv4Addr::new(203, 0, 113, 5);
        let port: u16 = 40000;

        // Build MAPPED-ADDRESS attribute value
        let mut attr_value = [0u8; 8];
        attr_value[0] = 0x00;
        attr_value[1] = 0x01; // IPv4
        attr_value[2..4].copy_from_slice(&port.to_be_bytes());
        attr_value[4..8].copy_from_slice(&ip.octets());

        let attr_header_len: u16 = 4 + 8;
        let mut response = Vec::with_capacity(STUN_HEADER_SIZE + attr_header_len as usize);

        response.extend_from_slice(&BINDING_RESPONSE.to_be_bytes());
        response.extend_from_slice(&attr_header_len.to_be_bytes());
        response.extend_from_slice(&MAGIC_COOKIE.to_be_bytes());
        response.extend_from_slice(&txn_id);

        response.extend_from_slice(&ATTR_MAPPED_ADDRESS.to_be_bytes());
        response.extend_from_slice(&8u16.to_be_bytes());
        response.extend_from_slice(&attr_value);

        let result = parse_binding_response(&response, &txn_id).unwrap();
        assert_eq!(result, SocketAddr::new(ip.into(), port));
    }

    #[test]
    fn test_symmetric_nat_detection_same_ports() {
        let results = vec![
            ReflexiveAddress {
                address: SocketAddr::new(Ipv4Addr::new(1, 2, 3, 4).into(), 5000),
                server: "server1".to_string(),
            },
            ReflexiveAddress {
                address: SocketAddr::new(Ipv4Addr::new(5, 6, 7, 8).into(), 5000),
                server: "server2".to_string(),
            },
        ];
        assert_eq!(
            detect_nat_behavior(&results),
            NatBehavior::EndpointIndependent
        );
    }

    #[test]
    fn test_symmetric_nat_detection_different_ports() {
        let results = vec![
            ReflexiveAddress {
                address: SocketAddr::new(Ipv4Addr::new(1, 2, 3, 4).into(), 5000),
                server: "server1".to_string(),
            },
            ReflexiveAddress {
                address: SocketAddr::new(Ipv4Addr::new(1, 2, 3, 4).into(), 6000),
                server: "server2".to_string(),
            },
        ];
        assert_eq!(detect_nat_behavior(&results), NatBehavior::Symmetric);
    }

    #[test]
    fn test_symmetric_nat_detection_single_result() {
        let results = vec![ReflexiveAddress {
            address: SocketAddr::new(Ipv4Addr::new(1, 2, 3, 4).into(), 5000),
            server: "server1".to_string(),
        }];
        assert_eq!(
            detect_nat_behavior(&results),
            NatBehavior::EndpointIndependent
        );
    }

    #[test]
    fn test_symmetric_nat_detection_empty() {
        let results: Vec<ReflexiveAddress> = vec![];
        assert_eq!(
            detect_nat_behavior(&results),
            NatBehavior::EndpointIndependent
        );
    }

    #[test]
    fn test_stun_client_new() {
        let servers = vec![StunServerConfig {
            address: "stun.example.com:3478".to_string(),
            label: Some("Test".to_string()),
        }];
        let client = StunClient::new(servers);
        assert_eq!(client.servers.len(), 1);
    }

    // ---- IPv6 tests ---------------------------------------------------------

    #[test]
    fn test_parse_xor_mapped_address_ipv6() {
        // Test reflexive address: [2001:db8::1]:54321
        let ip = Ipv6Addr::new(0x2001, 0x0db8, 0, 0, 0, 0, 0, 1);
        let port: u16 = 54321;
        let txn_id: [u8; 12] = [
            0xA1, 0xB2, 0xC3, 0xD4, 0xE5, 0xF6, 0x17, 0x28, 0x39, 0x4A, 0x5B, 0x6C,
        ];

        // XOR the port with high 16 bits of magic cookie
        let xored_port = port ^ (MAGIC_COOKIE >> 16) as u16;

        // XOR the IP: first 4 bytes with magic cookie, remaining 12 with txn_id
        let ip_bytes = ip.octets();
        let cookie_bytes = MAGIC_COOKIE.to_be_bytes();
        let mut xored_ip = [0u8; 16];
        for i in 0..16 {
            if i < 4 {
                xored_ip[i] = ip_bytes[i] ^ cookie_bytes[i];
            } else {
                xored_ip[i] = ip_bytes[i] ^ txn_id[i - 4];
            }
        }

        // Build attribute value: [reserved(1), family(1), port(2), ip(16)] = 20 bytes
        let mut value = [0u8; 20];
        value[0] = 0x00; // reserved
        value[1] = 0x02; // IPv6
        value[2..4].copy_from_slice(&xored_port.to_be_bytes());
        value[4..20].copy_from_slice(&xored_ip);

        let result = parse_xor_mapped_address(&value, &txn_id).unwrap();
        assert_eq!(result.ip(), ip);
        assert_eq!(result.port(), port);
    }

    #[test]
    fn test_parse_xor_mapped_address_ipv6_all_ones() {
        // Edge case: all-ones IPv6 address
        let ip = Ipv6Addr::new(
            0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF,
        );
        let port: u16 = 65535;
        let txn_id = [0xFF; 12];

        let xored_port = port ^ (MAGIC_COOKIE >> 16) as u16;
        let ip_bytes = ip.octets();
        let cookie_bytes = MAGIC_COOKIE.to_be_bytes();
        let mut xored_ip = [0u8; 16];
        for i in 0..16 {
            if i < 4 {
                xored_ip[i] = ip_bytes[i] ^ cookie_bytes[i];
            } else {
                xored_ip[i] = ip_bytes[i] ^ txn_id[i - 4];
            }
        }

        let mut value = [0u8; 20];
        value[0] = 0x00;
        value[1] = 0x02;
        value[2..4].copy_from_slice(&xored_port.to_be_bytes());
        value[4..20].copy_from_slice(&xored_ip);

        let result = parse_xor_mapped_address(&value, &txn_id).unwrap();
        assert_eq!(result.ip(), ip);
        assert_eq!(result.port(), port);
    }

    #[test]
    fn test_parse_mapped_address_ipv6() {
        let ip = Ipv6Addr::new(0x2001, 0x0db8, 0, 0, 0, 0, 0, 1);
        let port: u16 = 12345;

        let mut value = [0u8; 20];
        value[0] = 0x00; // reserved
        value[1] = 0x02; // IPv6
        value[2..4].copy_from_slice(&port.to_be_bytes());
        value[4..20].copy_from_slice(&ip.octets());

        let result = parse_mapped_address(&value).unwrap();
        assert_eq!(result.ip(), ip);
        assert_eq!(result.port(), port);
    }

    #[test]
    fn test_parse_binding_response_xor_mapped_ipv6() {
        // Full binding response with an IPv6 XOR-MAPPED-ADDRESS
        let ip = Ipv6Addr::new(0x2001, 0x0db8, 0xABCD, 0, 0, 0, 0, 0x42);
        let port: u16 = 9876;
        let txn_id: [u8; 12] = [
            0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC,
        ];

        let xored_port = port ^ (MAGIC_COOKIE >> 16) as u16;
        let ip_bytes = ip.octets();
        let cookie_bytes = MAGIC_COOKIE.to_be_bytes();
        let mut xored_ip = [0u8; 16];
        for i in 0..16 {
            if i < 4 {
                xored_ip[i] = ip_bytes[i] ^ cookie_bytes[i];
            } else {
                xored_ip[i] = ip_bytes[i] ^ txn_id[i - 4];
            }
        }

        // Build XOR-MAPPED-ADDRESS attribute value (20 bytes for IPv6)
        let mut attr_value = [0u8; 20];
        attr_value[0] = 0x00; // reserved
        attr_value[1] = 0x02; // IPv6
        attr_value[2..4].copy_from_slice(&xored_port.to_be_bytes());
        attr_value[4..20].copy_from_slice(&xored_ip);

        // Attribute: type(2) + length(2) + value(20) = 24 bytes
        let attr_header_len: u16 = 4 + 20;
        let mut response = Vec::with_capacity(STUN_HEADER_SIZE + attr_header_len as usize);

        // Header
        response.extend_from_slice(&BINDING_RESPONSE.to_be_bytes());
        response.extend_from_slice(&attr_header_len.to_be_bytes());
        response.extend_from_slice(&MAGIC_COOKIE.to_be_bytes());
        response.extend_from_slice(&txn_id);

        // XOR-MAPPED-ADDRESS attribute
        response.extend_from_slice(&ATTR_XOR_MAPPED_ADDRESS.to_be_bytes());
        response.extend_from_slice(&20u16.to_be_bytes()); // value length = 20 for IPv6
        response.extend_from_slice(&attr_value);

        let result = parse_binding_response(&response, &txn_id).unwrap();
        assert_eq!(result, SocketAddr::new(ip.into(), port));
    }

    #[test]
    fn test_parse_binding_response_mapped_address_ipv6_fallback() {
        // Test that MAPPED-ADDRESS (0x0001) with IPv6 is used as fallback
        let ip = Ipv6Addr::new(0xFE80, 0, 0, 0, 0, 0, 0, 1);
        let port: u16 = 40000;
        let txn_id = [0u8; 12];

        let mut attr_value = [0u8; 20];
        attr_value[0] = 0x00;
        attr_value[1] = 0x02; // IPv6
        attr_value[2..4].copy_from_slice(&port.to_be_bytes());
        attr_value[4..20].copy_from_slice(&ip.octets());

        let attr_header_len: u16 = 4 + 20;
        let mut response = Vec::with_capacity(STUN_HEADER_SIZE + attr_header_len as usize);

        response.extend_from_slice(&BINDING_RESPONSE.to_be_bytes());
        response.extend_from_slice(&attr_header_len.to_be_bytes());
        response.extend_from_slice(&MAGIC_COOKIE.to_be_bytes());
        response.extend_from_slice(&txn_id);

        response.extend_from_slice(&ATTR_MAPPED_ADDRESS.to_be_bytes());
        response.extend_from_slice(&20u16.to_be_bytes());
        response.extend_from_slice(&attr_value);

        let result = parse_binding_response(&response, &txn_id).unwrap();
        assert_eq!(result, SocketAddr::new(ip.into(), port));
    }

    #[test]
    fn test_symmetric_nat_detection_ipv6_same_ports() {
        let results = vec![
            ReflexiveAddress {
                address: SocketAddr::new(
                    Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1).into(),
                    5000,
                ),
                server: "server1".to_string(),
            },
            ReflexiveAddress {
                address: SocketAddr::new(
                    Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 2).into(),
                    5000,
                ),
                server: "server2".to_string(),
            },
        ];
        assert_eq!(
            detect_nat_behavior(&results),
            NatBehavior::EndpointIndependent
        );
    }

    #[test]
    fn test_symmetric_nat_detection_ipv6_different_ports() {
        let results = vec![
            ReflexiveAddress {
                address: SocketAddr::new(
                    Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1).into(),
                    5000,
                ),
                server: "server1".to_string(),
            },
            ReflexiveAddress {
                address: SocketAddr::new(
                    Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1).into(),
                    6000,
                ),
                server: "server2".to_string(),
            },
        ];
        assert_eq!(detect_nat_behavior(&results), NatBehavior::Symmetric);
    }

    #[test]
    fn test_parse_xor_mapped_address_ipv6_too_short() {
        // IPv6 attribute value needs 20 bytes; provide only 18
        let txn_id = [0u8; 12];
        let mut value = [0u8; 18];
        value[1] = 0x02; // IPv6 family
        assert!(parse_xor_mapped_address(&value, &txn_id).is_none());
    }

    #[test]
    fn test_parse_mapped_address_ipv6_too_short() {
        // IPv6 mapped address needs 20 bytes; provide only 10
        let mut value = [0u8; 10];
        value[1] = 0x02; // IPv6 family
        assert!(parse_mapped_address(&value).is_none());
    }

    // ---- STUN address classification (resolution-order helper) --------------

    #[test]
    fn test_classify_stun_address_ipv4_literal() {
        // A literal v4 ip:port short-circuits DNS.
        let addr = classify_stun_address("203.0.113.5:3478");
        assert_eq!(
            addr,
            StunAddress::Literal(SocketAddr::new(Ipv4Addr::new(203, 0, 113, 5).into(), 3478))
        );
    }

    #[test]
    fn test_classify_stun_address_ipv6_bracketed_literal() {
        // A bracketed v6 literal with port parses directly as SocketAddr.
        let addr = classify_stun_address("[2001:db8::1]:3478");
        assert_eq!(
            addr,
            StunAddress::Literal(SocketAddr::new(
                Ipv6Addr::new(0x2001, 0x0db8, 0, 0, 0, 0, 0, 1).into(),
                3478
            ))
        );
    }

    #[test]
    fn test_classify_stun_address_ipv6_loopback_bracketed_literal() {
        let addr = classify_stun_address("[::1]:5000");
        assert_eq!(
            addr,
            StunAddress::Literal(SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 5000))
        );
    }

    #[test]
    fn test_classify_stun_address_hostname_port() {
        // A hostname:port needs DNS resolution.
        let addr = classify_stun_address("stun.example.com:3478");
        assert_eq!(
            addr,
            StunAddress::HostPort {
                host: "stun.example.com".to_string(),
                port: 3478,
            }
        );
    }

    #[test]
    fn test_classify_stun_address_hostname_default_stun_port() {
        let addr = classify_stun_address("stun.l.google.com:19302");
        assert_eq!(
            addr,
            StunAddress::HostPort {
                host: "stun.l.google.com".to_string(),
                port: 19302,
            }
        );
    }

    #[test]
    fn test_classify_stun_address_no_port_is_invalid() {
        // Bare hostname with no port has no usable host:port shape.
        assert_eq!(
            classify_stun_address("stun.example.com"),
            StunAddress::Invalid
        );
    }

    #[test]
    fn test_classify_stun_address_unbracketed_ipv6_is_invalid() {
        // An unbracketed v6 literal is not a valid host:port string: it neither
        // parses as a SocketAddr nor splits cleanly on a single trailing port.
        assert_eq!(classify_stun_address("2001:db8::1"), StunAddress::Invalid);
    }

    #[test]
    fn test_classify_stun_address_non_numeric_port_is_invalid() {
        assert_eq!(
            classify_stun_address("stun.example.com:notaport"),
            StunAddress::Invalid
        );
    }

    #[test]
    fn test_classify_stun_address_empty_host_is_invalid() {
        // ":3478" has an empty host part.
        assert_eq!(classify_stun_address(":3478"), StunAddress::Invalid);
    }

    #[tokio::test]
    async fn test_resolve_stun_address_literal_skips_dns() {
        // A literal ip:port must resolve without any resolver (passing None for
        // the bypass and never reaching the host-resolver fallback proves the
        // literal short-circuit fires first).
        let resolved = resolve_stun_address("198.51.100.7:3478", None).await;
        assert_eq!(
            resolved,
            Some(SocketAddr::new(Ipv4Addr::new(198, 51, 100, 7).into(), 3478))
        );
    }

    #[tokio::test]
    async fn test_resolve_stun_address_bracketed_ipv6_literal_skips_dns() {
        let resolved = resolve_stun_address("[2001:db8::42]:19302", None).await;
        assert_eq!(
            resolved,
            Some(SocketAddr::new(
                Ipv6Addr::new(0x2001, 0x0db8, 0, 0, 0, 0, 0, 0x42).into(),
                19302
            ))
        );
    }

    #[tokio::test]
    async fn test_resolve_stun_address_invalid_returns_none() {
        // Unparseable address yields None without panicking or hitting the
        // network (no resolver supplied).
        assert!(resolve_stun_address("not-a-valid-address", None)
            .await
            .is_none());
    }
}