diff --git a/vnt/src/channel/channel.rs b/vnt/src/channel/channel.rs index 7d241b2..3a6365b 100644 --- a/vnt/src/channel/channel.rs +++ b/vnt/src/channel/channel.rs @@ -8,7 +8,7 @@ use std::time::{Duration, Instant}; use std::{io, thread}; use crossbeam_utils::atomic::AtomicCell; -use parking_lot::RwLock; +use parking_lot::{Mutex, RwLock}; use tokio::net::UdpSocket; use tokio::sync::watch::{channel, Receiver, Sender}; @@ -22,7 +22,7 @@ pub struct ContextInner { //udp用于打洞、服务端通信(可选) pub(crate) main_channel: Arc, //在udp的基础上，可以选择使用tcp和服务端通信 - pub(crate) main_tcp_channel: Option>>, + pub(crate) main_tcp_channel: Option>, pub(crate) route_table: RwLock)>>>, pub(crate) status_receiver: Receiver, pub(crate) status_sender: Sender, @@ -40,7 +40,7 @@ pub struct Context { impl Context { pub fn new( main_channel: Arc, - main_tcp_channel: Option>>, + main_tcp_channel: Option, current_device: Arc>, _channel_num: usize, first_latency: bool, @@ -48,6 +48,7 @@ impl Context { //当前版本只支持一个通道 let channel_num = 1; let (status_sender, status_receiver) = channel(Status::Cone); + let main_tcp_channel = main_tcp_channel.map(|e| Mutex::new(e)); let inner = Arc::new(ContextInner { main_channel, main_tcp_channel, @@ -79,7 +80,7 @@ impl Context { ); } if let Some(tcp) = &self.inner.main_tcp_channel { - let _ = tcp.send(vec![]); + tcp.lock().shutdown(Shutdown::Both)?; } Ok(()) } @@ -126,14 +127,32 @@ impl Context { } self.inner.main_channel.send_to(buf, addr) } + #[inline] + pub fn send_main_tcp(&self, buf: &[u8]) -> io::Result { + if let Some(sender) = &self.inner.main_tcp_channel { + let mut stream = sender.lock(); + let mut head = [0; 4]; + let len = buf.len(); + head[2] = (len >> 8) as u8; + head[3] = (len & 0xFF) as u8; + stream.write_all(&head)?; + stream.write_all(buf)?; + Ok(len) + } else { + return Err(io::Error::new(io::ErrorKind::NotFound, "tcp not found")); + } + } pub fn send_main(&self, buf: &[u8], addr: SocketAddr) -> io::Result { if let Some(sender) = &self.inner.main_tcp_channel { - if sender.try_send(buf.to_vec()).is_ok() { - Ok(buf.len()) - } else { - Err(io::Error::new(io::ErrorKind::Other, "send_main err")) - } + let mut stream = sender.lock(); + let mut head = [0; 4]; + let len = buf.len(); + head[2] = (len >> 8) as u8; + head[3] = (len & 0xFF) as u8; + stream.write_all(&head)?; + stream.write_all(buf)?; + Ok(len) } else { self.send_main_udp(buf, addr) } @@ -185,17 +204,7 @@ impl Context { pub async fn send_by_key(&self, buf: &[u8], route_key: &RouteKey) -> io::Result { match route_key.index { - TCP_ID => { - if let Some(sender) = &self.inner.main_tcp_channel { - if sender.send(buf.to_vec()).is_ok() { - Ok(buf.len()) - } else { - Err(io::Error::new(io::ErrorKind::Other, "send_by_key err")) - } - } else { - Err(io::Error::new(io::ErrorKind::Other, "send_by_key err")) - } - } + TCP_ID => self.send_main_tcp(buf), UDP_ID => self.send_main_udp(buf, route_key.addr), _ => { if let Some(udp) = self.get_udp_by_route(route_key) { @@ -207,17 +216,7 @@ impl Context { } pub fn try_send_by_key(&self, buf: &[u8], route_key: &RouteKey) -> io::Result { match route_key.index { - TCP_ID => { - if let Some(sender) = &self.inner.main_tcp_channel { - if sender.try_send(buf.to_vec()).is_ok() { - Ok(buf.len()) - } else { - Err(io::Error::new(io::ErrorKind::Other, "send_by_key err")) - } - } else { - Err(io::Error::new(io::ErrorKind::Other, "send_by_key err")) - } - } + TCP_ID => self.send_main_tcp(buf), UDP_ID => self.send_main_udp(buf, route_key.addr), _ => { if let Some(udp) = self.get_udp_by_route(route_key) { @@ -406,14 +405,17 @@ fn buf_channel_group(size: usize) -> (BufSenderGroup, BufReceiverGroup) { impl Channel { fn tcp_handle( tcp_r: &mut TcpStream, - context: Context, - handler: ChannelDataHandler, + context: &Context, + handler: &ChannelDataHandler, head_reserve: usize, ) -> io::Result<()> { let mut head = [0; 4]; let addr = tcp_r.peer_addr()?; let key = RouteKey::new(TCP_ID, addr); loop { + if context.is_close() { + return Ok(()); + } let mut buf = [0; 4096]; tcp_r.read_exact(&mut head)?; let len = (((head[2] as u16) << 8) | head[3] as u16) as usize; @@ -424,97 +426,57 @@ impl Channel { )); } tcp_r.read_exact(&mut buf[head_reserve..head_reserve + len])?; - handler.handle(&mut buf, head_reserve, head_reserve + len, key, &context); + handler.handle(&mut buf, head_reserve, head_reserve + len, key, context); } } fn start_tcp( - worker: VntWorker, mut tcp_stream: TcpStream, - receiver: std::sync::mpsc::Receiver>, context: Context, handler: ChannelDataHandler, head_reserve: usize, ) { 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(); - 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 { - let data = match receiver.recv() { - Ok(data) => data, - Err(_) => { - break; - } - }; - let len = data.len(); - if len == 0 { - break; + if let Err(e) = tcp_stream.set_nodelay(true) { + log::info!("set_nodelay:{:?}", e); } - 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 let Err(e) = tcp_stream.set_write_timeout(Some(Duration::from_secs(3))) { + log::info!("set_write_timeout:{:?}", e); } - if err { - if let Err(e) = tcp_stream.shutdown(Shutdown::Both) { - log::info!("tcp链接关闭异常:{:?}", e); + if let Err(e) = tcp_stream.set_read_timeout(Some(Duration::from_secs(10))) { + log::info!("set_read_timeout:{:?}", e); + } + if let Err(e) = Self::tcp_handle(&mut tcp_stream, &context, &handler, head_reserve) { + log::info!("tcp链接断开:{:?}", e); + } + if let Err(e) = tcp_stream.shutdown(Shutdown::Both) { + log::info!("tcp链接关闭异常:{:?}", e); + } + loop { + if context.is_close() { + return; } - match TcpStream::connect(current_device.load().connect_server) { + let device_info = current_device.load(); + match TcpStream::connect(device_info.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); - } - if let Err(e) = tcp_r.shutdown(Shutdown::Both) { - log::info!("重连后 tcp链接关闭异常:{:?}", e); - } - }) - .unwrap(); + tcp_stream = tcp.try_clone().unwrap(); + let mut guard = context.inner.main_tcp_channel.as_ref().unwrap().lock(); + *guard = tcp; + break; } Err(e) => { - log::info!("重连失败:{:?}", e); + log::info!("重连失败,{},{:?}", device_info.connect_server, e); + thread::sleep(Duration::from_secs(3)); } } } } - 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, std::sync::mpsc::Receiver>)>, + tcp: Option, head_reserve: usize, //头部预留字节 symmetric_channel_num: usize, //对称网络，则再加一组监听，提升打洞成功率 relay: bool, @@ -544,32 +506,26 @@ impl Channel { } else { None }; - if let Some((tcp_stream, receiver)) = tcp { + if let Some(tcp_stream) = tcp { 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()) + .name("channel_tcp".into()) .spawn(move || { - Self::start_tcp( - main_channel_tcp, - tcp_stream, - receiver, - context, - handler, - head_reserve, - ) + Self::start_tcp(tcp_stream, context, handler, head_reserve); + drop(main_channel_tcp) }) .unwrap(); } { - let worker = worker.worker("main_channel_1"); + let worker = worker.worker("main_channel_udp"); let context = context.clone(); let main_channel = main_channel.clone(); let handler = handler.clone(); let buf_sender = buf_sender.clone(); thread::Builder::new() - .name("ipv4-recv".into()) + .name("channel_udp".into()) .spawn(move || { log::info!("启动udp v4"); Self::main_start_( diff --git a/vnt/src/core/mod.rs b/vnt/src/core/mod.rs index 0997f35..e87e71d 100644 --- a/vnt/src/core/mod.rs +++ b/vnt/src/core/mod.rs @@ -236,9 +236,8 @@ 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) = std::sync::mpsc::sync_channel::>(100); - (Some(tcp_sender), Some((main_tcp_channel, tcp_receiver))) + let (tcp_sender, tcp_receiver) = if let Some(main_tcp_channel) = self.main_tcp_channel { + (Some(main_tcp_channel.try_clone()?), Some(main_tcp_channel)) } else { (None, None) }; @@ -383,7 +382,7 @@ impl VntUtil { let relay = config.relay; tokio::spawn(async move { channel - .start(channel_worker, tcp, 14, 65, relay, config.parallel) + .start(channel_worker, tcp_receiver, 14, 65, relay, config.parallel) .await }); }