增加安卓端支持、优化广播、增加停止监听

switch/Cargo.toml

@@ -13,21 +13,15 @@ libc = "0.2.137"
 crossbeam-utils = "0.8"
 crossbeam-skiplist = "0.1"
 parking_lot = "0.12.1"
-
-#rsa = "0.7.2"
 rand = "0.8.5"
 sha2 = { version = "0.10.6", features = ["oid"] }
 aes-gcm = "0.10.2"
 
 thiserror = "1.0.37"
-#chrono = "0.4.23"
-#lazy_static = "1.4.0"
-#moka = "0.9.6"
 protobuf = "3.2.0"
-#local-ip-address = "0.4.9"
 socket2 ={ version = "0.5.2", features = ["all"] }
 tokio = { version = "1.28.1", features = ["full"] }
-[target.'cfg(any(unix))'.dependencies]
+[target.'cfg(any(target_os = "linux",target_os = "macos"))'.dependencies]
 tun = { path = "./rust-tun" }
 
 [target.'cfg(target_os = "windows")'.dependencies]

switch/rust-tun/src/platform/mod.rs

@@ -27,15 +27,6 @@ pub mod macos;
 #[cfg(target_os = "macos")]
 pub use self::macos::{create, Configuration, Device, Queue};
 
-#[cfg(target_os = "ios")]
-pub mod ios;
-#[cfg(target_os = "ios")]
-pub use self::ios::{create, Configuration, Device, Queue};
-
-#[cfg(target_os = "android")]
-pub mod android;
-#[cfg(target_os = "android")]
-pub use self::android::{create, Configuration, Device, Queue};
 
 #[cfg(test)]
 mod test {

switch/src/channel/channel.rs

@@ -10,10 +10,11 @@ use tokio::net::UdpSocket;
 use tokio::sync::watch::{channel, Receiver, Sender};
 use crate::channel::{Route, RouteKey, Status};
 use crate::channel::punch::NatType;
+use crate::core::status::SwitchWorker;
 use crate::handle::recv_handler::ChannelDataHandler;
 
 pub struct ContextInner {
-    pub(crate) lock:Mutex<()>,
+    pub(crate) lock: Mutex<()>,
     pub(crate) count: AtomicUsize,
     pub(crate) main_channel: Arc<UdpSocket>,
     pub(crate) route_table: SkipMap<Ipv4Addr, Vec<Route>>,
@@ -35,7 +36,7 @@ impl Context {
         let channel_num = 1;
         let (status_sender, status_receiver) = channel(Status::Cone);
         let inner = Arc::new(ContextInner {
-            lock:Mutex::new(()),
+            lock: Mutex::new(()),
             count: AtomicUsize::new(0),
             main_channel,
             route_table: SkipMap::new(),
@@ -143,9 +144,9 @@ impl Context {
     fn add_route_(&self, id: Ipv4Addr, route: Route, only_if_absent: bool) {
         let key = route.route_key();
         let guard = self.inner.lock.lock();
-        let mut list = if let Some(entry) = self.inner.route_table.get(&id){
+        let mut list = if let Some(entry) = self.inner.route_table.get(&id) {
             entry.value().clone()
-        }else{
+        } else {
             Vec::with_capacity(4)
         };
         let mut exist = false;
@@ -178,7 +179,7 @@ impl Context {
                 list.truncate(max_len);
             }
         }
-        self.inner.route_table.insert(id,list);
+        self.inner.route_table.insert(id, list);
         self.inner.route_table_time.insert((key, id), AtomicCell::new(Instant::now()));
         drop(guard);
     }
@@ -250,7 +251,7 @@ impl Context {
             let mut routes = v.value().clone();
             drop(v);
             routes.retain(|x| x.route_key() != route_key);
-            self.inner.route_table.insert(*id,routes);
+            self.inner.route_table.insert(*id, routes);
             self.inner.route_table_time.remove(&(route_key, *id));
         }
         drop(guard);
@@ -294,54 +295,63 @@ impl Channel {
         }
     }
     pub async fn start(self,
+                       mut worker: SwitchWorker,
                        head_reserve: usize,//头部预留字节
                        symmetric_channel_num: usize,//对称网络，则再加一组监听，提升打洞成功率
     ) {
         let context = self.context;
         let main_channel = context.inner.main_channel.clone();
         let handler = self.handler.clone();
-        tokio::spawn(Self::start_(context.clone(), handler.clone(), main_channel.clone(), head_reserve, true));
+        tokio::spawn(Self::start_(worker.clone(), context.clone(), handler.clone(), main_channel.clone(), head_reserve, true));
-        tokio::spawn(Self::start_(context.clone(), handler, main_channel, head_reserve, true));
+        tokio::spawn(Self::start_(worker.clone(), context.clone(), handler, main_channel, head_reserve, true));
         let mut cur_status = Status::Cone;
         let mut status_receiver = context.inner.status_receiver.clone();
         loop {
-            match status_receiver.changed().await {
+            tokio::select! {
-                Ok(_) => {
+                _=worker.stop_wait()=>{
-                    match *status_receiver.borrow() {
+                    break;
-                        Status::Cone => {
+                }
-                            cur_status = Status::Cone;
+                rs=status_receiver.changed()=>{
-                        }
+                    match rs {
-                        Status::Symmetric => {
+                        Ok(_) => {
-                            if cur_status == Status::Symmetric {
+                            match *status_receiver.borrow() {
-                                continue;
+                                Status::Cone => {
-                            }
+                                    cur_status = Status::Cone;
-                            cur_status = Status::Symmetric;
+                                }
-                            for _ in 0..symmetric_channel_num {
+                                Status::Symmetric => {
-                                match UdpSocket::bind("0.0.0.0:0").await {
+                                    if cur_status == Status::Symmetric {
-                                    Ok(udp) => {
+                                        continue;
-                                        let udp = Arc::new(udp);
-                                        let context = context.clone();
-                                        let handler = self.handler.clone();
-                                        tokio::spawn(Self::start_(context, handler, udp, head_reserve, false));
                                     }
-                                    Err(e) => {
+                                    cur_status = Status::Symmetric;
-                                        log::error!("{}",e);
+                                    for _ in 0..symmetric_channel_num {
+                                        match UdpSocket::bind("0.0.0.0:0").await {
+                                            Ok(udp) => {
+                                                let udp = Arc::new(udp);
+                                                let context = context.clone();
+                                                let handler = self.handler.clone();
+                                                tokio::spawn(Self::start_(worker.clone(),context, handler, udp, head_reserve, false));
+                                            }
+                                            Err(e) => {
+                                                log::error!("{}",e);
+                                            }
+                                        }
                                     }
                                 }
+                                Status::Close => {
+                                    break;
+                                }
                             }
                         }
-                        Status::Close => {
+                        Err(_) => {
                             break;
                         }
                     }
                 }
-                Err(_) => {
-                    break;
-                }
             }
         }
+        worker.stop_all();
     }
-    async fn start_(context: Context,
+    async fn start_(mut worker: SwitchWorker, context: Context,
                     mut handler: ChannelDataHandler,
                     udp: Arc<UdpSocket>,
                     head_reserve: usize,
@@ -382,8 +392,14 @@ impl Channel {
                             }
                         }
                 }
+                _=worker.stop_wait()=>{
+                    break;
+                }
             }
         }
         context.inner.udp_map.remove(&id);
+        if is_core {
+            worker.stop_all();
+        }
     }
 }

switch/src/core/mod.rs

@@ -2,34 +2,44 @@ use std::{io, thread};
 use std::net::{Ipv4Addr, SocketAddr};
 use std::sync::Arc;
 use std::time::Duration;
-use aes_gcm::{Aes256Gcm, Key, KeyInit};
 
-use crossbeam_utils::atomic::AtomicCell;
+use aes_gcm::{Aes256Gcm, Key, KeyInit};
 use crossbeam_skiplist::SkipMap;
+use crossbeam_utils::atomic::AtomicCell;
 use parking_lot::Mutex;
+use sha2::Digest;
 use tokio::net::UdpSocket;
 use tokio::sync::mpsc::channel;
 
-
+use crate::channel::{Route, RouteKey};
 use crate::channel::channel::{Channel, Context};
 use crate::channel::idle::Idle;
 use crate::channel::punch::{NatInfo, Punch};
-use crate::channel::{Route, RouteKey};
 use crate::channel::sender::ChannelSender;
-
+use crate::core::status::SwitchStatusManger;
+use crate::error::Error;
 use crate::external_route::ExternalRoute;
 use crate::handle::{ConnectStatus, CurrentDeviceInfo, heartbeat_handler, PeerDeviceInfo, punch_handler, registration_handler};
 use crate::handle::recv_handler::ChannelDataHandler;
-use crate::handle::tun_tap::{tap_handler, tun_handler};
+use crate::handle::registration_handler::{RegResponse, ReqEnum};
+#[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
+use crate::handle::tun_tap::tap_handler;
+use crate::handle::tun_tap::tun_handler;
 use crate::igmp_server::IgmpServer;
 use crate::nat::NatTest;
 use crate::tun_tap_device;
-use crate::tun_tap_device::DeviceWriter;
+use crate::tun_tap_device::{DeviceReader, DeviceWriter};
+
+pub mod status;
+pub mod sync;
+
 
 pub struct Switch {
     name: String,
     current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
     context: Context,
+    switch_status_manager: SwitchStatusManger,
+    #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
     device_writer: DeviceWriter,
     /// 0. 机器纪元，每一次上线或者下线都会增1，用于感知网络中机器变化
     /// 服务端和客户端的不一致，则服务端会推送新的设备列表
@@ -40,37 +50,121 @@ pub struct Switch {
     peer_nat_info_map: Arc<SkipMap<Ipv4Addr, NatInfo>>,
 }
 
-impl Switch {
+pub struct SwitchUtil {
-    pub async fn start(config: Config) -> crate::Result<Switch> {
+    config: Config,
-        log::info!("config:{:?}",config);
+    main_channel: Arc<UdpSocket>,
+    response: Option<RegResponse>,
+    iface: Option<(DeviceWriter, DeviceReader)>,
+}
+
+impl SwitchUtil {
+    pub async fn new(config: Config) -> io::Result<SwitchUtil> {
+        let main_channel = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
+        Ok(SwitchUtil {
+            config,
+            main_channel,
+            response: None,
+            iface: None,
+        })
+    }
+    pub async fn connect(&mut self) -> Result<RegResponse, ReqEnum> {
+        match registration_handler::registration(&self.main_channel, self.config.server_address,
+                                                 self.config.token.clone(), self.config.device_id.clone(),
+                                                 self.config.name.clone()).await {
+            Ok(res) => {
+                let _ = self.response.insert(res.clone());
+                Ok(res)
+            }
+            Err(e) => {
+                Err(e)
+            }
+        }
+    }
+    #[cfg(any(target_os = "android"))]
+    pub fn create_iface(&mut self, vpn_fd: i32) {
+        let (device_writer, device_reader) = tun_tap_device::create(vpn_fd);
+        let _ = self.iface.insert((device_writer, device_reader));
+    }
+    #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
+    pub fn create_iface(&mut self) -> io::Result<tun_tap_device::DriverInfo> {
+        if self.iface.is_some() {
+            return Err(io::Error::from(io::ErrorKind::AlreadyExists));
+        }
+        let response = match &self.response {
+            None => {
+                return Err(io::Error::from(io::ErrorKind::AlreadyExists));
+            }
+            Some(res) => {
+                res
+            }
+        };
+        let device_type = if self.config.tap {
+            #[cfg(windows)]
+            {
+                //删除switch的tun网卡避免ip冲突，因为非正常退出会保留网卡
+                tun_tap_device::delete_device(tun_tap_device::DeviceType::Tun);
+            }
+            tun_tap_device::DeviceType::Tap
+        } else {
+            #[cfg(windows)]
+            {
+                //删除switch的tap网卡避免ip冲突，非正常退出会保留网卡
+                tun_tap_device::delete_device(tun_tap_device::DeviceType::Tap);
+            }
+            tun_tap_device::DeviceType::Tun
+        };
+        let in_ips = self.config.in_ips.iter().map(|(dest, mask, _)| { (Ipv4Addr::from(*dest & *mask), Ipv4Addr::from(*mask)) }).collect::<Vec<(Ipv4Addr, Ipv4Addr)>>();
+
+        let (device_writer, device_reader, driver_info) = tun_tap_device::create_device(device_type, response.virtual_ip, response.virtual_netmask, response.virtual_gateway, in_ips)?;
+        let _ = self.iface.insert((device_writer, device_reader));
+        Ok(driver_info)
+    }
+    pub async fn build(self) -> crate::Result<Switch> {
+        let response = match self.response {
+            None => {
+                return Err(Error::Stop("response None".to_string()));
+            }
+            Some(res) => {
+                res
+            }
+        };
+        let (device_writer, device_reader) = match self.iface {
+            None => {
+                return Err(Error::Stop("iface None".to_string()));
+            }
+            Some(res) => {
+                res
+            }
+        };
+        let config = self.config;
+        let switch_status_manager = SwitchStatusManger::new();
         let cipher = if let Some(key) = &config.key {
             let key: &Key<Aes256Gcm> = key.into();
             Some(Aes256Gcm::new(&key))
         } else {
             None
         };
-        let main_channel = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
-        let response = registration_handler::registration(&main_channel, config.server_address, config.token.clone(), config.device_id.clone(), config.name.clone()).await?;
         let (cone_sender, cone_receiver) = channel(3);
         let (symmetric_sender, symmetric_receiver) = channel(2);
-        let context = Context::new(main_channel, 1);
+        let context = Context::new(self.main_channel, 1);
         let punch = Punch::new(context.clone());
         let idle = Idle::new(Duration::from_secs(16), context.clone());
         let channel_sender = ChannelSender::new(context.clone());
 
-        let register = Arc::new(registration_handler::Register::new(channel_sender.clone(), config.server_address, config.token.clone(), config.device_id.clone(), config.name.clone()));
+        let register = Arc::new(registration_handler::Register::new(channel_sender.clone(),
+                                                                    config.server_address, config.token.clone(),
+                                                                    config.device_id.clone(), config.name.clone()));
         let device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>> = Arc::new(Mutex::new((0, Vec::new())));
         let peer_nat_info_map: Arc<SkipMap<Ipv4Addr, NatInfo>> = Arc::new(SkipMap::new());
         let connect_status = Arc::new(AtomicCell::new(ConnectStatus::Connected));
-        let virtual_ip = Ipv4Addr::from(response.virtual_ip);
+        let virtual_ip = response.virtual_ip;
-        let virtual_gateway = Ipv4Addr::from(response.virtual_gateway);
+        let virtual_gateway = response.virtual_gateway;
-        let virtual_netmask = Ipv4Addr::from(response.virtual_netmask);
+        let virtual_netmask = response.virtual_netmask;
 
         let local_ip = crate::nat::local_ip()?;
         let local_port = context.main_local_port()?;
         // NAT检测
-        let nat_test = NatTest::new(config.nat_test_server.clone(), Ipv4Addr::from(response.public_ip), response.public_port as u16, local_ip, local_port);
+        let nat_test = NatTest::new(config.nat_test_server.clone(), response.public_ip, response.public_port, local_ip, local_port);
-        let in_ips = config.in_ips.iter().map(|(dest, mask, _)| { (Ipv4Addr::from(*dest & *mask), Ipv4Addr::from(*mask)) }).collect::<Vec<(Ipv4Addr, Ipv4Addr)>>();
 
         let out_ips = config.out_ips.iter().map(|(_, _, ip)| *ip).collect::<Vec<Ipv4Addr>>();
         let out_external_route = ExternalRoute::new(config.out_ips);
@@ -80,45 +174,30 @@ impl Switch {
             Some(ExternalRoute::new(config.in_ips))
         };
         let current_device = Arc::new(AtomicCell::new(CurrentDeviceInfo::new(virtual_ip, virtual_gateway, virtual_netmask, config.server_address)));
-        let ip_proxy_map = if out_ips.is_empty(){
+        let (tcp_proxy, udp_proxy, ip_proxy_map) = if out_ips.is_empty() {
-            None
+            (None, None, None)
-        }else{
-            Some(crate::ip_proxy::init_proxy(channel_sender.clone(), out_ips, current_device.clone()).await?)
-        };
-        let (device_writer, igmp_server) = if config.tap {
-            #[cfg(windows)]
-            {
-                //删除switch的tun网卡避免ip冲突，因为非正常退出会保留网卡
-                tun_tap_device::delete_device(tun_tap_device::DeviceType::Tap);
-            }
-            let (tap_writer, tap_reader) = tun_tap_device::create_device(tun_tap_device::DeviceType::Tap, virtual_ip, virtual_netmask, virtual_gateway, in_ips)?;
-            let igmp_server = if config.simulate_multicast {
-                Some(IgmpServer::new(tap_writer.clone()))
-            } else {
-                None
-            };
-            //tap数据处理
-            tap_handler::start(channel_sender.clone(), tap_reader.clone(), tap_writer.clone(),
-                               igmp_server.clone(), current_device.clone(), in_external_route, ip_proxy_map.clone(), cipher.clone());
-            (tap_writer, igmp_server)
         } else {
-            #[cfg(windows)]
+            let (tcp_proxy, udp_proxy, ip_proxy_map) = crate::ip_proxy::init_proxy(channel_sender.clone(), out_ips, current_device.clone()).await?;
-            {
+            (Some(tcp_proxy), Some(udp_proxy), Some(ip_proxy_map))
-                //删除switch的tap网卡避免ip冲突，非正常退出会保留网卡
-                tun_tap_device::delete_device(tun_tap_device::DeviceType::Tap);
-            }
-            // tun通道
-            let (tun_writer, tun_reader) = tun_tap_device::create_device(tun_tap_device::DeviceType::Tun, virtual_ip, virtual_netmask, virtual_gateway, in_ips)?;
-            let igmp_server = if config.simulate_multicast {
-                Some(IgmpServer::new(tun_writer.clone()))
-            } else {
-                None
-            };
-            //tun数据接收处理
-            tun_handler::start(channel_sender.clone(), tun_reader.clone(), tun_writer.clone(),
-                               igmp_server.clone(), current_device.clone(), in_external_route, ip_proxy_map.clone(), cipher.clone());
-            (tun_writer, igmp_server)
         };
+
+        let igmp_server = if config.simulate_multicast {
+            Some(IgmpServer::new(device_writer.clone()))
+        } else {
+            None
+        };
+        #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
+        if config.tap {
+            tap_handler::start(switch_status_manager.worker(), channel_sender.clone(), device_reader, device_writer.clone(),
+                               igmp_server.clone(), current_device.clone(), in_external_route, ip_proxy_map.clone(), cipher.clone());
+        } else {
+            tun_handler::start(switch_status_manager.worker(), channel_sender.clone(), device_reader, device_writer.clone(),
+                               igmp_server.clone(), current_device.clone(), in_external_route, ip_proxy_map.clone(), cipher.clone());
+        }
+        #[cfg(any(target_os = "android"))]
+        tun_handler::start(switch_status_manager.worker(), channel_sender.clone(), device_reader, device_writer.clone(),
+                           igmp_server.clone(), current_device.clone(), in_external_route, ip_proxy_map.clone(), cipher.clone());
+
         //外部数据接收处理
         let channel_recv_handler = ChannelDataHandler::new(current_device.clone(), device_list.clone(),
                                                            register.clone(), nat_test.clone(), igmp_server,
@@ -126,27 +205,53 @@ impl Switch {
                                                            peer_nat_info_map.clone(), ip_proxy_map, out_external_route,
                                                            cone_sender, symmetric_sender, cipher);
         let channel = Channel::new(context.clone(), channel_recv_handler);
+        let channel_worker = switch_status_manager.worker();
+        //数据接收
         thread::spawn(move || {
             tokio::runtime::Builder::new_multi_thread()
                 .enable_all()
                 .build().unwrap()
-                .block_on(channel.start(14, 60));
+                .block_on(async move {
+                    if let Some(tcp_proxy) = tcp_proxy {
+                        tokio::spawn(tcp_proxy.start());
+                    }
+                    if let Some(udp_proxy) = udp_proxy {
+                        tokio::spawn(udp_proxy.start());
+                    }
+                    channel.start(channel_worker, 14, 65).await;
+                });
         });
+        {
+            let other_worker = switch_status_manager.worker();
+            let nat_test = nat_test.clone();
+            let device_list = device_list.clone();
+            let current_device = current_device.clone();
+            //其他任务处理
+            thread::spawn(move || {
+                tokio::runtime::Builder::new_multi_thread()
+                    .enable_all()
+                    .build().unwrap()
+                    .block_on(async move {
+                        // 定时心跳
+                        heartbeat_handler::start_heartbeat(other_worker.clone(), channel_sender.clone(), device_list.clone(), current_device.clone());
+                        // 空闲检查
+                        heartbeat_handler::start_idle(other_worker.clone(), idle, channel_sender.clone());
+                        // 打洞处理
+                        punch_handler::start(other_worker.clone(), cone_receiver, punch.clone(), current_device.clone());
+                        punch_handler::start(other_worker.clone(), symmetric_receiver, punch, current_device.clone());
+                        punch_handler::start_punch(other_worker.clone(), nat_test.clone(),
+                                                   device_list.clone(), channel_sender.clone(),
+                                                   current_device.clone()).await;
+                    });
+            });
+        }
         context.switch(nat_test.nat_info().nat_type);
-        // 定时心跳
-        heartbeat_handler::start_heartbeat(channel_sender.clone(), device_list.clone(), current_device.clone()).await;
-        // 空闲检查
-        heartbeat_handler::start_idle(idle, channel_sender.clone()).await;
-        // 打洞处理
-        punch_handler::start(cone_receiver, punch.clone(), current_device.clone()).await;
-        punch_handler::start(symmetric_receiver, punch, current_device.clone()).await;
-        punch_handler::start_punch(nat_test.clone(), device_list.clone(), channel_sender.clone(), current_device.clone()).await;
-
-        log::info!("switch启动成功");
         Ok(Switch {
             name: config.name,
             current_device,
             context,
+            switch_status_manager,
+            #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
             device_writer,
             nat_test,
             device_list,
@@ -189,9 +294,15 @@ impl Switch {
     }
     pub fn stop(&self) -> io::Result<()> {
         self.context.close();
+        self.switch_status_manager.stop_all();
+        #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
         self.device_writer.close()?;
         Ok(())
     }
+    pub async fn wait_stop(&mut self) {
+        self.switch_status_manager.wait().await;
+        let _ = self.stop();
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -208,7 +319,6 @@ pub struct Config {
     pub simulate_multicast: bool,
 }
 
-use sha2::Digest;
 
 impl Config {
     pub fn new(tap: bool, token: String,

switch/src/core/status.rs

@@ -0,0 +1,85 @@
+use std::sync::Arc;
+use tokio::sync::watch;
+use tokio::sync::watch::{Receiver, Sender};
+use crate::util::wait::WaitGroup;
+
+#[derive(Copy, Clone, Eq, PartialEq)]
+pub enum SwitchStatus {
+    Starting,
+    Stopping,
+}
+
+pub struct SwitchWorker {
+    wg: WaitGroup,
+    status_s: Arc<Sender<SwitchStatus>>,
+    status_r: Receiver<SwitchStatus>,
+}
+
+impl Clone for SwitchWorker {
+    fn clone(&self) -> Self {
+        self.wg.add();
+        SwitchWorker {
+            wg: self.wg.clone(),
+            status_s: self.status_s.clone(),
+            status_r: self.status_r.clone(),
+        }
+    }
+}
+
+impl Drop for SwitchWorker {
+    fn drop(&mut self) {
+        self.wg.done();
+    }
+}
+
+impl SwitchWorker {
+    pub fn stop_all(&self) {
+        let _ = self.status_s.send(SwitchStatus::Stopping);
+    }
+    pub async fn stop_wait(&mut self) {
+        loop {
+            if *self.status_r.borrow() == SwitchStatus::Stopping {
+                return;
+            }
+            match self.status_r.changed().await {
+                Ok(_) => {
+                    if *self.status_r.borrow() == SwitchStatus::Stopping {
+                        return;
+                    }
+                }
+                Err(_) => { return; }
+            }
+        }
+    }
+}
+
+pub struct SwitchStatusManger {
+    wg: WaitGroup,
+    status_s: Arc<Sender<SwitchStatus>>,
+    status_r: Receiver<SwitchStatus>,
+}
+
+impl SwitchStatusManger {
+    pub fn new() -> Self {
+        let (status_s, status_r) = watch::channel(SwitchStatus::Starting);
+        Self {
+            wg: WaitGroup::new(),
+            status_s: Arc::new(status_s),
+            status_r,
+        }
+    }
+    pub fn stop_all(&self) {
+        let _ = self.status_s.send(SwitchStatus::Stopping);
+    }
+    pub async fn wait(&mut self) {
+        self.wg.wait().await
+    }
+    pub fn worker(&self) -> SwitchWorker {
+        self.wg.add();
+        SwitchWorker {
+            wg: self.wg.clone(),
+            status_s: self.status_s.clone(),
+            status_r: self.status_r.clone(),
+        }
+    }
+}

switch/src/core/sync.rs

@@ -0,0 +1,64 @@
+use std::io;
+use std::ops::Deref;
+use std::time::Duration;
+use tokio::runtime::Runtime;
+use crate::core::{Config, Switch, SwitchUtil};
+use crate::handle::registration_handler::{RegResponse, ReqEnum};
+
+pub struct SwitchUtilSync {
+    switch_util: SwitchUtil,
+    runtime: Runtime,
+}
+
+pub struct SwitchSync {
+    switch: Switch,
+    runtime: Runtime,
+}
+
+impl SwitchUtilSync {
+    pub fn new(config: Config) -> io::Result<SwitchUtilSync> {
+        let runtime = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();
+        let switch_util = runtime.block_on(SwitchUtil::new(config))?;
+        Ok(SwitchUtilSync {
+            switch_util,
+            runtime,
+        })
+    }
+    pub fn connect(&mut self) -> Result<RegResponse, ReqEnum> {
+        self.runtime.block_on(self.switch_util.connect())
+    }
+    #[cfg(any(target_os = "android"))]
+    pub fn create_iface(&mut self, vpn_fd: i32) {
+        self.switch_util.create_iface(vpn_fd)
+    }
+    #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
+    pub fn create_iface(&mut self) -> io::Result<crate::tun_tap_device::DriverInfo> {
+        self.switch_util.create_iface()
+    }
+    pub fn build(self) -> crate::Result<SwitchSync> {
+        let runtime = self.runtime;
+        let switch = runtime.block_on(self.switch_util.build())?;
+        Ok(SwitchSync {
+            switch,
+            runtime,
+        })
+    }
+}
+
+impl SwitchSync {
+    pub fn wait_stop(&mut self) {
+        self.runtime.block_on(self.switch.wait_stop())
+    }
+    pub fn wait_stop_ms(&mut self, ms: u64) -> bool {
+        self.runtime.block_on(tokio::time::timeout(Duration::from_millis(ms),
+                                                   self.switch.wait_stop())).is_ok()
+    }
+}
+
+impl Deref for SwitchSync {
+    type Target = Switch;
+
+    fn deref(&self) -> &Self::Target {
+        &self.switch
+    }
+}

switch/src/handle/heartbeat_handler.rs

@@ -9,20 +9,26 @@ use rand::prelude::SliceRandom;
 use crate::channel::idle::Idle;
 use crate::channel::Route;
 use crate::channel::sender::ChannelSender;
+use crate::core::status::SwitchWorker;
 
 
 use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo};
 use crate::protocol::control_packet::PingPacket;
 use crate::protocol::{control_packet, NetPacket, Protocol, Version};
 
-pub async fn start_idle(idle: Idle, sender: ChannelSender) {
+pub fn start_idle(mut worker: SwitchWorker, idle: Idle, sender: ChannelSender) {
     tokio::spawn(async move {
-        match start_idle_(idle, sender).await {
+        tokio::select! {
-            Ok(_) => {}
+             _=worker.stop_wait()=>{
-            Err(e) => {
+                    return;
-                log::warn!("空闲检测任务停止:{:?}", e);
+             }
+             rs=start_idle_(idle, sender)=>{
+                 if let Err(e) = rs {
+                    log::warn!("空闲检测任务停止:{:?}", e);
+                }
             }
         }
+        worker.stop_all();
     });
 }
 
@@ -38,15 +44,24 @@ async fn start_idle_(idle: Idle, sender: ChannelSender) -> io::Result<()> {
     }
 }
 
-pub async fn start_heartbeat(
+pub fn start_heartbeat(
+    mut worker: SwitchWorker,
     sender: ChannelSender,
     device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
     current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
 ) {
     tokio::spawn(async move {
-        if let Err(e) = start_heartbeat_(sender, device_list, current_device).await {
+        tokio::select! {
-            log::warn!("心跳任务停止:{:?}", e);
+             _=worker.stop_wait()=>{
+                    return;
+             }
+            rs=start_heartbeat_(sender, device_list, current_device)=>{
+                if let Err(e) = rs {
+                    log::warn!("心跳任务停止:{:?}", e);
+                }
+            }
         }
+        worker.stop_all();
     });
 }
 
@@ -70,6 +85,9 @@ async fn start_heartbeat_(
     net_packet.first_set_ttl(2);
     let mut count = 0;
     loop {
+        if sender.is_close() {
+            return Ok(());
+        }
         let current_device = current_device.load();
         net_packet.set_source(current_device.virtual_ip());
         {
@@ -103,7 +121,7 @@ async fn start_heartbeat_(
                     }
                 } else {
                     //没有直连路由则发送到网关
-                    let _ = sender.try_send_main(net_packet.buffer(), current_device.connect_server);
+                    let _ = sender.send_main(net_packet.buffer(), current_device.connect_server).await;
                     continue;
                 }
 

switch/src/handle/mod.rs

@@ -38,7 +38,7 @@ impl PeerDeviceInfo {
     }
 }
 
-#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq,Ord, PartialOrd)]
 pub enum PeerDeviceStatus {
     Online,
     Offline,

switch/src/handle/punch_handler.rs

@@ -13,10 +13,17 @@ use std::io;
 use tokio::sync::mpsc::Receiver;
 use crate::channel::punch::{NatInfo, Punch};
 use crate::channel::sender::ChannelSender;
+use crate::core::status::SwitchWorker;
 
-pub async fn start(receiver: Receiver<(Ipv4Addr, NatInfo)>, punch: Punch, current_device: Arc<AtomicCell<CurrentDeviceInfo>>) {
+pub fn start(mut worker: SwitchWorker, receiver: Receiver<(Ipv4Addr, NatInfo)>, punch: Punch, current_device: Arc<AtomicCell<CurrentDeviceInfo>>) {
     tokio::spawn(async move {
-        start0(receiver, punch, current_device).await;
+        tokio::select! {
+            _=start0(receiver, punch, current_device)=>{}
+            _=worker.stop_wait()=>{
+                return;
+            }
+        }
+        worker.stop_all();
     });
 }
 
@@ -47,54 +54,62 @@ async fn start_(
 }
 
 pub async fn start_punch(
+    mut worker: SwitchWorker,
     nat_test: NatTest,
     device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
     sender: ChannelSender,
     current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
 ) {
-    tokio::spawn(async move {
-        if let Err(e) = start_punch_(nat_test, device_list, sender, current_device).await {
-            log::warn!("打洞处理任务停止 {:?}", e);
-        }
-    });
-}
-
-async fn start_punch_(
-    nat_test: NatTest,
-    device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
-    sender: ChannelSender,
-    current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
-) -> crate::Result<()> {
     let mut num = 0;
     let sleep_time = [3, 5, 7, 11, 13, 17, 19, 23, 29];
     loop {
         if sender.is_close() {
-            return Ok(());
+            break;
         }
-        let current_device = current_device.load();
+        tokio::select! {
-        let nat_info = nat_test.nat_info();
+            rs= start_punch_(Duration::from_secs(sleep_time[num % sleep_time.len()]),&nat_test, &device_list, &sender, &current_device)=>{
-        {
+                 if let Err(e) = rs {
-            let mut list = device_list.lock().clone().1;
+                    log::warn!("打洞处理任务异常 {:?}", e);
-            list.shuffle(&mut rand::thread_rng());
-            let mut count = 0;
-            for info in list {
-                if info.virtual_ip <= current_device.virtual_ip {
-                    continue;
                 }
-                if !sender.need_punch(&info.virtual_ip) {
+            }
-                    continue;
+           _=worker.stop_wait()=>{
-                }
+                break;
-                count += 1;
-                if count > 2 {
-                    break;
-                }
-                let buf = punch_packet(current_device.virtual_ip(), &nat_info, info.virtual_ip)?;
-                let _ = sender.send_main(&buf, current_device.connect_server).await;
             }
         }
         num += 1;
-        tokio::time::sleep(Duration::from_secs(sleep_time[num % sleep_time.len()])).await;
     }
+
+    worker.stop_all();
+}
+
+async fn start_punch_(
+    sleep_time: Duration,
+    nat_test: &NatTest,
+    device_list: &Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
+    sender: &ChannelSender,
+    current_device: &Arc<AtomicCell<CurrentDeviceInfo>>,
+) -> crate::Result<()> {
+    let current_device = current_device.load();
+    let nat_info = nat_test.nat_info();
+    let mut list = device_list.lock().clone().1;
+    list.shuffle(&mut rand::thread_rng());
+    let mut count = 0;
+    for info in list {
+        if info.virtual_ip <= current_device.virtual_ip {
+            continue;
+        }
+        if !sender.need_punch(&info.virtual_ip) {
+            continue;
+        }
+        count += 1;
+        if count > 2 {
+            break;
+        }
+        let buf = punch_packet(current_device.virtual_ip(), &nat_info, info.virtual_ip)?;
+        let _ = sender.send_main(&buf, current_device.connect_server).await;
+    }
+    tokio::time::sleep(sleep_time).await;
+    Ok(())
 }
 
 pub fn punch_packet(

switch/src/handle/recv_handler.rs

@@ -245,7 +245,7 @@ impl ChannelDataHandler {
                                             ipv4.set_destination_ip(destination);
                                             ipv4.update_checksum();
                                             ip_proxy_map.tcp_proxy_map.insert(SocketAddrV4::new(source, source_port),
-                                                                                   (SocketAddrV4::new(gate_way, 0), SocketAddrV4::new(dest_ip, dest_port)));
+                                                                              (SocketAddrV4::new(gate_way, 0), SocketAddrV4::new(dest_ip, dest_port)));
                                         }
                                         ipv4::protocol::Protocol::Udp => {
                                             let dest_ip = ipv4.destination_ip();
@@ -258,7 +258,7 @@ impl ChannelDataHandler {
                                             ipv4.set_destination_ip(destination);
                                             ipv4.update_checksum();
                                             ip_proxy_map.udp_proxy_map.insert(SocketAddrV4::new(source, source_port),
-                                                                                   (SocketAddrV4::new(gate_way, 0), SocketAddrV4::new(dest_ip, dest_port)));
+                                                                              (SocketAddrV4::new(gate_way, 0), SocketAddrV4::new(dest_ip, dest_port)));
                                         }
                                         ipv4::protocol::Protocol::Icmp => {
                                             let dest_ip = ipv4.destination_ip();
@@ -329,11 +329,14 @@ impl ChannelDataHandler {
                 if current_ip != new_ip {
                     // ip发生变化
                     log::info!("ip发生变化,old_ip:{:?},new_ip:{:?}",current_ip,new_ip);
-                    let old_netmask = current_device.virtual_netmask;
+                    #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
-                    let old_gateway = current_device.virtual_gateway();
+                        let old_netmask = current_device.virtual_netmask;
+                    #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
+                        let old_gateway = current_device.virtual_gateway();
                     let virtual_ip = Ipv4Addr::from(response.virtual_ip);
                     let virtual_gateway = Ipv4Addr::from(response.virtual_gateway);
                     let virtual_netmask = Ipv4Addr::from(response.virtual_netmask);
+                    #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
                     self.device_writer.change_ip(virtual_ip, virtual_netmask, virtual_gateway, old_netmask, old_gateway)?;
                     let new_current_device = CurrentDeviceInfo::new(virtual_ip, virtual_gateway,
                                                                     virtual_netmask, current_device.connect_server);
@@ -346,7 +349,7 @@ impl ChannelDataHandler {
             service_packet::Protocol::PollDeviceList => {}
             service_packet::Protocol::PushDeviceList => {
                 let device_list_t = DeviceList::parse_from_bytes(net_packet.payload())?;
-                let ip_list = device_list_t
+                let ip_list: Vec<PeerDeviceInfo> = device_list_t
                     .device_info_list
                     .into_iter()
                     .map(|info| {
@@ -357,6 +360,10 @@ impl ChannelDataHandler {
                         )
                     })
                     .collect();
+                let route = Route::from(*route_key, 2, 99);
+                for x in &ip_list {
+                    context.add_route_if_absent(x.virtual_ip, route);
+                }
                 let mut dev = self.device_list.lock();
                 if dev.0 != device_list_t.epoch as u16 {
                     dev.0 = device_list_t.epoch as u16;
@@ -434,7 +441,6 @@ impl ChannelDataHandler {
                 }
             }
             ControlPacket::PunchRequest => {
-                // log::info!("PunchRequest route_key:{:?}",route_key);
                 //回应
                 net_packet.set_transport_protocol(control_packet::Protocol::PunchResponse.into());
                 net_packet.set_source(current_device.virtual_ip());

switch/src/handle/registration_handler.rs

@@ -1,5 +1,5 @@
 use std::io;
-use std::net::SocketAddr;
+use std::net::{Ipv4Addr, SocketAddr};
 use std::time::{Duration, Instant};
 use crossbeam_utils::atomic::AtomicCell;
 
@@ -7,11 +7,28 @@ use protobuf::Message;
 use tokio::net::UdpSocket;
 use crate::channel::sender::ChannelSender;
 
-use crate::error::*;
 use crate::proto::message::{RegistrationRequest, RegistrationResponse};
 use crate::protocol::error_packet::InErrorPacket;
 use crate::protocol::{service_packet, NetPacket, Protocol, Version, MAX_TTL};
 
+pub enum ReqEnum {
+    TokenError,
+    AddressExhausted,
+    Timeout,
+    ServerError(String),
+    Other(String),
+}
+
+#[derive(Clone, Debug)]
+pub struct RegResponse {
+    pub virtual_ip: Ipv4Addr,
+    pub virtual_gateway: Ipv4Addr,
+    pub virtual_netmask: Ipv4Addr,
+    pub epoch: u32,
+    pub public_ip: Ipv4Addr,
+    pub public_port: u16,
+}
+
 ///向中继服务器注册，token标识一个虚拟网关，device_id防止多次注册时得到的ip不一致
 pub async fn registration(
     main_channel: &UdpSocket,
@@ -19,77 +36,90 @@ pub async fn registration(
     token: String,
     device_id: String,
     name: String,
-) -> Result<RegistrationResponse> {
+) -> Result<RegResponse, ReqEnum> {
     let request_packet =
-        registration_request_packet(token.clone(), device_id.clone(), name.clone(), false)?;
+        registration_request_packet(token.clone(), device_id.clone(), name.clone(), false).unwrap();
     let buf = request_packet.buffer();
     let mut recv_buf = [0u8; 10240];
-    let mut count = 0;
+    return match main_channel.send_to(buf, server_address).await {
-    loop {
+        Ok(_) => {
-        match main_channel.send_to(buf, server_address).await {
+            match tokio::time::timeout(Duration::from_millis(300), main_channel.recv_from(&mut recv_buf)).await {
-            Ok(_) => {
+                Ok(rs) => {
-                match tokio::time::timeout(Duration::from_millis(300), main_channel.recv_from(&mut recv_buf)).await {
+                    match rs {
-                    Ok(rs) => {
+                        Ok((len, addr)) => {
-                        match rs {
+                            if server_address == addr {
-                            Ok((len, addr)) => {
+                                let net_packet = match NetPacket::new(&recv_buf[..len]) {
-                                if server_address == addr {
+                                    Ok(net_packet) => {
-                                    let net_packet = NetPacket::new(&recv_buf[..len])?;
+                                        net_packet
-                                    match net_packet.protocol() {
+                                    }
-                                        Protocol::Service => {
+                                    Err(e) => {
-                                            match service_packet::Protocol::from(net_packet.transport_protocol()) {
+                                        return Err(ReqEnum::ServerError(format!("{}",e)))
-                                                service_packet::Protocol::RegistrationResponse => {
+                                    }
-                                                    let response = RegistrationResponse::parse_from_bytes(net_packet.payload())?;
+                                };
-                                                    return Ok(response);
+                                match net_packet.protocol() {
+                                    Protocol::Service => {
+                                        match service_packet::Protocol::from(net_packet.transport_protocol()) {
+                                            service_packet::Protocol::RegistrationResponse => {
+                                                match RegistrationResponse::parse_from_bytes(net_packet.payload()) {
+                                                    Ok(response) => {
+                                                        Ok(RegResponse {
+                                                            virtual_ip: Ipv4Addr::from(response.virtual_ip),
+                                                            virtual_gateway: Ipv4Addr::from(response.virtual_gateway),
+                                                            virtual_netmask: Ipv4Addr::from(response.virtual_netmask),
+                                                            epoch: response.epoch,
+                                                            public_ip: Ipv4Addr::from(response.public_ip),
+                                                            public_port: response.public_port as u16,
+                                                        })
+                                                    }
+                                                    Err(_) => {
+                                                        Err(ReqEnum::ServerError("invalid data".to_string()))
+                                                    }
                                                 }
-                                                _ => println!("响应数据错误"),
+                                            }
+                                            _ => {
+                                                Err(ReqEnum::ServerError("invalid data".to_string()))
                                             }
                                         }
-                                        Protocol::Error => {
+                                    }
-                                            match InErrorPacket::new(net_packet.transport_protocol(), net_packet.payload()) {
+                                    Protocol::Error => {
-                                                Ok(e) => match e {
+                                        match InErrorPacket::new(net_packet.transport_protocol(), net_packet.payload()) {
-                                                    InErrorPacket::TokenError => return Err(Error::Stop("token错误".to_string())),
+                                            Ok(e) => match e {
-                                                    InErrorPacket::Disconnect => {
+                                                InErrorPacket::TokenError => Err(ReqEnum::TokenError),
-                                                        println!("断开连接");
+                                                InErrorPacket::Disconnect => {
+                                                    Err(ReqEnum::ServerError("disconnect".to_string()))
+                                                }
+                                                InErrorPacket::AddressExhausted => {
+                                                    Err(ReqEnum::AddressExhausted)
+                                                }
+                                                InErrorPacket::OtherError(e) => match e.message() {
+                                                    Ok(str) => {
+                                                        Err(ReqEnum::ServerError(str))
                                                     }
-                                                    InErrorPacket::AddressExhausted => {
+                                                    Err(e) => Err(ReqEnum::Other(format!("{}", e))),
-                                                        println!("地址用尽");
-                                                        log::warn!("地址用尽");
-                                                    }
-                                                    InErrorPacket::OtherError(e) => match e.message() {
-                                                        Ok(str) => {
-                                                            println!("其他异常:{:?}", str);
-                                                            log::warn!("其他异常{:?}",str);
-                                                        }
-                                                        Err(e) => println!("其他异常:{:?}", e),
-                                                    },
                                                 },
-                                                Err(e) => println!("数据解析异常:{:?}", e),
+                                            },
-                                            }
+                                            Err(e) => Err(ReqEnum::Other(format!("{}", e))),
                                         }
-                                        _ => println!("响应数据错误"),
+                                    }
-                                    };
+                                    _ => Err(ReqEnum::ServerError("invalid data".to_string())),
                                 }
-                            }
+                            } else {
-                            Err(e) => {
+                                Err(ReqEnum::Other(format!("invalid data,from {}", addr)))
-                                println!("接收服务器数据失败:{:?}", e);
-                                log::warn!("接收服务器数据失败:{:?}",e);
                             }
                         }
-                    }
+                        Err(e) => {
-                    Err(_) => {
+                            Err(ReqEnum::Other(format!("receiver error:{}", e)))
-                        println!("接收超时");
+                        }
-                        log::warn!("接收超时");
                     }
                 }
-            }
+                Err(_) => {
-            Err(e) => {
+                    Err(ReqEnum::Timeout)
-                println!("发送数据到服务器失败:{:?}", e);
+                }
-                log::warn!("发送数据到服务器失败:{:?}",e);
             }
         }
-        count += 1;
+        Err(e) => {
-        println!("重试中(retrying)...");
+            Err(ReqEnum::Other(format!("send error:{}", e)))
-        std::thread::sleep(Duration::from_secs(count % 10 + 1));
+        }
     };
 }
 

switch/src/handle/tun_tap/mod.rs

@@ -18,41 +18,33 @@ use crate::protocol;
 use crate::protocol::ip_turn_packet::BroadcastPacketEnd;
 
 pub mod tun_handler;
+#[cfg(any(target_os = "linux", target_os = "macos",target_os = "windows"))]
 pub mod tap_handler;
 
 async fn broadcast(sender: &ChannelSender, net_packet: &mut NetPacket<&mut [u8]>, data_len: usize, current_device: &CurrentDeviceInfo) -> Result<()> {
     let mut peer_ips = Vec::with_capacity(8);
     let vec = sender.route_table_one();
     let mut relay_count = 0;
-    let mut last_peer = None;
+    const MAX_COUNT: usize = u8::MAX as usize;
     for (peer_ip, route) in vec {
         if peer_ip == current_device.virtual_gateway {
             continue;
         }
-        if peer_ips.len() < u8::MAX as usize && route.is_p2p()
+        if peer_ips.len() == MAX_COUNT {
+            break;
+        }
+        if route.is_p2p()
             && sender.send_by_key(&net_packet.buffer()[..data_len], &route.route_key()).await.is_ok() {
             peer_ips.push(peer_ip);
         } else {
             relay_count += 1;
-            if relay_count == 1 {
-                last_peer = Some((peer_ip, route));
-            }
-            if relay_count > 1 && peer_ips.len() == u8::MAX as usize {
-                break;
-            }
         }
     }
-    if relay_count == 0 && !peer_ips.is_empty() {
+    if relay_count == 0 && !peer_ips.is_empty() && peer_ips.len() != MAX_COUNT {
         //不需要转发
         return Ok(());
     }
-    if relay_count == 1 && !net_packet.is_encrypt() {
+
-        //只有一个目标，并且没加密
-        let (peer_ip, route) = last_peer.unwrap();
-        net_packet.set_destination(peer_ip);
-        sender.send_by_key(&net_packet.buffer()[..data_len], &route.route_key()).await?;
-        return Ok(());
-    }
     if peer_ips.is_empty() {
         sender.send_main(&net_packet.buffer()[..data_len], current_device.connect_server).await?;
     } else {
@@ -71,7 +63,7 @@ async fn multicast(igmp_server: &IgmpServer, multicast_addr: Ipv4Addr, sender: &
     let mut peer_ips = Vec::with_capacity(8);
     let vec = sender.route_table_one();
     let mut relay_count = 0;
-    let mut last_peer = None;
+    const MAX_COUNT: usize = u8::MAX as usize;
     if let Some(members) = igmp_server.load(&multicast_addr) {
         for (peer_ip, route) in vec {
             if peer_ip == current_device.virtual_gateway {
@@ -79,32 +71,22 @@ async fn multicast(igmp_server: &IgmpServer, multicast_addr: Ipv4Addr, sender: &
             }
             let is_send = { members.read().is_send(&peer_ip) };
             if is_send {
-                if peer_ips.len() < u8::MAX as usize && route.is_p2p()
+                if peer_ips.len() == MAX_COUNT {
+                    break;
+                }
+                if route.is_p2p()
                     && sender.send_by_key(&net_packet.buffer()[..data_len], &route.route_key()).await.is_ok() {
                     peer_ips.push(peer_ip);
                 } else {
                     relay_count += 1;
-                    if relay_count == 1 {
-                        last_peer = Some((peer_ip, route));
-                    }
-                    if relay_count > 1 && peer_ips.len() == u8::MAX as usize {
-                        break;
-                    }
                 }
             }
         }
     }
-    if relay_count == 0 && !peer_ips.is_empty() {
+    if relay_count == 0 && !peer_ips.is_empty() && peer_ips.len() != MAX_COUNT {
         //不需要转发
         return Ok(());
     }
-    if relay_count == 1 && !net_packet.is_encrypt() {
-        //只有一个目标，并且没加密
-        let (peer_ip, route) = last_peer.unwrap();
-        net_packet.set_destination(peer_ip);
-        sender.send_by_key(&net_packet.buffer()[..data_len], &route.route_key()).await?;
-        return Ok(());
-    }
     if peer_ips.is_empty() {
         sender.send_main(&net_packet.buffer()[..data_len], current_device.connect_server).await?;
     } else {

switch/src/handle/tun_tap/tap_handler.rs

@@ -10,13 +10,14 @@ use packet::icmp::Kind;
 use packet::ip::ipv4;
 use packet::ip::ipv4::packet::IpV4Packet;
 use crate::channel::sender::ChannelSender;
+use crate::core::status::SwitchWorker;
 use crate::external_route::ExternalRoute;
 use crate::handle::CurrentDeviceInfo;
 use crate::igmp_server::IgmpServer;
 use crate::ip_proxy::IpProxyMap;
 use crate::tun_tap_device::{DeviceReader, DeviceWriter};
 
-pub fn start(sender: ChannelSender,
+pub fn start(worker: SwitchWorker, sender: ChannelSender,
              device_reader: DeviceReader,
              device_writer: DeviceWriter,
              igmp_server: Option<IgmpServer>,
@@ -24,7 +25,7 @@ pub fn start(sender: ChannelSender,
              ip_route: Option<ExternalRoute>,
              ip_proxy_map: Option<IpProxyMap>,
              cipher: Option<Aes256Gcm>) {
-    thread::Builder::new().name("tap-handler".into()).spawn(move || {
+    thread::spawn(move || {
         tokio::runtime::Builder::new_current_thread()
             .enable_all().build().unwrap()
             .block_on(async move {
@@ -33,8 +34,9 @@ pub fn start(sender: ChannelSender,
                                        current_device, ip_route, ip_proxy_map, cipher).await {
                     log::warn!("tap:{:?}",e);
                 }
+                worker.stop_all();
             });
-    }).unwrap();
+    });
 }
 
 async fn start_(sender: ChannelSender,

switch/src/handle/tun_tap/tun_handler.rs

@@ -9,6 +9,7 @@ use packet::icmp::icmp::IcmpPacket;
 use packet::ip::ipv4;
 use packet::ip::ipv4::packet::IpV4Packet;
 use crate::channel::sender::ChannelSender;
+use crate::core::status::SwitchWorker;
 
 use crate::error::*;
 use crate::external_route::ExternalRoute;
@@ -36,7 +37,7 @@ fn icmp(device_writer: &DeviceWriter, mut ipv4_packet: IpV4Packet<&mut [u8]>) ->
 /// 接收tun数据，并且转发到udp上
 #[inline]
 async fn handle(sender: &ChannelSender, data: &mut [u8], len: usize, device_writer: &DeviceWriter, igmp_server: &Option<IgmpServer>, current_device: CurrentDeviceInfo,
-                ip_route: &Option<ExternalRoute>, proxy_map: &Option<IpProxyMap>,cipher: &Option<Aes256Gcm>) -> Result<()> {
+                ip_route: &Option<ExternalRoute>, proxy_map: &Option<IpProxyMap>, cipher: &Option<Aes256Gcm>) -> Result<()> {
     let ipv4_packet = if let Ok(ipv4_packet) = IpV4Packet::new(&mut data[12..len]) {
         ipv4_packet
     } else {
@@ -50,10 +51,10 @@ async fn handle(sender: &ChannelSender, data: &mut [u8], len: usize, device_writ
     if src_ip == dest_ip {
         return icmp(&device_writer, ipv4_packet);
     }
-    return crate::handle::tun_tap::base_handle(sender, data, len, igmp_server, current_device, ip_route, proxy_map,cipher).await;
+    return crate::handle::tun_tap::base_handle(sender, data, len, igmp_server, current_device, ip_route, proxy_map, cipher).await;
 }
 
-pub fn start(sender: ChannelSender,
+pub fn start(worker: SwitchWorker, sender: ChannelSender,
              device_reader: DeviceReader,
              device_writer: DeviceWriter,
              igmp_server: Option<IgmpServer>,
@@ -61,15 +62,16 @@ pub fn start(sender: ChannelSender,
              ip_route: Option<ExternalRoute>,
              ip_proxy_map: Option<IpProxyMap>,
              cipher: Option<Aes256Gcm>) {
-    thread::Builder::new().name("tun-handler".into()).spawn(move || {
+    thread::spawn(move || {
-        tokio::runtime::Builder::new_multi_thread()
+        tokio::runtime::Builder::new_current_thread()
             .enable_all().build().unwrap()
             .block_on(async move {
-                if let Err(e) = start_(sender, device_reader, device_writer, igmp_server, current_device, ip_route, ip_proxy_map,cipher).await {
+                if let Err(e) = start_(sender, device_reader, device_writer, igmp_server, current_device, ip_route, ip_proxy_map, cipher).await {
                     log::warn!("tun:{:?}",e);
                 }
+                worker.stop_all();
             })
-    }).unwrap();
+    });
 }
 
 async fn start_(sender: ChannelSender,
@@ -80,24 +82,17 @@ async fn start_(sender: ChannelSender,
                 ip_route: Option<ExternalRoute>,
                 ip_proxy_map: Option<IpProxyMap>,
                 cipher: Option<Aes256Gcm>) -> io::Result<()> {
+    let mut buf = [0; 4096];
     loop {
-        let mut buf = [0; 4096];
+        if sender.is_close() {
-        let sender = sender.clone();
+            return Ok(());
-        let device_writer = device_writer.clone();
+        }
-        let igmp_server = igmp_server.clone();
-        let ip_route = ip_route.clone();
-        let ip_proxy_map = ip_proxy_map.clone();
-        let cipher = cipher.clone();
         let len = device_reader.read(&mut buf[12..])? + 12;
-        let current_device = current_device.load();
+        match handle(&sender, &mut buf, len, &device_writer, &igmp_server,current_device.load(), &ip_route, &ip_proxy_map, &cipher).await {
-        tokio::spawn(async move {
+            Ok(_) => {}
-            match handle(&sender, &mut buf, len, &device_writer, &igmp_server, current_device, &ip_route, &ip_proxy_map,&cipher).await {
+            Err(e) => {
-                Ok(_) => {}
+                log::warn!("{:?}", e)
-                Err(e) => {
-                    log::warn!("{:?}", e)
-                }
             }
-        });
+        }
-
     }
 }

switch/src/ip_proxy/mod.rs

@@ -45,30 +45,17 @@ impl IpProxyMap {
     }
 }
 
-pub async fn init_proxy(sender: ChannelSender, bind_ips: Vec<Ipv4Addr>, current_device: Arc<AtomicCell<CurrentDeviceInfo>>) -> io::Result<IpProxyMap> {
+pub async fn init_proxy(sender: ChannelSender, bind_ips: Vec<Ipv4Addr>, current_device: Arc<AtomicCell<CurrentDeviceInfo>>) -> io::Result<(TcpProxy, UdpProxy, IpProxyMap)> {
     let mut icmp_sockets = HashMap::new();
     let tcp_proxy_map: Arc<SkipMap<SocketAddrV4, (SocketAddrV4, SocketAddrV4)>> = Arc::new(SkipMap::new());
     let udp_proxy_map: Arc<SkipMap<SocketAddrV4, (SocketAddrV4, SocketAddrV4)>> = Arc::new(SkipMap::new());
     let icmp_proxy_map: Arc<SkipMap<(Ipv4Addr, u16, u16), Ipv4Addr>> = Arc::new(SkipMap::new());
-    let (tcp_proxy_port, udp_proxy_port) = if !bind_ips.is_empty() {
+    let tcp_listener = TcpListener::bind("0.0.0.0:0").await?;
-        let tcp_listener = TcpListener::bind("0.0.0.0:0").await?;
+    let udp_socket = UdpSocket::bind("0.0.0.0:0").await?;
-        let udp_socket = UdpSocket::bind("0.0.0.0:0").await?;
+    let tcp_proxy_port = tcp_listener.local_addr()?.port();
-        let tcp_proxy_port = tcp_listener.local_addr()?.port();
+    let udp_proxy_port = udp_socket.local_addr()?.port();
-        let udp_proxy_port = udp_socket.local_addr()?.port();
+    let tcp_proxy = TcpProxy::new(tcp_listener, tcp_proxy_map.clone());
-        let tcp_proxy_map = tcp_proxy_map.clone();
+    let udp_proxy = UdpProxy::new(udp_socket, udp_proxy_map.clone());
-        tokio::spawn(async {
-            let tcp_proxy = TcpProxy::new(tcp_listener, tcp_proxy_map);
-            tcp_proxy.start().await
-        });
-        let udp_proxy_map = udp_proxy_map.clone();
-        tokio::spawn(async {
-            let udp_proxy = UdpProxy::new(udp_socket, udp_proxy_map);
-            udp_proxy.start().await
-        });
-        (tcp_proxy_port, udp_proxy_port)
-    } else {
-        (0, 0)
-    };
     for ip in bind_ips {
         let addr = SocketAddrV4::new(ip, 0);
         let icmp_proxy_map = icmp_proxy_map.clone();
@@ -79,12 +66,12 @@ pub async fn init_proxy(sender: ChannelSender, bind_ips: Vec<Ipv4Addr>, current_
         });
     }
 
-    Ok(IpProxyMap {
+    Ok((tcp_proxy, udp_proxy, IpProxyMap {
         tcp_proxy_port,
         udp_proxy_port,
         tcp_proxy_map,
         udp_proxy_map,
         icmp_proxy_map,
         icmp_sockets,
-    })
+    }))
 }

switch/src/lib.rs

@@ -13,3 +13,4 @@ pub mod igmp_server;
 pub mod tun_tap_device;
 pub mod core;
 pub mod channel;
+pub mod util;

switch/src/nat/check.rs

@@ -58,7 +58,6 @@ pub fn public_ip_list_(
     udp: &UdpSocket,
     addrs: &Vec<SocketAddr>,
 ) -> io::Result<(HashSet<Ipv4Addr>, u16, u16)> {
-    // println!("local port {:?}", udp.local_addr().unwrap().port());
     udp.set_read_timeout(Some(Duration::from_millis(300)))?;
     let mut buf = [0u8; 128];
     for addr in addrs {

switch/src/tun_tap_device/android.rs

@@ -0,0 +1,42 @@
+use std::io;
+use std::os::unix::io::RawFd;
+
+#[derive(Clone)]
+pub struct DeviceWriter(RawFd);
+
+pub struct DeviceReader(RawFd);
+
+impl DeviceWriter {
+    pub fn write_ipv4_tun(&self, buf: &[u8]) -> io::Result<()> {
+        unsafe {
+            let amount = libc::write(self.0, buf.as_ptr() as *const _, buf.len() );
+            if amount < 0 {
+                return Err(io::Error::last_os_error());
+            }
+            Ok(())
+        }
+    }
+    ///写入ipv4数据，为了兼容其他代码，头部空了14个字节
+    pub fn write_ipv4(&self, buf: &[u8]) -> io::Result<()> {
+        let buf = &buf[14..];
+        self.write_ipv4_tun(buf)
+    }
+}
+
+impl DeviceReader {
+    pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
+        unsafe {
+            let amount = libc::read(self.0, buf.as_mut_ptr() as *mut _, buf.len() );
+
+            if amount < 0 {
+                return Err(io::Error::last_os_error());
+            }
+
+            Ok(amount as usize)
+        }
+    }
+}
+
+pub fn create(fd: i32) -> (DeviceWriter, DeviceReader) {
+    (DeviceWriter(fd as _), DeviceReader(fd as _))
+}

switch/src/tun_tap_device/linux.rs

@@ -1,11 +1,11 @@
 use std::io;
 use std::net::Ipv4Addr;
-use crate::tun_tap_device::{DeviceReader, DeviceType, DeviceWriter};
+use crate::tun_tap_device::{DeviceReader, DeviceType, DeviceWriter, DriverInfo};
 use tun::Device;
 use parking_lot::Mutex;
 use std::process::Command;
 use std::sync::Arc;
-use crate::tun_tap_device::unix::DeviceW;
+use crate::tun_tap_device::linux_mac::DeviceW;
 
 impl DeviceWriter {
     pub fn change_ip(&self, address: Ipv4Addr, netmask: Ipv4Addr,
@@ -56,8 +56,7 @@ pub fn create_device(device_type: DeviceType,
                      netmask: Ipv4Addr,
                      gateway: Ipv4Addr,
                      in_ips: Vec<(Ipv4Addr, Ipv4Addr)>,
-) -> io::Result<(DeviceWriter, DeviceReader)> {
+) -> io::Result<(DeviceWriter, DeviceReader,DriverInfo)> {
-    println!("========网卡配置========");
     let mut config = tun::Configuration::default();
 
     config
@@ -79,7 +78,6 @@ pub fn create_device(device_type: DeviceType,
     let reader = queue.reader();
     let writer = queue.writer();
     let name = dev.name();
-    println!("name:{:?}", name);
     for (address, netmask) in &in_ips {
         add_route(name, *address, *netmask)?;
     }
@@ -111,10 +109,16 @@ pub fn create_device(device_type: DeviceType,
             DeviceW::Tap((writer, mac))
         }
     };
-    println!("========TUN网卡配置========");
+    let driver_info = DriverInfo {
+        device_type,
+        name:name.to_string(),
+        version:String::new(),
+        mac: None,
+    };
     Ok((
         DeviceWriter::new(device_w, Arc::new(Mutex::new(dev)), in_ips, address, packet_information),
         DeviceReader::new(reader),
+        driver_info,
     ))
 }
 

switch/src/tun_tap_device/linux_mac.rs

@@ -6,9 +6,9 @@ use tun::platform::posix::{Reader, Writer};
 use std::net::Ipv4Addr;
 use std::os::unix::io::AsRawFd;
 use crossbeam_utils::atomic::AtomicCell;
-#[cfg(any(target_os = "linux", target_os = "android"))]
+#[cfg(any(target_os = "linux"))]
 use tun::platform::linux::Device;
-#[cfg(any(target_os = "macos", target_os = "ios"))]
+#[cfg(any(target_os = "macos"))]
 use tun::platform::macos::Device;
 use parking_lot::Mutex;
 use packet::ethernet;
@@ -128,7 +128,6 @@ impl DeviceWriter {
     }
 }
 
-#[derive(Clone)]
 pub struct DeviceReader(Reader);
 
 impl DeviceReader {

switch/src/tun_tap_device/mac.rs

@@ -1,11 +1,11 @@
 use std::io;
 use std::net::Ipv4Addr;
-use crate::tun_tap_device::{DeviceReader, DeviceType, DeviceWriter};
+use crate::tun_tap_device::{DeviceReader, DeviceType, DeviceWriter, DriverInfo};
 use tun::Device;
 use parking_lot::Mutex;
 use std::process::Command;
 use std::sync::Arc;
-use crate::tun_tap_device::unix::DeviceW;
+use crate::tun_tap_device::linux_mac::DeviceW;
 
 impl DeviceWriter {
     pub fn change_ip(&self, address: Ipv4Addr, netmask: Ipv4Addr,
@@ -42,14 +42,13 @@ pub fn create_device(device_type: DeviceType,
                      netmask: Ipv4Addr,
                      gateway: Ipv4Addr,
                      in_ips: Vec<(Ipv4Addr, Ipv4Addr)>,
-) -> io::Result<(DeviceWriter, DeviceReader)> {
+) -> io::Result<(DeviceWriter, DeviceReader,DriverInfo)> {
     match device_type {
         DeviceType::Tun => {}
         DeviceType::Tap => {
             unimplemented!()
         }
     }
-    println!("========TUN网卡配置========");
     let mut config = tun::Configuration::default();
 
     config
@@ -74,11 +73,16 @@ pub fn create_device(device_type: DeviceType,
     let queue = dev.queue(0).unwrap();
     let reader = queue.reader();
     let writer = queue.writer();
-    println!("name:{:?}", name);
+    let driver_info = DriverInfo {
-    println!("========TUN网卡配置========");
+        device_type,
+        name:name.to_string(),
+        version:String::new(),
+        mac: None,
+    };
     Ok((
         DeviceWriter::new(DeviceW::Tun(writer), Arc::new(Mutex::new(dev)), in_ips, address, packet_information),
         DeviceReader::new(reader),
+        driver_info
     ))
 }
 

switch/src/tun_tap_device/mod.rs

@@ -1,19 +1,25 @@
 #[cfg(target_os = "windows")]
-pub mod windows;
+mod windows;
-#[cfg(any(target_os = "linux", target_os = "android"))]
+#[cfg(any(target_os = "linux"))]
-pub mod linux;
+mod linux;
 #[cfg(target_os = "macos")]
-pub mod mac;
+mod mac;
-#[cfg(any(unix))]
+#[cfg(any(target_os = "linux", target_os = "macos"))]
-pub mod unix;
+mod linux_mac;
+#[cfg(target_os = "android")]
+mod android;
 
-#[cfg(any(target_os = "linux", target_os = "android"))]
+
+#[cfg(any(target_os = "linux"))]
 pub use linux::create_device;
-#[cfg(any(target_os = "linux", target_os = "android"))]
+#[cfg(any(target_os = "linux"))]
 pub use linux::delete_device;
-#[cfg(any(unix))]
+#[cfg(target_os = "android")]
-pub use unix::{DeviceWriter, DeviceReader};
+pub use android::create;
-
+#[cfg(any(target_os = "linux", target_os = "macos"))]
+pub use linux_mac::{DeviceWriter, DeviceReader};
+#[cfg(target_os = "android")]
+pub use android::{DeviceWriter, DeviceReader};
 #[cfg(target_os = "macos")]
 pub use mac::create_device;
 #[cfg(target_os = "macos")]
@@ -26,7 +32,21 @@ pub use windows::delete_device;
 #[cfg(target_os = "windows")]
 pub use windows::{DeviceWriter, DeviceReader};
 
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
 pub enum DeviceType {
     Tun,
     Tap,
+}
+
+impl DeviceType {
+    pub fn is_tun(&self) -> bool {
+        *self == DeviceType::Tun
+    }
+}
+
+pub struct DriverInfo {
+    pub device_type: DeviceType,
+    pub name: String,
+    pub version: String,
+    pub mac: Option<String>,
 }

switch/src/tun_tap_device/windows.rs

@@ -8,7 +8,7 @@ use parking_lot::Mutex;
 use packet::ethernet;
 use packet::ethernet::packet::EthernetPacket;
 use win_tun_tap::{IFace, TapDevice, TunDevice};
-use crate::tun_tap_device::DeviceType;
+use crate::tun_tap_device::{DriverInfo, DeviceType};
 
 pub const TUN_INTERFACE_NAME: &str = "Switch-Tun-V1";
 pub const TUN_POOL_NAME: &str = "Switch-Tun-V1";
@@ -161,7 +161,6 @@ fn dest(ip: Ipv4Addr, mask: Ipv4Addr) -> Ipv4Addr {
     ])
 }
 
-#[derive(Clone)]
 pub struct DeviceReader {
     device: Arc<Device>,
 }
@@ -199,9 +198,8 @@ fn create_tun(
     netmask: Ipv4Addr,
     gateway: Ipv4Addr,
     in_ips: Vec<(Ipv4Addr, Ipv4Addr)>,
-) -> io::Result<(DeviceWriter, DeviceReader)> {
+) -> io::Result<(DeviceWriter, DeviceReader, DriverInfo)> {
     unsafe {
-        println!("========TUN网卡配置========");
         match Library::new("wintun.dll") {
             Ok(lib) => match TunDevice::delete_for_name(lib, TUN_INTERFACE_NAME) {
                 Ok(_) => {
@@ -210,7 +208,6 @@ fn create_tun(
                 Err(_) => {}
             },
             Err(e) => {
-                log::error!("wintun.dll not found");
                 return Err(io::Error::new(
                     io::ErrorKind::Other,
                     format!("wintun.dll not found {:?}", e),
@@ -240,8 +237,8 @@ fn create_tun(
                 }
             }
         };
-        println!("name:{:?}", tun_device.get_name()?);
+        let name = tun_device.get_name()?;
-        println!("version:{:?}", tun_device.version()?);
+        let version = format!("{:?}", tun_device.version()?);
         tun_device.set_ip(address, netmask)?;
         tun_device.set_metric(1)?;
         tun_device.set_mtu(1420)?;
@@ -256,14 +253,21 @@ fn create_tun(
         tun_device.add_route(Ipv4Addr::from([224, 0, 0, 0]), Ipv4Addr::from([240, 0, 0, 0]), gateway, 1)?;
         delete_cache();
         let device = Arc::new(Device::Tun(tun_device));
-        println!("========TUN网卡配置========");
+        let driver_info = DriverInfo {
+            device_type: DeviceType::Tun,
+            name,
+            version,
+            mac: None,
+        };
         Ok((
             DeviceWriter::new(device.clone(), in_ips, address),
             DeviceReader::new(device),
+            driver_info
         ))
     }
 }
-fn delete_cache(){
+
+fn delete_cache() {
     //清除路由缓存
     let delete_cache = "netsh interface ip delete destinationcache";
     let out = std::process::Command::new("cmd")
@@ -271,7 +275,7 @@ fn delete_cache(){
         .arg(delete_cache)
         .output()
         .unwrap();
-    if !out.status.success(){
+    if !out.status.success() {
         log::warn!("删除缓存失败:{:?}",out);
     }
 }
@@ -293,8 +297,7 @@ fn create_tap(
     netmask: Ipv4Addr,
     gateway: Ipv4Addr,
     in_ips: Vec<(Ipv4Addr, Ipv4Addr)>,
-) -> io::Result<(DeviceWriter, DeviceReader)> {
+) -> io::Result<(DeviceWriter, DeviceReader, DriverInfo)> {
-    println!("========TAP网卡配置========");
     let tap_device = match TapDevice::open(TAP_INTERFACE_NAME) {
         Ok(tap_device) => tap_device,
         Err(e) => {
@@ -305,9 +308,9 @@ fn create_tap(
         }
     };
     let mac = tap_device.get_mac()?;
-    println!("name:{:?}", tap_device.get_name()?);
+    let name = tap_device.get_name()?;
-    println!("version:{:x?}", tap_device.get_version()?);
+    let version = format!("{:?}", tap_device.get_version()?);
-    println!("mac:{:x?}", mac);
+    let mac_str = format!("mac:{:x?}", mac);
     tap_device.set_ip(address, netmask)?;
     tap_device.set_metric(1)?;
     tap_device.set_mtu(1420)?;
@@ -321,10 +324,16 @@ fn create_tap(
     tap_device.add_route(Ipv4Addr::from([224, 0, 0, 0]), Ipv4Addr::from([240, 0, 0, 0]), gateway, 1)?;
     delete_cache();
     let tap = Arc::new(Device::Tap((tap_device, mac)));
-    println!("========TAP网卡配置========");
+    let driver_info = DriverInfo {
+        device_type: DeviceType::Tap,
+        name,
+        version,
+        mac: Some(mac_str),
+    };
     Ok((
         DeviceWriter::new(tap.clone(), in_ips, address),
-        DeviceReader::new(tap)
+        DeviceReader::new(tap),
+        driver_info
     ))
 }
 
@@ -341,7 +350,7 @@ fn delete_tap() {
 pub fn create_device(device_type: DeviceType, address: Ipv4Addr,
                      netmask: Ipv4Addr,
                      gateway: Ipv4Addr,
-                     in_ips: Vec<(Ipv4Addr, Ipv4Addr)>, ) -> io::Result<(DeviceWriter, DeviceReader)> {
+                     in_ips: Vec<(Ipv4Addr, Ipv4Addr)>, ) -> io::Result<(DeviceWriter, DeviceReader, DriverInfo)> {
     match device_type {
         DeviceType::Tun => {
             create_tun(address, netmask, gateway, in_ips)

switch/src/util/mod.rs

@@ -0,0 +1 @@
+pub mod wait;

switch/src/util/wait.rs

@@ -0,0 +1,44 @@
+use std::sync::Arc;
+use std::sync::atomic::{AtomicIsize, Ordering};
+use tokio::sync::watch::{channel, Receiver, Sender};
+
+#[derive(Clone)]
+pub struct WaitGroup {
+    count: Arc<AtomicIsize>,
+    receiver: Receiver<usize>,
+    sender: Arc<Sender<usize>>,
+}
+
+impl WaitGroup {
+    pub fn new() -> Self {
+        let (sender, receiver) = channel(1);
+        Self {
+            count: Arc::new(Default::default()),
+            receiver,
+            sender: Arc::new(sender),
+        }
+    }
+    pub fn add(&self) {
+        let _ = self.count.fetch_add(1, Ordering::Relaxed);
+    }
+    pub fn done(&self) {
+        let i = self.count.fetch_sub(1, Ordering::Relaxed);
+        if i == 1 {
+            let _ = self.sender.send(0);
+        }
+    }
+    pub async fn wait(&mut self) {
+        loop {
+            if 0 == *self.receiver.borrow() {
+                return;
+            }
+            if self.receiver.changed().await.is_ok() {
+                if 0 == *self.receiver.borrow() {
+                    return;
+                }
+            } else {
+                return;
+            }
+        }
+    }
+}