257 lines
9.3 KiB
Rust
257 lines
9.3 KiB
Rust
|
use std::{
|
||
|
|
collections::{hash_map::Entry, HashMap},
|
||
|
|
net::SocketAddr,
|
||
|
|
sync::{
|
||
|
|
atomic::{AtomicBool, Ordering},
|
||
|
|
Arc,
|
||
|
|
},
|
||
|
|
time::Duration,
|
||
|
|
};
|
||
|
|
|
||
|
|
use byte_string::ByteStr;
|
||
|
|
use kcp::KcpResult;
|
||
|
|
use log::{error, trace};
|
||
|
|
use tokio::{
|
||
|
|
net::UdpSocket,
|
||
|
|
sync::{mpsc, Mutex},
|
||
|
|
time::{self, Instant},
|
||
|
|
};
|
||
|
|
|
||
|
|
use crate::plugins::kcp::{skcp::KcpSocket, KcpConfig};
|
||
|
|
|
||
|
|
pub struct KcpSession {
|
||
|
|
socket: Mutex<KcpSocket>,
|
||
|
|
closed: AtomicBool,
|
||
|
|
session_expire: Duration,
|
||
|
|
session_close_notifier: Option<mpsc::Sender<u32>>,
|
||
|
|
input_tx: mpsc::Sender<Vec<u8>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl KcpSession {
|
||
|
|
fn new(
|
||
|
|
socket: KcpSocket,
|
||
|
|
session_expire: Duration,
|
||
|
|
session_close_notifier: Option<mpsc::Sender<u32>>,
|
||
|
|
input_tx: mpsc::Sender<Vec<u8>>,
|
||
|
|
) -> KcpSession {
|
||
|
|
KcpSession {
|
||
|
|
socket: Mutex::new(socket),
|
||
|
|
closed: AtomicBool::new(false),
|
||
|
|
session_expire,
|
||
|
|
session_close_notifier,
|
||
|
|
input_tx,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn new_shared(
|
||
|
|
socket: KcpSocket,
|
||
|
|
session_expire: Duration,
|
||
|
|
session_close_notifier: Option<mpsc::Sender<u32>>,
|
||
|
|
) -> Arc<KcpSession> {
|
||
|
|
let is_client = session_close_notifier.is_none();
|
||
|
|
|
||
|
|
let (input_tx, mut input_rx) = mpsc::channel(64);
|
||
|
|
|
||
|
|
let udp_socket = socket.udp_socket().clone();
|
||
|
|
|
||
|
|
let session = Arc::new(KcpSession::new(
|
||
|
|
socket,
|
||
|
|
session_expire,
|
||
|
|
session_close_notifier,
|
||
|
|
input_tx,
|
||
|
|
));
|
||
|
|
|
||
|
|
{
|
||
|
|
let session = session.clone();
|
||
|
|
tokio::spawn(async move {
|
||
|
|
let mut input_buffer = [0u8; 65536];
|
||
|
|
let update_timer = time::sleep(Duration::from_millis(10));
|
||
|
|
tokio::pin!(update_timer);
|
||
|
|
|
||
|
|
loop {
|
||
|
|
tokio::select! {
|
||
|
|
// recv() then input()
|
||
|
|
// Drives the KCP machine forward
|
||
|
|
recv_result = udp_socket.recv(&mut input_buffer), if is_client => {
|
||
|
|
match recv_result {
|
||
|
|
Err(err) => {
|
||
|
|
error!("[SESSION] UDP recv failed, error: {}", err);
|
||
|
|
}
|
||
|
|
Ok(n) => {
|
||
|
|
let input_buffer = &input_buffer[..n];
|
||
|
|
trace!("[SESSION] UDP recv {} bytes, going to input {:?}", n, ByteStr::new(input_buffer));
|
||
|
|
|
||
|
|
let mut socket = session.socket.lock().await;
|
||
|
|
|
||
|
|
match socket.input(input_buffer) {
|
||
|
|
Ok(true) => {
|
||
|
|
trace!("[SESSION] UDP input {} bytes and waked sender/receiver", n);
|
||
|
|
}
|
||
|
|
Ok(false) => {}
|
||
|
|
Err(err) => {
|
||
|
|
error!("[SESSION] UDP input {} bytes error: {}, input buffer {:?}", n, err, ByteStr::new(input_buffer));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// bytes received from listener socket
|
||
|
|
input_opt = input_rx.recv() => {
|
||
|
|
if let Some(input_buffer) = input_opt {
|
||
|
|
let mut socket = session.socket.lock().await;
|
||
|
|
match socket.input(&input_buffer) {
|
||
|
|
Ok(..) => {
|
||
|
|
trace!("[SESSION] UDP input {} bytes from channel {:?}", input_buffer.len(), ByteStr::new(&input_buffer));
|
||
|
|
}
|
||
|
|
Err(err) => {
|
||
|
|
error!("[SESSION] UDP input {} bytes from channel failed, error: {}, input buffer {:?}",
|
||
|
|
input_buffer.len(), err, ByteStr::new(&input_buffer));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Call update() in period
|
||
|
|
_ = &mut update_timer => {
|
||
|
|
let mut socket = session.socket.lock().await;
|
||
|
|
|
||
|
|
let is_closed = session.closed.load(Ordering::Acquire);
|
||
|
|
if is_closed && socket.can_close() {
|
||
|
|
trace!("[SESSION] KCP session closed");
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
|
||
|
|
// server socket expires
|
||
|
|
if !is_client {
|
||
|
|
// If this is a server stream, close it automatically after a period of time
|
||
|
|
let last_update_time = socket.last_update_time();
|
||
|
|
let elapsed = last_update_time.elapsed();
|
||
|
|
|
||
|
|
if elapsed > session.session_expire {
|
||
|
|
if elapsed > session.session_expire * 2 {
|
||
|
|
// Force close. Client may have already gone.
|
||
|
|
trace!(
|
||
|
|
"[SESSION] force close inactive session, conv: {}, last_update: {}s ago",
|
||
|
|
socket.conv(),
|
||
|
|
elapsed.as_secs()
|
||
|
|
);
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
|
||
|
|
if !is_closed {
|
||
|
|
trace!(
|
||
|
|
"[SESSION] closing inactive session, conv: {}, last_update: {}s ago",
|
||
|
|
socket.conv(),
|
||
|
|
elapsed.as_secs()
|
||
|
|
);
|
||
|
|
session.closed.store(true, Ordering::Release);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
match socket.update() {
|
||
|
|
Ok(next_next) => {
|
||
|
|
update_timer.as_mut().reset(Instant::from_std(next_next));
|
||
|
|
}
|
||
|
|
Err(err) => {
|
||
|
|
error!("[SESSION] KCP update failed, error: {}", err);
|
||
|
|
update_timer.as_mut().reset(Instant::now() + Duration::from_millis(10));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
{
|
||
|
|
// Close the socket.
|
||
|
|
// Wake all pending tasks and let all send/recv return EOF
|
||
|
|
|
||
|
|
let mut socket = session.socket.lock().await;
|
||
|
|
socket.close();
|
||
|
|
}
|
||
|
|
|
||
|
|
if let Some(ref notifier) = session.session_close_notifier {
|
||
|
|
let socket = session.socket.lock().await;
|
||
|
|
let _ = notifier.send(socket.conv()).await;
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
session
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn kcp_socket(&self) -> &Mutex<KcpSocket> {
|
||
|
|
&self.socket
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn close(&self) {
|
||
|
|
self.closed.store(true, Ordering::Release);
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn input(&self, buf: &[u8]) {
|
||
|
|
self.input_tx
|
||
|
|
.send(buf.to_owned())
|
||
|
|
.await
|
||
|
|
.expect("input channel closed")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub struct KcpSessionManager {
|
||
|
|
sessions: HashMap<u32, Arc<KcpSession>>,
|
||
|
|
next_free_conv: u32,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl KcpSessionManager {
|
||
|
|
pub fn new() -> KcpSessionManager {
|
||
|
|
KcpSessionManager {
|
||
|
|
sessions: HashMap::new(),
|
||
|
|
next_free_conv: 0,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn close_conv(&mut self, conv: u32) {
|
||
|
|
self.sessions.remove(&conv);
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn alloc_conv(&mut self) -> u32 {
|
||
|
|
loop {
|
||
|
|
let (mut c, _) = self.next_free_conv.overflowing_add(1);
|
||
|
|
if c == 0 {
|
||
|
|
let (nc, _) = c.overflowing_add(1);
|
||
|
|
c = nc;
|
||
|
|
}
|
||
|
|
self.next_free_conv = c;
|
||
|
|
|
||
|
|
if self.sessions.get(&self.next_free_conv).is_none() {
|
||
|
|
let conv = self.next_free_conv;
|
||
|
|
return conv;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn get_or_create(
|
||
|
|
&mut self,
|
||
|
|
config: &KcpConfig,
|
||
|
|
conv: u32,
|
||
|
|
udp: &Arc<UdpSocket>,
|
||
|
|
peer_addr: SocketAddr,
|
||
|
|
session_close_notifier: &mpsc::Sender<u32>,
|
||
|
|
) -> KcpResult<(Arc<KcpSession>, bool)> {
|
||
|
|
match self.sessions.entry(conv) {
|
||
|
|
Entry::Occupied(occ) => Ok((occ.get().clone(), false)),
|
||
|
|
Entry::Vacant(vac) => {
|
||
|
|
let socket = KcpSocket::new(config, conv, udp.clone(), peer_addr, config.stream)?;
|
||
|
|
let session = KcpSession::new_shared(
|
||
|
|
socket,
|
||
|
|
config.session_expire,
|
||
|
|
Some(session_close_notifier.clone()),
|
||
|
|
);
|
||
|
|
trace!("created session for conv: {}, peer: {}", conv, peer_addr);
|
||
|
|
vac.insert(session.clone());
|
||
|
|
Ok((session, true))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|