local-ip-address 0.6.7

use std::collections::HashMap;
use std::ffi::{CStr, c_int};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

use neli::attr::Attribute;
use neli::consts::nl::{NlmF};
use neli::consts::socket::NlFamily;
use neli::consts::rtnl::{Ifa, Ifla, RtAddrFamily, RtScope, RtTable, Rta, Rtm, RtmF, Rtn, Rtprot};
use neli::err::RouterError;
use neli::nl::{NlPayload, Nlmsghdr};
use neli::router::synchronous::NlRouter;
use neli::rtnl::{
    Ifaddrmsg, IfaddrmsgBuilder, Ifinfomsg, IfinfomsgBuilder, RtattrBuilder, Rtmsg, RtmsgBuilder,
};
use neli::types::RtBuffer;
use neli::consts::rtnl::RtAddrFamily::{Inet, Inet6};
use neli::utils::Groups;

use crate::Error;

#[cfg(target_env = "gnu")]
const RTM_FLAGS_LOOKUP: &[RtmF] = &[RtmF::LOOKUPTABLE];
#[cfg(not(target_env = "gnu"))]
const RTM_FLAGS_LOOKUP: &[RtmF] = &[];

/// Retrieves the local IPv4 address for this system
pub fn local_ip() -> Result<IpAddr, Error> {
    local_ip_impl(Inet)
}

/// Retrieves the local IPv6 address for this system
pub fn local_ipv6() -> Result<IpAddr, Error> {
    local_ip_impl(Inet6)
}

/// Retrieves the local broadcast IPv4 address for this system
pub fn local_broadcast_ip() -> Result<IpAddr, Error> {
    local_broadcast_impl(Inet)
}

fn local_broadcast_impl(family: RtAddrFamily) -> Result<IpAddr, Error> {
    let (netlink_socket, _) = NlRouter::connect(NlFamily::Route, None, Groups::empty())
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    let pref_ip = local_ip()?;

    let ifaddrmsg = IfaddrmsgBuilder::default()
        .ifa_family(family)
        .build()
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    let recv = netlink_socket
        .send(
            Rtm::Getaddr,
            NlmF::REQUEST | NlmF::ROOT,
            NlPayload::Payload(ifaddrmsg),
        )
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    let mut broadcast_ip = None;
    for response in recv {
        let header: Nlmsghdr<Rtm, Ifaddrmsg> = response.map_err(|_| {
            Error::StrategyError(String::from(
                "An error occurred retrieving Netlink's socket response",
            ))
        })?;

        if let NlPayload::Empty = header.nl_payload() {
            continue;
        }

        if *header.nl_type() != Rtm::Newaddr {
            return Err(Error::StrategyError(String::from(
                "The Netlink header type is not the expected",
            )));
        }

        let p = header.get_payload().ok_or_else(|| {
            Error::StrategyError(String::from(
                "An error occurred getting Netlink's header payload",
            ))
        })?;

        if *p.ifa_scope() != RtScope::Universe {
            continue;
        }

        if *p.ifa_family() != family {
            Err(Error::StrategyError(format!(
                "Invalid family in Netlink payload: {:?}",
                p.ifa_family()
            )))?
        }

        let mut is_match = false;
        for rtattr in p.rtattrs().iter() {
            if *rtattr.rta_type() == Ifa::Local {
                if *p.ifa_family() == Inet {
                    let addr = Ipv4Addr::from(u32::from_be(
                        rtattr.get_payload_as::<u32>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    is_match = pref_ip == IpAddr::V4(addr);
                } else {
                    let addr = Ipv6Addr::from(u128::from_be(
                        rtattr.get_payload_as::<u128>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    is_match = pref_ip == IpAddr::V6(addr);
                }
            }
            if is_match && *rtattr.rta_type() == Ifa::Broadcast && *p.ifa_family() == Inet {
                let addr = Ipv4Addr::from(u32::from_be(
                    rtattr.get_payload_as::<u32>().map_err(|_| {
                        Error::StrategyError(String::from(
                            "An error occurred retrieving Netlink's route payload broadcast attribute",
                        ))
                    })?,
                ));
                return Ok(IpAddr::V4(addr));
            }
            if *rtattr.rta_type() == Ifa::Broadcast && *p.ifa_family() == Inet {
                let addr = Ipv4Addr::from(u32::from_be(rtattr.get_payload_as::<u32>().map_err(
                    |_| {
                        Error::StrategyError(String::from(
                            "An error occurred retrieving Netlink's route payload attribute",
                        ))
                    },
                )?));
                broadcast_ip = Some(IpAddr::V4(addr));
            }
        }
    }

    if let Some(broadcast_ip) = broadcast_ip {
        return Ok(broadcast_ip);
    }

    Err(Error::LocalIpAddressNotFound)
}

fn local_ip_impl(family: RtAddrFamily) -> Result<IpAddr, Error> {
    let (netlink_socket, _) = NlRouter::connect(NlFamily::Route, None, Groups::empty())
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    match local_ip_impl_route(family, &netlink_socket) {
        Ok(ip_addr) => Ok(ip_addr),
        Err(Error::LocalIpAddressNotFound) => local_ip_impl_addr(family, &netlink_socket),
        Err(e) => Err(e),
    }
}

fn local_ip_impl_route(family: RtAddrFamily, netlink_socket: &NlRouter) -> Result<IpAddr, Error> {
    let route_attr = match family {
        Inet => {
            let dstip = Ipv4Addr::new(192, 0, 2, 0); // reserved external IP
            let raw_dstip = u32::from(dstip).to_be();
            RtattrBuilder::default()
                .rta_type(Rta::Dst)
                .rta_payload(raw_dstip)
                .build()
        }
        Inet6 => {
            let dstip = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0); // reserved external IP
            let raw_dstip = u128::from(dstip).to_be();
            RtattrBuilder::default()
                .rta_type(Rta::Dst)
                .rta_payload(raw_dstip)
                .build()
        }
        _ => Err(Error::StrategyError(format!(
            "Invalid address family given: {:#?}",
            family
        )))?,
    };

    let mut ifroutemsg = RtmsgBuilder::default()
        .rtm_family(family)
        .rtm_src_len(0)
        .rtm_dst_len(0)
        .rtm_tos(0)
        .rtm_table(RtTable::Unspec)
        .rtm_protocol(Rtprot::Unspec)
        .rtm_scope(RtScope::Universe)
        .rtm_type(Rtn::Unspec);

    let route_attr = route_attr.map_err(|err| Error::StrategyError(err.to_string()))?;
    let mut route_payload = RtBuffer::new();
    route_payload.push(route_attr);
    ifroutemsg = ifroutemsg.rtattrs(route_payload);

    let rtm_flags = RTM_FLAGS_LOOKUP.iter().cloned().reduce(|a, b| a | b);
    if let Some(flags) = rtm_flags {
        ifroutemsg = ifroutemsg.rtm_flags(flags);
    }

    let ifroutemsg = ifroutemsg
        .build()
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    let recv = netlink_socket
        .send(Rtm::Getroute, NlmF::REQUEST, NlPayload::Payload(ifroutemsg))
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    for response in recv {
        let header: Nlmsghdr<Rtm, Rtmsg> = response.map_err(|err| {
            if let RouterError::Nlmsgerr(ref err) = err {
                if *err.error() == -libc::ENETUNREACH {
                    return Error::LocalIpAddressNotFound;
                }
            }
            Error::StrategyError(format!(
                "An error occurred retrieving Netlink's socket response: {err}",
            ))
        })?;

        if let NlPayload::Empty = *header.nl_payload() {
            continue;
        }

        if *header.nl_type() != Rtm::Newroute {
            return Err(Error::StrategyError(String::from(
                "The Netlink header type is not the expected",
            )));
        }

        let p = header.get_payload().ok_or_else(|| {
            Error::StrategyError(String::from(
                "An error occurred getting Netlink's header payload",
            ))
        })?;

        if *p.rtm_scope() != RtScope::Universe {
            continue;
        }

        if *p.rtm_family() != family {
            Err(Error::StrategyError(format!(
                "Invalid address family in Netlink payload: {:?}",
                p.rtm_family()
            )))?
        }

        for rtattr in p.rtattrs().iter() {
            if *rtattr.rta_type() == Rta::Prefsrc {
                if *p.rtm_family() == Inet {
                    let addr = Ipv4Addr::from(u32::from_be(
                        rtattr.get_payload_as::<u32>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    return Ok(IpAddr::V4(addr));
                } else {
                    let addr = Ipv6Addr::from(u128::from_be(
                        rtattr.get_payload_as::<u128>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    return Ok(IpAddr::V6(addr));
                }
            }
        }
    }
    Err(Error::LocalIpAddressNotFound)
}

fn local_ip_impl_addr(family: RtAddrFamily, netlink_socket: &NlRouter) -> Result<IpAddr, Error> {
    let ifaddrmsg = IfaddrmsgBuilder::default()
        .ifa_family(family)
        .build()
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    let recv = netlink_socket
        .send(
            Rtm::Getaddr,
            NlmF::REQUEST | NlmF::ROOT,
            NlPayload::Payload(ifaddrmsg),
        )
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    for response in recv {
        let header: Nlmsghdr<Rtm, Ifaddrmsg> = response.map_err(|_| {
            Error::StrategyError(String::from(
                "An error occurred retrieving Netlink's socket response",
            ))
        })?;

        if let NlPayload::Empty = *header.nl_payload() {
            continue;
        }

        if *header.nl_type() != Rtm::Newaddr {
            return Err(Error::StrategyError(String::from(
                "The Netlink header type is not the expected",
            )));
        }

        let p = header.get_payload().ok_or_else(|| {
            Error::StrategyError(String::from(
                "An error occurred getting Netlink's header payload",
            ))
        })?;

        if *p.ifa_scope() != RtScope::Universe {
            continue;
        }

        if *p.ifa_family() != family {
            Err(Error::StrategyError(format!(
                "Invalid family in Netlink payload: {:?}",
                p.ifa_family()
            )))?
        }

        for rtattr in p.rtattrs().iter() {
            if *rtattr.rta_type() == Ifa::Local {
                if *p.ifa_family() == Inet {
                    let addr = Ipv4Addr::from(u32::from_be(
                        rtattr.get_payload_as::<u32>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    return Ok(IpAddr::V4(addr));
                } else {
                    let addr = Ipv6Addr::from(u128::from_be(
                        rtattr.get_payload_as::<u128>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    return Ok(IpAddr::V6(addr));
                }
            }
        }
    }

    Err(Error::LocalIpAddressNotFound)
}

/// Perform a search over the system's network interfaces using Netlink Route information,
/// retrieved network interfaces belonging to both socket address families
/// `AF_INET` and `AF_INET6` are retrieved along with the interface address name.
///
/// # Example
///
/// ```
/// use std::net::IpAddr;
/// use local_ip_address::list_afinet_netifas;
///
/// let ifas = list_afinet_netifas().unwrap();
///
/// if let Some((_, ipaddr)) = ifas
/// .iter()
/// .find(|(name, ipaddr)| *name == "en0" && matches!(ipaddr, IpAddr::V4(_))) {
///     // This is your local IP address: 192.168.1.111
///     println!("This is your local IP address: {:?}", ipaddr);
/// }
/// ```
pub fn list_afinet_netifas() -> Result<Vec<(String, IpAddr)>, Error> {
    let (netlink_socket, _) = NlRouter::connect(NlFamily::Route, None, Groups::empty())
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    // First get list of interfaces via RTM_GETLINK

    let ifinfomsg = IfinfomsgBuilder::default()
        .ifi_family(RtAddrFamily::Unspecified)
        .build()
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    let recv = netlink_socket
        .send(
            Rtm::Getlink,
            NlmF::REQUEST | NlmF::DUMP,
            NlPayload::Payload(ifinfomsg),
        )
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    let mut if_indexes = HashMap::new();
    for response in recv {
        let header: Nlmsghdr<Rtm, Ifinfomsg> = response.map_err(|_| {
            Error::StrategyError(String::from(
                "An error occurred retrieving Netlink's socket response",
            ))
        })?;

        if let NlPayload::Empty = *header.nl_payload() {
            continue;
        }

        if *header.nl_type() != Rtm::Newlink {
            continue;
        }

        let p = header.get_payload().ok_or_else(|| {
            Error::StrategyError(String::from(
                "An error occurred getting Netlink's header payload",
            ))
        })?;

        for rtattr in p.rtattrs().iter() {
            if *rtattr.rta_type() == Ifla::Ifname {
                let ifname = parse_ifname(rtattr.payload().as_ref())?;
                if_indexes.insert(*p.ifi_index(), ifname);
                break;
            }
        }
    }

    // Secondly get addresses of interfaces via RTM_GETADDR

    let ifaddrmsg = IfaddrmsgBuilder::default()
        .ifa_family(RtAddrFamily::Unspecified)
        .ifa_prefixlen(0)
        .ifa_scope(RtScope::Universe)
        .ifa_index(0)
        .build()
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    let recv = netlink_socket
        .send(
            Rtm::Getaddr,
            NlmF::REQUEST | NlmF::DUMP,
            NlPayload::Payload(ifaddrmsg),
        )
        .map_err(|err| Error::StrategyError(err.to_string()))?;

    let mut interfaces = Vec::new();
    for response in recv {
        let header: Nlmsghdr<Rtm, Ifaddrmsg> = response.map_err(|err| {
            Error::StrategyError(format!(
                "An error occurred retrieving Netlink's socket response: {err}"
            ))
        })?;

        if let NlPayload::Empty = header.nl_payload() {
            continue;
        }

        if *header.nl_type() != Rtm::Newaddr {
            continue;
        }

        let p = header.get_payload().ok_or_else(|| {
            Error::StrategyError(String::from(
                "An error occurred getting Netlink's header payload",
            ))
        })?;

        if *p.ifa_family() != Inet6 && *p.ifa_family() != Inet {
            Err(Error::StrategyError(format!(
                "Netlink payload has unsupported family: {:?}",
                p.ifa_family()
            )))?
        }

        let mut ipaddr = None;
        let mut label = None;

        for rtattr in p.rtattrs().iter() {
            if *rtattr.rta_type() == Ifa::Label {
                let ifname = parse_ifname(rtattr.payload().as_ref())?;
                label = Some(ifname);
            } else if *rtattr.rta_type() == Ifa::Address {
                if ipaddr.is_some() {
                    // do not override IFA_LOCAL
                    continue;
                }
                if *p.ifa_family() == Inet6 {
                    let rtaddr = Ipv6Addr::from(u128::from_be(
                        rtattr.get_payload_as::<u128>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    ipaddr = Some(IpAddr::V6(rtaddr));
                } else {
                    let rtaddr = Ipv4Addr::from(u32::from_be(
                        rtattr.get_payload_as::<u32>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    ipaddr = Some(IpAddr::V4(rtaddr));
                }
            } else if *rtattr.rta_type() == Ifa::Local {
                if *p.ifa_family() == Inet6 {
                    let rtlocal = Ipv6Addr::from(u128::from_be(
                        rtattr.get_payload_as::<u128>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    ipaddr = Some(IpAddr::V6(rtlocal));
                } else {
                    let rtlocal = Ipv4Addr::from(u32::from_be(
                        rtattr.get_payload_as::<u32>().map_err(|_| {
                            Error::StrategyError(String::from(
                                "An error occurred retrieving Netlink's route payload attribute",
                            ))
                        })?,
                    ));
                    ipaddr = Some(IpAddr::V4(rtlocal));
                }
            }
        }

        if let Some(ipaddr) = ipaddr {
            if let Some(ifname) = label {
                interfaces.push((ifname, ipaddr));
            } else if let Some(ifname) = if_indexes.get(&(*p.ifa_index() as c_int)) {
                interfaces.push((ifname.clone(), ipaddr));
            }
        }
    }

    Ok(interfaces)
}

/// Parse network interface name of slice type to string type.
/// If the slice is suffixed with '\0', this suffix will be removed when parsing.
fn parse_ifname(bytes: &[u8]) -> Result<String, Error> {
    let ifname = if bytes.ends_with(&[0u8]) {
        CStr::from_bytes_with_nul(bytes)
            .map_err(|err| {
                Error::StrategyError(format!(
                    "An error occurred converting interface name to string: {err}",
                ))
            })?
            .to_string_lossy()
            .to_string()
    } else {
        String::from_utf8_lossy(bytes).to_string()
    };

    Ok(ifname)
}

#[cfg(test)]
mod tests {
    use crate::linux::parse_ifname;

    #[test]
    fn parse_ifname_without_nul() {
        let expected = "hello, world";
        let bytes = [104, 101, 108, 108, 111, 44, 32, 119, 111, 114, 108, 100];
        let res = parse_ifname(&bytes);
        assert!(res.is_ok());
        assert_eq!(res.unwrap(), expected);
    }

    #[test]
    fn parse_ifname_with_nul() {
        let expected = "hello, world";
        let bytes = [104, 101, 108, 108, 111, 44, 32, 119, 111, 114, 108, 100, 0];
        let res = parse_ifname(&bytes);
        assert!(res.is_ok());
        assert_eq!(res.unwrap(), expected);
    }

    #[test]
    fn parse_ifname_only_nul() {
        let expected = "";
        let bytes = [0u8];
        let res = parse_ifname(&bytes);
        assert!(res.is_ok());
        assert_eq!(res.unwrap(), expected);
    }
}