From d81ac6368ea2adc7fc0c34df0c571f81651d2c2e Mon Sep 17 00:00:00 2001 From: lbl8603 <49143209+lbl8603@users.noreply.github.com> Date: Sat, 6 Jul 2024 16:27:08 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E6=B5=81=E9=87=8F=E7=BB=9F?= =?UTF-8?q?=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.lock | 4 +- common/src/cli.rs | 21 ++++ common/src/command/client.rs | 13 ++- common/src/command/entity.rs | 23 +++- common/src/command/mod.rs | 104 ++++++++++++++++- common/src/command/server.rs | 20 +++- common/src/config/file_config.rs | 3 + common/src/console_out/mod.rs | 120 +++++++++++++++++++- vn-link-cli/src/main.rs | 2 +- vnt-cli/src/main.rs | 2 +- vnt/src/channel/context.rs | 46 ++++++-- vnt/src/channel/mod.rs | 9 +- vnt/src/channel/sender.rs | 4 +- vnt/src/core/conn.rs | 65 ++++++----- vnt/src/core/mod.rs | 3 + vnt/src/handle/handshaker.rs | 2 +- vnt/src/handle/maintain/heartbeat.rs | 10 +- vnt/src/handle/maintain/punch.rs | 2 +- vnt/src/handle/maintain/up_status.rs | 37 ++---- vnt/src/handle/recv_data/client.rs | 10 +- vnt/src/handle/recv_data/mod.rs | 11 +- vnt/src/handle/recv_data/server.rs | 10 +- vnt/src/handle/recv_data/turn.rs | 2 +- vnt/src/handle/tun_tap/mod.rs | 2 +- vnt/src/handle/tun_tap/tun_handler.rs | 19 ++-- vnt/src/handle/tun_tap/unix.rs | 7 +- vnt/src/handle/tun_tap/windows.rs | 6 +- vnt/src/ip_proxy/icmp_proxy.rs | 2 +- vnt/src/tun_tap_device/tun_create_helper.rs | 8 +- vnt/src/util/mod.rs | 6 +- vnt/tun/src/windows/tun/mod.rs | 2 +- 31 files changed, 425 insertions(+), 150 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1d63b81..3329c85 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2019,9 +2019,9 @@ dependencies = [ [[package]] name = "unicode-width" -version = "0.1.11" +version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" +checksum = "0336d538f7abc86d282a4189614dfaa90810dfc2c6f6427eaf88e16311dd225d" [[package]] name = "universal-hash" diff --git a/common/src/cli.rs b/common/src/cli.rs index 9092fe4..f5c5c8b 100644 --- a/common/src/cli.rs +++ b/common/src/cli.rs @@ -74,12 +74,15 @@ pub fn parse_args_config() -> anyhow::Result, bool)> opts.optmulti("", "vnt-mapping", "vnt-mapping", ""); opts.optopt("f", "", "配置文件", ""); opts.optopt("", "compressor", "压缩算法", ""); + opts.optflag("", "disable-stats", "关闭流量统计"); //"后台运行时,查看其他设备列表" opts.optflag("", "add", "后台运行时,添加地址"); opts.optflag("", "list", "后台运行时,查看其他设备列表"); opts.optflag("", "all", "后台运行时,查看其他设备完整信息"); opts.optflag("", "info", "后台运行时,查看当前设备信息"); opts.optflag("", "route", "后台运行时,查看数据转发路径"); + opts.optflag("", "chart_a", "后台运行时,查看流量统计"); + opts.optopt("", "chart_b", "后台运行时,查看流量统计", ""); opts.optflag("", "stop", "停止后台运行"); opts.optflag("h", "help", "帮助"); let matches = match opts.parse(&args[1..]) { @@ -110,6 +113,13 @@ pub fn parse_args_config() -> anyhow::Result, bool)> } else if matches.opt_present("all") { command::command(command::CommandEnum::All); return Ok(None); + } else if matches.opt_present("chart_a") { + command::command(command::CommandEnum::ChartA); + return Ok(None); + } + if let Some(v) = matches.opt_str("chart_b") { + command::command(command::CommandEnum::ChartB(v)); + return Ok(None); } let conf = matches.opt_str("f"); let (config, vnt_link_config, cmd) = if conf.is_some() { @@ -267,6 +277,7 @@ pub fn parse_args_config() -> anyhow::Result, bool)> #[cfg(feature = "port_mapping")] let port_mapping_list = matches.opt_strs("mapping"); let vnt_mapping_list = matches.opt_strs("vnt-mapping"); + let disable_stats = matches.opt_present("disable-stats"); let compressor = if let Some(compressor) = matches.opt_str("compressor").as_ref() { Compressor::from_str(compressor) .map_err(|e| anyhow!("{}", e)) @@ -306,6 +317,7 @@ pub fn parse_args_config() -> anyhow::Result, bool)> #[cfg(feature = "port_mapping")] port_mapping_list, compressor, + !disable_stats, ) { Ok(config) => config, Err(e) => { @@ -419,6 +431,7 @@ fn print_usage(program: &str, _opts: Options) { .to_string() ) ); + println!(" --disable-stats 关闭流量统计"); println!(); #[cfg(feature = "command")] { @@ -443,6 +456,14 @@ fn print_usage(program: &str, _opts: Options) { " --route {}", yellow("后台运行时,查看数据转发路径".to_string()) ); + println!( + " --chart_a {}", + yellow("后台运行时,查看所有IP的流量统计".to_string()) + ); + println!( + " --chart_b {}", + yellow("后台运行时,查看单个IP的历史流量".to_string()) + ); println!( " --stop {}", yellow("停止后台运行".to_string()) diff --git a/common/src/command/client.rs b/common/src/command/client.rs index 8143fb6..e6caab9 100644 --- a/common/src/command/client.rs +++ b/common/src/command/client.rs @@ -4,7 +4,7 @@ use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket}; use std::str::FromStr; use std::time::Duration; -use crate::command::entity::{DeviceItem, Info, RouteItem}; +use crate::command::entity::{ChartA, ChartB, DeviceItem, Info, RouteItem}; pub struct CommandClient { buf: Vec, @@ -53,6 +53,17 @@ impl CommandClient { pub fn info(&mut self) -> io::Result { self.send_cmd(b"info") } + pub fn chart_a(&mut self) -> io::Result { + self.send_cmd(b"chart_a") + } + pub fn chart_b(&mut self, input: &str) -> io::Result { + let cmd = if input.is_empty() { + "chart_b".to_string() + } else { + format!("chart_b:{}", input) + }; + self.send_cmd(cmd.as_bytes()) + } fn send_cmd<'a, V: Deserialize<'a>>(&'a mut self, cmd: &[u8]) -> io::Result { self.udp.send(cmd)?; let len = self.udp.recv(&mut self.buf)?; diff --git a/common/src/command/entity.rs b/common/src/command/entity.rs index b741488..b42f054 100644 --- a/common/src/command/entity.rs +++ b/common/src/command/entity.rs @@ -1,5 +1,7 @@ use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use std::net::{Ipv4Addr, SocketAddr}; + #[derive(Serialize, Deserialize, Debug)] pub struct Info { pub name: String, @@ -12,8 +14,6 @@ pub struct Info { pub public_ips: String, pub local_addr: String, pub ipv6_addr: String, - pub up: u64, - pub down: u64, pub port_mapping_list: Vec<(bool, SocketAddr, String)>, pub in_ips: Vec<(u32, u32, Ipv4Addr)>, pub out_ips: Vec<(u32, u32)>, @@ -46,3 +46,22 @@ pub struct DeviceItem { pub current_client_secret: bool, pub current_client_secret_hash: Vec, } + +#[derive(Serialize, Deserialize, Debug, Default)] +pub struct ChartA { + pub disable_stats: bool, + pub up_total: u64, + pub down_total: u64, + pub up_map: HashMap, + pub down_map: HashMap, +} + +#[derive(Serialize, Deserialize, Debug, Default)] +pub struct ChartB { + pub disable_stats: bool, + pub ip: Option, + pub up_total: u64, + pub up_list: Vec, + pub down_total: u64, + pub down_list: Vec, +} diff --git a/common/src/command/mod.rs b/common/src/command/mod.rs index 0f2f1d1..b408b88 100644 --- a/common/src/command/mod.rs +++ b/common/src/command/mod.rs @@ -1,8 +1,10 @@ +use std::collections::HashSet; use std::io; +use std::net::Ipv4Addr; use vnt::channel::ConnectProtocol; use vnt::core::Vnt; -use crate::command::entity::{DeviceItem, Info, RouteItem}; +use crate::command::entity::{ChartA, ChartB, DeviceItem, Info, RouteItem}; use crate::console_out; pub mod client; @@ -14,6 +16,8 @@ pub enum CommandEnum { List, All, Info, + ChartA, + ChartB(String), Stop, } @@ -21,7 +25,9 @@ pub fn command_str(cmd: &str, vnt: &Vnt) -> bool { if cmd.is_empty() { return false; } - match cmd.to_lowercase().trim() { + let cmd = cmd.to_lowercase(); + let cmd = cmd.trim(); + match cmd { "list" => { let list = command_list(&vnt); console_out::console_device_list(list); @@ -38,12 +44,24 @@ pub fn command_str(cmd: &str, vnt: &Vnt) -> bool { let list = command_list(&vnt); console_out::console_device_list_all(list); } + "chart_a" => { + let chart = command_chart_a(&vnt); + console_out::console_chart_a(chart); + } "stop" => { let _ = vnt.stop(); return false; } _ => {} } + if let Some(ip) = cmd.strip_prefix("chart_b") { + let chart = if ip.is_empty() { + command_chart_b(&vnt, &vnt.current_device().virtual_gateway.to_string()) + } else { + command_chart_b(&vnt, &ip[1..]) + }; + console_out::console_chart_b(chart); + } println!(); return true; } @@ -73,6 +91,14 @@ fn command_(cmd: CommandEnum) -> io::Result<()> { let info = command_client.info()?; console_out::console_info(info); } + CommandEnum::ChartA => { + let chart = command_client.chart_a()?; + console_out::console_chart_a(chart); + } + CommandEnum::ChartB(input) => { + let chart = command_client.chart_b(&input)?; + console_out::console_chart_b(chart); + } CommandEnum::Stop => { command_client.stop()?; } @@ -224,8 +250,6 @@ pub fn command_info(vnt: &Vnt) -> Info { .ipv6() .map(|v| v.to_string()) .unwrap_or("None".to_string()); - let up = vnt.up_stream(); - let down = vnt.down_stream(); #[cfg(feature = "port_mapping")] let port_mapping_list = vnt.config().port_mapping_list.clone(); #[cfg(not(feature = "port_mapping"))] @@ -249,8 +273,6 @@ pub fn command_info(vnt: &Vnt) -> Info { public_ips, local_addr, ipv6_addr, - up, - down, port_mapping_list, in_ips, out_ips, @@ -258,3 +280,73 @@ pub fn command_info(vnt: &Vnt) -> Info { tcp_listen_addr, } } + +pub fn command_chart_a(vnt: &Vnt) -> ChartA { + let disable_stats = !vnt.config().enable_traffic; + if disable_stats { + let mut chart = ChartA::default(); + chart.disable_stats = true; + return chart; + } + let (up_total, up_map) = vnt.up_stream_all().unwrap_or_default(); + let (down_total, down_map) = vnt.down_stream_all().unwrap_or_default(); + ChartA { + disable_stats, + up_total, + down_total, + up_map, + down_map, + } +} + +pub fn command_chart_b(vnt: &Vnt, input_str: &str) -> ChartB { + let disable_stats = !vnt.config().enable_traffic; + if disable_stats { + let mut chart = ChartB::default(); + chart.disable_stats = true; + return chart; + } + let (_, up_map) = vnt.up_stream_history().unwrap_or_default(); + let (_, down_map) = vnt.down_stream_history().unwrap_or_default(); + let up_keys: HashSet<_> = up_map.keys().cloned().collect(); + let down_keys: HashSet<_> = down_map.keys().cloned().collect(); + let mut keys: Vec = up_keys.union(&down_keys).cloned().collect(); + keys.sort(); + if let Some(ip) = find_matching_ipv4_address(input_str, &keys) { + let (up_total, up_list) = up_map.get(&ip).cloned().unwrap_or_default(); + let (down_total, down_list) = down_map.get(&ip).cloned().unwrap_or_default(); + ChartB { + disable_stats, + ip: Some(ip), + up_total, + up_list, + down_total, + down_list, + } + } else { + ChartB::default() + } +} + +fn match_from_end(input_str: &str, ip: &str) -> bool { + let mut input_chars = input_str.chars().rev(); + let mut ip_chars = ip.chars().rev(); + + while let (Some(ic), Some(pc)) = (input_chars.next(), ip_chars.next()) { + if ic != pc { + return false; + } + } + + input_chars.next().is_none() // Ensure all input characters matched +} + +fn find_matching_ipv4_address(input_str: &str, ip_addresses: &[Ipv4Addr]) -> Option { + for &ip in ip_addresses { + let ip_str = ip.to_string(); + if match_from_end(input_str, &ip_str) { + return Some(ip); + } + } + None +} diff --git a/common/src/command/server.rs b/common/src/command/server.rs index 3ac61a2..38a7c14 100644 --- a/common/src/command/server.rs +++ b/common/src/command/server.rs @@ -1,3 +1,4 @@ +use crate::command::command_chart_b; use std::io; use std::io::Write; use std::net::UdpSocket; @@ -62,15 +63,26 @@ fn command(cmd: &str, vnt: &Vnt) -> io::Result { .unwrap_or_else(|e| format!("error {:?}", e)), "info" => serde_yaml::to_string(&crate::command::command_info(vnt)) .unwrap_or_else(|e| format!("error {:?}", e)), + "chart_a" => serde_yaml::to_string(&crate::command::command_chart_a(vnt)) + .unwrap_or_else(|e| format!("error {:?}", e)), "stop" => { vnt.stop(); "stopped".to_string() } _ => { - format!( - "command '{}' not found. Try to enter: 'route'/'list'/'stop' \n", - cmd - ) + if let Some(ip) = cmd.strip_prefix("chart_b") { + let chart = if ip.is_empty() { + command_chart_b(&vnt, &vnt.current_device().virtual_gateway.to_string()) + } else { + command_chart_b(&vnt, &ip[1..]) + }; + serde_yaml::to_string(&chart).unwrap_or_else(|e| format!("error {:?}", e)) + } else { + format!( + "command '{}' not found. Try to enter: 'route'/'list'/'stop' \n", + cmd + ) + } } }; Ok(out_str) diff --git a/common/src/config/file_config.rs b/common/src/config/file_config.rs index 4af6075..2c9b24e 100644 --- a/common/src/config/file_config.rs +++ b/common/src/config/file_config.rs @@ -45,6 +45,7 @@ pub struct FileConfig { pub mapping: Vec, pub compressor: Option, pub vnt_mapping: Vec, + pub disable_stats: bool, } impl Default for FileConfig { @@ -88,6 +89,7 @@ impl Default for FileConfig { mapping: vec![], compressor: None, vnt_mapping: vec![], + disable_stats: false, } } } @@ -174,6 +176,7 @@ pub fn read_config(file_path: &str) -> anyhow::Result<(Config, Vec, bool #[cfg(feature = "port_mapping")] file_conf.mapping, compressor, + !file_conf.disable_stats, )?; Ok((config, file_conf.vnt_mapping, file_conf.cmd)) diff --git a/common/src/console_out/mod.rs b/common/src/console_out/mod.rs index 641767f..0d0ec94 100644 --- a/common/src/console_out/mod.rs +++ b/common/src/console_out/mod.rs @@ -1,7 +1,8 @@ use console::{style, Style}; +use std::collections::HashSet; use std::net::Ipv4Addr; -use crate::command::entity::{DeviceItem, Info, RouteItem}; +use crate::command::entity::{ChartA, ChartB, DeviceItem, Info, RouteItem}; pub mod table; @@ -29,8 +30,6 @@ pub fn console_info(status: Info) { println!("Public ips: {}", style(status.public_ips).green()); println!("Local addr: {}", style(status.local_addr).green()); println!("IPv6: {}", style(status.ipv6_addr).green()); - println!("Up: {}", style(convert(status.up)).green()); - println!("Down: {}", style(convert(status.down)).green()); if !status.port_mapping_list.is_empty() { println!("------------------------------------------"); @@ -242,3 +241,118 @@ pub fn console_device_list_all(mut list: Vec) { } table::println_table(out_list) } + +pub fn console_chart_a(chart_a: ChartA) { + if chart_a.disable_stats { + println!("Traffic statistics not enabled"); + return; + } + println!(); + println!("-----------------------------------------------------------------"); + println!( + "Upload total = {}", + style(convert(chart_a.up_total)).green() + ); + println!( + "Download total = {}", + style(convert(chart_a.down_total)).green() + ); + println!("-----------------------------------------------------------------"); + let up_keys: HashSet<_> = chart_a.up_map.keys().cloned().collect(); + let down_keys: HashSet<_> = chart_a.down_map.keys().cloned().collect(); + let mut keys: Vec = up_keys.union(&down_keys).cloned().collect(); + // 排序 + keys.sort(); + + // 找到最大的值，用于缩放条形图长度 + let up_max_value = *chart_a.up_map.values().max().unwrap_or(&0); + let down_max_value = *chart_a.down_map.values().max().unwrap_or(&0); + let max_value = up_max_value.max(down_max_value); + let max_height = 50; + // 打印条形图 + for key in &keys { + if let Some(&value) = chart_a.up_map.get(key) { + let bar = "█".repeat(((value as f64 / max_value as f64) * max_height as f64) as usize); + println!( + "{:<10} | {} upload {}", + key, + bar, + style(convert(value)).green() + ); + } + if let Some(&value) = chart_a.down_map.get(key) { + let bar = "█".repeat(((value as f64 / max_value as f64) * max_height as f64) as usize); + println!( + "{:<10} | {} download {}", + key, + bar, + style(convert(value)).green() + ); + } + println!("-"); + } +} + +pub fn console_chart_b(chart_b: ChartB) { + if chart_b.disable_stats { + println!("Traffic statistics not enabled"); + return; + } + let ip = if let Some(ip) = chart_b.ip { + ip + } else { + println!("Ip: None"); + return; + }; + println!("---------------------------- upload ----------------------------"); + println!("IP: {}", ip); + println!("Upload total: {}", style(convert(chart_b.up_total)).green()); + println!( + "Max: {}", + style(convert( + chart_b + .up_list + .iter() + .max() + .cloned() + .map_or(0, |v| v as u64) + )) + .green() + ); + console_chart_b_list(chart_b.up_list); + println!("---------------------------- download ----------------------------"); + println!("IP: {}", ip); + println!( + "Download total: {}", + style(convert(chart_b.down_total)).green() + ); + println!( + "Max: {}", + style(convert( + chart_b + .down_list + .iter() + .max() + .cloned() + .map_or(0, |v| v as u64) + )) + .green() + ); + console_chart_b_list(chart_b.down_list); +} +fn console_chart_b_list(list: Vec) { + let max_value = *list.iter().max().unwrap_or(&0); + let max_height = max_value.min(20); + // 遍历从最大高度到0 + for i in (0..=max_height).rev() { + for &value in &list { + let scaled_value = (value as f64 / max_value as f64 * max_height as f64) as usize; + if scaled_value >= i { + print!("█"); + } else { + print!(" "); + } + } + println!(); + } +} diff --git a/vn-link-cli/src/main.rs b/vn-link-cli/src/main.rs index 78336a3..33d526d 100644 --- a/vn-link-cli/src/main.rs +++ b/vn-link-cli/src/main.rs @@ -65,7 +65,7 @@ async fn main0(config: Config, vn_link_config: VnLinkConfig, _show_cmd: bool) { let mut reader = tokio::io::BufReader::new(tokio::io::stdin()); loop { cmd.clear(); - println!("======== input:list,info,route,all,stop ========"); + println!("======== input:list,info,route,all,stop,chart_a,chart_b[:ip] ========"); match reader.read_line(&mut cmd).await { Ok(len) => { if !common::command::command_str(&cmd[..len], vnt_c) { diff --git a/vnt-cli/src/main.rs b/vnt-cli/src/main.rs index 6a0dd49..4c5a23f 100644 --- a/vnt-cli/src/main.rs +++ b/vnt-cli/src/main.rs @@ -77,7 +77,7 @@ fn main0(config: Config, _show_cmd: bool) { let mut cmd = String::new(); loop { cmd.clear(); - println!("======== input:list,info,route,all,stop ========"); + println!("======== input:list,info,route,all,stop,chart_a,chart_b[:ip] ========"); match std::io::stdin().read_line(&mut cmd) { Ok(len) => { if !common::command::command_str(&cmd[..len], &vnt_util) { diff --git a/vnt/src/channel/context.rs b/vnt/src/channel/context.rs index 6c5da52..a78c999 100644 --- a/vnt/src/channel/context.rs +++ b/vnt/src/channel/context.rs @@ -13,6 +13,8 @@ use rand::Rng; use crate::channel::punch::NatType; use crate::channel::sender::{AcceptSocketSender, PacketSender}; use crate::channel::{ConnectProtocol, Route, RouteKey, UseChannelType, DEFAULT_RT}; +use crate::protocol::NetPacket; +use crate::util::limit::TrafficMeterMultiAddress; /// 传输通道上下文，持有udp socket、tcp socket和路由信息 #[derive(Clone)] @@ -29,6 +31,8 @@ impl ChannelContext { packet_loss_rate: Option, packet_delay: u32, use_ipv6: bool, + up_traffic_meter: Option, + down_traffic_meter: Option, ) -> Self { let channel_num = main_udp_socket.len(); assert_ne!(channel_num, 0, "not channel"); @@ -52,6 +56,8 @@ impl ChannelContext { packet_delay, main_index: AtomicUsize::new(0), use_ipv6, + up_traffic_meter, + down_traffic_meter, }; Self { inner: Arc::new(inner), @@ -88,6 +94,8 @@ pub struct ContextInner { packet_delay: u32, main_index: AtomicUsize, use_ipv6: bool, + pub(crate) up_traffic_meter: Option, + pub(crate) down_traffic_meter: Option, } impl ContextInner { @@ -179,12 +187,20 @@ impl ContextInner { Ok(()) } /// 将数据发送到默认通道，一般发往服务器才用此方法 - pub fn send_default(&self, buf: &[u8], addr: SocketAddr) -> io::Result<()> { + pub fn send_default>( + &self, + buf: &NetPacket, + addr: SocketAddr, + ) -> io::Result<()> { if self.protocol.is_udp() { - self.send_main_udp(self.main_index.load(Ordering::Relaxed), buf, addr) + self.send_main_udp(self.main_index.load(Ordering::Relaxed), buf.buffer(), addr)? } else { - self.send_tcp(buf, addr) + self.send_tcp(buf.buffer(), addr)? } + if let Some(up_traffic_meter) = &self.up_traffic_meter { + up_traffic_meter.add_traffic(buf.destination(), buf.data_len()); + } + Ok(()) } pub fn change_main_index(&self) { @@ -209,9 +225,9 @@ impl ContextInner { } } /// 发送网络数据 - pub fn send_ipv4_by_id( + pub fn send_ipv4_by_id>( &self, - buf: &[u8], + buf: &NetPacket, id: &Ipv4Addr, server_addr: SocketAddr, send_default: bool, @@ -221,6 +237,7 @@ impl ContextInner { return Ok(()); } } + if self.packet_delay > 0 { thread::sleep(Duration::from_millis(self.packet_delay as _)); } @@ -237,7 +254,7 @@ impl ContextInner { Ok(()) } /// 将数据发到指定id - pub fn send_by_id(&self, buf: &[u8], id: &Ipv4Addr) -> io::Result<()> { + pub fn send_by_id>(&self, buf: &NetPacket, id: &Ipv4Addr) -> io::Result<()> { let mut c = 0; loop { let route = self.route_table.get_route_by_id(c, id)?; @@ -257,28 +274,35 @@ impl ContextInner { } } /// 将数据发到指定路由 - pub fn send_by_key(&self, buf: &[u8], route_key: RouteKey) -> io::Result<()> { + pub fn send_by_key>( + &self, + buf: &NetPacket, + route_key: RouteKey, + ) -> io::Result<()> { match route_key.protocol() { ConnectProtocol::UDP => { if let Some(main_udp) = self.main_udp_socket.get(route_key.index) { - main_udp.send_to(buf, route_key.addr)?; + main_udp.send_to(buf.buffer(), route_key.addr)?; } else { if let Some(udp) = self .sub_udp_socket .read() .get(route_key.index - self.main_udp_socket.len()) { - udp.send_to(buf, route_key.addr)?; + udp.send_to(buf.buffer(), route_key.addr)?; } else { Err(io::Error::from(io::ErrorKind::NotFound))? } } - Ok(()) } ConnectProtocol::TCP | ConnectProtocol::WS | ConnectProtocol::WSS => { - self.send_tcp(buf, route_key.addr) + self.send_tcp(buf.buffer(), route_key.addr)? } } + if let Some(up_traffic_meter) = &self.up_traffic_meter { + up_traffic_meter.add_traffic(buf.destination(), buf.data_len()); + } + Ok(()) } pub fn remove_route(&self, ip: &Ipv4Addr, route_key: RouteKey) { self.route_table.remove_route(ip, route_key) diff --git a/vnt/src/channel/mod.rs b/vnt/src/channel/mod.rs index 3da109b..639ab4c 100644 --- a/vnt/src/channel/mod.rs +++ b/vnt/src/channel/mod.rs @@ -10,6 +10,7 @@ use crate::channel::tcp_channel::tcp_listen; use crate::channel::udp_channel::udp_listen; #[cfg(feature = "ws")] use crate::channel::ws_channel::ws_connect_accept; +use crate::util::limit::TrafficMeterMultiAddress; use crate::util::StopManager; pub mod context; @@ -200,6 +201,8 @@ pub(crate) fn init_context( protocol: ConnectProtocol, packet_loss_rate: Option, packet_delay: u32, + up_traffic_meter: Option, + down_traffic_meter: Option, ) -> anyhow::Result<(ChannelContext, std::net::TcpListener)> { assert!(!ports.is_empty(), "not channel"); let mut udps = Vec::with_capacity(ports.len()); @@ -227,7 +230,9 @@ pub(crate) fn init_context( address, ) }; - + if let Err(e) = socket.set_recv_buffer_size(2 * 1024 * 1024) { + log::warn!("set_recv_buffer_size {:?}", e); + } socket .bind(&address.into()) .with_context(|| format!("bind failed: {}", &address))?; @@ -242,6 +247,8 @@ pub(crate) fn init_context( packet_loss_rate, packet_delay, use_ipv6, + up_traffic_meter, + down_traffic_meter, ); let port = context.main_local_udp_port()?[0]; diff --git a/vnt/src/channel/sender.rs b/vnt/src/channel/sender.rs index e5fbe39..08e9252 100644 --- a/vnt/src/channel/sender.rs +++ b/vnt/src/channel/sender.rs @@ -87,7 +87,7 @@ impl IpPacketSender { if dest_ip.is_broadcast() { //走服务端广播 self.context - .send_default(net_packet.buffer(), device_info.connect_server)?; + .send_default(&net_packet, device_info.connect_server)?; return Ok(()); } @@ -96,7 +96,7 @@ impl IpPacketSender { return Ok(()); } self.context.send_ipv4_by_id( - net_packet.buffer(), + &net_packet, &dest_ip, device_info.connect_server, device_info.status.online(), diff --git a/vnt/src/core/conn.rs b/vnt/src/core/conn.rs index 7ccdd87..dfd35ad 100644 --- a/vnt/src/core/conn.rs +++ b/vnt/src/core/conn.rs @@ -27,13 +27,15 @@ use crate::nat::NatTest; #[cfg(feature = "integrated_tun")] use crate::tun_tap_device::tun_create_helper::{DeviceAdapter, TunDeviceHelper}; use crate::tun_tap_device::vnt_device::DeviceWrite; -use crate::util::{Scheduler, StopManager, U64Adder, WatchU64Adder}; +use crate::util::limit::TrafficMeterMultiAddress; +use crate::util::{Scheduler, StopManager}; use crate::{nat, VntCallback}; #[derive(Clone)] pub struct Vnt { inner: Arc, } + impl Vnt { #[cfg(feature = "integrated_tun")] pub fn new(config: Config, callback: Call) -> anyhow::Result { @@ -50,6 +52,7 @@ impl Vnt { Ok(Self { inner }) } } + impl Deref for Vnt { type Target = VntInner; @@ -57,6 +60,7 @@ impl Deref for Vnt { &self.inner } } + pub struct VntInner { stop_manager: StopManager, config: Config, @@ -65,12 +69,12 @@ pub struct VntInner { device_list: Arc)>>, context: Arc>>, peer_nat_info_map: Arc>>, - down_count_watcher: WatchU64Adder, - up_count_watcher: WatchU64Adder, client_secret_hash: Option<[u8; 16]>, compressor: Compressor, client_cipher: Cipher, external_route: ExternalRoute, + up_traffic_meter: Option, + down_traffic_meter: Option, } impl VntInner { @@ -91,7 +95,15 @@ impl VntInner { callback: Call, device: Device, ) -> anyhow::Result { - log::info!("config.toml:{:?}", config); + log::info!("config: {:?}", config); + let (up_traffic_meter, down_traffic_meter) = if config.enable_traffic { + ( + Some(TrafficMeterMultiAddress::default()), + Some(TrafficMeterMultiAddress::default()), + ) + } else { + (None, None) + }; //服务端非对称加密 #[cfg(feature = "server_encrypt")] let rsa_cipher: Arc>> = Arc::new(Mutex::new(None)); @@ -165,6 +177,8 @@ impl VntInner { config.protocol, config.packet_loss_rate, config.packet_delay, + up_traffic_meter.clone(), + down_traffic_meter.clone(), )?; let local_ipv4 = nat::local_ipv4(); let local_ipv6 = nat::local_ipv6(); @@ -199,14 +213,10 @@ impl VntInner { let (punch_sender, punch_receiver) = maintain::punch_channel(); let peer_nat_info_map: Arc>> = Arc::new(RwLock::new(HashMap::with_capacity(16))); - let down_counter = U64Adder::default(); - let down_count_watcher = down_counter.watch(); let handshake = Handshake::new( #[cfg(feature = "server_encrypt")] rsa_cipher.clone(), ); - let up_counter = U64Adder::default(); - let up_count_watcher = up_counter.watch(); #[cfg(feature = "integrated_tun")] let tun_device_helper = { TunDeviceHelper::new( @@ -218,7 +228,6 @@ impl VntInner { proxy_map.clone(), client_cipher.clone(), server_cipher.clone(), - up_counter, device_list.clone(), config.compressor, device.clone().into_device_adapter(), @@ -243,7 +252,6 @@ impl VntInner { #[cfg(feature = "ip_proxy")] #[cfg(feature = "integrated_tun")] proxy_map.clone(), - down_counter, handshake.clone(), #[cfg(feature = "integrated_tun")] tun_device_helper, @@ -280,8 +288,6 @@ impl VntInner { let context = context.clone(); let nat_test = nat_test.clone(); let device_list = device_list.clone(); - let down_count_watcher = down_count_watcher.clone(); - let up_count_watcher = up_count_watcher.clone(); let config_info = config_info.clone(); let current_device = current_device.clone(); if !config.use_channel_type.is_only_relay() { @@ -308,8 +314,6 @@ impl VntInner { config_info, punch, callback, - down_count_watcher, - up_count_watcher, ); }); } @@ -322,12 +326,12 @@ impl VntInner { device_list, context: Arc::new(Mutex::new(Some(context))), peer_nat_info_map, - down_count_watcher, - up_count_watcher, client_secret_hash: config_info.client_secret_hash, compressor, client_cipher, external_route, + up_traffic_meter, + down_traffic_meter, }) } } @@ -344,8 +348,6 @@ pub fn start( config_info: BaseConfigInfo, punch: Punch, callback: Call, - down_count_watcher: WatchU64Adder, - up_count_watcher: WatchU64Adder, ) { // 定时心跳 maintain::heartbeat( @@ -399,13 +401,7 @@ pub fn start( punch, ); } - maintain::up_status( - scheduler, - context.clone(), - current_device.clone(), - down_count_watcher, - up_count_watcher, - ) + maintain::up_status(scheduler, context.clone(), current_device.clone()) } impl VntInner { @@ -463,10 +459,24 @@ impl VntInner { } } pub fn up_stream(&self) -> u64 { - self.up_count_watcher.get() + self.up_traffic_meter.as_ref().map_or(0, |v| v.total()) + } + pub fn up_stream_all(&self) -> Option<(u64, HashMap)> { + self.up_traffic_meter.as_ref().map(|v| v.get_all()) + } + pub fn up_stream_history(&self) -> Option<(u64, HashMap)>)> { + self.up_traffic_meter.as_ref().map(|v| v.get_all_history()) } pub fn down_stream(&self) -> u64 { - self.down_count_watcher.get() + self.down_traffic_meter.as_ref().map_or(0, |v| v.total()) + } + pub fn down_stream_all(&self) -> Option<(u64, HashMap)> { + self.down_traffic_meter.as_ref().map(|v| v.get_all()) + } + pub fn down_stream_history(&self) -> Option<(u64, HashMap)>)> { + self.down_traffic_meter + .as_ref() + .map(|v| v.get_all_history()) } pub fn stop(&self) { //退出协助回收资源 @@ -505,6 +515,7 @@ impl VntInner { } } } + impl Drop for VntInner { fn drop(&mut self) { self.stop(); diff --git a/vnt/src/core/mod.rs b/vnt/src/core/mod.rs index 04c4a6b..44238fa 100644 --- a/vnt/src/core/mod.rs +++ b/vnt/src/core/mod.rs @@ -50,6 +50,7 @@ pub struct Config { #[cfg(feature = "port_mapping")] pub port_mapping_list: Vec<(bool, SocketAddr, String)>, pub compressor: Compressor, + pub enable_traffic: bool, } impl Config { @@ -86,6 +87,7 @@ impl Config { // 例如 [udp:127.0.0.1:80->10.26.0.10:8080,tcp:127.0.0.1:80->10.26.0.10:8080] #[cfg(feature = "port_mapping")] port_mapping_list: Vec, compressor: Compressor, + enable_traffic: bool, ) -> anyhow::Result { for x in stun_server.iter_mut() { if !x.contains(":") { @@ -177,6 +179,7 @@ impl Config { #[cfg(feature = "port_mapping")] port_mapping_list, compressor, + enable_traffic, }) } } diff --git a/vnt/src/handle/handshaker.rs b/vnt/src/handle/handshaker.rs index 9259072..bd8a401 100644 --- a/vnt/src/handle/handshaker.rs +++ b/vnt/src/handle/handshaker.rs @@ -47,7 +47,7 @@ impl Handshake { } let request_packet = self.handshake_request_packet(secret)?; log::info!("发送握手请求,secret={},{:?}", secret, addr); - context.send_default(request_packet.buffer(), addr)?; + context.send_default(&request_packet, addr)?; self.time.store(Instant::now()); Ok(()) } diff --git a/vnt/src/handle/maintain/heartbeat.rs b/vnt/src/handle/maintain/heartbeat.rs index e5b4a94..f18ac8d 100644 --- a/vnt/src/handle/maintain/heartbeat.rs +++ b/vnt/src/handle/maintain/heartbeat.rs @@ -59,8 +59,7 @@ fn heartbeat0( let mut is_send_gateway = false; match heartbeat_packet_server(device_list, server_cipher, src_ip, gateway_ip) { Ok(net_packet) => { - if let Err(e) = context.send_default(net_packet.buffer(), current_device.connect_server) - { + if let Err(e) = context.send_default(&net_packet, current_device.connect_server) { log::warn!("heartbeat err={:?}", e) } else { is_send_gateway = true @@ -88,7 +87,7 @@ fn heartbeat0( } }; for route in routes { - if let Err(e) = context.send_by_key(net_packet.buffer(), route.route_key()) { + if let Err(e) = context.send_by_key(&net_packet, route.route_key()) { log::warn!("heartbeat err={:?}", e) } } @@ -113,8 +112,7 @@ fn heartbeat0( continue; } }; - if let Err(e) = context.send_default(net_packet.buffer(), current_device.connect_server) - { + if let Err(e) = context.send_default(&net_packet, current_device.connect_server) { log::error!("heartbeat_packet send_default err={:?}", e); } } @@ -195,7 +193,7 @@ fn client_relay0( if current_device.is_gateway(ip) { continue; } - if let Err(e) = context.send_by_key(client_packet.buffer(), route.route_key()) { + if let Err(e) = context.send_by_key(&client_packet, route.route_key()) { log::error!("{:?}", e); } if index >= 2 { diff --git a/vnt/src/handle/maintain/punch.rs b/vnt/src/handle/maintain/punch.rs index ee5490d..69c5ff8 100644 --- a/vnt/src/handle/maintain/punch.rs +++ b/vnt/src/handle/maintain/punch.rs @@ -289,7 +289,7 @@ fn punch0( punch_count, total_count, ); - context.send_default(packet.buffer(), current_device.connect_server)?; + context.send_default(&packet, current_device.connect_server)?; break; } } diff --git a/vnt/src/handle/maintain/up_status.rs b/vnt/src/handle/maintain/up_status.rs index 1f68ade..25b51e7 100644 --- a/vnt/src/handle/maintain/up_status.rs +++ b/vnt/src/handle/maintain/up_status.rs @@ -3,7 +3,7 @@ use crate::handle::CurrentDeviceInfo; use crate::proto::message::{ClientStatusInfo, PunchNatType, RouteItem}; use crate::protocol::body::ENCRYPTION_RESERVED; use crate::protocol::{service_packet, NetPacket, Protocol, HEAD_LEN, MAX_TTL}; -use crate::util::{Scheduler, WatchU64Adder}; +use crate::util::Scheduler; use crossbeam_utils::atomic::AtomicCell; use protobuf::Message; use std::io; @@ -15,17 +15,9 @@ pub fn up_status( scheduler: &Scheduler, context: ChannelContext, current_device_info: Arc>, - down_count_watcher: WatchU64Adder, - up_count_watcher: WatchU64Adder, ) { let _ = scheduler.timeout(Duration::from_secs(60), move |x| { - up_status0( - x, - context, - current_device_info, - down_count_watcher, - up_count_watcher, - ) + up_status0(x, context, current_device_info) }); } @@ -33,25 +25,12 @@ fn up_status0( scheduler: &Scheduler, context: ChannelContext, current_device_info: Arc>, - down_count_watcher: WatchU64Adder, - up_count_watcher: WatchU64Adder, ) { - if let Err(e) = send_up_status_packet( - &context, - ¤t_device_info, - &down_count_watcher, - &up_count_watcher, - ) { + if let Err(e) = send_up_status_packet(&context, ¤t_device_info) { log::warn!("{:?}", e) } let rs = scheduler.timeout(Duration::from_secs(10 * 60), move |x| { - up_status0( - x, - context, - current_device_info, - down_count_watcher, - up_count_watcher, - ) + up_status0(x, context, current_device_info) }); if !rs { log::info!("定时任务停止"); @@ -61,8 +40,6 @@ fn up_status0( fn send_up_status_packet( context: &ChannelContext, current_device_info: &AtomicCell, - down_count_watcher: &WatchU64Adder, - up_count_watcher: &WatchU64Adder, ) -> io::Result<()> { let device_info = current_device_info.load(); if device_info.status.offline() { @@ -79,8 +56,8 @@ fn send_up_status_packet( item.next_ip = ip.into(); message.p2p_list.push(item); } - message.up_stream = up_count_watcher.get(); - message.down_stream = down_count_watcher.get(); + message.up_stream = context.up_traffic_meter.as_ref().map_or(0, |v| v.total()); + message.down_stream = context.down_traffic_meter.as_ref().map_or(0, |v| v.total()); message.nat_type = protobuf::EnumOrUnknown::new(if context.is_cone() { PunchNatType::Cone } else { @@ -99,6 +76,6 @@ fn send_up_status_packet( net_packet.set_source(device_info.virtual_ip); net_packet.set_destination(device_info.virtual_gateway); net_packet.set_payload(&buf)?; - context.send_default(net_packet.buffer(), device_info.connect_server)?; + context.send_default(&net_packet, device_info.connect_server)?; Ok(()) } diff --git a/vnt/src/handle/recv_data/client.rs b/vnt/src/handle/recv_data/client.rs index 5970867..54a7b93 100644 --- a/vnt/src/handle/recv_data/client.rs +++ b/vnt/src/handle/recv_data/client.rs @@ -140,7 +140,7 @@ impl ClientPacketHandler { net_packet.set_destination(source); //不管加不加密，和接收到的数据长度都一致 self.client_cipher.encrypt_ipv4(&mut net_packet)?; - context.send_by_key(net_packet.buffer(), route_key)?; + context.send_by_key(&net_packet, route_key)?; return Ok(()); } } @@ -217,7 +217,7 @@ impl ClientPacketHandler { net_packet.set_destination(source); net_packet.first_set_ttl(MAX_TTL); self.client_cipher.encrypt_ipv4(&mut net_packet)?; - context.send_by_key(net_packet.buffer(), route_key)?; + context.send_by_key(&net_packet, route_key)?; let route = Route::from_default_rt(route_key, metric); context.route_table.add_route_if_absent(source, route); } @@ -249,7 +249,7 @@ impl ClientPacketHandler { net_packet.set_destination(source); net_packet.first_set_ttl(1); self.client_cipher.encrypt_ipv4(&mut net_packet)?; - context.send_by_key(net_packet.buffer(), route_key)?; + context.send_by_key(&net_packet, route_key)?; // 收到PunchRequest就添加路由，会导致单向通信的问题，删掉试试 // let route = Route::from_default_rt(route_key, 1); // context.route_table.add_route_if_absent(source, route); @@ -281,7 +281,7 @@ impl ClientPacketHandler { addr_packet.set_ipv4(ipv4); addr_packet.set_port(route_key.addr.port()); self.client_cipher.encrypt_ipv4(&mut packet)?; - context.send_by_key(packet.buffer(), route_key)?; + context.send_by_key(&packet, route_key)?; } std::net::IpAddr::V6(_) => {} }, @@ -377,7 +377,7 @@ impl ClientPacketHandler { punch_packet.set_payload(&bytes)?; self.client_cipher.encrypt_ipv4(&mut punch_packet)?; if self.punch_sender.send(true, source, peer_nat_info) { - context.send_by_key(punch_packet.buffer(), route_key)?; + context.send_by_key(&punch_packet, route_key)?; } } else { self.punch_sender.send(false, source, peer_nat_info); diff --git a/vnt/src/handle/recv_data/mod.rs b/vnt/src/handle/recv_data/mod.rs index a23b185..e2bd5d6 100644 --- a/vnt/src/handle/recv_data/mod.rs +++ b/vnt/src/handle/recv_data/mod.rs @@ -26,7 +26,6 @@ use crate::ip_proxy::IpProxyMap; use crate::nat::NatTest; use crate::protocol::{NetPacket, HEAD_LEN}; use crate::tun_tap_device::vnt_device::DeviceWrite; -use crate::util::U64Adder; mod client; mod server; @@ -38,7 +37,6 @@ pub struct RecvDataHandler { turn: TurnPacketHandler, client: ClientPacketHandler, server: ServerPacketHandler, - counter: U64Adder, nat_test: NatTest, } @@ -93,7 +91,6 @@ impl RecvDataHandler { #[cfg(feature = "integrated_tun")] #[cfg(feature = "ip_proxy")] ip_proxy_map: Option, - counter: U64Adder, handshake: Handshake, #[cfg(feature = "integrated_tun")] tun_device_helper: crate::tun_tap_device::tun_create_helper::TunDeviceHelper, @@ -130,7 +127,6 @@ impl RecvDataHandler { turn, client, server, - counter, nat_test, } } @@ -141,9 +137,8 @@ impl RecvDataHandler { route_key: RouteKey, context: &ChannelContext, ) -> anyhow::Result<()> { - // 统计流量 - self.counter.add(buf.len() as _); let net_packet = NetPacket::new(buf)?; + let extend = NetPacket::unchecked(extend); if net_packet.ttl() == 0 || net_packet.source_ttl() < net_packet.ttl() { log::warn!("丢弃过时包:{:?} {}", net_packet.head(), route_key.addr); @@ -158,6 +153,10 @@ impl RecvDataHandler { || dest.is_unspecified() || dest == current_device.broadcast_ip { + // 统计流量 + if let Some(down_traffic_meter) = &context.down_traffic_meter { + down_traffic_meter.add_traffic(net_packet.source(), net_packet.data_len()) + } //发给自己的包 if net_packet.is_gateway() { //服务端-客户端包 diff --git a/vnt/src/handle/recv_data/server.rs b/vnt/src/handle/recv_data/server.rs index 8452bec..66ecfc0 100644 --- a/vnt/src/handle/recv_data/server.rs +++ b/vnt/src/handle/recv_data/server.rs @@ -140,7 +140,7 @@ impl PacketHandler for ServerPacketHandl self.config_info.token.clone(), key, )?; - context.send_by_key(packet.buffer(), route_key)?; + context.send_by_key(&packet, route_key)?; } } } @@ -164,7 +164,7 @@ impl PacketHandler for ServerPacketHandl key, )?; drop(guard); - context.send_by_key(packet.buffer(), route_key)?; + context.send_by_key(&packet, route_key)?; return Ok(()); } log::warn!( @@ -199,7 +199,7 @@ impl PacketHandler for ServerPacketHandl self.config_info.token.clone(), key, )?; - context.send_by_key(packet.buffer(), route_key)?; + context.send_by_key(&packet, route_key)?; self.rsa_cipher.lock().replace(rsa_cipher); } return Ok(()); @@ -484,7 +484,7 @@ impl ServerPacketHandler { )?; log::info!("发送注册请求，{:?}", self.config_info); //注册请求只发送到默认通道 - context.send_default(response.buffer(), current_device.connect_server)?; + context.send_default(&response, current_device.connect_server)?; Ok(()) } fn error( @@ -568,7 +568,7 @@ impl ServerPacketHandler { .set_transport_protocol(service_packet::Protocol::PullDeviceList.into()); self.server_cipher.encrypt_ipv4(&mut poll_device)?; //发送到默认服务端即可 - context.send_default(poll_device.buffer(), current_device.connect_server)?; + context.send_default(&poll_device, current_device.connect_server)?; } } ControlPacket::AddrResponse(addr_packet) => { diff --git a/vnt/src/handle/recv_data/turn.rs b/vnt/src/handle/recv_data/turn.rs index c356e68..59f82d9 100644 --- a/vnt/src/handle/recv_data/turn.rs +++ b/vnt/src/handle/recv_data/turn.rs @@ -40,7 +40,7 @@ impl PacketHandler for TurnPacketHandler { } if route.metric <= ttl { return context - .send_by_key(net_packet.buffer(), route.route_key()) + .send_by_key(&net_packet, route.route_key()) .context("转发失败"); } } diff --git a/vnt/src/handle/tun_tap/mod.rs b/vnt/src/handle/tun_tap/mod.rs index 9dd4e9e..abbe944 100644 --- a/vnt/src/handle/tun_tap/mod.rs +++ b/vnt/src/handle/tun_tap/mod.rs @@ -37,7 +37,7 @@ impl DeviceStop { pub fn stopped(&self) { self.stopped.store(true); } - pub fn is_stop(&self) -> bool { + pub fn is_stopped(&self) -> bool { self.stopped.load() } } diff --git a/vnt/src/handle/tun_tap/tun_handler.rs b/vnt/src/handle/tun_tap/tun_handler.rs index 52bf5bc..6aee8f2 100644 --- a/vnt/src/handle/tun_tap/tun_handler.rs +++ b/vnt/src/handle/tun_tap/tun_handler.rs @@ -26,7 +26,7 @@ use crate::protocol; use crate::protocol::body::ENCRYPTION_RESERVED; use crate::protocol::ip_turn_packet::BroadcastPacket; use crate::protocol::{ip_turn_packet, NetPacket, MAX_TTL}; -use crate::util::{StopManager, U64Adder}; +use crate::util::StopManager; fn icmp(device_writer: &Device, mut ipv4_packet: IpV4Packet<&mut [u8]>) -> anyhow::Result<()> { if ipv4_packet.protocol() == Protocol::Icmp { @@ -53,7 +53,6 @@ pub fn start( #[cfg(feature = "ip_proxy")] ip_proxy_map: Option, client_cipher: Cipher, server_cipher: Cipher, - up_counter: U64Adder, device_list: Arc)>>, compressor: Compressor, device_stop: DeviceStop, @@ -71,7 +70,6 @@ pub fn start( ip_proxy_map, client_cipher, server_cipher, - &up_counter, device_list, compressor, device_stop, @@ -107,10 +105,7 @@ fn broadcast( break; } if let Some(route) = sender.route_table.route_one_p2p(&peer_ip) { - if sender - .send_by_key(net_packet.buffer(), route.route_key()) - .is_ok() - { + if sender.send_by_key(&net_packet, route.route_key()).is_ok() { p2p_ips.push(peer_ip); continue; } @@ -125,7 +120,7 @@ fn broadcast( if p2p_ips.is_empty() { //都没有p2p则直接由服务器转发 if current_device.status.online() { - sender.send_default(net_packet.buffer(), current_device.connect_server)?; + sender.send_default(&net_packet, current_device.connect_server)?; } return Ok(()); } @@ -135,7 +130,7 @@ fn broadcast( //非直连的广播要改变目的地址，不然服务端收到了会再次广播 net_packet.set_destination(peer_ip); sender.send_ipv4_by_id( - net_packet.buffer(), + &net_packet, &peer_ip, current_device.connect_server, current_device.status.online(), @@ -163,7 +158,7 @@ fn broadcast( broadcast.set_address(&p2p_ips)?; broadcast.set_data(net_packet.buffer())?; server_cipher.encrypt_ipv4(&mut server_packet)?; - sender.send_default(server_packet.buffer(), current_device.connect_server)?; + sender.send_default(&server_packet, current_device.connect_server)?; Ok(()) } @@ -211,7 +206,7 @@ pub(crate) fn handle( if protocol == Protocol::Icmp { net_packet.set_gateway_flag(true); server_cipher.encrypt_ipv4(&mut net_packet)?; - context.send_default(net_packet.buffer(), current_device.connect_server)?; + context.send_default(&net_packet, current_device.connect_server)?; } return Ok(()); } @@ -269,7 +264,7 @@ pub(crate) fn handle( client_cipher.encrypt_ipv4(&mut net_packet)?; context.send_ipv4_by_id( - net_packet.buffer(), + &net_packet, &dest_ip, current_device.connect_server, current_device.status.online(), diff --git a/vnt/src/handle/tun_tap/unix.rs b/vnt/src/handle/tun_tap/unix.rs index a71a23a..53d11f7 100644 --- a/vnt/src/handle/tun_tap/unix.rs +++ b/vnt/src/handle/tun_tap/unix.rs @@ -7,7 +7,7 @@ use crate::handle::tun_tap::DeviceStop; use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo}; #[cfg(feature = "ip_proxy")] use crate::ip_proxy::IpProxyMap; -use crate::util::{StopManager, U64Adder}; +use crate::util::StopManager; use crossbeam_utils::atomic::AtomicCell; use mio::event::Source; use mio::unix::SourceFd; @@ -30,7 +30,6 @@ pub(crate) fn start_simple( #[cfg(feature = "ip_proxy")] ip_proxy_map: Option, client_cipher: Cipher, server_cipher: Cipher, - up_counter: &U64Adder, device_list: Arc)>>, compressor: Compressor, device_stop: DeviceStop, @@ -62,7 +61,6 @@ pub(crate) fn start_simple( ip_proxy_map, client_cipher, server_cipher, - up_counter, device_list, compressor, ) { @@ -85,7 +83,6 @@ fn start_simple0( #[cfg(feature = "ip_proxy")] ip_proxy_map: Option, client_cipher: Cipher, server_cipher: Cipher, - up_counter: &U64Adder, device_list: Arc)>>, compressor: Compressor, ) -> anyhow::Result<()> { @@ -118,8 +115,6 @@ fn start_simple0( Err(e)? } }; - //单线程的 - up_counter.add(len as u64); // buf是重复利用的，需要重置头部 buf[..12].fill(0); match crate::handle::tun_tap::tun_handler::handle( diff --git a/vnt/src/handle/tun_tap/windows.rs b/vnt/src/handle/tun_tap/windows.rs index c462d1f..0df3e4a 100644 --- a/vnt/src/handle/tun_tap/windows.rs +++ b/vnt/src/handle/tun_tap/windows.rs @@ -7,7 +7,7 @@ use crate::handle::tun_tap::DeviceStop; use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo}; #[cfg(feature = "ip_proxy")] use crate::ip_proxy::IpProxyMap; -use crate::util::{StopManager, U64Adder}; +use crate::util::StopManager; use crossbeam_utils::atomic::AtomicCell; use parking_lot::Mutex; use std::sync::Arc; @@ -23,7 +23,6 @@ pub(crate) fn start_simple( #[cfg(feature = "ip_proxy")] ip_proxy_map: Option, client_cipher: Cipher, server_cipher: Cipher, - up_counter: &U64Adder, device_list: Arc)>>, compressor: Compressor, device_stop: DeviceStop, @@ -55,7 +54,6 @@ pub(crate) fn start_simple( ip_proxy_map, client_cipher, server_cipher, - up_counter, device_list, compressor, ) { @@ -76,7 +74,6 @@ fn start_simple0( #[cfg(feature = "ip_proxy")] ip_proxy_map: Option, client_cipher: Cipher, server_cipher: Cipher, - up_counter: &U64Adder, device_list: Arc)>>, compressor: Compressor, ) -> anyhow::Result<()> { @@ -85,7 +82,6 @@ fn start_simple0( loop { let len = device.read(&mut buf[12..])? + 12; //单线程的 - up_counter.add(len as u64); // buf是重复利用的，需要重置头部 buf[..12].fill(0); match crate::handle::tun_tap::tun_handler::handle( diff --git a/vnt/src/ip_proxy/icmp_proxy.rs b/vnt/src/ip_proxy/icmp_proxy.rs index 93557cf..ffb7d2e 100644 --- a/vnt/src/ip_proxy/icmp_proxy.rs +++ b/vnt/src/ip_proxy/icmp_proxy.rs @@ -155,7 +155,7 @@ fn recv_handle( return; } if let Err(e) = context.send_ipv4_by_id( - net_packet.buffer(), + &net_packet, &dest_ip, current_device.connect_server, current_device.status.online(), diff --git a/vnt/src/tun_tap_device/tun_create_helper.rs b/vnt/src/tun_tap_device/tun_create_helper.rs index 09a1a97..76babcf 100644 --- a/vnt/src/tun_tap_device/tun_create_helper.rs +++ b/vnt/src/tun_tap_device/tun_create_helper.rs @@ -16,7 +16,7 @@ use crate::handle::{CurrentDeviceInfo, PeerDeviceInfo}; #[cfg(feature = "ip_proxy")] use crate::ip_proxy::IpProxyMap; use crate::tun_tap_device::vnt_device::DeviceWrite; -use crate::util::{StopManager, U64Adder}; +use crate::util::StopManager; #[repr(transparent)] #[derive(Clone, Default)] @@ -67,7 +67,6 @@ struct TunDeviceHelperInner { ip_proxy_map: Option, client_cipher: Cipher, server_cipher: Cipher, - up_counter: U64Adder, device_list: Arc)>>, compressor: Compressor, } @@ -81,7 +80,6 @@ impl TunDeviceHelper { #[cfg(feature = "ip_proxy")] ip_proxy_map: Option, client_cipher: Cipher, server_cipher: Cipher, - up_counter: U64Adder, device_list: Arc)>>, compressor: Compressor, device_adapter: DeviceAdapter, @@ -95,7 +93,6 @@ impl TunDeviceHelper { ip_proxy_map, client_cipher, server_cipher, - up_counter, device_list, compressor, }; @@ -113,7 +110,7 @@ impl TunDeviceHelper { device_stop.stop(); std::thread::sleep(std::time::Duration::from_millis(300)); //确保停止了 - if device_stop.is_stop() { + if device_stop.is_stopped() { break; } } @@ -136,7 +133,6 @@ impl TunDeviceHelper { inner.ip_proxy_map, inner.client_cipher, inner.server_cipher, - inner.up_counter, inner.device_list, inner.compressor, device_stop, diff --git a/vnt/src/util/mod.rs b/vnt/src/util/mod.rs index 40a4cec..fade277 100644 --- a/vnt/src/util/mod.rs +++ b/vnt/src/util/mod.rs @@ -3,8 +3,8 @@ mod scheduler; pub use notify::{StopManager, Worker}; pub use scheduler::Scheduler; -mod counter; -pub use counter::*; +// mod counter; +// pub use counter::*; mod dns_query; pub use dns_query::*; @@ -13,3 +13,5 @@ pub use dns_query::*; mod upnp; #[cfg(feature = "upnp")] pub use upnp::*; + +pub mod limit; diff --git a/vnt/tun/src/windows/tun/mod.rs b/vnt/tun/src/windows/tun/mod.rs index f67f6f0..6d48f87 100644 --- a/vnt/tun/src/windows/tun/mod.rs +++ b/vnt/tun/src/windows/tun/mod.rs @@ -101,7 +101,7 @@ impl Device { )); } // 开启session - let session = win_tun.WintunStartSession(adapter, 128 * 1024); + let session = win_tun.WintunStartSession(adapter, 4 * 1024 * 1024); if session.is_null() { log::error!("session.is_null {:?}", io::Error::last_os_error()); return Err(io::Error::new(