diff --git a/vnt-cli/src/main.rs b/vnt-cli/src/main.rs index 6cfa7b2..4cd09e6 100644 --- a/vnt-cli/src/main.rs +++ b/vnt-cli/src/main.rs @@ -27,13 +27,7 @@ pub fn app_home() -> io::Result { Ok(path) } -#[tokio::main] -async fn main() { - main0().await; - std::process::exit(0); -} - -async fn main0() { +fn main() { let _ = log4rs::init_file("log4rs.yaml", Default::default()); let args: Vec = std::env::args().collect(); let program = args[0].clone(); @@ -52,8 +46,10 @@ async fn main0() { opts.optflag("m", "", "模拟组播"); opts.optopt("u", "", "自定义mtu(默认为1430)", ""); opts.optflag("", "tcp", "tcp"); - opts.optopt("", "ip", "指定虚拟ip", ""); + opts.optopt("", "ip", "指定虚拟ip", ""); opts.optflag("", "relay", "仅使用服务器转发"); + opts.optopt("", "par", "任务并行度(必须为正整数)", ""); + opts.optopt("", "thread", "线程数(必须为正整数)", ""); //"后台运行时,查看其他设备列表" opts.optflag("", "list", "后台运行时,查看其他设备列表"); opts.optflag("", "all", "后台运行时,查看其他设备完整信息"); @@ -199,13 +195,30 @@ async fn main0() { } let tcp_channel = matches.opt_present("tcp"); let relay = matches.opt_present("relay"); + let parallel = matches.opt_get::("par").unwrap().unwrap_or(2); + if parallel == 0 { + println!("--par invalid"); + return; + } + let thread_num = matches.opt_get::("thread").unwrap().unwrap_or(std::thread::available_parallelism().unwrap().get() * 2); + if thread_num == 0 { + println!("--thread invalid"); + return; + } println!("version 1.2.0"); let config = Config::new(tap, token, device_id, name, server_address, server_address_str, stun_server, in_ip, out_ip, password, simulate_multicast, mtu, - tcp_channel, virtual_ip, relay, server_encrypt); + tcp_channel, virtual_ip, relay, server_encrypt, parallel); + let runtime = tokio::runtime::Builder::new_multi_thread().enable_all().worker_threads(thread_num).build().unwrap(); + runtime.block_on(main0(config,!unused_cmd)); + std::process::exit(0); +} + +async fn main0(config: Config,show_cmd:bool) { + let server_encrypt = config.server_encrypt; let mut vnt_util = VntUtil::new(config).await.unwrap(); let mut conn_count = 0; let response = loop { @@ -320,7 +333,7 @@ async fn main0() { println!("command error :{}", e); } }); - if !unused_cmd { + if show_cmd { let stdin = tokio::io::stdin(); let mut cmd = String::new(); let mut reader = BufReader::new(stdin); @@ -437,8 +450,10 @@ fn print_usage(program: &str, _opts: Options) { println!(" -m 模拟组播,默认情况下组播数据会被当作广播发送,开启后会模拟真实组播的数据发送"); println!(" -u 自定义mtu(不加密默认为1430，加密默认为1410)"); println!(" --tcp 和服务端使用tcp通信,默认使用udp,遇到udp qos时可指定使用tcp"); - println!(" --ip 指定虚拟ip,指定的ip不能和其他设备重复,必须有效并且在服务端所属网段下,默认情况由服务端分配"); + println!(" --ip 指定虚拟ip,指定的ip不能和其他设备重复,必须有效并且在服务端所属网段下,默认情况由服务端分配"); println!(" --relay 仅使用服务器转发,不使用p2p,默认情况允许使用p2p"); + println!(" --par 任务并行度(必须为正整数),默认值为2"); + println!(" --thread 线程数(必须为正整数),默认为核心数乘2"); println!(); println!(" --list {}", yellow("后台运行时,查看其他设备列表".to_string())); println!(" --all {}", yellow("后台运行时,查看其他设备完整信息".to_string())); diff --git a/vnt-jni/src/vnt_util.rs b/vnt-jni/src/vnt_util.rs index 1db665d..0052861 100644 --- a/vnt-jni/src/vnt_util.rs +++ b/vnt-jni/src/vnt_util.rs @@ -80,7 +80,7 @@ fn new_sync(env: &mut JNIEnv, config: JObject) -> Result { token, device_id, name, server_address, server_address_str, stun_server, vec![], - vec![], password, false, None, false, None, false,false); + vec![], password, false, None, false, None, false,false,1); match VntUtilSync::new(config) { Ok(vnt_util) => { Ok(vnt_util) diff --git a/vnt/Cargo.toml b/vnt/Cargo.toml index 2af95d8..35ba07d 100644 --- a/vnt/Cargo.toml +++ b/vnt/Cargo.toml @@ -13,6 +13,8 @@ libc = "0.2.137" crossbeam-utils = "0.8" dashmap = "5.5.1" parking_lot = "0.12.1" +byte-pool = "0.2.4" +lazy_static = "1.4.0" rand = "0.8.5" sha2 = { version = "0.10.6", features = ["oid"] } thiserror = "1.0.37" diff --git a/vnt/src/channel/channel.rs b/vnt/src/channel/channel.rs index 0b82474..5ea7a05 100644 --- a/vnt/src/channel/channel.rs +++ b/vnt/src/channel/channel.rs @@ -13,7 +13,10 @@ use crate::channel::punch::NatType; use crate::core::status::VntWorker; use crate::handle::CurrentDeviceInfo; use crate::handle::recv_handler::ChannelDataHandler; - +use byte_pool::{Block, BytePool}; +lazy_static::lazy_static! { + static ref POOL:BytePool = BytePool::new(); +} pub struct ContextInner { //udp用于打洞、服务端通信(可选) pub(crate) main_channel: Arc, @@ -65,7 +68,7 @@ impl Context { pub fn close(&self) { let _ = self.inner.status_sender.send(Status::Close); } - pub fn is_main_tcp(&self)->bool{ + pub fn is_main_tcp(&self) -> bool { self.inner.main_tcp_channel.is_some() } pub fn switch(&self, nat_type: NatType) { @@ -204,7 +207,7 @@ impl Context { } fn add_route_(&self, id: Ipv4Addr, route: Route, only_if_absent: bool) { let key = route.route_key(); - let mut list = self.inner.route_table.entry(id).or_insert_with(||Vec::with_capacity(4)); + let mut list = self.inner.route_table.entry(id).or_insert_with(|| Vec::with_capacity(4)); let mut exist = false; for x in list.iter_mut() { if x.metric < route.metric { @@ -293,7 +296,7 @@ impl Context { v } pub fn remove_route_all(&self, id: &Ipv4Addr) { - if let Some((_,routes)) = self.inner.route_table.remove(id) { + if let Some((_, routes)) = self.inner.route_table.remove(id) { for x in routes { self.inner.route_table_time.remove(&(x.route_key(), *id)); } @@ -331,12 +334,12 @@ impl Channel { } #[derive(Clone)] -struct BufSenderGroup(usize, Vec, usize, usize, RouteKey)>>); +struct BufSenderGroup(usize, Vec, usize, usize, RouteKey)>>); -struct BufReceiverGroup(Vec, usize, usize, RouteKey)>>); +struct BufReceiverGroup(Vec, usize, usize, RouteKey)>>); impl BufSenderGroup { - pub async fn send(&mut self, val: (Vec, usize, usize, RouteKey)) -> bool { + pub async fn send(&mut self, val: (Block<'static>, usize, usize, RouteKey)) -> bool { let index = self.0 % self.1.len(); self.0 = self.0.wrapping_add(1); self.1[index].send(val).await.is_ok() @@ -347,7 +350,7 @@ fn buf_channel_group(size: usize) -> (BufSenderGroup, BufReceiverGroup) { let mut buf_sender_group = Vec::with_capacity(size); let mut buf_receiver_group = Vec::with_capacity(size); for _ in 0..size { - let (buf_sender, buf_receiver) = tokio::sync::mpsc::channel::<(Vec, usize, usize, RouteKey)>(10); + let (buf_sender, buf_receiver) = tokio::sync::mpsc::channel::<(Block<'static, Vec>, usize, usize, RouteKey)>(10); buf_sender_group.push(buf_sender); buf_receiver_group.push(buf_receiver); } @@ -360,7 +363,7 @@ impl Channel { let addr = tcp_r.peer_addr()?; let key = RouteKey::new(0, addr); loop { - let mut buf = vec![0; 4096]; + let mut buf = POOL.alloc(4096); tcp_r.read_exact(&mut head).await?; let len = (((head[2] as u16) << 8) | head[3] as u16) as usize; if len < 12 || len > buf.len() { @@ -439,8 +442,9 @@ impl Channel { head_reserve: usize,//头部预留字节 symmetric_channel_num: usize,//对称网络，则再加一组监听，提升打洞成功率 relay: bool, + parallel: usize, ) { - let (buf_sender, buf_receiver) = buf_channel_group(std::thread::available_parallelism().unwrap().get()); + let (buf_sender, buf_receiver) = buf_channel_group(parallel); for mut buf_receiver in buf_receiver.0 { let context = self.context.clone(); let handler = self.handler.clone(); @@ -523,7 +527,7 @@ impl Channel { let id = 1 + udp.as_raw_fd() as usize; context.inner.udp_map.insert(id, udp.clone()); loop { - let mut buf = vec![0; 4096]; + let mut buf = POOL.alloc(4096); tokio::select! { rs=udp.recv_from(&mut buf[head_reserve..])=>{ match rs { diff --git a/vnt/src/core/mod.rs b/vnt/src/core/mod.rs index cdbc185..baa73f8 100644 --- a/vnt/src/core/mod.rs +++ b/vnt/src/core/mod.rs @@ -98,13 +98,13 @@ impl VntUtil { } /// 加密握手 用于同步密钥 pub async fn secret_handshake(&mut self) -> Result<(), HandshakeEnum> { - handshake_handler::secret_handshake(&self.main_channel, self.main_tcp_channel.as_mut(), self.config.server_address, self.rsa_cipher.as_ref().unwrap(), &self.server_cipher,self.config.token.clone()).await + handshake_handler::secret_handshake(&self.main_channel, self.main_tcp_channel.as_mut(), self.config.server_address, self.rsa_cipher.as_ref().unwrap(), &self.server_cipher, self.config.token.clone()).await } /// 注册 pub async fn register(&mut self) -> Result { match registration_handler::registration(&self.main_channel, self.main_tcp_channel.as_mut(), &self.server_cipher, self.config.server_address, self.config.token.clone(), self.config.device_id.clone(), - self.config.name.clone(), self.config.ip.unwrap_or(Ipv4Addr::UNSPECIFIED),self.config.password.is_some()).await { + self.config.name.clone(), self.config.ip.unwrap_or(Ipv4Addr::UNSPECIFIED), self.config.password.is_some()).await { Ok(res) => { let _ = self.response.insert(res.clone()); Ok(res) @@ -149,9 +149,9 @@ impl VntUtil { }; let mtu = match self.config.mtu { None => { - if self.config.password.is_none(){ + if self.config.password.is_none() { 1430 - }else{ + } else { 1410 } } @@ -206,7 +206,7 @@ impl VntUtil { let register = Arc::new(registration_handler::Register::new(self.server_cipher.clone(), channel_sender.clone(), config.server_address, config.token.clone(), - config.device_id.clone(), config.name.clone(),config.password.is_some())); + config.device_id.clone(), config.name.clone(), config.password.is_some())); let device_list: Arc)>> = Arc::new(Mutex::new((response.epoch, response.device_info_list))); let peer_nat_info_map: Arc> = Arc::new(DashMap::new()); let connect_status = Arc::new(AtomicCell::new(ConnectStatus::Connected)); @@ -237,14 +237,16 @@ impl VntUtil { #[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))] if config.tap { tap_handler::start(vnt_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(), client_cipher.clone(), self.server_cipher.clone()); + igmp_server.clone(), current_device.clone(), in_external_route, ip_proxy_map.clone(), + client_cipher.clone(), self.server_cipher.clone(),config.parallel); } else { tun_handler::start(vnt_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(), client_cipher.clone(), self.server_cipher.clone()).await; + igmp_server.clone(), current_device.clone(), in_external_route, ip_proxy_map.clone(), + client_cipher.clone(), self.server_cipher.clone(),config.parallel).await; } #[cfg(any(target_os = "android"))] tun_handler::start(vnt_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()).await; + igmp_server.clone(), current_device.clone(), in_external_route, ip_proxy_map.clone(), cipher.clone(),config.parallel).await; //外部数据接收处理 let channel_recv_handler = ChannelDataHandler::new(current_device.clone(), device_list.clone(), @@ -252,7 +254,7 @@ impl VntUtil { device_writer.clone(), connect_status.clone(), peer_nat_info_map.clone(), ip_proxy_map, out_external_route, cone_sender, symmetric_sender, client_cipher.clone(), - self.server_cipher.clone(), self.rsa_cipher.clone(), config.relay,config.token.clone()); + self.server_cipher.clone(), self.rsa_cipher.clone(), config.relay, config.token.clone()); { let channel = Channel::new(context.clone(), channel_recv_handler); let channel_worker = vnt_status_manager.worker("channel_worker"); @@ -264,7 +266,7 @@ impl VntUtil { tokio::spawn(udp_proxy.start()); } tokio::spawn(async move { - channel.start(channel_worker, tcp, 14, 65, relay).await + channel.start(channel_worker, tcp, 14, 65, relay,config.parallel).await }); } { @@ -287,7 +289,7 @@ impl VntUtil { } context.switch(nat_test.nat_info().nat_type); Ok(Vnt { - config:self.config, + config: self.config, current_device, context, vnt_status_manager, @@ -387,6 +389,7 @@ pub struct Config { pub ip: Option, pub relay: bool, pub server_encrypt: bool, + pub parallel: usize, } @@ -400,7 +403,7 @@ impl Config { in_ips: Vec<(u32, u32, Ipv4Addr)>, out_ips: Vec<(u32, u32)>, password: Option, simulate_multicast: bool, mtu: Option, tcp: bool, ip: Option, - relay: bool, server_encrypt: bool, ) -> Self { + relay: bool, server_encrypt: bool, parallel: usize, ) -> Self { for x in stun_server.iter_mut() { if !x.contains(":") { x.push_str(":3478"); @@ -423,6 +426,7 @@ impl Config { ip, relay, server_encrypt, + parallel, } } } \ No newline at end of file diff --git a/vnt/src/handle/tun_tap/channel_group.rs b/vnt/src/handle/tun_tap/channel_group.rs index 9d896bc..0745645 100644 --- a/vnt/src/handle/tun_tap/channel_group.rs +++ b/vnt/src/handle/tun_tap/channel_group.rs @@ -1,10 +1,12 @@ -#[derive(Clone)] -pub struct BufSenderGroup(usize, Vec, usize, usize)>>); +use byte_pool::Block; -pub struct BufReceiverGroup(pub Vec, usize, usize)>>); +#[derive(Clone)] +pub struct BufSenderGroup(usize, Vec, usize, usize)>>); + +pub struct BufReceiverGroup(pub Vec, usize, usize)>>); impl BufSenderGroup { - pub async fn send(&mut self, val: (Vec, usize, usize)) -> bool { + pub async fn send(&mut self, val: (Block<'static>, usize, usize)) -> bool { let index = self.0 % self.1.len(); self.0 = self.0.wrapping_add(1); self.1[index].send(val).await.is_ok() @@ -15,7 +17,7 @@ pub fn buf_channel_group(size: usize) -> (BufSenderGroup, BufReceiverGroup) { let mut buf_sender_group = Vec::with_capacity(size); let mut buf_receiver_group = Vec::with_capacity(size); for _ in 0..size { - let (buf_sender, buf_receiver) = tokio::sync::mpsc::channel::<(Vec, usize, usize)>(10); + let (buf_sender, buf_receiver) = tokio::sync::mpsc::channel::<(Block<'static>, usize, usize)>(10); buf_sender_group.push(buf_sender); buf_receiver_group.push(buf_receiver); } diff --git a/vnt/src/handle/tun_tap/tap_handler.rs b/vnt/src/handle/tun_tap/tap_handler.rs index 4440f85..62bca57 100644 --- a/vnt/src/handle/tun_tap/tap_handler.rs +++ b/vnt/src/handle/tun_tap/tap_handler.rs @@ -1,7 +1,9 @@ use std::{io, thread}; use std::sync::Arc; +use byte_pool::BytePool; use crossbeam_utils::atomic::AtomicCell; +use lazy_static::lazy_static; use packet::arp::arp::ArpPacket; use packet::ethernet; @@ -28,8 +30,8 @@ pub fn start(worker: VntWorker, sender: ChannelSender, current_device: Arc>, ip_route: Option, ip_proxy_map: Option, - client_cipher: Cipher, server_cipher: Cipher) { - let (buf_sender, buf_receiver) = buf_channel_group(thread::available_parallelism().unwrap().get()); + client_cipher: Cipher, server_cipher: Cipher,parallel:usize) { + let (buf_sender, buf_receiver) = buf_channel_group(parallel); for mut buf_receiver in buf_receiver.0 { let sender = sender.clone(); let device_writer = device_writer.clone(); @@ -62,12 +64,15 @@ pub fn start(worker: VntWorker, sender: ChannelSender, }); }).unwrap(); } - +lazy_static!{ + static ref POOL:BytePool> = BytePool::>::new(); +} async fn start_(sender: ChannelSender, device_reader: DeviceReader, mut buf_sender: BufSenderGroup) -> io::Result<()> { + loop { - let mut buf = vec![0; 4096]; + let mut buf = POOL.alloc(4096); if sender.is_close() { return Ok(()); } diff --git a/vnt/src/handle/tun_tap/tun_handler.rs b/vnt/src/handle/tun_tap/tun_handler.rs index 8cf9c0c..e097b30 100644 --- a/vnt/src/handle/tun_tap/tun_handler.rs +++ b/vnt/src/handle/tun_tap/tun_handler.rs @@ -1,5 +1,6 @@ use std::{io, thread}; use std::sync::Arc; +use byte_pool::BytePool; use crossbeam_utils::atomic::AtomicCell; @@ -18,7 +19,9 @@ use crate::handle::tun_tap::channel_group::{buf_channel_group, BufSenderGroup}; use crate::igmp_server::IgmpServer; use crate::ip_proxy::IpProxyMap; use crate::tun_tap_device::{DeviceReader, DeviceWriter}; - +lazy_static::lazy_static! { + static ref POOL:BytePool> = BytePool::>::new(); +} fn icmp(device_writer: &DeviceWriter, mut ipv4_packet: IpV4Packet<&mut [u8]>) -> Result<()> { if ipv4_packet.protocol() == ipv4::protocol::Protocol::Icmp { let mut icmp = IcmpPacket::new(ipv4_packet.payload_mut())?; @@ -65,8 +68,8 @@ pub async fn start(worker: VntWorker, sender: ChannelSender, current_device: Arc>, ip_route: Option, ip_proxy_map: Option, - client_cipher: Cipher, server_cipher: Cipher) { - let (buf_sender, buf_receiver) = buf_channel_group(thread::available_parallelism().unwrap().get()); + client_cipher: Cipher, server_cipher: Cipher, parallel: usize) { + let (buf_sender, buf_receiver) = buf_channel_group(parallel); for mut buf_receiver in buf_receiver.0 { let sender = sender.clone(); let device_writer = device_writer.clone(); @@ -104,7 +107,7 @@ pub async fn start(worker: VntWorker, sender: ChannelSender, async fn start_(sender: ChannelSender, device_reader: DeviceReader, mut buf_sender: BufSenderGroup) -> io::Result<()> { loop { - let mut buf = vec![0; 4096]; + let mut buf = POOL.alloc(4096); if sender.is_close() { return Ok(()); }