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

[Fix] Stability

Browse source
This commit is contained in:
Alex 2026-07-27 20:36:23 +02:00
commit 4f7260419a
20 changed files with 563 additions and 464 deletions
Download patch file Download diff file Expand all files Collapse all files

59
src/transport/handlers/presence.rs
View file

@ -21,7 +21,18 @@ pub async fn user_connected(
.and_then(UserStatus::from_str)
.unwrap_or(UserStatus::user_online);
if let Ok(user_id) = i64::try_from(user_id) {
user_online_tracker::track_user_status(user_id, status, omikron_id);
if let Some(session_id) = value
.get_data(DataType::SessionId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
{
user_online_tracker::track_user_session_status(
user_id, session_id, status, omikron_id,
);
} else {
user_online_tracker::track_user_status(user_id, status, omikron_id);
}
}
}
Ok(())
@ -34,11 +45,55 @@ pub async fn user_disconnected(
) -> OmikronResult<()> {
log_in!(crate::util::logger::PrintType::Omega, "User disconnected");
if let Some(user_id) = value.get_data(DataType::UserId).as_number() {
user_online_tracker::untrack_user_status(user_id as i64, omikron_id);
if let Some(session_id) = value
.get_data(DataType::SessionId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
{
user_online_tracker::untrack_user_session_status(
user_id as i64,
session_id,
omikron_id,
);
} else {
user_online_tracker::untrack_user_status(user_id as i64, omikron_id);
}
}
Ok(())
}
pub async fn client_changed(
_: Arc<OmikronConnection>,
value: CommunicationValue,
_: i64,
) -> OmikronResult<()> {
let Some(user_id) = value
.get_data(DataType::UserId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
else {
return Ok(());
};
let Some(status) = value
.get_data(DataType::UserState)
.as_str()
.and_then(UserStatus::from_str)
else {
return Ok(());
};
// Connectivity is derived from routes. Clients may choose only public
// presence preferences, never server/offline states.
if matches!(
status,
UserStatus::user_offline | UserStatus::iota_offline | UserStatus::iota_online
) {
return Ok(());
}
user_online_tracker::update_user_session_status(user_id, status);
Ok(())
}
pub async fn iota_connected(
connection: Arc<OmikronConnection>,
value: CommunicationValue,

119
src/transport/handlers/register.rs
View file

@ -13,13 +13,25 @@ pub async fn get_register(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let register_id = user_repo::get_register_id().await?;
let iota_id = value
.get_data(DataType::IotaId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| user_repo::valid_protocol_id(*id));
let Some(iota_id) = iota_id else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.await;
};
let (register_id, registration_token) =
user_repo::allocate_registration(IotaId::from(iota_id), value.get_id()).await?;
let response = CommunicationValue::new(CommunicationType::GetRegister)
.with_id(value.get_id())
.add_typed_default(
DataType::UserId,
DataValue::SignedNumber(register_id.0.into()),
);
)
.add_typed_default(DataType::RegisterId, DataValue::Str(registration_token));
connection.send(&response).await
}
@ -27,10 +39,6 @@ pub async fn complete_iota(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let iota_id = value
.get_data(DataType::IotaId)
.as_number()
.map(|id| id as i64);
let public_key = value
.get_data(DataType::PublicKey)
.as_str()
@ -40,57 +48,25 @@ pub async fn complete_iota(
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.await;
};
match iota_id {
Some(iota_id) => {
match iota_repo::register_complete_iota(IotaId::from(iota_id), public_key).await {
Ok(()) => {
connection
.send(
&CommunicationValue::new(CommunicationType::Success)
.with_id(value.get_id()),
)
.await
}
Err(error) => {
connection
.send(
&CommunicationValue::new(CommunicationType::ErrorInternal)
.with_id(value.get_id())
.add_typed_default(
DataType::ErrorType,
DataValue::Str(error.to_string()),
),
)
.await
}
}
match iota_repo::create_new_iota(public_key).await {
Ok(id) => {
connection
.send(
&CommunicationValue::new(CommunicationType::CompleteRegisterIota)
.with_id(value.get_id())
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(id.0.into())),
)
.await
}
Err(error) => {
connection
.send(
&CommunicationValue::new(CommunicationType::ErrorInternal)
.with_id(value.get_id())
.add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string())),
)
.await
}
None => match iota_repo::create_new_iota(public_key).await {
Ok(id) => {
connection
.send(
&CommunicationValue::new(CommunicationType::CompleteRegisterIota)
.with_id(value.get_id())
.add_typed_default(
DataType::IotaId,
DataValue::SignedNumber(id.0.into()),
),
)
.await
}
Err(error) => {
connection
.send(
&CommunicationValue::new(CommunicationType::ErrorInternal)
.with_id(value.get_id())
.add_typed_default(
DataType::ErrorType,
DataValue::Str(error.to_string()),
),
)
.await
}
},
}
}
@ -101,7 +77,8 @@ pub async fn complete_user(
let user_id = value
.get_data(DataType::UserId)
.as_number()
.map(|id| id as i64);
.and_then(|id| i64::try_from(id).ok())
.filter(|id| user_repo::valid_protocol_id(*id));
let username = value
.get_data(DataType::Username)
.as_str()
@ -114,22 +91,44 @@ pub async fn complete_user(
.get_data(DataType::ResetToken)
.as_str()
.map(str::to_owned);
let Some((user_id, username, public_key, reset_token)) = user_id
let registration_token = value
.get_data(DataType::RegisterId)
.as_str()
.filter(|token| uuid::Uuid::parse_str(token).is_ok())
.map(str::to_owned);
let Some((user_id, username, public_key, reset_token, registration_token)) = user_id
.zip(username)
.zip(public_key)
.zip(reset_token)
.map(|(((id, name), key), token)| (id, name, key, token))
.zip(registration_token)
.map(|((((id, name), key), token), registration_token)| {
(id, name, key, token, registration_token)
})
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.await;
};
// Omikron supplies the authenticated Iota ID in the payload. The lease
// check below binds completion to that Iota rather than trusting sender.
let iota_id = value
.get_data(DataType::IotaId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| user_repo::valid_protocol_id(*id));
let Some(iota_id) = iota_id else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.await;
};
match user_repo::register_complete_user(
UserId::from(user_id),
username,
public_key,
IotaId::from(value.get_sender() as i64),
IotaId::from(iota_id),
reset_token,
registration_token,
)
.await
{

94
src/transport/omikron_connection.rs
View file

@ -91,10 +91,29 @@ impl OmikronConnection {
.retain(|_, task| task.inserted_at.elapsed() < MAX_WAITING_AGE);
}
}));
while let Ok(value) = receiver.receive().await {
if let Err(error) = self.clone().process_message(value).await {
log_err!(0, PrintType::Omega, "Error processing message: {}", error);
if matches!(error, crate::error::OmegaError::NotConnected) {
loop {
match receiver.receive().await {
Ok(value) => {
if let Err(error) = self.clone().process_message(value).await {
log_err!(
self.id as i64,
PrintType::Omega,
"Error processing Omikron message: {}",
error
);
if matches!(error, crate::error::OmegaError::NotConnected) {
break;
}
}
}
Err(error) => {
log_err!(
self.id as i64,
PrintType::Omega,
"Omikron receive loop ended: {}; transport close reason: {:?}",
error,
receiver.close_reason()
);
break;
}
}
@ -133,6 +152,9 @@ impl OmikronConnection {
Some(CommunicationType::UserDisconnected) => {
crate::transport::handlers::presence::user_disconnected(self, value, id).await
}
Some(CommunicationType::ClientChanged) => {
crate::transport::handlers::presence::client_changed(self, value, id).await
}
Some(CommunicationType::IotaConnected) => {
crate::transport::handlers::presence::iota_connected(self, value, id).await
}
@ -259,40 +281,38 @@ pub async fn complete_register(_: PublicKeyBundle, _: Option<String>) -> u64 {
pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
let cert_pem = load_file_vec("certs", "cert.pem")?;
let key_pem = load_file_vec("certs", "key.pem")?;
let web_config = server::server::build_web_config()?;
let host_config = HostConfig::new(
IpAddr::from(Ipv4Addr::new(0, 0, 0, 0)),
port,
cert_pem,
key_pem,
)
.with_policy(Policy {
send_mode: SendMode::SingleStreamPerMessage,
max_message_size: 1_000_000_000,
handshake_max_message_size: 1_000_000,
close_frame_len: u32::MAX,
application_close_code: 0,
open_stream_timeout: Duration::from_millis(2_000),
write_timeout: Duration::from_millis(2_000),
accept_stream_timeout: Duration::from_millis(10_000),
read_timeout: Duration::from_millis(30_000),
keep_alive_interval: Some(Duration::from_secs(6)),
max_idle_timeout: Some(Duration::from_secs(30)),
force_close_delay: Duration::from_millis(300),
receiver_queue_capacity: 1000,
max_concurrent_stream_tasks: 10,
persistent_stream_max_retries: 5,
persistent_stream_retry_backoff: Duration::from_secs(5),
max_frames_per_stream: None,
})
.with_authentication(
load_keyring(),
Box::new(|id, description| Box::pin(get_by_omikron_id(id, description))),
Box::new(|key, description| Box::pin(complete_register(key, description))),
)
.with_authentication_policy(AuthenticationPolicy::ForceAuthentication);
let web_config = server::server::build_web_config()?
.serve_tcp_https(true)
.max_tcp_connections(256);
let ip = IpAddr::from(Ipv4Addr::new(0, 0, 0, 0));
let host_config = HostConfig::new(ip, port, cert_pem, key_pem)
.with_policy(Policy {
send_mode: SendMode::SingleStreamPerMessage,
max_message_size: 1_000_000_000,
handshake_max_message_size: 1_000_000,
close_frame_len: u32::MAX,
application_close_code: 0,
open_stream_timeout: Duration::from_millis(5_000),
write_timeout: Duration::from_millis(5_000),
accept_stream_timeout: Duration::from_millis(10_000),
read_timeout: Duration::from_millis(30_000),
keep_alive_interval: Some(Duration::from_secs(6)),
max_idle_timeout: Some(Duration::from_secs(30)),
force_close_delay: Duration::from_millis(300),
receiver_queue_capacity: 1000,
max_concurrent_stream_tasks: 64,
persistent_stream_max_retries: 5,
persistent_stream_retry_backoff: Duration::from_secs(5),
max_frames_per_stream: None,
})
.with_authentication(
load_keyring(),
Box::new(|id, description| Box::pin(get_by_omikron_id(id, description))),
Box::new(|key, description| Box::pin(complete_register(key, description))),
)
.with_authentication_policy(AuthenticationPolicy::ForceAuthentication);
let mut server = MTPWebServer::new(host_config, web_config).await?;
log!("OmegaServer listening on port {}", port);
log!("OmegaServer listening on {}:{}", ip.to_string(), port);
loop {
let mut conn = match server.accept().await {
Ok(Some(conn)) => conn,

10
src/transport/omikron_manager.rs
View file

@ -26,12 +26,18 @@ pub async fn remove_omikron(omikron_id: i64) {
OMIKRON_CONNECTIONS.remove(&omikron_id);
}
pub fn get_connected_omikron(omikron_id: i64) -> Option<Arc<OmikronConnection>> {
OMIKRON_CONNECTIONS
.get(&omikron_id)
.map(|connection| connection.clone())
}
pub async fn get_random_omikron() -> Result<Arc<OmikronConnection>, ()> {
let keys: Vec<_> = OMIKRON_CONNECTIONS.iter().map(|e| *e.key()).collect();
if let Some(key) = keys.into_iter().choose(&mut rand::rng()) {
if let Some(entry) = OMIKRON_CONNECTIONS.get(&key) {
return Ok(entry.clone());
if let Some(connection) = get_connected_omikron(key) {
return Ok(connection);
}
}