archive/tuic
1
0
Fork 0
Code

moving next_assoc_id&udp_sessionsintoServer

Browse Source
This commit is contained in:
EAimTY 2023-02-02 20:55:13 +09:00
parent 7d395ca825
commit 61477f5094
1 changed files with 23 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
tuic-client/src/socks5.rs
View File

@ -1,6 +1,6 @@
use crate::{config::Local, connection::Connection as TuicConnection, error::Error}; use crate::{config::Local, connection::Connection as TuicConnection, error::Error};
use bytes::Bytes; use bytes::Bytes;
use once_cell::sync::{Lazy, OnceCell}; use once_cell::sync::OnceCell;
use parking_lot::Mutex; use parking_lot::Mutex;
use socket2::{Domain, Protocol, SockAddr, Socket, Type}; use socket2::{Domain, Protocol, SockAddr, Socket, Type};
use socks5_proto::{Address, Reply}; use socks5_proto::{Address, Reply};
@ -26,15 +26,14 @@ use tokio_util::compat::FuturesAsyncReadCompatExt;
use tuic::Address as TuicAddress; use tuic::Address as TuicAddress;
static SERVER: OnceCell<Server> = OnceCell::new(); static SERVER: OnceCell<Server> = OnceCell::new();
static NEXT_ASSOCIATE_ID: AtomicU16 = AtomicU16::new(0);
static UDP_SESSIONS: Lazy<Mutex<HashMap<u16, Arc<AssociatedUdpSocket>>>> =
Lazy::new(|| Mutex::new(HashMap::new()));
pub struct Server { pub struct Server {
inner: Socks5Server, inner: Socks5Server,
addr: SocketAddr, addr: SocketAddr,
dual_stack: Option<bool>, dual_stack: Option<bool>,
max_packet_size: usize, max_pkt_size: usize,
next_assoc_id: AtomicU16,
udp_sessions: Mutex<HashMap<u16, Arc<AssociatedUdpSocket>>>,
} }
impl Server { impl Server {
@ -68,7 +67,9 @@ impl Server {
inner: Socks5Server::new(socket, auth), inner: Socks5Server::new(socket, auth),
addr: cfg.server, addr: cfg.server,
dual_stack: cfg.dual_stack, dual_stack: cfg.dual_stack,
max_packet_size: cfg.max_packet_size, max_pkt_size: cfg.max_packet_size,
next_assoc_id: AtomicU16::new(0),
udp_sessions: Mutex::new(HashMap::new()),
}; };
SERVER SERVER
@ -129,7 +130,7 @@ impl Server {
let socket = AssociatedUdpSocket::from(( let socket = AssociatedUdpSocket::from((
UdpSocket::from_std(StdUdpSocket::from(socket))?, UdpSocket::from_std(StdUdpSocket::from(socket))?,
SERVER.get().unwrap().max_packet_size, SERVER.get().unwrap().max_pkt_size,
)); ));
let addr = socket.local_addr()?; let addr = socket.local_addr()?;
@ -206,8 +207,19 @@ impl Server {
mut assoc: Associate<associate::Ready>, mut assoc: Associate<associate::Ready>,
assoc_socket: Arc<AssociatedUdpSocket>, assoc_socket: Arc<AssociatedUdpSocket>,
) -> Result<(), Error> { ) -> Result<(), Error> {
let assoc_id = NEXT_ASSOCIATE_ID.fetch_add(1, Ordering::AcqRel); let assoc_id = SERVER
UDP_SESSIONS.lock().insert(assoc_id, assoc_socket.clone()); .get()
.unwrap()
.next_assoc_id
.fetch_add(1, Ordering::AcqRel);
SERVER
.get()
.unwrap()
.udp_sessions
.lock()
.insert(assoc_id, assoc_socket.clone());
let mut connected = None; let mut connected = None;
async fn accept_pkt( async fn accept_pkt(
@ -257,7 +269,7 @@ impl Server {
}; };
let _ = assoc.shutdown().await; let _ = assoc.shutdown().await;
UDP_SESSIONS.lock().remove(&assoc_id); SERVER.get().unwrap().udp_sessions.lock().remove(&assoc_id);
match TuicConnection::get().await { match TuicConnection::get().await {
Ok(conn) => match conn.dissociate(assoc_id).await { Ok(conn) => match conn.dissociate(assoc_id).await {
@ -271,7 +283,7 @@ impl Server {
} }
pub async fn recv_pkt(pkt: Bytes, addr: Address, assoc_id: u16) -> Result<(), Error> { pub async fn recv_pkt(pkt: Bytes, addr: Address, assoc_id: u16) -> Result<(), Error> {
let sessions = UDP_SESSIONS.lock(); let sessions = SERVER.get().unwrap().udp_sessions.lock();
let Some(assoc_socket) = sessions.get(&assoc_id) else { unreachable!() }; let Some(assoc_socket) = sessions.get(&assoc_id) else { unreachable!() };
assoc_socket.send(pkt, 0, addr).await?; assoc_socket.send(pkt, 0, addr).await?;
Ok(()) Ok(())