Store required attributes in attributes instead of vec

Author: ties

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::*;
@@ -339,12 +337,41 @@ 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> {
@@ -452,26 +479,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()
     }
 }
 

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)
         );
     }

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));
         }