From af5d28447094cd2ee619b3c9eb3b1a69415d5a15 Mon Sep 17 00:00:00 2001 From: lubeilin <1791778603@qq.com> Date: Thu, 29 Feb 2024 22:06:20 +0800 Subject: [PATCH] =?UTF-8?q?[mio]=20=E6=94=AF=E6=8C=81=E5=A4=9A=E9=80=9A?= =?UTF-8?q?=E9=81=93=E4=BC=A0=E8=BE=93=EF=BC=8C=E4=BD=BF=E7=94=A8mio?= =?UTF-8?q?=E4=BB=A3=E6=9B=BFtokio?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- vnt/src/channel/channel.rs | 764 --------------------------------- vnt/src/channel/context.rs | 420 ++++++++++++++++++ vnt/src/channel/handler.rs | 6 + vnt/src/channel/idle.rs | 47 +- vnt/src/channel/mod.rs | 102 ++++- vnt/src/channel/notify.rs | 126 ++++++ vnt/src/channel/punch.rs | 156 ++++--- vnt/src/channel/sender.rs | 86 +++- vnt/src/channel/tcp_channel.rs | 456 ++++++++++++++++++++ vnt/src/channel/udp_channel.rs | 219 ++++++++++ 10 files changed, 1518 insertions(+), 864 deletions(-) delete mode 100644 vnt/src/channel/channel.rs create mode 100644 vnt/src/channel/context.rs create mode 100644 vnt/src/channel/handler.rs create mode 100644 vnt/src/channel/notify.rs create mode 100644 vnt/src/channel/tcp_channel.rs create mode 100644 vnt/src/channel/udp_channel.rs diff --git a/vnt/src/channel/channel.rs b/vnt/src/channel/channel.rs deleted file mode 100644 index fde3e89..0000000 --- a/vnt/src/channel/channel.rs +++ /dev/null @@ -1,764 +0,0 @@ -use std::collections::HashMap; -use std::io::{Read, Write}; -use std::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr}; -use std::net::{SocketAddrV6, TcpStream}; -use std::net::{TcpListener, UdpSocket as StdUdpSocket}; -#[cfg(any(unix))] -use std::os::fd::AsRawFd; -#[cfg(target_os = "windows")] -use std::os::windows::io::AsRawSocket; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::Arc; -use std::time::{Duration, Instant}; -use std::{io, thread}; - -use crossbeam_utils::atomic::AtomicCell; -use parking_lot::{Mutex, RwLock}; -use tokio::net::UdpSocket; -use tokio::sync::watch::{channel, Receiver, Sender}; - -use crate::channel::punch::NatType; -use crate::channel::{Route, RouteKey, Status, TCP_ID, UDP_ID}; -use crate::core::status::VntWorker; -use crate::handle::recv_handler::ChannelDataHandler; -use crate::handle::CurrentDeviceInfo; - -pub struct ContextInner { - //udp用于打洞、服务端通信(可选) - pub(crate) main_channel: StdUdpSocket, - //在udp的基础上，可以选择使用tcp和服务端通信 - pub(crate) main_tcp_channel: Option>, - pub(crate) route_table: RwLock)>>>, - 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, - tcp_port: u16, -} - -#[derive(Clone)] -pub struct Context { - pub(crate) inner: Arc, -} - -impl Context { - pub fn new( - main_channel: StdUdpSocket, - main_tcp_channel: Option, - current_device: Arc>, - _channel_num: usize, - first_latency: bool, - tcp_port: u16, - ) -> Self { - //当前版本只支持一个通道 - 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, - route_table: RwLock::new(HashMap::with_capacity(16)), - 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), - tcp_port, - }); - Self { inner } - } -} - -impl Context { - pub fn is_close(&self) -> bool { - 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<()> { - let last = self.is_close(); - 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() { - 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 { - if let Err(e) = tcp.lock().shutdown(Shutdown::Both) { - log::error!("发送停止消息到tcp失败:{:?}", e); - } - } - if !last { - for (_, tcp) in self.inner.tcp_map.read().clone() { - if let Err(e) = tcp.lock().shutdown(Shutdown::Both) { - log::error!("发送停止消息到tcp失败:{:?}", e); - } - } - if let Err(e) = TcpStream::connect_timeout( - &SocketAddr::V6(SocketAddrV6::new( - Ipv6Addr::LOCALHOST, - self.inner.tcp_port, - 0, - 0, - )), - Duration::from_secs(1), - ) { - log::error!("发送停止消息到tcp_listener失败:{:?}", e); - } - } - Ok(()) - } - pub fn is_main_tcp(&self) -> bool { - self.inner.main_tcp_channel.is_some() - } - pub fn is_first_latency(&self) -> bool { - self.inner.first_latency - } - pub fn switch(&self, nat_type: NatType) { - match nat_type { - NatType::Symmetric => { - self.switch_to_symmetric(); - } - NatType::Cone => { - self.switch_to_cone(); - } - } - } - pub fn switch_to_cone(&self) { - let _ = self.inner.status_sender.send(Status::Cone); - } - pub fn switch_to_symmetric(&self) { - let _ = self.inner.status_sender.send(Status::Symmetric); - } - pub fn main_local_udp_port(&self) -> io::Result { - self.inner.main_channel.local_addr().map(|k| k.port()) - } - fn insert_udp(&self, id: usize, udp: Arc) { - self.inner.udp_map.write().insert(id, udp); - } - fn remove_udp(&self, id: usize) { - self.inner.udp_map.write().remove(&id); - } - #[inline] - pub fn send_main_udp(&self, buf: &[u8], mut addr: SocketAddr) -> io::Result { - if let SocketAddr::V4(ipv4) = addr { - addr = SocketAddr::V6(SocketAddrV6::new( - ipv4.ip().to_ipv6_mapped(), - ipv4.port(), - 0, - 0, - )); - } - 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 { - Self::send_tcp(sender, buf) - } else { - return Err(io::Error::new(io::ErrorKind::NotFound, "tcp not found")); - } - } - pub fn send_tcp(sender: &Mutex, buf: &[u8]) -> io::Result { - let mut stream = sender.lock(); - send_tcp(&mut stream, buf) - } - - pub fn send_main(&self, buf: &[u8], addr: SocketAddr) -> 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 { - self.send_main_udp(buf, addr) - } - } - - pub(crate) fn try_send_all(&self, buf: &[u8], addr: SocketAddr) -> io::Result<()> { - let table = self.inner.udp_map.read(); - if table.is_empty() { - log::error!("udp列表为空,addr={}", addr); - return Ok(()); - } - for (_, udp) in table.iter() { - //使用ipv6的udp发送ipv4报文会出错 - if let Err(e) = udp.try_send_to(buf, addr) { - log::error!("{:?}", e); - } - } - Ok(()) - } - - pub async fn send_by_id(&self, buf: &[u8], id: &Ipv4Addr) -> io::Result { - let route = self.get_route_by_id(id)?; - self.send_by_key(buf, &route.route_key()).await - } - pub fn try_send_by_id(&self, buf: &[u8], id: &Ipv4Addr) -> io::Result { - let route = self.get_route_by_id(id)?; - self.try_send_by_key(buf, &route.route_key()) - } - fn get_route_by_id(&self, id: &Ipv4Addr) -> io::Result { - if let Some(v) = self.inner.route_table.read().get(id) { - if v.is_empty() { - return Err(io::Error::new(io::ErrorKind::NotFound, "route not found")); - } - let (route, time) = &v[0]; - if route.rt == 199 { - //这通常是刚加入路由，直接放弃使用,避免抖动 - return Err(io::Error::new(io::ErrorKind::NotFound, "route not found")); - } - if !route.is_p2p() { - //借道传输时，长时间不通信的通道不使用 - if time.load().elapsed() > Duration::from_secs(6) { - return Err(io::Error::new(io::ErrorKind::NotFound, "route time out")); - } - } - return Ok(*route); - } - Err(io::Error::new(io::ErrorKind::NotFound, "route not found")) - } - - pub async fn send_by_key(&self, buf: &[u8], route_key: &RouteKey) -> io::Result { - match route_key.index { - TCP_ID => self.send_main_tcp(buf), - UDP_ID => self.send_main_udp(buf, route_key.addr), - _ => { - if route_key.is_tcp { - if let Some(tcp) = self.get_tcp_by_route(route_key) { - return Self::send_tcp(&tcp, buf); - } - } else { - if let Some(udp) = self.get_udp_by_route(route_key) { - return udp.send_to(buf, route_key.addr).await; - } - } - Err(io::Error::new(io::ErrorKind::NotFound, "route not found")) - } - } - } - pub fn try_send_by_key(&self, buf: &[u8], route_key: &RouteKey) -> io::Result { - match route_key.index { - TCP_ID => self.send_main_tcp(buf), - UDP_ID => self.send_main_udp(buf, route_key.addr), - _ => { - if route_key.is_tcp { - if let Some(tcp) = self.get_tcp_by_route(route_key) { - return Self::send_tcp(&tcp, buf); - } - } else { - if let Some(udp) = self.get_udp_by_route(route_key) { - return udp.try_send_to(buf, route_key.addr); - } - } - Err(io::Error::new(io::ErrorKind::NotFound, "route not found")) - } - } - } - #[inline] - fn get_udp_by_route(&self, route_key: &RouteKey) -> Option> { - self.inner.udp_map.read().get(&route_key.index).cloned() - } - #[inline] - fn get_tcp_by_route(&self, route_key: &RouteKey) -> Option>> { - self.inner.tcp_map.read().get(&route_key.index).cloned() - } - - pub fn add_route_if_absent(&self, id: Ipv4Addr, route: Route) { - self.add_route_(id, route, true) - } - pub fn add_route(&self, id: Ipv4Addr, route: Route) { - self.add_route_(id, route, false) - } - fn add_route_(&self, id: Ipv4Addr, route: Route, only_if_absent: bool) { - let key = route.route_key(); - let mut route_table = self.inner.route_table.write(); - let list = route_table - .entry(id) - .or_insert_with(|| Vec::with_capacity(4)); - let mut exist = false; - for (x, time) in list.iter_mut() { - if x.metric < route.metric && !self.inner.first_latency { - //非优先延迟的情况下 不能比当前的路径更长 - return; - } - if x.route_key() == key { - if only_if_absent { - return; - } - x.metric = route.metric; - x.rt = route.rt; - exist = true; - time.store(Instant::now()); - break; - } - } - if exist { - list.sort_by_key(|(k, _)| k.rt); - } else { - let max_len = if self.inner.first_latency { - self.inner.channel_num + 1 - } else { - if route.metric == 1 { - //非优先延迟的情况下 添加了直连的则排除非直连的 - list.retain(|(k, _)| k.metric == 1); - } - self.inner.channel_num - }; - list.sort_by_key(|(k, _)| k.rt); - if list.len() > max_len { - list.truncate(max_len); - } - list.push((route, AtomicCell::new(Instant::now()))); - } - } - pub fn route(&self, id: &Ipv4Addr) -> Option> { - if let Some(v) = self.inner.route_table.read().get(id) { - Some(v.iter().map(|(i, _)| *i).collect()) - } else { - None - } - } - pub fn route_one(&self, id: &Ipv4Addr) -> Option { - if let Some(v) = self.inner.route_table.read().get(id) { - v.first().map(|(i, _)| *i) - } else { - None - } - } - pub fn route_to_id(&self, route_key: &RouteKey) -> Option { - let table = self.inner.route_table.read(); - for (k, v) in table.iter() { - for (route, _) in v { - if &route.route_key() == route_key && route.is_p2p() { - return Some(*k); - } - } - } - None - } - pub fn need_punch(&self, id: &Ipv4Addr) -> bool { - if let Some(v) = self.inner.route_table.read().get(id) { - if v.iter().filter(|(k, _)| k.is_p2p()).count() >= self.inner.channel_num { - return false; - } - } - true - } - pub fn route_table(&self) -> Vec<(Ipv4Addr, Vec)> { - let table = self.inner.route_table.read(); - table - .iter() - .map(|(k, v)| (k.clone(), v.iter().map(|(i, _)| *i).collect())) - .collect() - } - pub fn route_table_one(&self) -> Vec<(Ipv4Addr, Route)> { - let mut list = Vec::with_capacity(8); - let table = self.inner.route_table.read(); - for (k, v) in table.iter() { - if let Some((route, _)) = v.first() { - list.push((*k, *route)); - } - } - list - } - pub fn direct_route_table_one(&self) -> Vec<(Ipv4Addr, Route)> { - let mut list = Vec::with_capacity(8); - let table = self.inner.route_table.read(); - for (k, v) in table.iter() { - if let Some((route, _)) = v.first() { - if route.metric == 1 { - list.push((*k, *route)); - } - } - } - list - } - - pub fn remove_route(&self, id: &Ipv4Addr, route_key: RouteKey) { - if let Some(routes) = self.inner.route_table.write().get_mut(id) { - routes.retain(|(x, _)| x.route_key() != route_key); - } else { - return; - } - } - pub fn update_read_time(&self, id: &Ipv4Addr, route_key: &RouteKey) { - if let Some(routes) = self.inner.route_table.read().get(id) { - for (route, time) in routes { - if &route.route_key() == route_key { - time.store(Instant::now()); - break; - } - } - } - } -} - -pub struct Channel { - context: Context, - handler: ChannelDataHandler, - tcp_listener: TcpListener, -} - -impl Channel { - pub fn new(context: Context, handler: ChannelDataHandler, tcp_listener: TcpListener) -> Self { - Self { - context, - handler, - tcp_listener, - } - } -} - -impl Channel { - fn start_tcp(mut tcp_stream: TcpStream, context: Context, handler: ChannelDataHandler) { - let current_device = context.inner.current_device.clone(); - loop { - if let Err(e) = tcp_stream.set_nodelay(true) { - log::info!("set_nodelay:{:?}", e); - } - if let Err(e) = tcp_stream.set_write_timeout(Some(Duration::from_secs(5))) { - log::info!("set_write_timeout:{:?}", 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) = tcp_handle(TCP_ID, &mut tcp_stream, &context, &handler) { - log::info!("tcp链接断开:{:?}", e); - } - if let Err(e) = tcp_stream.shutdown(Shutdown::Both) { - log::info!("tcp链接关闭异常:{:?}", e); - } - loop { - if context.is_close() { - return; - } - let device_info = current_device.load(); - match TcpStream::connect(device_info.connect_server) { - Ok(tcp) => { - 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!("重连失败,{},{:?}", device_info.connect_server, e); - thread::sleep(Duration::from_secs(3)); - } - } - } - } - } - fn start_tcp_listen( - worker: VntWorker, - context: Context, - handler: ChannelDataHandler, - tcp_listener: TcpListener, - ) { - let counter = Arc::new(AtomicUsize::new(0)); - for stream in tcp_listener.incoming() { - if context.is_close() { - break; - } - if counter.load(Ordering::Relaxed) > 20 { - continue; - } - match stream { - Ok(stream) => { - let context = context.clone(); - let handler = handler.clone(); - let counter = counter.clone(); - counter.fetch_add(1, Ordering::Relaxed); - thread::spawn(move || { - if let Err(e) = start_tcp_handle(stream, context, handler) { - log::error!("{:?}", e); - } - counter.fetch_sub(1, Ordering::Relaxed); - }); - } - Err(e) => { - log::error!("connection failed {:?}", e); - } - } - } - for (_, tcp) in context.inner.tcp_map.read().clone() { - if let Err(e) = tcp.lock().shutdown(Shutdown::Both) { - log::error!("发送停止消息到tcp失败:{:?}", e); - } - } - worker.stop_all(); - } - - pub async fn start( - self, - mut worker: VntWorker, - tcp: Option, - symmetric_channel_num: usize, //对称网络，则再加一组监听，提升打洞成功率 - relay: bool, - ) { - let handler = self.handler.clone(); - let context = self.context; - let main_channel = context.inner.main_channel.try_clone().unwrap(); - 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("channel_tcp".into()) - .spawn(move || { - Self::start_tcp(tcp_stream, context, handler); - drop(main_channel_tcp) - }) - .unwrap(); - } - { - let worker = worker.worker("main_channel_udp"); - let context = context.clone(); - let main_channel = main_channel.try_clone().unwrap(); - let handler = handler.clone(); - thread::Builder::new() - .name("channel_udp".into()) - .spawn(move || { - log::info!("启动udp v4"); - Self::main_start_(worker, context, UDP_ID, main_channel, handler) - }) - .unwrap(); - } - if relay { - worker.stop_wait().await; - return; - } - { - let context = context.clone(); - let handler = handler.clone(); - let tcp_listener = self.tcp_listener; - let worker = worker.worker("tcp_listener"); - thread::Builder::new() - .name("tcp_listener".into()) - .spawn(move || { - log::info!("启动tcp"); - Self::start_tcp_listen(worker, context, handler, tcp_listener) - }) - .unwrap(); - } - let mut cur_status = Status::Cone; - let mut status_receiver = context.inner.status_receiver.clone(); - let channel_num = context.inner.channel_num; - loop { - tokio::select! { - _=worker.stop_wait()=>{ - break; - } - rs=status_receiver.changed()=>{ - match rs { - Ok(_) => { - let s = status_receiver.borrow().clone(); - match s { - Status::Cone => { - cur_status = Status::Cone; - } - Status::Symmetric => { - if cur_status == Status::Symmetric { - continue; - } - cur_status = Status::Symmetric; - for _ in 0..symmetric_channel_num - channel_num { - match UdpSocket::bind("0.0.0.0:0").await { - Ok(udp) => { - let udp = Arc::new(udp); - let context = context.clone(); - tokio::spawn(Self::start_(worker.worker("symmetric_channel"),context, udp,handler.clone())); - } - Err(e) => { - log::error!("{}",e); - } - } - } - } - Status::Close => { - break; - } - } - } - Err(_) => { - break; - } - } - } - } - } - worker.stop_all(); - } - fn main_start_( - worker: VntWorker, - context: Context, - id: usize, - udp: StdUdpSocket, - handler: ChannelDataHandler, - ) { - let mut buf = [0; 4096]; - let head_reserve = handler.head_reserve; - 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); - } - } - } - - worker.stop_all(); - } - async fn start_( - mut worker: VntWorker, - context: Context, - udp: Arc, - handler: ChannelDataHandler, - ) { - let mut status_receiver = context.inner.status_receiver.clone(); - - #[cfg(target_os = "windows")] - let id = 3 + udp.as_raw_socket() as usize; - #[cfg(any(unix))] - let id = 3 + udp.as_raw_fd() as usize; - - context.insert_udp(id, udp.clone()); - let mut buf = [0; 4096]; - let head_reserve = handler.head_reserve; - loop { - tokio::select! { - rs=udp.recv_from(&mut buf[head_reserve..])=>{ - match rs { - Ok((len, addr)) => { - handler.handle(&mut buf, head_reserve, head_reserve + len, RouteKey::new(false,id, addr), &context); - } - Err(e) => { - log::error!("{:?}",e) - } - } - } - changed=status_receiver.changed()=>{ - match changed { - Ok(_) => { - match *status_receiver.borrow() { - Status::Cone => { - break; - } - Status::Close=>{ - break; - } - Status::Symmetric => {} - } - } - Err(_) => { - break; - } - } - } - _=worker.stop_wait()=>{ - break; - } - } - } - context.remove_udp(id); - } -} -pub fn start_tcp_handle( - mut stream: TcpStream, - context: Context, - handler: ChannelDataHandler, -) -> io::Result<()> { - stream.set_write_timeout(Some(Duration::from_secs(5)))?; - stream.set_read_timeout(Some(Duration::from_secs(10)))?; - if let Err(e) = stream.set_nodelay(true) { - log::error!("设置nodelay失败 {:?}", e); - } - let writer = stream.try_clone()?; - #[cfg(target_os = "windows")] - let id = 3 + stream.as_raw_socket() as usize; - #[cfg(any(unix))] - let id = 3 + stream.as_raw_fd() as usize; - context - .inner - .tcp_map - .write() - .insert(id, Arc::new(Mutex::new(writer))); - if let Err(e) = tcp_handle(id, &mut stream, &context, &handler) { - log::error!("tcp_handle {:?}", e); - } - context.inner.tcp_map.write().remove(&id); - Ok(()) -} -pub fn tcp_handle( - id: usize, - tcp_r: &mut TcpStream, - context: &Context, - handler: &ChannelDataHandler, -) -> io::Result<()> { - let mut head = [0; 4]; - let addr = tcp_r.peer_addr()?; - let key = RouteKey::new(true, id, addr); - let head_reserve = handler.head_reserve; - 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; - if len < 12 || len > buf.len() { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "length overflow", - )); - } - tcp_r.read_exact(&mut buf[head_reserve..head_reserve + len])?; - handler.handle(&mut buf, head_reserve, head_reserve + len, key, context); - } -} -pub fn send_tcp(stream: &mut TcpStream, buf: &[u8]) -> io::Result { - 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) -} diff --git a/vnt/src/channel/context.rs b/vnt/src/channel/context.rs new file mode 100644 index 0000000..29ffcc9 --- /dev/null +++ b/vnt/src/channel/context.rs @@ -0,0 +1,420 @@ +use std::collections::HashMap; +use std::io; +use std::net::{Ipv4Addr, SocketAddr, SocketAddrV6, UdpSocket}; +use std::ops::Deref; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use crossbeam_utils::atomic::AtomicCell; +use parking_lot::RwLock; + +use crate::channel::punch::NatType; +use crate::channel::sender::{AcceptSocketSender, ChannelSender, PacketSender}; +use crate::channel::{Route, RouteKey}; +use crate::handle::{ConnectStatus, CurrentDeviceInfo}; + +/// 传输通道上下文，持有udp socket、tcp socket和路由信息 +#[derive(Clone)] +pub struct Context { + inner: Arc, +} + +impl Context { + pub fn new(main_udp_socket: Vec, first_latency: bool, is_tcp: bool) -> Self { + let channel_num = main_udp_socket.len(); + assert_ne!(channel_num, 0, "not channel"); + let inner = ContextInner { + main_udp_socket, + sub_udp_socket: RwLock::new(Vec::with_capacity(64)), + tcp_map: RwLock::new(HashMap::with_capacity(64)), + route_table: RouteTable::new(first_latency, channel_num), + is_tcp, + }; + Self { + inner: Arc::new(inner), + } + } + pub fn sender(&self) -> ChannelSender { + ChannelSender::new(self.clone()) + } +} + +impl Deref for Context { + type Target = ContextInner; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +/// 对称网络增加的udp socket数目，有助于增加打洞成功率 +pub const SYMMETRIC_CHANNEL_NUM: usize = 64; + +pub struct ContextInner { + // 核心udp socket + pub(crate) main_udp_socket: Vec, + // 对称网络增加的udp socket + sub_udp_socket: RwLock>, + // tcp数据发送器 + pub(crate) tcp_map: RwLock>, + // 路由信息 + pub route_table: RouteTable, + // 是否使用tcp连接服务器 + is_tcp: bool, +} + +impl ContextInner { + /// 通过sub_udp_socket是否为空来判断是否为锥形网络 + pub fn is_cone(&self) -> bool { + self.sub_udp_socket.read().is_empty() + } + pub fn is_main_tcp(&self) -> bool { + self.is_tcp + } + pub fn is_udp_main(&self, route_key: &RouteKey) -> bool { + !route_key.is_tcp() && route_key.index < self.main_udp_socket.len() + } + pub fn first_latency(&self) -> bool { + self.route_table.first_latency + } + /// 切换NAT类型，不同的nat打洞模式会有不同 + pub fn switch( + &self, + nat_type: NatType, + udp_socket_sender: &AcceptSocketSender>>, + ) -> io::Result<()> { + let mut write_guard = self.sub_udp_socket.write(); + match nat_type { + NatType::Symmetric => { + if !write_guard.is_empty() { + return Ok(()); + } + let mut vec = Vec::with_capacity(SYMMETRIC_CHANNEL_NUM); + for _ in 0..SYMMETRIC_CHANNEL_NUM { + let udp = UdpSocket::bind("0.0.0.0:0")?; + //副通道使用异步io + udp.set_nonblocking(true)?; + vec.push(udp); + } + let mut mio_vec = Vec::with_capacity(SYMMETRIC_CHANNEL_NUM); + for udp in vec.iter() { + let udp_socket = mio::net::UdpSocket::from_std(udp.try_clone()?); + mio_vec.push(udp_socket); + } + udp_socket_sender.try_add_socket(Some(mio_vec))?; + *write_guard = vec; + } + NatType::Cone => { + if write_guard.is_empty() { + return Ok(()); + } + udp_socket_sender.try_add_socket(None)?; + *write_guard = Vec::new(); + } + } + Ok(()) + } + pub fn change_status( + &self, + current_device: &AtomicCell, + ) -> CurrentDeviceInfo { + let mut cur = current_device.load(); + loop { + let status = if self.route_table.route_one(&cur.virtual_gateway).is_some() { + //已连接 + if cur.status == ConnectStatus::Connected { + return cur; + } + //状态变为已连接 + ConnectStatus::Connected + } else { + //未连接 + if cur.status == ConnectStatus::Connecting { + return cur; + } + //状态变为未连接 + ConnectStatus::Connecting + }; + let mut new_info = cur; + new_info.status = status; + match current_device.compare_exchange(cur, new_info) { + Ok(_) => { + return new_info; + } + Err(c) => { + cur = c; + } + } + } + } + pub fn channel_num(&self) -> usize { + self.main_udp_socket.len() + } + /// 获取核心udp监听的端口，用于其他客户端连接 + pub fn main_local_udp_port(&self) -> io::Result> { + let mut ports = Vec::new(); + for udp in self.main_udp_socket.iter() { + ports.push(udp.local_addr()?.port()) + } + Ok(ports) + } + pub fn send_tcp(&self, buf: &[u8], addr: SocketAddr) -> io::Result<()> { + if let Some(tcp) = self.tcp_map.read().get(&addr) { + tcp.try_send(buf) + } else { + Err(io::Error::from(io::ErrorKind::NotFound)) + } + } + pub fn send_main_udp(&self, index: usize, buf: &[u8], mut addr: SocketAddr) -> io::Result<()> { + //核心udp socket都是ipv6模式,如果是v4地址则需要转换成v6 + //只有服务器地址可能需要这样转换 + if let SocketAddr::V4(ipv4) = addr { + addr = SocketAddr::V6(SocketAddrV6::new( + ipv4.ip().to_ipv6_mapped(), + ipv4.port(), + 0, + 0, + )); + } + self.main_udp_socket[index].send_to(buf, addr)?; + Ok(()) + } + /// 将数据发送到默认通道，一般发往服务器才用此方法 + pub fn send_default(&self, buf: &[u8], addr: SocketAddr) -> io::Result<()> { + if self.is_tcp { + //服务端地址只在重连时检测变化 + self.send_tcp(buf, addr) + } else { + self.send_main_udp(0, buf, addr) + } + } + /// 此方法仅用于对称网络打洞 + pub fn try_send_all(&self, buf: &[u8], addr: SocketAddr) { + self.try_send_all_main(buf, addr); + for udp in self.sub_udp_socket.read().iter() { + if let Err(e) = udp.send_to(buf, addr) { + log::warn!("{:?},add={:?}", e, addr); + } + } + } + pub fn try_send_all_main(&self, buf: &[u8], mut addr: SocketAddr) { + if let SocketAddr::V4(ipv4) = addr { + addr = SocketAddr::V6(SocketAddrV6::new( + ipv4.ip().to_ipv6_mapped(), + ipv4.port(), + 0, + 0, + )); + } + for udp in &self.main_udp_socket { + if let Err(e) = udp.send_to(buf, addr) { + log::warn!("{:?},add={:?}", e, addr); + } + } + } + /// 将数据发到指定id + pub fn send_by_id(&self, buf: &[u8], id: &Ipv4Addr) -> io::Result<()> { + let route = self.route_table.get_route_by_id(id)?; + self.send_by_key(buf, route.route_key()) + } + /// 将数据发到指定路由 + pub fn send_by_key(&self, buf: &[u8], route_key: RouteKey) -> io::Result<()> { + if route_key.is_tcp { + self.send_tcp(buf, route_key.addr) + } else { + if let Some(main_udp) = self.main_udp_socket.get(route_key.index) { + main_udp.send_to(buf, 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)?; + } else { + Err(io::Error::from(io::ErrorKind::NotFound))? + } + } + Ok(()) + } + } +} + +pub struct RouteTable { + pub(crate) route_table: + RwLock)>)>>, + first_latency: bool, + channel_num: usize, +} + +impl RouteTable { + fn new(first_latency: bool, channel_num: usize) -> Self { + Self { + route_table: RwLock::new(HashMap::with_capacity(64)), + first_latency, + channel_num, + } + } +} + +impl RouteTable { + fn get_route_by_id(&self, id: &Ipv4Addr) -> io::Result { + if let Some((count, v)) = self.route_table.read().get(id) { + if v.is_empty() { + return Err(io::Error::new(io::ErrorKind::NotFound, "route not found")); + } + if self.channel_num > 1 { + //多通道的，则轮流使用 + let index = count.fetch_add(1, Ordering::Relaxed); + if let Some((route, _time)) = v.get(index) { + if route.is_p2p() && route.rt != 199 { + return Ok(*route); + } + } + } + let (route, time) = &v[0]; + if route.rt == 199 { + //这通常是刚加入路由，直接放弃使用,避免抖动 + return Err(io::Error::new(io::ErrorKind::NotFound, "route not found")); + } + if !route.is_p2p() { + //借道传输时，长时间不通信的通道不使用 + if time.load().elapsed() > Duration::from_secs(5) { + return Err(io::Error::new(io::ErrorKind::NotFound, "route time out")); + } + } + return Ok(*route); + } + Err(io::Error::new(io::ErrorKind::NotFound, "route not found")) + } + pub fn add_route_if_absent(&self, id: Ipv4Addr, route: Route) { + self.add_route_(id, route, true) + } + pub fn add_route(&self, id: Ipv4Addr, route: Route) { + self.add_route_(id, route, false) + } + fn add_route_(&self, id: Ipv4Addr, route: Route, only_if_absent: bool) { + let key = route.route_key(); + let mut route_table = self.route_table.write(); + let (_, list) = route_table + .entry(id) + .or_insert_with(|| (AtomicUsize::new(0), Vec::with_capacity(4))); + let mut exist = false; + for (x, time) in list.iter_mut() { + if x.metric < route.metric && !self.first_latency { + //非优先延迟的情况下 不能比当前的路径更长 + return; + } + if x.route_key() == key { + if only_if_absent { + return; + } + x.metric = route.metric; + x.rt = route.rt; + exist = true; + time.store(Instant::now()); + break; + } + } + if exist { + list.sort_by_key(|(k, _)| k.rt); + } else { + let max_len = if self.first_latency { + self.channel_num + 1 + } else { + if route.metric == 1 { + //非优先延迟的情况下 添加了直连的则排除非直连的 + list.retain(|(k, _)| k.metric == 1); + } + self.channel_num + }; + list.sort_by_key(|(k, _)| k.rt); + if list.len() > max_len { + list.truncate(max_len); + } + list.push((route, AtomicCell::new(Instant::now()))); + } + } + pub fn route(&self, id: &Ipv4Addr) -> Option> { + if let Some((_, v)) = self.route_table.read().get(id) { + Some(v.iter().map(|(i, _)| *i).collect()) + } else { + None + } + } + pub fn route_one(&self, id: &Ipv4Addr) -> Option { + if let Some((_, v)) = self.route_table.read().get(id) { + v.first().map(|(i, _)| *i) + } else { + None + } + } + pub fn route_to_id(&self, route_key: &RouteKey) -> Option { + let table = self.route_table.read(); + for (k, (_, v)) in table.iter() { + for (route, _) in v { + if &route.route_key() == route_key && route.is_p2p() { + return Some(*k); + } + } + } + None + } + pub fn need_punch(&self, id: &Ipv4Addr) -> bool { + if let Some((_, v)) = self.route_table.read().get(id) { + if v.iter().filter(|(k, _)| k.is_p2p()).count() >= self.channel_num { + return false; + } + } + true + } + /// 返回所有路由 + pub fn route_table(&self) -> Vec<(Ipv4Addr, Vec)> { + let table = self.route_table.read(); + table + .iter() + .map(|(k, (_, v))| (k.clone(), v.iter().map(|(i, _)| *i).collect())) + .collect() + } + pub fn route_table_p2p(&self) -> Vec<(Ipv4Addr, Route)> { + let table = self.route_table.read(); + let mut list = Vec::with_capacity(8); + for (ip, (_, routes)) in table.iter() { + if let Some((route, _)) = routes.first() { + if route.is_p2p() { + list.push((*ip, *route)); + } + } + } + list + } + pub fn route_table_one(&self) -> Vec<(Ipv4Addr, Route)> { + let mut list = Vec::with_capacity(8); + let table = self.route_table.read(); + for (k, (_, v)) in table.iter() { + if let Some((route, _)) = v.first() { + list.push((*k, *route)); + } + } + list + } + pub fn remove_route(&self, id: &Ipv4Addr, route_key: RouteKey) { + if let Some((_, routes)) = self.route_table.write().get_mut(id) { + routes.retain(|(x, _)| x.route_key() != route_key); + } else { + return; + } + } + /// 更新路由入栈包的时刻，长时间没有收到数据的路由将会被剔除 + pub fn update_read_time(&self, id: &Ipv4Addr, route_key: &RouteKey) { + if let Some((_, routes)) = self.route_table.read().get(id) { + for (route, time) in routes { + if &route.route_key() == route_key { + time.store(Instant::now()); + break; + } + } + } + } +} diff --git a/vnt/src/channel/handler.rs b/vnt/src/channel/handler.rs new file mode 100644 index 0000000..16c1017 --- /dev/null +++ b/vnt/src/channel/handler.rs @@ -0,0 +1,6 @@ +use crate::channel::context::Context; +use crate::channel::RouteKey; + +pub trait RecvChannelHandler: Clone + Send + 'static { + fn handle(&mut self, buf: &mut [u8], route_key: RouteKey, context: &Context); +} diff --git a/vnt/src/channel/idle.rs b/vnt/src/channel/idle.rs index 45501bf..8e400c3 100644 --- a/vnt/src/channel/idle.rs +++ b/vnt/src/channel/idle.rs @@ -1,9 +1,7 @@ -use std::io; -use std::io::{Error, ErrorKind}; use std::net::Ipv4Addr; use std::time::Duration; -use crate::channel::channel::Context; +use crate::channel::context::Context; use crate::channel::RouteKey; pub struct Idle { @@ -17,32 +15,31 @@ impl Idle { } } +pub enum IdleType { + Timeout(Ipv4Addr, RouteKey), + Sleep(Duration), + None, +} + impl Idle { /// 获取空闲路由 - pub async fn next_idle(&self) -> io::Result<(Ipv4Addr, RouteKey)> { - loop { - let mut max = Duration::from_secs(0); - { - for (ip, routes) in self.context.inner.route_table.read().iter() { - for (route, time) in routes { - let last_read = time.load().elapsed(); - if last_read >= self.read_idle { - return Ok((*ip, route.route_key())); - } else { - if max < last_read { - max = last_read; - } - } - } + pub fn next_idle(&self) -> IdleType { + let mut max = Duration::from_secs(0); + let read_guard = self.context.route_table.route_table.read(); + if read_guard.is_empty() { + return IdleType::None; + } + for (ip, (_, routes)) in read_guard.iter() { + for (route, time) in routes { + let last_read = time.load().elapsed(); + if last_read >= self.read_idle { + return IdleType::Timeout(*ip, route.route_key()); + } else if max < last_read { + max = last_read; } } - if self.read_idle > max { - let sleep_time = self.read_idle - max; - tokio::time::sleep(sleep_time).await; - } - if self.context.is_close() { - return Err(Error::new(ErrorKind::Other, "closed")); - } } + let sleep_time = self.read_idle - max; + return IdleType::Sleep(sleep_time); } } diff --git a/vnt/src/channel/mod.rs b/vnt/src/channel/mod.rs index 02d0267..a263ae4 100644 --- a/vnt/src/channel/mod.rs +++ b/vnt/src/channel/mod.rs @@ -1,12 +1,24 @@ -use std::net::SocketAddr; +use std::io; +use std::net::{SocketAddr, UdpSocket}; +use std::time::Duration; -pub mod channel; +use crate::channel::context::Context; +use crate::channel::handler::RecvChannelHandler; +use crate::channel::sender::AcceptSocketSender; +use crate::channel::tcp_channel::tcp_listen; +use crate::channel::udp_channel::udp_listen; +use crate::util::{io_convert, StopManager}; + +pub mod context; +pub mod handler; pub mod idle; +pub mod notify; pub mod punch; pub mod sender; +pub mod tcp_channel; +pub mod udp_channel; -const TCP_ID: usize = 0; -const UDP_ID: usize = 1; +const BUFFER_SIZE: usize = 1024 * 16; #[derive(Copy, Clone, Eq, PartialEq)] pub enum Status { @@ -83,6 +95,86 @@ impl RouteKey { } } pub fn is_tcp(&self) -> bool { - self.index == TCP_ID + self.is_tcp + } + pub fn index(&self) -> usize { + self.index } } + +pub fn init_context( + ports: Vec, + first_latency: bool, + is_tcp: bool, +) -> io::Result<(Context, mio::net::TcpListener)> { + assert!(!ports.is_empty(), "not channel"); + let mut udps = Vec::with_capacity(ports.len()); + for port in &ports { + //监听v6+v4双栈，主通道使用同步io + let address: SocketAddr = format!("[::]:{}", port).parse().unwrap(); + let socket = socket2::Socket::new(socket2::Domain::IPV6, socket2::Type::DGRAM, None)?; + io_convert(socket.set_only_v6(false), |_| { + format!("set_only_v6 failed: {}", &address) + })?; + io_convert(socket.bind(&address.into()), |_| { + format!("bind failed: {}", &address) + })?; + let main_channel: UdpSocket = socket.into(); + main_channel.set_write_timeout(Some(Duration::from_secs(5)))?; + udps.push(main_channel); + } + let context = Context::new(udps, first_latency, is_tcp); + + let port = context.main_local_udp_port()?[0]; + //监听v6+v4双栈，tcp通道使用异步io + let address: SocketAddr = format!("[::]:{}", port).parse().unwrap(); + let socket = socket2::Socket::new(socket2::Domain::IPV6, socket2::Type::STREAM, None)?; + io_convert(socket.set_only_v6(false), |_| { + format!("set_only_v6 failed: {}", &address) + })?; + + if let Err(e) = socket.bind(&address.into()) { + if ports[0] == 0 { + //端口可能冲突，则使用任意端口 + log::warn!("监听tcp端口失败 {:?},重试一次", address); + let address: SocketAddr = format!("[::]:{}", 0).parse().unwrap(); + io_convert(socket.bind(&address.into()), |_| { + format!("bind failed: {}", &address) + })?; + } else { + //手动指定的ip,直接报错 + io_convert(Err(e), |_| format!("bind failed: {}", &address))?; + } + } + socket.listen(2)?; + socket.set_nonblocking(true)?; + socket.set_nodelay(false)?; + let tcp_listener = mio::net::TcpListener::from_std(socket.into()); + Ok((context, tcp_listener)) +} + +pub fn init_channel( + tcp_listener: mio::net::TcpListener, + context: Context, + stop_manager: StopManager, + recv_handler: H, +) -> io::Result<( + AcceptSocketSender>>, + AcceptSocketSender<(mio::net::TcpStream, SocketAddr, Option>)>, +)> +where + H: RecvChannelHandler, +{ + // udp监听，udp_socket_sender 用于NAT类型切换 + let udp_socket_sender = + udp_listen(stop_manager.clone(), recv_handler.clone(), context.clone())?; + // 建立tcp监听，tcp_socket_sender 用于tcp 直连 + let tcp_socket_sender = tcp_listen( + tcp_listener, + stop_manager.clone(), + recv_handler.clone(), + context.clone(), + )?; + + Ok((udp_socket_sender, tcp_socket_sender)) +} diff --git a/vnt/src/channel/notify.rs b/vnt/src/channel/notify.rs new file mode 100644 index 0000000..5a6dc57 --- /dev/null +++ b/vnt/src/channel/notify.rs @@ -0,0 +1,126 @@ +use mio::{Token, Waker}; +use parking_lot::Mutex; +use std::io; +use std::ops::Deref; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +#[derive(Clone)] +pub struct WritableNotify { + inner: Arc, +} + +impl WritableNotify { + pub fn new(waker: Waker) -> Self { + Self { + inner: Arc::new(WritableNotifyInner { + waker, + state: AtomicUsize::new(0), + tokens: Mutex::new(Vec::with_capacity(8)), + }), + } + } +} + +impl Deref for WritableNotify { + type Target = WritableNotifyInner; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +pub struct WritableNotifyInner { + waker: Waker, + state: AtomicUsize, + tokens: Mutex>, +} + +impl WritableNotifyInner { + pub fn notify(&self, token: Token, state: bool) -> io::Result<()> { + { + let mut guard = self.tokens.lock(); + if guard.is_empty() || !guard.contains(&(token, state)) { + guard.push((token, state)); + } + drop(guard); + } + self.need_write() + } + + pub fn stop(&self) -> io::Result<()> { + self.state.store(0b001, Ordering::Release); + self.waker.wake() + } + pub fn need_write(&self) -> io::Result<()> { + self.state.fetch_or(0b010, Ordering::AcqRel); + self.waker.wake() + } + pub fn add_socket(&self) -> io::Result<()> { + self.state.fetch_or(0b100, Ordering::AcqRel); + self.waker.wake() + } + pub fn take_all(&self) -> Option> { + let mut guard = self.tokens.lock(); + if guard.is_empty() { + None + } else { + Some(guard.drain(..).collect()) + } + } + pub fn is_stop(&self) -> bool { + self.state.load(Ordering::Acquire) & 0b001 == 0b001 + } + pub fn is_need_write(&self) -> bool { + self.state.fetch_and(!0b010, Ordering::AcqRel) & 0b010 == 0b010 + } + pub fn is_add_socket(&self) -> bool { + self.state.fetch_and(!0b100, Ordering::AcqRel) & 0b100 == 0b100 + } +} + +#[derive(Clone)] +pub struct AcceptNotify { + inner: Arc, +} + +impl AcceptNotify { + pub fn new(waker: Waker) -> Self { + Self { + inner: Arc::new(AcceptNotifyInner { + waker, + state: AtomicUsize::new(0), + }), + } + } +} + +impl Deref for AcceptNotify { + type Target = AcceptNotifyInner; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +pub struct AcceptNotifyInner { + waker: Waker, + state: AtomicUsize, +} + +impl AcceptNotifyInner { + pub fn is_stop(&self) -> bool { + self.state.load(Ordering::Acquire) & 0b001 == 0b001 + } + pub fn is_add_socket(&self) -> bool { + self.state.fetch_and(!0b100, Ordering::AcqRel) & 0b100 == 0b100 + } + pub fn stop(&self) -> io::Result<()> { + self.state.store(0b001, Ordering::Release); + self.waker.wake() + } + pub fn add_socket(&self) -> io::Result<()> { + self.state.fetch_or(0b100, Ordering::AcqRel); + self.waker.wake() + } +} diff --git a/vnt/src/channel/punch.rs b/vnt/src/channel/punch.rs index 71bc121..a0d2aaf 100644 --- a/vnt/src/channel/punch.rs +++ b/vnt/src/channel/punch.rs @@ -1,13 +1,14 @@ use std::collections::HashMap; -use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, TcpStream}; +use std::io; +use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; use std::str::FromStr; use std::time::Duration; -use std::{io, thread}; +use mio::net::TcpStream; use rand::prelude::SliceRandom; -use crate::channel::channel::{send_tcp, start_tcp_handle, Context}; -use crate::handle::recv_handler::ChannelDataHandler; +use crate::channel::context::Context; +use crate::channel::sender::AcceptSocketSender; #[derive(Copy, Clone, Eq, PartialEq, Debug)] pub enum PunchModel { @@ -28,15 +29,21 @@ impl FromStr for PunchModel { } } +impl Default for PunchModel { + fn default() -> Self { + PunchModel::All + } +} + #[derive(Clone, Debug)] pub struct NatInfo { pub public_ips: Vec, - pub public_port: u16, + pub public_ports: Vec, pub public_port_range: u16, pub nat_type: NatType, pub(crate) local_ipv4: Option, pub(crate) ipv6: Option, - pub(crate) udp_port: u16, + pub(crate) udp_ports: Vec, pub tcp_port: u16, } @@ -49,11 +56,11 @@ pub enum NatType { impl NatInfo { pub fn new( mut public_ips: Vec, - public_port: u16, + public_ports: Vec, public_port_range: u16, mut local_ipv4: Option, mut ipv6: Option, - udp_port: u16, + udp_ports: Vec, tcp_port: u16, mut nat_type: NatType, ) -> Self { @@ -79,16 +86,16 @@ impl NatInfo { } Self { public_ips, - public_port, + public_ports, public_port_range, local_ipv4, ipv6, - udp_port, + udp_ports, tcp_port, nat_type, } } - pub fn update_addr(&mut self, ip: Ipv4Addr, port: u16) { + pub fn update_addr(&mut self, index: usize, ip: Ipv4Addr, port: u16) { if !ip.is_multicast() && !ip.is_broadcast() && !ip.is_unspecified() @@ -96,7 +103,9 @@ impl NatInfo { && !ip.is_private() && port != 0 { - self.public_port = port; + if let Some(public_port) = self.public_ports.get_mut(index) { + *public_port = port; + } if !self.public_ips.contains(&ip) { self.public_ips.push(ip); } @@ -108,22 +117,32 @@ impl NatInfo { pub fn ipv6(&self) -> Option { self.ipv6 } - pub fn local_udp_ipv4addr(&self) -> Option { - if self.udp_port == 0 { + pub fn local_udp_ipv4addr(&self, index: usize) -> Option { + let len = self.udp_ports.len(); + if len == 0 { return None; } if let Some(local_ipv4) = self.local_ipv4 { - Some(SocketAddr::V4(SocketAddrV4::new(local_ipv4, self.udp_port))) + Some(SocketAddr::V4(SocketAddrV4::new( + local_ipv4, + self.udp_ports[index % len], + ))) } else { None } } - pub fn local_udp_ipv6addr(&self) -> Option { - if self.udp_port == 0 { + pub fn local_udp_ipv6addr(&self, index: usize) -> Option { + let len = self.udp_ports.len(); + if len == 0 { return None; } if let Some(ipv6) = self.ipv6 { - Some(SocketAddr::V6(SocketAddrV6::new(ipv6, self.udp_port, 0, 0))) + Some(SocketAddr::V6(SocketAddrV6::new( + ipv6, + self.udp_ports[index % len], + 0, + 0, + ))) } else { None } @@ -158,7 +177,7 @@ pub struct Punch { port_index: HashMap, punch_model: PunchModel, is_tcp: bool, - handler: ChannelDataHandler, + tcp_socket_sender: AcceptSocketSender<(TcpStream, SocketAddr, Option>)>, } impl Punch { @@ -166,7 +185,7 @@ impl Punch { context: Context, punch_model: PunchModel, is_tcp: bool, - handler: ChannelDataHandler, + tcp_socket_sender: AcceptSocketSender<(TcpStream, SocketAddr, Option>)>, ) -> Self { let mut port_vec: Vec = (1..65535).collect(); port_vec.push(65535); @@ -178,30 +197,23 @@ impl Punch { port_index: HashMap::new(), punch_model, is_tcp, - handler, + tcp_socket_sender, } } } impl Punch { - fn connect_tcp(&self, buf: &[u8], addr: &SocketAddr) -> bool { - match TcpStream::connect_timeout(&addr, Duration::from_secs(1)) { - Ok(mut tcp_stream) => { - let context = self.context.clone(); - let handler = self.handler.clone(); - match send_tcp(&mut tcp_stream, buf) { - Ok(_) => {} - Err(e) => { - log::warn!("发送到tcp失败,addr={},err={}", addr, e); - return false; - } + fn connect_tcp(&self, buf: &[u8], addr: SocketAddr) -> bool { + // mio是非阻塞的，不能立马判断是否能连接成功，所以用标准库的tcp + match std::net::TcpStream::connect_timeout(&addr, Duration::from_secs(1)) { + Ok(tcp_stream) => { + if tcp_stream.set_nonblocking(true).is_err() { + return false; } - thread::spawn(move || { - if let Err(e) = start_tcp_handle(tcp_stream, context, handler) { - log::error!("{:?}", e); - } - }); - return true; + return self + .tcp_socket_sender + .try_add_socket((TcpStream::from_std(tcp_stream), addr, Some(buf.to_vec()))) + .is_ok(); } Err(e) => { log::warn!("连接到tcp失败,addr={},err={}", addr, e); @@ -209,41 +221,48 @@ impl Punch { } false } - pub async fn punch(&mut self, buf: &[u8], id: Ipv4Addr, nat_info: NatInfo) -> io::Result<()> { - if !self.context.need_punch(&id) { + pub fn punch(&mut self, buf: &[u8], id: Ipv4Addr, nat_info: NatInfo) -> io::Result<()> { + if !self.context.route_table.need_punch(&id) { return Ok(()); } + log::info!("nat_info={:?}", nat_info); + if self.is_tcp { //向tcp发起连接 if let Some(ipv6_addr) = nat_info.local_tcp_ipv6addr() { - if self.connect_tcp(buf, &ipv6_addr) { + if self.connect_tcp(buf, ipv6_addr) { return Ok(()); } } - log::info!("local_tcp_ipv4addr={:?}", nat_info.local_tcp_ipv4addr()); //向tcp发起连接 if let Some(ipv4_addr) = nat_info.local_tcp_ipv4addr() { - if self.connect_tcp(buf, &ipv4_addr) { + if self.connect_tcp(buf, ipv4_addr) { return Ok(()); } } if nat_info.nat_type == NatType::Cone && nat_info.public_ips.len() == 1 { let addr = SocketAddr::V4(SocketAddrV4::new(nat_info.public_ips[0], nat_info.tcp_port)); - if self.connect_tcp(buf, &addr) { + if self.connect_tcp(buf, addr) { return Ok(()); } } } - if let Some(ipv4_addr) = nat_info.local_udp_ipv4addr() { - let _ = self.context.send_main_udp(buf, ipv4_addr); + let channel_num = self.context.channel_num(); + for index in 0..channel_num { + if let Some(ipv4_addr) = nat_info.local_udp_ipv4addr(index) { + let _ = self.context.send_main_udp(index, buf, ipv4_addr); + } } + if self.punch_model != PunchModel::IPv4 { - if let Some(ipv6_addr) = nat_info.local_udp_ipv6addr() { - let rs = self.context.send_main_udp(buf, ipv6_addr); - log::info!("发送到ipv6地址:{:?},rs={:?}", ipv6_addr, rs); - if rs.is_ok() && self.punch_model == PunchModel::IPv6 { - return Ok(()); + for index in 0..channel_num { + if let Some(ipv6_addr) = nat_info.local_udp_ipv6addr(index) { + let rs = self.context.send_main_udp(index, buf, ipv6_addr); + log::info!("发送到ipv6地址:{:?},rs={:?}", ipv6_addr, rs); + if rs.is_ok() && self.punch_model == PunchModel::IPv6 { + return Ok(()); + } } } } @@ -258,16 +277,15 @@ impl Punch { let max_k1 = 60; //全局最多发送max_k2个包 let max_k2 = 800; + let port = nat_info.public_ports.get(0).map(|e| *e).unwrap_or(0); if nat_info.public_port_range < max_k1 * 3 { //端口变化不大时，在预测的范围内随机发送 - let min_port = if nat_info.public_port > nat_info.public_port_range { - nat_info.public_port - nat_info.public_port_range + let min_port = if port > nat_info.public_port_range { + port - nat_info.public_port_range } else { 1 }; - let (max_port, overflow) = nat_info - .public_port - .overflowing_add(nat_info.public_port_range); + let (max_port, overflow) = port.overflowing_add(nat_info.public_port_range); let max_port = if overflow { 65535 } else { max_port }; let k = if max_port - min_port + 1 > max_k1 { max_k1 as usize @@ -280,8 +298,7 @@ impl Punch { let mut rng = rand::thread_rng(); nums.shuffle(&mut rng); } - self.punch_symmetric(&nums[..k], buf, &nat_info.public_ips, max_k1 as usize) - .await?; + self.punch_symmetric(&nums[..k], buf, &nat_info.public_ips, max_k1 as usize)?; } let start = *self.port_index.entry(id.clone()).or_insert(0); let mut end = start + max_k2; @@ -295,21 +312,23 @@ impl Punch { buf, &nat_info.public_ips, max_k2, - ) - .await?; + )?; self.port_index.insert(id, index); } NatType::Cone => { - if nat_info.public_port != 0 { + for index in 0..channel_num { let is_cone = self.context.is_cone(); - for ip in nat_info.public_ips { - let addr = SocketAddr::V4(SocketAddrV4::new(ip, nat_info.public_port)); - self.context.send_main_udp(buf, addr)?; + let len = nat_info.public_ports.len(); + for ip in &nat_info.public_ips { + let addr = SocketAddr::V4(SocketAddrV4::new( + *ip, + nat_info.public_ports[index % len], + )); + self.context.send_main_udp(index, buf, addr)?; if !is_cone { //只有一方是对称，则对称方要使用全部端口发送数据，符合上述计算的概率 - self.context.try_send_all(buf, addr)?; + self.context.try_send_all(buf, addr); } - tokio::time::sleep(Duration::from_millis(2)).await; } } } @@ -317,7 +336,7 @@ impl Punch { Ok(()) } - async fn punch_symmetric( + fn punch_symmetric( &self, ports: &[u16], buf: &[u8], @@ -332,8 +351,7 @@ impl Punch { return Ok(()); } let addr = SocketAddr::V4(SocketAddrV4::new(*pub_ip, *port)); - self.context.send_main_udp(buf, addr)?; - tokio::time::sleep(Duration::from_millis(2)).await; + self.context.send_main_udp(0, buf, addr)?; } } Ok(()) diff --git a/vnt/src/channel/sender.rs b/vnt/src/channel/sender.rs index 83b42d6..0ec1c79 100644 --- a/vnt/src/channel/sender.rs +++ b/vnt/src/channel/sender.rs @@ -1,5 +1,12 @@ -use crate::channel::channel::Context; +use std::io; use std::ops::Deref; +use std::sync::mpsc::{SyncSender, TrySendError}; +use std::sync::Arc; + +use mio::Token; + +use crate::channel::context::Context; +use crate::channel::notify::{AcceptNotify, WritableNotify}; #[derive(Clone)] pub struct ChannelSender { @@ -19,3 +26,80 @@ impl Deref for ChannelSender { &self.context } } +pub struct AcceptSocketSender { + sender: SyncSender, + notify: AcceptNotify, +} + +impl Clone for AcceptSocketSender { + fn clone(&self) -> Self { + Self { + sender: self.sender.clone(), + notify: self.notify.clone(), + } + } +} +impl AcceptSocketSender { + pub fn new(notify: AcceptNotify, sender: SyncSender) -> Self { + Self { sender, notify } + } + pub fn try_add_socket(&self, t: T) -> io::Result<()> { + match self.sender.try_send(t) { + Ok(_) => self.notify.add_socket(), + Err(e) => match e { + TrySendError::Full(_) => Err(io::Error::from(io::ErrorKind::WouldBlock)), + TrySendError::Disconnected(_) => Err(io::Error::from(io::ErrorKind::WriteZero)), + }, + } + } +} + +#[derive(Clone)] +pub struct PacketSender { + inner: Arc, +} + +impl PacketSender { + pub fn new(notify: WritableNotify, buffer: SyncSender>, token: Token) -> Self { + Self { + inner: Arc::new(PacketSenderInner { + token, + notify, + buffer, + }), + } + } + #[inline] + pub fn try_send(&self, buf: &[u8]) -> io::Result<()> { + self.inner.try_send(buf) + } + pub fn shutdown(&self) -> io::Result<()> { + self.inner.shutdown() + } +} + +pub struct PacketSenderInner { + token: Token, + notify: WritableNotify, + buffer: SyncSender>, +} + +impl PacketSenderInner { + #[inline] + fn try_send(&self, buf: &[u8]) -> io::Result<()> { + let len = buf.len(); + let mut buf_vec = Vec::with_capacity(buf.len() + 4); + buf_vec.extend_from_slice(&[0, 0, (len >> 8) as u8, (len & 0xFF) as u8]); + buf_vec.extend_from_slice(buf); + match self.buffer.try_send(buf_vec) { + Ok(_) => self.notify.notify(self.token, true), + Err(e) => match e { + TrySendError::Disconnected(_) => Err(io::Error::from(io::ErrorKind::WriteZero)), + TrySendError::Full(_) => Err(io::Error::from(io::ErrorKind::WouldBlock)), + }, + } + } + fn shutdown(&self) -> io::Result<()> { + self.notify.notify(self.token, false) + } +} diff --git a/vnt/src/channel/tcp_channel.rs b/vnt/src/channel/tcp_channel.rs new file mode 100644 index 0000000..05f1ec0 --- /dev/null +++ b/vnt/src/channel/tcp_channel.rs @@ -0,0 +1,456 @@ +use std::collections::HashMap; +use std::io::{Read, Write}; +use std::net::{Shutdown, SocketAddr}; +#[cfg(any(unix))] +use std::os::fd::FromRawFd; +#[cfg(any(unix))] +use std::os::fd::IntoRawFd; +#[cfg(windows)] +use std::os::windows::io::FromRawSocket; +#[cfg(windows)] +use std::os::windows::io::IntoRawSocket; +use std::sync::mpsc::{sync_channel, Receiver, SyncSender, TryRecvError, TrySendError}; +use std::{io, thread}; + +use mio::net::{TcpListener, TcpStream}; +use mio::{Events, Interest, Poll, Registry, Token, Waker}; + +use crate::channel::context::Context; +use crate::channel::handler::RecvChannelHandler; +use crate::channel::notify::{AcceptNotify, WritableNotify}; +use crate::channel::sender::{AcceptSocketSender, PacketSender}; +use crate::channel::{RouteKey, BUFFER_SIZE}; +use crate::util::StopManager; + +const SERVER: Token = Token(0); +const NOTIFY: Token = Token(1); + +/// 监听tcp端口，等待客户端连接 +pub fn tcp_listen( + tcp_server: TcpListener, + stop_manager: StopManager, + recv_handler: H, + context: Context, +) -> io::Result>)>> +where + H: RecvChannelHandler, +{ + let (tcp_sender, tcp_receiver) = sync_channel(64); + let poll = Poll::new()?; + let waker = AcceptNotify::new(Waker::new(poll.registry(), NOTIFY)?); + let accept = AcceptSocketSender::new(waker.clone(), tcp_sender); + let worker = { + let waker = waker.clone(); + stop_manager.add_listener("tcp_listen".into(), move || { + if let Err(e) = waker.stop() { + log::error!("{:?}", e); + } + })? + }; + + thread::Builder::new() + .name("tcp读事件处理线程".into()) + .spawn(move || { + if let Err(e) = tcp_listen0( + poll, + tcp_server, + &stop_manager, + waker, + tcp_receiver, + recv_handler, + context, + ) { + log::error!("{:?}", e); + } + worker.stop_all(); + })?; + Ok(accept) +} + +fn tcp_listen0( + mut poll: Poll, + mut tcp_server: TcpListener, + stop_manager: &StopManager, + accept_notify: AcceptNotify, + accept_tcp_receiver: Receiver<(TcpStream, SocketAddr, Option>)>, + mut recv_handler: H, + context: Context, +) -> io::Result<()> +where + H: RecvChannelHandler, +{ + let (tcp_sender, tcp_receiver) = sync_channel(64); + let write_waker = init_writable_handler(tcp_receiver, stop_manager.clone(), context.clone())?; + poll.registry() + .register(&mut tcp_server, SERVER, Interest::READABLE)?; + let mut events = Events::with_capacity(1024); + + let mut read_map: HashMap, usize)> = + HashMap::with_capacity(32); + loop { + poll.poll(&mut events, None)?; + for event in events.iter() { + match event.token() { + SERVER => loop { + match tcp_server.accept() { + Ok((stream, addr)) => { + accept_handle( + stream, + addr, + None, + &write_waker, + &mut read_map, + &tcp_sender, + poll.registry(), + )?; + } + Err(e) => { + if e.kind() == io::ErrorKind::WouldBlock { + break; + } + return Err(e); + } + } + }, + NOTIFY => { + if accept_notify.is_stop() { + return Ok(()); + } + if accept_notify.is_add_socket() { + while let Ok((stream, addr, init_buf)) = accept_tcp_receiver.try_recv() { + accept_handle( + stream, + addr, + init_buf, + &write_waker, + &mut read_map, + &tcp_sender, + poll.registry(), + )?; + } + } + } + token => { + if event.is_readable() { + if let Err(e) = + readable_handle(&token, &mut read_map, &mut recv_handler, &context) + { + closed_handle_r(&token, &mut read_map); + log::warn!("{:?}", e); + if let Err(e) = write_waker.notify(token, false) { + log::warn!("{:?}", e); + } + } + } else { + closed_handle_r(&token, &mut read_map); + if let Err(e) = write_waker.notify(token, false) { + log::warn!("{:?}", e); + } + } + } + } + } + } +} + +/// 处理写事件 + +fn init_writable_handler( + receiver: Receiver<(TcpStream, Token, SocketAddr, Option>)>, + stop_manager: StopManager, + context: Context, +) -> io::Result { + let poll = Poll::new()?; + let writable_notify = WritableNotify::new(Waker::new(poll.registry(), NOTIFY)?); + let worker = { + let writable_notify = writable_notify.clone(); + stop_manager.add_listener("tcp_writable_handler".into(), move || { + if let Err(e) = writable_notify.stop() { + log::error!("{:?}", e); + } + })? + }; + { + let writable_notify = writable_notify.clone(); + thread::Builder::new() + .name("tcp-writeable-listen".into()) + .spawn(move || { + if let Err(e) = tcp_writable_listen(receiver, poll, writable_notify, &context) { + log::error!("{:?}", e); + } + worker.stop_all(); + })?; + } + + Ok(writable_notify) +} + +/// 处理写事件 +fn tcp_writable_listen( + receiver: Receiver<(TcpStream, Token, SocketAddr, Option>)>, + mut poll: Poll, + writable_notify: WritableNotify, + context: &Context, +) -> io::Result<()> { + let mut events = Events::with_capacity(1024); + let mut write_map: HashMap< + Token, + ( + TcpStream, + SocketAddr, + Receiver>, + Option<(Vec, usize)>, + ), + > = HashMap::with_capacity(32); + loop { + poll.poll(&mut events, None)?; + for event in events.iter() { + match event.token() { + NOTIFY => { + if writable_notify.is_stop() { + //服务停止 + return Ok(()); + } + if writable_notify.is_need_write() { + // 需要写入数据 + if let Some(tokens) = writable_notify.take_all() { + for (token, state) in tokens { + if !state { + closed_handle_w(&token, &mut write_map, &context); + continue; + } + if let Err(e) = writable_handle(&token, &mut write_map) { + closed_handle_w(&token, &mut write_map, &context); + log::warn!("{:?}", e); + } + } + } + } + if writable_notify.is_add_socket() { + //添加tcp连接，并监听写事件 + while let Ok((mut stream, token, addr, init_buf)) = receiver.try_recv() { + if let Err(e) = stream.set_nodelay(true) { + log::warn!("set_nodelay err={:?}", e); + } + if let Err(e) = + poll.registry() + .register(&mut stream, token, Interest::WRITABLE) + { + log::warn!("registry err={:?}", e); + continue; + } + let (sender, receiver) = sync_channel(128); + let packet_sender = + PacketSender::new(writable_notify.clone(), sender, token); + if let Some(init_buf) = init_buf { + packet_sender.try_send(&init_buf)?; + } + + context.tcp_map.write().insert(addr, packet_sender); + write_map.insert(token, (stream, addr, receiver, None)); + } + } + } + token => { + if event.is_writable() { + if let Err(e) = writable_handle(&token, &mut write_map) { + closed_handle_w(&token, &mut write_map, &context); + log::warn!("{:?}", e); + } + } else { + closed_handle_w(&token, &mut write_map, &context); + } + } + } + } + } +} + +fn accept_handle( + stream: TcpStream, + addr: SocketAddr, + init_buf: Option>, + write_waker: &WritableNotify, + read_map: &mut HashMap, usize)>, + tcp_sender: &SyncSender<(TcpStream, Token, SocketAddr, Option>)>, + registry: &Registry, +) -> io::Result<()> { + #[cfg(windows)] + let (tcp_stream, index) = unsafe { + let fd = stream.into_raw_socket(); + (std::net::TcpStream::from_raw_socket(fd), fd as usize) + }; + #[cfg(any(unix))] + let (tcp_stream, index) = unsafe { + let fd = stream.into_raw_fd(); + (std::net::TcpStream::from_raw_fd(fd), fd as usize) + }; + if index == 0 || index == 1 { + log::error!("index err={:?}", addr); + return Ok(()); + } + let token = Token(index); + match tcp_stream.try_clone() { + Ok(tcp_writer) => { + match tcp_sender.try_send((TcpStream::from_std(tcp_writer), token, addr, init_buf)) { + Ok(_) => { + if let Err(e) = write_waker.add_socket() { + log::error!("write_waker,err={:?},addr={:?}", e, addr); + return Ok(()); + } + } + Err(e) => { + return match e { + TrySendError::Full(_) => { + log::error!("Full,addr={:?}", addr); + Ok(()) + } + TrySendError::Disconnected(_) => { + Err(io::Error::new(io::ErrorKind::Other, "write thread exit")) + } + }; + } + } + } + Err(e) => { + log::error!("try_clone err={:?},addr={:?}", e, addr); + return Ok(()); + } + } + let mut stream = TcpStream::from_std(tcp_stream); + if let Err(e) = registry.register(&mut stream, token, Interest::READABLE) { + log::error!("registry err={:?},addr={:?}", e, addr); + return Ok(()); + } + read_map.insert( + token, + ( + RouteKey::new(true, index, addr), + stream, + Box::new([0; BUFFER_SIZE]), + 0, + ), + ); + Ok(()) +} + +fn readable_handle( + token: &Token, + map: &mut HashMap, usize)>, + recv_handler: &mut H, + context: &Context, +) -> io::Result<()> +where + H: RecvChannelHandler, +{ + if let Some((route_key, stream, buf, begin)) = map.get_mut(token) { + loop { + let end = if *begin >= 4 { + 4 + (((buf[2] as u16) << 8) | buf[3] as u16) as usize + } else { + 4 + }; + if end > BUFFER_SIZE { + return Err(io::Error::from(io::ErrorKind::InvalidData)); + } + match stream.read(&mut buf[*begin..end]) { + Ok(len) => { + if len == 0 { + return Err(io::Error::from(io::ErrorKind::UnexpectedEof)); + } + *begin += len; + if end > 4 && *begin == end { + recv_handler.handle(&mut buf[4..end], *route_key, context); + *begin = 0; + } + } + Err(e) => { + if e.kind() == io::ErrorKind::WouldBlock { + break; + } + return Err(e); + } + } + } + } + Ok(()) +} + +fn writable_handle( + token: &Token, + map: &mut HashMap< + Token, + ( + TcpStream, + SocketAddr, + Receiver>, + Option<(Vec, usize)>, + ), + >, +) -> io::Result<()> { + if let Some((stream, _, receiver, last)) = map.get_mut(token) { + loop { + if let Some((buf, begin)) = last { + match stream.write(&buf[*begin..]) { + Ok(len) => { + if len == 0 { + return Err(io::Error::from(io::ErrorKind::WriteZero)); + } + if len + *begin == buf.len() { + *last = None; + } else { + *begin += len; + continue; + } + } + Err(e) => { + if e.kind() == io::ErrorKind::WouldBlock { + break; + } + return Err(e); + } + } + } + match receiver.try_recv() { + Ok(buf) => *last = Some((buf, 0)), + Err(e) => match e { + TryRecvError::Empty => { + break; + } + TryRecvError::Disconnected => { + return Err(io::Error::from(io::ErrorKind::Other)); + } + }, + } + } + } + Ok(()) +} + +fn closed_handle_r( + token: &Token, + map: &mut HashMap, usize)>, +) { + if let Some((_, tcp, _, _)) = map.remove(token) { + let _ = tcp.shutdown(Shutdown::Both); + } +} + +fn closed_handle_w( + token: &Token, + map: &mut HashMap< + Token, + ( + TcpStream, + SocketAddr, + Receiver>, + Option<(Vec, usize)>, + ), + >, + context: &Context, +) { + if let Some((tcp, addr, _, _)) = map.remove(token) { + context.tcp_map.write().remove(&addr); + let _ = tcp.shutdown(Shutdown::Both); + } +} diff --git a/vnt/src/channel/udp_channel.rs b/vnt/src/channel/udp_channel.rs new file mode 100644 index 0000000..2cb123d --- /dev/null +++ b/vnt/src/channel/udp_channel.rs @@ -0,0 +1,219 @@ +use std::collections::HashMap; +use std::net::{IpAddr, UdpSocket as StdUdpSocket}; +use std::net::{Ipv4Addr, SocketAddr}; +use std::sync::mpsc::{sync_channel, Receiver}; +use std::{io, thread}; + +use mio::event::Source; +use mio::net::UdpSocket; +use mio::{Events, Interest, Poll, Token, Waker}; + +use crate::channel::context::Context; +use crate::channel::handler::RecvChannelHandler; +use crate::channel::notify::AcceptNotify; +use crate::channel::sender::AcceptSocketSender; +use crate::channel::{RouteKey, BUFFER_SIZE}; +use crate::util::StopManager; + +pub fn udp_listen( + stop_manager: StopManager, + recv_handler: H, + context: Context, +) -> io::Result>>> +where + H: RecvChannelHandler, +{ + //根据通道数创建对应线程进行读取 + for index in 0..context.channel_num() { + main_udp_listen( + index, + stop_manager.clone(), + recv_handler.clone(), + context.clone(), + )?; + } + sub_udp_listen(stop_manager, recv_handler, context) +} + +const NOTIFY: Token = Token(0); + +fn sub_udp_listen( + stop_manager: StopManager, + recv_handler: H, + context: Context, +) -> io::Result>>> +where + H: RecvChannelHandler, +{ + let (udp_sender, udp_receiver) = sync_channel(64); + let poll = Poll::new()?; + let waker = AcceptNotify::new(Waker::new(poll.registry(), NOTIFY)?); + let worker = { + let waker = waker.clone(); + stop_manager.add_listener("sub_udp_listen".into(), move || { + if let Err(e) = waker.stop() { + log::error!("{:?}", e); + } + })? + }; + let accept = AcceptSocketSender::new(waker.clone(), udp_sender); + thread::Builder::new() + .name("sub_udp读事件处理线程".into()) + .spawn(move || { + if let Err(e) = sub_udp_listen0(poll, recv_handler, context, waker, udp_receiver) { + log::error!("{:?}", e); + } + worker.stop_all(); + })?; + Ok(accept) +} + +fn sub_udp_listen0( + mut poll: Poll, + mut recv_handler: H, + context: Context, + accept_notify: AcceptNotify, + accept_receiver: Receiver>>, +) -> io::Result<()> +where + H: RecvChannelHandler, +{ + let mut events = Events::with_capacity(1024); + let mut buf = [0; BUFFER_SIZE]; + let mut read_map: HashMap = HashMap::with_capacity(32); + loop { + poll.poll(&mut events, None)?; + for event in events.iter() { + match event.token() { + NOTIFY => { + if accept_notify.is_stop() { + return Ok(()); + } + if accept_notify.is_add_socket() { + while let Ok(option) = accept_receiver.try_recv() { + match option { + None => { + log::info!("切换成锥形模式"); + for (_, mut udp_socket) in read_map.drain() { + if let Err(e) = udp_socket.deregister(poll.registry()) { + log::error!("{:?}", e); + } + } + } + Some(socket_list) => { + log::info!("切换成对称模式 监听端口数：{}", socket_list.len()); + for (index, mut udp_socket) in + socket_list.into_iter().enumerate() + { + let token = Token(index + context.channel_num()); + poll.registry().register( + &mut udp_socket, + token, + Interest::READABLE, + )?; + read_map.insert(token, udp_socket); + } + } + } + } + } + } + token => { + if let Some(udp_socket) = read_map.get(&token) { + loop { + match udp_socket.recv_from(&mut buf) { + Ok((len, addr)) => { + recv_handler.handle( + &mut buf[..len], + RouteKey::new(false, token.0, addr), + &context, + ); + } + Err(e) => { + if e.kind() == io::ErrorKind::WouldBlock { + break; + } + log::error!("{:?}", e); + } + } + } + } + } + } + } + } +} + +/// 阻塞监听 +fn main_udp_listen( + index: usize, + stop_manager: StopManager, + recv_handler: H, + context: Context, +) -> io::Result<()> +where + H: RecvChannelHandler, +{ + let port = context.main_udp_socket[index].local_addr()?.port(); + let worker = stop_manager.add_listener(format!("main_udp_listen-{}", index), move || { + 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); + } + } + })?; + thread::Builder::new() + .name("main_udp读事件处理线程".into()) + .spawn(move || { + if let Err(e) = main_udp_listen0(index, recv_handler, context) { + log::error!("{:?}", e); + } + worker.stop_all(); + })?; + Ok(()) +} + +pub fn main_udp_listen0(index: usize, mut recv_handler: H, context: Context) -> io::Result<()> +where + H: RecvChannelHandler, +{ + let mut buf = [0; BUFFER_SIZE]; + let udp_socket = &context.main_udp_socket[index]; + loop { + match udp_socket.recv_from(&mut buf) { + Ok((len, addr)) => { + if &buf[..len] == b"stop" { + match addr.ip() { + IpAddr::V4(ip) => { + if ip.is_loopback() { + return Ok(()); + } + } + IpAddr::V6(ip) => { + if ip.is_loopback() { + return Ok(()); + } + if let Some(ip) = ip.to_ipv4_mapped() { + if ip.is_loopback() { + return Ok(()); + } + } + } + } + } + recv_handler.handle(&mut buf[..len], RouteKey::new(false, index, addr), &context); + } + Err(e) => { + log::error!("main_udp_listen0={:?}", e); + } + } + } +}