need encrypt rpc if dst is in peer map (#1151)

easytier/src/gateway/icmp_proxy.rs

@@ -306,7 +306,7 @@ impl IcmpProxy {
             return Err(anyhow::anyhow!("peer manager is gone").into());
         };
 
-        pm.add_packet_process_pipeline(self.clone()).await;
+        pm.add_packet_process_pipeline(Box::new(self.clone())).await;
         Ok(())
     }
 

easytier/src/gateway/kcp_proxy.rs

@@ -341,13 +341,13 @@ impl KcpProxySrc {
 
     pub async fn start(&self) {
         self.peer_manager
-            .add_nic_packet_process_pipeline(Arc::new(self.tcp_proxy.clone()))
+            .add_nic_packet_process_pipeline(Box::new(self.tcp_proxy.clone()))
             .await;
         self.peer_manager
-            .add_packet_process_pipeline(Arc::new(self.tcp_proxy.0.clone()))
+            .add_packet_process_pipeline(Box::new(self.tcp_proxy.0.clone()))
             .await;
         self.peer_manager
-            .add_packet_process_pipeline(Arc::new(KcpEndpointFilter {
+            .add_packet_process_pipeline(Box::new(KcpEndpointFilter {
                 kcp_endpoint: self.kcp_endpoint.clone(),
                 is_src: true,
             }))
@@ -484,7 +484,7 @@ impl KcpProxyDst {
     pub async fn start(&mut self) {
         self.run_accept_task().await;
         self.peer_manager
-            .add_packet_process_pipeline(Arc::new(KcpEndpointFilter {
+            .add_packet_process_pipeline(Box::new(KcpEndpointFilter {
                 kcp_endpoint: self.kcp_endpoint.clone(),
                 is_src: false,
             }))

easytier/src/gateway/quic_proxy.rs

@@ -227,10 +227,10 @@ impl QUICProxySrc {
 
     pub async fn start(&self) {
         self.peer_manager
-            .add_nic_packet_process_pipeline(Arc::new(self.tcp_proxy.clone()))
+            .add_nic_packet_process_pipeline(Box::new(self.tcp_proxy.clone()))
             .await;
         self.peer_manager
-            .add_packet_process_pipeline(Arc::new(self.tcp_proxy.0.clone()))
+            .add_packet_process_pipeline(Box::new(self.tcp_proxy.0.clone()))
             .await;
         self.tcp_proxy.0.start(false).await.unwrap();
     }

easytier/src/gateway/socks5.rs

@@ -621,7 +621,7 @@ impl Socks5Server {
 
         if need_start {
             self.peer_manager
-                .add_packet_process_pipeline(self.clone())
+                .add_packet_process_pipeline(Box::new(self.clone()))
                 .await;
 
             self.run_net_update_task().await;

easytier/src/gateway/tcp_proxy.rs

@@ -476,10 +476,10 @@ impl<C: NatDstConnector> TcpProxy<C> {
         self.run_listener().await?;
         if add_pipeline {
             self.peer_manager
-                .add_packet_process_pipeline(self.clone())
+                .add_packet_process_pipeline(Box::new(self.clone()))
                 .await;
             self.peer_manager
-                .add_nic_packet_process_pipeline(self.clone())
+                .add_nic_packet_process_pipeline(Box::new(self.clone()))
                 .await;
         }
         join_joinset_background(self.tasks.clone(), "TcpProxy".to_owned());

easytier/src/gateway/udp_proxy.rs

@@ -404,7 +404,7 @@ impl UdpProxy {
 
     pub async fn start(self: &Arc<Self>) -> Result<(), Error> {
         self.peer_manager
-            .add_packet_process_pipeline(self.clone())
+            .add_packet_process_pipeline(Box::new(self.clone()))
             .await;
 
         // clean up nat table

easytier/src/instance/dns_server/server_instance.rs

@@ -404,7 +404,9 @@ impl MagicDnsServerInstance {
             .register(MagicDnsServerRpcServer::new(data.clone()), "");
         rpc_server.set_hook(data.clone());
 
-        peer_mgr.add_nic_packet_process_pipeline(data.clone()).await;
+        peer_mgr
+            .add_nic_packet_process_pipeline(Box::new(data.clone()))
+            .await;
 
         let data_clone = data.clone();
         tokio::task::spawn_blocking(move || data_clone.do_system_config(DEFAULT_ET_DNS_ZONE))

easytier/src/peers/mod.rs

@@ -23,8 +23,6 @@ pub mod peer_task;
 #[cfg(test)]
 pub mod tests;
 
-use std::sync::Arc;
-
 use crate::tunnel::packet_def::ZCPacket;
 
 #[async_trait::async_trait]
@@ -45,8 +43,8 @@ pub trait NicPacketFilter {
     }
 }
 
-type BoxPeerPacketFilter = Arc<dyn PeerPacketFilter + Send + Sync>;
-type BoxNicPacketFilter = Arc<dyn NicPacketFilter + Send + Sync>;
+type BoxPeerPacketFilter = Box<dyn PeerPacketFilter + Send + Sync>;
+type BoxNicPacketFilter = Box<dyn NicPacketFilter + Send + Sync>;
 
 // pub type PacketRecvChan = tachyonix::Sender<ZCPacket>;
 // pub type PacketRecvChanReceiver = tachyonix::Receiver<ZCPacket>;

easytier/src/peers/peer_manager.rs

@@ -6,7 +6,6 @@ use std::{
 };
 
 use anyhow::Context;
-use arc_swap::ArcSwap;
 use async_trait::async_trait;
 
 use dashmap::DashMap;
@@ -14,7 +13,7 @@ use dashmap::DashMap;
 use tokio::{
     sync::{
         mpsc::{self, UnboundedReceiver, UnboundedSender},
-        Mutex,
+        Mutex, RwLock,
     },
     task::JoinSet,
 };
@@ -132,8 +131,8 @@ pub struct PeerManager {
     peer_rpc_mgr: Arc<PeerRpcManager>,
     peer_rpc_tspt: Arc<RpcTransport>,
 
-    peer_packet_process_pipeline: Arc<ArcSwap<Vec<BoxPeerPacketFilter>>>,
-    nic_packet_process_pipeline: ArcSwap<Vec<BoxNicPacketFilter>>,
+    peer_packet_process_pipeline: Arc<RwLock<Vec<BoxPeerPacketFilter>>>,
+    nic_packet_process_pipeline: Arc<RwLock<Vec<BoxNicPacketFilter>>>,
 
     route_algo_inst: RouteAlgoInst,
 
@@ -262,8 +261,8 @@ impl PeerManager {
             peer_rpc_mgr,
             peer_rpc_tspt: rpc_tspt,
 
-            peer_packet_process_pipeline: Arc::new(ArcSwap::from(Arc::new(Vec::new()))),
-            nic_packet_process_pipeline: ArcSwap::from(Arc::new(Vec::new())),
+            peer_packet_process_pipeline: Arc::new(RwLock::new(Vec::new())),
+            nic_packet_process_pipeline: Arc::new(RwLock::new(Vec::new())),
 
             route_algo_inst,
 
@@ -648,7 +647,7 @@ impl PeerManager {
                     let mut processed = false;
                     let mut zc_packet = Some(ret);
                     let mut idx = 0;
-                    for pipeline in pipe_line.load().iter().rev() {
+                    for pipeline in pipe_line.read().await.iter().rev() {
                         tracing::trace!(?zc_packet, ?idx, "try_process_packet_from_peer");
                         idx += 1;
                         zc_packet = pipeline
@@ -670,20 +669,18 @@ impl PeerManager {
 
     pub async fn add_packet_process_pipeline(&self, pipeline: BoxPeerPacketFilter) {
         // newest pipeline will be executed first
-        let current = self.peer_packet_process_pipeline.load();
-        let mut new_pipelines = (*(*current)).iter().map(|x| x.clone()).collect::<Vec<_>>();
-        new_pipelines.push(pipeline);
         self.peer_packet_process_pipeline
-            .swap(Arc::new(new_pipelines));
+            .write()
+            .await
+            .push(pipeline);
     }
 
     pub async fn add_nic_packet_process_pipeline(&self, pipeline: BoxNicPacketFilter) {
         // newest pipeline will be executed first
-        let current = self.nic_packet_process_pipeline.load();
-        let mut new_pipelines = (*current).iter().map(|x| x.clone()).collect::<Vec<_>>();
-        new_pipelines.push(pipeline);
         self.nic_packet_process_pipeline
-            .swap(Arc::new(new_pipelines));
+            .write()
+            .await
+            .push(pipeline);
     }
 
     async fn init_packet_process_pipeline(&self) {
@@ -705,7 +702,7 @@ impl PeerManager {
                 }
             }
         }
-        self.add_packet_process_pipeline(Arc::new(NicPacketProcessor {
+        self.add_packet_process_pipeline(Box::new(NicPacketProcessor {
             nic_channel: self.nic_channel.clone(),
         }))
         .await;
@@ -730,7 +727,7 @@ impl PeerManager {
                 }
             }
         }
-        self.add_packet_process_pipeline(Arc::new(PeerRpcPacketProcessor {
+        self.add_packet_process_pipeline(Box::new(PeerRpcPacketProcessor {
             peer_rpc_tspt_sender: self.peer_rpc_tspt.peer_rpc_tspt_sender.clone(),
         }))
         .await;
@@ -738,8 +735,12 @@ impl PeerManager {
 
     pub async fn add_route<T>(&self, route: T)
     where
-        T: Route + Send + Sync + Clone + 'static,
+        T: Route + PeerPacketFilter + Send + Sync + Clone + 'static,
     {
+        // for route
+        self.add_packet_process_pipeline(Box::new(route.clone()))
+            .await;
+
         struct Interface {
             my_peer_id: PeerId,
             peers: Weak<PeerMap>,
@@ -865,19 +866,15 @@ impl PeerManager {
             return;
         }
 
-        let pipelines = self.nic_packet_process_pipeline.load();
-        for pipeline in pipelines.iter().rev() {
+        for pipeline in self.nic_packet_process_pipeline.read().await.iter().rev() {
             let _ = pipeline.try_process_packet_from_nic(data).await;
         }
     }
 
     pub async fn remove_nic_packet_process_pipeline(&self, id: String) -> Result<(), Error> {
-        let current = self.nic_packet_process_pipeline.load();
-        let mut new_pipelines = (*current).iter().map(|x| x.clone()).collect::<Vec<_>>();
-        if let Some(pos) = new_pipelines.iter().position(|x| x.id() == id) {
-            new_pipelines.remove(pos);
-            self.nic_packet_process_pipeline
-                .swap(Arc::new(new_pipelines));
+        let mut pipelines = self.nic_packet_process_pipeline.write().await;
+        if let Some(pos) = pipelines.iter().position(|x| x.id() == id) {
+            pipelines.remove(pos);
             Ok(())
         } else {
             Err(Error::NotFound)
@@ -1209,9 +1206,10 @@ impl PeerManager {
     }
 
     pub async fn clear_resources(&self) {
-        self.peer_packet_process_pipeline
-            .store(Arc::new(Vec::new()));
-        self.nic_packet_process_pipeline.store(Arc::new(Vec::new()));
+        let mut peer_pipeline = self.peer_packet_process_pipeline.write().await;
+        peer_pipeline.clear();
+        let mut nic_pipeline = self.nic_packet_process_pipeline.write().await;
+        nic_pipeline.clear();
 
         self.peer_rpc_mgr.rpc_server().registry().unregister_all();
     }

easytier/src/peers/peer_map.rs

@@ -4,8 +4,8 @@ use std::{
 };
 
 use anyhow::Context;
-use arc_swap::ArcSwap;
 use dashmap::{DashMap, DashSet};
+use tokio::sync::RwLock;
 
 use crate::{
     common::{
@@ -32,7 +32,7 @@ pub struct PeerMap {
     my_peer_id: PeerId,
     peer_map: DashMap<PeerId, Arc<Peer>>,
     packet_send: PacketRecvChan,
-    route: ArcSwap<Option<ArcRoute>>,
+    routes: RwLock<Vec<ArcRoute>>,
     alive_conns: Arc<DashMap<(PeerId, PeerConnId), PeerConnInfo>>,
 }
 
@@ -43,7 +43,7 @@ impl PeerMap {
             my_peer_id,
             peer_map: DashMap::new(),
             packet_send,
-            route: ArcSwap::from(Arc::new(None)),
+            routes: RwLock::new(Vec::new()),
             alive_conns: Arc::new(DashMap::new()),
         }
     }
@@ -154,7 +154,7 @@ impl PeerMap {
         }
 
         // get route info
-        if let Some(route) = self.route.load().as_ref() {
+        for route in self.routes.read().await.iter() {
             if let Some(gateway_peer_id) = route
                 .get_next_hop_with_policy(dst_peer_id, policy.clone())
                 .await
@@ -171,13 +171,14 @@ impl PeerMap {
         &self,
         network_identity: &NetworkIdentity,
     ) -> Vec<PeerId> {
-        if let Some(route) = self.route.load().as_ref() {
-            route
+        let mut ret = Vec::new();
+        for route in self.routes.read().await.iter() {
+            let peers = route
                 .list_peers_own_foreign_network(&network_identity)
-                .await
-        } else {
-            Vec::new()
+                .await;
+            ret.extend(peers);
         }
+        ret
     }
 
     pub async fn send_msg(
@@ -198,42 +199,49 @@ impl PeerMap {
     }
 
     pub async fn get_peer_id_by_ipv4(&self, ipv4: &Ipv4Addr) -> Option<PeerId> {
-        if let Some(route) = self.route.load().as_ref() {
-            route.get_peer_id_by_ipv4(ipv4).await
-        } else {
-            None
+        for route in self.routes.read().await.iter() {
+            let peer_id = route.get_peer_id_by_ipv4(ipv4).await;
+            if peer_id.is_some() {
+                return peer_id;
             }
         }
+        None
+    }
 
     pub async fn get_peer_id_by_ipv6(&self, ipv6: &Ipv6Addr) -> Option<PeerId> {
-        if let Some(route) = self.route.load().as_ref() {
-            route.get_peer_id_by_ipv6(ipv6).await
-        } else {
-            None
+        for route in self.routes.read().await.iter() {
+            let peer_id = route.get_peer_id_by_ipv6(ipv6).await;
+            if peer_id.is_some() {
+                return peer_id;
             }
         }
+        None
+    }
 
     pub async fn get_route_peer_info(&self, peer_id: PeerId) -> Option<RoutePeerInfo> {
-        if let Some(route) = self.route.load().as_ref() {
-            route.get_peer_info(peer_id).await
-        } else {
-            None
+        for route in self.routes.read().await.iter() {
+            if let Some(info) = route.get_peer_info(peer_id).await {
+                return Some(info);
             }
         }
+        None
+    }
 
     pub async fn get_origin_my_peer_id(
         &self,
         network_name: &str,
         foreign_my_peer_id: PeerId,
     ) -> Option<PeerId> {
-        if let Some(route) = self.route.load().as_ref() {
-            route
+        for route in self.routes.read().await.iter() {
+            let origin_peer_id = route
                 .get_origin_my_peer_id(network_name, foreign_my_peer_id)
-                .await
-        } else {
-            None
+                .await;
+            if origin_peer_id.is_some() {
+                return origin_peer_id;
             }
         }
+        None
+    }
 
     pub fn is_empty(&self) -> bool {
         self.peer_map.is_empty()
@@ -301,7 +309,8 @@ impl PeerMap {
     }
 
     pub async fn add_route(&self, route: ArcRoute) {
-        self.route.store(Arc::new(Some(route)));
+        let mut routes = self.routes.write().await;
+        routes.insert(0, route);
     }
 
     pub async fn clean_peer_without_conn(&self) {
@@ -321,7 +330,7 @@ impl PeerMap {
 
     pub async fn list_routes(&self) -> DashMap<PeerId, PeerId> {
         let route_map = DashMap::new();
-        if let Some(route) = self.route.load().as_ref() {
+        for route in self.routes.read().await.iter() {
             for item in route.list_routes().await.iter() {
                 route_map.insert(item.peer_id, item.next_hop_peer_id);
             }
@@ -330,11 +339,10 @@ impl PeerMap {
     }
 
     pub async fn list_route_infos(&self) -> Vec<cli::Route> {
-        if let Some(route) = self.route.load().as_ref() {
-            route.list_routes().await
-        } else {
-            vec![]
+        for route in self.routes.read().await.iter() {
+            return route.list_routes().await;
         }
+        vec![]
     }
 
     pub async fn need_relay_by_foreign_network(&self, dst_peer_id: PeerId) -> Result<bool, Error> {
@@ -375,42 +383,3 @@ impl Drop for PeerMap {
         );
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::common::global_ctx::tests::get_mock_global_ctx;
-    use tokio::sync::mpsc;
-
-    #[tokio::test]
-    async fn test_peer_map_route_arcswap() {
-        let (packet_send, _packet_recv) = mpsc::channel(128);
-        let global_ctx = get_mock_global_ctx();
-        let my_peer_id = 1;
-
-        let peer_map = PeerMap::new(packet_send, global_ctx, my_peer_id);
-
-        // Initially, no route should be set
-        assert!(peer_map.route.load().is_none());
-
-        // Test that methods return None/empty when no route is set
-        assert_eq!(
-            peer_map
-                .get_gateway_peer_id(2, NextHopPolicy::LeastHop)
-                .await,
-            None
-        );
-        assert_eq!(
-            peer_map
-                .get_peer_id_by_ipv4(&"192.168.1.1".parse().unwrap())
-                .await,
-            None
-        );
-        assert_eq!(peer_map.get_route_peer_info(2).await, None);
-        assert_eq!(peer_map.list_route_infos().await.len(), 0);
-
-        // The route field should be accessible and work with ArcSwap
-        let route_loaded = peer_map.route.load();
-        assert!(route_loaded.is_none());
-    }
-}

easytier/src/peers/peer_ospf_route.rs

@@ -55,6 +55,7 @@ use super::{
         DefaultRouteCostCalculator, ForeignNetworkRouteInfoMap, NextHopPolicy, RouteCostCalculator,
         RouteCostCalculatorInterface,
     },
+    PeerPacketFilter,
 };
 
 static SERVICE_ID: u32 = 7;
@@ -2368,6 +2369,8 @@ impl Route for PeerRoute {
     }
 }
 
+impl PeerPacketFilter for Arc<PeerRoute> {}
+
 #[cfg(test)]
 mod tests {
     use std::{

easytier/src/vpn_portal/wireguard.rs

@@ -203,7 +203,7 @@ impl WireGuardImpl {
         }
 
         self.peer_mgr
-            .add_packet_process_pipeline(Arc::new(PeerPacketFilterForVpnPortal {
+            .add_packet_process_pipeline(Box::new(PeerPacketFilterForVpnPortal {
                 wg_peer_ip_table: self.wg_peer_ip_table.clone(),
             }))
             .await;