From d4774fd270c5c60d8402e62e40927377d3aec603 Mon Sep 17 00:00:00 2001
From: yuanchao <yuanchao@yunshan.net>
Date: Fri, 18 Oct 2024 15:12:45 +0800
Subject: [PATCH] feat: modify the epc of multicast address

---
 agent/src/policy/labeler.rs | 36 ++++++++++++++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/agent/src/policy/labeler.rs b/agent/src/policy/labeler.rs
index c1939f4d67a..9988190b432 100644
--- a/agent/src/policy/labeler.rs
+++ b/agent/src/policy/labeler.rs
@@ -655,10 +655,18 @@ impl Labeler {
         }
     }
 
+    fn is_multicast(ip: &IpAddr) -> bool {
+        match ip {
+            IpAddr::V4(a) => a.is_broadcast() || a.is_multicast(),
+            IpAddr::V6(a) => a.is_multicast(),
+        }
+    }
+
     fn is_intranet_address(ip: &IpAddr) -> bool {
+        let is_multicast = Self::is_multicast(ip);
         match ip {
-            IpAddr::V4(a) => a.is_link_local() || a.is_private(),
-            IpAddr::V6(a) => is_unicast_link_local(a),
+            IpAddr::V4(a) => a.is_link_local() || a.is_private() || is_multicast,
+            IpAddr::V6(a) => is_unicast_link_local(a) || is_multicast,
         }
     }
 
@@ -666,6 +674,12 @@ impl Labeler {
         let src_data = &mut endpoint.src_info;
         let dst_data = &mut endpoint.dst_info;
 
+        if src_data.l3_epc_id == 0 && dst_data.l3_epc_id > 0 && Self::is_multicast(ip_src) {
+            src_data.l3_epc_id = dst_data.l3_epc_id;
+        } else if src_data.l3_epc_id > 0 && dst_data.l3_epc_id == 0 && Self::is_multicast(ip_dst) {
+            dst_data.l3_epc_id = src_data.l3_epc_id;
+        }
+
         if src_data.l3_epc_id == 0 && !Self::is_intranet_address(ip_src) {
             src_data.l3_epc_id = EPC_INTERNET;
         }
@@ -1228,6 +1242,24 @@ mod tests {
         labeler.modify_internet_epc(&ip, &ip, &mut endpoints);
         assert_eq!(endpoints.dst_info.l3_epc_id, 0);
         assert_eq!(endpoints.src_info.l3_epc_id, 0);
+
+        endpoints.src_info.l3_epc_id = 10;
+        endpoints.dst_info.l3_epc_id = 0;
+        let multicast_ip = IpAddr::V4(Ipv4Addr::new(224, 0, 0, 10));
+        labeler.modify_internet_epc(&ip, &multicast_ip, &mut endpoints);
+        assert_eq!(endpoints.dst_info.l3_epc_id, 10);
+
+        endpoints.src_info.l3_epc_id = 10;
+        endpoints.dst_info.l3_epc_id = 0;
+        let multicast_ip = IpAddr::V4(Ipv4Addr::new(123, 0, 0, 10));
+        labeler.modify_internet_epc(&ip, &multicast_ip, &mut endpoints);
+        assert_eq!(endpoints.dst_info.l3_epc_id, -2);
+
+        endpoints.src_info.l3_epc_id = 0;
+        endpoints.dst_info.l3_epc_id = 0;
+        let multicast_ip = IpAddr::V6(Ipv6Addr::new(0xff00, 0, 0, 0, 0, 0, 0, 1));
+        labeler.modify_internet_epc(&multicast_ip, &ip, &mut endpoints);
+        assert_eq!(endpoints.src_info.l3_epc_id, 0);
     }
 
     #[test]