yeyinlaoying/vnt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[mio] 支持多通道传输，使用mio代替tokio

Browse Source
This commit is contained in:
lubeilin
2024-02-29 22:06:20 +08:00
parent e7f1165287
commit af5d284470
10 changed files with 1518 additions and 864 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vnt/src/channel/channel.rs
-764
View File
@@ -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<Mutex<TcpStream>>,
pub(crate) route_table: RwLock<HashMap<Ipv4Addr, Vec<(Route, AtomicCell<Instant>)>>>,
pub(crate) status_receiver: Receiver<Status>,
pub(crate) status_sender: Sender<Status>,
pub(crate) udp_map: RwLock<HashMap<usize, Arc<UdpSocket>>>,
pub(crate) tcp_map: RwLock<HashMap<usize, Arc<Mutex<TcpStream>>>>,
pub(crate) channel_num: usize,
current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
first_latency: bool,
is_close: AtomicBool,
tcp_port: u16,
}
#[derive(Clone)]
pub struct Context {
pub(crate) inner: Arc<ContextInner>,
}
impl Context {
pub fn new(
main_channel: StdUdpSocket,
main_tcp_channel: Option<TcpStream>,
current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
_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<u16> {
self.inner.main_channel.local_addr().map(|k| k.port())
}
fn insert_udp(&self, id: usize, udp: Arc<UdpSocket>) {
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<usize> {
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<usize> {
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<TcpStream>, buf: &[u8]) -> io::Result<usize> {
let mut stream = sender.lock();
send_tcp(&mut stream, buf)
}
pub fn send_main(&self, buf: &[u8], addr: SocketAddr) -> io::Result<usize> {
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<usize> {
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<usize> {
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<Route> {
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<usize> {
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<usize> {
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<Arc<UdpSocket>> {
self.inner.udp_map.read().get(&route_key.index).cloned()
}
#[inline]
fn get_tcp_by_route(&self, route_key: &RouteKey) -> Option<Arc<Mutex<TcpStream>>> {
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<Vec<Route>> {
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<Route> {
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<Ipv4Addr> {
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<Route>)> {
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<TcpStream>,
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<UdpSocket>,
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<usize> {
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)
}
vnt/src/channel/context.rs
+420
View File
@@ -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<ContextInner>,
}
impl Context {
pub fn new(main_udp_socket: Vec<UdpSocket>, 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<UdpSocket>,
// 对称网络增加的udp socket
sub_udp_socket: RwLock<Vec<UdpSocket>>,
// tcp数据发送器
pub(crate) tcp_map: RwLock<HashMap<SocketAddr, PacketSender>>,
// 路由信息
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<Option<Vec<mio::net::UdpSocket>>>,
) -> 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>,
) -> 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<Vec<u16>> {
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<HashMap<Ipv4Addr, (AtomicUsize, Vec<(Route, AtomicCell<Instant>)>)>>,
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<Route> {
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<Vec<Route>> {
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<Route> {
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<Ipv4Addr> {
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<Route>)> {
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;
}
}
}
}
}
vnt/src/channel/handler.rs
+6
View File
@@ -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);
}
vnt/src/channel/idle.rs
+22 -25
View File
@@ -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);
}
}
vnt/src/channel/mod.rs
+97 -5
View File
@@ -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<u16>,
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<H>(
tcp_listener: mio::net::TcpListener,
context: Context,
stop_manager: StopManager,
recv_handler: H,
) -> io::Result<(
AcceptSocketSender<Option<Vec<mio::net::UdpSocket>>>,
AcceptSocketSender<(mio::net::TcpStream, SocketAddr, Option<Vec<u8>>)>,
)>
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))
}
vnt/src/channel/notify.rs
+126
View File
@@ -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<WritableNotifyInner>,
}
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<Vec<(Token, bool)>>,
}
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<Vec<(Token, bool)>> {
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<AcceptNotifyInner>,
}
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()
}
}
vnt/src/channel/punch.rs
+87 -69
View File
@@ -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<Ipv4Addr>,
pub public_port: u16,
pub public_ports: Vec<u16>,
pub public_port_range: u16,
pub nat_type: NatType,
pub(crate) local_ipv4: Option<Ipv4Addr>,
pub(crate) ipv6: Option<Ipv6Addr>,
pub(crate) udp_port: u16,
pub(crate) udp_ports: Vec<u16>,
pub tcp_port: u16,
}
@@ -49,11 +56,11 @@ pub enum NatType {
impl NatInfo {
pub fn new(
mut public_ips: Vec<Ipv4Addr>,
public_port: u16,
public_ports: Vec<u16>,
public_port_range: u16,
mut local_ipv4: Option<Ipv4Addr>,
mut ipv6: Option<Ipv6Addr>,
udp_port: u16,
udp_ports: Vec<u16>,
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<Ipv6Addr> {
self.ipv6
}
pub fn local_udp_ipv4addr(&self) -> Option<SocketAddr> {
if self.udp_port == 0 {
pub fn local_udp_ipv4addr(&self, index: usize) -> Option<SocketAddr> {
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<SocketAddr> {
if self.udp_port == 0 {
pub fn local_udp_ipv6addr(&self, index: usize) -> Option<SocketAddr> {
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<Ipv4Addr, usize>,
punch_model: PunchModel,
is_tcp: bool,
handler: ChannelDataHandler,
tcp_socket_sender: AcceptSocketSender<(TcpStream, SocketAddr, Option<Vec<u8>>)>,
}
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<Vec<u8>>)>,
) -> Self {
let mut port_vec: Vec<u16> = (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(())
vnt/src/channel/sender.rs
+85 -1
View File
@@ -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<T> {
sender: SyncSender<T>,
notify: AcceptNotify,
}
impl<T> Clone for AcceptSocketSender<T> {
fn clone(&self) -> Self {
Self {
sender: self.sender.clone(),
notify: self.notify.clone(),
}
}
}
impl<T> AcceptSocketSender<T> {
pub fn new(notify: AcceptNotify, sender: SyncSender<T>) -> 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<PacketSenderInner>,
}
impl PacketSender {
pub fn new(notify: WritableNotify, buffer: SyncSender<Vec<u8>>, 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<Vec<u8>>,
}
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)
}
}
vnt/src/channel/tcp_channel.rs
+456
View File
@@ -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<H>(
tcp_server: TcpListener,
stop_manager: StopManager,
recv_handler: H,
context: Context,
) -> io::Result<AcceptSocketSender<(TcpStream, SocketAddr, Option<Vec<u8>>)>>
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<H>(
mut poll: Poll,
mut tcp_server: TcpListener,
stop_manager: &StopManager,
accept_notify: AcceptNotify,
accept_tcp_receiver: Receiver<(TcpStream, SocketAddr, Option<Vec<u8>>)>,
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<Token, (RouteKey, TcpStream, Box<[u8; BUFFER_SIZE]>, 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<Vec<u8>>)>,
stop_manager: StopManager,
context: Context,
) -> io::Result<WritableNotify> {
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<Vec<u8>>)>,
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<Vec<u8>>,
Option<(Vec<u8>, 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<Vec<u8>>,
write_waker: &WritableNotify,
read_map: &mut HashMap<Token, (RouteKey, TcpStream, Box<[u8; BUFFER_SIZE]>, usize)>,
tcp_sender: &SyncSender<(TcpStream, Token, SocketAddr, Option<Vec<u8>>)>,
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<H>(
token: &Token,
map: &mut HashMap<Token, (RouteKey, TcpStream, Box<[u8; BUFFER_SIZE]>, 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<Vec<u8>>,
Option<(Vec<u8>, 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<Token, (RouteKey, TcpStream, Box<[u8; BUFFER_SIZE]>, 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<Vec<u8>>,
Option<(Vec<u8>, usize)>,
),
>,
context: &Context,
) {
if let Some((tcp, addr, _, _)) = map.remove(token) {
context.tcp_map.write().remove(&addr);
let _ = tcp.shutdown(Shutdown::Both);
}
}
vnt/src/channel/udp_channel.rs
+219
View File
@@ -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<H>(
stop_manager: StopManager,
recv_handler: H,
context: Context,
) -> io::Result<AcceptSocketSender<Option<Vec<UdpSocket>>>>
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<H>(
stop_manager: StopManager,
recv_handler: H,
context: Context,
) -> io::Result<AcceptSocketSender<Option<Vec<UdpSocket>>>>
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<H>(
mut poll: Poll,
mut recv_handler: H,
context: Context,
accept_notify: AcceptNotify,
accept_receiver: Receiver<Option<Vec<UdpSocket>>>,
) -> io::Result<()>
where
H: RecvChannelHandler,
{
let mut events = Events::with_capacity(1024);
let mut buf = [0; BUFFER_SIZE];
let mut read_map: HashMap<Token, UdpSocket> = 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<H>(
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<H>(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);
}
}
}
}