From cca91d433186e139d9505a029fc1d01a52b1d685 Mon Sep 17 00:00:00 2001 From: lubeilin <1791778603@qq.com> Date: Fri, 6 Oct 2023 22:21:29 +0800 Subject: [PATCH] =?UTF-8?q?tcp=E6=94=B9=E4=B8=BA=E4=BD=BF=E7=94=A8?= =?UTF-8?q?=E5=90=8C=E6=AD=A5=E6=96=B9=E6=B3=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- vnt-cli/src/main.rs | 10 +- vnt/src/channel/channel.rs | 172 ++++++++++++++----------- vnt/src/core/mod.rs | 25 ++-- vnt/src/core/sync.rs | 10 +- vnt/src/handle/handshake_handler.rs | 29 ++--- vnt/src/handle/registration_handler.rs | 12 +- 6 files changed, 139 insertions(+), 119 deletions(-) diff --git a/vnt-cli/src/main.rs b/vnt-cli/src/main.rs index 938f61c..45d7323 100644 --- a/vnt-cli/src/main.rs +++ b/vnt-cli/src/main.rs @@ -292,23 +292,23 @@ fn main() { #[tokio::main] 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 vnt_util = VntUtil::new(config).unwrap(); let mut conn_count = 0; let response = loop { if conn_count > 0 { tokio::time::sleep(std::time::Duration::from_secs(2)).await; } conn_count += 1; - if let Err(e) = vnt_util.connect().await { + if let Err(e) = vnt_util.connect() { println!("connect server failed {}", e); return; } - match vnt_util.handshake().await { + match vnt_util.handshake() { Ok(response) => { if server_encrypt { let finger = response.unwrap().finger().unwrap(); println!("{}{}", green("server fingerprint:".to_string()), finger); - match vnt_util.secret_handshake().await { + match vnt_util.secret_handshake() { Ok(_) => {} Err(e) => { match e { @@ -328,7 +328,7 @@ async fn main0(config: Config, show_cmd: bool) { } } } - match vnt_util.register().await { + match vnt_util.register() { Ok(response) => { break response; } diff --git a/vnt/src/channel/channel.rs b/vnt/src/channel/channel.rs index 130d421..026cf68 100644 --- a/vnt/src/channel/channel.rs +++ b/vnt/src/channel/channel.rs @@ -1,18 +1,18 @@ use std::collections::HashMap; -use std::io; -use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::net::UdpSocket as StdUdpSocket; +use std::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr}; use std::ops::Sub; use std::sync::atomic::Ordering; use std::sync::Arc; use std::time::{Duration, Instant}; +use std::{io, thread}; use crossbeam_epoch::{Atomic, Owned}; use crossbeam_utils::atomic::AtomicCell; use dashmap::DashMap; -use std::net::UdpSocket as StdUdpSocket; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::tcp::OwnedReadHalf; -use tokio::net::{TcpStream, UdpSocket}; +use tokio::net::UdpSocket; use tokio::sync::watch::{channel, Receiver, Sender}; use crate::channel::punch::NatType; @@ -27,7 +27,7 @@ pub struct ContextInner { pub(crate) main_channel: Arc, pub(crate) main_channel_ipv6: Option>, //在udp的基础上，可以选择使用tcp和服务端通信 - pub(crate) main_tcp_channel: Option>>, + pub(crate) main_tcp_channel: Option>>, pub(crate) route_table: Atomic>>, pub(crate) route_table_time: DashMap<(RouteKey, Ipv4Addr), Instant>, pub(crate) status_receiver: Receiver, @@ -46,7 +46,7 @@ impl Context { pub fn new( main_channel: Arc, main_channel_ipv6: Option>, - main_tcp_channel: Option>>, + main_tcp_channel: Option>>, current_device: Arc>, _channel_num: usize, ) -> Self { @@ -90,6 +90,9 @@ impl Context { SocketAddr::V6(std::net::SocketAddrV6::new(Ipv6Addr::LOCALHOST, port, 0, 0)), ); } + if let Some(tcp) = &self.inner.main_tcp_channel { + let _ = tcp.send(vec![]); + } Ok(()) } pub fn is_main_tcp(&self) -> bool { @@ -248,7 +251,7 @@ impl Context { match route_key.index { TCP_ID => { if let Some(sender) = &self.inner.main_tcp_channel { - if sender.send(buf.to_vec()).await.is_ok() { + if sender.send(buf.to_vec()).is_ok() { Ok(buf.len()) } else { Err(io::Error::new(io::ErrorKind::Other, "send_by_key err")) @@ -569,8 +572,8 @@ fn buf_channel_group(size: usize) -> (BufSenderGroup, BufReceiverGroup) { } impl Channel { - async fn tcp_handle( - mut tcp_r: OwnedReadHalf, + fn tcp_handle( + tcp_r: &mut TcpStream, context: Context, handler: ChannelDataHandler, head_reserve: usize, @@ -580,7 +583,7 @@ impl Channel { let key = RouteKey::new(TCP_ID, addr); loop { let mut buf = [0; 4096]; - tcp_r.read_exact(&mut head).await?; + tcp_r.read_exact(&mut head)?; let len = (((head[2] as u16) << 8) | head[3] as u16) as usize; if len < 12 || len > buf.len() { return Err(io::Error::new( @@ -588,82 +591,98 @@ impl Channel { "length overflow", )); } - tcp_r - .read_exact(&mut buf[head_reserve..head_reserve + len]) - .await?; + tcp_r.read_exact(&mut buf[head_reserve..head_reserve + len])?; handler.handle(&mut buf, head_reserve, head_reserve + len, key, &context); } } - async fn start_tcp( - mut worker: VntWorker, - tcp_stream: TcpStream, - mut receiver: tokio::sync::mpsc::Receiver>, - current_device: Arc>, + fn start_tcp( + worker: VntWorker, + mut tcp_stream: TcpStream, + receiver: std::sync::mpsc::Receiver>, context: Context, handler: ChannelDataHandler, head_reserve: usize, ) { - let (tcp_r, mut tcp_w) = tcp_stream.into_split(); + let current_device = context.inner.current_device.clone(); { + let mut tcp_r = tcp_stream.try_clone().unwrap(); let context = context.clone(); let handler = handler.clone(); - tokio::spawn(async move { - if let Err(e) = Self::tcp_handle(tcp_r, context, handler, head_reserve).await { - log::info!("tcp链接断开:{:?}", e); - } - }); + thread::Builder::new() + .name("tcp_reader".into()) + .spawn(move || { + if let Err(e) = Self::tcp_handle(&mut tcp_r, context, handler, head_reserve) { + log::info!("tcp链接断开:{:?}", e); + } + if let Err(e) = tcp_r.shutdown(Shutdown::Both) { + log::info!("tcp链接关闭异常:{:?}", e); + } + }) + .unwrap(); } let mut head = [0; 4]; loop { - tokio::select! { - _=worker.stop_wait()=>{ + let data = match receiver.recv() { + Ok(data) => data, + Err(_) => { break; } - rs=receiver.recv()=>{ - if let Some(data) = rs{ - let len = data.len(); - head[2] = (len >> 8) as u8; - head[3] = (len & 0xFF) as u8; - let mut err = false; - if let Err(e) = tcp_w.write_all(&head).await{ - err = true; - log::info!("发送失败,需要重连:{:?}",e); - }else if let Err(e) = tcp_w.write_all(&data).await{ - err = true; - log::info!("发送失败,需要重连:{:?}",e); - } - if err { - let _ = tcp_w.shutdown().await; - match TcpStream::connect(current_device.load().connect_server).await { - Ok(tcp_stream) => { - let (r, w) = tcp_stream.into_split(); - tcp_w = w; - let context = context.clone(); - let handler = handler.clone(); - tokio::spawn(async move { - if let Err(e) = Self::tcp_handle(r, context,handler, head_reserve).await { - log::info!("tcp 链接断开:{:?}",e); - } - }); + }; + let len = data.len(); + if len == 0 { + break; + } + head[2] = (len >> 8) as u8; + head[3] = (len & 0xFF) as u8; + let mut err = false; + if let Err(e) = tcp_stream.write_all(&head) { + err = true; + log::info!("发送失败,需要重连:{:?}", e); + } else if let Err(e) = tcp_stream.write_all(&data) { + err = true; + log::info!("发送失败,需要重连:{:?}", e); + } + if err { + if let Err(e) = tcp_stream.shutdown(Shutdown::Both) { + log::info!("tcp链接关闭异常:{:?}", e); + } + match TcpStream::connect(current_device.load().connect_server) { + Ok(tcp) => { + tcp.set_read_timeout(Some(Duration::from_secs(10))).unwrap(); + tcp_stream = tcp; + let mut tcp_r = tcp_stream.try_clone().unwrap(); + let context = context.clone(); + let handler = handler.clone(); + thread::Builder::new() + .name("tcp_reader".into()) + .spawn(move || { + if let Err(e) = + Self::tcp_handle(&mut tcp_r, context, handler, head_reserve) + { + log::info!("重连后 tcp链接断开:{:?}", e); } - Err(e) => { - log::info!("重连失败:{:?}",e); + if let Err(e) = tcp_r.shutdown(Shutdown::Both) { + log::info!("重连后 tcp链接关闭异常:{:?}", e); } - }; - } - }else{ - break; + }) + .unwrap(); + } + Err(e) => { + log::info!("重连失败:{:?}", e); } } } } + if let Err(e) = tcp_stream.shutdown(Shutdown::Both) { + log::info!("tcp链接关闭异常:{:?}", e); + } worker.stop_all(); } pub async fn start( self, mut worker: VntWorker, - tcp: Option<(TcpStream, tokio::sync::mpsc::Receiver>)>, + tcp: Option<(TcpStream, std::sync::mpsc::Receiver>)>, head_reserve: usize, //头部预留字节 symmetric_channel_num: usize, //对称网络，则再加一组监听，提升打洞成功率 relay: bool, @@ -678,7 +697,7 @@ impl Channel { for buf_receiver in buf_receiver.0 { let context = context.clone(); let handler = handler.clone(); - std::thread::Builder::new() + thread::Builder::new() .name(format!("recv-handler-{}", num)) .spawn(move || { while let Ok((mut buf, start, end, route_key)) = buf_receiver.recv() { @@ -694,15 +713,22 @@ impl Channel { None }; if let Some((tcp_stream, receiver)) = tcp { - tokio::spawn(Self::start_tcp( - worker.worker("main_channel_tcp"), - tcp_stream, - receiver, - context.inner.current_device.clone(), - context.clone(), - handler.clone(), - head_reserve, - )); + let context = context.clone(); + let handler = handler.clone(); + let main_channel_tcp = worker.worker("main_channel_tcp"); + thread::Builder::new() + .name("main_channel_tcp".into()) + .spawn(move || { + Self::start_tcp( + main_channel_tcp, + tcp_stream, + receiver, + context, + handler, + head_reserve, + ) + }) + .unwrap(); } if let Some(main_channel_ipv6) = &context.inner.main_channel_ipv6 { let worker = worker.worker("main_channel_ipv6"); @@ -710,7 +736,7 @@ impl Channel { let main_channel_ipv6 = main_channel_ipv6.clone(); let handler = handler.clone(); let buf_sender = buf_sender.clone(); - std::thread::Builder::new() + thread::Builder::new() .name("ipv6-recv".into()) .spawn(move || { log::info!("启动udp v6"); @@ -732,7 +758,7 @@ impl Channel { let main_channel = main_channel.clone(); let handler = handler.clone(); let buf_sender = buf_sender.clone(); - std::thread::Builder::new() + thread::Builder::new() .name("ipv4-recv".into()) .spawn(move || { log::info!("启动udp v4"); diff --git a/vnt/src/core/mod.rs b/vnt/src/core/mod.rs index fb16c36..a140638 100644 --- a/vnt/src/core/mod.rs +++ b/vnt/src/core/mod.rs @@ -7,8 +7,8 @@ use crossbeam_utils::atomic::AtomicCell; use dashmap::DashMap; use parking_lot::Mutex; use rand::Rng; +use std::net::TcpStream; use std::net::UdpSocket; -use tokio::net::TcpStream; use tokio::sync::mpsc::channel; use crate::channel::channel::{Channel, Context}; @@ -67,7 +67,7 @@ pub struct VntUtil { } impl VntUtil { - pub async fn new(config: Config) -> io::Result { + pub fn new(config: Config) -> io::Result { //单个udp用同步的性能更好，但是代理和多端口监听用异步更方便，这里将两者结合起来 let main_channel = UdpSocket::bind(format!("0.0.0.0:{}", config.port))?; main_channel.set_write_timeout(Some(Duration::from_secs(5)))?; @@ -105,28 +105,28 @@ impl VntUtil { }) } ///链接 - pub async fn connect(&mut self) -> io::Result<()> { + pub fn connect(&mut self) -> io::Result<()> { if self.config.tcp { - let tcp = TcpStream::connect(self.config.server_address).await?; + let tcp = TcpStream::connect(self.config.server_address)?; + tcp.set_read_timeout(Some(Duration::from_secs(10)))?; let _ = self.main_tcp_channel.insert(tcp); } Ok(()) } ///握手 用于获取公钥 - pub async fn handshake(&mut self) -> Result, HandshakeEnum> { + pub fn handshake(&mut self) -> Result, HandshakeEnum> { let rsa_cipher = handshake_handler::handshake( &self.main_channel, self.main_tcp_channel.as_mut(), self.config.server_address, self.config.server_encrypt, - ) - .await?; + )?; self.rsa_cipher = rsa_cipher.clone(); Ok(rsa_cipher) } /// 加密握手 用于同步密钥 - pub async fn secret_handshake(&mut self) -> Result<(), HandshakeEnum> { + pub fn secret_handshake(&mut self) -> Result<(), HandshakeEnum> { handshake_handler::secret_handshake( &self.main_channel, self.main_tcp_channel.as_mut(), @@ -135,10 +135,9 @@ impl VntUtil { &self.server_cipher, self.config.token.clone(), ) - .await } /// 注册 - pub async fn register(&mut self) -> Result { + pub fn register(&mut self) -> Result { match registration_handler::registration( &self.main_channel, self.main_tcp_channel.as_mut(), @@ -149,9 +148,7 @@ impl VntUtil { 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) @@ -253,7 +250,7 @@ impl VntUtil { let (cone_sender, cone_receiver) = channel(3); let (symmetric_sender, symmetric_receiver) = channel(2); let (tcp_sender, tcp) = if let Some(main_tcp_channel) = self.main_tcp_channel { - let (tcp_sender, tcp_receiver) = channel::>(100); + let (tcp_sender, tcp_receiver) = std::sync::mpsc::sync_channel::>(100); (Some(tcp_sender), Some((main_tcp_channel, tcp_receiver))) } else { (None, None) diff --git a/vnt/src/core/sync.rs b/vnt/src/core/sync.rs index e7f1f22..386792d 100644 --- a/vnt/src/core/sync.rs +++ b/vnt/src/core/sync.rs @@ -22,20 +22,20 @@ impl VntUtilSync { let runtime = tokio::runtime::Builder::new_multi_thread() .enable_all() .build()?; - let vnt_util = runtime.block_on(VntUtil::new(config))?; + let vnt_util = VntUtil::new(config)?; Ok(VntUtilSync { vnt_util, runtime }) } pub fn connect(&mut self) -> io::Result<()> { - self.runtime.block_on(self.vnt_util.connect()) + self.vnt_util.connect() } pub fn handshake(&mut self) -> Result, HandshakeEnum> { - self.runtime.block_on(self.vnt_util.handshake()) + self.vnt_util.handshake() } pub fn secret_handshake(&mut self) -> Result<(), HandshakeEnum> { - self.runtime.block_on(self.vnt_util.secret_handshake()) + self.vnt_util.secret_handshake() } pub fn register(&mut self) -> Result { - self.runtime.block_on(self.vnt_util.register()) + self.vnt_util.register() } #[cfg(any(target_os = "android"))] pub fn create_iface(&mut self, vpn_fd: i32) { diff --git a/vnt/src/handle/handshake_handler.rs b/vnt/src/handle/handshake_handler.rs index 5c8100e..0a8e48e 100644 --- a/vnt/src/handle/handshake_handler.rs +++ b/vnt/src/handle/handshake_handler.rs @@ -1,16 +1,15 @@ use std::net::SocketAddr; -use protobuf::Message; -use std::net::UdpSocket; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::TcpStream; - use crate::channel::channel::Context; use crate::channel::RouteKey; use crate::cipher::{Cipher, RsaCipher}; use crate::proto::message::{HandshakeRequest, HandshakeResponse, SecretHandshakeRequest}; use crate::protocol::body::RSA_ENCRYPTION_RESERVED; use crate::protocol::{service_packet, NetPacket, Protocol, Version, MAX_TTL}; +use protobuf::Message; +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::net::UdpSocket; pub enum HandshakeEnum { NotSecret, @@ -59,7 +58,7 @@ fn secret_handshake_request_packet( } /// 第一次握手，拿到公钥 -pub async fn handshake( +pub fn handshake( main_channel: &UdpSocket, main_tcp_channel: Option<&mut TcpStream>, server_address: SocketAddr, @@ -74,8 +73,7 @@ pub async fn handshake( server_address, send_buf, &mut recv_buf, - ) - .await?; + )?; let net_packet = match NetPacket::new(&recv_buf[..len]) { Ok(net_packet) => net_packet, Err(e) => { @@ -140,7 +138,7 @@ pub async fn handshake( } } -async fn send_recv( +fn send_recv( main_channel: &UdpSocket, main_tcp_channel: Option<&mut TcpStream>, server_address: SocketAddr, @@ -152,20 +150,20 @@ async fn send_recv( let len = send_buf.len(); head[2] = (len >> 8) as u8; head[3] = (len & 0xFF) as u8; - if let Err(e) = main_tcp_channel.write_all(&head).await { + if let Err(e) = main_tcp_channel.write_all(&head) { return Err(HandshakeEnum::Other(format!("send error:{}", e))); } - if let Err(e) = main_tcp_channel.write_all(send_buf).await { + if let Err(e) = main_tcp_channel.write_all(send_buf) { return Err(HandshakeEnum::Other(format!("send error:{}", e))); } - if let Err(e) = main_tcp_channel.read_exact(&mut head).await { + if let Err(e) = main_tcp_channel.read_exact(&mut head) { return Err(HandshakeEnum::Other(format!("read error:{}", e))); } let len = (((head[2] as u16) << 8) | head[3] as u16) as usize; if len > recv_buf.len() { return Err(HandshakeEnum::Other("too long".to_string())); } - if let Err(e) = main_tcp_channel.read_exact(&mut recv_buf[..len]).await { + if let Err(e) = main_tcp_channel.read_exact(&mut recv_buf[..len]) { return Err(HandshakeEnum::Other(format!("read error:{}", e))); } Ok(len) @@ -187,7 +185,7 @@ async fn send_recv( } /// 第二次握手，同步对称密钥，后续将使用对称加密 -pub async fn secret_handshake( +pub fn secret_handshake( main_channel: &UdpSocket, main_tcp_channel: Option<&mut TcpStream>, server_address: SocketAddr, @@ -213,8 +211,7 @@ pub async fn secret_handshake( server_address, send_buf, &mut recv_buf, - ) - .await?; + )?; let mut net_packet = match NetPacket::new(&mut recv_buf[..len]) { Ok(net_packet) => net_packet, Err(e) => { diff --git a/vnt/src/handle/registration_handler.rs b/vnt/src/handle/registration_handler.rs index 5ebcf99..bcc455d 100644 --- a/vnt/src/handle/registration_handler.rs +++ b/vnt/src/handle/registration_handler.rs @@ -1,4 +1,5 @@ use crossbeam_utils::atomic::AtomicCell; +use std::io::{Read, Write}; use std::net::{Ipv4Addr, SocketAddr}; use std::time::{Duration, Instant}; @@ -6,9 +7,8 @@ use crate::channel::sender::ChannelSender; use crate::cipher::Cipher; use crate::handle::PeerDeviceInfo; use protobuf::Message; +use std::net::TcpStream; use std::net::UdpSocket; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::TcpStream; use crate::proto::message::{RegistrationRequest, RegistrationResponse}; use crate::protocol::body::ENCRYPTION_RESERVED; @@ -37,7 +37,7 @@ pub struct RegResponse { } ///向中继服务器注册，token标识一个虚拟网关，device_id防止多次注册时得到的ip不一致 -pub async fn registration( +pub fn registration( main_channel: &UdpSocket, main_tcp_channel: Option<&mut TcpStream>, server_cipher: &Cipher, @@ -67,17 +67,17 @@ pub async fn registration( vec[2] = (len >> 8) as u8; vec[3] = (len & 0xFF) as u8; vec[4..].copy_from_slice(buf); - if let Err(e) = main_tcp_channel.write_all(&vec).await { + if let Err(e) = main_tcp_channel.write_all(&vec) { return Err(ReqEnum::Other(format!("send error:{}", e))); } - if let Err(e) = main_tcp_channel.read_exact(&mut recv_buf[..4]).await { + if let Err(e) = main_tcp_channel.read_exact(&mut recv_buf[..4]) { return Err(ReqEnum::Other(format!("read error:{}", e))); } let len = 4 + (((recv_buf[2] as u16) << 8) | recv_buf[3] as u16) as usize; if len > recv_buf.len() { return Err(ReqEnum::Other("too long".to_string())); } - if let Err(e) = main_tcp_channel.read_exact(&mut recv_buf[4..len]).await { + if let Err(e) = main_tcp_channel.read_exact(&mut recv_buf[4..len]) { return Err(ReqEnum::Other(format!("read error:{}", e))); } &mut recv_buf[4..len]