Fixed withdraws encoding

wladwm · May 23, 2021 · 2db5497 · 2db5497
1 parent 791af48
commit 2db5497
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 7 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "zettabgp"
-version = "0.1.1"
+version = "0.1.2"
 authors = ["Vladimir Melnikov <[email protected]>"]
 edition = "2018"
 license = "MIT OR Apache-2.0"

diff --git a/src/message/update/mod.rs b/src/message/update/mod.rs
@@ -102,6 +102,30 @@ impl BgpUpdateMessage {
         }
         None
     }
+    /// returns MPUpdates
+    pub fn get_mpupdates(&self) -> Option<&BgpMPUpdates> {
+        for i in self.attrs.iter() {
+            match i {
+                BgpAttrItem::MPUpdates(n) => {
+                    return Some(&n);
+                }
+                _ => {}
+            }
+        }
+        None
+    }
+    /// returns MPWithdraws
+    pub fn get_mpwithdraws(&self) -> Option<&BgpMPWithdraws> {
+        for i in self.attrs.iter() {
+            match i {
+                BgpAttrItem::MPWithdraws(n) => {
+                    return Some(&n);
+                }
+                _ => {}
+            }
+        }
+        None
+    }
 }
 impl BgpMessage for BgpUpdateMessage {
     fn decode_from(&mut self, peer: &BgpSessionParams, buf: &[u8]) -> Result<(), BgpError> {
@@ -173,19 +197,25 @@ impl BgpMessage for BgpUpdateMessage {
         match peer.peer_mode {
             BgpTransportMode::IPv4 => {
                 if let BgpAddrs::IPV4U(wdrw) = &self.withdraws {
-                    setn_u16(wdrw.len() as u16, &mut buf[curpos..]);
-                    curpos += 2;
-                    curpos += encode_bgpitems_to(&wdrw, &mut buf[curpos..])?;
+                    let wlen = encode_bgpitems_to(&wdrw, &mut buf[curpos + 2..])?;
+                    if wlen > 65535 {
+                        return Err(BgpError::too_many_data());
+                    }
+                    setn_u16(wlen as u16, &mut buf[curpos..]);
+                    curpos += 2 + wlen;
                 } else {
                     setn_u16(0, buf);
                     curpos = 2;
                 }
             }
             BgpTransportMode::IPv6 => {
                 if let BgpAddrs::IPV6U(wdrw) = &self.withdraws {
-                    setn_u16(wdrw.len() as u16, &mut buf[curpos..]);
-                    curpos += 2;
-                    curpos += encode_bgpitems_to(&wdrw, &mut buf[curpos..])?;
+                    let wlen = encode_bgpitems_to(&wdrw, &mut buf[curpos + 2..])?;
+                    if wlen > 65535 {
+                        return Err(BgpError::too_many_data());
+                    }
+                    setn_u16(wlen as u16, &mut buf[curpos..]);
+                    curpos += 2 + wlen;
                 } else {
                     setn_u16(0, buf);
                     curpos = 2;