From d4f941d4e72888c4ffaec7369c40e2f885bc3df3 Mon Sep 17 00:00:00 2001 From: lubeilin <1791778603@qq.com> Date: Wed, 4 Jan 2023 22:09:59 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=91=BD=E4=BB=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.toml | 18 +----- src/handle/heartbeat_handler.rs | 14 ++--- src/handle/mod.rs | 20 +++---- src/handle/registration_handler.rs | 4 +- src/handle/tun_handler.rs | 2 +- src/handle/udp_recv_handler.rs | 4 +- src/main.rs | 89 +++++++++++++++++++++++++----- src/protocol/control_packet.rs | 18 ++++++ src/tun_device/windows.rs | 1 + 9 files changed, 119 insertions(+), 51 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 78a54ce..d245643 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "virtual_network" +name = "switch" version = "0.1.0" edition = "2021" @@ -18,7 +18,7 @@ parking_lot = "0.12.1" rsa = "0.7.2" rand = "0.8.5" sha2 = { version = "0.10.6", features = ["oid"] } -colored = "2.0.0" +#colored = "2.0.0" thiserror = "1.0.37" chrono = "0.4.23" @@ -29,7 +29,6 @@ protobuf = "3.2.0" console = "0.15.2" mac_address = "1.1.4" clap = { version = "4.0.32", features = ["derive"] } -#clap_derive = "4.0.21" [target.'cfg(any(unix))'.dependencies] tun = { version = "0.5" } sudo = "0.6.0" @@ -40,19 +39,6 @@ wintun = "0.2.1" libloading = "0.7.4" runas = "0.2.1" -#[dependencies.windows] -# -#libloading = "0.7.4" -# -#version = "0.37.0" -#features = [ -# "alloc", -# "Win32_Foundation", -# "Win32_NetworkManagement_IpHelper", -# "Win32_Networking_WinSock", -# "Win32_UI_WindowsAndMessaging", -# "Win32_System_IO" -#] [build-dependencies] protobuf-codegen = "3.2.0" protoc-bin-vendored = "3.0.0" \ No newline at end of file diff --git a/src/handle/heartbeat_handler.rs b/src/handle/heartbeat_handler.rs index a97e1bc..614ce5a 100644 --- a/src/handle/heartbeat_handler.rs +++ b/src/handle/heartbeat_handler.rs @@ -4,15 +4,15 @@ use std::time::Duration; use chrono::Local; +use crate::DEVICE_LIST; use crate::error::*; use crate::handle::DIRECT_ROUTE_TABLE; -use crate::protocol::control_packet::PingPacket; use crate::protocol::{control_packet, NetPacket, Protocol, Version}; -use crate::DEVICE_LIST; +use crate::protocol::control_packet::PingPacket; pub fn handle_loop(udp: UdpSocket, server_addr: SocketAddr) -> Result<()> { const INTERVAL: u64 = 3000; - const MAX_INTERVAL: i64 = 3000 * 5; + const MAX_INTERVAL: i64 = 3000 * 3; let mut buf = [0u8; (4 + 8 + 4)]; let mut net_packet = NetPacket::new(&mut buf)?; net_packet.set_version(Version::V1); @@ -20,23 +20,23 @@ pub fn handle_loop(udp: UdpSocket, server_addr: SocketAddr) -> Result<()> { net_packet.set_transport_protocol(control_packet::Protocol::Ping.into()); net_packet.set_ttl(255); loop { - let current_time = Local::now().timestamp(); + let current_time = Local::now().timestamp_millis(); { let mut ping = PingPacket::new(net_packet.payload_mut())?; ping.set_time(current_time); let epoch = { DEVICE_LIST.lock().0 }; ping.set_epoch(epoch); } - udp.send_to(net_packet.buffer(), server_addr)?; + let _ = udp.send_to(net_packet.buffer(), server_addr); for x in DIRECT_ROUTE_TABLE.iter() { let virtual_ip = x.key().clone(); let route = x.value().clone(); drop(x); if current_time - route.recv_time <= MAX_INTERVAL { - udp.send_to(net_packet.buffer(), route.address)?; + let _ = udp.send_to(net_packet.buffer(), route.address); } else { DIRECT_ROUTE_TABLE.remove_if(&virtual_ip, |_, route| { - current_time - route.recv_time > MAX_INTERVAL + current_time - route.recv_time <= MAX_INTERVAL }); } } diff --git a/src/handle/mod.rs b/src/handle/mod.rs index 7554c6b..e59b538 100644 --- a/src/handle/mod.rs +++ b/src/handle/mod.rs @@ -82,15 +82,15 @@ pub fn init_nat_info(public_ip: u32, public_port: u16) { #[derive(Clone, Debug)] pub struct CurrentDeviceInfo { - virtual_ip: Ipv4Addr, - virtual_gateway: Ipv4Addr, - virtual_netmask: Ipv4Addr, + pub(crate) virtual_ip: Ipv4Addr, + pub(crate) virtual_gateway: Ipv4Addr, + pub(crate) virtual_netmask: Ipv4Addr, //网络地址 - virtual_network: Ipv4Addr, + pub(crate) virtual_network: Ipv4Addr, //直接广播地址 - broadcast_address: Ipv4Addr, + pub(crate) broadcast_address: Ipv4Addr, //链接的服务器地址 - connect_server: SocketAddr, + pub(crate) connect_server: SocketAddr, } impl CurrentDeviceInfo { @@ -114,11 +114,11 @@ impl CurrentDeviceInfo { #[derive(Clone,Debug)] pub struct Route { - address: SocketAddr, + pub(crate) address: SocketAddr, //用心跳探测延迟，收包时更新 - delay: i64, + pub(crate) delay: i64, //收包时更新，如果太久没有收到消息则剔除 - recv_time: i64, + pub(crate) recv_time: i64, } impl Route { @@ -126,7 +126,7 @@ impl Route { Self { address, delay: -1, - recv_time: Local::now().timestamp(), + recv_time: Local::now().timestamp_millis(), } } } diff --git a/src/handle/registration_handler.rs b/src/handle/registration_handler.rs index b0d195e..707f9e1 100644 --- a/src/handle/registration_handler.rs +++ b/src/handle/registration_handler.rs @@ -91,7 +91,7 @@ fn registration_request_packet(token: String, mac_address: String) -> Result Result<()> { let last = REGISTRATION_TIME.load(Ordering::Relaxed); - let new = Local::now().timestamp(); + let new = Local::now().timestamp_millis(); if new - last < 2000 || REGISTRATION_TIME .compare_exchange(last, new, Ordering::Relaxed, Ordering::Relaxed) @@ -106,7 +106,7 @@ pub fn fast_registration(udp: &UdpSocket, server_address: SocketAddr) -> Result< if let Some((token, mac_address)) = option { let request_packet = registration_request_packet(token, mac_address)?; udp.send_to(request_packet.buffer(), server_address)?; - REGISTRATION_TIME.store(Local::now().timestamp(), Ordering::Relaxed); + REGISTRATION_TIME.store(Local::now().timestamp_millis(), Ordering::Relaxed); return Ok(()); } return Err(Error::Stop("注册信息不存在".to_string())); diff --git a/src/handle/tun_handler.rs b/src/handle/tun_handler.rs index 140f8d1..99d8cc0 100644 --- a/src/handle/tun_handler.rs +++ b/src/handle/tun_handler.rs @@ -71,7 +71,7 @@ fn handle( ipv4_turn_packet.set_payload(ipv4_packet.buffer); //优先发到直连到地址 if let Some(route) = DIRECT_ROUTE_TABLE.get(&dest_ip) { - let current_time = Local::now().timestamp(); + let current_time = Local::now().timestamp_millis(); if current_time - route.recv_time < 3_000 { udp.send_to(&net_packet.buffer()[..(4 + 8 + data_len)], route.address)?; return Ok(()); diff --git a/src/handle/udp_recv_handler.rs b/src/handle/udp_recv_handler.rs index d7d2831..d548c77 100644 --- a/src/handle/udp_recv_handler.rs +++ b/src/handle/udp_recv_handler.rs @@ -200,12 +200,12 @@ fn other_handle( } Protocol::Control => { match ControlPacket::new(net_packet.transport_protocol(), net_packet.payload())? { - ControlPacket::PingPacket(_) => { + ControlPacket::PingPacket(ping) => { net_packet.set_transport_protocol(control_packet::Protocol::Pong.into()); udp.send_to(&net_packet.buffer()[..12], peer_addr)?; } ControlPacket::PongPacket(pong_packet) => { - let current_time = Local::now().timestamp(); + let current_time = Local::now().timestamp_millis(); let rt = current_time - pong_packet.time(); if rt >= 0 { if peer_addr == server_addr { diff --git a/src/main.rs b/src/main.rs index c28c912..7d06b16 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,11 @@ use std::{io, thread}; use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket}; +use std::sync::atomic::Ordering; use clap::Parser; use console::style; -use crossbeam::sync::Parker; -use crate::handle::{CurrentDeviceInfo, DEVICE_LIST, NAT_INFO, NatInfo}; +use crate::handle::{CurrentDeviceInfo, DEVICE_LIST, DIRECT_ROUTE_TABLE, NAT_INFO, NatInfo, SERVER_RT}; use crate::handle::registration_handler::registration; use crate::tun_device::create_tun; @@ -96,16 +96,14 @@ fn main() { println!("virtual_gateway:{:?}", virtual_gateway); println!("virtual_netmask:{:?}", virtual_netmask); println!("当前设备ip(virtual_ip):{}", style(virtual_ip).green()); - let parker = Parker::new(); //心跳线程 { - let un_parker = parker.unparker().clone(); let udp = udp.try_clone().unwrap(); let _ = thread::spawn(move || { if let Err(e) = handle::heartbeat_handler::handle_loop(udp, server_address) { println!("心跳线程停止:{:?}", e); } - un_parker.unpark(); + std::process::exit(1); }); } //初始化nat数据 @@ -117,7 +115,6 @@ fn main() { let (punch_sender, cone_receiver, req_symmetric_receiver, res_symmetric_receiver) = handle::punch_handler::bounded(); //udp数据处理 { - let un_parker = parker.unparker().clone(); // 低优先级的udp数据通道 let (sender, receiver) = crossbeam::channel::bounded(100); let udp1 = udp.try_clone().unwrap(); @@ -132,16 +129,15 @@ fn main() { ) { println!("udp数据处理线程停止:{:?}", e); } - un_parker.unpark(); + std::process::exit(1); }); let udp1 = udp.try_clone().unwrap(); - let un_parker = parker.unparker().clone(); let _ = thread::spawn(move || { let current_device = CurrentDeviceInfo::new(virtual_ip, virtual_gateway, virtual_netmask, server_address); if let Err(e) = handle::udp_recv_handler::other_loop(udp1, receiver, current_device, punch_sender) { println!("udp数据处理线程停止:{:?}", e); } - un_parker.unpark(); + std::process::exit(1); }); } //打洞处理 @@ -171,15 +167,82 @@ fn main() { //tun数据处理 { let udp = udp.try_clone().unwrap(); - let un_parker = parker.unparker().clone(); let _ = thread::spawn(move || { let current_device = CurrentDeviceInfo::new(virtual_ip, virtual_gateway, virtual_netmask, server_address); if let Err(e) = handle::tun_handler::handle_loop(udp, tun_reader, current_device) { println!("tun数据处理线程停止:{:?}", e); } - un_parker.unpark(); + std::process::exit(1); }); } - parker.park(); - std::process::exit(1); + use console::Term; + let term = Term::stdout(); + let current_device = CurrentDeviceInfo::new(virtual_ip, virtual_gateway, virtual_netmask, server_address); + loop { + println!("{}", style("Please enter the command (Usage: list,status,exit,help):").color256(102)); + match term.read_line() { + Ok(cmd) => { + command(cmd.trim(), ¤t_device); + } + Err(e) => { + println!("read_line:{:?}", e); + std::process::exit(1); + } + } + } +} + +fn command(cmd: &str, current_device: &CurrentDeviceInfo) { + match cmd { + "list" => { + let server_delay = SERVER_RT.load(Ordering::Relaxed); + let device_list_lock = DEVICE_LIST.lock(); + let (_epoch, device_list) = device_list_lock.clone(); + drop(device_list_lock); + if device_list.is_empty() { + println!("No other devices found"); + return; + } + for ip in device_list { + if let Some(route_ref) = DIRECT_ROUTE_TABLE.get(&ip) { + let str = if route_ref.value().delay >= 0 { + format!("{}(p2p delay:{}ms)", ip, route_ref.value().delay) + } else { + format!("{}(p2p)", ip) + }; + drop(route_ref); + println!("{}", style(str).green()); + } else { + let str = if server_delay >= 0 { + format!("{}(relay delay:{}ms)", ip, server_delay * 2) + } else { + format!("{}(relay)", ip) + }; + println!("{}", style(str).blue()); + } + } + } + "status" => { + let server_delay = SERVER_RT.load(Ordering::Relaxed); + println!("Virtual ip:{}", style(current_device.virtual_ip).green()); + println!("Virtual gateway:{}", style(current_device.virtual_gateway).green()); + println!("Relay server :{}", style(current_device.connect_server).green()); + if server_delay >= 0 { + println!("Delay of relay server :{}", style(server_delay).green()); + } + } + "help" | "h" => { + println!("Options: "); + println!("{} , Query the virtual IP of other devices", style("list").green()); + println!("{} , View current device status", style("status").green()); + println!("{} , Exit the program", style("exit").green()); + } + "exit" => { + std::process::exit(1); + } + _ => { + println!("command {} not fount. ", style(cmd).red()); + println!("Try to enter: '{}'", style("help").green()); + } + } } diff --git a/src/protocol/control_packet.rs b/src/protocol/control_packet.rs index d1af643..499ec26 100644 --- a/src/protocol/control_packet.rs +++ b/src/protocol/control_packet.rs @@ -70,6 +70,7 @@ pub struct PongPacket { buffer: B, } + impl> PingPacket { pub fn new(buffer: B) -> Result> { let len = buffer.as_ref().len(); @@ -98,6 +99,15 @@ impl + AsMut<[u8]>> PingPacket { } } +impl> fmt::Debug for PingPacket { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PingPacket") + .field("time", &self.time()) + .field("epoch", &self.epoch()) + .finish() + } +} + impl> PongPacket { pub fn new(buffer: B) -> Result> { let len = buffer.as_ref().len(); @@ -120,6 +130,14 @@ impl + AsMut<[u8]>> PongPacket { } } +impl> fmt::Debug for PongPacket { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PongPacket") + .field("time", &self.time()) + .finish() + } +} + pub type TurnPongPacket = TurnPingPacket; /// 探测目标延迟 diff --git a/src/tun_device/windows.rs b/src/tun_device/windows.rs index 2e59f11..524d79d 100644 --- a/src/tun_device/windows.rs +++ b/src/tun_device/windows.rs @@ -51,6 +51,7 @@ pub fn create_tun( } }, Err(e) => { + println!("{}", console::style("wintun.dll not found").red()); return Err(Error::Stop(format!("{:?}", e))); } }