From 1deda333ab51a3acad79ae3f3b45482eb5e7ed79 Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Sat, 14 Feb 2026 11:17:15 +0100 Subject: [PATCH] [WIP] Migration to AXUM & HTTP2 --- Cargo.lock | 62 ++++- Cargo.toml | 4 +- src/server/api.rs | 426 ++++++++++++++------------------ src/server/server.rs | 564 ++++++++++--------------------------------- src/server/socket.rs | 8 +- 5 files changed, 367 insertions(+), 697 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 286f75d..b3f2d5a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -91,6 +91,15 @@ dependencies = [ "derive_arbitrary", ] +[[package]] +name = "arc-swap" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ded5f9a03ac8f24d1b8a25101ee812cd32cdc8c50a4c50237de2c4915850e73" +dependencies = [ + "rustversion", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -164,11 +173,12 @@ dependencies = [ [[package]] name = "axum" -version = "0.8.7" +version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b098575ebe77cb6d14fc7f32749631a6e44edbef6b796f89b020e99ba20d425" +checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8" dependencies = [ "axum-core", + "base64", "bytes", "form_urlencoded", "futures-util", @@ -187,8 +197,10 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_urlencoded", + "sha1", "sync_wrapper", "tokio", + "tokio-tungstenite", "tower", "tower-layer", "tower-service", @@ -214,6 +226,28 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum-server" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1df331683d982a0b9492b38127151e6453639cd34926eb9c07d4cd8c6d22bfc" +dependencies = [ + "arc-swap", + "bytes", + "either", + "fs-err", + "http", + "http-body", + "hyper", + "hyper-util", + "pin-project-lite", + "rustls", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tower-service", +] + [[package]] name = "backtrace" version = "0.3.76" @@ -985,6 +1019,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs-err" +version = "3.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73fde052dbfc920003cfd2c8e2c6e6d4cc7c1091538c3a24226cec0665ab08c0" +dependencies = [ + "autocfg", + "tokio", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -1550,6 +1594,7 @@ dependencies = [ "async-tungstenite", "aws-lc-rs", "axum", + "axum-server", "base64", "bytes", "chrono", @@ -1581,6 +1626,7 @@ dependencies = [ "rustls", "rustls-pemfile", "serde", + "serde_json", "sha1", "sha2", "strum 0.27.2", @@ -3016,15 +3062,15 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.145" +version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "402a6f66d8c709116cf22f558eab210f5a50187f702eb4d7e5ef38d9a7f1c79c" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ "itoa", "memchr", - "ryu", "serde", "serde_core", + "zmij", ] [[package]] @@ -4615,6 +4661,12 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51f936044d677be1a1168fae1d03b583a285a5dd9d8cbf7b24c23aa1fc775235" +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + [[package]] name = "zopfli" version = "0.8.3" diff --git a/Cargo.toml b/Cargo.toml index 6e24e98..16d54e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ edition = "2024" [dependencies] aes-gcm = "*" -axum = "*" +axum = { version = "0.8.8", features = ["ws", "http2"] } base64 = "0.22.1" bytes = "*" cmake = "*" @@ -60,3 +60,5 @@ ratatui = "0.30.0" ratatui_input = "0.1.3" open = "5.3.3" chrono = "0.4.43" +axum-server = { version = "0.8.0", features = ["tls-rustls"] } +serde_json = "1.0.149" diff --git a/src/server/api.rs b/src/server/api.rs index c9637a0..2e6b7aa 100755 --- a/src/server/api.rs +++ b/src/server/api.rs @@ -1,255 +1,187 @@ +use crate::server::server::is_local_network; +use axum::routing::{get, post}; +use axum::{ + extract::{ConnectInfo, Json, Path, State}, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, +}; +use json::JsonValue; +use serde_json::{Value, json}; +use std::net::SocketAddr; use std::sync::Arc; -use crate::communities::community::Community; -use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; -use crate::gui::log_panel::log_message; -use crate::util::file_util::delete_file; -use crate::{RELOAD, SHUTDOWN}; -use axum::http::HeaderValue; -use http_body_util::Full; -use hyper::body::Bytes; -use hyper::{HeaderMap, Response as HttpResponse, StatusCode}; -use json::JsonValue; - -use crate::util::config_util::CONFIG; -use crate::{APP_STATE, communities::community_manager, users::user_manager}; - -pub async fn handle( - path: &str, - is_local: &bool, - headers: HeaderMap, - body_string: Option, -) -> HttpResponse> { - if !is_local { - return HttpResponse::builder() - .status(StatusCode::FORBIDDEN) - .body(Full::new(Bytes::from("403 Forbidden".to_string()))) - .unwrap(); +pub fn api_router() -> axum::Router { + axum::Router::new() + .route("/shutdown", post(shutdown)) + .route("/reload", post(reload)) + .route("/users/add", post(users_add)) + .route("/users/remove", post(users_remove)) + .route("/users/get", get(users_get)) + .route("/communities/add", post(communities_add)) + .route("/communities/get", get(communities_get)) + .route("/settings/set", post(settings_set)) + .route("/settings/get", get(settings_get)) +} +pub async fn settings_set( + ConnectInfo(addr): ConnectInfo, + State(ssl): State, + headers: HeaderMap, +) -> impl IntoResponse { + if !is_allowed(addr, ssl) { + return forbidden(); } - let path_parts: Vec<&str> = path.split("/").collect(); - let body: Option = if body_string.is_some() { - if let Ok(body_json) = json::parse(&body_string.unwrap()) { - Some(body_json) - } else { - None + let key = headers.get("key").and_then(|v| v.to_str().ok()); + let value = headers.get("value").and_then(|v| v.to_str().ok()); + + match (key, value) { + (Some(k), Some(v)) => { + crate::util::config_util::CONFIG + .write() + .await + .config + .insert(k, v); + + success() } - } else { - None + _ => error(), + } +} + +pub async fn settings_get( + ConnectInfo(addr): ConnectInfo, + State(ssl): State, +) -> impl IntoResponse { + if !is_allowed(addr, ssl) { + return forbidden(); + } + + (StatusCode::OK, Json(CONFIG.read().await.config.clone())).into_response() +} +pub async fn communities_get( + ConnectInfo(addr): ConnectInfo, + State(ssl): State, +) -> impl IntoResponse { + if !is_allowed(addr, ssl) { + return forbidden(); + } + + let communities = crate::communities::community_manager::get_communities().await; + + let mut list = JsonValue::new_array(); + + for c in communities { + list.push(c.frontend().await); + } + + (StatusCode::OK, Json(list)).into_response() +} +pub async fn communities_add( + ConnectInfo(addr): ConnectInfo, + State(ssl): State, + Json(payload): Json, +) -> impl IntoResponse { + if !is_allowed(addr, ssl) { + return forbidden(); + } + + let name = payload["name"].as_str().unwrap_or("").to_string(); + let owner = payload["owner"].as_i64().unwrap_or(0); + + let community = Arc::new(crate::communities::community::Community::create(name, owner).await); + + crate::communities::community_manager::add_community(community).await; + + success() +} +pub async fn users_get( + ConnectInfo(addr): ConnectInfo, + State(ssl): State, +) -> impl IntoResponse { + if !is_allowed(addr, ssl) { + return forbidden(); + } + + let users = crate::users::user_manager::get_users(); + + let list: Vec<_> = users.into_iter().map(|u| u.frontend()).collect(); + + Json(list) +} +pub async fn users_remove( + ConnectInfo(addr): ConnectInfo, + State(ssl): State, + Json(payload): Json, +) -> impl IntoResponse { + if !is_allowed(addr, ssl) { + return forbidden(); + } + + let uuid = payload.get("uuid").and_then(|v| v.as_i64()).unwrap_or(0); + + crate::users::user_manager::remove_user(uuid); + crate::users::user_manager::save_users(); + + success() +} +pub async fn users_add( + ConnectInfo(addr): ConnectInfo, + State(ssl): State, + Json(payload): Json, +) -> impl IntoResponse { + if !is_allowed(addr, ssl) { + return forbidden(); + } + + let username = match payload.get("username").and_then(|v| v.as_str()) { + Some(u) => u, + None => return error().into_response(), }; - let (status, content, body_text) = if path_parts.len() >= 3 { - match path_parts[2] { - "shutdown" => { - *SHUTDOWN.write().await = true; - ( - StatusCode::OK, - "application/json", - "{\"type\":\"success\"}".to_string(), - ) - } - "reload" => { - *SHUTDOWN.write().await = true; - *RELOAD.write().await = true; - ( - StatusCode::OK, - "application/json", - "{\"type\":\"success\"}".to_string(), - ) - } - "app_state" => (StatusCode::OK, "application/json", { - let with = headers - .get("size") - .unwrap_or(&HeaderValue::from_static("50")) - .to_str() - .unwrap() - .to_string(); - let json = APP_STATE - .lock() - .unwrap() - .with_width(with.parse::().unwrap_or(50)); - json.to_json().to_string() - }), - "users" => (StatusCode::OK, "application/json", { - if path_parts.len() >= 4 { - match path_parts[3] { - "add" => { - if body.is_none() { - "{\"type\":\"error\"}".to_string() - } else { - let username = - body.unwrap()["username"].as_str().unwrap().to_string(); - if let (Some(user), Some(_private_key)) = - user_manager::create_user(&username).await - { - let cv = - CommunicationValue::new(CommunicationType::create_user) - .add_data(DataTypes::user, user.frontend()); - cv.to_json().to_string() - } else { - "{\"type\":\"error\"}".to_string() - } - } - } - "remove_tu" => { - if body.is_none() { - "{\"type\":\"error\"}".to_string() - } else { - let username = - body.unwrap()["username"].as_str().unwrap().to_string(); - - delete_file("", &format!("{}.tu", username)); - - "{\"type\":\"success\"}".to_string() - } - } - "remove" => { - if body.is_none() { - "{\"type\":\"error\"}".to_string() - } else { - let uuid = body.unwrap()["uuid"].as_i64().unwrap_or(0); - // TODO MOVE TO OMIKRON CONNECTION - /*unregister_user( - uuid, - &user_manager::get_user(uuid).unwrap().reset_token, - ) - .await;*/ - user_manager::remove_user(uuid); - user_manager::save_users(); - "{}".to_string() - } - } - "get" => { - let users = user_manager::get_users(); - let mut json = JsonValue::new_array(); - for user in users { - let _ = json.push(user.frontend()); - } - json.to_string() - } - _ => "{\"type\":\"error\"}".to_string(), - } - } else { - let users = user_manager::get_users(); - let mut json = JsonValue::new_array(); - for user in users { - let _ = json.push(user.to_json()); - } - json.to_string() - } - }), - "communities" => (StatusCode::OK, "application/json", { - if path_parts.len() >= 4 { - match path_parts[3] { - "add" => { - if let Some(body) = body { - let name = body["name"].as_str().unwrap().to_string(); - let user_id = body["owner"].as_i64().unwrap_or(0); - let community = Arc::new(Community::create(name, user_id).await); - community_manager::add_community(community).await; - "{\"type\":\"success\"}".to_string() - } else { - "{\"type\":\"error\"}".to_string() - } - } - "remove" => { - if let Some(body) = body { - let name = body["name"].as_str().unwrap().to_string(); - community_manager::remove_community(&name).await; - "{\"type\":\"success\"}".to_string() - } else { - "{\"type\":\"error\"}".to_string() - } - } - "change" => { - if path_parts.len() >= 5 { - match path_parts[4] { - "owner" => { - if let Some(body) = body { - let name = body["name"].as_str().unwrap().to_string(); - let owner = body["owner"].as_i64().unwrap_or(0); - community_manager::get_community(&name) - .await - .unwrap() - .set_owner(owner) - .await; - "{\"type\":\"success\"}".to_string() - } else { - "{\"type\":\"error\"}".to_string() - } - } - _ => "{\"type\":\"error\"}".to_string(), - } - } else { - "{\"type\":\"error\"}".to_string() - } - } - "get" => { - let communities = community_manager::get_communities().await; - let mut json = JsonValue::new_array(); - for community in communities { - let _ = json.push(community.frontend().await); - } - json.to_string() - } - _ => "{\"type\":\"error\"}".to_string(), - } - } else { - let communities = community_manager::get_communities().await; - let mut json = JsonValue::new_array(); - for community in communities { - let _ = json.push(community.to_json().await); - } - json.to_string() - } - }), - "settings" => (StatusCode::OK, "application/json", { - if path_parts.len() >= 4 { - match path_parts[3] { - "set" => { - if let Some(key) = headers.get("key") { - if let Some(value) = headers.get("value") { - let _ = CONFIG - .write() - .await - .config - .insert(key.to_str().unwrap(), value.to_str().unwrap()); - "{\"type\":\"success\"}".to_string() - } else { - "{\"type\":\"error\"}".to_string() - } - } else { - "{\"type\":\"error\"}".to_string() - } - } - "get" => CONFIG.read().await.config.to_string(), - _ => "{\"type\":\"error\"}".to_string(), - } - } else { - CONFIG.read().await.config.to_string() - } - }), - - _ => { - log_message(format!("Unknown API endpoint: {}", path)); - ( - StatusCode::NOT_FOUND, - "application/json", - "404 Not Found".to_string(), - ) - } - } + + if let (Some(user), Some(_)) = crate::users::user_manager::create_user(username).await { + (StatusCode::OK, Json(user.frontend())).into_response() } else { - log_message(format!("Invalid API path: {}", path)); - ( - StatusCode::NOT_FOUND, - "application/json", - "404 Not Found".to_string(), - ) - }; - let body = Full::new(Bytes::from(body_text.to_string())); - HttpResponse::builder() - .header("Content-Type", content) - .status(status) - .body(body) - .unwrap() + error() + } +} +pub async fn shutdown( + ConnectInfo(addr): ConnectInfo, + State(ssl): State, +) -> impl IntoResponse { + if !is_allowed(addr, ssl) { + return forbidden(); + } + + *crate::SHUTDOWN.write().await = true; + success() +} + +pub async fn reload( + ConnectInfo(addr): ConnectInfo, + State(ssl): State, +) -> impl IntoResponse { + if !is_allowed(addr, ssl) { + return forbidden(); + } + + *crate::SHUTDOWN.write().await = true; + *crate::RELOAD.write().await = true; + + success() +} + +fn forbidden() -> Response { + (StatusCode::FORBIDDEN, "403 Forbidden").into_response() +} + +fn success() -> Response { + Json(json!({ "type": "success" })).into_response() +} + +fn error() -> Response { + Json(json!({ "type": "error" })).into_response() +} + +fn is_allowed(addr: SocketAddr, ssl: bool) -> bool { + is_local_network(addr.ip()) || ssl } diff --git a/src/server/server.rs b/src/server/server.rs index b4449f2..8a26652 100644 --- a/src/server/server.rs +++ b/src/server/server.rs @@ -1,492 +1,164 @@ use crate::gui::log_panel::log_message; -use crate::server::api; +use crate::server::api::{self, api_router}; use crate::server::socket::handle; use crate::util::file_util::{load_file_buf, load_file_vec}; use crate::{ACTIVE_TASKS, SHUTDOWN}; -use base64::Engine; -use base64::engine::general_purpose::STANDARD; -use bytes::Bytes; -use futures::StreamExt; -use futures_util::TryFutureExt; -use http_body_util::BodyExt; -use http_body_util::Full; -use hyper::{Method, StatusCode}; -use hyper::{ - Request as HttpRequest, Response as HttpResponse, body::Incoming, server::conn::http1, upgrade, +use axum::{ + Router, + extract::{ + ConnectInfo, Path, + ws::{WebSocket, WebSocketUpgrade}, + }, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + routing::{any, get}, }; -use hyper_util::rt::tokio::TokioIo; -use hyper_util::service::TowerToHyperService; -use pnet::datalink::NetworkInterface; +use axum_server::tls_rustls::RustlsConfig; +use bytes::Bytes; +use futures_util::{SinkExt, StreamExt}; use rustls::ServerConfig; use rustls::pki_types::{CertificateDer, PrivateKeyDer}; -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 std::{ + error::Error, + io::{self, BufReader, ErrorKind}, + net::{IpAddr, SocketAddr}, + sync::Arc, + time::Duration, +}; use tokio::net::TcpListener; -use tokio::sync::broadcast; // Import broadcast for the kill switch -use tokio_rustls::TlsAcceptor; -use tokio_tungstenite::WebSocketStream; -use tower::Service; -#[derive(Clone)] -struct HttpService { - peer_addr: SocketAddr, - ssl: bool, +use tower::ServiceBuilder; + +fn build_router(ssl: bool) -> Router { + Router::new() + .route("/ws/*path", get(ws_handler)) + .nest("/api/*path", api_router()) + .route("/*path", any(static_handler)) + .layer(ServiceBuilder::new()) + .with_state(ssl) } -impl Service> for HttpService { - type Response = HttpResponse>; - type Error = io::Error; - type Future = Pin> + Send>>; +async fn ws_handler( + ws: WebSocketUpgrade, + Path(path): Path, + ConnectInfo(addr): ConnectInfo, +) -> impl IntoResponse { + log_message(format!("WS connection from {}", addr)); - fn poll_ready( - &mut self, - _cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - std::task::Poll::Ready(std::io::Result::Ok(())) - } - - fn call(&mut self, req: HttpRequest) -> Self::Future { - let peer_ip = self.peer_addr.ip(); - - let is_acceptable = is_local_network(peer_ip) || self.ssl; - - let (parts, body) = req.into_parts(); - - 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_message("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); - - match upgrades.await { - std::result::Result::Ok(upgraded_stream) => { - let raw_stream = TokioIo::new(upgraded_stream); - - let handshake_result = WebSocketStream::from_raw_socket( - raw_stream, - tungstenite::protocol::Role::Server, - None, - ) - .await; - log_message(format!("Handling WebSocket connection",)); - let (writer, reader) = handshake_result.split(); - handle(path.clone(), writer, reader); - } - Err(e) => { - log_message(format!( - "WebSocket upgrade failed after response: {:?}", - e - )); - } - } - Ok(response) - } else { - log_message("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 = match body.collect().await { - Ok(collected) => collected, - Err(e) => { - log_message(format!("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 = match String::from_utf8(bytes.to_vec()) { - Ok(s) => Some(s), - Err(_) => None, - }; - - Ok(api::handle(&path, &is_acceptable, headers.clone(), body_string).await) - } else { - let mut path_parts: Vec<&str> = path.split("/").collect(); - let name = path_parts.remove(path_parts.len() - 1); - let name = if name.is_empty() { - "index.html" - } else if name.contains(".") && name.contains("?") { - name.split("?").next().unwrap() - } else if name.contains(".") { - name - } else { - &format!("{}.html", name) - }; - let code = if let Some(ext) = name.split(".").last() { - match ext { - "html" => "text/html", - "css" => "text/css", - "ico" => "image/x-icon", - "png" => "image/png", - "js" => "application/javascript", - "json" => "application/json", - _ => "application/octet-stream", - } - } else { - "application/octet-stream" - }; - let (status, content, body_text): (StatusCode, &str, Vec) = { - let content = load_file_vec(&format!("web{}/", path_parts.join("/")), name); - if content.is_empty() { - let content = load_file_vec("web", "404.html"); - if content.is_empty() { - ( - StatusCode::NOT_FOUND, - code, - include_str!("../../static/web/404.html") - .as_bytes() - .to_vec(), - ) - } else { - (StatusCode::OK, code, content) - } - } else { - (StatusCode::OK, code, content) - } - }; - - let body = Full::new(Bytes::from(body_text.to_vec())); - let response = HttpResponse::builder() - .header("Content-Type", content) - .status(status) - .body(body) - .unwrap(); - Ok(response) - } - }; - - Box::pin(fut.map_err(|err: color_eyre::eyre::ErrReport| { - io::Error::new( - io::ErrorKind::Other, - format!("Error in request handling: {}", err), - ) - })) - } + ws.on_upgrade(move |socket| async move { + handle_ws(socket, path).await; + }) } -pub fn is_local_network(addr: IpAddr) -> bool { - match addr { - IpAddr::V4(v4) => { - let octets = v4.octets(); - if octets[0] == 10 { - return true; - } - 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; - } +async fn handle_ws(socket: WebSocket, path: String) { + let (mut sender, mut receiver) = socket.split(); + handle(path, sender, receiver); +} - false - } +async fn static_handler(Path(path): Path) -> impl IntoResponse { + let mut parts: Vec<&str> = path.split('/').collect(); + let name = parts.pop().unwrap_or("index.html"); - IpAddr::V6(v6) => { - let segments = v6.segments(); - if (segments[0] & 0xfe00) == 0xfc00 { - return true; - } - if (segments[0] & 0xffc0) == 0xfe80 { - return true; - } - if v6.is_loopback() { - return true; - } + let name = if name.is_empty() { + "index.html" + } else if name.contains('.') { + name + } else { + &format!("{}.html", name) + }; + let content = load_file_vec(&format!("web{}/", parts.join("/")), name); + + if content.is_empty() { + return (StatusCode::NOT_FOUND, load_file_vec("web", "404.html")); + } + + let mime = match name.split('.').last().unwrap_or("") { + "html" => "text/html", + "css" => "text/css", + "js" => "application/javascript", + "json" => "application/json", + "png" => "image/png", + "ico" => "image/x-icon", + _ => "application/octet-stream", + }; + + (StatusCode::OK, content) +} + +pub async fn start(port: u16) -> bool { + match load_tls_config() { + Ok(Some(tls)) => run_tls_server(port, tls).await, + Ok(_) => run_http_server(port).await, + Err(e) => { + log_message(format!("TLS config error: {}", e)); false } } } async fn run_http_server(port: u16) -> bool { - let mut ip = "0.0.0.0".to_string(); - for iface in pnet::datalink::interfaces() { - let iface: NetworkInterface = iface; - if iface.ips.len() > 0 { - let ipsv = format!("{}", iface.ips[0]); - let ips: &str = ipsv.split('/').next().unwrap(); - 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_message(format!("Failed to bind to port {}: {:?}", port, e)); - return false; - } - let listener = listener.unwrap(); - log_message(format!( - "Standard Server listening for HTTP and WS on {}:{}", - ip, port - )); + let addr = SocketAddr::from(([0, 0, 0, 0], port)); + let router = build_router(false); - // Create a broadcast channel for graceful shutdown signal - let (shutdown_tx, _) = broadcast::channel::<()>(1); + log_message(format!("HTTP Server running on {}", addr)); - ACTIVE_TASKS.lock().unwrap().push("WebServer".to_string()); + ACTIVE_TASKS.lock().unwrap().push("WebServer".into()); tokio::spawn(async move { - loop { - tokio::select! { - // Monitor for shutdown signal - _ = async { - loop { - if *SHUTDOWN.read().await { return; } - tokio::time::sleep(Duration::from_millis(100)).await; - } - } => { - log_message("Standard Server received shutdown signal."); - // Send kill signal to all active connection tasks - let _ = shutdown_tx.send(()); - break; - } + let listener = TcpListener::bind(addr).await.unwrap(); - // Accept new connections - accepted = listener.accept() => { - match accepted { - std::result::Result::Ok((stream, addr)) => { - let service = HttpService { ssl: false, peer_addr: addr }; - let io = TokioIo::new(stream); + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(wait_for_shutdown()) + .await + .unwrap(); - // Subscribe to the shutdown signal for this specific connection - let mut rx = shutdown_tx.subscribe(); - - tokio::spawn(async move { - // Prepare the 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::()) { - if io_err.kind() != io::ErrorKind::ConnectionReset - && io_err.kind() != io::ErrorKind::BrokenPipe - { - log_message(format!("Error serving connection: {:?}", err)); - } - } - } - } - _ = rx.recv() => { - // Shutdown signal received. - // Dropping the 'conn' future here closes the socket immediately. - } - } - }); - } - Err(e) => { - log_message(format!("Error accepting connection: {:?}", e)); - tokio::time::sleep(Duration::from_millis(500)).await; - } - } - } - } - } - - ACTIVE_TASKS - .lock() - .unwrap() - .retain(|t| !t.eq(&"WebServer".to_string())); - log_message("Standard Server shutdown complete."); + ACTIVE_TASKS.lock().unwrap().retain(|t| t != "WebServer"); + log_message("HTTP Server shutdown complete."); }); true } -/// Runs the encrypted HTTPS/WSS server loop using the provided TLS config. -async fn run_tls_server(port: u16, tls_config: Arc) -> 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]); - let ips: &str = ipsv.split('/').next().unwrap(); - log_message(ips); - if format!("{}", ips).starts_with("10.") { - ip = ips.to_string(); - } - } +async fn run_tls_server(port: u16, tls: Arc) -> bool { + let addr = SocketAddr::from(([0, 0, 0, 0], port)); + let router = build_router(true); - let acceptor = TlsAcceptor::from(tls_config); + let tls_config = RustlsConfig::from_config(tls); - let listener = TcpListener::bind(format!("0.0.0.0:{}", port)).await; - if let Err(e) = listener { - log_message(format!("Failed to bind to port {}: {:?}", port, e)); - return false; - } - let listener = listener.unwrap(); - log_message(format!( - "Encrypted Server listening for HTTPS and WSS on {}:{}", - ip, port - )); + log_message(format!("HTTPS (HTTP/2) Server running on {}", addr)); - // Create a broadcast channel for graceful shutdown signal - let (shutdown_tx, _) = broadcast::channel::<()>(1); - - ACTIVE_TASKS.lock().unwrap().push("WebServer".to_string()); + ACTIVE_TASKS.lock().unwrap().push("WebServer".into()); tokio::spawn(async move { - loop { - tokio::select! { - // Monitor for shutdown signal - _ = async { - loop { - if *SHUTDOWN.read().await { return; } - tokio::time::sleep(Duration::from_millis(100)).await; - } - } => { - log_message("Encrypted Server received shutdown signal."); - // Send kill signal to all active connection tasks - let _ = shutdown_tx.send(()); - break; - } + axum_server::bind_rustls(addr, tls_config) + .serve(router.into_make_service_with_connect_info::()) + .await + .unwrap(); - // Accept new connections - accepted = listener.accept() => { - match accepted { - std::result::Result::Ok((stream, addr)) => { - let service = HttpService { ssl: true, 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_message(format!("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::()) { - if io_err.kind() != io::ErrorKind::ConnectionReset - && io_err.kind() != io::ErrorKind::BrokenPipe - { - log_message(format!("Error serving connection: {:?}", err)); - } - } - } - } - _ = rx.recv() => { - // Shutdown signal received. - // Dropping the 'conn' future here closes the socket immediately. - } - } - }); - } - Err(e) => { - log_message(format!("Error accepting connection: {:?}", e)); - tokio::time::sleep(Duration::from_millis(500)).await; - } - } - } - } - } - - ACTIVE_TASKS - .lock() - .unwrap() - .retain(|t| !t.eq(&"WebServer".to_string())); - log_message("Encrypted Server shutdown complete."); + ACTIVE_TASKS.lock().unwrap().retain(|t| t != "WebServer"); + log_message("HTTPS 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_message(format!("Fatal error during TLS config load: {}", e)); - false +async fn wait_for_shutdown() { + loop { + if *SHUTDOWN.read().await { + log_message("Shutdown signal received."); + break; } + tokio::time::sleep(Duration::from_millis(100)).await; } } -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>, Box> { 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 => { @@ -543,3 +215,19 @@ fn load_tls_config() -> Result>, Box> { Ok(Some(Arc::new(config))) } +pub fn is_local_network(addr: IpAddr) -> bool { + match addr { + IpAddr::V4(v4) => { + let o = v4.octets(); + o[0] == 10 + || (o[0] == 172 && (16..=31).contains(&o[1])) + || (o[0] == 192 && o[1] == 168) + || o[0] == 127 + || (o[0] == 169 && o[1] == 254) + } + IpAddr::V6(v6) => { + let s = v6.segments(); + (s[0] & 0xfe00) == 0xfc00 || (s[0] & 0xffc0) == 0xfe80 || v6.is_loopback() + } + } +} diff --git a/src/server/socket.rs b/src/server/socket.rs index 65531e4..2171668 100644 --- a/src/server/socket.rs +++ b/src/server/socket.rs @@ -3,19 +3,15 @@ use crate::communities::{community_connection::CommunityConnection, community_ma use crate::gui::log_panel::log_message; use crate::omikron::omikron_connection::OmikronConnection; +use axum::extract::ws::{Message, Utf8Bytes, WebSocket}; use futures::StreamExt; use futures::stream::SplitSink; use futures::stream::SplitStream; use hyper::upgrade::Upgraded; use hyper_util::rt::TokioIo; use std::sync::Arc; -use tungstenite::Message; -pub fn handle( - path: String, - writer: SplitSink>, Message>, - reader: SplitStream>>, -) { +pub fn handle(path: String, writer: SplitSink, reader: SplitStream) { tokio::spawn(async move { if path.starts_with("/ws/users/") { OmikronConnection::client(writer, reader).await;