Cleanup and implementation of Update / Notification parser / serializers.

crates/primitives/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "bgp-primitives"
+
+description = "Primitive types for working with routing information"
+edition.workspace = true
+homepage.workspace = true
+license.workspace = true
+repository.workspace = true
+version.workspace = true
+
+[dependencies]
+bitfield.workspace = true
+bytes.workspace = true
+eyre.workspace = true
+nom.workspace = true
+serde.workspace = true
+serde_repr.workspace = true
+thiserror.workspace = true
+strum.workspace = true

crates/primitives/src/constants.rs

@@ -0,0 +1,80 @@
+use std::fmt::Display;
+
+use eyre::bail;
+use serde_repr::{Deserialize_repr, Serialize_repr};
+
+#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
+#[repr(u16)]
+pub enum AddressFamilyId {
+    Ipv4 = 1,
+    Ipv6 = 2,
+}
+
+impl TryFrom<u16> for AddressFamilyId {
+    type Error = eyre::Error;
+
+    fn try_from(value: u16) -> Result<Self, Self::Error> {
+        Ok(match value {
+            x if x == Self::Ipv4 as u16 => Self::Ipv4,
+            x if x == Self::Ipv6 as u16 => Self::Ipv6,
+            other => bail!("Unknown AddressFamily: {}", other),
+        })
+    }
+}
+
+impl Display for AddressFamilyId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            AddressFamilyId::Ipv4 => write!(f, "Ipv4"),
+            AddressFamilyId::Ipv6 => write!(f, "Ipv6"),
+        }
+    }
+}
+
+/// Represents a Subsequent Address Family Identifier.
+#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
+#[repr(u8)]
+pub enum SubsequentAfi {
+    Unicast = 1,
+    Multicast = 2,
+    NlriWithMpls = 4,
+    MplsLabelledVpn = 128,
+    MulticastMplsVpn = 129,
+}
+
+impl TryFrom<u8> for SubsequentAfi {
+    type Error = eyre::Error;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        Ok(match value {
+            x if x == Self::Unicast as u8 => Self::Unicast,
+            x if x == Self::Multicast as u8 => Self::Multicast,
+            x if x == Self::NlriWithMpls as u8 => Self::NlriWithMpls,
+            x if x == Self::MplsLabelledVpn as u8 => Self::MplsLabelledVpn,
+            x if x == Self::MulticastMplsVpn as u8 => Self::MulticastMplsVpn,
+            other => bail!("Unknown SubsequentAfi: {}", other),
+        })
+    }
+}
+
+/// The implementation try_from for u16 just reuses the u8 implementation since
+/// we don't have any cases where the SAFI is longer than u8::MAX yet.
+impl TryFrom<u16> for SubsequentAfi {
+    type Error = eyre::Error;
+
+    fn try_from(value: u16) -> Result<Self, Self::Error> {
+        TryFrom::<u8>::try_from(value as u8)
+    }
+}
+
+impl Display for SubsequentAfi {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            SubsequentAfi::Unicast => write!(f, "Unicast"),
+            SubsequentAfi::Multicast => write!(f, "Multicast"),
+            SubsequentAfi::NlriWithMpls => write!(f, "NlriWithMpls"),
+            SubsequentAfi::MplsLabelledVpn => write!(f, "MplsLabelledVpn"),
+            SubsequentAfi::MulticastMplsVpn => write!(f, "MulticastMplsVpn"),
+        }
+    }
+}

crates/primitives/src/ip_prefix.rs

@@ -0,0 +1,215 @@
+use std::fmt::Display;
+use std::net::{Ipv4Addr, Ipv6Addr};
+use std::str::FromStr;
+
+use bytes::{Buf, BufMut, BytesMut};
+use eyre::{Result, bail};
+use serde::{Deserialize, Serialize};
+
+use crate::constants::AddressFamilyId;
+use crate::parser::{ParserContext, ToWireError};
+
+/// IpPrefix represents some IP address prefix, for a specific AddressFamilyId.
+#[derive(Debug, PartialEq, Eq, Clone, Hash, Serialize, Deserialize)]
+pub struct IpPrefix {
+    pub address_family: AddressFamilyId,
+    pub prefix: Vec<u8>,
+    pub length: u8,
+}
+
+impl IpPrefix {
+    pub fn new(address_family: AddressFamilyId, prefix: Vec<u8>, length: u8) -> Result<Self> {
+        // Ensure that the prefix we are given contains the right number of bytes corresponding to the prefix length.
+        if prefix.len() < ((length + 7) / 8).into() {
+            bail!(
+                "Mismatched prefix {:?} for given prefix length: {}",
+                prefix,
+                length
+            );
+        }
+        Ok(Self {
+            address_family,
+            prefix,
+            length,
+        })
+    }
+
+    pub fn to_wire(&self, _: &ParserContext, out: &mut BytesMut) -> Result<(), ToWireError> {
+        // Verify that there is enough space to write the IpPrefix.
+        if out.remaining() < (self.prefix.len() + 1) {
+            Err(ToWireError::OutBufferOverflow)?;
+        }
+
+        // Write length and prefix.
+        out.put_u8(self.length);
+        out.put(self.prefix.as_slice());
+
+        Ok(())
+    }
+}
+
+impl TryFrom<&str> for IpPrefix {
+    type Error = eyre::Error;
+
+    fn try_from(value: &str) -> Result<Self, Self::Error> {
+        let parts: Vec<&str> = value.split("/").collect();
+
+        if parts.len() != 2 {
+            bail!(
+                "Expected IpPrefix in format prefix/length but got: {}",
+                value
+            );
+        }
+
+        let length: u8 = u8::from_str_radix(parts[1], 10).map_err(eyre::Error::from)?;
+        let mut octets;
+        let afi: AddressFamilyId;
+
+        if parts[0].contains(":") {
+            afi = AddressFamilyId::Ipv6;
+            let addr: Ipv6Addr = Ipv6Addr::from_str(parts[0]).map_err(eyre::Error::from)?;
+            octets = addr.octets().to_vec();
+        } else if parts[0].contains(".") {
+            afi = AddressFamilyId::Ipv4;
+            let addr: Ipv4Addr = Ipv4Addr::from_str(parts[0]).map_err(eyre::Error::from)?;
+            octets = addr.octets().to_vec();
+        } else {
+            bail!("Could not figure out address type")
+        }
+
+        // Truncate the octets we have to the right number of bytes to match the prefix length..
+        if length % 8 == 0 {
+            // We can cleanly truncate the number of bytes since we are at a byte boundary.
+            octets.truncate((length / 8).into());
+        } else {
+            // We need to keep length % 8 bits of the last byte.
+            let num_bytes = (length / 8) + 1;
+            let mask = u8::MAX << (8 - (length % 8));
+            octets.truncate(num_bytes.into());
+            // Fix up the last byte.
+            let last_pos = octets.len() - 1;
+            octets[last_pos] &= mask;
+        }
+
+        Ok(IpPrefix {
+            address_family: afi,
+            prefix: octets,
+            length,
+        })
+    }
+}
+
+impl Display for IpPrefix {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self.address_family {
+            AddressFamilyId::Ipv4 => {
+                // Pad the prefix with 0 bytes if it's less than 4 bytes long.
+                let bytes = &mut self.prefix.clone();
+                if bytes.len() < 4 {
+                    bytes.extend(std::iter::repeat(0).take(4 - bytes.len()));
+                }
+                let four_bytes: [u8; 4] = bytes
+                    .as_slice()
+                    .try_into()
+                    .map_err(|_| std::fmt::Error {})?;
+                let ipv4_addr = Ipv4Addr::from(four_bytes);
+                write!(f, "{}/{}", ipv4_addr, self.length)
+            }
+            AddressFamilyId::Ipv6 => {
+                let bytes = &mut self.prefix.clone();
+                if bytes.len() < 16 {
+                    bytes.extend(std::iter::repeat(0).take(16 - bytes.len()));
+                }
+                let sixteen_bytes: [u8; 16] = bytes
+                    .as_slice()
+                    .try_into()
+                    .map_err(|_| std::fmt::Error {})?;
+                let ipv6_addr = Ipv6Addr::from(sixteen_bytes);
+                write!(f, "{}/{}", ipv6_addr, self.length)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::constants::AddressFamilyId;
+
+    use super::IpPrefix;
+
+    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.prefix, $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
+    );
+
+    verify_roundtrip!(
+        verify_roundtrip_ipv6_0,
+        "::/0",
+        AddressFamilyId::Ipv6,
+        vec![],
+        0
+    );
+}

crates/primitives/src/lib.rs

@@ -0,0 +1,4 @@
+pub mod constants;
+pub mod ip_prefix;
+pub mod parser;
+pub mod poptrie;

crates/primitives/src/parser.rs

@@ -0,0 +1,61 @@
+use std::fmt::Display;
+
+use nom::error::{ErrorKind, ParseError};
+
+use crate::constants::AddressFamilyId;
+
+#[derive(Debug, Default)]
+pub struct ParserContext {
+    /// Whether thi parser is being run with a peer that is RFC6793 compliant.
+    pub four_octet_asn: Option<bool>,
+    /// Which address family should be parsed by default with this parser.
+    pub address_family: Option<AddressFamilyId>,
+}
+
+#[derive(Debug)]
+pub enum ToWireError {
+    /// There was not enough space in the output buffer to serialize the data into.
+    OutBufferOverflow,
+    /// Another error.
+    Other(eyre::Error),
+}
+
+impl Display for ToWireError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ToWireError::OutBufferOverflow => write!(f, "OutBufferOverflow"),
+            ToWireError::Other(report) => report.fmt(f),
+        }
+    }
+}
+
+impl std::error::Error for ToWireError {}
+
+// Custom error type for the parser.
+#[derive(Debug)]
+pub enum BgpParserError<I> {
+    /// ProtocolError encapsulates protocol level semantics problems that require
+    /// action from the caller of the parser.
+    ProtocolError(ProtocolErrorValues),
+    CustomText(&'static str),
+    Eyre(eyre::ErrReport),
+    Nom(I, ErrorKind),
+}
+
+impl<I> ParseError<I> for BgpParserError<I> {
+    fn from_error_kind(input: I, kind: ErrorKind) -> Self {
+        BgpParserError::Nom(input, kind)
+    }
+    fn append(_: I, _: ErrorKind, other: Self) -> Self {
+        other
+    }
+}
+
+#[derive(Debug)]
+pub enum ProtocolErrorValues {
+    /// `UnsupportedVersion` is returned when the BGP Open message received is for a
+    /// version other than one we support. The unsupported version number is the value.
+    UnsupportedVersion(u8),
+    /// `UnknownOpenOption` is returned with the open option type code that is unknown.
+    UnknownOpenOption(u8),
+}

crates/primitives/src/poptrie/mod.rs

@@ -0,0 +1,176 @@
+use std::collections::BTreeMap;
+
+#[derive(Debug)]
+pub struct Poptrie<T> {
+    pub inners: Vec<Inner>,
+    pub leaves: Vec<Leaf<T>>,
+}
+
+#[derive(Debug)]
+pub struct Inner {
+    /// A bit vector that describes the child internal nodes.
+    pub internal_v: u64,
+    /// A bit vector that describes the child leaf nodes.
+    pub leaf_v: u64,
+    /// Starting offset where the child interior nodes of this node begins.
+    pub inner_offset: u64,
+    /// Starting offset where the child leaf nodes of this node begins.
+    pub leaf_offset: u64,
+}
+
+#[derive(Debug)]
+pub struct Leaf<T> {
+    pub data: T,
+}
+
+#[inline]
+fn extract(width: u8, offset: u8, v: u32, bits: u8) -> u8 {
+    let shift = bits.saturating_sub(width * (offset + 1));
+    let mask = 0b111111 << shift;
+    ((v & mask) >> shift) as u8
+}
+
+fn extract_32(width: u8, offset: u8, v: u32) -> u8 {
+    extract(width, offset, v, 32)
+}
+
+fn extract_128(width: u8, offset: u8, v: u32) -> u8 {
+    extract(width, offset, v, 128)
+}
+
+impl<T: Clone> Poptrie<T> {
+    pub fn match_u32(&self, addr: u32) -> Option<&T> {
+        let mut i = 0_u64;
+        let mut v = self.inners.get(i as usize)?.internal_v;
+        let mut offset = 0;
+        let mut n = extract(6, offset, addr, 32);
+
+        let mut result = None;
+
+        while (v & (1 << n)) != 0 {
+            let base = self.inners.get(i as usize)?.inner_offset;
+            let arg = v & ((2 << n) - 1);
+            let bc = arg.count_ones() as u64;
+            i = base + bc - 1;
+            v = self.inners.get(i as usize)?.internal_v;
+
+            offset += 1;
+            n = extract(6, offset, addr, 32);
+
+            let base = self.inners.get(i as usize)?.leaf_offset;
+            let v = self.inners.get(i as usize)?.leaf_v;
+            if (v & (1 << n)) != 0 {
+                let i = base - 1;
+                result = Some(&self.leaves.get(i as usize)?.data);
+            }
+        }
+
+        let base = self.inners.get(i as usize)?.leaf_offset;
+        let v = self.inners.get(i as usize)?.leaf_v;
+        if (v & (1 << n)) != 0 {
+            i = base - 1;
+            result = Some(&self.leaves.get(i as usize)?.data);
+        }
+
+        result
+    }
+
+    fn new(&mut self, rt: BTreeMap<(u32, u8), T>) {
+        const DEPTH: u8 = 6;
+        let mut forest = vec![(0, rt)];
+
+        let mut i_offset = 1;
+        for _ in 0..DEPTH {
+            let mut subforest = Vec::<(u8, BTreeMap<(u32, u8), T>)>::new();
+            let mut children = 0;
+            let mut siblings = 0;
+            for (_, tree) in &forest {
+                let mut iv = 0_u64;
+                let mut lv = 0_u64;
+
+                let mut subsubforest = Vec::<(u8, BTreeMap<(u32, u8), T>)>::new();
+                for (r, e) in tree {
+                    let k = extract(6, DEPTH, r.0, 32);
+                    let consumed = core::cmp::min((DEPTH + 1) * 6, 32);
+                    if r.1 < consumed {
+                        if ((1 << k) & iv) == 0 {
+                            lv |= 1 << k;
+                            self.leaves.push(Leaf { data: e.clone() });
+                        }
+
+                        // If the prefix of the route entry is less than but not equal
+                        // to the consumed number of bits, we need to add those bits
+                        // to the bitvec.
+                        if r.1 != consumed {
+                            // Shift the extra bits and add to the bitvec for this
+                            // internal node.
+                            let extra = 1_u32 << (consumed - r.1);
+                            for i in 1..(extra).min(64 - k as u32) {
+                                lv |= 1_u64 << (k as u32 + i);
+                                self.leaves.push(Leaf { data: e.clone() });
+                            }
+                        }
+                        continue;
+                    }
+
+                    iv |= 1 << k;
+                    match subsubforest.iter_mut().find(|x| x.0 == k) {
+                        Some(ref mut entry) => {
+                            entry.1.insert(*r, e.clone());
+                        }
+                        None => {
+                            let mut tbl = BTreeMap::new();
+                            tbl.insert(*r, e.clone());
+                            subsubforest.push((k, tbl));
+                            if iv > 0 {
+                                children += 1;
+                            }
+                        }
+                    }
+                }
+
+                if iv > 0 || lv > 0 {
+                    self.inners.push(Inner {
+                        internal_v: iv,
+                        leaf_v: lv,
+                        inner_offset: if iv > 0 { i_offset + siblings } else { 0 },
+                        leaf_offset: self.leaves.len() as u64,
+                    });
+                    if iv > 0 {
+                        siblings += 1;
+                    }
+                }
+
+                subforest.extend_from_slice(subsubforest.as_slice());
+            }
+            i_offset += children;
+            forest = subforest;
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{collections::BTreeMap, net::Ipv4Addr};
+
+    use crate::poptrie::Poptrie;
+
+    #[test]
+    fn test_poptrie() {
+        let mut map: BTreeMap<(u32, u8), Ipv4Addr> = Default::default();
+        map.insert(
+            (Ipv4Addr::new(127, 0, 0, 0).into(), 8),
+            Ipv4Addr::new(10, 0, 0, 1),
+        );
+
+        let mut poptrie: Poptrie<Ipv4Addr> = Poptrie {
+            inners: vec![],
+            leaves: vec![],
+        };
+
+        poptrie.new(map);
+        println!("Success");
+
+        dbg!(poptrie);
+    }
+}