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 4becc01f3a
7 changed files with 1432 additions and 961 deletions
Download patch file Download diff file Expand all files Collapse all files

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.");
}