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

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 async fn start(config: Config) -> crate::Result<Switch> {
-        log::info!("config:{:?}",config);
+pub struct SwitchUtil {
+    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_gateway = Ipv4Addr::from(response.virtual_gateway);
-        let virtual_netmask = Ipv4Addr::from(response.virtual_netmask);
+        let virtual_ip = response.virtual_ip;
+        let virtual_gateway = response.virtual_gateway;
+        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 in_ips = config.in_ips.iter().map(|(dest, mask, _)| { (Ipv4Addr::from(*dest & *mask), Ipv4Addr::from(*mask)) }).collect::<Vec<(Ipv4Addr, Ipv4Addr)>>();
+        let nat_test = NatTest::new(config.nat_test_server.clone(), response.public_ip, response.public_port, local_ip, local_port);
 
         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(){
-            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)
+        let (tcp_proxy, udp_proxy, ip_proxy_map) = if out_ips.is_empty() {
+            (None, None, None)
         } else {
-            #[cfg(windows)]
-            {
-                //删除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 (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))
         };
+
+        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
+    }
+}