Add fields for mandatory BGP attributes

src/models/bgp/attributes/mod.rs

@@ -5,12 +5,11 @@ mod origin;
 
 use crate::models::network::*;
 use bitflags::bitflags;
+use itertools::chain;
 use num_enum::{FromPrimitive, IntoPrimitive};
 use std::cmp::Ordering;
-use std::iter::{FromIterator, Map};
+use std::iter::FromIterator;
 use std::net::IpAddr;
-use std::slice::Iter;
-use std::vec::IntoIter;
 
 use crate::error::BgpValidationWarning;
 use crate::models::*;
@@ -138,25 +137,90 @@ pub fn get_deprecated_attr_type(attr_type: u8) -> Option<&'static str> {
 pub struct Attributes {
     // Black box type to allow for later changes/optimizations. The most common attributes could be
     // added as fields to allow for easier lookup.
-    pub(crate) inner: Vec<Attribute>,
+    inner: Vec<Attribute>,
+
+    // Direct fields for well-known mandatory path attributes.
+    // ORIGIN is a well-known mandatory attribute
+    origin: Option<Attribute>,
+    // AS_PATH is a well-known mandatory attribute
+    as_path: Option<Attribute>,
+    // NEXT_HOP is a well-known mandatory attribute
+    next_hop: Option<Attribute>,
+
     /// RFC 7606 validation warnings collected during parsing
     pub(crate) validation_warnings: Vec<BgpValidationWarning>,
 }
 
+struct WellKnownMandatoryAndOtherAttributes {
+    origin: Option<Attribute>,
+    as_path: Option<Attribute>,
+    next_hop: Option<Attribute>,
+    inner: Vec<Attribute>,
+}
+
+impl FromIterator<Attribute> for WellKnownMandatoryAndOtherAttributes {
+    fn from_iter<T: IntoIterator<Item = Attribute>>(iter: T) -> Self {
+        let iter = iter.into_iter();
+
+        let mut origin = None;
+        let mut as_path = None;
+        let mut next_hop = None;
+        let mut inner = Vec::with_capacity(iter.size_hint().0.max(253));
+
+        for attr in iter {
+            match attr.value.attr_type() {
+                AttrType::ORIGIN => origin = Some(attr),
+                AttrType::AS_PATH => as_path = Some(attr),
+                AttrType::NEXT_HOP => next_hop = Some(attr),
+                _ => inner.push(attr),
+            }
+        }
+
+        WellKnownMandatoryAndOtherAttributes {
+            origin,
+            as_path,
+            next_hop,
+            inner,
+        }
+    }
+}
+
+impl From<Vec<Attribute>> for WellKnownMandatoryAndOtherAttributes {
+    fn from(attributes: Vec<Attribute>) -> Self {
+        attributes.into_iter().collect()
+    }
+}
+
 impl Attributes {
     pub fn has_attr(&self, ty: AttrType) -> bool {
-        self.inner.iter().any(|x| x.value.attr_type() == ty)
+        match ty {
+            AttrType::ORIGIN => self.origin.is_some(),
+            AttrType::AS_PATH => self.as_path.is_some(),
+            AttrType::NEXT_HOP => self.next_hop.is_some(),
+            _ => self.inner.iter().any(|x| x.value.attr_type() == ty),
+        }
     }
 
     pub fn get_attr(&self, ty: AttrType) -> Option<Attribute> {
-        self.inner
-            .iter()
-            .find(|x| x.value.attr_type() == ty)
-            .cloned()
+        match ty {
+            AttrType::ORIGIN => self.origin.clone(),
+            AttrType::AS_PATH => self.as_path.clone(),
+            AttrType::NEXT_HOP => self.next_hop.clone(),
+            _ => self
+                .inner
+                .iter()
+                .find(|x| x.value.attr_type() == ty)
+                .cloned(),
+        }
     }
 
     pub fn add_attr(&mut self, attr: Attribute) {
-        self.inner.push(attr);
+        match attr.value.attr_type() {
+            AttrType::ORIGIN => self.origin = Some(attr),
+            AttrType::AS_PATH => self.as_path = Some(attr),
+            AttrType::NEXT_HOP => self.next_hop = Some(attr),
+            _ => self.inner.push(attr),
+        }
     }
 
     /// Add a validation warning to the attributes
@@ -177,9 +241,9 @@ impl Attributes {
     /// Get the `ORIGIN` attribute. In the event that this attribute is not present,
     /// [Origin::INCOMPLETE] will be returned instead.
     pub fn origin(&self) -> Origin {
-        self.inner
-            .iter()
-            .find_map(|x| match &x.value {
+        self.origin
+            .as_ref()
+            .and_then(|x| match &x.value {
                 AttributeValue::Origin(x) => Some(*x),
                 _ => None,
             })
@@ -199,7 +263,7 @@ impl Attributes {
     /// **Note**: Even when this attribute is not present, the next hop address may still be
     /// attainable from the `MP_REACH_NLRI` attribute.
     pub fn next_hop(&self) -> Option<IpAddr> {
-        self.inner.iter().find_map(|x| match &x.value {
+        self.next_hop.as_ref().and_then(|x| match &x.value {
             AttributeValue::NextHop(x) => Some(*x),
             _ => None,
         })
@@ -250,9 +314,7 @@ impl Attributes {
 
     // These implementations are horribly inefficient, but they were super easy to write and use
     pub fn as_path(&self) -> Option<&AsPath> {
-        // Begin searching at the end of the attributes to increase the odds of finding an AS4
-        // attribute first.
-        self.inner.iter().rev().find_map(|x| match &x.value {
+        self.as_path.as_ref().and_then(|x| match &x.value {
             AttributeValue::AsPath { path, .. } => Some(path),
             _ => None,
         })
@@ -279,16 +341,54 @@ impl Attributes {
         }
     }
 
+    /// Count the total number of attributes (including well-known mandatory attributes)
+    pub fn len(&self) -> usize {
+        self.origin.iter().count()
+            + self.as_path.iter().count()
+            + self.next_hop.iter().count()
+            + self.inner.len()
+    }
+
+    /// Check if there are no attributes
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// Get the first attribute if any exists
+    pub fn first(&self) -> Option<&Attribute> {
+        self.origin
+            .as_ref()
+            .or(self.as_path.as_ref())
+            .or(self.next_hop.as_ref())
+            .or(self.inner.first())
+    }
+
     /// Get an iterator over the held [AttributeValue]s. If you also need attribute flags, consider
     /// using [Attributes::into_attributes_iter] instead.
     pub fn iter(&self) -> <&'_ Self as IntoIterator>::IntoIter {
         self.into_iter()
     }
 
+    /// Get an iterator over references to the held [Attribute]s. If you do not need attribute flags,
+    /// consider using [Attributes::iter] instead.
+    pub fn attributes_iter(&self) -> impl Iterator<Item = &Attribute> {
+        chain!(
+            self.as_path.iter(),
+            self.origin.iter(),
+            self.next_hop.iter(),
+            self.inner.iter(),
+        )
+    }
+
     /// Get an iterator over the held [Attribute]s. If you do no not need attribute flags, consider
     /// using [Attributes::iter] instead.
     pub fn into_attributes_iter(self) -> impl Iterator<Item = Attribute> {
-        self.inner.into_iter()
+        chain!(
+            self.as_path.into_iter(),
+            self.origin.into_iter(),
+            self.next_hop.into_iter(),
+            self.inner.into_iter(),
+        )
     }
 }
 
@@ -326,17 +426,27 @@ impl Iterator for MetaCommunitiesIter<'_> {
 
 impl FromIterator<Attribute> for Attributes {
     fn from_iter<T: IntoIterator<Item = Attribute>>(iter: T) -> Self {
+        let attrs = WellKnownMandatoryAndOtherAttributes::from_iter(iter);
+
         Attributes {
-            inner: iter.into_iter().collect(),
+            origin: attrs.origin,
+            as_path: attrs.as_path,
+            next_hop: attrs.next_hop,
+            inner: attrs.inner,
             validation_warnings: Vec::new(),
         }
     }
 }
 
 impl From<Vec<Attribute>> for Attributes {
     fn from(value: Vec<Attribute>) -> Self {
+        let attrs = WellKnownMandatoryAndOtherAttributes::from_iter(value.into_iter());
+
         Attributes {
-            inner: value,
+            origin: attrs.origin,
+            as_path: attrs.as_path,
+            next_hop: attrs.next_hop,
+            inner: attrs.inner,
             validation_warnings: Vec::new(),
         }
     }
@@ -356,34 +466,55 @@ impl Extend<AttributeValue> for Attributes {
 
 impl FromIterator<AttributeValue> for Attributes {
     fn from_iter<T: IntoIterator<Item = AttributeValue>>(iter: T) -> Self {
-        Attributes {
-            inner: iter
-                .into_iter()
-                .map(|value| Attribute {
+        let attrs =
+            WellKnownMandatoryAndOtherAttributes::from_iter(iter.into_iter().map(|value| {
+                Attribute {
                     value,
                     flag: AttrFlags::empty(),
-                })
-                .collect(),
+                }
+            }));
+
+        Attributes {
+            origin: attrs.origin,
+            as_path: attrs.as_path,
+            next_hop: attrs.next_hop,
+            inner: attrs.inner,
             validation_warnings: Vec::new(),
         }
     }
 }
 
 impl IntoIterator for Attributes {
     type Item = AttributeValue;
-    type IntoIter = Map<IntoIter<Attribute>, fn(Attribute) -> AttributeValue>;
+    type IntoIter = std::vec::IntoIter<AttributeValue>;
 
     fn into_iter(self) -> Self::IntoIter {
-        self.inner.into_iter().map(|x| x.value)
+        chain!(
+            self.as_path.into_iter(),
+            self.origin.into_iter(),
+            self.next_hop.into_iter(),
+            self.inner.into_iter()
+        )
+        .map(|x| x.value)
+        .collect::<Vec<_>>()
+        .into_iter()
     }
 }
 
 impl<'a> IntoIterator for &'a Attributes {
     type Item = &'a AttributeValue;
-    type IntoIter = Map<Iter<'a, Attribute>, fn(&Attribute) -> &AttributeValue>;
+    type IntoIter = std::vec::IntoIter<&'a AttributeValue>;
 
     fn into_iter(self) -> Self::IntoIter {
-        self.inner.iter().map(|x| &x.value)
+        chain!(
+            self.as_path.iter(),
+            self.origin.iter(),
+            self.next_hop.iter(),
+            self.inner.iter()
+        )
+        .map(|x| &x.value)
+        .collect::<Vec<_>>()
+        .into_iter()
     }
 }
 
@@ -406,8 +537,14 @@ mod serde_impl {
         where
             D: Deserializer<'de>,
         {
+            let attributes = <Vec<Attribute>>::deserialize(deserializer)?;
+            let well_known_and_others = WellKnownMandatoryAndOtherAttributes::from(attributes);
+
             Ok(Attributes {
-                inner: <Vec<Attribute>>::deserialize(deserializer)?,
+                inner: well_known_and_others.inner,
+                origin: well_known_and_others.origin,
+                as_path: well_known_and_others.as_path,
+                next_hop: well_known_and_others.next_hop,
                 validation_warnings: Vec::new(),
             })
         }

src/parser/bgp/attributes/mod.rs

@@ -459,7 +459,7 @@ impl Attribute {
 impl Attributes {
     pub fn encode(&self, asn_len: AsnLength) -> Bytes {
         let mut bytes = BytesMut::new();
-        for attr in &self.inner {
+        for attr in self.attributes_iter() {
             bytes.extend(attr.encode(asn_len));
         }
         bytes.freeze()
@@ -481,9 +481,9 @@ mod tests {
         let attributes = parse_attributes(data, &asn_len, add_path, afi, safi, prefixes);
         assert!(attributes.is_ok());
         let attributes = attributes.unwrap();
-        assert_eq!(attributes.inner.len(), 1);
+        assert_eq!(attributes.len(), 1);
         assert_eq!(
-            attributes.inner[0].value.attr_type(),
+            attributes.first().unwrap().value.attr_type(),
             AttrType::Unknown(254)
         );
     }

src/parser/bgp/messages.rs

@@ -476,21 +476,21 @@ impl BgpUpdateMessage {
             return false;
         }
 
-        if self.attributes.inner.is_empty() {
+        if self.attributes.is_empty() {
             // no attributes, no prefixes:
             // case 1 end-of-rib
             return true;
         }
 
         // has some attributes, it can only be withdrawal with no prefixes
 
-        if self.attributes.inner.len() > 1 {
+        if self.attributes.len() > 1 {
             // has more than one attributes, not end-of-rib
             return false;
         }
 
         // has only one attribute, check if it is withdrawal attribute
-        if let AttributeValue::MpUnreachNlri(nlri) = &self.attributes.inner.first().unwrap().value {
+        if let AttributeValue::MpUnreachNlri(nlri) = &self.attributes.first().unwrap().value {
             if nlri.prefixes.is_empty() {
                 // the only attribute is MP_UNREACH_NLRI with no prefixes:
                 // case 2 end-of-rib

src/parser/bmp/messages/route_monitoring.rs

@@ -75,9 +75,9 @@ mod tests {
             bgp_message: BgpMessage::Update(msg),
         };
         #[cfg(feature = "parser")]
-        let expected = "RouteMonitoring { bgp_message: Update(BgpUpdateMessage { withdrawn_prefixes: [], attributes: Attributes { inner: [], validation_warnings: [] }, announced_prefixes: [] }) }";
+        let expected = "RouteMonitoring { bgp_message: Update(BgpUpdateMessage { withdrawn_prefixes: [], attributes: Attributes { inner: [], origin: None, as_path: None, next_hop: None, validation_warnings: [] }, announced_prefixes: [] }) }";
         #[cfg(not(feature = "parser"))]
-        let expected = "RouteMonitoring { bgp_message: Update(BgpUpdateMessage { withdrawn_prefixes: [], attributes: Attributes { inner: [] }, announced_prefixes: [] }) }";
+        let expected = "RouteMonitoring { bgp_message: Update(BgpUpdateMessage { withdrawn_prefixes: [], attributes: Attributes { inner: [], origin: None, as_path: None, next_hop: None }, announced_prefixes: [] }) }";
 
         assert_eq!(format!("{mon_msg:?}"), expected);
     }

src/parser/mrt/messages/table_dump.rs

@@ -145,7 +145,7 @@ impl TableDumpMessage {
 
         // encode attributes
         let mut attr_bytes = BytesMut::new();
-        for attr in &self.attributes.inner {
+        for attr in self.attributes.attributes_iter() {
             // asn_len always 16 bites
             attr_bytes.extend(attr.encode(AsnLength::Bits16));
         }