[Fix] Now using Axum instead of Hyper and Http2 instead of 1

Cargo.lock

@@ -390,6 +390,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8"
 dependencies = [
  "axum-core",
+ "base64 0.22.1",
  "bytes",
  "form_urlencoded",
  "futures-util",
@@ -408,8 +409,10 @@ dependencies = [
  "serde_json",
  "serde_path_to_error",
  "serde_urlencoded",
+ "sha1",
  "sync_wrapper",
  "tokio",
+ "tokio-tungstenite",
  "tower",
  "tower-layer",
  "tower-service",

Cargo.toml

@@ -39,7 +39,7 @@ async-tungstenite = { version = "0.32.0", features = [
     "verbose-logging",
     "webpki-roots",
 ] }
-axum = "0.8.8"
+axum = { version = "0.8.8", features = [ "ws" ] }
 base64 = "0.22.1"
 bytes = "1.11.1"
 color-eyre = "0.6.5"

src/server/api.rs

@@ -10,14 +10,12 @@ use crate::{
 };
 use axum::http::HeaderValue;
 use base64::Engine as _;
-use http_body_util::{Full, StreamBody};
+use http_body_util::Full;
-use hyper::body::{Body, Bytes, Frame};
+use hyper::body::Bytes;
 use hyper::header::{CONTENT_DISPOSITION, CONTENT_LENGTH, CONTENT_TYPE};
 use hyper::{HeaderMap, Response as HttpResponse, StatusCode};
 use json::JsonValue;
 use json::number::Number;
-use tokio::fs::File;
-use tokio_util::io::ReaderStream;
 
 pub async fn handle(
     path: &str,

src/server/omikron_connection.rs

@@ -6,29 +6,26 @@ use crate::sql::sql::{self, get_by_user_id, get_by_username, get_iota_by_id, get
 use crate::sql::user_online_tracker::{self};
 use crate::util::crypto_helper::encrypt;
 use crate::util::logger::PrintType;
-use crate::{get_private_key, get_public_key, log_in, log_out};
+use crate::{get_private_key, get_public_key, log_err, log_in, log_out};
+use axum::extract::ws::WebSocket;
+use axum::extract::ws::{Message, Utf8Bytes};
 use base64::{Engine as _, engine::general_purpose::STANDARD};
 use dashmap::DashMap;
-use futures::SinkExt;
 use futures::stream::SplitSink;
 use futures::stream::SplitStream;
-use hyper::upgrade::Upgraded;
+use futures_util::SinkExt;
-use hyper_util::rt::TokioIo;
 use json::JsonValue;
 use json::number::Number;
 use rand::Rng;
 use rand::distributions::Alphanumeric;
 use std::sync::Arc;
 use tokio::sync::RwLock;
-use tokio_tungstenite::WebSocketStream;
-use tungstenite::Message;
-use tungstenite::Utf8Bytes;
 use uuid::Uuid;
 use x448::PublicKey;
 
 pub struct OmikronConnection {
-    pub sender: Arc<RwLock<SplitSink<WebSocketStream<TokioIo<Upgraded>>, Message>>>,
+    pub sender: Arc<RwLock<SplitSink<WebSocket, Message>>>,
-    pub receiver: Arc<RwLock<SplitStream<WebSocketStream<TokioIo<Upgraded>>>>>,
+    pub receiver: Arc<RwLock<SplitStream<WebSocket>>>,
     pub omikron_id: Arc<RwLock<i64>>,
     pub pub_key: Arc<RwLock<Option<Vec<u8>>>>,
     identified: Arc<RwLock<bool>>,
@@ -43,8 +40,8 @@ pub struct OmikronConnection {
 
 impl OmikronConnection {
     pub fn new(
-        sender: SplitSink<WebSocketStream<TokioIo<Upgraded>>, Message>,
+        sender: SplitSink<WebSocket, Message>,
-        receiver: SplitStream<WebSocketStream<TokioIo<Upgraded>>>,
+        receiver: SplitStream<WebSocket>,
     ) -> Arc<Self> {
         Arc::new(Self {
             sender: Arc::new(RwLock::new(sender)),
@@ -66,10 +63,17 @@ impl OmikronConnection {
                 *self.omikron_id.read().await,
                 PrintType::Omikron,
                 "{}",
-                message_text
+                cv.to_json().to_string()
+            );
+        }
+        if let Err(e) = sender.send(message_text).await {
+            log_err!(
+                *self.omikron_id.read().await,
+                PrintType::Omikron,
+                "WebSocket send error: {}",
+                e
             );
         }
-        let _ = sender.send(message_text).await;
     }
     pub async fn get_omikron_id(&self) -> i64 {
         *self.omikron_id.read().await
@@ -77,9 +81,6 @@ impl OmikronConnection {
     pub async fn is_identified(&self) -> bool {
         *self.identified.read().await && *self.challenged.read().await
     }
-    pub async fn get_public_key(&self) -> PublicKey {
-        PublicKey::from_bytes(self.pub_key.read().await.as_ref().unwrap()).unwrap()
-    }
 
     pub async fn handle_message(self: Arc<Self>, message: String) {
         let cv = CommunicationValue::from_json(&message);
@@ -619,7 +620,6 @@ impl OmikronConnection {
 
             if let Some(public_key) = cv.get_data(DataTypes::public_key).and_then(|v| v.as_str()) {
                 if let Some(iota_id) = iota_id_opt {
-                    // Existing logic to update iota
                     match sql::register_complete_iota(iota_id, public_key.to_string()).await {
                         Ok(_) => {
                             let response = CommunicationValue::new(CommunicationType::success)
@@ -792,7 +792,7 @@ impl OmikronConnection {
             ) {
                 match sql::get_by_user_id(user_id).await {
                     Ok(user) => {
-                        let current_token = user.11; // token is the 12th element (index 11)
+                        let current_token = user.11;
                         if current_token == reset_token {
                             let mut success = true;
                             let mut error_message = String::new();
@@ -824,7 +824,6 @@ impl OmikronConnection {
                         } else {
                             self.send_error_response(
                                 &cv.get_id(),
-                                // Using this for invalid token
                                 CommunicationType::error_invalid_challenge,
                             )
                             .await;

src/server/omikron_manager.rs

@@ -24,11 +24,3 @@ pub async fn get_random_omikron() -> Result<Arc<OmikronConnection>, ()> {
         return Err(());
     }
 }
-
-pub async fn get_omikron(omikron_id: i64) -> Option<Arc<OmikronConnection>> {
-    if let Some(omikron) = OMIKRON_CONNECTIONS.get(&omikron_id) {
-        Some(omikron.clone())
-    } else {
-        None
-    }
-}

src/server/server.rs

@@ -1,486 +1,106 @@
-use base64::Engine;
+use axum::{
-use base64::engine::general_purpose::STANDARD;
+    Router,
-use bytes::Bytes;
+    body::Body,
-use futures::StreamExt;
+    extract::{ConnectInfo, OriginalUri, Path, ws::WebSocketUpgrade},
-use futures_util::TryFutureExt;
+    response::{IntoResponse, Redirect},
-use http_body_util::BodyExt;
+    routing::get,
-use http_body_util::Full;
-use hyper::server::conn::http2;
-use hyper::{Method, StatusCode};
-use hyper::{
-    Request as HttpRequest, Response as HttpResponse, body::Incoming, server::conn::http1, upgrade,
 };
-use hyper_util::rt::TokioExecutor;
+
-use hyper_util::rt::tokio::TokioIo;
-use hyper_util::service::TowerToHyperService;
 use pnet::datalink::NetworkInterface;
-use rustls::ServerConfig;
+use std::net::SocketAddr;
-use rustls::pki_types::{CertificateDer, PrivateKeyDer};
+use std::time::Duration;
-use sha1::{Digest, Sha1};
-use std::error::Error;
-use std::io::ErrorKind;
-use std::io::{self, BufReader};
-use std::net::{IpAddr, SocketAddr};
-use std::result::Result::Ok;
-use std::sync::Arc;
-use std::{future::Future, pin::Pin, time::Duration};
 use tokio::net::TcpListener;
-use tokio::sync::broadcast;
-use tokio_rustls::TlsAcceptor;
-use tower::Service;
 
 use crate::log;
 use crate::server::api;
 use crate::server::short_link::get_short_link;
 use crate::server::socket;
-use crate::util::file_util::load_file_buf;
+
-#[derive(Clone)]
+pub async fn start(port: u16) -> bool {
-struct HttpService {
+    let app = Router::new()
-    peer_addr: SocketAddr,
+        .route("/ws/omikron", get(ws_handler))
+        .route("/direct/{short}", get(direct_handler))
+        .fallback(fallback_handler);
+    run_http_server(port, app).await
 }
 
-impl Service<HttpRequest<Incoming>> for HttpService {
+async fn ws_handler(
-    type Response = HttpResponse<Full<Bytes>>;
+    ws: WebSocketUpgrade,
-    type Error = io::Error;
+    OriginalUri(uri): OriginalUri,
-    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+    ConnectInfo(_): ConnectInfo<SocketAddr>,
+) -> impl IntoResponse {
+    log!("Attempting WebSocket upgrade on {}", uri.path());
+    let path = uri.path().to_string();
 
-    fn poll_ready(
+    ws.on_upgrade(async move |socket| socket::handle(path, socket))
-        &mut self,
+}
-        _cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Result<(), Self::Error>> {
-        std::task::Poll::Ready(std::io::Result::Ok(()))
-    }
 
-    fn call(&mut self, req: HttpRequest<Incoming>) -> Self::Future {
+async fn direct_handler(Path(short): Path<String>) -> impl IntoResponse {
-        let peer_ip = self.peer_addr.ip();
+    match get_short_link(&short).await {
-
+        Ok(long) => Redirect::temporary(&long),
-        let (parts, body) = req.into_parts();
+        Err(_) => Redirect::temporary("https://tensamin.net"),
-
-        let method = parts.method.clone();
-        let path = parts.uri.path().to_string();
-        let headers = parts.headers.clone();
-
-        let fut = async move {
-            let is_websocket_upgrade = path.starts_with("/ws")
-                && method == Method::GET
-                && headers
-                    .get("connection")
-                    .map(|v| v.to_str().unwrap_or("").contains("Upgrade"))
-                    .unwrap_or(false)
-                && headers.get("upgrade").map(|v| v.to_str().unwrap_or("")) == Some("websocket");
-
-            if is_websocket_upgrade {
-                log!("Attempting WebSocket upgrade on /ws");
-
-                if let Some(sec_websocket_key) = headers.get("sec-websocket-key") {
-                    let sec_websocket_key = sec_websocket_key.to_str().unwrap_or("").to_string();
-                    let sec_websocket_accept = calculate_accept_key(&sec_websocket_key);
-
-                    let response = HttpResponse::builder()
-                        .status(StatusCode::SWITCHING_PROTOCOLS)
-                        .header("Upgrade", "websocket")
-                        .header("Connection", "Upgrade")
-                        .header("Sec-WebSocket-Accept", sec_websocket_accept)
-                        .body(Full::new(Bytes::from("")))
-                        .unwrap();
-                    let req_for_upgrade = HttpRequest::from_parts(parts, body);
-                    let upgrades = upgrade::on(req_for_upgrade);
-
-                    socket::handle(path, upgrades);
-                    Ok(response)
-                } else {
-                    log!("No Sec-WebSocket-Key found in request headers");
-                    let response = HttpResponse::builder()
-                        .status(StatusCode::BAD_REQUEST)
-                        .body(Full::new(Bytes::from("Missing Sec-WebSocket-Key")))
-                        .unwrap();
-                    Ok(response)
-                }
-            } else if path.starts_with("/api") {
-                let whole_body =
-                    tokio::time::timeout(Duration::from_secs(10), body.collect()).await??;
-                let bytes = whole_body.to_bytes();
-
-                let body_string: Option<String> = match String::from_utf8(bytes.to_vec()) {
-                    Ok(s) => Some(s),
-                    Err(_) => None,
-                };
-
-                Ok(api::handle(&path, headers.clone(), body_string).await)
-            } else if path.starts_with("/direct") {
-                let short = path.replace("/direct/", "");
-                if let Ok(long) = get_short_link(&short).await {
-                    let response = HttpResponse::builder()
-                        .status(StatusCode::FOUND)
-                        .header("Location", long)
-                        .body(Full::new(Bytes::from("")))
-                        .unwrap();
-                    Ok(response)
-                } else {
-                    let response = HttpResponse::builder()
-                        .status(StatusCode::NOT_FOUND)
-                        .body(Full::new(Bytes::from("Short link not found")))
-                        .unwrap();
-                    Ok(response)
-                }
-            } else {
-                let whole_body = match body.collect().await {
-                    Ok(collected) => collected,
-                    Err(e) => {
-                        log!("Error collecting body: {}", e);
-                        return Ok(HttpResponse::builder()
-                            .status(StatusCode::INTERNAL_SERVER_ERROR)
-                            .body(Full::new(Bytes::from(format!(
-                                "Failed to read body: {}",
-                                e
-                            ))))
-                            .unwrap());
-                    }
-                };
-                let bytes = whole_body.to_bytes();
-
-                let body_string: Option<String> = match String::from_utf8(bytes.to_vec()) {
-                    Ok(s) => Some(s),
-                    Err(_) => None,
-                };
-
-                Ok(api::handle(&path, headers, body_string).await)
-            }
-        };
-
-        Box::pin(fut.map_err(|err: color_eyre::eyre::ErrReport| {
-            io::Error::new(
-                io::ErrorKind::Other,
-                format!("Error in request handling: {}", err),
-            )
-        }))
     }
 }
 
-pub fn is_local_network(addr: IpAddr) -> bool {
+async fn fallback_handler(
-    match addr {
+    OriginalUri(uri): OriginalUri,
-        IpAddr::V4(v4) => {
+    headers: axum::http::HeaderMap,
-            let octets = v4.octets();
+    body: Body,
-            if octets[0] == 10 {
+) -> impl IntoResponse {
-                return true;
+    let path = uri.path().to_string();
-            }
-            if octets[0] == 172 && (16..=31).contains(&octets[1]) {
-                return true;
-            }
-            if octets[0] == 192 && octets[1] == 168 {
-                return true;
-            }
-            if octets[0] == 127 {
-                return true;
-            }
-            if octets[0] == 169 && octets[1] == 254 {
-                return true;
-            }
 
-            false
+    let whole_body = tokio::time::timeout(
-        }
+        Duration::from_secs(10),
+        axum::body::to_bytes(body, 1024 * 1024 * 10),
+    )
+    .await;
 
-        IpAddr::V6(v6) => {
+    let body_string = match whole_body {
-            let segments = v6.segments();
+        Ok(Ok(bytes)) => String::from_utf8(bytes.to_vec()).ok(),
-            if (segments[0] & 0xfe00) == 0xfc00 {
+        _ => None,
-                return true;
+    };
-            }
-            if (segments[0] & 0xffc0) == 0xfe80 {
-                return true;
-            }
-            if v6.is_loopback() {
-                return true;
-            }
 
-            false
+    api::handle(&path, headers, body_string).await
-        }
-    }
 }
 
-async fn run_http_server(port: u16) -> bool {
+async fn run_http_server(port: u16, app: Router) -> bool {
-    let mut ip = "0.0.0.0".to_string();
+    let ip = find_local_ip();
-    for iface in pnet::datalink::interfaces() {
+    let listener = match TcpListener::bind(format!("0.0.0.0:{}", port)).await {
-        let iface: NetworkInterface = iface;
+        Ok(l) => l,
-        if iface.ips.len() > 0 {
+        Err(e) => {
-            let ipsv = format!("{}", iface.ips[0]);
+            log!("Failed to bind to port {}: {:?}", port, e);
-            let ips: &str = ipsv.split('/').next().unwrap();
+            return false;
-            if format!("{}", ips).starts_with("10.") || format!("{}", ips).starts_with("192.") {
-                ip = ips.to_string();
-            }
         }
-    }
+    };
-    let listener = TcpListener::bind(format!("0.0.0.0:{}", port)).await;
+
-    if let Err(e) = listener {
-        log!("Failed to bind to port {}: {:?}", port, e);
-        return false;
-    }
-    let listener = listener.unwrap();
     log!(
         "Standard Server listening for HTTP and WS on {}:{}",
         ip,
         port
     );
 
-    // Create a broadcast channel for graceful shutdown signal
+    axum::serve(
-    let (shutdown_tx, _) = broadcast::channel::<()>(1);
+        listener,
-
+        app.into_make_service_with_connect_info::<SocketAddr>(),
-    tokio::spawn(async move {
+    )
-        loop {
+    .await
-            tokio::select! {
+    .map(|_| true)
-                // Monitor for shutdown signal
+    .unwrap_or_else(|e| {
-                _ = async {
+        log!("Server error: {}", e);
-                    loop {
+        false
-                        tokio::time::sleep(Duration::from_millis(100)).await;
+    })
-                    }
-                } => {
-                    log!("Standard Server received shutdown signal.");
-                    // Send kill signal to all active connection tasks
-                    let _ = shutdown_tx.send(());
-                    break;
-                }
-
-                // Accept new connections
-                accepted = listener.accept() => {
-                    match accepted {
-                        std::result::Result::Ok((stream, addr)) => {
-                            let service = HttpService {  peer_addr: addr };
-                            let io = TokioIo::new(stream);
-
-                            // Subscribe to the shutdown signal for this specific connection
-                            let mut rx = shutdown_tx.subscribe();
-
-                            tokio::spawn(async move {
-                                use hyper::server::conn::http2;
-
-                                let conn = http2::Builder::new(TokioExecutor::new())
-                                    .serve_connection(io, TowerToHyperService::new(service));
-
-
-                                // Wait for either the connection to finish naturally OR the shutdown signal
-                                tokio::select! {
-                                    res = conn => {
-                                        if let Err(err) = res {
-                                            if let Some(io_err) = err.source().and_then(|e| e.downcast_ref::<io::Error>()) {
-                                                if io_err.kind() != io::ErrorKind::ConnectionReset
-                                                   && io_err.kind() != io::ErrorKind::BrokenPipe
-                                                {
-                                                    log!("Error serving connection: {:?}", err);
-                                                }
-                                            }
-                                        }
-                                    }
-                                    _ = rx.recv() => {
-                                        // Shutdown signal received.
-                                        // Dropping the 'conn' future here closes the socket immediately.
-                                    }
-                                }
-                            });
-                        }
-                        Err(e) => {
-                            log!("Error accepting connection: {:?}", e);
-                            tokio::time::sleep(Duration::from_millis(500)).await;
-                        }
-                    }
-                }
-            }
-        }
-
-        log!("Standard Server shutdown complete.");
-    });
-
-    true
 }
 
-/// Runs the encrypted HTTPS/WSS server loop using the provided TLS config.
+fn find_local_ip() -> String {
-async fn run_tls_server(port: u16, tls_config: Arc<ServerConfig>) -> bool {
-    let mut ip = "0.0.0.0".to_string();
     for iface in pnet::datalink::interfaces() {
         let iface: NetworkInterface = iface;
-        let ipsv = format!("{}", iface.ips[0]);
+        if !iface.ips.is_empty() {
-        let ips: &str = ipsv.split('/').next().unwrap();
+            let ipsv = format!("{}", iface.ips[0]);
-        log!("{}", ips);
+            let ips: &str = ipsv.split('/').next().unwrap();
-        if format!("{}", ips).starts_with("10.") {
+            if ips.starts_with("10.") || ips.starts_with("192.") {
-            ip = ips.to_string();
+                return ips.to_string();
-        }
-    }
-
-    let acceptor = TlsAcceptor::from(tls_config);
-
-    let listener = TcpListener::bind(format!("0.0.0.0:{}", port)).await;
-    if let Err(e) = listener {
-        log!("Failed to bind to port {}: {:?}", port, e);
-        return false;
-    }
-    let listener = listener.unwrap();
-    log!(
-        "Encrypted Server listening for HTTPS and WSS on {}:{}",
-        ip,
-        port
-    );
-
-    // Create a broadcast channel for graceful shutdown signal
-    let (shutdown_tx, _) = broadcast::channel::<()>(1);
-
-    tokio::spawn(async move {
-        loop {
-            tokio::select! {
-                // Monitor for shutdown signal
-                _ = async {
-                    loop {
-                        tokio::time::sleep(Duration::from_millis(100)).await;
-                    }
-                } => {
-                    log!("Encrypted Server received shutdown signal.");
-                    // Send kill signal to all active connection tasks
-                    let _ = shutdown_tx.send(());
-                    break;
-                }
-
-                // Accept new connections
-                accepted = listener.accept() => {
-                    match accepted {
-                        std::result::Result::Ok((stream, addr)) => {
-                            let service = HttpService { peer_addr: addr };
-                            let acceptor = acceptor.clone();
-
-                            // Subscribe to the shutdown signal for this specific connection
-                            let mut rx = shutdown_tx.subscribe();
-
-                            tokio::spawn(async move {
-                                // Perform TLS handshake
-                                let tls_stream = match acceptor.accept(stream).await {
-                                    Ok(s) => s,
-                                    Err(e) => {
-                                        if e.kind() != io::ErrorKind::Interrupted {
-                                            log!("TLS Handshake error: {:?}", e);
-                                        }
-                                        return;
-                                    }
-                                };
-                                let io = TokioIo::new(tls_stream);
-
-                                // Prepare connection future
-                                let conn = http1::Builder::new()
-                                    .preserve_header_case(true)
-                                    .title_case_headers(true)
-                                    .serve_connection(io, TowerToHyperService::new(service))
-                                    .with_upgrades();
-
-                                // Wait for either the connection to finish naturally OR the shutdown signal
-                                tokio::select! {
-                                    res = conn => {
-                                        if let Err(err) = res {
-                                            if let Some(io_err) = err.source().and_then(|e| e.downcast_ref::<io::Error>()) {
-                                                if io_err.kind() != io::ErrorKind::ConnectionReset
-                                                   && io_err.kind() != io::ErrorKind::BrokenPipe
-                                                {
-                                                    log!("Error serving connection: {:?}", err);
-                                                }
-                                            }
-                                        }
-                                    }
-                                    _ = rx.recv() => {
-                                        // Shutdown signal received.
-                                        // Dropping the 'conn' future here closes the socket immediately.
-                                    }
-                                }
-                            });
-                        }
-                        Err(e) => {
-                            log!("Error accepting connection: {:?}", e);
-                            tokio::time::sleep(Duration::from_millis(500)).await;
-                        }
-                    }
-                }
             }
         }
-
-        log!("Encrypted Server shutdown complete.");
-    });
-    true
-}
-
-pub async fn start(port: u16) -> bool {
-    let tls_result = load_tls_config();
-
-    match tls_result {
-        Ok(Some(tls_config)) => run_tls_server(port, tls_config).await,
-        Ok(_) => run_http_server(port).await,
-        Err(e) => {
-            log!("Fatal error during TLS config load: {}", e);
-            false
-        }
     }
-}
+    "0.0.0.0".to_string()
-fn calculate_accept_key(key: &str) -> String {
-    let websocket_guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
-    let mut sha1 = Sha1::new();
-    sha1.update(key.as_bytes());
-    sha1.update(websocket_guid.as_bytes());
-    let result = sha1.finalize();
-    STANDARD.encode(result) // Base64 encode the result
-}
-
-/// Loads TLS config. Returns Ok(None) if cert files are not found, and an error if parsing fails.
-fn load_tls_config() -> Result<Option<Arc<ServerConfig>>, Box<dyn Error>> {
-    let cert_file_res = load_file_buf("certs", "cert.pem");
-    let key_file_res = load_file_buf("certs", "cert.key");
-
-    // Check if certificate files are present. If not, return None.
-    let cert_file_buf = match cert_file_res {
-        Ok(b) => b,
-        Err(e) if e.kind() == ErrorKind::NotFound => {
-            log!("TLS certificate 'certs/cert.pem' not found.");
-            return Ok(None);
-        }
-        Err(e) => return Err(e.into()), // Other IO error
-    };
-
-    let key_file_buf = match key_file_res {
-        Ok(b) => b,
-        Err(e) if e.kind() == ErrorKind::NotFound => {
-            log!("TLS key 'certs/cert.key' not found.");
-            return Ok(None);
-        }
-        Err(e) => return Err(e.into()), // Other IO error
-    };
-
-    // Continue with configuration if both files were found
-    let mut cert_reader = BufReader::new(cert_file_buf);
-    let cert_ders = rustls_pemfile::certs(&mut cert_reader)
-        .collect::<Result<Vec<CertificateDer>, io::Error>>()?;
-
-    // PKCS8
-    let mut key_reader = BufReader::new(key_file_buf);
-    let mut key_ders = rustls_pemfile::pkcs8_private_keys(&mut key_reader)
-        .map(|r| r.map(Into::into)) // Explicit conversion
-        .collect::<Result<Vec<PrivateKeyDer>, io::Error>>()?;
-
-    if key_ders.is_empty() {
-        // RSA
-        key_reader = BufReader::new(load_file_buf("certs", "cert.key")?); // Re-read key file
-        key_ders = rustls_pemfile::rsa_private_keys(&mut key_reader)
-            .map(|r| r.map(Into::into))
-            .collect::<Result<Vec<PrivateKeyDer>, io::Error>>()?;
-    }
-
-    if key_ders.is_empty() {
-        // EC
-        key_reader = BufReader::new(load_file_buf("certs", "cert.key")?); // Re-read key file
-        key_ders = rustls_pemfile::ec_private_keys(&mut key_reader)
-            .map(|r| r.map(Into::into))
-            .collect::<Result<Vec<PrivateKeyDer>, io::Error>>()?;
-    }
-
-    if key_ders.is_empty() {
-        return Err("No private keys found in key file. (Tried PKCS8, RSA, and EC)".into());
-    }
-
-    let mut config = rustls::ServerConfig::builder()
-        .with_no_client_auth()
-        .with_single_cert(cert_ders, key_ders.remove(0))?;
-
-    config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
-
-    Ok(Some(Arc::new(config)))
 }

src/server/socket.rs

@@ -1,51 +1,31 @@
 use std::sync::Arc;
 
+use axum::extract::ws::{Message, Utf8Bytes, WebSocket};
 use futures::StreamExt;
-use hyper::upgrade::OnUpgrade;
-use hyper_util::rt::TokioIo;
-use tokio_tungstenite::WebSocketStream;
-use tungstenite::{Message, Utf8Bytes};
 
 use crate::data::communication::{CommunicationType, CommunicationValue};
 use crate::log;
 use crate::server::omikron_connection::OmikronConnection;
 
-pub fn handle(path: String, upgrades: OnUpgrade) {
+pub fn handle(path: String, upgrades: WebSocket) {
     tokio::spawn(async move {
         log!(
             "[ws] Spawning new task to handle WebSocket upgrade for path: {}",
             path
         );
-        match upgrades.await {
+        log!("[ws] WebSocket upgrade successful for path: {}", path);
-            Ok(upgraded_stream) => {
-                log!("[ws] WebSocket upgrade successful for path: {}", path);
-                let raw_stream = TokioIo::new(upgraded_stream);
 
-                let ws_stream = WebSocketStream::from_raw_socket(
+        log!(
-                    raw_stream,
+            "[ws] WebSocket handshake successful, handling connection for {}",
-                    tungstenite::protocol::Role::Server,
+            path
-                    None,
+        );
-                )
-                .await;
-                log!(
-                    "[ws] WebSocket handshake successful, handling connection for {}",
-                    path
-                );
 
-                let (writer, reader) = ws_stream.split();
+        let (writer, reader) = upgrades.split();
-                if path == "/ws/omikron" {
+        if path == "/ws/omikron" {
-                    let connection = OmikronConnection::new(writer, reader);
+            let connection = OmikronConnection::new(writer, reader);
-                    tokio::spawn(start_connecteable_handler(connection));
+            tokio::spawn(start_connecteable_handler(connection));
-                }
-            }
-            Err(e) => {
-                log!(
-                    "[ERROR] WebSocket upgrade failed for path {}: {:?}",
-                    path,
-                    e
-                );
-            }
         }
+
         log!(
             "[ws] WebSocket handling task for path: {} is finished.",
             path
@@ -89,7 +69,7 @@ pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
                 log!("[ERROR] WS Error: {}. Breaking loop.", e);
                 break;
             }
-            Ok(None) => {
+            Ok(_) => {
                 log!("[ws_handler] WebSocket stream closed by peer. Breaking loop.");
                 break;
             }