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

增加tcp通道，去除并行逻辑

Browse Source
This commit is contained in:
lubeilin
2024-01-06 12:06:21 +08:00
parent 57b903ea29
commit 7098111ad1
5 changed files with 85 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vnt/proto/message.proto
+1
View File
@@ -56,6 +56,7 @@ message PunchInfo{
uint32 local_port = 8;
bytes ipv6 = 9;
uint32 ipv6_port = 10;
uint32 tcp_port = 11;
}
enum PunchNatType{
Symmetric = 0;
vnt/src/channel/channel.rs
+49 -110
View File
@@ -3,6 +3,7 @@ use std::io::{Read, Write};
use std::net::UdpSocket as StdUdpSocket;
use std::net::{Ipv4Addr, Shutdown, SocketAddr};
use std::net::{SocketAddrV6, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::{io, thread};
@@ -27,9 +28,11 @@ pub struct ContextInner {
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,
}
#[derive(Clone)]
@@ -56,9 +59,11 @@ impl Context {
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),
});
Self { inner }
}
@@ -66,21 +71,38 @@ impl Context {
impl Context {
pub fn is_close(&self) -> bool {
*self.inner.status_receiver.borrow() == Status::Close
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<()> {
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() {
let _ = StdUdpSocket::bind("127.0.0.1:0")?.send_to(
b"stop",
SocketAddr::V4(std::net::SocketAddrV4::new(Ipv4Addr::LOCALHOST, 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 {
tcp.lock().shutdown(Shutdown::Both)?;
if let Err(e) = tcp.lock().shutdown(Shutdown::Both) {
log::error!("发送停止消息到tcp失败:{:?}", e);
}
}
for (_, tcp) in self.inner.tcp_map.read().clone() {
if let Err(e) = tcp.lock().shutdown(Shutdown::Both) {
log::error!("发送停止消息到tcp失败:{:?}", e);
}
}
Ok(())
}
@@ -371,37 +393,6 @@ impl Channel {
}
}
#[derive(Clone)]
struct BufSenderGroup(
usize,
Vec<std::sync::mpsc::SyncSender<(Vec<u8>, usize, usize, RouteKey)>>,
);
struct BufReceiverGroup(Vec<std::sync::mpsc::Receiver<(Vec<u8>, usize, usize, RouteKey)>>);
impl BufSenderGroup {
pub fn send(&mut self, val: (Vec<u8>, usize, usize, RouteKey)) -> bool {
let index = self.0 % self.1.len();
self.0 = self.0.wrapping_add(1);
self.1[index].send(val).is_ok()
}
}
fn buf_channel_group(size: usize) -> (BufSenderGroup, BufReceiverGroup) {
let mut buf_sender_group = Vec::with_capacity(size);
let mut buf_receiver_group = Vec::with_capacity(size);
for _ in 0..size {
let (buf_sender, buf_receiver) =
std::sync::mpsc::sync_channel::<(Vec<u8>, usize, usize, RouteKey)>(1);
buf_sender_group.push(buf_sender);
buf_receiver_group.push(buf_receiver);
}
(
BufSenderGroup(0, buf_sender_group),
BufReceiverGroup(buf_receiver_group),
)
}
impl Channel {
fn tcp_handle(
tcp_r: &mut TcpStream,
@@ -411,7 +402,7 @@ impl Channel {
) -> io::Result<()> {
let mut head = [0; 4];
let addr = tcp_r.peer_addr()?;
let key = RouteKey::new(TCP_ID, addr);
let key = RouteKey::new(true, TCP_ID, addr);
loop {
if context.is_close() {
return Ok(());
@@ -480,32 +471,10 @@ impl Channel {
head_reserve: usize, //头部预留字节
symmetric_channel_num: usize, //对称网络，则再加一组监听，提升打洞成功率
relay: bool,
parallel: usize,
) {
let handler = self.handler.clone();
let context = self.context;
let main_channel = context.inner.main_channel.try_clone().unwrap();
let buf_sender = if parallel > 1 {
let (buf_sender, buf_receiver) = buf_channel_group(parallel);
let mut num = 0;
for buf_receiver in buf_receiver.0 {
let context = context.clone();
let handler = handler.clone();
thread::Builder::new()
.name(format!("recv-handler-{}", num))
.spawn(move || {
while let Ok((mut buf, start, end, route_key)) = buf_receiver.recv() {
handler.handle(&mut buf, start, end, route_key, &context);
}
log::warn!("异步处理停止");
})
.unwrap();
num += 1;
}
Some(buf_sender)
} else {
None
};
if let Some(tcp_stream) = tcp {
let context = context.clone();
let handler = handler.clone();
@@ -527,15 +496,7 @@ impl Channel {
.name("channel_udp".into())
.spawn(move || {
log::info!("启动udp v4");
Self::main_start_(
worker,
context,
UDP_ID,
main_channel,
handler,
buf_sender,
head_reserve,
)
Self::main_start_(worker, context, UDP_ID, main_channel, handler, head_reserve)
})
.unwrap();
}
@@ -597,52 +558,30 @@ impl Channel {
id: usize,
udp: StdUdpSocket,
handler: ChannelDataHandler,
buf_sender: Option<BufSenderGroup>,
head_reserve: usize,
) {
match buf_sender {
None => {
let mut buf = [0; 4096];
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(id, addr),
&context,
);
}
Err(e) => {
log::error!("udp :{:?}", e);
let mut buf = [0; 4096];
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);
}
}
Some(mut buf_sender) => loop {
let mut buf = vec![0; 4096];
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;
}
}
buf_sender.send((buf, head_reserve, end, RouteKey::new(id, addr)));
}
Err(e) => {
log::error!("udp :{:?}", e);
}
}
},
}
worker.stop_all();
@@ -671,7 +610,7 @@ impl Channel {
rs=udp.recv_from(&mut buf[head_reserve..])=>{
match rs {
Ok((len, addr)) => {
handler.handle(&mut buf, head_reserve, head_reserve + len, RouteKey::new(id, addr), &context);
handler.handle(&mut buf, head_reserve, head_reserve + len, RouteKey::new(false,id, addr), &context);
}
Err(e) => {
log::error!("{:?}",e)
vnt/src/channel/mod.rs
+12 -3
View File
@@ -17,6 +17,7 @@ pub enum Status {
#[derive(Copy, Clone, Debug)]
pub struct Route {
is_tcp: bool,
index: usize,
pub addr: SocketAddr,
pub metric: u8,
@@ -30,8 +31,9 @@ pub struct RouteSortKey {
}
impl Route {
pub fn new(index: usize, addr: SocketAddr, metric: u8, rt: i64) -> Self {
pub fn new(is_tcp: bool, index: usize, addr: SocketAddr, metric: u8, rt: i64) -> Self {
Self {
is_tcp,
index,
addr,
metric,
@@ -40,6 +42,7 @@ impl Route {
}
pub fn from(route_key: RouteKey, metric: u8, rt: i64) -> Self {
Self {
is_tcp: route_key.is_tcp,
index: route_key.index,
addr: route_key.addr,
metric,
@@ -48,6 +51,7 @@ impl Route {
}
pub fn route_key(&self) -> RouteKey {
RouteKey {
is_tcp: self.is_tcp,
index: self.index,
addr: self.addr,
}
@@ -65,13 +69,18 @@ impl Route {
#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
pub struct RouteKey {
is_tcp: bool,
index: usize,
pub addr: SocketAddr,
}
impl RouteKey {
pub(crate) fn new(index: usize, addr: SocketAddr) -> Self {
Self { index, addr }
pub(crate) fn new(is_tcp: bool, index: usize, addr: SocketAddr) -> Self {
Self {
is_tcp,
index,
addr,
}
}
pub fn is_tcp(&self) -> bool {
self.index == TCP_ID
vnt/src/core/mod.rs
+1 -1
View File
@@ -382,7 +382,7 @@ impl VntUtil {
let relay = config.relay;
tokio::spawn(async move {
channel
.start(channel_worker, tcp_receiver, 14, 65, relay, config.parallel)
.start(channel_worker, tcp_receiver, 14, 65, relay)
.await
});
}
vnt/src/proto/message.rs
+22 -4
View File
@@ -1317,6 +1317,8 @@ pub struct PunchInfo {
pub ipv6: ::std::vec::Vec<u8>,
// @@protoc_insertion_point(field:PunchInfo.ipv6_port)
pub ipv6_port: u32,
// @@protoc_insertion_point(field:PunchInfo.tcp_port)
pub tcp_port: u32,
// special fields
// @@protoc_insertion_point(special_field:PunchInfo.special_fields)
pub special_fields: ::protobuf::SpecialFields,
@@ -1334,7 +1336,7 @@ impl PunchInfo {
}
fn generated_message_descriptor_data() -> ::protobuf::reflect::GeneratedMessageDescriptorData {
let mut fields = ::std::vec::Vec::with_capacity(9);
let mut fields = ::std::vec::Vec::with_capacity(10);
let mut oneofs = ::std::vec::Vec::with_capacity(0);
fields.push(::protobuf::reflect::rt::v2::make_vec_simpler_accessor::<_, _>(
"public_ip_list",
@@ -1381,6 +1383,11 @@ impl PunchInfo {
|m: &PunchInfo| { &m.ipv6_port },
|m: &mut PunchInfo| { &mut m.ipv6_port },
));
fields.push(::protobuf::reflect::rt::v2::make_simpler_field_accessor::<_, _>(
"tcp_port",
|m: &PunchInfo| { &m.tcp_port },
|m: &mut PunchInfo| { &mut m.tcp_port },
));
::protobuf::reflect::GeneratedMessageDescriptorData::new_2::<PunchInfo>(
"PunchInfo",
fields,
@@ -1429,6 +1436,9 @@ impl ::protobuf::Message for PunchInfo {
80 => {
self.ipv6_port = is.read_uint32()?;
},
88 => {
self.tcp_port = is.read_uint32()?;
},
tag => {
::protobuf::rt::read_unknown_or_skip_group(tag, is, self.special_fields.mut_unknown_fields())?;
},
@@ -1466,6 +1476,9 @@ impl ::protobuf::Message for PunchInfo {
if self.ipv6_port != 0 {
my_size += ::protobuf::rt::uint32_size(10, self.ipv6_port);
}
if self.tcp_port != 0 {
my_size += ::protobuf::rt::uint32_size(11, self.tcp_port);
}
my_size += ::protobuf::rt::unknown_fields_size(self.special_fields.unknown_fields());
self.special_fields.cached_size().set(my_size as u32);
my_size
@@ -1499,6 +1512,9 @@ impl ::protobuf::Message for PunchInfo {
if self.ipv6_port != 0 {
os.write_uint32(10, self.ipv6_port)?;
}
if self.tcp_port != 0 {
os.write_uint32(11, self.tcp_port)?;
}
os.write_unknown_fields(self.special_fields.unknown_fields())?;
::std::result::Result::Ok(())
}
@@ -1525,6 +1541,7 @@ impl ::protobuf::Message for PunchInfo {
self.local_port = 0;
self.ipv6.clear();
self.ipv6_port = 0;
self.tcp_port = 0;
self.special_fields.clear();
}
@@ -1539,6 +1556,7 @@ impl ::protobuf::Message for PunchInfo {
local_port: 0,
ipv6: ::std::vec::Vec::new(),
ipv6_port: 0,
tcp_port: 0,
special_fields: ::protobuf::SpecialFields::new(),
};
&instance
@@ -1643,7 +1661,7 @@ static file_descriptor_proto_data: &'static [u8] = b"\
\n\rdevice_status\x18\x03\x20\x01(\rR\x0cdeviceStatus\x12#\n\rclient_sec\
ret\x18\x04\x20\x01(\x08R\x0cclientSecret\"Y\n\nDeviceList\x12\x14\n\x05\
epoch\x18\x01\x20\x01(\rR\x05epoch\x125\n\x10device_info_list\x18\x02\
\x20\x03(\x0b2\x0b.DeviceInfoR\x0edeviceInfoList\"\xa9\x02\n\tPunchInfo\
\x20\x03(\x0b2\x0b.DeviceInfoR\x0edeviceInfoList\"\xc4\x02\n\tPunchInfo\
\x12$\n\x0epublic_ip_list\x18\x02\x20\x03(\x07R\x0cpublicIpList\x12\x1f\
\n\x0bpublic_port\x18\x03\x20\x01(\rR\npublicPort\x12*\n\x11public_port_\
range\x18\x04\x20\x01(\rR\x0fpublicPortRange\x12(\n\x08nat_type\x18\x05\
@@ -1651,8 +1669,8 @@ static file_descriptor_proto_data: &'static [u8] = b"\
\x01(\x08R\x05reply\x12\x19\n\x08local_ip\x18\x07\x20\x01(\x07R\x07local\
Ip\x12\x1d\n\nlocal_port\x18\x08\x20\x01(\rR\tlocalPort\x12\x12\n\x04ipv\
6\x18\t\x20\x01(\x0cR\x04ipv6\x12\x1b\n\tipv6_port\x18\n\x20\x01(\rR\x08\
ipv6Port*'\n\x0cPunchNatType\x12\r\n\tSymmetric\x10\0\x12\x08\n\x04Cone\
\x10\x01b\x06proto3\
ipv6Port\x12\x19\n\x08tcp_port\x18\x0b\x20\x01(\rR\x07tcpPort*'\n\x0cPun\
chNatType\x12\r\n\tSymmetric\x10\0\x12\x08\n\x04Cone\x10\x01b\x06proto3\
";
/// `FileDescriptorProto` object which was a source for this generated file