vnt/switch/src/core/mod.rs

use std::io;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::sync::Arc;
use std::time::Duration;

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::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::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::{DeviceReader, DeviceWriter};

pub mod status;
pub mod sync;


#[derive(Clone)]
pub struct Switch {
    name: String,
    current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
    context: Context,
    switch_status_manager: SwitchStatusManger,
    device_writer: DeviceWriter,
    /// 0. 机器纪元，每一次上线或者下线都会增1，用于感知网络中机器变化
    /// 服务端和客户端的不一致，则服务端会推送新的设备列表
    /// 1. 网络中的虚拟ip列表
    device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>>,
    nat_test: NatTest,
    connect_status: Arc<AtomicCell<ConnectStatus>>,
    peer_nat_info_map: Arc<SkipMap<Ipv4Addr, NatInfo>>,
}

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 mtu = self.config.mtu.unwrap_or(1430);
        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, mtu)?;
        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 (cone_sender, cone_receiver) = channel(3);
        let (symmetric_sender, symmetric_receiver) = channel(2);
        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 device_list: Arc<Mutex<(u16, Vec<PeerDeviceInfo>)>> = Arc::new(Mutex::new((response.epoch, response.device_info_list)));
        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 = 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(), 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);
        let in_external_route = if config.in_ips.is_empty() {
            None
        } else {
            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 (tcp_proxy, udp_proxy, ip_proxy_map) = if out_ips.is_empty() {
            (None, None, None)
        } else {
            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("tap_handler"), 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("tun_handler"), 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("android tun_handler"), 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,
                                                           device_writer.clone(), connect_status.clone(),
                                                           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("channel_worker");
            if let Some(tcp_proxy) = tcp_proxy {
                tokio::spawn(tcp_proxy.start());
            }
            if let Some(udp_proxy) = udp_proxy {
                tokio::spawn(udp_proxy.start());
            }
            tokio::spawn(async move {
                channel.start(channel_worker, 14, 65).await
            });
        }
        {
            let other_worker = switch_status_manager.worker("punch_handler");
            let nat_test = nat_test.clone();
            let device_list = device_list.clone();
            let current_device = current_device.clone();
            // 定时心跳
            heartbeat_handler::start_heartbeat(other_worker.worker("heartbeat"), channel_sender.clone(), device_list.clone(), current_device.clone(), config.server_address_str);
            // 空闲检查
            heartbeat_handler::start_idle(other_worker.worker("idle"), idle, channel_sender.clone());
            // 打洞处理
            punch_handler::start(other_worker.worker("cone_receiver"), cone_receiver, punch.clone(), current_device.clone());
            punch_handler::start(other_worker.worker("symmetric_receiver"), symmetric_receiver, punch, current_device.clone());
            tokio::spawn(punch_handler::start_punch(other_worker, nat_test,
                                                    device_list, channel_sender, current_device));
        }
        context.switch(nat_test.nat_info().nat_type);
        Ok(Switch {
            name: config.name,
            current_device,
            context,
            switch_status_manager,
            device_writer,
            nat_test,
            device_list,
            connect_status,
            peer_nat_info_map,
        })
    }
}

impl Switch {
    pub fn name(&self) -> &str {
        &self.name
    }
    pub fn current_device(&self) -> CurrentDeviceInfo {
        self.current_device.load()
    }
    pub fn peer_nat_info(&self, ip: &Ipv4Addr) -> Option<NatInfo> {
        self.peer_nat_info_map.get(ip).map(|e| e.value().clone())
    }
    pub fn connection_status(&self) -> ConnectStatus {
        self.connect_status.load()
    }
    pub fn nat_info(&self) -> NatInfo {
        self.nat_test.nat_info()
    }
    pub fn device_list(&self) -> Vec<PeerDeviceInfo> {
        let device_list_lock = self.device_list.lock();
        let (_epoch, device_list) = device_list_lock.clone();
        drop(device_list_lock);
        device_list
    }
    pub fn route(&self, ip: &Ipv4Addr) -> Option<Route> {
        self.context.route_one(ip)
    }
    pub fn route_key(&self, route_key: &RouteKey) -> Option<Ipv4Addr> {
        self.context.route_to_id(route_key)
    }
    pub fn route_table(&self) -> Vec<(Ipv4Addr, Route)> {
        self.context.route_table_one()
    }
    pub fn stop(&self) -> io::Result<()> {
        self.context.close();
        self.switch_status_manager.stop_all();
        self.device_writer.close()?;
        let virtual_gateway = self.current_device.load().virtual_gateway;
        let _ = std::net::UdpSocket::bind("0.0.0.0:0")?.send_to(&[0],
                                                                SocketAddr::V4(SocketAddrV4::new(virtual_gateway, 10000)));
        Ok(())
    }
    pub async fn wait_stop(&mut self) {
        self.switch_status_manager.wait().await;
        let _ = self.stop();
    }
    pub async fn wait_stop_ms(&mut self, ms: Duration) -> bool {
        tokio::select! {
            _=self.switch_status_manager.wait()=>{
                let _ = self.stop();
                return true;
            }
            _=tokio::time::sleep(ms)=>{
                return false;
            }
        }
    }
}

impl Drop for Switch {
    fn drop(&mut self) {
        let _ = self.stop();
    }
}

#[derive(Clone, Debug)]
pub struct Config {
    pub tap: bool,
    pub token: String,
    pub device_id: String,
    pub name: String,
    pub server_address: SocketAddr,
    pub server_address_str: String,
    pub nat_test_server: Vec<SocketAddr>,
    pub in_ips: Vec<(u32, u32, Ipv4Addr)>,
    pub out_ips: Vec<(u32, u32, Ipv4Addr)>,
    pub key: Option<[u8; 32]>,
    pub simulate_multicast: bool,
    pub mtu: Option<u16>,
}


impl Config {
    pub fn new(tap: bool, token: String,
               device_id: String,
               name: String,
               server_address: SocketAddr,
               server_address_str: String,
               nat_test_server: Vec<SocketAddr>,
               in_ips: Vec<(u32, u32, Ipv4Addr)>, out_ips: Vec<(u32, u32, Ipv4Addr)>,
               password: Option<String>, simulate_multicast: bool, mtu: Option<u16>, ) -> Self {
        let key = if let Some(password) = password {
            let mut hasher = sha2::Sha256::new();
            hasher.update(password.as_bytes());
            let key: [u8; 32] = hasher.finalize().into();
            Some(key)
        } else {
            None
        };
        Self {
            tap,
            token,
            device_id,
            name,
            server_address,
            server_address_str,
            nat_test_server,
            in_ips,
            out_ips,
            key,
            simulate_multicast,
            mtu,
        }
    }
}