1
0

add auth timeout & gc on server

This commit is contained in:
EAimTY 2023-02-04 16:29:02 +09:00
parent 7781f5c62a
commit 8cf357012d

View File

@ -15,6 +15,7 @@ use std::{
Arc, Arc,
}, },
task::{Context, Poll, Waker}, task::{Context, Poll, Waker},
time::Duration,
}; };
use tokio::{ use tokio::{
io::{self, AsyncWriteExt}, io::{self, AsyncWriteExt},
@ -23,6 +24,7 @@ use tokio::{
oneshot::{self, Receiver, Sender}, oneshot::{self, Receiver, Sender},
Mutex as AsyncMutex, Mutex as AsyncMutex,
}, },
time,
}; };
use tokio_util::compat::FuturesAsyncReadCompatExt; use tokio_util::compat::FuturesAsyncReadCompatExt;
use tuic::Address; use tuic::Address;
@ -35,6 +37,9 @@ pub struct Server {
token: Arc<[u8]>, token: Arc<[u8]>,
udp_relay_ipv6: bool, udp_relay_ipv6: bool,
zero_rtt_handshake: bool, zero_rtt_handshake: bool,
auth_timeout: Duration,
gc_interval: Duration,
gc_lifetime: Duration,
} }
impl Server { impl Server {
@ -45,11 +50,15 @@ impl Server {
pub async fn start(&self) { pub async fn start(&self) {
loop { loop {
let conn = self.ep.accept().await.unwrap(); let conn = self.ep.accept().await.unwrap();
tokio::spawn(Connection::init(
tokio::spawn(Connection::new(
conn, conn,
self.token.clone(), self.token.clone(),
self.udp_relay_ipv6, self.udp_relay_ipv6,
self.zero_rtt_handshake, self.zero_rtt_handshake,
self.auth_timeout,
self.gc_interval,
self.gc_lifetime,
)); ));
} }
} }
@ -71,14 +80,21 @@ struct Connection {
} }
impl Connection { impl Connection {
pub async fn init( async fn new(
conn: Connecting, conn: Connecting,
token: Arc<[u8]>, token: Arc<[u8]>,
udp_relay_ipv6: bool, udp_relay_ipv6: bool,
zero_rtt_handshake: bool, zero_rtt_handshake: bool,
auth_timeout: Duration,
gc_interval: Duration,
gc_lifetime: Duration,
) { ) {
match Self::handshake(conn, token, udp_relay_ipv6, zero_rtt_handshake).await { match Self::init(conn, token, udp_relay_ipv6, zero_rtt_handshake).await {
Ok(conn) => loop { Ok(conn) => {
tokio::spawn(conn.clone().handle_auth_timeout(auth_timeout));
tokio::spawn(conn.clone().collect_garbage(gc_interval, gc_lifetime));
loop {
if conn.is_closed() { if conn.is_closed() {
break; break;
} }
@ -87,12 +103,13 @@ impl Connection {
Ok(()) => {} Ok(()) => {}
Err(err) => eprintln!("{err}"), Err(err) => eprintln!("{err}"),
} }
}, }
}
Err(err) => eprintln!("{err}"), Err(err) => eprintln!("{err}"),
} }
} }
async fn handshake( async fn init(
conn: Connecting, conn: Connecting,
token: Arc<[u8]>, token: Arc<[u8]>,
udp_relay_ipv6: bool, udp_relay_ipv6: bool,
@ -352,6 +369,26 @@ impl Connection {
Ok(()) Ok(())
} }
async fn handle_auth_timeout(self, timeout: Duration) {
time::sleep(timeout).await;
if !self.is_authed() {
self.close();
}
}
async fn collect_garbage(self, gc_interval: Duration, gc_lifetime: Duration) {
loop {
time::sleep(gc_interval).await;
if self.is_closed() {
break;
}
self.model.collect_garbage(gc_lifetime);
}
}
fn set_authed(&self) { fn set_authed(&self) {
self.is_authed.set_authed(); self.is_authed.set_authed();
} }