From b2cae01bf891de991bfbdaa0c33c383e2e02d7c0 Mon Sep 17 00:00:00 2001
From: restitux <restitux@ohea.xyz>
Date: Tue, 21 May 2024 00:22:41 -0600
Subject: [PATCH] Add debug logging and global connection map

---
 Cargo.lock  |  1 +
 Cargo.toml  |  1 +
 src/main.rs | 34 +++++++++++++++++++++++++++++-----
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index af6f605..e50243d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -381,6 +381,7 @@ name = "mumble-webtransport-proxy"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "lazy_static",
  "tokio",
  "tracing",
  "tracing-subscriber",
diff --git a/Cargo.toml b/Cargo.toml
index 255ffb0..93f151c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021"
 
 [dependencies]
 anyhow = "1.0.86"
+lazy_static = "1.4.0"
 tokio = { version = "1.37.0", features = ["full"] }
 tracing = { version = "0.1.40", features = ["async-await"] }
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
diff --git a/src/main.rs b/src/main.rs
index 34a3bc5..80d7155 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,20 +9,29 @@ use tracing::Instrument;
 use tracing_subscriber::filter::LevelFilter;
 use tracing_subscriber::EnvFilter;
 use wtransport::endpoint::IncomingSession;
-use wtransport::Endpoint;
+use wtransport::{connection, Endpoint};
 use wtransport::Identity;
 use wtransport::ServerConfig;
 use tokio::io::AsyncReadExt;
 
+use lazy_static::lazy_static;
+use std::collections::HashMap;
+use std::sync::Mutex;
+
+type GlobalMap = Mutex<HashMap<usize, wtransport::Connection>>;
+
+lazy_static! {
+    static ref DATA_MAP: GlobalMap = Mutex::new(HashMap::new());
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     init_logging();
 
     let config = ServerConfig::builder()
         .with_bind_default(4433)
-        //.with_identity(&Identity::self_signed(["localhost"]).unwrap())
         .with_identity(&Identity::load_pemfiles("cert.pem", "key.pem").await?)
-        .keep_alive_interval(Some(Duration::from_secs(3)))
+        .keep_alive_interval(Some(Duration::from_secs(20)))
         .build();
 
     let server = Endpoint::server(config)?;
@@ -38,8 +47,10 @@ async fn main() -> Result<()> {
 }
 
 async fn handle_connection(incoming_session: IncomingSession, id: usize) {
-    let _result = handle_connection_impl(incoming_session, id).await;
-    //error!("{:?}", result);
+    // Wrapper to handle connection establishment failures
+    if let Err(e) = handle_connection_impl(incoming_session, id).await {
+        error!("{:?}", e);
+    }
 }
 
 async fn handle_connection_impl(incoming_session: IncomingSession, id: usize) -> Result<()> {
@@ -58,11 +69,22 @@ async fn handle_connection_impl(incoming_session: IncomingSession, id: usize) ->
     let connection = session_request.accept().await?;
     let stream = connection.accept_bi().await?;
 
+    info!("Connecting to corresponding Mumble server...");
+
     let server_tcp = TcpStream::connect("127.0.0.1:64738").await?.into_split();
 
+    info!("Connected to Mumble Server!");
+
+    // Store connection in global map to prevent it getting dropped
+    DATA_MAP.lock().unwrap().insert(id, connection);
+
+    info!("Spawing jobs...");
+    // Spawn tasks to handle transmitting data between the WebTransport client and Mumble TCP Server
     tokio::spawn(handle_client_to_server(stream.1, server_tcp.1).instrument(info_span!("Handler", "Client to server")));
     tokio::spawn(handle_server_to_client(stream.0, server_tcp.0).instrument(info_span!("Handler", "Server to client")));
 
+    info!("Spawned jobs.");
+
     Ok(())
 }
 
@@ -81,11 +103,13 @@ async fn client_to_server_loop(
 ) -> Result<()> {
     let mut buffer = vec![0; 65536].into_boxed_slice();
     loop {
+        info!("Reading Data");
         let _bytes_read = match client_stream.read(&mut buffer).await? {
             Some(bytes_read) => bytes_read,
             None => continue,
         };
 
+        info!("Writing data");
         server_stream.try_write(&buffer)?;
     }
 }