ferrix/crates/primitives/src/ip_prefix.rs

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

use bytes::{Buf, BufMut, BytesMut};
use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::constants::AddressFamilyId;
use crate::parser::{EncodeError, ParserContext};

/// Why something is not a valid prefix.
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum PrefixError {
    #[error("prefix {0:?} is not in address/length form")]
    Format(String),
    #[error("prefix {0:?} has an address that is neither IPv4 nor IPv6")]
    Address(String),
    #[error("prefix length {length} is longer than the {max} bits of an address")]
    LengthTooLong { length: u8, max: u8 },
    #[error("prefix length {length} needs {needed} address bytes, got {got}")]
    TooFewBytes {
        length: u8,
        needed: usize,
        got: usize,
    },
}

/// An IPv4 or IPv6 prefix.
///
/// It is a small value with no allocation behind it, cheap to copy, compare
/// and hash: a routing table holds millions of them and looks them up by
/// value. A prefix has one representation: the bits past its length are
/// always zero, so two prefixes are equal exactly when they are the same
/// prefix.
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, Serialize, Deserialize)]
pub struct IpPrefix {
    address_family: AddressFamilyId,
    length: u8,
    /// The address, left-aligned, zero past `length` bits.
    address: [u8; 16],
}

impl IpPrefix {
    /// Makes a prefix from the leading bytes of its address. `prefix` must
    /// hold at least the bytes that `length` bits cover; anything past
    /// those bits is ignored (RFC 4271 §4.3: such bits are irrelevant).
    pub fn new(
        address_family: AddressFamilyId,
        prefix: impl AsRef<[u8]>,
        length: u8,
    ) -> Result<Self, PrefixError> {
        let prefix = prefix.as_ref();
        let max = address_family.bits();
        if length > max {
            return Err(PrefixError::LengthTooLong { length, max });
        }
        let needed = usize::from(length.div_ceil(8));
        if prefix.len() < needed {
            return Err(PrefixError::TooFewBytes {
                length,
                needed,
                got: prefix.len(),
            });
        }
        let mut address = [0u8; 16];
        address[..needed].copy_from_slice(&prefix[..needed]);
        if let (Some(last), rem @ 1..) = (address[..needed].last_mut(), length % 8) {
            *last &= u8::MAX << (8 - rem);
        }
        Ok(Self {
            address_family,
            length,
            address,
        })
    }

    pub fn address_family(&self) -> AddressFamilyId {
        self.address_family
    }

    /// The prefix length in bits.
    pub fn length(&self) -> u8 {
        self.length
    }

    /// The bytes the prefix length covers, as they go on the wire: none for
    /// a default route, three for a /24.
    pub fn bytes(&self) -> &[u8] {
        &self.address[..usize::from(self.length.div_ceil(8))]
    }

    /// Returns true if `other` is equal to or more specific than `self`, i.e.
    /// `self` covers `other`. Prefixes of different address families never
    /// contain each other.
    pub fn contains(&self, other: &IpPrefix) -> bool {
        if self.address_family != other.address_family || other.length < self.length {
            return false;
        }
        let full_bytes = usize::from(self.length / 8);
        if self.address[..full_bytes] != other.address[..full_bytes] {
            return false;
        }
        let rem_bits = self.length % 8;
        if rem_bits == 0 {
            return true;
        }
        let mask = u8::MAX << (8 - rem_bits);
        (self.address[full_bytes] & mask) == (other.address[full_bytes] & mask)
    }

    pub fn to_wire(&self, _: &ParserContext, out: &mut BytesMut) -> Result<(), EncodeError> {
        let bytes = self.bytes();
        // Verify that there is enough space to write the IpPrefix.
        if out.remaining() < (bytes.len() + 1) {
            return Err(EncodeError::OutBufferOverflow);
        }

        // Write length and prefix.
        out.put_u8(self.length);
        out.put(bytes);

        Ok(())
    }
}

impl TryFrom<&str> for IpPrefix {
    type Error = PrefixError;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        let format = || PrefixError::Format(value.to_string());
        let address = || PrefixError::Address(value.to_string());
        let (addr, length) = value.split_once('/').ok_or_else(format)?;
        let length: u8 = length.parse().map_err(|_| format())?;

        let (afi, octets) = if addr.contains(':') {
            let addr = Ipv6Addr::from_str(addr).map_err(|_| address())?;
            (AddressFamilyId::Ipv6, addr.octets().to_vec())
        } else {
            let addr = Ipv4Addr::from_str(addr).map_err(|_| address())?;
            (AddressFamilyId::Ipv4, addr.octets().to_vec())
        };
        IpPrefix::new(afi, octets, length)
    }
}

impl Display for IpPrefix {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.address_family {
            AddressFamilyId::Ipv4 => {
                let [a, b, c, d, ..] = self.address;
                write!(f, "{}/{}", Ipv4Addr::new(a, b, c, d), self.length)
            }
            AddressFamilyId::Ipv6 => {
                write!(f, "{}/{}", Ipv6Addr::from(self.address), self.length)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::constants::AddressFamilyId;

    use proptest::prelude::*;

    use super::{IpPrefix, PrefixError};

    macro_rules! verify_roundtrip {
        ($name:ident, $prefix_str:expr, $afi:expr, $bytes:expr, $length:expr) => {
            #[test]
            fn $name() {
                let ip_prefix = IpPrefix::try_from($prefix_str).unwrap();

                assert_eq!(ip_prefix.address_family(), $afi);
                assert_eq!(ip_prefix.bytes(), $bytes);
                assert_eq!(ip_prefix.length(), $length);

                let to_str: &str = &ip_prefix.to_string();
                assert_eq!(IpPrefix::try_from(to_str).unwrap(), ip_prefix);
            }
        };
    }

    verify_roundtrip!(
        verify_roundtrip_ipv4_24,
        "10.1.2.0/24",
        AddressFamilyId::Ipv4,
        vec![10, 1, 2],
        24
    );

    // Verify truncation.
    verify_roundtrip!(
        verify_roundtrip_ipv4_19,
        "10.245.123.0/19",
        AddressFamilyId::Ipv4,
        vec![10, 245, 96],
        19
    );

    // Verify truncation.
    verify_roundtrip!(
        verify_roundtrip_ipv4_3,
        "192.168.1.0/3",
        AddressFamilyId::Ipv4,
        vec![192],
        3
    );

    // Verify default address.
    verify_roundtrip!(
        verify_roundtrip_ipv4_0,
        "0.0.0.0/0",
        AddressFamilyId::Ipv4,
        vec![],
        0
    );

    verify_roundtrip!(
        verify_roundtrip_ipv6_48,
        "2001:db8:cafe::/48",
        AddressFamilyId::Ipv6,
        vec![32, 1, 13, 184, 202, 254],
        48
    );

    // Verify truncation.
    verify_roundtrip!(
        verify_roundtrip_ipv6_32,
        "2001:db8:cafe::/32",
        AddressFamilyId::Ipv6,
        vec![32, 1, 13, 184],
        32
    );

    #[test]
    fn text_that_is_not_a_prefix_is_rejected_with_the_reason() {
        for (text, expected) in [
            ("10.0.0.0", PrefixError::Format("10.0.0.0".into())),
            ("10.0.0.0/x", PrefixError::Format("10.0.0.0/x".into())),
            ("10.0.0/8", PrefixError::Address("10.0.0/8".into())),
            ("nonsense/8", PrefixError::Address("nonsense/8".into())),
            (
                "10.0.0.0/33",
                PrefixError::LengthTooLong {
                    length: 33,
                    max: 32,
                },
            ),
            (
                "2001:db8::/129",
                PrefixError::LengthTooLong {
                    length: 129,
                    max: 128,
                },
            ),
        ] {
            assert_eq!(IpPrefix::try_from(text), Err(expected), "{text}");
        }
        assert!(IpPrefix::try_from("0.0.0.0/0").is_ok());
        assert!(IpPrefix::try_from("2001:db8::1/128").is_ok());
        assert_eq!(
            IpPrefix::new(AddressFamilyId::Ipv4, vec![10], 16),
            Err(PrefixError::TooFewBytes {
                length: 16,
                needed: 2,
                got: 1
            })
        );
    }

    #[test]
    fn a_prefix_has_one_representation() {
        // Extra bytes and stray bits past the length are dropped.
        let noisy = IpPrefix::new(AddressFamilyId::Ipv4, vec![10, 0xff, 3, 4], 12).unwrap();
        assert_eq!(noisy, IpPrefix::try_from("10.240.0.0/12").unwrap());
        assert_eq!(noisy.bytes(), [10, 0xf0]);
        assert_eq!(
            IpPrefix::new(AddressFamilyId::Ipv4, vec![0xff; 13], 99),
            Err(PrefixError::LengthTooLong {
                length: 99,
                max: 32
            })
        );
    }

    #[test]
    fn a_prefix_is_a_small_value() {
        assert!(size_of::<IpPrefix>() <= 20);
        let prefix = IpPrefix::try_from("10.1.2.0/24").unwrap();
        let copy = prefix;
        assert_eq!(copy, prefix);
        assert_eq!(prefix.bytes(), [10, 1, 2]);
        assert_eq!(IpPrefix::try_from("0.0.0.0/0").unwrap().bytes(), [0u8; 0]);
    }

    proptest! {
        #[test]
        fn text_round_trips(version: bool, bytes: [u8; 16], length in 0u8..=128) {
            let version = if version { AddressFamilyId::Ipv6 } else { AddressFamilyId::Ipv4 };
            let length = length.min(version.bits());
            let prefix = IpPrefix::new(version, bytes, length).unwrap();
            let text = prefix.to_string();
            prop_assert_eq!(IpPrefix::try_from(text.as_str()), Ok(prefix));
            prop_assert!(prefix.contains(&prefix));
        }

        #[test]
        fn text_that_is_noise_is_an_error_or_a_prefix_that_round_trips(text in "\\PC{0,40}") {
            if let Ok(prefix) = IpPrefix::try_from(text.as_str()) {
                let shown = prefix.to_string();
                prop_assert_eq!(IpPrefix::try_from(shown.as_str()), Ok(prefix));
            }
        }

        /// `contains` agrees with comparing the leading bits by hand.
        #[test]
        fn contains_is_a_comparison_of_leading_bits(
            a: [u8; 4], a_len in 0u8..=32, b: [u8; 4], b_len in 0u8..=32,
        ) {
            let outer = IpPrefix::new(AddressFamilyId::Ipv4, a, a_len).unwrap();
            let inner = IpPrefix::new(AddressFamilyId::Ipv4, b, b_len).unwrap();
            let mask = if a_len == 0 { 0 } else { u32::MAX << (32 - a_len) };
            let expected = b_len >= a_len
                && (u32::from_be_bytes(a) & mask) == (u32::from_be_bytes(b) & mask);
            prop_assert_eq!(outer.contains(&inner), expected);
        }
    }

    #[test]
    fn contains_covers_more_specifics() {
        let p = |s: &str| IpPrefix::try_from(s).unwrap();
        assert!(p("10.100.0.0/16").contains(&p("10.100.7.0/24")));
        assert!(p("10.100.0.0/16").contains(&p("10.100.0.0/16")));
        assert!(!p("10.100.0.0/16").contains(&p("10.101.0.0/24")));
        assert!(!p("10.100.7.0/24").contains(&p("10.100.0.0/16")));
        assert!(p("0.0.0.0/0").contains(&p("192.0.2.0/24")));
        // Non-byte-aligned boundary.
        assert!(p("10.100.0.0/19").contains(&p("10.100.31.0/24")));
        assert!(!p("10.100.0.0/19").contains(&p("10.100.32.0/24")));
        // Cross-family never matches.
        assert!(!p("::/0").contains(&p("10.0.0.0/8")));
        assert!(p("2001:db8::/32").contains(&p("2001:db8:cafe::/48")));
    }

    verify_roundtrip!(
        verify_roundtrip_ipv6_0,
        "::/0",
        AddressFamilyId::Ipv6,
        vec![],
        0
    );
}