tetratto/crates/core/src/model/addr.rs

use std::net::SocketAddr;

/// How many bytes should be taken as the prefix (from the begining of the address).
pub(crate) const IPV6_PREFIX_BYTES: usize = 16;

/// The protocol of a [`RemoteAddr`].
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AddrProto {
    IPV4,
    IPV6,
}

/// A representation of a remote IP address.
#[derive(Clone, Debug)]
pub struct RemoteAddr(pub SocketAddr, pub AddrProto);

impl From<&str> for RemoteAddr {
    fn from(value: &str) -> Self {
        if value.len() >= 16 {
            // ipv6 (16 bytes; 128 bits)
            if !(value.starts_with("[") | value.contains("]:1000")) {
                Self(format!("[{value}]:1000").parse().unwrap(), AddrProto::IPV6)
            } else {
                Self(value.parse().unwrap(), AddrProto::IPV6)
            }
        } else {
            // ipv4 (4 bytes; 32 bits)
            if !value.contains(":1000") {
                Self(
                    format!("{value}:1000")
                        .parse()
                        .unwrap_or("0.0.0.0:1000".parse().unwrap()),
                    AddrProto::IPV4,
                )
            } else {
                Self(
                    value.parse().unwrap_or("0.0.0.0:1000".parse().unwrap()),
                    AddrProto::IPV4,
                )
            }
        }
    }
}

impl RemoteAddr {
    /// Get the address' prefix (returns the entire address for IPV4 addresses).
    ///
    ///  Operates on the IP address **without** colon characters.
    pub fn prefix(&self, chop: Option<usize>) -> String {
        if self.1 == AddrProto::IPV4 {
            return self.0.to_string();
        }

        // we're adding 2 bytes to the chop bytes because we also
        // need to remove 2 colons
        let as_string = self.ipv6_inner();
        (&as_string[0..(match chop {
            Some(c) => c,
            None => IPV6_PREFIX_BYTES,
        } + (IPV6_PREFIX_BYTES / 4))])
            .to_string()
    }

    /// [`Self::prefix`], but it returns another [`RemoteAddr`].
    pub fn to_prefix(self, chop: Option<usize>) -> Self {
        Self::from(self.prefix(chop).as_str())
    }

    /// Get the innter content from the address (as ipv6).
    pub fn ipv6_inner(&self) -> String {
        self.0.to_string().replace("[", "").replace("]:1000", "")
    }
}