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

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

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