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

Merge branch 'main' of ssh://github.com/Tensamin/Omega

Browse source
This commit is contained in:
Alex Emmet 2026-02-16 10:54:16 +01:00
commit b3008e1e05
7 changed files with 1491 additions and 690 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

14
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",
@ -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"

4
src/main.rs
View file

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

172
src/server/api.rs
View file

@ -3,25 +3,25 @@ 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::{
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::http::{StatusCode, header};
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 {
if path == "OPTIONS" {
return HttpResponse::Ok()
.insert_header(("Access-Control-Allow-Origin", "*"))
.insert_header(("Access-Control-Allow-Methods", "GET, POST, OPTIONS"))
.insert_header(("Access-Control-Allow-Headers", "*"))
.finish();
}
let path_parts: Vec<&str> = path.split("/").filter(|s| !s.is_empty()).collect();
let _body: Option<JsonValue> = if body_string.is_some() {
@ -33,36 +33,30 @@ 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,
let (status, body_text) = match path_parts.as_slice() {
["api", "download", "iota_frontend"] => {
let file_path = "downloads/[iota_frontend].zip";
match std::fs::read(file_path) {
Ok(file_bytes) => {
return HttpResponse::Ok()
.insert_header(("Access-Control-Allow-Origin", "*"))
.insert_header(("Content-Type", "application/zip"))
.insert_header((
"Content-Disposition",
"attachment; filename=\"iota_frontend.zip\"",
)
.header(CONTENT_LENGTH, len)
.body(body)
.unwrap();
return response;
} else {
bad_request()
))
.body(file_bytes);
}
Err(_) => {
return HttpResponse::NotFound()
.insert_header(("Access-Control-Allow-Origin", "*"))
.body("File not found");
}
}
"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 {
}
["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
@ -88,8 +82,9 @@ pub async fn handle(
"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 {
not_found()
} else if let Ok((public_key, ip_address)) = get_omikron_by_id(id).await {
@ -101,9 +96,7 @@ pub async fn handle(
),
)
} 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
{
if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await {
(
StatusCode::OK,
format!(
@ -114,13 +107,10 @@ pub async fn handle(
} else {
not_found()
}
} else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) =
get_by_user_id(id).await
} 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
{
if let Ok((public_key, ip_address)) = get_omikron_by_id(omikron_id).await {
(
StatusCode::OK,
format!(
@ -137,20 +127,11 @@ pub async fn handle(
} else {
not_found()
}
} else {
bad_request()
}
}
// get/id/<username>
"id" => {
if path_parts.len() != 4 {
bad_request()
} else {
let username = path_parts[3];
["api", "get", "id", username] => {
if username.is_empty() {
not_found()
} else {
if let Ok((
} else if let Ok((
id,
iota_id,
username,
@ -168,22 +149,13 @@ pub async fn handle(
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::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)),
);
.add_data(DataTypes::sub_end, JsonValue::Number(Number::from(sub_end)));
(StatusCode::OK, cv.to_json().to_string())
} else {
(
@ -194,19 +166,12 @@ pub async fn handle(
)
}
}
}
}
"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];
["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((
} else if let Ok((
id,
iota_id,
username,
@ -224,22 +189,13 @@ pub async fn handle(
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::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)),
);
.add_data(DataTypes::sub_end, JsonValue::Number(Number::from(sub_end)));
if let Some(display) = display {
cv = cv.add_data_str(DataTypes::display, display);
}
@ -265,26 +221,22 @@ pub async fn handle(
)
}
}
}
}
_ => {
let id = path_parts[2];
let id: i64 = id.parse().unwrap_or(0);
if id == 0 {
bad_request()
} else {
bad_request()
}
}
},
_ => not_found(),
}
} else {
not_found()
_ => (
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.into_bytes();
HttpResponse::build(status)
.insert_header((header::ACCESS_CONTROL_ALLOW_ORIGIN, "*"))
.insert_header((header::ACCESS_CONTROL_ALLOW_HEADERS, "*"))
.insert_header((header::ACCESS_CONTROL_ALLOW_METHODS, "GET, POST, OPTIONS"))
.body(body_bytes)
}
pub fn bad_request() -> (StatusCode, String) {
(StatusCode::BAD_REQUEST, "400 Bad Request".to_string())
}

62
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
}
@ -226,6 +213,7 @@ impl OmikronConnection {
// ONLINE STATUS TRACKING
if cv.is_type(CommunicationType::user_connected) {
log_in!(PrintType::Omega, "User connected");
if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) {
user_online_tracker::track_user_status(
user_id,
@ -236,6 +224,7 @@ impl OmikronConnection {
return;
}
if cv.is_type(CommunicationType::user_disconnected) {
log_in!(PrintType::Omega, "User disconnected");
if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) {
if let Some(status) = user_online_tracker::get_user_status(user_id) {
user_online_tracker::track_user_status(
@ -277,6 +266,7 @@ impl OmikronConnection {
return;
}
if cv.is_type(CommunicationType::iota_disconnected) {
log_in!(PrintType::Omega, "IOTA disconnected");
if let Some(iota_id) = cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()) {
let iota_offline =
user_online_tracker::untrack_iota_connection(iota_id, omikron_id);
@ -337,7 +327,7 @@ impl OmikronConnection {
let mut response =
CommunicationValue::new(CommunicationType::get_user_data)
.with_id(cv.get_id())
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::username, username.clone())
.add_data_str(DataTypes::public_key, public_key)
.add_data(DataTypes::user_id, JsonValue::Number(Number::from(id)))
.add_data(
@ -354,14 +344,22 @@ impl OmikronConnection {
);
if let Some(display) = display {
if display.is_empty() {
response = response.add_data_str(DataTypes::display, username);
} else {
response = response.add_data_str(DataTypes::display, display);
}
}
if let Some(status) = status {
if !status.is_empty() {
response = response.add_data_str(DataTypes::status, status);
}
}
if let Some(about) = about {
if !about.is_empty() {
response = response.add_data_str(DataTypes::about, about);
}
}
if let Some(avatar) = avatar {
response =
response.add_data_str(DataTypes::avatar, STANDARD.encode(avatar));
@ -422,7 +420,7 @@ impl OmikronConnection {
let mut response =
CommunicationValue::new(CommunicationType::get_user_data)
.with_id(cv.get_id())
.add_data_str(DataTypes::username, username)
.add_data_str(DataTypes::username, username.clone())
.add_data_str(DataTypes::public_key, public_key)
.add_data(DataTypes::user_id, JsonValue::Number(Number::from(id)))
.add_data(
@ -439,14 +437,22 @@ impl OmikronConnection {
);
if let Some(display) = display {
if display.is_empty() {
response = response.add_data_str(DataTypes::display, username);
} else {
response = response.add_data_str(DataTypes::display, display);
}
}
if let Some(status) = status {
if !status.is_empty() {
response = response.add_data_str(DataTypes::status, status);
}
}
if let Some(about) = about {
if !about.is_empty() {
response = response.add_data_str(DataTypes::about, about);
}
}
if let Some(avatar) = avatar {
response =
response.add_data_str(DataTypes::avatar, STANDARD.encode(avatar));
@ -952,7 +958,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 {
@ -961,7 +966,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 {

155
src/server/server.rs Executable file → Normal file
View file

@ -1,106 +1,83 @@
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;
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};
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
}
use std::fs::File;
use std::io::BufReader;
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();
pub async fn start(port: u16) -> anyhow::Result<()> {
let mut cert_reader =
BufReader::new(File::open(format!("{}/certs/cert.pem", get_directory()))?);
ws.on_upgrade(async move |socket| socket::handle(path, socket))
}
let mut key_reader = BufReader::new(File::open(format!("{}/certs/key.pem", get_directory()))?);
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"),
}
}
let cert_chain: Vec<CertificateDer<'static>> =
certs(&mut cert_reader).collect::<Result<_, _>>()?;
async fn fallback_handler(
OriginalUri(uri): OriginalUri,
headers: axum::http::HeaderMap,
body: Body,
) -> impl IntoResponse {
let path = uri.path().to_string();
let mut keys: Vec<PrivateKeyDer<'static>> = pkcs8_private_keys(&mut key_reader)
.map(|res| res.map(Into::into))
.collect::<Result<_, _>>()?;
let whole_body = tokio::time::timeout(
Duration::from_secs(10),
axum::body::to_bytes(body, 1024 * 1024 * 10),
)
.await;
let key = keys.remove(0);
let body_string = match whole_body {
Ok(Ok(bytes)) => String::from_utf8(bytes.to_vec()).ok(),
_ => None,
};
let mut config = ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(cert_chain, key)?;
api::handle(&path, headers, body_string).await
}
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
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;
}
};
let addr = format!("0.0.0.0:{port}");
log!(" Server on {}", addr);
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()
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
}

178
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);
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.",
path
self.path
);
if let Some(conn) = &self.omikron {
let conn = conn.clone();
actix_rt::spawn(async move {
conn.handle_close().await;
});
}
}
}
pub async fn start_connecteable_handler(connection: Arc<OmikronConnection>) {
use futures::SinkExt;
use tokio::time::Duration;
const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
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);
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 {
Message::Text(text) => {
let conn_clone = connection.clone();
tokio::spawn(async move {
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;
});
}
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.");
}