From 167ee20542d945a2c1c8cb0f973464d3236fe541 Mon Sep 17 00:00:00 2001
From: Yureka <yuka@yuka.dev>
Date: Wed, 13 Mar 2024 10:28:42 +0100
Subject: [PATCH] add support for RD and local peer types

---
 src/bmp_collector.rs | 44 +++++++++++++++++++++++++++++++++-----------
 src/store.rs         |  8 ++++++++
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/src/bmp_collector.rs b/src/bmp_collector.rs
index 7ac1fa8..0ba6dcc 100644
--- a/src/bmp_collector.rs
+++ b/src/bmp_collector.rs
@@ -1,4 +1,6 @@
-use crate::store::{Client, RouteState, Session, SessionId, Store, TableSelector};
+use crate::store::{
+    Client, PeerDistinguisher, RouteState, Session, SessionId, Store, TableSelector,
+};
 use bitvec::prelude::Msb0;
 use bitvec::view::BitView;
 use futures_util::future::join_all;
@@ -19,16 +21,36 @@ fn table_selector_for_peer(
     client_addr: SocketAddr,
     peer: &BmpMessagePeerHeader,
 ) -> Option<TableSelector> {
-    match (peer.peertype, peer.flags.view_bits::<Msb0>()[1]) {
-        (0, false) => Some(TableSelector::PrePolicyAdjIn(SessionId {
-            from_client: client_addr,
-            peer_address: peer.peeraddress,
-        })),
-        (0, true) => Some(TableSelector::PostPolicyAdjIn(SessionId {
-            from_client: client_addr,
-            peer_address: peer.peeraddress,
-        })),
-        (3, _) => Some(TableSelector::LocRib {
+    let peer_distinguisher = match peer.peertype {
+        0 => {
+            if peer.peerdistinguisher.is_zero() {
+                Some(PeerDistinguisher::Global)
+            } else {
+                warn!("Peer type global but peer distinguisher is not empty");
+                None
+            }
+        }
+        1 => Some(PeerDistinguisher::RD(
+            peer.peerdistinguisher.rdh,
+            peer.peerdistinguisher.rdl,
+        )),
+        2 => Some(PeerDistinguisher::Local(
+            peer.peerdistinguisher.rdh,
+            peer.peerdistinguisher.rdl,
+        )),
+        _ => None,
+    };
+
+    let session_id = peer_distinguisher.map(|peer_distinguisher| SessionId {
+        from_client: client_addr,
+        peer_distinguisher,
+        peer_address: peer.peeraddress,
+    });
+
+    match (peer.peertype, session_id, peer.flags.view_bits::<Msb0>()[1]) {
+        (0 | 1 | 2, Some(session), false) => Some(TableSelector::PrePolicyAdjIn(session)),
+        (0 | 1 | 2, Some(session), true) => Some(TableSelector::PostPolicyAdjIn(session)),
+        (3, _, _) => Some(TableSelector::LocRib {
             from_client: client_addr,
             route_state: RouteState::Selected,
         }),
diff --git a/src/store.rs b/src/store.rs
index 4025a1c..392e7c8 100644
--- a/src/store.rs
+++ b/src/store.rs
@@ -28,10 +28,18 @@ pub struct RouteAttrs {
     pub nexthop: Option<IpAddr>,
 }
 
+#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
+pub enum PeerDistinguisher {
+    Global,
+    RD(u32, u32),
+    Local(u32, u32),
+}
+
 #[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct SessionId {
     pub from_client: SocketAddr,
+    pub peer_distinguisher: PeerDistinguisher,
     pub peer_address: IpAddr,
 }