[FIX] Websocket & webserver is now actix, because that should be better

than hyper & axum, i guess we'll see in the long run :},

ToDo: IF THIS WORKS: implement on the iota
This commit is contained in:
Alex Emmet 2026-02-15 00:37:42 +01:00
commit cb7579d86d
7 changed files with 1432 additions and 961 deletions

1280
Cargo.lock generated

File diff suppressed because it is too large Load diff

View file

@ -4,8 +4,13 @@ version = "0.1.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]
actix = "0.13.5"
actix-rt = "2.11.0"
actix-web = { version = "4.12.1", features = ["rustls-0_23"] }
actix-web-actors = "4.3.1"
aes-gcm = "*" aes-gcm = "*"
ansi_term = "0.12.1" ansi_term = "0.12.1"
anyhow = "1.0.101"
async-trait = "0.1.89" async-trait = "0.1.89"
async-tungstenite = { version = "0.32.0", features = [ async-tungstenite = { version = "0.32.0", features = [
"futures-03-sink", "futures-03-sink",
@ -50,15 +55,16 @@ futures-util = "0.3.31"
hex = "0.4.3" hex = "0.4.3"
hkdf = "0.12.4" hkdf = "0.12.4"
http-body-util = "0.1.3" http-body-util = "0.1.3"
hyper = { version = "1.8.1", features = ["full"] } hyper = { version = "1.8.1", features = ["http2", "full"] }
hyper-util = "0.1.20" hyper-rustls = { version = "0.27.7", features = ["http2"] }
hyper-util = { version = "0.1.20", features = ["full"] }
json = "0.12.4" json = "0.12.4"
once_cell = "1.21.3" once_cell = "1.21.3"
pnet = "0.35.0" pnet = "0.35.0"
rand = "0.8" rand = "0.8"
rand_core = { version = "0.6", features = ["getrandom", "std"] } rand_core = { version = "0.6", features = ["getrandom", "std"] }
reqwest = "0.12.28" reqwest = "0.13.2"
rustls = "0.23.35" rustls = "0.23.36"
rustls-pemfile = "2.2.0" rustls-pemfile = "2.2.0"
sha1 = "0.10.6" sha1 = "0.10.6"
sha2 = "0.10.9" sha2 = "0.10.9"

View file

@ -43,8 +43,7 @@ async fn main() {
} else { } else {
log!(" Users"); log!(" Users");
} }
server::server::start(9187, 9188).await; let _ = server::server::start(9187).await;
log!(" API Server on 9187");
log!(" WS Server on 9188");
tokio::signal::ctrl_c().await.unwrap(); tokio::signal::ctrl_c().await.unwrap();
} }

View file

@ -1,27 +1,21 @@
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::get_public_key;
use crate::server::omikron_manager::get_random_omikron; use crate::server::omikron_manager::get_random_omikron;
use crate::sql::sql; use crate::sql::sql;
use crate::sql::user_online_tracker::get_iota_primary_omikron_connection; use crate::sql::user_online_tracker::get_iota_primary_omikron_connection;
use crate::util::file_util::load_file_vec; use crate::util::file_util::load_file;
use crate::{get_public_key, log};
use crate::{ use crate::{
sql::sql::{get_by_user_id, get_omikron_by_id}, sql::sql::{get_by_user_id, get_omikron_by_id},
util::crypto_helper::public_key_to_base64, util::crypto_helper::public_key_to_base64,
}; };
use axum::http::HeaderValue; use actix_web::HttpResponse;
use actix_web::body::BoxBody;
use actix_web::http::StatusCode;
use base64::Engine as _; use base64::Engine as _;
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::JsonValue;
use json::number::Number; use json::number::Number;
pub async fn handle( pub async fn handle(path: &str, body_string: Option<String>) -> HttpResponse {
path: &str,
_headers: HeaderMap<HeaderValue>,
body_string: Option<String>,
) -> HttpResponse<Full<Bytes>> {
let path_parts: Vec<&str> = path.split("/").filter(|s| !s.is_empty()).collect(); let path_parts: Vec<&str> = path.split("/").filter(|s| !s.is_empty()).collect();
let _body: Option<JsonValue> = if body_string.is_some() { let _body: Option<JsonValue> = if body_string.is_some() {
@ -33,36 +27,13 @@ pub async fn handle(
} else { } else {
None None
}; };
let (status, body_text) = if path_parts.len() >= 2 { log!("Path parts: {:?}", path_parts);
match path_parts[1] { let (status, body_text) = match path_parts.as_slice() {
"download" => { ["api", "download", "iota_frontend"] => {
if path_parts.len() == 3 && path_parts[2] == "iota_frontend" { let file = load_file("downloads", "iota_frontend.zip");
let file: Bytes = load_file_vec("downloads", "iota_frontend.zip") (StatusCode::OK, file)
.unwrap()
.into();
let len = file.len();
let body = Full::new(file);
let response = HttpResponse::builder()
.status(StatusCode::OK)
.header(CONTENT_TYPE, "application/zip")
.header(
CONTENT_DISPOSITION,
"attachment; filename=\"iota_frontend.zip\"",
)
.header(CONTENT_LENGTH, len)
.body(body)
.unwrap();
return response;
} else {
bad_request()
} }
} ["api", "get", "omikron"] => {
"get" => match path_parts[2] {
// api/get/omikron -> any omikron
// api/get/omikron/<id> -> omikron for id (user / iota / omikron)
"omikron" => {
if path_parts.len() == 3 {
if let Ok(omikron_conn) = get_random_omikron().await { if let Ok(omikron_conn) = get_random_omikron().await {
if let Ok((public_key, ip_address)) = if let Ok((public_key, ip_address)) =
sql::get_omikron_by_id(omikron_conn.get_omikron_id().await).await sql::get_omikron_by_id(omikron_conn.get_omikron_id().await).await
@ -88,8 +59,9 @@ pub async fn handle(
"couldn't find online omikron".to_string(), "couldn't find online omikron".to_string(),
) )
} }
} else if path_parts.len() == 4 { }
let id = path_parts[3].parse::<i64>().unwrap_or(0); ["api", "get", "omikron", id] => {
let id = id.parse::<i64>().unwrap_or(0);
if id == 0 { if id == 0 {
not_found() not_found()
} else if let Ok((public_key, ip_address)) = get_omikron_by_id(id).await { } else if let Ok((public_key, ip_address)) = get_omikron_by_id(id).await {
@ -101,9 +73,7 @@ pub async fn handle(
), ),
) )
} else if let Some(omikron_id) = get_iota_primary_omikron_connection(id) { } else if let Some(omikron_id) = get_iota_primary_omikron_connection(id) {
if let Ok((public_key, ip_address)) = if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await {
get_omikron_by_id(omikron_id).await
{
( (
StatusCode::OK, StatusCode::OK,
format!( format!(
@ -114,13 +84,10 @@ pub async fn handle(
} else { } else {
not_found() not_found()
} }
} else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) = } else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) = get_by_user_id(id).await
get_by_user_id(id).await
{ {
if let Some(omikron_id) = get_iota_primary_omikron_connection(iota_id) { if let Some(omikron_id) = get_iota_primary_omikron_connection(iota_id) {
if let Ok((public_key, ip_address)) = if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await {
get_omikron_by_id(omikron_id).await
{
( (
StatusCode::OK, StatusCode::OK,
format!( format!(
@ -137,20 +104,11 @@ pub async fn handle(
} else { } else {
not_found() not_found()
} }
} else {
bad_request()
} }
} ["api", "get", "id", username] => {
// get/id/<username>
"id" => {
if path_parts.len() != 4 {
bad_request()
} else {
let username = path_parts[3];
if username.is_empty() { if username.is_empty() {
not_found() not_found()
} else { } else if let Ok((
if let Ok((
id, id,
iota_id, iota_id,
username, username,
@ -168,22 +126,13 @@ pub async fn handle(
let cv = CommunicationValue::new(CommunicationType::success) let cv = CommunicationValue::new(CommunicationType::success)
.add_data_str(DataTypes::username, username) .add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key) .add_data_str(DataTypes::public_key, public_key)
.add_data( .add_data(DataTypes::user_id, JsonValue::Number(Number::from(id)))
DataTypes::user_id, .add_data(DataTypes::iota_id, JsonValue::Number(Number::from(iota_id)))
JsonValue::Number(Number::from(id)),
)
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
)
.add_data( .add_data(
DataTypes::sub_level, DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)), JsonValue::Number(Number::from(sub_level)),
) )
.add_data( .add_data(DataTypes::sub_end, JsonValue::Number(Number::from(sub_end)));
DataTypes::sub_end,
JsonValue::Number(Number::from(sub_end)),
);
(StatusCode::OK, cv.to_json().to_string()) (StatusCode::OK, cv.to_json().to_string())
} else { } else {
( (
@ -194,19 +143,12 @@ pub async fn handle(
) )
} }
} }
} ["api", "get", "public_key"] => (StatusCode::OK, public_key_to_base64(&get_public_key())),
} ["api", "get", "user", id] => {
"public_key" => (StatusCode::OK, public_key_to_base64(&get_public_key())),
"user" => {
if path_parts.len() != 4 {
bad_request()
} else {
let id = path_parts[3];
let id: i64 = id.parse().unwrap_or(0); let id: i64 = id.parse().unwrap_or(0);
if id == 0 { if id == 0 {
bad_request() bad_request()
} else { } else if let Ok((
if let Ok((
id, id,
iota_id, iota_id,
username, username,
@ -224,22 +166,13 @@ pub async fn handle(
let mut cv = CommunicationValue::new(CommunicationType::success) let mut cv = CommunicationValue::new(CommunicationType::success)
.add_data_str(DataTypes::username, username) .add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key) .add_data_str(DataTypes::public_key, public_key)
.add_data( .add_data(DataTypes::user_id, JsonValue::Number(Number::from(id)))
DataTypes::user_id, .add_data(DataTypes::iota_id, JsonValue::Number(Number::from(iota_id)))
JsonValue::Number(Number::from(id)),
)
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
)
.add_data( .add_data(
DataTypes::sub_level, DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)), JsonValue::Number(Number::from(sub_level)),
) )
.add_data( .add_data(DataTypes::sub_end, JsonValue::Number(Number::from(sub_end)));
DataTypes::sub_end,
JsonValue::Number(Number::from(sub_end)),
);
if let Some(display) = display { if let Some(display) = display {
cv = cv.add_data_str(DataTypes::display, display); cv = cv.add_data_str(DataTypes::display, display);
} }
@ -265,26 +198,18 @@ pub async fn handle(
) )
} }
} }
} _ => (
} StatusCode::INTERNAL_SERVER_ERROR,
_ => { CommunicationValue::new(CommunicationType::error)
let id = path_parts[2]; .to_json()
let id: i64 = id.parse().unwrap_or(0); .to_string(),
if id == 0 { ),
bad_request()
} else {
bad_request()
}
}
},
_ => not_found(),
}
} else {
not_found()
}; };
let body = Full::new(Bytes::from(body_text.to_string())); let body_bytes = body_text.to_string().into_bytes();
HttpResponse::builder().status(status).body(body).unwrap() let body = BoxBody::new(body_bytes.clone());
HttpResponse::new(status).set_body(body)
} }
pub fn bad_request() -> (StatusCode, String) { pub fn bad_request() -> (StatusCode, String) {
(StatusCode::BAD_REQUEST, "400 Bad Request".to_string()) (StatusCode::BAD_REQUEST, "400 Bad Request".to_string())
} }

View file

@ -1,19 +1,16 @@
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes}; use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::server::omikron_manager; use crate::server::omikron_manager;
use crate::server::short_link::add_short_link; use crate::server::short_link::add_short_link;
use crate::server::socket::{WsSendMessage, WsSession};
use crate::sql::connection_status::UserStatus; use crate::sql::connection_status::UserStatus;
use crate::sql::sql::{self, get_by_user_id, get_by_username, get_iota_by_id, get_omikron_by_id}; use crate::sql::sql::{self, get_by_user_id, get_by_username, get_iota_by_id, get_omikron_by_id};
use crate::sql::user_online_tracker::{self}; use crate::sql::user_online_tracker::{self};
use crate::util::crypto_helper::encrypt; use crate::util::crypto_helper::encrypt;
use crate::util::logger::PrintType; use crate::util::logger::PrintType;
use crate::{get_private_key, get_public_key, log_err, log_in, log_out}; use crate::{get_private_key, get_public_key, log_in, log_out};
use axum::extract::ws::WebSocket; use actix::Addr;
use axum::extract::ws::{Message, Utf8Bytes};
use base64::{Engine as _, engine::general_purpose::STANDARD}; use base64::{Engine as _, engine::general_purpose::STANDARD};
use dashmap::DashMap; use dashmap::DashMap;
use futures::stream::SplitSink;
use futures::stream::SplitStream;
use futures_util::SinkExt;
use json::JsonValue; use json::JsonValue;
use json::number::Number; use json::number::Number;
use rand::Rng; use rand::Rng;
@ -24,8 +21,7 @@ use uuid::Uuid;
use x448::PublicKey; use x448::PublicKey;
pub struct OmikronConnection { pub struct OmikronConnection {
pub sender: Arc<RwLock<SplitSink<WebSocket, Message>>>, pub ws_addr: Arc<RwLock<Addr<WsSession>>>,
pub receiver: Arc<RwLock<SplitStream<WebSocket>>>,
pub omikron_id: Arc<RwLock<i64>>, pub omikron_id: Arc<RwLock<i64>>,
pub pub_key: Arc<RwLock<Option<Vec<u8>>>>, pub pub_key: Arc<RwLock<Option<Vec<u8>>>>,
identified: Arc<RwLock<bool>>, identified: Arc<RwLock<bool>>,
@ -39,13 +35,9 @@ pub struct OmikronConnection {
} }
impl OmikronConnection { impl OmikronConnection {
pub fn new( pub fn new(ws_addr: Addr<WsSession>) -> Arc<Self> {
sender: SplitSink<WebSocket, Message>,
receiver: SplitStream<WebSocket>,
) -> Arc<Self> {
Arc::new(Self { Arc::new(Self {
sender: Arc::new(RwLock::new(sender)), ws_addr: Arc::new(RwLock::new(ws_addr)),
receiver: Arc::new(RwLock::new(receiver)),
omikron_id: Arc::new(RwLock::new(0)), omikron_id: Arc::new(RwLock::new(0)),
pub_key: Arc::new(RwLock::new(None)), pub_key: Arc::new(RwLock::new(None)),
identified: Arc::new(RwLock::new(false)), identified: Arc::new(RwLock::new(false)),
@ -56,25 +48,20 @@ impl OmikronConnection {
}) })
} }
pub async fn send_message(&self, cv: &CommunicationValue) { pub async fn send_message(&self, cv: &CommunicationValue) {
let mut sender = self.sender.write().await; let text = cv.to_json().to_string();
let message_text = Message::Text(Utf8Bytes::from(cv.to_json().to_string()));
if !cv.is_type(CommunicationType::pong) { if !cv.is_type(CommunicationType::pong) {
log_out!( log_out!(
*self.omikron_id.read().await, *self.omikron_id.read().await,
PrintType::Omikron, PrintType::Omikron,
"{}", "{}",
cv.to_json().to_string() text
);
}
if let Err(e) = sender.send(message_text).await {
log_err!(
*self.omikron_id.read().await,
PrintType::Omikron,
"WebSocket send error: {}",
e
); );
} }
self.ws_addr.read().await.do_send(WsSendMessage(text));
} }
pub async fn get_omikron_id(&self) -> i64 { pub async fn get_omikron_id(&self) -> i64 {
*self.omikron_id.read().await *self.omikron_id.read().await
} }
@ -948,7 +935,6 @@ impl OmikronConnection {
self.send_message(&error).await; self.send_message(&error).await;
} }
pub async fn close(&self) { pub async fn close(&self) {
let mut sender = self.sender.write().await;
if self.is_identified().await { if self.is_identified().await {
let omikron_id = self.get_omikron_id().await; let omikron_id = self.get_omikron_id().await;
if omikron_id != 0 { if omikron_id != 0 {
@ -957,7 +943,6 @@ impl OmikronConnection {
user_online_tracker::untrack_omikron(omikron_id).await; user_online_tracker::untrack_omikron(omikron_id).await;
} }
} }
let _ = sender.close().await;
} }
pub async fn handle_close(self: Arc<Self>) { pub async fn handle_close(self: Arc<Self>) {
if self.is_identified().await { if self.is_identified().await {

View file

@ -1,389 +1,83 @@
<<<<<<< HEAD use crate::{
use base64::Engine; log,
use base64::engine::general_purpose::STANDARD; server::{api, short_link::get_short_link, socket},
use bytes::Bytes; util::file_util::get_directory,
use futures_util::TryFutureExt;
use http_body_util::BodyExt;
use http_body_util::Full;
use hyper::server::conn::{http1, http2};
use hyper::{Method, StatusCode};
use hyper::{Request as HttpRequest, Response as HttpResponse, body::Incoming, upgrade};
use hyper_util::rt::TokioExecutor;
use hyper_util::rt::tokio::TokioIo;
use hyper_util::service::TowerToHyperService;
use sha1::{Digest, Sha1};
use std::io;
use std::net::SocketAddr;
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 actix_web::{App, HttpRequest, HttpResponse, HttpServer, Responder, http::header, web};
use std::net::SocketAddr; use actix_web_actors::ws;
use std::time::Duration;
use tokio::net::TcpListener;
>>>>>>> 7f78c8669b36cbe39755d69cccd6971e56e10290
use crate::log; use rustls::ServerConfig;
use crate::server::api; use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use crate::server::short_link::get_short_link; use rustls_pemfile::{certs, pkcs8_private_keys};
use crate::server::socket;
<<<<<<< HEAD use std::fs::File;
// --- ApiService for HTTP/2 --- use std::io::BufReader;
#[derive(Clone)] pub async fn start(port: u16) -> anyhow::Result<()> {
struct ApiService { let mut cert_reader =
_peer_addr: SocketAddr, BufReader::new(File::open(format!("{}/certs/cert.pem", get_directory()))?);
}
impl Service<HttpRequest<Incoming>> for ApiService { let mut key_reader = BufReader::new(File::open(format!("{}/certs/key.pem", get_directory()))?);
type Response = HttpResponse<Full<Bytes>>;
type Error = io::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + 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
}
async fn ws_handler( let cert_chain: Vec<CertificateDer<'static>> =
ws: WebSocketUpgrade, certs(&mut cert_reader).collect::<Result<_, _>>()?;
OriginalUri(uri): OriginalUri,
ConnectInfo(_): ConnectInfo<SocketAddr>,
) -> 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)) let mut keys: Vec<PrivateKeyDer<'static>> = pkcs8_private_keys(&mut key_reader)
} .map(|res| res.map(Into::into))
.collect::<Result<_, _>>()?;
<<<<<<< HEAD let key = keys.remove(0);
fn call(&mut self, req: HttpRequest<Incoming>) -> Self::Future {
let (parts, body) = req.into_parts();
let path = parts.uri.path().to_string();
let headers = parts.headers.clone();
let fut = async move { let mut config = ServerConfig::builder()
if path.starts_with("/api") { .with_no_client_auth()
let whole_body = .with_single_cert(cert_chain, key)?;
tokio::time::timeout(Duration::from_secs(10), body.collect()).await??;
let bytes = whole_body.to_bytes();
let body_string: Option<String> = String::from_utf8(bytes.to_vec()).ok();
Ok(api::handle(&path, headers, 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 {
// Not a WebSocket, /api, or /direct, so it's a 404
let response = HttpResponse::builder()
.status(StatusCode::NOT_FOUND)
.body(Full::new(Bytes::from("Not Found")))
.unwrap();
Ok(response)
}
};
Box::pin(fut.map_err(|err: color_eyre::eyre::ErrReport| { config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
io::Error::new(
io::ErrorKind::Other,
format!("Error in API request: {}", err),
)
}))
}
}
// --- WebSocketService for HTTP/1.1 --- let addr = format!("0.0.0.0:{port}");
log!(" Server on {}", addr);
#[derive(Clone)] HttpServer::new(move || {
struct WebSocketService { App::new()
_peer_addr: SocketAddr, .route("/api/{path:.*}", web::to(api_handler))
} .route("/direct/{path:.*}", web::to(direct_handler))
.route("/ws/{path:.*}", web::get().to(ws_handler))
impl Service<HttpRequest<Incoming>> for WebSocketService {
type Response = HttpResponse<Full<Bytes>>;
type Error = io::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(
&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 {
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!("Attempting WebSocket upgrade on /ws");
if let Some(sec_websocket_key) = headers.get("sec-websocket-key") {
let sec_websocket_key_str =
sec_websocket_key.to_str().unwrap_or("").to_string();
let sec_websocket_accept = calculate_accept_key(&sec_websocket_key_str);
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 {
let response = HttpResponse::builder()
.status(StatusCode::NOT_FOUND)
.body(Full::new(Bytes::from(
"Not Found: Use the API server for non-WebSocket requests.",
)))
.unwrap();
Ok(response)
}
};
Box::pin(fut.map_err(|_err: color_eyre::eyre::ErrReport| {
io::Error::new(io::ErrorKind::Other, "Error in WebSocket request")
}))
}
}
async fn run_api_server(port: u16) -> bool {
let mut ip = "0.0.0.0".to_string();
for iface in pnet::datalink::interfaces() {
if let Some(ip_net) = iface.ips.iter().find(|ip_net| {
ip_net.is_ipv4()
&& (ip_net.ip().to_string().starts_with("10.")
|| ip_net.ip().to_string().starts_with("192."))
}) {
ip = ip_net.ip().to_string();
break;
}
}
let addr = SocketAddr::new(ip.parse().unwrap(), port);
let listener_res = TcpListener::bind(format!("0.0.0.0:{}", port)).await;
if let Err(e) = listener_res {
log!(
"[FATAL] Failed to bind API server to port {}: {:?}",
port,
e
);
return false;
}
let listener = listener_res.unwrap();
log!("API Server (HTTP/2) listening on {}:{}", ip, port);
tokio::spawn(async move {
loop {
if let Ok((stream, addr)) = listener.accept().await {
let service = ApiService { _peer_addr: addr };
let io = TokioIo::new(stream);
tokio::spawn(async move {
let conn = http2::Builder::new(TokioExecutor::new())
.serve_connection(io, TowerToHyperService::new(service));
if let Err(err) = conn.await {
if !err
.to_string()
.starts_with("error shutting down connection")
{
log!("API server connection error: {:?}", err);
}
}
});
}
}
});
true
}
async fn run_websocket_server(port: u16) -> bool {
let mut ip = "0.0.0.0".to_string();
for iface in pnet::datalink::interfaces() {
if let Some(ip_net) = iface.ips.iter().find(|ip_net| {
ip_net.is_ipv4()
&& (ip_net.ip().to_string().starts_with("10.")
|| ip_net.ip().to_string().starts_with("192."))
}) {
ip = ip_net.ip().to_string();
break;
}
}
let addr = SocketAddr::new(ip.parse().unwrap(), port);
let listener_res = TcpListener::bind(format!("0.0.0.0:{}", port)).await;
if let Err(e) = listener_res {
log!(
"[FATAL] Failed to bind WebSocket server to port {}: {:?}",
port,
e
);
return false;
}
let listener = listener_res.unwrap();
log!("WebSocket Server (HTTP/1.1) listening on {}:{}", ip, port);
tokio::spawn(async move {
loop {
if let Ok((stream, addr)) = listener.accept().await {
let service = WebSocketService { _peer_addr: addr };
let io = TokioIo::new(stream);
tokio::spawn(async move {
let conn = http1::Builder::new()
.preserve_header_case(true)
.title_case_headers(true)
.serve_connection(io, TowerToHyperService::new(service))
.with_upgrades();
if let Err(err) = conn.await {
if !err
.to_string()
.starts_with("error shutting down connection")
{
log!("WebSocket server connection error: {:?}", err);
}
}
});
}
}
});
true
}
// --- Main Start function ---
pub async fn start(api_port: u16, ws_port: u16) {
// We are not using TLS for this setup as per the request's focus on splitting protocols.
// The previous TLS loading logic is removed for simplicity.
tokio::spawn(run_api_server(api_port));
tokio::spawn(run_websocket_server(ws_port));
}
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)
=======
async fn direct_handler(Path(short): Path<String>) -> 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::<SocketAddr>(),
)
.await
.map(|_| true)
.unwrap_or_else(|e| {
log!("Server error: {}", e);
false
}) })
.bind_rustls_0_23(addr, config)?
.run()
.await?;
Ok(())
}
async fn direct_handler(req: HttpRequest) -> impl Responder {
let path = req.uri().path().to_string();
let short = path.replace("/direct/", "");
if let Ok(long) = get_short_link(&short).await {
HttpResponse::TemporaryRedirect()
.append_header((header::LOCATION, long))
.finish()
} else {
HttpResponse::TemporaryRedirect()
.append_header((header::LOCATION, "https://tensamin.net"))
.finish()
}
}
async fn ws_handler(
req: HttpRequest,
stream: web::Payload,
path: web::Path<String>,
) -> Result<HttpResponse, actix_web::Error> {
let path = path.into_inner();
println!("WS handler reached: {}", path);
ws::start(socket::WsSession::new(path), &req, stream)
} }
fn find_local_ip() -> String { async fn api_handler(req: HttpRequest, body: web::Bytes) -> HttpResponse {
for iface in pnet::datalink::interfaces() { let path = req.uri().path().to_string();
let iface: NetworkInterface = iface; let body_string = String::from_utf8_lossy(&body).to_string();
if !iface.ips.is_empty() {
let ipsv = format!("{}", iface.ips[0]); api::handle(&path, Some(body_string)).await
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
} }

188
src/server/socket.rs Normal file → Executable file
View file

@ -1,99 +1,129 @@
use std::sync::Arc; use std::sync::Arc;
use std::time::{Duration, Instant};
use axum::extract::ws::{Message, Utf8Bytes, WebSocket}; use actix::{Actor, ActorContext, AsyncContext, StreamHandler};
use futures::StreamExt; use actix_web_actors::ws;
use crate::data::communication::{CommunicationType, CommunicationValue}; use crate::data::communication::{CommunicationType, CommunicationValue};
use crate::log; use crate::log;
use crate::server::omikron_connection::OmikronConnection; use crate::server::omikron_connection::OmikronConnection;
pub fn handle(path: String, upgrades: WebSocket) {
tokio::spawn(async move {
log!(
"[ws] Spawning new task to handle WebSocket upgrade for path: {}",
path
);
log!("[ws] WebSocket upgrade successful for path: {}", path);
log!(
"[ws] WebSocket handshake successful, handling connection for {}",
path
);
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
);
});
}
pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
use futures::SinkExt;
use tokio::time::Duration;
const IDLE_TIMEOUT: Duration = Duration::from_secs(30); const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
log!("[ws_handler] Starting connection handler loop."); use actix::Message;
loop {
let mut receiver_guard = connection.receiver.write().await;
match tokio::time::timeout(IDLE_TIMEOUT, receiver_guard.next()).await { #[derive(Message)]
Ok(Some(Ok(msg))) => { #[rtype(result = "()")]
drop(receiver_guard); pub struct WsSendMessage(pub String);
impl actix::Handler<WsSendMessage> for WsSession {
type Result = ();
fn handle(&mut self, msg: WsSendMessage, ctx: &mut Self::Context) {
ctx.text(msg.0);
}
}
pub struct WsSession {
path: String,
last_heartbeat: Instant,
omikron: Option<Arc<OmikronConnection>>,
}
impl WsSession {
pub fn new(path: String) -> Self {
Self {
path,
last_heartbeat: Instant::now(),
omikron: None,
}
}
fn start_heartbeat(&self, ctx: &mut ws::WebsocketContext<Self>) {
ctx.run_interval(Duration::from_secs(5), |act, ctx| {
if Instant::now().duration_since(act.last_heartbeat) > IDLE_TIMEOUT {
log!("[ws_handler] Heartbeat failed. Disconnecting.");
ctx.close(None);
ctx.stop();
return;
}
let ping = CommunicationValue::new(CommunicationType::ping)
.to_json()
.to_string();
ctx.text(ping);
});
}
}
impl Actor for WsSession {
type Context = ws::WebsocketContext<Self>;
fn started(&mut self, ctx: &mut Self::Context) {
self.start_heartbeat(ctx);
if self.path == "omikron" {
let addr = ctx.address();
let connection = OmikronConnection::new(addr);
self.omikron = Some(connection);
}
}
fn stopped(&mut self, _: &mut Self::Context) {
log!(
"[ws] WebSocket handling task for path: {} is finished.",
self.path
);
if let Some(conn) = &self.omikron {
let conn = conn.clone();
actix_rt::spawn(async move {
conn.handle_close().await;
});
}
}
}
impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WsSession {
fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
match msg { match msg {
Message::Text(text) => { Ok(ws::Message::Text(text)) => {
let conn_clone = connection.clone(); self.last_heartbeat = Instant::now();
tokio::spawn(async move {
if let Some(conn) = &self.omikron {
let conn_clone = conn.clone();
actix_rt::spawn(async move {
conn_clone.handle_message(text.to_string()).await; conn_clone.handle_message(text.to_string()).await;
}); });
} }
Message::Close(_) => {
log!("[ws_handler] Received 'Close' message. Breaking loop.");
break;
} }
Message::Pong(_) => {
log!("[ws_handler] Received 'Pong'. Connection is alive."); Ok(ws::Message::Ping(msg)) => {
self.last_heartbeat = Instant::now();
ctx.pong(&msg);
} }
_ => {
log!("[ws_handler] Received unhandled message type."); Ok(ws::Message::Pong(_)) => {
self.last_heartbeat = Instant::now();
log!("[ws_handler] Received Pong. Connection alive.");
}
Ok(ws::Message::Close(reason)) => {
log!("[ws_handler] Received Close. Disconnecting.");
ctx.close(reason);
ctx.stop();
}
Ok(ws::Message::Binary(_)) => {
log!("[ws_handler] Binary message ignored.");
}
Err(e) => {
log!("[ERROR] WS Error: {}. Closing.", e);
ctx.stop();
}
_ => {}
} }
} }
} }
Ok(Some(Err(e))) => {
log!("[ERROR] WS Error: {}. Breaking loop.", e);
break;
}
Ok(_) => {
log!("[ws_handler] WebSocket stream closed by peer. Breaking loop.");
break;
}
Err(_) => {
drop(receiver_guard);
log!("[ws_handler] Timeout: Dropped receiver lock. Sending a ping.");
let mut sender = connection.sender.write().await;
log!("[ws_handler] Acquired sender lock for ping.");
if let Err(e) = sender
.send(Message::Text(Utf8Bytes::from(
CommunicationValue::new(CommunicationType::ping)
.to_json()
.to_string(),
)))
.await
{
log!("[ERROR] Failed to send ping: {}. Closing connection.", e);
break;
}
log!("[ws_handler] Ping sent successfully.");
}
}
}
log!("[ws_handler] Connection handler loop finished.");
connection.handle_close().await;
log!("[ws_handler] Connection closed.");
}