refactor: migrate net operation to KCP

gameserver/src/net/gateway.rs

@@ -1,28 +1,143 @@
+use std::{
+    collections::HashMap,
+    net::SocketAddr,
+    sync::{
+        atomic::{AtomicU32, Ordering},
+        Arc,
+    },
+};
+
 use anyhow::Result;
-use tokio::net::TcpListener;
-use tracing::{info_span, Instrument};
+use rand::RngCore;
 
-use crate::{log_error, net::PlayerSession};
+use crate::net::PlayerSession;
 
-pub async fn listen(host: &str, port: u16) -> Result<()> {
-    let listener = TcpListener::bind(format!("{host}:{port}")).await?;
-    tracing::info!("Listening at {host}:{port}");
+use tokio::{
+    net::UdpSocket,
+    sync::{Mutex, RwLock},
+};
 
-    loop {
-        let Ok((client_socket, client_addr)) = listener.accept().await else {
-            continue;
+use crate::net::packet::NetOperation;
+
+const MAX_PACKET_SIZE: usize = 1200;
+
+pub struct Gateway {
+    socket: Arc<UdpSocket>,
+    id_counter: AtomicU32,
+    sessions: Mutex<HashMap<u32, Arc<RwLock<PlayerSession>>>>,
+}
+
+impl Gateway {
+    pub async fn new(host: &str, port: u16) -> Result<Self> {
+        let socket = Arc::new(UdpSocket::bind(format!("{host}:{port}")).await?);
+
+        Ok(Self {
+            socket,
+            id_counter: AtomicU32::new(0),
+            sessions: Mutex::new(HashMap::new()),
+        })
+    }
+
+    pub async fn listen(&mut self) -> Result<()> {
+        tracing::info!(
+            "KCP Gateway is listening at {}",
+            self.socket.local_addr().unwrap()
+        );
+
+        let mut buf = [0; MAX_PACKET_SIZE];
+        loop {
+            let Ok((len, addr)) = self.socket.recv_from(&mut buf).await else {
+                continue;
+            };
+
+            match len {
+                20 => self.process_net_operation(buf[..len].into(), addr).await?,
+                28.. => self.process_kcp_payload(buf[..len].into(), addr).await,
+                _ => {
+                    tracing::warn!("unk data len {len}")
+                }
+            }
+        }
+    }
+
+    async fn process_net_operation(&mut self, op: NetOperation, addr: SocketAddr) -> Result<()> {
+        match (op.head, op.tail) {
+            (0xFF, 0xFFFFFFFF) => self.establish_kcp_session(op.data, addr).await?,
+            (0x194, 0x19419494) => self.drop_kcp_session(op.param1, op.param2, addr).await,
+            _ => tracing::warn!("Unknown magic pair received {:X}-{:X}", op.head, op.tail),
+        }
+
+        Ok(())
+    }
+
+    async fn establish_kcp_session(&mut self, data: u32, addr: SocketAddr) -> Result<()> {
+        tracing::info!("New connection from {addr}");
+        let (conv_id, session_token) = self.next_conv_pair();
+
+        self.sessions.lock().await.insert(
+            conv_id,
+            Arc::new(RwLock::new(PlayerSession::new(
+                self.socket.clone(),
+                addr,
+                conv_id,
+                session_token,
+            ))),
+        );
+
+        self.socket
+            .send_to(
+                &Vec::from(NetOperation {
+                    head: 0x145,
+                    param1: conv_id,
+                    param2: session_token,
+                    data,
+                    tail: 0x14514545,
+                }),
+                addr,
+            )
+            .await?;
+
+        Ok(())
+    }
+
+    async fn drop_kcp_session(&mut self, conv_id: u32, token: u32, addr: SocketAddr) {
+        tracing::info!("drop_kcp_session {conv_id} {token}");
+        let Some(session) = self.sessions.lock().await.get(&conv_id).cloned() else {
+            tracing::warn!("drop_kcp_session failed, no session with conv_id {conv_id} was found");
+            return;
         };
 
-        let mut session = PlayerSession::new(client_socket);
-        tokio::spawn(
-            async move {
-                log_error!(
-                    "Session from {client_addr} disconnected",
-                    format!("An error occurred while processing session ({client_addr})"),
-                    Box::pin(session.run()).await
-                );
+        if session.read().await.token == token {
+            self.sessions.lock().await.remove(&conv_id);
+            tracing::info!("Client from {addr} disconnected");
+        }
+    }
+
+    async fn process_kcp_payload(&mut self, data: Box<[u8]>, addr: SocketAddr) {
+        let conv_id = mhy_kcp::get_conv(&data);
+
+        let Some(session) = self
+            .sessions
+            .lock()
+            .await
+            .get_mut(&conv_id)
+            .map(|s| s.clone())
+        else {
+            tracing::warn!("Session with conv_id {conv_id} not found!");
+            return;
+        };
+
+        tokio::spawn(async move {
+            if let Err(err) = Box::pin(session.write().await.consume(&data)).await {
+                tracing::error!("An error occurred while processing session ({addr}): {err}");
             }
-            .instrument(info_span!("session", addr = %client_addr)),
-        );
+        });
+    }
+
+    fn next_conv_pair(&mut self) -> (u32, u32) {
+        (
+            self.id_counter.fetch_add(1, Ordering::SeqCst) + 1,
+            rand::thread_rng().next_u32(),
+        )
     }
 }