121 lines
3.5 KiB
Rust
121 lines
3.5 KiB
Rust
use std::collections::HashMap;
|
|
use std::io;
|
|
use std::net::SocketAddr;
|
|
use std::sync::LazyLock;
|
|
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
use tokio::sync::{Mutex, mpsc};
|
|
use tokio::task::JoinHandle;
|
|
|
|
use super::SERVER_ADDR;
|
|
|
|
pub(crate) static TCP_SERVER: LazyLock<Mutex<TcpServer>> =
|
|
LazyLock::new(|| Mutex::new(TcpServer::new()));
|
|
|
|
pub(crate) struct TcpServer {
|
|
is_running: bool,
|
|
clients: HashMap<SocketAddr, JoinHandle<()>>,
|
|
shutdown_tx: Option<mpsc::Sender<()>>,
|
|
}
|
|
|
|
impl TcpServer {
|
|
fn new() -> Self {
|
|
Self {
|
|
is_running: false,
|
|
clients: HashMap::new(),
|
|
shutdown_tx: None,
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn start(&mut self) {
|
|
if self.is_running {
|
|
log::warn!("TCP server is already running");
|
|
return;
|
|
}
|
|
|
|
match TcpListener::bind(SERVER_ADDR).await {
|
|
Ok(listener) => {
|
|
let (shutdown_tx, shutdown_rx) = mpsc::channel(1);
|
|
self.is_running = true;
|
|
self.shutdown_tx = Some(shutdown_tx);
|
|
|
|
tokio::spawn(async move {
|
|
Self::listen_to_clients(listener, shutdown_rx).await;
|
|
});
|
|
|
|
log::info!("TCP Server started on {}", SERVER_ADDR);
|
|
}
|
|
Err(e) => log::error!("Failed to bind to address: {e}"),
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn stop(&mut self) {
|
|
if !self.is_running {
|
|
return;
|
|
}
|
|
|
|
self.is_running = false;
|
|
|
|
if let Some(shutdown_tx) = self.shutdown_tx.take() {
|
|
_ = shutdown_tx.send(()).await;
|
|
}
|
|
|
|
for (addr, connection) in self.clients.drain() {
|
|
log::info!("Closing connection to {}", addr);
|
|
connection.abort();
|
|
}
|
|
}
|
|
|
|
async fn listen_to_clients(listener: TcpListener, mut shutdown_rx: mpsc::Receiver<()>) {
|
|
loop {
|
|
tokio::select! {
|
|
result = listener.accept() => {
|
|
match result {
|
|
Ok((socket, addr)) => {
|
|
log::info!("New client connected: {addr}");
|
|
|
|
let task_handle = tokio::spawn(async move {
|
|
if let Err(e) = Self::handle_client(socket, addr).await {
|
|
log::error!("Client {addr} error: {e}");
|
|
}
|
|
log::info!("Client {addr} disconnected");
|
|
});
|
|
|
|
let mut server = TCP_SERVER.lock().await;
|
|
server.clients.insert(addr, task_handle);
|
|
}
|
|
Err(e) => log::error!("Couldn't get client: {e}"),
|
|
}
|
|
}
|
|
|
|
_ = shutdown_rx.recv() => {
|
|
log::info!("TCP Server shutting down");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_client(mut socket: TcpStream, addr: SocketAddr) -> io::Result<()> {
|
|
let mut buffer = [0; 1024];
|
|
|
|
loop {
|
|
let len = socket.read(&mut buffer).await?;
|
|
if len == 0 {
|
|
break;
|
|
}
|
|
|
|
log::debug!("Received {} bytes from {}", len, addr);
|
|
|
|
// TODO: Deserialize data
|
|
socket.write_all(&buffer[..len]).await?;
|
|
}
|
|
|
|
let mut server = TCP_SERVER.lock().await;
|
|
server.clients.remove(&addr);
|
|
|
|
Ok(())
|
|
}
|
|
}
|