bgpkit_parser/encoder/
rib_encoder.rs

1//! MRT Encoder module
2//!
3//! `mrt_encoder` module handles serializing BGP/MRT messages back to MRT binary files. The main
4//! difficulty part of this process is the handling of TableDumpV2 RIB dumps, which requires
5//! reconstructing the peer index table before encoding all other contents.
6
7use crate::error::EncodingError;
8use crate::models::{
9    Attributes, BgpElem, CommonHeader, EntryType, MrtMessage, NetworkPrefix, Peer, PeerIndexTable,
10    RibAfiEntries, RibEntry, TableDumpV2Message, TableDumpV2Type,
11};
12use crate::utils::convert_timestamp;
13use bytes::{Bytes, BytesMut};
14use ipnet::IpNet;
15use std::collections::HashMap;
16use std::net::{IpAddr, Ipv4Addr};
17
18fn rib_type_for_prefix(prefix: &IpNet, has_add_path: bool) -> TableDumpV2Type {
19    match (prefix.addr().is_ipv6(), has_add_path) {
20        (true, true) => TableDumpV2Type::RibIpv6UnicastAddPath,
21        (true, false) => TableDumpV2Type::RibIpv6Unicast,
22        (false, true) => TableDumpV2Type::RibIpv4UnicastAddPath,
23        (false, false) => TableDumpV2Type::RibIpv4Unicast,
24    }
25}
26
27#[derive(Default)]
28pub struct MrtRibEncoder {
29    index_table: PeerIndexTable,
30
31    per_prefix_entries_map: HashMap<IpNet, HashMap<u16, RibEntry>>,
32
33    timestamp: f64,
34}
35
36impl MrtRibEncoder {
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    pub fn reset(&mut self) {
42        self.index_table = PeerIndexTable::default();
43        self.per_prefix_entries_map = HashMap::default();
44        self.timestamp = 0.0;
45    }
46
47    /// Processes a BgpElem and updates the internal data structures.
48    ///
49    /// Returns [`EncodingError::ValueTooLarge`] if the element's peer would
50    /// exceed the 65535-peer capacity of the PEER_INDEX_TABLE; the encoder's
51    /// state is left unchanged in that case.
52    ///
53    /// # Arguments
54    ///
55    /// * `elem` - A reference to a BgpElem that contains the information to be processed.
56    pub fn process_elem(&mut self, elem: &BgpElem) -> Result<(), EncodingError> {
57        if self.timestamp == 0.0 {
58            self.timestamp = elem.timestamp;
59        }
60        let bgp_identifier = match elem.peer_ip {
61            IpAddr::V4(ip) => ip,
62            IpAddr::V6(_ip) => Ipv4Addr::from(0),
63        };
64        let peer = Peer::new(bgp_identifier, elem.peer_ip, elem.peer_asn);
65        let peer_index = self.index_table.add_peer(peer)?;
66        let path_id = elem.prefix.path_id;
67        let prefix = elem.prefix.prefix;
68
69        let entries_map = self.per_prefix_entries_map.entry(prefix).or_default();
70        let entry = RibEntry {
71            peer_index,
72            path_id,
73            originated_time: elem.timestamp as u32,
74            attributes: Attributes::from(elem),
75        };
76        entries_map.insert(peer_index, entry);
77        Ok(())
78    }
79
80    /// Export the data stored in the struct to a byte array.
81    ///
82    /// The function first encodes the peer-index-table data into a `MrtMessage` and appends it to the `BytesMut` object.
83    /// Then, for each prefix in the `per_prefix_entries_map`, it creates a `RibAfiEntries` object and encodes it as a `MrtMessage`.
84    /// The resulting `BytesMut` object is then converted to an immutable `Bytes` object using `freeze()` and returned.
85    ///
86    /// # Return
87    /// Returns a `Bytes` object containing the exported data as a byte array,
88    /// or an [EncodingError] if any element cannot be represented in wire format.
89    pub fn export_bytes(&mut self) -> Result<Bytes, EncodingError> {
90        let mut bytes = BytesMut::new();
91
92        // encode peer-index-table
93        let mrt_message = MrtMessage::TableDumpV2Message(TableDumpV2Message::PeerIndexTable(
94            self.index_table.clone(),
95        ));
96        let (seconds, _microseconds) = convert_timestamp(self.timestamp);
97        let subtype = TableDumpV2Type::PeerIndexTable as u16;
98        let data_bytes = mrt_message.encode(subtype)?;
99        let header = CommonHeader {
100            timestamp: seconds,
101            microsecond_timestamp: None,
102            entry_type: EntryType::TABLE_DUMP_V2,
103            entry_subtype: subtype,
104            length: data_bytes.len() as u32,
105        };
106        let header_bytes = header.encode();
107        bytes.extend(header_bytes);
108        bytes.extend(data_bytes);
109
110        // encode each RibAfiEntries
111        for (entry_count, (prefix, entries_map)) in self.per_prefix_entries_map.iter().enumerate() {
112            let has_add_path = entries_map.values().any(|entry| entry.path_id.is_some());
113            let rib_type = rib_type_for_prefix(prefix, has_add_path);
114
115            let mut prefix_rib_entry = RibAfiEntries {
116                rib_type,
117                sequence_number: entry_count as u32,
118                prefix: NetworkPrefix::new(*prefix, None),
119                rib_entries: vec![],
120            };
121            for entry in entries_map.values() {
122                prefix_rib_entry.rib_entries.push(entry.clone());
123            }
124
125            let mrt_message =
126                MrtMessage::TableDumpV2Message(TableDumpV2Message::RibAfi(prefix_rib_entry));
127
128            let (seconds, _microseconds) = convert_timestamp(self.timestamp);
129            let subtype = rib_type as u16;
130            let data_bytes = mrt_message.encode(subtype)?;
131            let header_bytes = CommonHeader {
132                timestamp: seconds,
133                microsecond_timestamp: None,
134                entry_type: EntryType::TABLE_DUMP_V2,
135                entry_subtype: subtype,
136                length: data_bytes.len() as u32,
137            }
138            .encode();
139            bytes.extend(header_bytes);
140            bytes.extend(data_bytes);
141        }
142
143        self.reset();
144
145        Ok(bytes.freeze())
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use crate::models::Asn;
153    use crate::parse_mrt_record;
154    use bytes::Buf;
155    use std::io::Cursor;
156
157    #[test]
158    fn test_encoding_rib() {
159        let mut encoder = MrtRibEncoder::new();
160        let mut elem = BgpElem {
161            peer_ip: IpAddr::V4("10.0.0.1".parse().unwrap()),
162            peer_asn: Asn::from(65000),
163            ..Default::default()
164        };
165        elem.prefix.prefix = "10.250.0.0/24".parse().unwrap();
166        encoder.process_elem(&elem).unwrap();
167        elem.prefix.prefix = "10.251.0.0/24".parse().unwrap();
168        encoder.process_elem(&elem).unwrap();
169        let bytes = encoder.export_bytes().unwrap();
170
171        let mut cursor = Cursor::new(bytes.clone());
172        while cursor.has_remaining() {
173            let _parsed = parse_mrt_record(&mut cursor).unwrap();
174        }
175
176        // v6
177        let mut encoder = MrtRibEncoder::new();
178        let mut elem = BgpElem {
179            peer_ip: IpAddr::V6("::1".parse().unwrap()),
180            peer_asn: Asn::from(65000),
181            ..Default::default()
182        };
183        // ipv6 prefix
184        elem.prefix.prefix = "2001:db8::/32".parse().unwrap();
185        encoder.process_elem(&elem).unwrap();
186        let bytes = encoder.export_bytes().unwrap();
187
188        let mut cursor = Cursor::new(bytes.clone());
189        while cursor.has_remaining() {
190            let _parsed = parse_mrt_record(&mut cursor).unwrap();
191        }
192    }
193
194    #[test]
195    fn test_encoding_rib_with_add_path() {
196        let mut encoder = MrtRibEncoder::new();
197        let mut elem = BgpElem {
198            peer_ip: IpAddr::V4("10.0.0.1".parse().unwrap()),
199            peer_asn: Asn::from(65000),
200            ..Default::default()
201        };
202        elem.prefix = NetworkPrefix::new("10.250.0.0/24".parse().unwrap(), Some(42));
203        encoder.process_elem(&elem).unwrap();
204
205        let bytes = encoder.export_bytes().unwrap();
206        let mut cursor = Cursor::new(bytes);
207        let _peer_table = parse_mrt_record(&mut cursor).unwrap();
208        let parsed = parse_mrt_record(&mut cursor).unwrap();
209
210        match parsed.message {
211            MrtMessage::TableDumpV2Message(TableDumpV2Message::RibAfi(rib)) => {
212                assert_eq!(rib.rib_type, TableDumpV2Type::RibIpv4UnicastAddPath);
213                assert_eq!(rib.rib_entries.len(), 1);
214                assert_eq!(rib.rib_entries[0].path_id, Some(42));
215            }
216            other => panic!("unexpected MRT message: {other:?}"),
217        }
218    }
219}
bgpkit_parser/encoder/rib_encoder.rs

bgpkit_parser/encoder/
rib_encoder.rs
