diff --git a/Cargo.lock b/Cargo.lock index 4b58524..ba72ade 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 22121da..f89a8a4 100755 --- a/Cargo.toml +++ b/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" diff --git a/src/server/api.rs b/src/server/api.rs index 35ee9ec..ae30a24 100644 --- a/src/server/api.rs +++ b/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 hyper::body::{Body, Bytes, Frame}; +use http_body_util::Full; +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, diff --git a/src/server/omikron_connection.rs b/src/server/omikron_connection.rs index 7ce8f8b..3eeb22b 100644 --- a/src/server/omikron_connection.rs +++ b/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 hyper_util::rt::TokioIo; +use futures_util::SinkExt; 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>, Message>>>, - pub receiver: Arc>>>>, + pub sender: Arc>>, + pub receiver: Arc>>, pub omikron_id: Arc>, pub pub_key: Arc>>>, identified: Arc>, @@ -43,8 +40,8 @@ pub struct OmikronConnection { impl OmikronConnection { pub fn new( - sender: SplitSink>, Message>, - receiver: SplitStream>>, + sender: SplitSink, + receiver: SplitStream, ) -> Arc { 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, 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; diff --git a/src/server/omikron_manager.rs b/src/server/omikron_manager.rs index 8440835..864c9d6 100644 --- a/src/server/omikron_manager.rs +++ b/src/server/omikron_manager.rs @@ -24,11 +24,3 @@ pub async fn get_random_omikron() -> Result, ()> { return Err(()); } } - -pub async fn get_omikron(omikron_id: i64) -> Option> { - if let Some(omikron) = OMIKRON_CONNECTIONS.get(&omikron_id) { - Some(omikron.clone()) - } else { - None - } -} diff --git a/src/server/server.rs b/src/server/server.rs old mode 100755 new mode 100644 index 22c24c6..afa0244 --- a/src/server/server.rs +++ b/src/server/server.rs @@ -1,3 +1,4 @@ +<<<<<<< HEAD use base64::Engine; use base64::engine::general_purpose::STANDARD; use bytes::Bytes; @@ -17,12 +18,27 @@ use std::result::Result::Ok; use std::{future::Future, pin::Pin, time::Duration}; use tokio::net::TcpListener; use tower::Service; +======= +use axum::{ + Router, + body::Body, + extract::{ConnectInfo, OriginalUri, Path, ws::WebSocketUpgrade}, + response::{IntoResponse, Redirect}, + routing::get, +}; + +use pnet::datalink::NetworkInterface; +use std::net::SocketAddr; +use std::time::Duration; +use tokio::net::TcpListener; +>>>>>>> 7f78c8669b36cbe39755d69cccd6971e56e10290 use crate::log; use crate::server::api; use crate::server::short_link::get_short_link; use crate::server::socket; +<<<<<<< HEAD // --- ApiService for HTTP/2 --- #[derive(Clone)] @@ -34,14 +50,28 @@ impl Service> for ApiService { type Response = HttpResponse>; type Error = io::Error; type Future = Pin> + Send>>; +======= +pub async fn start(port: u16) -> bool { + let app = Router::new() + .route("/ws/omikron", get(ws_handler)) + .route("/direct/{short}", get(direct_handler)) + .fallback(fallback_handler); + run_http_server(port, app).await +} - fn poll_ready( - &mut self, - _cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - std::task::Poll::Ready(std::io::Result::Ok(())) - } +async fn ws_handler( + ws: WebSocketUpgrade, + OriginalUri(uri): OriginalUri, + ConnectInfo(_): ConnectInfo, +) -> impl IntoResponse { + log!("Attempting WebSocket upgrade on {}", uri.path()); + let path = uri.path().to_string(); +>>>>>>> 7f78c8669b36cbe39755d69cccd6971e56e10290 + ws.on_upgrade(async move |socket| socket::handle(path, socket)) +} + +<<<<<<< HEAD fn call(&mut self, req: HttpRequest) -> Self::Future { let (parts, body) = req.into_parts(); let path = parts.uri.path().to_string(); @@ -286,4 +316,74 @@ fn calculate_accept_key(key: &str) -> String { sha1.update(websocket_guid.as_bytes()); let result = sha1.finalize(); STANDARD.encode(result) +======= +async fn direct_handler(Path(short): Path) -> impl IntoResponse { + match get_short_link(&short).await { + Ok(long) => Redirect::temporary(&long), + Err(_) => Redirect::temporary("https://tensamin.net"), + } +} + +async fn fallback_handler( + OriginalUri(uri): OriginalUri, + headers: axum::http::HeaderMap, + body: Body, +) -> impl IntoResponse { + let path = uri.path().to_string(); + + let whole_body = tokio::time::timeout( + Duration::from_secs(10), + axum::body::to_bytes(body, 1024 * 1024 * 10), + ) + .await; + + let body_string = match whole_body { + Ok(Ok(bytes)) => String::from_utf8(bytes.to_vec()).ok(), + _ => None, + }; + + api::handle(&path, headers, body_string).await +} + +async fn run_http_server(port: u16, app: Router) -> bool { + let ip = find_local_ip(); + let listener = match TcpListener::bind(format!("0.0.0.0:{}", port)).await { + Ok(l) => l, + Err(e) => { + log!("Failed to bind to port {}: {:?}", port, e); + return false; + } + }; + + log!( + "Standard Server listening for HTTP and WS on {}:{}", + ip, + port + ); + + axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .await + .map(|_| true) + .unwrap_or_else(|e| { + log!("Server error: {}", e); + false + }) +} + +fn find_local_ip() -> String { + for iface in pnet::datalink::interfaces() { + let iface: NetworkInterface = iface; + if !iface.ips.is_empty() { + let ipsv = format!("{}", iface.ips[0]); + let ips: &str = ipsv.split('/').next().unwrap(); + if ips.starts_with("10.") || ips.starts_with("192.") { + return ips.to_string(); + } + } + } + "0.0.0.0".to_string() +>>>>>>> 7f78c8669b36cbe39755d69cccd6971e56e10290 } diff --git a/src/server/socket.rs b/src/server/socket.rs index 5a8b815..25a00f7 100644 --- a/src/server/socket.rs +++ b/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 { - Ok(upgraded_stream) => { - log!("[ws] WebSocket upgrade successful for path: {}", path); - let raw_stream = TokioIo::new(upgraded_stream); + log!("[ws] WebSocket upgrade successful for path: {}", path); - let ws_stream = WebSocketStream::from_raw_socket( - raw_stream, - tungstenite::protocol::Role::Server, - None, - ) - .await; - log!( - "[ws] WebSocket handshake successful, handling connection for {}", - path - ); + log!( + "[ws] WebSocket handshake successful, handling connection for {}", + path + ); - let (writer, reader) = ws_stream.split(); - if path == "/ws/omikron" { - let connection = OmikronConnection::new(writer, reader); - tokio::spawn(start_connecteable_handler(connection)); - } - } - Err(e) => { - log!( - "[ERROR] WebSocket upgrade failed for path {}: {:?}", - path, - e - ); - } + let (writer, reader) = upgrades.split(); + if path == "/ws/omikron" { + let connection = OmikronConnection::new(writer, reader); + tokio::spawn(start_connecteable_handler(connection)); } + log!( "[ws] WebSocket handling task for path: {} is finished.", path @@ -89,7 +69,7 @@ pub async fn start_connecteable_handler(connection: Arc) { log!("[ERROR] WS Error: {}. Breaking loop.", e); break; } - Ok(None) => { + Ok(_) => { log!("[ws_handler] WebSocket stream closed by peer. Breaking loop."); break; }