From 7098111ad1a2825163f9dd272776d5ef55c81be4 Mon Sep 17 00:00:00 2001 From: lubeilin <1791778603@qq.com> Date: Sat, 6 Jan 2024 12:06:21 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0tcp=E9=80=9A=E9=81=93?= =?UTF-8?q?=EF=BC=8C=E5=8E=BB=E9=99=A4=E5=B9=B6=E8=A1=8C=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- vnt/proto/message.proto | 1 + vnt/src/channel/channel.rs | 159 ++++++++++++------------------------- vnt/src/channel/mod.rs | 15 +++- vnt/src/core/mod.rs | 2 +- vnt/src/proto/message.rs | 26 +++++- 5 files changed, 85 insertions(+), 118 deletions(-) diff --git a/vnt/proto/message.proto b/vnt/proto/message.proto index 3a0e6b8..cc62cd3 100644 --- a/vnt/proto/message.proto +++ b/vnt/proto/message.proto @@ -56,6 +56,7 @@ message PunchInfo{ uint32 local_port = 8; bytes ipv6 = 9; uint32 ipv6_port = 10; + uint32 tcp_port = 11; } enum PunchNatType{ Symmetric = 0; diff --git a/vnt/src/channel/channel.rs b/vnt/src/channel/channel.rs index a00ab47..792c20c 100644 --- a/vnt/src/channel/channel.rs +++ b/vnt/src/channel/channel.rs @@ -3,6 +3,7 @@ use std::io::{Read, Write}; use std::net::UdpSocket as StdUdpSocket; use std::net::{Ipv4Addr, Shutdown, SocketAddr}; use std::net::{SocketAddrV6, TcpStream}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use std::{io, thread}; @@ -27,9 +28,11 @@ pub struct ContextInner { pub(crate) status_receiver: Receiver, pub(crate) status_sender: Sender, pub(crate) udp_map: RwLock>>, + pub(crate) tcp_map: RwLock>>>, pub(crate) channel_num: usize, current_device: Arc>, first_latency: bool, + is_close: AtomicBool, } #[derive(Clone)] @@ -56,9 +59,11 @@ impl Context { status_receiver, status_sender, udp_map: RwLock::new(HashMap::with_capacity(16)), + tcp_map: RwLock::new(HashMap::with_capacity(16)), channel_num, current_device, first_latency, + is_close: AtomicBool::new(false), }); Self { inner } } @@ -66,21 +71,38 @@ impl Context { impl Context { pub fn is_close(&self) -> bool { - *self.inner.status_receiver.borrow() == Status::Close + self.inner.is_close.load(Ordering::Relaxed) } pub fn is_cone(&self) -> bool { *self.inner.status_receiver.borrow() == Status::Cone } pub fn close(&self) -> io::Result<()> { + self.inner.is_close.store(true, Ordering::Release); let _ = self.inner.status_sender.send(Status::Close); if let Ok(port) = self.main_local_udp_port() { - let _ = StdUdpSocket::bind("127.0.0.1:0")?.send_to( - b"stop", - SocketAddr::V4(std::net::SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)), - ); + match StdUdpSocket::bind("127.0.0.1:0") { + Ok(udp) => { + if let Err(e) = udp.send_to( + b"stop", + SocketAddr::V4(std::net::SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)), + ) { + log::error!("发送停止消息到udp失败:{:?}", e); + } + } + Err(e) => { + log::error!("发送停止-绑定udp失败:{:?}", e); + } + } } if let Some(tcp) = &self.inner.main_tcp_channel { - tcp.lock().shutdown(Shutdown::Both)?; + if let Err(e) = tcp.lock().shutdown(Shutdown::Both) { + log::error!("发送停止消息到tcp失败:{:?}", e); + } + } + for (_, tcp) in self.inner.tcp_map.read().clone() { + if let Err(e) = tcp.lock().shutdown(Shutdown::Both) { + log::error!("发送停止消息到tcp失败:{:?}", e); + } } Ok(()) } @@ -371,37 +393,6 @@ impl Channel { } } -#[derive(Clone)] -struct BufSenderGroup( - usize, - Vec, usize, usize, RouteKey)>>, -); - -struct BufReceiverGroup(Vec, usize, usize, RouteKey)>>); - -impl BufSenderGroup { - pub fn send(&mut self, val: (Vec, usize, usize, RouteKey)) -> bool { - let index = self.0 % self.1.len(); - self.0 = self.0.wrapping_add(1); - self.1[index].send(val).is_ok() - } -} - -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) = - std::sync::mpsc::sync_channel::<(Vec, usize, usize, RouteKey)>(1); - buf_sender_group.push(buf_sender); - buf_receiver_group.push(buf_receiver); - } - ( - BufSenderGroup(0, buf_sender_group), - BufReceiverGroup(buf_receiver_group), - ) -} - impl Channel { fn tcp_handle( tcp_r: &mut TcpStream, @@ -411,7 +402,7 @@ impl Channel { ) -> io::Result<()> { let mut head = [0; 4]; let addr = tcp_r.peer_addr()?; - let key = RouteKey::new(TCP_ID, addr); + let key = RouteKey::new(true, TCP_ID, addr); loop { if context.is_close() { return Ok(()); @@ -480,32 +471,10 @@ impl Channel { head_reserve: usize, //头部预留字节 symmetric_channel_num: usize, //对称网络，则再加一组监听，提升打洞成功率 relay: bool, - parallel: usize, ) { let handler = self.handler.clone(); let context = self.context; let main_channel = context.inner.main_channel.try_clone().unwrap(); - let buf_sender = if parallel > 1 { - let (buf_sender, buf_receiver) = buf_channel_group(parallel); - let mut num = 0; - for buf_receiver in buf_receiver.0 { - let context = context.clone(); - let handler = handler.clone(); - thread::Builder::new() - .name(format!("recv-handler-{}", num)) - .spawn(move || { - while let Ok((mut buf, start, end, route_key)) = buf_receiver.recv() { - handler.handle(&mut buf, start, end, route_key, &context); - } - log::warn!("异步处理停止"); - }) - .unwrap(); - num += 1; - } - Some(buf_sender) - } else { - None - }; if let Some(tcp_stream) = tcp { let context = context.clone(); let handler = handler.clone(); @@ -527,15 +496,7 @@ impl Channel { .name("channel_udp".into()) .spawn(move || { log::info!("启动udp v4"); - Self::main_start_( - worker, - context, - UDP_ID, - main_channel, - handler, - buf_sender, - head_reserve, - ) + Self::main_start_(worker, context, UDP_ID, main_channel, handler, head_reserve) }) .unwrap(); } @@ -597,52 +558,30 @@ impl Channel { id: usize, udp: StdUdpSocket, handler: ChannelDataHandler, - buf_sender: Option, head_reserve: usize, ) { - match buf_sender { - None => { - let mut buf = [0; 4096]; - loop { - match udp.recv_from(&mut buf[head_reserve..]) { - Ok((len, addr)) => { - let end = head_reserve + len; - if &buf[head_reserve..end] == b"stop" { - if context.is_close() { - break; - } - } - handler.handle( - &mut buf, - head_reserve, - end, - RouteKey::new(id, addr), - &context, - ); - } - Err(e) => { - log::error!("udp :{:?}", e); + let mut buf = [0; 4096]; + loop { + match udp.recv_from(&mut buf[head_reserve..]) { + Ok((len, addr)) => { + let end = head_reserve + len; + if &buf[head_reserve..end] == b"stop" { + if context.is_close() { + break; } } + handler.handle( + &mut buf, + head_reserve, + end, + RouteKey::new(false, id, addr), + &context, + ); + } + Err(e) => { + log::error!("udp :{:?}", e); } } - Some(mut buf_sender) => loop { - let mut buf = vec![0; 4096]; - match udp.recv_from(&mut buf[head_reserve..]) { - Ok((len, addr)) => { - let end = head_reserve + len; - if &buf[head_reserve..end] == b"stop" { - if context.is_close() { - break; - } - } - buf_sender.send((buf, head_reserve, end, RouteKey::new(id, addr))); - } - Err(e) => { - log::error!("udp :{:?}", e); - } - } - }, } worker.stop_all(); @@ -671,7 +610,7 @@ impl Channel { rs=udp.recv_from(&mut buf[head_reserve..])=>{ match rs { Ok((len, addr)) => { - handler.handle(&mut buf, head_reserve, head_reserve + len, RouteKey::new(id, addr), &context); + handler.handle(&mut buf, head_reserve, head_reserve + len, RouteKey::new(false,id, addr), &context); } Err(e) => { log::error!("{:?}",e) diff --git a/vnt/src/channel/mod.rs b/vnt/src/channel/mod.rs index 0909a47..7bdca0c 100644 --- a/vnt/src/channel/mod.rs +++ b/vnt/src/channel/mod.rs @@ -17,6 +17,7 @@ pub enum Status { #[derive(Copy, Clone, Debug)] pub struct Route { + is_tcp: bool, index: usize, pub addr: SocketAddr, pub metric: u8, @@ -30,8 +31,9 @@ pub struct RouteSortKey { } impl Route { - pub fn new(index: usize, addr: SocketAddr, metric: u8, rt: i64) -> Self { + pub fn new(is_tcp: bool, index: usize, addr: SocketAddr, metric: u8, rt: i64) -> Self { Self { + is_tcp, index, addr, metric, @@ -40,6 +42,7 @@ impl Route { } pub fn from(route_key: RouteKey, metric: u8, rt: i64) -> Self { Self { + is_tcp: route_key.is_tcp, index: route_key.index, addr: route_key.addr, metric, @@ -48,6 +51,7 @@ impl Route { } pub fn route_key(&self) -> RouteKey { RouteKey { + is_tcp: self.is_tcp, index: self.index, addr: self.addr, } @@ -65,13 +69,18 @@ impl Route { #[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)] pub struct RouteKey { + is_tcp: bool, index: usize, pub addr: SocketAddr, } impl RouteKey { - pub(crate) fn new(index: usize, addr: SocketAddr) -> Self { - Self { index, addr } + pub(crate) fn new(is_tcp: bool, index: usize, addr: SocketAddr) -> Self { + Self { + is_tcp, + index, + addr, + } } pub fn is_tcp(&self) -> bool { self.index == TCP_ID diff --git a/vnt/src/core/mod.rs b/vnt/src/core/mod.rs index c4b8988..c07c4d1 100644 --- a/vnt/src/core/mod.rs +++ b/vnt/src/core/mod.rs @@ -382,7 +382,7 @@ impl VntUtil { let relay = config.relay; tokio::spawn(async move { channel - .start(channel_worker, tcp_receiver, 14, 65, relay, config.parallel) + .start(channel_worker, tcp_receiver, 14, 65, relay) .await }); } diff --git a/vnt/src/proto/message.rs b/vnt/src/proto/message.rs index 686d930..451e054 100644 --- a/vnt/src/proto/message.rs +++ b/vnt/src/proto/message.rs @@ -1317,6 +1317,8 @@ pub struct PunchInfo { pub ipv6: ::std::vec::Vec, // @@protoc_insertion_point(field:PunchInfo.ipv6_port) pub ipv6_port: u32, + // @@protoc_insertion_point(field:PunchInfo.tcp_port) + pub tcp_port: u32, // special fields // @@protoc_insertion_point(special_field:PunchInfo.special_fields) pub special_fields: ::protobuf::SpecialFields, @@ -1334,7 +1336,7 @@ impl PunchInfo { } fn generated_message_descriptor_data() -> ::protobuf::reflect::GeneratedMessageDescriptorData { - let mut fields = ::std::vec::Vec::with_capacity(9); + let mut fields = ::std::vec::Vec::with_capacity(10); let mut oneofs = ::std::vec::Vec::with_capacity(0); fields.push(::protobuf::reflect::rt::v2::make_vec_simpler_accessor::<_, _>( "public_ip_list", @@ -1381,6 +1383,11 @@ impl PunchInfo { |m: &PunchInfo| { &m.ipv6_port }, |m: &mut PunchInfo| { &mut m.ipv6_port }, )); + fields.push(::protobuf::reflect::rt::v2::make_simpler_field_accessor::<_, _>( + "tcp_port", + |m: &PunchInfo| { &m.tcp_port }, + |m: &mut PunchInfo| { &mut m.tcp_port }, + )); ::protobuf::reflect::GeneratedMessageDescriptorData::new_2::( "PunchInfo", fields, @@ -1429,6 +1436,9 @@ impl ::protobuf::Message for PunchInfo { 80 => { self.ipv6_port = is.read_uint32()?; }, + 88 => { + self.tcp_port = is.read_uint32()?; + }, tag => { ::protobuf::rt::read_unknown_or_skip_group(tag, is, self.special_fields.mut_unknown_fields())?; }, @@ -1466,6 +1476,9 @@ impl ::protobuf::Message for PunchInfo { if self.ipv6_port != 0 { my_size += ::protobuf::rt::uint32_size(10, self.ipv6_port); } + if self.tcp_port != 0 { + my_size += ::protobuf::rt::uint32_size(11, self.tcp_port); + } my_size += ::protobuf::rt::unknown_fields_size(self.special_fields.unknown_fields()); self.special_fields.cached_size().set(my_size as u32); my_size @@ -1499,6 +1512,9 @@ impl ::protobuf::Message for PunchInfo { if self.ipv6_port != 0 { os.write_uint32(10, self.ipv6_port)?; } + if self.tcp_port != 0 { + os.write_uint32(11, self.tcp_port)?; + } os.write_unknown_fields(self.special_fields.unknown_fields())?; ::std::result::Result::Ok(()) } @@ -1525,6 +1541,7 @@ impl ::protobuf::Message for PunchInfo { self.local_port = 0; self.ipv6.clear(); self.ipv6_port = 0; + self.tcp_port = 0; self.special_fields.clear(); } @@ -1539,6 +1556,7 @@ impl ::protobuf::Message for PunchInfo { local_port: 0, ipv6: ::std::vec::Vec::new(), ipv6_port: 0, + tcp_port: 0, special_fields: ::protobuf::SpecialFields::new(), }; &instance @@ -1643,7 +1661,7 @@ static file_descriptor_proto_data: &'static [u8] = b"\ \n\rdevice_status\x18\x03\x20\x01(\rR\x0cdeviceStatus\x12#\n\rclient_sec\ ret\x18\x04\x20\x01(\x08R\x0cclientSecret\"Y\n\nDeviceList\x12\x14\n\x05\ epoch\x18\x01\x20\x01(\rR\x05epoch\x125\n\x10device_info_list\x18\x02\ - \x20\x03(\x0b2\x0b.DeviceInfoR\x0edeviceInfoList\"\xa9\x02\n\tPunchInfo\ + \x20\x03(\x0b2\x0b.DeviceInfoR\x0edeviceInfoList\"\xc4\x02\n\tPunchInfo\ \x12$\n\x0epublic_ip_list\x18\x02\x20\x03(\x07R\x0cpublicIpList\x12\x1f\ \n\x0bpublic_port\x18\x03\x20\x01(\rR\npublicPort\x12*\n\x11public_port_\ range\x18\x04\x20\x01(\rR\x0fpublicPortRange\x12(\n\x08nat_type\x18\x05\ @@ -1651,8 +1669,8 @@ static file_descriptor_proto_data: &'static [u8] = b"\ \x01(\x08R\x05reply\x12\x19\n\x08local_ip\x18\x07\x20\x01(\x07R\x07local\ Ip\x12\x1d\n\nlocal_port\x18\x08\x20\x01(\rR\tlocalPort\x12\x12\n\x04ipv\ 6\x18\t\x20\x01(\x0cR\x04ipv6\x12\x1b\n\tipv6_port\x18\n\x20\x01(\rR\x08\ - ipv6Port*'\n\x0cPunchNatType\x12\r\n\tSymmetric\x10\0\x12\x08\n\x04Cone\ - \x10\x01b\x06proto3\ + ipv6Port\x12\x19\n\x08tcp_port\x18\x0b\x20\x01(\rR\x07tcpPort*'\n\x0cPun\ + chNatType\x12\r\n\tSymmetric\x10\0\x12\x08\n\x04Cone\x10\x01b\x06proto3\ "; /// `FileDescriptorProto` object which was a source for this generated file