diff --git a/switch/src/channel/channel.rs b/switch/src/channel/channel.rs new file mode 100644 index 0000000..193ae5e --- /dev/null +++ b/switch/src/channel/channel.rs @@ -0,0 +1,398 @@ +use std::io; +use std::net::{Ipv4Addr, SocketAddr}; +use std::sync::Arc; +use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering}; +use crossbeam_skiplist::SkipMap; +use dashmap::DashMap; +use tokio::net::UdpSocket; +use tokio::sync::Notify; +use tokio::sync::watch::{channel, Receiver, Sender}; +use crate::channel::{Route, RouteKey, Status}; +use crate::channel::punch::NatType; +use crate::handle::recv_handler::ChannelDataHandler; + +#[derive(Clone)] +pub struct Context { + pub(crate) count: Arc, + pub(crate) main_channel: Arc, + pub(crate) route_table: Arc>>, + pub(crate) route_table_time: Arc>, + pub(crate) status_receiver: Receiver, + pub(crate) status_sender: Arc>, + pub(crate) udp_map: Arc>>, + pub(crate) channel_num: usize, + pub(crate) notify: Arc, +} + +impl Context { + pub fn new(main_channel: Arc, _channel_num: usize) -> Self { + //当前版本只支持一个通道 + let channel_num = 1; + let (status_sender, status_receiver) = channel(Status::Cone); + let status_sender = Arc::new(status_sender); + Self { + count: Arc::new(AtomicUsize::new(0)), + main_channel, + route_table: Arc::new(DashMap::with_capacity(16)), + route_table_time: Arc::new(SkipMap::new()), + status_receiver, + status_sender, + udp_map: Arc::new(SkipMap::new()), + channel_num, + notify: Arc::new(Notify::new()), + } + } +} + +impl Context { + pub fn is_close(&self) -> bool { + *self.status_receiver.borrow() == Status::Close + } + pub fn is_cone(&self) -> bool { + *self.status_receiver.borrow() == Status::Cone + } + pub fn close(&self) { + let _ = self.status_sender.send(Status::Close); + } + 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.status_sender.send(Status::Cone); + } + pub fn switch_to_symmetric(&self) { + let _ = self.status_sender.send(Status::Symmetric); + } + pub fn main_local_port(&self) -> io::Result { + self.main_channel.local_addr().map(|k| k.port()) + } + pub async fn send_main(&self, buf: &[u8], addr: SocketAddr) -> io::Result { + self.main_channel.send_to(buf, addr).await + } + + pub(crate) async fn send_all(&self, buf: &[u8], addr: SocketAddr) -> io::Result<()> { + for udp in self.udp_map.iter() { + udp.value().send_to(buf, addr).await?; + } + Ok(()) + } + pub fn try_send_main(&self, buf: &[u8], addr: SocketAddr) -> io::Result { + self.main_channel.try_send_to(buf, addr) + } + pub async fn send_by_id(&self, buf: &[u8], id: &Ipv4Addr) -> io::Result { + if let Some(v) = self.route_table.get(id) { + let route = match v.len() { + 0 => { + return Err(io::Error::new(io::ErrorKind::NotFound, "route not found")); + } + 1 => &v[0], + len => &v[self.count.fetch_add(1, Ordering::Relaxed) % len] + }; + if let Some(udp) = self.udp_map.get(&route.index) { + return udp.value().send_to(buf, route.addr).await; + } + } + Err(io::Error::new(io::ErrorKind::NotFound, "route not found")) + } + pub fn try_send_by_id(&self, buf: &[u8], id: &Ipv4Addr) -> io::Result { + if let Some(v) = self.route_table.get(id) { + if v.is_empty() { + return Err(io::Error::new(io::ErrorKind::NotFound, "route not found")); + } + let route = &v[self.count.fetch_add(1, Ordering::Relaxed) % v.len()]; + if let Some(udp) = self.udp_map.get(&route.index) { + return udp.value().try_send_to(buf, route.addr); + } + } + Err(io::Error::new(io::ErrorKind::NotFound, "route not found")) + } + pub async fn send_by_key(&self, buf: &[u8], route_key: &RouteKey) -> io::Result { + if let Some(udp) = self.udp_map.get(&route_key.index) { + return udp.value().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 { + if let Some(udp) = self.udp_map.get(&route_key.index) { + return udp.value().try_send_to(buf, route_key.addr); + } + 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 ref_mut = self.route_table.entry(id.clone()).or_insert(Vec::with_capacity(4)); + let mut exist = false; + for x in ref_mut.iter_mut() { + if x.metric < route.metric { + //不能比当前的路径更长 + return; + } + if x.route_key() == key { + if only_if_absent { + return; + } + x.metric = route.metric; + x.rt = route.rt; + exist = true; + break; + } + } + if !exist { + if route.metric == 1 { + //添加了直连的则排除非直连的 + ref_mut.retain(|k| k.metric == 1); + } + ref_mut.push(route); + let max_len = self.channel_num; + if ref_mut.len() > max_len { + ref_mut.sort_by_key(|k| k.sort_key()); + ref_mut.truncate(max_len); + } + } + self.route_table_time.insert((key, id), AtomicI64::new(chrono::Local::now().timestamp_millis())); + self.notify.notify_one(); + } + pub fn route(&self, id: &Ipv4Addr) -> Option> { + if let Some(v) = self.route_table.get(id) { + Some(v.value().clone()) + } else { + None + } + } + pub fn route_one(&self, id: &Ipv4Addr) -> Option { + if let Some(v) = self.route_table.get(id) { + v.value().iter().max_by_key(|k| k.sort_key()).map(|k| *k) + } else { + None + } + } + pub fn route_to_id(&self, route_key: &RouteKey) -> Option { + for x in self.route_table_time.iter() { + if &x.key().0 == route_key { + return Some(x.key().1); + } + } + None + } + pub fn need_punch(&self, id: &Ipv4Addr) -> bool { + if let Some(v) = self.route_table.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)> { + self.route_table.iter().map(|k| (k.key().clone(), k.value().clone())).collect() + } + pub fn route_table_one(&self) -> Vec<(Ipv4Addr, Route)> { + let mut v = Vec::with_capacity(8); + for x in self.route_table.iter() { + if let Some(route) = x.value().iter().max_by_key(|k| k.sort_key()) { + v.push((*x.key(), *route)); + } + } + v + } + pub fn direct_route_table_one(&self) -> Vec<(Ipv4Addr, Route)> { + let mut v = Vec::with_capacity(8); + for x in self.route_table.iter() { + if let Some(route) = x.value().iter().max_by_key(|k| k.sort_key()) { + if route.metric == 1 { + v.push((*x.key(), *route)); + } + } + } + v + } + pub fn remove_route_all(&self, id: &Ipv4Addr) { + if let Some((_, v)) = self.route_table.remove(id) { + for x in v { + self.route_table_time.remove(&(x.route_key(), id.clone())); + } + } + } + pub fn remove_route(&self, id: &Ipv4Addr, route_key: RouteKey) { + if let Some(mut v) = self.route_table.get_mut(id) { + v.retain(|x| x.route_key() != route_key); + self.route_table_time.remove(&(route_key, id.clone())); + } + } + pub fn update_read_time(&self, id: &Ipv4Addr, route_key: &RouteKey) { + if let Some(time) = self.route_table_time.get(&(*route_key, *id)) { + time.value().store(chrono::Local::now().timestamp_millis(), Ordering::Relaxed); + } + } +} + +pub struct Channel { + context: Context, + handler: ChannelDataHandler, +} + +impl Channel { + pub fn new(context: Context, + handler: ChannelDataHandler, ) -> Self { + Self { + context, + handler, + } + } +} + +impl Channel { + async fn handle(handler: &mut ChannelDataHandler, + udp: &Arc, + context: &Context, + id: usize, + result: io::Result<(usize, SocketAddr)>, + buf: &mut [u8], start: usize) { + match result { + Ok((len, addr)) => { + handler.handle(buf, start, start + len, RouteKey::new(id, addr), &udp, context).await; + } + Err(e) => { + log::error!("{:?}",e) + } + } + } + pub async fn start(self, + head_reserve: usize,//头部预留字节 + symmetric_channel_num: usize,//对称网络,则再加一组监听,提升打洞成功率 + ) { + let mut context = self.context; + let main_channel = context.main_channel.clone(); + let handler = self.handler.clone(); + tokio::spawn(Self::start_(context.clone(), handler, main_channel, head_reserve, true)); + let mut cur_status = Status::Cone; + loop { + match context.status_receiver.changed().await { + Ok(_) => { + match *context.status_receiver.borrow() { + Status::Cone => { + cur_status = Status::Cone; + } + Status::Symmetric => { + if cur_status == Status::Symmetric { + continue; + } + cur_status = Status::Symmetric; + for _ in 0..symmetric_channel_num { + match UdpSocket::bind("0.0.0.0:0").await { + Ok(udp) => { + let udp = Arc::new(udp); + let context = context.clone(); + let handler = self.handler.clone(); + tokio::spawn(Self::start_(context, handler, udp, head_reserve, false)); + } + Err(e) => { + log::error!("{}",e); + } + } + } + } + Status::Close => { + break; + } + } + } + Err(_) => { + break; + } + } + } + } + async fn start_(context: Context, + mut handler: ChannelDataHandler, + udp: Arc, + head_reserve: usize, + is_core: bool) { + let mut status_receiver = context.status_receiver.clone(); + #[cfg(target_os = "windows")] + use std::os::windows::io::AsRawSocket; + #[cfg(target_os = "windows")] + let id = udp.as_raw_socket() as usize; + #[cfg(any(unix))] + use std::os::fd::AsRawFd; + #[cfg(any(unix))] + let id = udp.as_raw_fd() as usize; + context.udp_map.insert(id, udp.clone()); + let mut buf = [0; 65546]; + loop { + tokio::select! { + rs=udp.recv_from(&mut buf[head_reserve..])=>{ + Self::handle(&mut handler,&udp,&context,id,rs,&mut buf,head_reserve).await; + } + changed=status_receiver.changed()=>{ + match changed { + Ok(_) => { + match *status_receiver.borrow() { + Status::Cone => { + if !is_core{ + break; + } + } + Status::Close=>{ + break; + } + Status::Symmetric => {} + } + } + Err(_) => { + break; + } + } + } + } + } + context.udp_map.remove(&id); + } +} + + +// pub async fn start(mut handler: H, +// mut status_receiver: Receiver, +// head_reserve: usize, +// core_channel_num: usize, +// symmetric_channel_num: usize) -> io::Result<()> { +// for _ in 0..core_channel_num { +// let d = channel(1); +// } +// let udp = UdpSocket::bind("0.0.0.0:0").await?; +// let d = status_receiver.changed().await; +// match d { +// Ok(_) => { +// match *status_receiver.borrow() { +// Status::Cone => {} +// Status::Symmetric => {} +// Status::Close => {} +// } +// } +// Err(_) => {} +// } +// let mut buf = [0; 65546]; +// let result = udp.recv_from(&mut buf[head_reserve..]).await; +// match result { +// Ok((len, addr)) => {} +// Err(e) => {} +// } +// Ok(()) +// } +// +// pub struct Channel { +// handler: H, +// +// } \ No newline at end of file diff --git a/switch/src/channel/idle.rs b/switch/src/channel/idle.rs new file mode 100644 index 0000000..46584b4 --- /dev/null +++ b/switch/src/channel/idle.rs @@ -0,0 +1,62 @@ +use std::io; +use std::io::{Error, ErrorKind}; +use std::net::Ipv4Addr; +use std::sync::atomic::Ordering; +use std::time::Duration; +use crate::channel::channel::Context; +use crate::channel::RouteKey; + + +pub struct Idle { + read_idle: i64, + context: Context, +} + +impl Idle { + pub fn new(read_idle: i64, + context: Context, ) -> Self { + Self { + read_idle, + context, + } + } +} + +impl Idle { + /// 获取空闲路由 + pub async fn next_idle(&self) -> io::Result<(Ipv4Addr, RouteKey)> { + loop { + let now = chrono::Local::now().timestamp_millis(); + let last_read_idle = now - self.read_idle; + let mut min = i64::MAX; + for entry in self.context.route_table_time.iter() { + let mut is_read_idle = false; + if self.read_idle > 0 { + let last_read = entry.value().load(Ordering::Relaxed); + if last_read < last_read_idle { + is_read_idle = true; + } else { + if min > last_read { + min = last_read; + } + } + } + if is_read_idle { + return Ok((entry.key().1.clone(), entry.key().0.clone())); + } + } + if self.context.route_table_time.is_empty() { + self.context.notify.notified().await; + } else { + let sleep_time = chrono::Local::now().timestamp_millis() - min; + if sleep_time > 0 { + tokio::time::sleep(Duration::from_millis(sleep_time as u64)).await; + // let _ = tokio::time::timeout(Duration::from_millis(sleep_time as u64), self.context.notify.notified()).await; + } + } + if self.context.is_close() { + return Err(Error::new(ErrorKind::Other, "closed")); + } + } + } +} \ No newline at end of file diff --git a/switch/src/channel/mod.rs b/switch/src/channel/mod.rs new file mode 100644 index 0000000..28304f9 --- /dev/null +++ b/switch/src/channel/mod.rs @@ -0,0 +1,78 @@ +use std::net::SocketAddr; + +pub mod channel; +pub mod punch; +pub mod idle; +pub mod sender; + +#[derive(Copy, Clone, Eq, PartialEq)] +pub enum Status { + Cone, + Symmetric, + Close, +} + +#[derive(Copy, Clone, Debug)] +pub struct Route { + index: usize, + pub addr: SocketAddr, + pub metric: u8, + pub rt: i64, +} + +#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)] +pub struct RouteSortKey { + pub metric: u8, + pub rt: i64, +} + +impl Route { + pub fn new(index: usize, + addr: SocketAddr, metric: u8, rt: i64, ) -> Self { + Self { + index, + addr, + metric, + rt, + } + } + pub fn from(route_key: RouteKey, metric: u8, rt: i64) -> Self { + Self { + index: route_key.index, + addr: route_key.addr, + metric, + rt, + } + } + pub fn route_key(&self) -> RouteKey { + RouteKey { + index: self.index, + addr: self.addr, + } + } + pub fn sort_key(&self) -> RouteSortKey { + RouteSortKey { + metric: self.metric, + rt: self.rt, + } + } + pub fn is_p2p(&self) -> bool { + self.metric == 1 + } +} + +#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)] +pub struct RouteKey { + index: usize, + pub addr: SocketAddr, +} + +impl RouteKey { + pub(crate) fn new(index: usize, + addr: SocketAddr, ) -> Self { + Self { + index, + addr, + } + } +} \ No newline at end of file diff --git a/switch/src/channel/punch.rs b/switch/src/channel/punch.rs new file mode 100644 index 0000000..3c57d9b --- /dev/null +++ b/switch/src/channel/punch.rs @@ -0,0 +1,152 @@ +use std::collections::HashMap; +use std::io; +use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; +use std::time::Duration; + +use rand::prelude::SliceRandom; + +use crate::channel::channel::Context; + +#[derive(Clone, Debug)] +pub struct NatInfo { + pub public_ips: Vec, + pub public_port: u16, + pub public_port_range: u16, + pub local_ip: Ipv4Addr, + pub local_port: u16, + pub nat_type: NatType, +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)] +pub enum NatType { + Symmetric, + Cone, +} + +impl NatInfo { + pub fn new(public_ips: Vec, + public_port: u16, + public_port_range: u16, + local_ip: Ipv4Addr, + local_port: u16, + nat_type: NatType, ) -> Self { + Self { + public_ips, + public_port, + public_port_range, + local_ip, + local_port, + nat_type, + } + } +} + +#[derive(Clone)] +pub struct Punch { + context: Context, + port_vec: Vec, + port_index: HashMap, +} + +impl Punch { + pub fn new(context: Context) -> Self { + let mut port_vec: Vec = (1..65535).collect(); + port_vec.push(65535); + let mut rng = rand::thread_rng(); + port_vec.shuffle(&mut rng); + Punch { + context, + port_vec, + port_index: HashMap::new(), + } + } +} + +impl Punch { + pub async fn punch(&mut self, buf: &[u8], id: Ipv4Addr, nat_info: NatInfo) -> io::Result<()> { + if !self.context.need_punch(&id) { + return Ok(()); + } + if !nat_info.local_ip.is_unspecified() || nat_info.local_port != 0 { + let _ = self.context.send_main(buf, SocketAddr::V4(SocketAddrV4::new(nat_info.local_ip, nat_info.local_port))).await; + } + match nat_info.nat_type { + NatType::Symmetric => { + // 假设对方绑定n个端口,通过NAT对外映射出n个 公网ip:公网端口,自己随机尝试k次的情况下 + // 猜中的概率 p = 1-((65535-n)/65535)*((65535-n-1)/(65535-1))*...*((65535-n-k+1)/(65535-k+1)) + // n取76,k取600,猜中的概率就超过50%了 + // 前提 自己是锥形网络,否则猜中了也通信不了 + + //预测范围内最多发送max_k1个包 + let max_k1 = 60; + //全局最多发送max_k2个包 + let max_k2 = 800; + 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 + } else { + 1 + }; + let (max_port, overflow) = nat_info.public_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 + } else { + (max_port - min_port + 1) as usize + }; + let mut nums: Vec = (min_port..max_port).collect(); + nums.push(max_port); + { + 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?; + } + let start = *self.port_index.entry(id.clone()).or_insert(0); + let mut end = start + max_k2; + let mut index = end; + if end >= self.port_vec.len() { + end = self.port_vec.len(); + index = 0 + } + self.punch_symmetric(&self.port_vec[start..end], buf, &nat_info.public_ips, max_k2).await?; + self.port_index.insert(id, index); + } + NatType::Cone => { + 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)); + if is_cone { + self.context.send_main(buf, addr).await?; + } else { + //只有一方是对称,则对称方要使用全部端口发送数据,符合上述计算的概率 + self.context.send_all(buf, addr).await?; + } + tokio::time::sleep(Duration::from_millis(2)).await; + } + } + } + Ok(()) + } + + async fn punch_symmetric(&self, ports: &[u16], buf: &[u8], ips: &Vec, max: usize) -> io::Result<()> { + let mut count = 0; + for port in ports { + for pub_ip in ips { + count += 1; + if count == max { + return Ok(()); + } + let addr = SocketAddr::V4(SocketAddrV4::new(*pub_ip, *port)); + self.context.send_main(buf, addr).await?; + tokio::time::sleep(Duration::from_millis(2)).await; + } + } + Ok(()) + } +} diff --git a/switch/src/channel/sender.rs b/switch/src/channel/sender.rs new file mode 100644 index 0000000..e906401 --- /dev/null +++ b/switch/src/channel/sender.rs @@ -0,0 +1,23 @@ +use std::ops::Deref; +use crate::channel::channel::Context; + +#[derive(Clone)] +pub struct ChannelSender { + context: Context, +} + +impl ChannelSender { + pub fn new(context: Context) -> Self { + Self { + context, + } + } +} + +impl Deref for ChannelSender { + type Target = Context; + + fn deref(&self) -> &Self::Target { + &self.context + } +}