[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:
parent
5cfdadd291
commit
cb7579d86d
7 changed files with 1432 additions and 961 deletions
1280
Cargo.lock
generated
1280
Cargo.lock
generated
File diff suppressed because it is too large
Load diff
14
Cargo.toml
14
Cargo.toml
|
|
@ -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"
|
||||||
|
|
|
||||||
|
|
@ -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();
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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())
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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 {
|
||||||
|
|
|
||||||
|
|
@ -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
188
src/server/socket.rs
Normal file → Executable 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.");
|
|
||||||
}
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue