tensamin/omega
Watch
6
0
Fork
You've already forked omega
 0
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

1280
Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

16
Cargo.toml
View file

@ -4,8 +4,13 @@ version = "0.1.0"
edition = "2024"
[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 = "*"
ansi_term = "0.12.1"
anyhow = "1.0.101"
async-trait = "0.1.89"
async-tungstenite = { version = "0.32.0", features = [
"futures-03-sink",
@ -39,7 +44,7 @@ async-tungstenite = { version = "0.32.0", features = [
"verbose-logging",
"webpki-roots",
] }
axum = { version = "0.8.8", features = [ "ws" ] }
axum = { version = "0.8.8", features = ["ws"] }
base64 = "0.22.1"
bytes = "1.11.1"
color-eyre = "0.6.5"
@ -50,15 +55,16 @@ futures-util = "0.3.31"
hex = "0.4.3"
hkdf = "0.12.4"
http-body-util = "0.1.3"
hyper = { version = "1.8.1", features = ["full"] }
hyper-util = "0.1.20"
hyper = { version = "1.8.1", features = ["http2", "full"] }
hyper-rustls = { version = "0.27.7", features = ["http2"] }
hyper-util = { version = "0.1.20", features = ["full"] }
json = "0.12.4"
once_cell = "1.21.3"
pnet = "0.35.0"
rand = "0.8"
rand_core = { version = "0.6", features = ["getrandom", "std"] }
reqwest = "0.12.28"
rustls = "0.23.35"
reqwest = "0.13.2"
rustls = "0.23.36"
rustls-pemfile = "2.2.0"
sha1 = "0.10.6"
sha2 = "0.10.9"

5
src/main.rs
View file

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

447
src/server/api.rs
View file

@ -1,27 +1,21 @@
use crate::data::communication::{CommunicationType, CommunicationValue, DataTypes};
use crate::get_public_key;
use crate::server::omikron_manager::get_random_omikron;
use crate::sql::sql;
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::{
sql::sql::{get_by_user_id, get_omikron_by_id},
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 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;
pub async fn handle(
path: &str,
_headers: HeaderMap<HeaderValue>,
body_string: Option<String>,
) -> HttpResponse<Full<Bytes>> {
pub async fn handle(path: &str, body_string: Option<String>) -> HttpResponse {
let path_parts: Vec<&str> = path.split("/").filter(|s| !s.is_empty()).collect();
let _body: Option<JsonValue> = if body_string.is_some() {
@ -33,258 +27,189 @@ pub async fn handle(
} else {
None
};
let (status, body_text) = if path_parts.len() >= 2 {
match path_parts[1] {
"download" => {
if path_parts.len() == 3 && path_parts[2] == "iota_frontend" {
let file: Bytes = load_file_vec("downloads", "iota_frontend.zip")
.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()
}
}
"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((public_key, ip_address)) =
sql::get_omikron_by_id(omikron_conn.get_omikron_id().await).await
{
(
StatusCode::OK,
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
omikron_conn.get_omikron_id().await,
public_key,
ip_address
),
)
} else {
(
StatusCode::INTERNAL_SERVER_ERROR,
"selected an invalid omikron".to_string(),
)
}
} else {
(
StatusCode::NOT_FOUND,
"couldn't find online omikron".to_string(),
)
}
} else if path_parts.len() == 4 {
let id = path_parts[3].parse::<i64>().unwrap_or(0);
if id == 0 {
not_found()
} else if let Ok((public_key, ip_address)) = get_omikron_by_id(id).await {
(
StatusCode::OK,
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
id, public_key, ip_address
),
)
} else if let Some(omikron_id) = get_iota_primary_omikron_connection(id) {
if let Ok((public_key, ip_address)) =
get_omikron_by_id(omikron_id).await
{
(
StatusCode::OK,
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
omikron_id, public_key, ip_address
),
)
} else {
not_found()
}
} else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) =
get_by_user_id(id).await
{
if let Some(omikron_id) = get_iota_primary_omikron_connection(iota_id) {
if let Ok((public_key, ip_address)) =
get_omikron_by_id(omikron_id).await
{
(
StatusCode::OK,
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
omikron_id, public_key, ip_address
),
)
} else {
not_found()
}
} else {
not_found()
}
} else {
not_found()
}
} else {
bad_request()
}
}
// get/id/<username>
"id" => {
if path_parts.len() != 4 {
bad_request()
} else {
let username = path_parts[3];
if username.is_empty() {
not_found()
} else {
if let Ok((
id,
iota_id,
username,
_,
_,
_,
_,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_username(username).await
{
let cv = CommunicationValue::new(CommunicationType::success)
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key)
.add_data(
DataTypes::user_id,
JsonValue::Number(Number::from(id)),
)
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
)
.add_data(
DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)),
)
.add_data(
DataTypes::sub_end,
JsonValue::Number(Number::from(sub_end)),
);
(StatusCode::OK, cv.to_json().to_string())
} else {
(
StatusCode::OK,
CommunicationValue::new(CommunicationType::error_not_found)
.to_json()
.to_string(),
)
}
}
}
}
"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);
if id == 0 {
bad_request()
} else {
if let Ok((
id,
iota_id,
username,
display,
status,
about,
avatar,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_user_id(id).await
{
let mut cv = CommunicationValue::new(CommunicationType::success)
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key)
.add_data(
DataTypes::user_id,
JsonValue::Number(Number::from(id)),
)
.add_data(
DataTypes::iota_id,
JsonValue::Number(Number::from(iota_id)),
)
.add_data(
DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)),
)
.add_data(
DataTypes::sub_end,
JsonValue::Number(Number::from(sub_end)),
);
if let Some(display) = display {
cv = cv.add_data_str(DataTypes::display, display);
}
if let Some(status) = status {
cv = cv.add_data_str(DataTypes::status, status);
}
if let Some(about) = about {
cv = cv.add_data_str(DataTypes::about, about);
}
if let Some(avatar) = avatar {
cv = cv.add_data_str(
DataTypes::avatar,
base64::engine::general_purpose::STANDARD.encode(avatar),
);
}
(StatusCode::OK, cv.to_json().to_string())
} else {
(
StatusCode::OK,
CommunicationValue::new(CommunicationType::error_not_found)
.to_json()
.to_string(),
)
}
}
}
}
_ => {
let id = path_parts[2];
let id: i64 = id.parse().unwrap_or(0);
if id == 0 {
bad_request()
} else {
bad_request()
}
}
},
_ => not_found(),
log!("Path parts: {:?}", path_parts);
let (status, body_text) = match path_parts.as_slice() {
["api", "download", "iota_frontend"] => {
let file = load_file("downloads", "iota_frontend.zip");
(StatusCode::OK, file)
}
} else {
not_found()
["api", "get", "omikron"] => {
if let Ok(omikron_conn) = get_random_omikron().await {
if let Ok((public_key, ip_address)) =
sql::get_omikron_by_id(omikron_conn.get_omikron_id().await).await
{
(
StatusCode::OK,
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
omikron_conn.get_omikron_id().await,
public_key,
ip_address
),
)
} else {
(
StatusCode::INTERNAL_SERVER_ERROR,
"selected an invalid omikron".to_string(),
)
}
} else {
(
StatusCode::NOT_FOUND,
"couldn't find online omikron".to_string(),
)
}
}
["api", "get", "omikron", id] => {
let id = id.parse::<i64>().unwrap_or(0);
if id == 0 {
not_found()
} else if let Ok((public_key, ip_address)) = get_omikron_by_id(id).await {
(
StatusCode::OK,
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
id, public_key, ip_address
),
)
} else if let Some(omikron_id) = get_iota_primary_omikron_connection(id) {
if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await {
(
StatusCode::OK,
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
omikron_id, public_key, ip_address
),
)
} else {
not_found()
}
} else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) = get_by_user_id(id).await
{
if let Some(omikron_id) = get_iota_primary_omikron_connection(iota_id) {
if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await {
(
StatusCode::OK,
format!(
"{{\"id\": {}, \"public_key\": \"{}\", \"ip_address\": \"{}\"}}",
omikron_id, public_key, ip_address
),
)
} else {
not_found()
}
} else {
not_found()
}
} else {
not_found()
}
}
["api", "get", "id", username] => {
if username.is_empty() {
not_found()
} else if let Ok((
id,
iota_id,
username,
_,
_,
_,
_,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_username(username).await
{
let cv = CommunicationValue::new(CommunicationType::success)
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key)
.add_data(DataTypes::user_id, JsonValue::Number(Number::from(id)))
.add_data(DataTypes::iota_id, JsonValue::Number(Number::from(iota_id)))
.add_data(
DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)),
)
.add_data(DataTypes::sub_end, JsonValue::Number(Number::from(sub_end)));
(StatusCode::OK, cv.to_json().to_string())
} else {
(
StatusCode::OK,
CommunicationValue::new(CommunicationType::error_not_found)
.to_json()
.to_string(),
)
}
}
["api", "get", "public_key"] => (StatusCode::OK, public_key_to_base64(&get_public_key())),
["api", "get", "user", id] => {
let id: i64 = id.parse().unwrap_or(0);
if id == 0 {
bad_request()
} else if let Ok((
id,
iota_id,
username,
display,
status,
about,
avatar,
sub_level,
sub_end,
public_key,
_,
_,
)) = sql::get_by_user_id(id).await
{
let mut cv = CommunicationValue::new(CommunicationType::success)
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::public_key, public_key)
.add_data(DataTypes::user_id, JsonValue::Number(Number::from(id)))
.add_data(DataTypes::iota_id, JsonValue::Number(Number::from(iota_id)))
.add_data(
DataTypes::sub_level,
JsonValue::Number(Number::from(sub_level)),
)
.add_data(DataTypes::sub_end, JsonValue::Number(Number::from(sub_end)));
if let Some(display) = display {
cv = cv.add_data_str(DataTypes::display, display);
}
if let Some(status) = status {
cv = cv.add_data_str(DataTypes::status, status);
}
if let Some(about) = about {
cv = cv.add_data_str(DataTypes::about, about);
}
if let Some(avatar) = avatar {
cv = cv.add_data_str(
DataTypes::avatar,
base64::engine::general_purpose::STANDARD.encode(avatar),
);
}
(StatusCode::OK, cv.to_json().to_string())
} else {
(
StatusCode::OK,
CommunicationValue::new(CommunicationType::error_not_found)
.to_json()
.to_string(),
)
}
}
_ => (
StatusCode::INTERNAL_SERVER_ERROR,
CommunicationValue::new(CommunicationType::error)
.to_json()
.to_string(),
),
};
let body = Full::new(Bytes::from(body_text.to_string()));
HttpResponse::builder().status(status).body(body).unwrap()
let body_bytes = body_text.to_string().into_bytes();
let body = BoxBody::new(body_bytes.clone());
HttpResponse::new(status).set_body(body)
}
pub fn bad_request() -> (StatusCode, String) {
(StatusCode::BAD_REQUEST, "400 Bad Request".to_string())
}

39
src/server/omikron_connection.rs
View file

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

438
src/server/server.rs
View file

@ -1,389 +1,83 @@
<<<<<<< HEAD
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use bytes::Bytes;
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 crate::{
log,
server::{api, short_link::get_short_link, socket},
util::file_util::get_directory,
};
use pnet::datalink::NetworkInterface;
use std::net::SocketAddr;
use std::time::Duration;
use tokio::net::TcpListener;
>>>>>>> 7f78c8669b36cbe39755d69cccd6971e56e10290
use actix_web::{App, HttpRequest, HttpResponse, HttpServer, Responder, http::header, web};
use actix_web_actors::ws;
use crate::log;
use crate::server::api;
use crate::server::short_link::get_short_link;
use crate::server::socket;
use rustls::ServerConfig;
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use rustls_pemfile::{certs, pkcs8_private_keys};
<<<<<<< HEAD
// --- ApiService for HTTP/2 ---
use std::fs::File;
use std::io::BufReader;
#[derive(Clone)]
struct ApiService {
_peer_addr: SocketAddr,
}
pub async fn start(port: u16) -> anyhow::Result<()> {
let mut cert_reader =
BufReader::new(File::open(format!("{}/certs/cert.pem", get_directory()))?);
impl Service<HttpRequest<Incoming>> for ApiService {
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
}
let mut key_reader = BufReader::new(File::open(format!("{}/certs/key.pem", get_directory()))?);
async fn ws_handler(
ws: WebSocketUpgrade,
OriginalUri(uri): OriginalUri,
ConnectInfo(_): ConnectInfo<SocketAddr>,
) -> impl IntoResponse {
log!("Attempting WebSocket upgrade on {}", uri.path());
let path = uri.path().to_string();
>>>>>>> 7f78c8669b36cbe39755d69cccd6971e56e10290
let cert_chain: Vec<CertificateDer<'static>> =
certs(&mut cert_reader).collect::<Result<_, _>>()?;
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
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 key = keys.remove(0);
let fut = async move {
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> = 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)
}
};
let mut config = ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(cert_chain, key)?;
Box::pin(fut.map_err(|err: color_eyre::eyre::ErrReport| {
io::Error::new(
io::ErrorKind::Other,
format!("Error in API request: {}", err),
)
}))
}
}
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
// --- WebSocketService for HTTP/1.1 ---
let addr = format!("0.0.0.0:{port}");
log!(" Server on {}", addr);
#[derive(Clone)]
struct WebSocketService {
_peer_addr: SocketAddr,
}
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
HttpServer::new(move || {
App::new()
.route("/api/{path:.*}", web::to(api_handler))
.route("/direct/{path:.*}", web::to(direct_handler))
.route("/ws/{path:.*}", web::get().to(ws_handler))
})
.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 {
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
async fn api_handler(req: HttpRequest, body: web::Bytes) -> HttpResponse {
let path = req.uri().path().to_string();
let body_string = String::from_utf8_lossy(&body).to_string();
api::handle(&path, Some(body_string)).await
}

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

@ -1,99 +1,129 @@
use std::sync::Arc;
use std::time::{Duration, Instant};
use axum::extract::ws::{Message, Utf8Bytes, WebSocket};
use futures::StreamExt;
use actix::{Actor, ActorContext, AsyncContext, StreamHandler};
use actix_web_actors::ws;
use crate::data::communication::{CommunicationType, CommunicationValue};
use crate::log;
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);
const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
log!(
"[ws] WebSocket handshake successful, handling connection for {}",
path
);
use actix::Message;
let (writer, reader) = upgrades.split();
if path == "/ws/omikron" {
let connection = OmikronConnection::new(writer, reader);
tokio::spawn(start_connecteable_handler(connection));
}
#[derive(Message)]
#[rtype(result = "()")]
pub struct WsSendMessage(pub String);
log!(
"[ws] WebSocket handling task for path: {} is finished.",
path
);
});
impl actix::Handler<WsSendMessage> for WsSession {
type Result = ();
fn handle(&mut self, msg: WsSendMessage, ctx: &mut Self::Context) {
ctx.text(msg.0);
}
}
pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
use futures::SinkExt;
use tokio::time::Duration;
const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
pub struct WsSession {
path: String,
last_heartbeat: Instant,
omikron: Option<Arc<OmikronConnection>>,
}
log!("[ws_handler] Starting connection handler loop.");
loop {
let mut receiver_guard = connection.receiver.write().await;
match tokio::time::timeout(IDLE_TIMEOUT, receiver_guard.next()).await {
Ok(Some(Ok(msg))) => {
drop(receiver_guard);
match msg {
Message::Text(text) => {
let conn_clone = connection.clone();
tokio::spawn(async move {
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.");
}
_ => {
log!("[ws_handler] Received unhandled message type.");
}
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 {
Ok(ws::Message::Text(text)) => {
self.last_heartbeat = Instant::now();
if let Some(conn) = &self.omikron {
let conn_clone = conn.clone();
actix_rt::spawn(async move {
conn_clone.handle_message(text.to_string()).await;
});
}
}
Ok(Some(Err(e))) => {
log!("[ERROR] WS Error: {}. Breaking loop.", e);
break;
Ok(ws::Message::Ping(msg)) => {
self.last_heartbeat = Instant::now();
ctx.pong(&msg);
}
Ok(_) => {
log!("[ws_handler] WebSocket stream closed by peer. Breaking loop.");
break;
Ok(ws::Message::Pong(_)) => {
self.last_heartbeat = Instant::now();
log!("[ws_handler] Received Pong. Connection alive.");
}
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.");
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();
}
_ => {}
}
}
log!("[ws_handler] Connection handler loop finished.");
connection.handle_close().await;
log!("[ws_handler] Connection closed.");
}