Store required attributes in attributes instead of vec

ties · ties · commit c55feaf4c230 · 2025-10-25T21:22:43.000+02:00
diff --git a/src/models/bgp/attributes/mod.rs b/src/models/bgp/attributes/mod.rs
@@ -8,10 +8,8 @@ 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::*;
@@ -187,13 +185,19 @@ impl FromIterator<Attribute> for WellKnownMandatoryAndOtherAttributes {
     }
 }
 
+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 {
         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)
+            _ => self.inner.iter().any(|x| x.value.attr_type() == ty),
         }
     }
 
@@ -202,10 +206,11 @@ impl Attributes {
             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()
+            _ => self
+                .inner
+                .iter()
+                .find(|x| x.value.attr_type() == ty)
+                .cloned(),
         }
     }
 
@@ -236,7 +241,8 @@ 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.origin.as_ref()
+        self.origin
+            .as_ref()
             .and_then(|x| match &x.value {
                 AttributeValue::Origin(x) => Some(*x),
                 _ => None,
@@ -257,12 +263,10 @@ 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.next_hop
-            .as_ref()
-            .and_then(|x| match &x.value {
-                AttributeValue::NextHop(x) => Some(*x),
-                _ => None,
-            })
+        self.next_hop.as_ref().and_then(|x| match &x.value {
+            AttributeValue::NextHop(x) => Some(*x),
+            _ => None,
+        })
     }
 
     pub fn multi_exit_discriminator(&self) -> Option<u32> {
@@ -310,12 +314,10 @@ impl Attributes {
 
     // These implementations are horribly inefficient, but they were super easy to write and use
     pub fn as_path(&self) -> Option<&AsPath> {
-        self.as_path
-            .as_ref()
-            .and_then(|x| match &x.value {
-                AttributeValue::AsPath { path, .. } => Some(path),
-                _ => None,
-            })
+        self.as_path.as_ref().and_then(|x| match &x.value {
+            AttributeValue::AsPath { path, .. } => Some(path),
+            _ => None,
+        })
     }
 
     pub fn get_reachable_nlri(&self) -> Option<&Nlri> {
@@ -339,12 +341,45 @@ 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> {
@@ -431,14 +466,13 @@ impl Extend<AttributeValue> for Attributes {
 
 impl FromIterator<AttributeValue> for Attributes {
     fn from_iter<T: IntoIterator<Item = AttributeValue>>(iter: T) -> Self {
-        let attrs = WellKnownMandatoryAndOtherAttributes::from_iter(
-            iter
-                .into_iter()
-                .map(|value| Attribute {
+        let attrs =
+            WellKnownMandatoryAndOtherAttributes::from_iter(iter.into_iter().map(|value| {
+                Attribute {
                     value,
                     flag: AttrFlags::empty(),
-                })
-        );
+                }
+            }));
 
         Attributes {
             origin: attrs.origin,
@@ -452,26 +486,35 @@ impl FromIterator<AttributeValue> for Attributes {
 
 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 {
-        let tmp: Vec<_> = chain!(
+        chain!(
             self.as_path.iter(),
             self.origin.iter(),
             self.next_hop.iter(),
             self.inner.iter()
-        ).collect();
-        
-        tmp.as_slice().into_iter().map(|x| &x.value)
+        )
+        .map(|x| &x.value)
+        .collect::<Vec<_>>()
+        .into_iter()
     }
 }
 
@@ -494,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(),
             })
         }
diff --git a/src/parser/bgp/attributes/mod.rs b/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.iter() {
+        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)
         );
     }
diff --git a/src/parser/bgp/messages.rs b/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
diff --git a/src/parser/bmp/messages/route_monitoring.rs b/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);
     }
diff --git a/src/parser/mrt/messages/table_dump.rs b/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));
         }

Original file line number	Diff line number	Diff line change
`@@ -145,7 +145,7 @@ impl TableDumpMessage {`
`145`	`145`
`146`	`146`	`// encode attributes`
`147`	`147`	`let mut attr_bytes = BytesMut::new();`
`148`		`- for attr in &self.attributes.inner {`
	`148`	`+ for attr in self.attributes.attributes_iter() {`
`149`	`149`	`// asn_len always 16 bites`
`150`	`150`	`attr_bytes.extend(attr.encode(AsnLength::Bits16));`
`151`	`151`	`}`