bgpkit_parser/parser/mrt/messages/table_dump_v2/
peer_index_table.rs

1use crate::encoder::sink::put_u16_len_slice;
2use crate::error::{check_max, EncodingError};
3use crate::models::{Afi, AsnLength, Peer, PeerIndexTable, PeerType};
4use crate::parser::ReadUtils;
5use crate::ParserError;
6use bytes::{BufMut, Bytes, BytesMut};
7use std::collections::HashMap;
8use std::net::{IpAddr, Ipv4Addr};
9
10/// Parses a byte slice into a [PeerIndexTable].
11///
12/// RFC: https://www.rfc-editor.org/rfc/rfc6396#section-4.3.1
13///
14/// # Arguments
15///
16/// * `data` - The byte slice to parse.
17///
18/// # Returns
19///
20/// - `Ok(PeerIndexTable)` if the parsing is successful.
21/// - `Err(ParserError)` if an error occurs during parsing.
22pub fn parse_peer_index_table(data: &mut Bytes) -> Result<PeerIndexTable, ParserError> {
23    let collector_bgp_id = Ipv4Addr::from(data.read_u32()?);
24    // read and ignore view name
25    let view_name_length = data.read_u16()?;
26    let view_name =
27        String::from_utf8(data.read_n_bytes(view_name_length as usize)?).unwrap_or("".to_string());
28
29    let peer_count = data.read_u16()?;
30    let mut peers = vec![];
31    for _index in 0..peer_count {
32        let peer_type = PeerType::from_bits_retain(data.read_u8()?);
33        let afi = match peer_type.contains(PeerType::ADDRESS_FAMILY_IPV6) {
34            true => Afi::Ipv6,
35            false => Afi::Ipv4,
36        };
37        let asn_len = match peer_type.contains(PeerType::AS_SIZE_32BIT) {
38            true => AsnLength::Bits32,
39            false => AsnLength::Bits16,
40        };
41
42        let peer_bgp_id = Ipv4Addr::from(data.read_u32()?);
43        let peer_ip: IpAddr = data.read_address(&afi)?;
44        let peer_asn = data.read_asn(asn_len)?;
45        peers.push(Peer {
46            peer_type,
47            peer_bgp_id,
48            peer_ip,
49            peer_asn,
50        })
51    }
52
53    let mut id_peer_map = HashMap::new();
54    let mut peer_ip_id_map = HashMap::new();
55
56    for (id, p) in peers.into_iter().enumerate() {
57        id_peer_map.insert(id as u16, p);
58        peer_ip_id_map.insert(p.peer_ip, id as u16);
59    }
60
61    Ok(PeerIndexTable {
62        collector_bgp_id,
63        view_name,
64        id_peer_map,
65        peer_ip_id_map,
66    })
67}
68
69impl PeerIndexTable {
70    /// Add peer to peer index table and return peer id.
71    ///
72    /// The PEER_INDEX_TABLE wire format uses a 16-bit peer count, so at most
73    /// 65535 distinct peers can be stored. Adding a peer beyond that returns
74    /// [`EncodingError::ValueTooLarge`] and leaves the table unmodified —
75    /// previously the id silently wrapped, aliasing routes to the wrong peer.
76    pub fn add_peer(&mut self, peer: Peer) -> Result<u16, EncodingError> {
77        match self.peer_ip_id_map.get(&peer.peer_ip) {
78            Some(id) => Ok(*id),
79            None => {
80                let next_id = self.peer_ip_id_map.len();
81                check_max("PeerIndexTable peer count", next_id + 1, u16::MAX as usize)?;
82                let peer_id = next_id as u16;
83                self.peer_ip_id_map.insert(peer.peer_ip, peer_id);
84                self.id_peer_map.insert(peer_id, peer);
85                Ok(peer_id)
86            }
87        }
88    }
89
90    /// Returns the peer associated with the given peer ID.
91    ///
92    /// # Arguments
93    ///
94    /// * `peer_id` - A reference to the peer ID.
95    ///
96    /// # Returns
97    ///
98    /// An `Option` containing a reference to the [Peer] if found, otherwise `None`.
99    pub fn get_peer_by_id(&self, peer_id: &u16) -> Option<&Peer> {
100        self.id_peer_map.get(peer_id)
101    }
102
103    /// Returns the peer ID associated with the given IP address.
104    ///
105    /// # Arguments
106    ///
107    /// * `peer_ip` - The IP address of the peer.
108    ///
109    /// # Returns
110    ///
111    /// An optional `u16` representing the peer ID. Returns `None` if the IP address is not found.
112    ///
113    /// # Examples
114    ///
115    /// ```
116    /// use std::net::IpAddr;
117    /// use std::str::FromStr;
118    /// use bgpkit_parser::models::PeerIndexTable;
119    ///
120    /// let index_table = PeerIndexTable::default();
121    /// let peer_ip = IpAddr::from_str("127.0.0.1").unwrap();
122    /// let peer_id = index_table.get_peer_id_by_addr(&peer_ip);
123    /// ```
124    pub fn get_peer_id_by_addr(&self, peer_ip: &IpAddr) -> Option<u16> {
125        self.peer_ip_id_map.get(peer_ip).copied()
126    }
127
128    /// Encode the data in the struct into a byte array.
129    ///
130    /// # Returns
131    ///
132    /// A `Bytes` object containing the encoded data.
133    ///
134    /// # Example
135    ///
136    /// ```
137    /// use std::collections::HashMap;
138    /// use std::net::Ipv4Addr;
139    /// use bgpkit_parser::models::PeerIndexTable;
140    ///
141    /// let data = PeerIndexTable {
142    ///     collector_bgp_id: Ipv4Addr::from(1234),
143    ///     view_name: String::from("example"),
144    ///     id_peer_map: HashMap::new(),
145    ///     peer_ip_id_map: Default::default(),
146    /// };
147    ///
148    /// let encoded = data.encode().unwrap();
149    /// ```
150    pub fn encode(&self) -> Result<Bytes, EncodingError> {
151        let mut buf = BytesMut::new();
152
153        // Encode collector_bgp_id
154        buf.put_u32(self.collector_bgp_id.into());
155
156        // Encode view_name_length and view_name
157        put_u16_len_slice(
158            &mut buf,
159            "PeerIndexTable view name length",
160            self.view_name.as_bytes(),
161        )?;
162
163        // Encode peer_count
164        let peer_count = self.id_peer_map.len();
165        check_max("PeerIndexTable peer count", peer_count, u16::MAX as usize)?;
166        buf.put_u16(peer_count as u16);
167
168        // Encode peers
169        let mut peer_ids: Vec<_> = self.id_peer_map.keys().collect();
170        peer_ids.sort();
171        for id in peer_ids {
172            let peer = self.id_peer_map.get(id).unwrap();
173            // Encode PeerType
174            buf.put_u8(peer.peer_type.bits());
175
176            // Encode peer_bgp_id
177            buf.put_u32(peer.peer_bgp_id.into());
178
179            // Encode peer_ip
180            match peer.peer_ip {
181                IpAddr::V4(ipv4) => {
182                    buf.put_slice(&ipv4.octets());
183                }
184                IpAddr::V6(ipv6) => {
185                    buf.put_slice(&ipv6.octets());
186                }
187            };
188
189            // Encode peer_asn
190            match peer.peer_type.contains(PeerType::AS_SIZE_32BIT) {
191                true => buf.put_u32(peer.peer_asn.to_u32()),
192                false => buf.put_u16(peer.peer_asn.to_u32() as u16),
193            };
194        }
195
196        // Return Bytes
197        Ok(buf.freeze())
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use crate::models::Asn;
205    use std::str::FromStr;
206
207    #[test]
208    fn test_peer_index_table_encode() {
209        let mut index_table = PeerIndexTable {
210            collector_bgp_id: Ipv4Addr::from(1234),
211            view_name: String::from("example"),
212            id_peer_map: HashMap::new(),
213            peer_ip_id_map: Default::default(),
214        };
215
216        index_table
217            .add_peer(Peer::new(
218                Ipv4Addr::from(1234),
219                IpAddr::from_str("192.168.1.1").unwrap(),
220                Asn::new_32bit(1234),
221            ))
222            .unwrap();
223        index_table
224            .add_peer(Peer::new(
225                Ipv4Addr::from(12345),
226                IpAddr::from_str("192.168.1.2").unwrap(),
227                Asn::new_32bit(12345),
228            ))
229            .unwrap();
230
231        let encoded = index_table.encode().unwrap();
232        let parsed_index_table = parse_peer_index_table(&mut encoded.clone()).unwrap();
233        assert_eq!(index_table, parsed_index_table);
234    }
235
236    #[test]
237    fn test_get_peer_by_id() {
238        let mut index_table = PeerIndexTable {
239            collector_bgp_id: Ipv4Addr::from(1234),
240            view_name: String::from("example"),
241            id_peer_map: HashMap::new(),
242            peer_ip_id_map: Default::default(),
243        };
244
245        let peer1 = Peer::new(
246            Ipv4Addr::from(1234),
247            IpAddr::from_str("10.0.0.1").unwrap(),
248            Asn::new_32bit(1234),
249        );
250        let peer2 = Peer::new(
251            Ipv4Addr::from(12345),
252            IpAddr::from_str("10.0.0.2").unwrap(),
253            Asn::new_32bit(12345),
254        );
255
256        let peer1_id = index_table.add_peer(peer1).unwrap();
257        let peer2_id = index_table.add_peer(peer2).unwrap();
258
259        assert_eq!(
260            index_table.get_peer_by_id(&peer1_id),
261            Some(&Peer::new(
262                Ipv4Addr::from(1234),
263                IpAddr::from_str("10.0.0.1").unwrap(),
264                Asn::new_32bit(1234),
265            ))
266        );
267        assert_eq!(
268            index_table.get_peer_by_id(&peer2_id),
269            Some(&Peer::new(
270                Ipv4Addr::from(12345),
271                IpAddr::from_str("10.0.0.2").unwrap(),
272                Asn::new_32bit(12345),
273            ))
274        );
275    }
276
277    #[test]
278    fn test_add_peer_rejects_overflow_without_corruption() {
279        let mut index_table = PeerIndexTable::default();
280
281        // fill the table to its 16-bit wire capacity of 65535 peers
282        for i in 0..(u16::MAX as u32) {
283            let ip = IpAddr::from(Ipv4Addr::from(i + 1));
284            index_table
285                .add_peer(Peer::new(Ipv4Addr::from(1), ip, Asn::new_32bit(i)))
286                .unwrap();
287        }
288        assert_eq!(index_table.id_peer_map.len(), u16::MAX as usize);
289
290        // the 65536th peer must be rejected, not aliased onto an existing id
291        let overflow_ip = IpAddr::from(Ipv4Addr::from(u16::MAX as u32 + 1));
292        let overflow_peer = Peer::new(Ipv4Addr::from(1), overflow_ip, Asn::new_32bit(65536));
293        let err = index_table.add_peer(overflow_peer).unwrap_err();
294        assert_eq!(
295            err,
296            EncodingError::ValueTooLarge {
297                field: "PeerIndexTable peer count",
298                actual: u16::MAX as usize + 1,
299                max: u16::MAX as usize
300            }
301        );
302
303        // the failed insert must not have modified the table
304        assert_eq!(index_table.id_peer_map.len(), u16::MAX as usize);
305        assert_eq!(index_table.get_peer_id_by_addr(&overflow_ip), None);
306
307        // adding an existing peer still returns its id without error
308        let existing_ip = IpAddr::from(Ipv4Addr::from(1u32));
309        let existing = Peer::new(Ipv4Addr::from(1), existing_ip, Asn::new_32bit(0));
310        assert_eq!(index_table.add_peer(existing).unwrap(), 0);
311
312        // the full table still encodes successfully
313        index_table.encode().unwrap();
314    }
315}
bgpkit_parser/parser/mrt/messages/table_dump_v2/peer_index_table.rs

bgpkit_parser/parser/mrt/messages/table_dump_v2/
peer_index_table.rs
