handle close peer conn correctly (#1082)

easytier/src/peers/peer_manager.rs

@@ -1128,6 +1128,35 @@ impl PeerManager {
 
         self.peer_rpc_mgr.rpc_server().registry().unregister_all();
     }
+
+    pub async fn close_peer_conn(
+        &self,
+        peer_id: PeerId,
+        conn_id: &PeerConnId,
+    ) -> Result<(), Error> {
+        let ret = self.peers.close_peer_conn(peer_id, conn_id).await;
+        tracing::info!("close_peer_conn in peer map: {:?}", ret);
+        if ret.is_ok() || !matches!(ret.as_ref().unwrap_err(), Error::NotFound) {
+            return ret;
+        }
+
+        let ret = self
+            .foreign_network_client
+            .get_peer_map()
+            .close_peer_conn(peer_id, conn_id)
+            .await;
+        tracing::info!("close_peer_conn in foreign network client: {:?}", ret);
+        if ret.is_ok() || !matches!(ret.as_ref().unwrap_err(), Error::NotFound) {
+            return ret;
+        }
+
+        let ret = self
+            .foreign_network_manager
+            .close_peer_conn(peer_id, conn_id)
+            .await;
+        tracing::info!("close_peer_conn in foreign network manager done: {:?}", ret);
+        ret
+    }
 }
 
 #[cfg(test)]
@@ -1147,7 +1176,10 @@ mod tests {
             peer_manager::RouteAlgoType,
             peer_rpc::tests::register_service,
             route_trait::NextHopPolicy,
-            tests::{connect_peer_manager, wait_route_appear, wait_route_appear_with_cost},
+            tests::{
+                connect_peer_manager, create_mock_peer_manager_with_name, wait_route_appear,
+                wait_route_appear_with_cost,
+            },
         },
         proto::common::{CompressionAlgoPb, NatType, PeerFeatureFlag},
         tunnel::{
@@ -1427,4 +1459,127 @@ mod tests {
         )
         .await;
     }
+
+    #[tokio::test]
+    async fn close_conn_in_peer_map() {
+        let peer_mgr_a = create_mock_peer_manager_with_mock_stun(NatType::Unknown).await;
+        let peer_mgr_b = create_mock_peer_manager_with_mock_stun(NatType::Unknown).await;
+        connect_peer_manager(peer_mgr_a.clone(), peer_mgr_b.clone()).await;
+        wait_route_appear(peer_mgr_a.clone(), peer_mgr_b.clone())
+            .await
+            .unwrap();
+
+        let conns = peer_mgr_a
+            .get_peer_map()
+            .list_peer_conns(peer_mgr_b.my_peer_id)
+            .await;
+        assert!(conns.is_some());
+        let conn_info = conns.as_ref().unwrap().first().unwrap();
+
+        peer_mgr_a
+            .close_peer_conn(peer_mgr_b.my_peer_id, &conn_info.conn_id.parse().unwrap())
+            .await
+            .unwrap();
+
+        wait_for_condition(
+            || async {
+                let peers = peer_mgr_a.list_peers().await;
+                peers.is_empty()
+            },
+            Duration::from_secs(10),
+        )
+        .await;
+        // a is client, b is server
+    }
+
+    #[tokio::test]
+    async fn close_conn_in_foreign_network_client() {
+        let peer_mgr_server = create_mock_peer_manager_with_name("server".to_string()).await;
+        let peer_mgr_client = create_mock_peer_manager_with_name("client".to_string()).await;
+        connect_peer_manager(peer_mgr_client.clone(), peer_mgr_server.clone()).await;
+        wait_for_condition(
+            || async {
+                peer_mgr_client
+                    .get_foreign_network_client()
+                    .list_public_peers()
+                    .await
+                    .len()
+                    == 1
+            },
+            Duration::from_secs(3),
+        )
+        .await;
+
+        let peer_id = peer_mgr_client
+            .foreign_network_client
+            .list_public_peers()
+            .await[0];
+        let conns = peer_mgr_client
+            .foreign_network_client
+            .get_peer_map()
+            .list_peer_conns(peer_id)
+            .await;
+        assert!(conns.is_some());
+        let conn_info = conns.as_ref().unwrap().first().unwrap();
+        peer_mgr_client
+            .close_peer_conn(peer_id, &conn_info.conn_id.parse().unwrap())
+            .await
+            .unwrap();
+
+        wait_for_condition(
+            || async {
+                peer_mgr_client
+                    .get_foreign_network_client()
+                    .list_public_peers()
+                    .await
+                    .len()
+                    == 0
+            },
+            Duration::from_secs(10),
+        )
+        .await;
+    }
+
+    #[tokio::test]
+    async fn close_conn_in_foreign_network_manager() {
+        let peer_mgr_server = create_mock_peer_manager_with_name("server".to_string()).await;
+        let peer_mgr_client = create_mock_peer_manager_with_name("client".to_string()).await;
+        connect_peer_manager(peer_mgr_client.clone(), peer_mgr_server.clone()).await;
+        wait_for_condition(
+            || async {
+                peer_mgr_client
+                    .get_foreign_network_client()
+                    .list_public_peers()
+                    .await
+                    .len()
+                    == 1
+            },
+            Duration::from_secs(3),
+        )
+        .await;
+
+        let conns = peer_mgr_server
+            .foreign_network_manager
+            .list_foreign_networks()
+            .await;
+        let client_info = conns.foreign_networks["client"].peers[0].clone();
+        let conn_info = client_info.conns[0].clone();
+        peer_mgr_server
+            .close_peer_conn(client_info.peer_id, &conn_info.conn_id.parse().unwrap())
+            .await
+            .unwrap();
+
+        wait_for_condition(
+            || async {
+                peer_mgr_client
+                    .get_foreign_network_client()
+                    .list_public_peers()
+                    .await
+                    .len()
+                    == 0
+            },
+            Duration::from_secs(10),
+        )
+        .await;
+    }
 }