1use bytes::{Bytes, BytesMut};
2
3use crate::constants::HEADER_LEN;
4use crate::error::{DecodeError, EncodeError};
5use crate::header::{BgpHeader, MessageType};
6use crate::keepalive;
7use crate::notification_msg::NotificationMessage;
8use crate::open::OpenMessage;
9use crate::route_refresh::RouteRefreshMessage;
10use crate::update::UpdateMessage;
11
12#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum Message {
15 Open(OpenMessage),
17 Update(UpdateMessage),
19 Notification(NotificationMessage),
21 Keepalive,
23 RouteRefresh(RouteRefreshMessage),
25}
26
27impl Message {
28 #[must_use]
30 pub fn message_type(&self) -> MessageType {
31 match self {
32 Self::Open(_) => MessageType::Open,
33 Self::Update(_) => MessageType::Update,
34 Self::Notification(_) => MessageType::Notification,
35 Self::Keepalive => MessageType::Keepalive,
36 Self::RouteRefresh(_) => MessageType::RouteRefresh,
37 }
38 }
39}
40
41impl std::fmt::Display for Message {
42 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43 match self {
44 Self::Open(o) => write!(f, "OPEN AS={} hold={}", o.my_as, o.hold_time),
45 Self::Update(_) => write!(f, "UPDATE"),
46 Self::Notification(n) => write!(f, "NOTIFICATION {}/{}", n.code, n.subcode),
47 Self::Keepalive => write!(f, "KEEPALIVE"),
48 Self::RouteRefresh(rr) => write!(f, "{rr}"),
49 }
50 }
51}
52
53pub fn decode_message(buf: &mut Bytes, max_message_len: u16) -> Result<Message, DecodeError> {
70 let header = BgpHeader::decode(buf, max_message_len)?;
71 let body_len = usize::from(header.length) - HEADER_LEN;
72
73 match header.message_type {
74 MessageType::Keepalive => {
75 keepalive::validate_keepalive(&header)?;
76 Ok(Message::Keepalive)
77 }
78 MessageType::Notification => {
79 let msg = NotificationMessage::decode(buf, body_len)?;
80 Ok(Message::Notification(msg))
81 }
82 MessageType::Open => {
83 let msg = OpenMessage::decode(buf, body_len)?;
84 Ok(Message::Open(msg))
85 }
86 MessageType::Update => {
87 let msg = UpdateMessage::decode(buf, body_len)?;
88 Ok(Message::Update(msg))
89 }
90 MessageType::RouteRefresh => {
91 let msg = RouteRefreshMessage::decode(buf, body_len)?;
92 Ok(Message::RouteRefresh(msg))
93 }
94 }
95}
96
97pub fn encode_message(msg: &Message) -> Result<BytesMut, EncodeError> {
106 let mut buf = BytesMut::with_capacity(match msg {
107 Message::Keepalive => keepalive::KEEPALIVE_LEN,
108 Message::Notification(n) => n.encoded_len(),
109 Message::Open(o) => o.encoded_len(),
110 Message::Update(u) => u.encoded_len(),
111 Message::RouteRefresh(rr) => rr.encoded_len(),
112 });
113
114 match msg {
115 Message::Keepalive => {
116 keepalive::encode_keepalive(&mut buf);
117 }
118 Message::Notification(n) => {
119 n.encode(&mut buf)?;
120 }
121 Message::Open(o) => {
122 o.encode(&mut buf)?;
123 }
124 Message::Update(u) => {
125 u.encode(&mut buf)?;
126 }
127 Message::RouteRefresh(rr) => {
128 rr.encode(&mut buf)?;
129 }
130 }
131
132 Ok(buf)
133}
134
135pub fn encode_message_with_limit(
145 msg: &Message,
146 max_message_len: u16,
147) -> Result<BytesMut, EncodeError> {
148 let mut buf = BytesMut::with_capacity(match msg {
149 Message::Keepalive => keepalive::KEEPALIVE_LEN,
150 Message::Notification(n) => n.encoded_len(),
151 Message::Open(o) => o.encoded_len(),
152 Message::Update(u) => u.encoded_len(),
153 Message::RouteRefresh(rr) => rr.encoded_len(),
154 });
155
156 match msg {
157 Message::Keepalive => {
158 keepalive::encode_keepalive(&mut buf);
159 }
160 Message::Notification(n) => {
161 n.encode(&mut buf)?;
162 }
163 Message::Open(o) => {
164 o.encode(&mut buf)?;
165 }
166 Message::Update(u) => {
167 u.encode_with_limit(&mut buf, max_message_len)?;
168 }
169 Message::RouteRefresh(rr) => {
170 rr.encode(&mut buf)?;
171 }
172 }
173
174 Ok(buf)
175}
176
177#[cfg(test)]
178mod tests {
179 use std::net::Ipv4Addr;
180
181 use super::*;
182 use crate::capability::{Afi, Capability, Safi};
183 use crate::constants::{BGP_VERSION, MAX_MESSAGE_LEN};
184 use crate::notification::NotificationCode;
185
186 #[test]
187 fn roundtrip_keepalive() {
188 let encoded = encode_message(&Message::Keepalive).unwrap();
189 let mut bytes = encoded.freeze();
190 let decoded = decode_message(&mut bytes, MAX_MESSAGE_LEN).unwrap();
191 assert_eq!(decoded, Message::Keepalive);
192 }
193
194 #[test]
195 fn roundtrip_notification() {
196 let msg = Message::Notification(NotificationMessage::new(
197 NotificationCode::Cease,
198 2,
199 Bytes::from_static(&[0x01, 0x02]),
200 ));
201 let encoded = encode_message(&msg).unwrap();
202 let mut bytes = encoded.freeze();
203 let decoded = decode_message(&mut bytes, MAX_MESSAGE_LEN).unwrap();
204 assert_eq!(decoded, msg);
205 }
206
207 #[test]
208 fn roundtrip_open_minimal() {
209 let msg = Message::Open(OpenMessage {
210 version: BGP_VERSION,
211 my_as: 65001,
212 hold_time: 90,
213 bgp_identifier: Ipv4Addr::new(10, 0, 0, 1),
214 capabilities: vec![],
215 });
216 let encoded = encode_message(&msg).unwrap();
217 let mut bytes = encoded.freeze();
218 let decoded = decode_message(&mut bytes, MAX_MESSAGE_LEN).unwrap();
219 assert_eq!(decoded, msg);
220 }
221
222 #[test]
223 fn roundtrip_open_with_caps() {
224 let msg = Message::Open(OpenMessage {
225 version: BGP_VERSION,
226 my_as: 23456,
227 hold_time: 180,
228 bgp_identifier: Ipv4Addr::new(192, 168, 1, 1),
229 capabilities: vec![
230 Capability::MultiProtocol {
231 afi: Afi::Ipv4,
232 safi: Safi::Unicast,
233 },
234 Capability::FourOctetAs { asn: 4_200_000_001 },
235 Capability::Unknown {
236 code: 128,
237 data: Bytes::from_static(&[0xAA]),
238 },
239 ],
240 });
241 let encoded = encode_message(&msg).unwrap();
242 let mut bytes = encoded.freeze();
243 let decoded = decode_message(&mut bytes, MAX_MESSAGE_LEN).unwrap();
244 assert_eq!(decoded, msg);
245 }
246
247 #[test]
248 fn roundtrip_update_minimal() {
249 let msg = Message::Update(UpdateMessage {
250 withdrawn_routes: Bytes::new(),
251 path_attributes: Bytes::new(),
252 nlri: Bytes::new(),
253 });
254 let encoded = encode_message(&msg).unwrap();
255 let mut bytes = encoded.freeze();
256 let decoded = decode_message(&mut bytes, MAX_MESSAGE_LEN).unwrap();
257 assert_eq!(decoded, msg);
258 }
259
260 #[test]
261 fn roundtrip_update_with_data() {
262 let msg = Message::Update(UpdateMessage {
263 withdrawn_routes: Bytes::from_static(&[0x18, 0x0A, 0x00]),
264 path_attributes: Bytes::from_static(&[0x40, 0x01, 0x00]),
265 nlri: Bytes::from_static(&[0x18, 0xC0, 0xA8]),
266 });
267 let encoded = encode_message(&msg).unwrap();
268 let mut bytes = encoded.freeze();
269 let decoded = decode_message(&mut bytes, MAX_MESSAGE_LEN).unwrap();
270 assert_eq!(decoded, msg);
271 }
272
273 #[test]
274 fn message_type_accessor() {
275 assert_eq!(Message::Keepalive.message_type(), MessageType::Keepalive);
276 assert_eq!(
277 Message::Notification(NotificationMessage::new(
278 NotificationCode::Cease,
279 0,
280 Bytes::new()
281 ))
282 .message_type(),
283 MessageType::Notification
284 );
285 }
286
287 #[test]
288 fn decode_rejects_garbage() {
289 let mut buf = Bytes::from_static(&[0x00; 19]);
290 assert!(decode_message(&mut buf, MAX_MESSAGE_LEN).is_err());
291 }
292
293 #[test]
294 fn decode_rejects_truncated() {
295 let mut buf = Bytes::from_static(&[0xFF; 10]);
296 assert!(matches!(
297 decode_message(&mut buf, MAX_MESSAGE_LEN),
298 Err(DecodeError::Incomplete { .. })
299 ));
300 }
301}