rialo-types 0.11.0

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::{
    borrow::Cow,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
};

pub use ::multiaddr::{Error, Protocol};
use anyhow::{anyhow, Result};
use tracing::error;

/// Per-epoch UDP/TCP port offset: for epoch N, the new port is
/// `base_port + (N % EPOCH_PORT_COUNT) * EPOCH_PORT_OFFSET`. Gives each
/// epoch a different rotation slot from its predecessor; safe to wrap once
/// GC releases the prior occupant of a slot before the rotation comes back
/// to it. Sized to fit the per-validator open port-range on the firewall.
/// SUB-1666 tracks relocating the rotation decision into the
/// epoch-selection program.
///
/// A single, feature-independent value: the genesis tooling, the node, and
/// every harness derive reserved port sets from this constant, and they are
/// not always compiled in one cargo invocation — a feature-gated value would
/// let two binaries of one network disagree about which ports a validator
/// reserved, which the node's startup validation rejects. Tests need no
/// wider stride than production: reserved listen windows keep harness ports
/// disjoint from every validator's window.
pub const EPOCH_PORT_OFFSET: u16 = 10;

/// Number of distinct port slots the per-epoch rotation cycles through.
pub const EPOCH_PORT_COUNT: u64 = 10;

#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct Multiaddr(::multiaddr::Multiaddr);

impl Multiaddr {
    pub fn empty() -> Self {
        Self(::multiaddr::Multiaddr::empty())
    }

    #[allow(dead_code)]
    pub fn new_internal(inner: ::multiaddr::Multiaddr) -> Self {
        Self(inner)
    }

    pub fn iter(&self) -> ::multiaddr::Iter<'_> {
        self.0.iter()
    }

    pub fn pop<'a>(&mut self) -> Option<Protocol<'a>> {
        self.0.pop()
    }

    pub fn push(&mut self, p: Protocol<'_>) -> &mut Self {
        self.0.push(p);
        self
    }

    pub fn replace<'a, F>(&self, at: usize, by: F) -> Option<Multiaddr>
    where
        F: FnOnce(&Protocol<'_>) -> Option<Protocol<'a>>,
    {
        self.0.replace(at, by).map(Self)
    }

    pub fn len(&self) -> usize {
        self.0.len()
    }

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

    pub fn udp_multiaddr_to_listen_address(&self) -> Option<std::net::SocketAddr> {
        let mut iter = self.iter();

        match (iter.next(), iter.next()) {
            (Some(Protocol::Ip4(ipaddr)), Some(Protocol::Udp(port))) => Some((ipaddr, port).into()),
            (Some(Protocol::Ip6(ipaddr)), Some(Protocol::Udp(port))) => Some((ipaddr, port).into()),

            (Some(Protocol::Dns(_)), Some(Protocol::Udp(port))) => {
                Some((std::net::Ipv4Addr::UNSPECIFIED, port).into())
            }

            _ => None,
        }
    }

    // Converts a /ip{4,6}/-/tcp/-[/-] Multiaddr to SocketAddr.
    // Useful when an external library only accepts SocketAddr, e.g. to start a local server.
    // See `client::endpoint_from_multiaddr()` for converting to Endpoint for clients.
    pub fn to_socket_addr(&self) -> Result<SocketAddr> {
        let mut iter = self.iter();
        let ip = match iter.next().ok_or_else(|| {
            anyhow!("failed to convert to SocketAddr: Multiaddr does not contain IP")
        })? {
            Protocol::Ip4(ip4_addr) => IpAddr::V4(ip4_addr),
            Protocol::Ip6(ip6_addr) => IpAddr::V6(ip6_addr),
            unsupported => return Err(anyhow!("unsupported protocol {unsupported}")),
        };
        let tcp_port = parse_tcp(&mut iter)?;
        Ok(SocketAddr::new(ip, tcp_port))
    }

    // Returns true if the third component in the multiaddr is `Protocol::Tcp`
    pub fn is_loosely_valid_tcp_addr(&self) -> bool {
        let mut iter = self.iter();
        iter.next(); // Skip the ip/dns part
        match iter.next() {
            Some(Protocol::Tcp(_)) => true,
            _ => false, // including `None` and `Some(other)`
        }
    }

    /// Set the ip address to `0.0.0.0`. For instance, it converts the following address
    /// `/ip4/155.138.174.208/tcp/1500/http` into `/ip4/0.0.0.0/tcp/1500/http`.
    /// This is useful when starting a server and you want to listen on all interfaces.
    pub fn with_zero_ip(&self) -> Self {
        let mut new_address = self.0.clone();
        let Some(protocol) = new_address.iter().next() else {
            error!("Multiaddr is empty");
            return Self(new_address);
        };
        match protocol {
            multiaddr::Protocol::Ip4(_)
            | multiaddr::Protocol::Dns(_)
            | multiaddr::Protocol::Dns4(_) => {
                new_address = new_address
                    .replace(0, |_| Some(multiaddr::Protocol::Ip4(Ipv4Addr::UNSPECIFIED)))
                    .unwrap();
            }
            multiaddr::Protocol::Ip6(_) | multiaddr::Protocol::Dns6(_) => {
                new_address = new_address
                    .replace(0, |_| Some(multiaddr::Protocol::Ip6(Ipv6Addr::UNSPECIFIED)))
                    .unwrap();
            }
            p => {
                error!("Unsupported protocol {} in Multiaddr {}!", p, new_address);
            }
        }
        Self(new_address)
    }

    /// Set the ip address to `127.0.0.1`. For instance, it converts the following address
    /// `/ip4/155.138.174.208/tcp/1500/http` into `/ip4/127.0.0.1/tcp/1500/http`.
    pub fn with_localhost_ip(&self) -> Self {
        let mut new_address = self.0.clone();
        let Some(protocol) = new_address.iter().next() else {
            error!("Multiaddr is empty");
            return Self(new_address);
        };
        match protocol {
            multiaddr::Protocol::Ip4(_)
            | multiaddr::Protocol::Dns(_)
            | multiaddr::Protocol::Dns4(_) => {
                new_address = new_address
                    .replace(0, |_| Some(multiaddr::Protocol::Ip4(Ipv4Addr::LOCALHOST)))
                    .unwrap();
            }
            multiaddr::Protocol::Ip6(_) | multiaddr::Protocol::Dns6(_) => {
                new_address = new_address
                    .replace(0, |_| Some(multiaddr::Protocol::Ip6(Ipv6Addr::LOCALHOST)))
                    .unwrap();
            }
            p => {
                error!("Unsupported protocol {} in Multiaddr {}!", p, new_address);
            }
        }
        Self(new_address)
    }

    pub fn is_localhost_ip(&self) -> bool {
        let Some(protocol) = self.0.iter().next() else {
            error!("Multiaddr is empty");
            return false;
        };
        match protocol {
            multiaddr::Protocol::Ip4(addr) => addr == Ipv4Addr::LOCALHOST,
            multiaddr::Protocol::Ip6(addr) => addr == Ipv6Addr::LOCALHOST,
            _ => false,
        }
    }

    pub fn hostname(&self) -> Option<String> {
        for component in self.iter() {
            match component {
                Protocol::Ip4(ip) => return Some(ip.to_string()),
                Protocol::Ip6(ip) => return Some(ip.to_string()),
                Protocol::Dns(dns) => return Some(dns.to_string()),
                _ => (),
            }
        }
        None
    }

    pub fn port(&self) -> Option<u16> {
        for component in self.iter() {
            match component {
                Protocol::Udp(port) | Protocol::Tcp(port) => return Some(port),
                _ => (),
            }
        }
        None
    }

    pub fn rewrite_udp_to_tcp(&self) -> Self {
        let mut new = Self::empty();

        for component in self.iter() {
            if let Protocol::Udp(port) = component {
                new.push(Protocol::Tcp(port));
            } else {
                new.push(component);
            }
        }

        new
    }

    /// Rewrites the port component of the multiaddr with a new port.
    ///
    /// For instance, it converts `/ip4/127.0.0.1/tcp/8000` to `/ip4/127.0.0.1/tcp/9000`
    /// when called with `new_port = 9000`. Works with both TCP and UDP protocols.
    pub fn rewrite_port(&self, new_port: u16) -> Self {
        let mut new = Self::empty();

        for component in self.iter() {
            match component {
                Protocol::Tcp(_) => new.push(Protocol::Tcp(new_port)),
                Protocol::Udp(_) => new.push(Protocol::Udp(new_port)),
                _ => new.push(component),
            };
        }

        new
    }

    /// Returns the multiaddr with its port shifted by
    /// `(epoch % EPOCH_PORT_COUNT) * EPOCH_PORT_OFFSET`.
    pub fn with_epoch_port_offset(&self, epoch: u64) -> Result<Self> {
        let base_port = self
            .port()
            .ok_or_else(|| anyhow!("Multiaddr has no port: {}", self))?;
        // `checked_*` (not `saturating_*` / `wrapping_*`): silent overflow
        // would re-collide at `u16::MAX` and re-introduce the `AddrInUse`
        // bug this avoids.
        let epoch_u16 = u16::try_from(epoch % EPOCH_PORT_COUNT).map_err(|_| {
            anyhow!(
                "epoch slot {} (= {} % {}) exceeds u16::MAX for port-offset arithmetic",
                epoch % EPOCH_PORT_COUNT,
                epoch,
                EPOCH_PORT_COUNT
            )
        })?;
        let offset = epoch_u16.checked_mul(EPOCH_PORT_OFFSET).ok_or_else(|| {
            anyhow!(
                "epoch {} * EPOCH_PORT_OFFSET ({}) overflows u16",
                epoch,
                EPOCH_PORT_OFFSET
            )
        })?;
        let new_port = base_port.checked_add(offset).ok_or_else(|| {
            anyhow!(
                "base port {} + offset {} (epoch {}) overflows u16",
                base_port,
                offset,
                epoch
            )
        })?;
        Ok(self.rewrite_port(new_port))
    }

    /// The full per-epoch rotation of listen addresses for this base: the
    /// `EPOCH_PORT_COUNT` addresses `self.with_epoch_port_offset(0..EPOCH_PORT_COUNT)`.
    ///
    /// A node populates its reserved active listen-port sets with this rotation
    /// at config-creation, binds and holds all of them for the process lifetime,
    /// and lends the current epoch's address by port.
    ///
    /// TODO(SUB-1666): the per-epoch listen port is properly a decision of the
    /// epoch-selection program (the committee-assignment path in
    /// `rialo-tokenomics-governance-program`), made from the explicit port set
    /// each validator advertises in its config — not re-derived from a fixed
    /// `% EPOCH_PORT_COUNT` stride here. This hard-coded rotation should be
    /// removed from `multiaddr` once the program assigns ports from the
    /// validator-provided set. SUB-1666 tracks the broader removal of the
    /// port-offset workaround.
    pub fn epoch_rotation_addrs(&self) -> Result<Vec<Self>> {
        (0..EPOCH_PORT_COUNT)
            .map(|slot| self.with_epoch_port_offset(slot))
            .collect()
    }

    /// The reserved "on-the-fly" (passive) listen addresses for this base:
    /// `count` ports laid out immediately past the active rotation's span and
    /// disjoint from it. The rotation occupies `base_port + {0..EPOCH_PORT_COUNT-1}
    /// · EPOCH_PORT_OFFSET` (multiples of the offset); these take
    /// `base_port + (EPOCH_PORT_COUNT-1)·EPOCH_PORT_OFFSET + 1 + j`. A node holds
    /// these and lends them any-available to passive (observing) slots.
    ///
    /// Shares the `TODO(SUB-1666)` disposition of [`Self::epoch_rotation_addrs`]:
    /// the passive-port layout belongs with the epoch-selection program's
    /// port-set decision, not a fixed stride in `multiaddr`.
    pub fn on_the_fly_addrs(&self, count: u64) -> Result<Vec<Self>> {
        if count == 0 {
            return Ok(Vec::new());
        }
        let base_port = self
            .port()
            .ok_or_else(|| anyhow!("Multiaddr has no port: {}", self))?;
        let rotation_span = (EPOCH_PORT_COUNT as u16 - 1)
            .checked_mul(EPOCH_PORT_OFFSET)
            .ok_or_else(|| anyhow!("rotation span overflows u16 for {}", self))?;
        (0..count)
            .map(|j| {
                let offset = rotation_span
                    .checked_add(
                        1 + u16::try_from(j)
                            .map_err(|_| anyhow!("on-the-fly index {j} too large"))?,
                    )
                    .ok_or_else(|| anyhow!("on-the-fly offset overflows u16 for {}", self))?;
                let port = base_port
                    .checked_add(offset)
                    .ok_or_else(|| anyhow!("on-the-fly port overflows u16 for {}", self))?;
                Ok(self.rewrite_port(port))
            })
            .collect()
    }

    pub fn to_vec(&self) -> Vec<u8> {
        self.0.to_vec()
    }

    /// Create a Multiaddr from a SocketAddr for UDP protocol.
    ///
    /// This is useful when a server binds to an ephemeral port (0.0.0.0:0) and needs
    /// to discover and advertise the actual bound address to peers.
    pub fn from_socket_addr_udp(addr: SocketAddr) -> Self {
        let mut multiaddr = ::multiaddr::Multiaddr::empty();
        match addr.ip() {
            IpAddr::V4(ip) => multiaddr.push(Protocol::Ip4(ip)),
            IpAddr::V6(ip) => multiaddr.push(Protocol::Ip6(ip)),
        }
        multiaddr.push(Protocol::Udp(addr.port()));
        Self(multiaddr)
    }
}

impl From<IpAddr> for Multiaddr {
    fn from(ip_address: IpAddr) -> Self {
        Multiaddr(ip_address.into())
    }
}

impl From<Ipv4Addr> for Multiaddr {
    fn from(ip_address: Ipv4Addr) -> Self {
        Multiaddr(ip_address.into())
    }
}

impl From<Ipv6Addr> for Multiaddr {
    fn from(ip_address: Ipv6Addr) -> Self {
        Multiaddr(ip_address.into())
    }
}

impl std::fmt::Display for Multiaddr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(&self.0, f)
    }
}

impl std::str::FromStr for Multiaddr {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        ::multiaddr::Multiaddr::from_str(s).map(Self)
    }
}

impl<'a> TryFrom<&'a str> for Multiaddr {
    type Error = Error;

    fn try_from(value: &'a str) -> Result<Self, Self::Error> {
        value.parse()
    }
}

impl TryFrom<String> for Multiaddr {
    type Error = Error;

    fn try_from(value: String) -> Result<Self, Self::Error> {
        value.parse()
    }
}

impl TryFrom<Vec<u8>> for Multiaddr {
    type Error = Error;

    fn try_from(v: Vec<u8>) -> multiaddr::Result<Self> {
        Ok(Self(multiaddr::Multiaddr::try_from(v)?))
    }
}

impl serde::Serialize for Multiaddr {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_str(&self.0.to_string())
    }
}

impl<'de> serde::Deserialize<'de> for Multiaddr {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        s.parse()
            .map(Self)
            .map_err(|e| serde::de::Error::custom(e.to_string()))
    }
}

impl std::net::ToSocketAddrs for Multiaddr {
    type Iter = Box<dyn Iterator<Item = SocketAddr>>;

    fn to_socket_addrs(&self) -> std::io::Result<Self::Iter> {
        let mut iter = self.iter();

        match (iter.next(), iter.next()) {
            (Some(Protocol::Ip4(ip4)), Some(Protocol::Tcp(port) | Protocol::Udp(port))) => {
                (ip4, port)
                    .to_socket_addrs()
                    .map(|iter| Box::new(iter) as _)
            }
            (Some(Protocol::Ip6(ip6)), Some(Protocol::Tcp(port) | Protocol::Udp(port))) => {
                (ip6, port)
                    .to_socket_addrs()
                    .map(|iter| Box::new(iter) as _)
            }
            (Some(Protocol::Dns(hostname)), Some(Protocol::Tcp(port) | Protocol::Udp(port))) => {
                (hostname.as_ref(), port)
                    .to_socket_addrs()
                    .map(|iter| Box::new(iter) as _)
            }
            _ => Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "unable to convert Multiaddr to SocketAddr",
            )),
        }
    }
}

pub(crate) fn parse_tcp<'a, T: Iterator<Item = Protocol<'a>>>(protocols: &mut T) -> Result<u16> {
    if let Protocol::Tcp(port) = protocols
        .next()
        .ok_or_else(|| anyhow!("unexpected end of multiaddr"))?
    {
        Ok(port)
    } else {
        Err(anyhow!("expected tcp protocol"))
    }
}

#[allow(dead_code)]
pub(crate) fn parse_http_https<'a, T: Iterator<Item = Protocol<'a>>>(
    protocols: &mut T,
) -> Result<&'static str> {
    match protocols.next() {
        Some(Protocol::Http) => Ok("http"),
        Some(Protocol::Https) => Ok("https"),
        _ => Ok("http"),
    }
}

#[allow(dead_code)]
pub(crate) fn parse_end<'a, T: Iterator<Item = Protocol<'a>>>(protocols: &mut T) -> Result<()> {
    if protocols.next().is_none() {
        Ok(())
    } else {
        Err(anyhow!("expected end of multiaddr"))
    }
}

// Parse a full /dns/-/tcp/-/{http,https} address
#[allow(dead_code)]
pub fn parse_dns(address: &Multiaddr) -> Result<(Cow<'_, str>, u16, &'static str)> {
    let mut iter = address.iter();

    let dns_name = match iter
        .next()
        .ok_or_else(|| anyhow!("unexpected end of multiaddr"))?
    {
        Protocol::Dns(dns_name) => dns_name,
        other => return Err(anyhow!("expected dns found {other}")),
    };
    let tcp_port = parse_tcp(&mut iter)?;
    let http_or_https = parse_http_https(&mut iter)?;
    parse_end(&mut iter)?;
    Ok((dns_name, tcp_port, http_or_https))
}

// Parse a full /dns4/-/tcp/-/{http,https} address
#[allow(dead_code)]
pub fn parse_dns4(address: &Multiaddr) -> Result<(Cow<'_, str>, u16, &'static str)> {
    let mut iter = address.iter();

    let dns_name = match iter
        .next()
        .ok_or_else(|| anyhow!("unexpected end of multiaddr"))?
    {
        Protocol::Dns4(dns_name) => dns_name,
        other => return Err(anyhow!("expected dns4 found {other}")),
    };
    let tcp_port = parse_tcp(&mut iter)?;
    let http_or_https = parse_http_https(&mut iter)?;
    parse_end(&mut iter)?;
    Ok((dns_name, tcp_port, http_or_https))
}

// Parse a full /dns6/-/tcp/-/{http,https} address
#[allow(dead_code)]
pub fn parse_dns6(address: &Multiaddr) -> Result<(Cow<'_, str>, u16, &'static str)> {
    let mut iter = address.iter();

    let dns_name = match iter
        .next()
        .ok_or_else(|| anyhow!("unexpected end of multiaddr"))?
    {
        Protocol::Dns6(dns_name) => dns_name,
        other => return Err(anyhow!("expected dns6 found {other}")),
    };
    let tcp_port = parse_tcp(&mut iter)?;
    let http_or_https = parse_http_https(&mut iter)?;
    parse_end(&mut iter)?;
    Ok((dns_name, tcp_port, http_or_https))
}

/// Returns true iff `addr` is a TCP-host multiaddr in the shape accepted by
/// the codebase's existing parser family (`parse_ip4`/`parse_ip6`/
/// `parse_dns`/`parse_dns4`/`parse_dns6`):
/// `/{ip4,ip6,dns,dns4,dns6}/<host>/tcp/<port>` with **at most one**
/// trailing protocol after the TCP port. The trailing slot is canonically
/// `/http` or `/https`, but the codebase's existing parsers also silently
/// accept any other single trailing protocol there and treat it as if it
/// were `http` (see `parse_http_https`'s `_ => Ok("http")` arm). This
/// helper preserves that contract.
///
/// Intended for ingestion-time validation of `network_service_address`
/// (the TCP snapshot-service multiaddr on `ValidatorInfo`/`Authority`).
///
/// Invariant: if `is_tcp_host_multiaddr(addr)` returns `true`, then
/// `multiaddr_to_http_uri(addr)` in `rialo-statemanager-network` produces
/// `Some(_)`. A colocated tripwire test in
/// `rialo-statemanager-network::client::tests` pins this forward direction
/// so any drift is loud at CI time.
///
// TODO: tighten acceptance to reject specific trailing-protocol
// shapes that are semantically wrong for the snapshot gRPC service.
// Rejection candidates the operator almost certainly did not mean to
// configure as the long-lived TCP snapshot endpoint:
//
//   * /tcp/<port>/ws         — WebSocket-over-TCP (libp2p WS)
//   * /tcp/<port>/tls/ws     — WebSocket Secure variant (TLS + WS)
//   * /tcp/<port>/wss        — WebSocket Secure marker
//   * /tcp/<port>/quic-v1    — QUIC marker following TCP (malformed but parses)
//   * /tcp/<port>/noise      — libp2p Noise handshake marker
//   * /tcp/<port>/p2p/...    — explicit libp2p peer-id suffix
//
// Tightening at this site requires either (a) fixing `parse_http_https`
// to reject non-http/https trailing protocols (also tightens
// `endpoint_from_multiaddr`'s acceptance, which other gRPC channel
// establishment sites depend on — risk surface), or (b) writing a strict
// body inline that does not delegate to `parse_*` and explicitly checks
// `iter.next()` is None / Http / Https. Path (b) is the smaller-blast-radius
// option but creates a second multiaddr validator in the codebase with
// different semantics from the parser family. Defer until there is a
// concrete operator misconfiguration that motivates the cost.
pub fn is_tcp_host_multiaddr(addr: &Multiaddr) -> bool {
    parse_ip4(addr).is_ok()
        || parse_ip6(addr).is_ok()
        || parse_dns(addr).is_ok()
        || parse_dns4(addr).is_ok()
        || parse_dns6(addr).is_ok()
}

// Parse a full /ip4/-/tcp/-/{http,https} address
#[allow(dead_code)]
pub fn parse_ip4(address: &Multiaddr) -> Result<(SocketAddr, &'static str)> {
    let mut iter = address.iter();

    let ip_addr = match iter
        .next()
        .ok_or_else(|| anyhow!("unexpected end of multiaddr"))?
    {
        Protocol::Ip4(ip4_addr) => IpAddr::V4(ip4_addr),
        other => return Err(anyhow!("expected ip4 found {other}")),
    };
    let tcp_port = parse_tcp(&mut iter)?;
    let http_or_https = parse_http_https(&mut iter)?;
    parse_end(&mut iter)?;
    let socket_addr = SocketAddr::new(ip_addr, tcp_port);

    Ok((socket_addr, http_or_https))
}

// Parse a full /ip6/-/tcp/-/{http,https} address
#[allow(dead_code)]
pub fn parse_ip6(address: &Multiaddr) -> Result<(SocketAddr, &'static str)> {
    let mut iter = address.iter();

    let ip_addr = match iter
        .next()
        .ok_or_else(|| anyhow!("unexpected end of multiaddr"))?
    {
        Protocol::Ip6(ip6_addr) => IpAddr::V6(ip6_addr),
        other => return Err(anyhow!("expected ip6 found {other}")),
    };
    let tcp_port = parse_tcp(&mut iter)?;
    let http_or_https = parse_http_https(&mut iter)?;
    parse_end(&mut iter)?;
    let socket_addr = SocketAddr::new(ip_addr, tcp_port);

    Ok((socket_addr, http_or_https))
}

#[cfg(test)]
mod test {
    use multiaddr::multiaddr;

    use super::Multiaddr;

    #[test]
    fn test_to_socket_addr_basic() {
        let multi_addr_ipv4 = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Tcp(10500u16)));
        let socket_addr_ipv4 = multi_addr_ipv4
            .to_socket_addr()
            .expect("Couldn't convert to socket addr");
        assert_eq!(socket_addr_ipv4.to_string(), "127.0.0.1:10500");

        let multi_addr_ipv6 = Multiaddr(multiaddr!(Ip6([172, 0, 0, 1, 1, 1, 1, 1]), Tcp(10500u16)));
        let socket_addr_ipv6 = multi_addr_ipv6
            .to_socket_addr()
            .expect("Couldn't convert to socket addr");
        assert_eq!(socket_addr_ipv6.to_string(), "[ac::1:1:1:1:1]:10500");
    }

    #[test]
    fn test_to_socket_addr_unsupported_protocol() {
        let multi_addr_dns = Multiaddr(multiaddr!(Dnsaddr("mysten.sui"), Tcp(10500u16)));
        let _ = multi_addr_dns
            .to_socket_addr()
            .expect_err("DNS is unsupported");
    }

    #[test]
    fn test_is_loosely_valid_tcp_addr() {
        let multi_addr_ipv4 = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Tcp(10500u16)));
        assert!(multi_addr_ipv4.is_loosely_valid_tcp_addr());
        let multi_addr_ipv6 = Multiaddr(multiaddr!(Ip6([172, 0, 0, 1, 1, 1, 1, 1]), Tcp(10500u16)));
        assert!(multi_addr_ipv6.is_loosely_valid_tcp_addr());
        let multi_addr_dns = Multiaddr(multiaddr!(Dnsaddr("mysten.sui"), Tcp(10500u16)));
        assert!(multi_addr_dns.is_loosely_valid_tcp_addr());

        let multi_addr_ipv4 = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Udp(10500u16)));
        assert!(!multi_addr_ipv4.is_loosely_valid_tcp_addr());
        let multi_addr_ipv6 = Multiaddr(multiaddr!(Ip6([172, 0, 0, 1, 1, 1, 1, 1]), Udp(10500u16)));
        assert!(!multi_addr_ipv6.is_loosely_valid_tcp_addr());
        let multi_addr_dns = Multiaddr(multiaddr!(Dnsaddr("mysten.sui"), Udp(10500u16)));
        assert!(!multi_addr_dns.is_loosely_valid_tcp_addr());

        let invalid_multi_addr_ipv4 = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1])));
        assert!(!invalid_multi_addr_ipv4.is_loosely_valid_tcp_addr());
    }

    #[test]
    fn test_get_hostname_port() {
        let multi_addr_ip4 = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Tcp(10500u16)));
        assert_eq!(Some("127.0.0.1".to_string()), multi_addr_ip4.hostname());
        assert_eq!(Some(10500u16), multi_addr_ip4.port());

        let multi_addr_dns = Multiaddr(multiaddr!(Dns("mysten.sui"), Tcp(10501u16)));
        assert_eq!(Some("mysten.sui".to_string()), multi_addr_dns.hostname());
        assert_eq!(Some(10501u16), multi_addr_dns.port());
    }

    #[test]
    fn test_with_zero_ip() {
        let multi_addr_ip4 =
            Multiaddr(multiaddr!(Ip4([15, 15, 15, 1]), Tcp(10500u16))).with_zero_ip();
        assert_eq!(Some("0.0.0.0".to_string()), multi_addr_ip4.hostname());
        assert_eq!(Some(10500u16), multi_addr_ip4.port());

        let multi_addr_ip6 = Multiaddr(multiaddr!(
            Ip6([15, 15, 15, 15, 15, 15, 15, 1]),
            Tcp(10500u16)
        ))
        .with_zero_ip();
        assert_eq!(Some("::".to_string()), multi_addr_ip6.hostname());
        assert_eq!(Some(10500u16), multi_addr_ip4.port());

        let multi_addr_dns = Multiaddr(multiaddr!(Dns("mysten.sui"), Tcp(10501u16))).with_zero_ip();
        assert_eq!(Some("0.0.0.0".to_string()), multi_addr_dns.hostname());
        assert_eq!(Some(10501u16), multi_addr_dns.port());
    }

    #[test]
    fn test_with_localhost_ip() {
        let multi_addr_ip4 =
            Multiaddr(multiaddr!(Ip4([15, 15, 15, 1]), Tcp(10500u16))).with_localhost_ip();
        assert_eq!(Some("127.0.0.1".to_string()), multi_addr_ip4.hostname());
        assert_eq!(Some(10500u16), multi_addr_ip4.port());

        let multi_addr_ip6 = Multiaddr(multiaddr!(
            Ip6([15, 15, 15, 15, 15, 15, 15, 1]),
            Tcp(10500u16)
        ))
        .with_localhost_ip();
        assert_eq!(Some("::1".to_string()), multi_addr_ip6.hostname());
        assert_eq!(Some(10500u16), multi_addr_ip4.port());

        let multi_addr_dns =
            Multiaddr(multiaddr!(Dns("mysten.sui"), Tcp(10501u16))).with_localhost_ip();
        assert_eq!(Some("127.0.0.1".to_string()), multi_addr_dns.hostname());
        assert_eq!(Some(10501u16), multi_addr_dns.port());
    }

    #[test]
    fn test_with_epoch_port_offset() {
        let addr = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Udp(53203u16)));

        // Slot 0 at epoch 0: no offset.
        // Expected port for slot `s`: base + s * EPOCH_PORT_OFFSET.
        // Parameterised by the constant so the test holds under both the
        // production (`+10`) and testing (`+100`) gates.
        let expected = |slot: u16| 53203 + slot * super::EPOCH_PORT_OFFSET;

        let e0 = addr.with_epoch_port_offset(0).unwrap();
        assert_eq!(e0.port(), Some(expected(0)));

        // Slot 1 at epoch 1.
        let e1 = addr.with_epoch_port_offset(1).unwrap();
        assert_eq!(e1.port(), Some(expected(1)));

        // Slot 9 at epoch 9, the largest slot.
        let e9 = addr.with_epoch_port_offset(9).unwrap();
        assert_eq!(e9.port(), Some(expected(9)));

        // Rotation wraps at epoch == EPOCH_PORT_COUNT: epoch 10 → slot 0.
        let e10 = addr.with_epoch_port_offset(10).unwrap();
        assert_eq!(e10.port(), Some(expected(0)));

        // Epoch 11 → slot 1, same port as epoch 1.
        let e11 = addr.with_epoch_port_offset(11).unwrap();
        assert_eq!(e11.port(), Some(expected(1)));

        // Hostname is preserved across rotation.
        assert_eq!(e9.hostname(), Some("127.0.0.1".to_string()));

        // Epochs above u16::MAX round-trip safely: `% EPOCH_PORT_COUNT`
        // collapses them into 0..=9.
        let big = addr.with_epoch_port_offset(u64::MAX).unwrap();
        // u64::MAX % 10 == 5 → slot 5.
        assert_eq!(big.port(), Some(expected(5)));
        let very_big = addr.with_epoch_port_offset(65_536).unwrap();
        // 65_536 % 10 == 6 → slot 6.
        assert_eq!(very_big.port(), Some(expected(6)));

        // Overflow surfaces when `base + 9 * EPOCH_PORT_OFFSET` exceeds u16.
        // Pick a base one above the boundary so the test stays correct under
        // both gates.
        let overflow_base: u16 = u16::MAX - 9 * super::EPOCH_PORT_OFFSET + 1;
        let too_high = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Udp(overflow_base)));
        assert!(
            too_high.with_epoch_port_offset(9).is_err(),
            "base + slot-9 offset overflows u16; must surface as Err"
        );
        // Slot 0 at the same base still succeeds (offset is 0).
        assert_eq!(
            too_high.with_epoch_port_offset(0).unwrap().port(),
            Some(overflow_base)
        );

        // Multiaddr without a port is rejected (no UDP/TCP component).
        let no_port = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1])));
        assert!(no_port.with_epoch_port_offset(1).is_err());
    }

    #[test]
    fn test_epoch_rotation_addrs() {
        let base = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Udp(53203u16)));
        let rotation = base.epoch_rotation_addrs().unwrap();

        // One distinct address per epoch slot, in slot order.
        assert_eq!(rotation.len(), super::EPOCH_PORT_COUNT as usize);
        for (slot, addr) in rotation.iter().enumerate() {
            assert_eq!(
                addr.port(),
                Some(53203 + slot as u16 * super::EPOCH_PORT_OFFSET),
                "slot {slot} must equal base + slot * EPOCH_PORT_OFFSET"
            );
        }
        // Slot 0 is the base itself (offset 0).
        assert_eq!(rotation[0], base);

        // Every epoch's lend key is a member of the rotation (the load-bearing
        // pool-builder ⇄ lender agreement, checked over a representative spread
        // of epochs including past the wrap point).
        let ports: std::collections::HashSet<u16> =
            rotation.iter().filter_map(|a| a.port()).collect();
        for epoch in [0u64, 1, 9, 10, 11, 19, 100, 999, u64::MAX] {
            let key = base.with_epoch_port_offset(epoch).unwrap().port().unwrap();
            assert!(
                ports.contains(&key),
                "epoch {epoch} lend key {key} not in rotation"
            );
        }
    }

    #[test]
    fn test_on_the_fly_addrs() {
        let base = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Udp(53203u16)));

        // count == 0 yields an empty set.
        assert!(base.on_the_fly_addrs(0).unwrap().is_empty());

        let on_the_fly = base.on_the_fly_addrs(3).unwrap();
        assert_eq!(on_the_fly.len(), 3);

        // Laid out immediately past the rotation's span: base + (N-1)*OFFSET + 1 + j.
        let span = (super::EPOCH_PORT_COUNT as u16 - 1) * super::EPOCH_PORT_OFFSET;
        for (j, addr) in on_the_fly.iter().enumerate() {
            assert_eq!(addr.port(), Some(53203 + span + 1 + j as u16));
        }

        // Disjoint from the active rotation (which only uses multiples of OFFSET).
        let rotation: std::collections::HashSet<u16> = base
            .epoch_rotation_addrs()
            .unwrap()
            .iter()
            .filter_map(|a| a.port())
            .collect();
        for addr in &on_the_fly {
            assert!(
                !rotation.contains(&addr.port().unwrap()),
                "on-the-fly port {:?} collides with the active rotation",
                addr.port()
            );
        }

        // Hostname is preserved.
        assert_eq!(on_the_fly[0].hostname(), Some("127.0.0.1".to_string()));

        // A base without a port is rejected.
        let no_port = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1])));
        assert!(no_port.on_the_fly_addrs(1).is_err());
    }

    #[test]
    fn test_rewrite_port() {
        let multi_addr_tcp = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Tcp(8000u16)));
        let rewritten_tcp = multi_addr_tcp.rewrite_port(9000);
        assert_eq!(Some(9000u16), rewritten_tcp.port());
        assert_eq!(Some("127.0.0.1".to_string()), rewritten_tcp.hostname());

        let multi_addr_udp = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Udp(8000u16)));
        let rewritten_udp = multi_addr_udp.rewrite_port(9000);
        assert_eq!(Some(9000u16), rewritten_udp.port());
        assert_eq!(Some("127.0.0.1".to_string()), rewritten_udp.hostname());

        let multi_addr_dns = Multiaddr(multiaddr!(Dns("example.com"), Tcp(8080u16)));
        let rewritten_dns = multi_addr_dns.rewrite_port(8443);
        assert_eq!(Some(8443u16), rewritten_dns.port());
        assert_eq!(Some("example.com".to_string()), rewritten_dns.hostname());
    }

    #[test]
    fn test_from_socket_addr_udp_ipv4() {
        use std::net::SocketAddr;
        let socket_addr: SocketAddr = "127.0.0.1:12345".parse().unwrap();
        let multiaddr = Multiaddr::from_socket_addr_udp(socket_addr);
        assert_eq!(multiaddr.to_string(), "/ip4/127.0.0.1/udp/12345");
        assert_eq!(multiaddr.hostname(), Some("127.0.0.1".to_string()));
        assert_eq!(multiaddr.port(), Some(12345));
    }

    #[test]
    fn test_from_socket_addr_udp_ipv6() {
        use std::net::SocketAddr;
        let socket_addr: SocketAddr = "[::1]:54321".parse().unwrap();
        let multiaddr = Multiaddr::from_socket_addr_udp(socket_addr);
        assert_eq!(multiaddr.to_string(), "/ip6/::1/udp/54321");
        assert_eq!(multiaddr.hostname(), Some("::1".to_string()));
        assert_eq!(multiaddr.port(), Some(54321));
    }

    #[test]
    fn test_from_socket_addr_udp_roundtrip() {
        use std::net::SocketAddr;
        // Test that we can convert SocketAddr -> Multiaddr -> SocketAddr
        let original: SocketAddr = "192.168.1.100:8080".parse().unwrap();
        let multiaddr = Multiaddr::from_socket_addr_udp(original);
        let recovered = multiaddr.udp_multiaddr_to_listen_address().unwrap();
        assert_eq!(original, recovered);
    }

    // ---- is_tcp_host_multiaddr ----------------------------------------------
    //
    // Positive: every host class × {bare /tcp, /tcp/http, /tcp/https} accepted.
    // Negative: UDP, missing port, missing host, empty, trailing junk
    // (`/tcp/.../ws`, `/tcp/.../quic-v1`, `/tcp/.../tls/ws`).

    use std::str::FromStr;

    use super::is_tcp_host_multiaddr;

    fn ma(s: &str) -> Multiaddr {
        Multiaddr::from_str(s).expect("valid multiaddr literal for test")
    }

    #[test]
    fn test_is_tcp_host_multiaddr_ip4_tcp_bare() {
        assert!(is_tcp_host_multiaddr(&ma("/ip4/127.0.0.1/tcp/4300")));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_ip4_tcp_http() {
        assert!(is_tcp_host_multiaddr(&ma("/ip4/127.0.0.1/tcp/4300/http")));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_ip4_tcp_https() {
        assert!(is_tcp_host_multiaddr(&ma("/ip4/127.0.0.1/tcp/4300/https")));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_ip6_tcp_bare() {
        assert!(is_tcp_host_multiaddr(&ma("/ip6/::1/tcp/4300")));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_ip6_tcp_http() {
        assert!(is_tcp_host_multiaddr(&ma("/ip6/::1/tcp/4300/http")));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_dns_tcp_bare() {
        assert!(is_tcp_host_multiaddr(&ma("/dns/host.example.com/tcp/4300")));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_dns_tcp_http() {
        assert!(is_tcp_host_multiaddr(&ma(
            "/dns/host.example.com/tcp/4300/http"
        )));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_dns4_tcp_bare() {
        assert!(is_tcp_host_multiaddr(&ma(
            "/dns4/host.example.com/tcp/4300"
        )));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_dns4_tcp_http() {
        assert!(is_tcp_host_multiaddr(&ma(
            "/dns4/host.example.com/tcp/4300/http"
        )));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_dns6_tcp_bare() {
        assert!(is_tcp_host_multiaddr(&ma(
            "/dns6/host.example.com/tcp/4300"
        )));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_dns6_tcp_https() {
        assert!(is_tcp_host_multiaddr(&ma(
            "/dns6/host.example.com/tcp/4300/https"
        )));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_rejects_udp() {
        assert!(!is_tcp_host_multiaddr(&ma("/ip4/127.0.0.1/udp/4300")));
        assert!(!is_tcp_host_multiaddr(&ma("/ip6/::1/udp/4300")));
        assert!(!is_tcp_host_multiaddr(&ma(
            "/dns/host.example.com/udp/4300"
        )));
        assert!(!is_tcp_host_multiaddr(&ma(
            "/dns4/host.example.com/udp/4300"
        )));
        assert!(!is_tcp_host_multiaddr(&ma(
            "/dns6/host.example.com/udp/4300"
        )));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_rejects_missing_port() {
        assert!(!is_tcp_host_multiaddr(&ma("/ip4/127.0.0.1")));
        assert!(!is_tcp_host_multiaddr(&ma("/dns/host.example.com")));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_rejects_missing_host() {
        // A multiaddr starting with /tcp has no host class first;
        // libp2p's parser accepts the string but the helper rejects it.
        assert!(!is_tcp_host_multiaddr(&ma("/tcp/4300")));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_rejects_empty() {
        assert!(!is_tcp_host_multiaddr(&Multiaddr::empty()));
    }

    /// Current acceptance: ONE trailing protocol after `/tcp/<port>` is
    /// silently tolerated by the existing `parse_*` parser family
    /// (`parse_http_https` returns `Ok("http")` for anything other than
    /// `Http`/`Https`). The validator inherits that contract — see the
    /// `is_tcp_host_multiaddr` rustdoc TODO listing the rejection
    /// candidates for a future tightening pass.
    ///
    /// This test pins today's behaviour so a future tightening can flip
    /// the assertions deliberately, with reviewer awareness.
    #[test]
    fn test_is_tcp_host_multiaddr_accepts_single_trailing_protocol_today() {
        // /ws after /tcp/<port> is libp2p WebSocket — silently treated
        // as "http" by the existing parser. Today's helper accepts it;
        // the TODO in `is_tcp_host_multiaddr`'s rustdoc lists it as a
        // future rejection candidate.
        assert!(is_tcp_host_multiaddr(&ma("/ip4/127.0.0.1/tcp/4300/ws")));
        // /quic-v1 after /tcp/<port> is malformed (QUIC is UDP-only) but
        // parses; same leniency applies.
        assert!(is_tcp_host_multiaddr(&ma(
            "/ip4/127.0.0.1/tcp/4300/quic-v1"
        )));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_rejects_multi_trailing_junk() {
        // Two trailing protocols after /tcp/<port> exceed the one-slot
        // tolerance — `parse_end` rejects, so the validator rejects.
        assert!(!is_tcp_host_multiaddr(&ma(
            "/ip4/127.0.0.1/tcp/4300/tls/ws"
        )));
    }

    #[test]
    fn test_is_tcp_host_multiaddr_rejects_udp_quic() {
        // UDP at the port slot is rejected regardless of any trailing
        // protocol — the host class + TCP-at-second-slot gate fires
        // first.
        assert!(!is_tcp_host_multiaddr(&ma(
            "/ip4/127.0.0.1/udp/4300/quic-v1"
        )));
    }
}