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

[Updt] Mtp 0.3.0

Browse source
This commit is contained in:
Alex 2026-08-20 17:05:37 +02:00
commit b3441a8902
33 changed files with 1480 additions and 1531 deletions
Download patch file Download diff file Expand all files Collapse all files

154
src/transport/connection.rs
View file

@ -2,24 +2,94 @@ pub(crate) use super::omikron_connection::{OmikronConnection, OmikronResult};
use mtp::codec::{CommunicationValue, DataValue};
pub(crate) trait MtpValueCompat {
fn get_id(&self) -> u32;
fn get_sender(&self) -> u64;
fn get_receiver(&self) -> u64;
pub(crate) trait RequiredMtpFields {
/// Require a nonzero request correlation ID. `None` and `Some(0)` are
/// distinct MTP wire states, but both are invalid for Omega requests.
fn require_id(&self) -> OmikronResult<u32>;
/// Require a nonzero authenticated peer identity.
fn require_sender(&self) -> OmikronResult<u64>;
/// Require a nonzero application routing target.
fn require_receiver(&self) -> OmikronResult<u64>;
fn require_sender_i64(&self) -> OmikronResult<i64>;
fn require_receiver_i64(&self) -> OmikronResult<i64>;
}
impl MtpValueCompat for CommunicationValue {
fn get_id(&self) -> u32 {
self.id().unwrap_or_default()
impl RequiredMtpFields for CommunicationValue {
fn require_id(&self) -> OmikronResult<u32> {
self.id().filter(|id| *id != 0).ok_or_else(|| {
crate::OmegaError::Validation("MTP message is missing request id".into())
})
}
fn get_sender(&self) -> u64 {
self.sender().unwrap_or_default()
fn require_sender(&self) -> OmikronResult<u64> {
self.sender()
.filter(|sender| *sender != 0)
.ok_or_else(|| crate::OmegaError::Validation("MTP message is missing sender".into()))
}
fn get_receiver(&self) -> u64 {
self.receiver().unwrap_or_default()
fn require_receiver(&self) -> OmikronResult<u64> {
self.receiver()
.filter(|receiver| *receiver != 0)
.ok_or_else(|| crate::OmegaError::Validation("MTP message is missing receiver".into()))
}
fn require_sender_i64(&self) -> OmikronResult<i64> {
positive_i64(self.require_sender()?, "sender")
}
fn require_receiver_i64(&self) -> OmikronResult<i64> {
positive_i64(self.require_receiver()?, "receiver")
}
}
pub(crate) fn positive_i64(value: impl TryInto<i128>, field: &str) -> OmikronResult<i64> {
let value = value
.try_into()
.map_err(|_| crate::OmegaError::Validation(format!("invalid {field}")))?;
let value = i64::try_from(value)
.map_err(|_| crate::OmegaError::Validation(format!("invalid {field}")))?;
if value <= 0 {
return Err(crate::OmegaError::Validation(format!("invalid {field}")));
}
Ok(value)
}
pub(crate) fn validate_dispatch_fields(value: &CommunicationValue) -> OmikronResult<()> {
let Some(message_type) = value.get_comm_type_enum() else {
return Err(crate::OmegaError::Validation(
"MTP message has an unknown communication type".into(),
));
};
if !matches!(
message_type,
mtp::codec::CommunicationType::ClientChanged
| mtp::codec::CommunicationType::PushNotification
) {
value.require_id()?;
}
if matches!(
message_type,
mtp::codec::CommunicationType::GetUserData
| mtp::codec::CommunicationType::ChangeUserData
| mtp::codec::CommunicationType::ChangeIotaData
| mtp::codec::CommunicationType::DeleteUser
| mtp::codec::CommunicationType::AttachUserBegin
| mtp::codec::CommunicationType::AttachUserComplete
| mtp::codec::CommunicationType::DeleteUserCredentialBegin
| mtp::codec::CommunicationType::DeleteUserCredentialComplete
| mtp::codec::CommunicationType::EraseHostedUserDataAck
| mtp::codec::CommunicationType::ReleaseUserFromIota
| mtp::codec::CommunicationType::DeleteIota
| mtp::codec::CommunicationType::GetNotifications
| mtp::codec::CommunicationType::ReadNotification
| mtp::codec::CommunicationType::StateSubscribe
) {
value.require_sender()?;
}
Ok(())
}
pub(crate) trait OptionalDataValueCompat {
@ -46,3 +116,65 @@ impl OptionalDataValueCompat for Option<&DataValue> {
self.and_then(|value| value.as_bytes())
}
}
#[cfg(test)]
mod tests {
use super::{RequiredMtpFields, validate_dispatch_fields};
use mtp::codec::{CommunicationType, CommunicationValue};
#[test]
fn missing_request_id_is_rejected_without_a_zero_fallback() {
let value = CommunicationValue::new(CommunicationType::GetUserData).without_id();
assert!(value.require_id().is_err());
assert!(validate_dispatch_fields(&value).is_err());
}
#[test]
fn zero_routing_values_are_rejected_by_omega_contract() {
let value = CommunicationValue::new(CommunicationType::GetUserData)
.with_id(0)
.with_sender(0)
.with_receiver(0);
assert!(value.require_id().is_err());
assert!(value.require_sender().is_err());
assert!(value.require_receiver().is_err());
assert!(validate_dispatch_fields(&value).is_err());
}
#[test]
fn security_sensitive_messages_require_sender() {
for message_type in [
CommunicationType::GetUserData,
CommunicationType::ChangeUserData,
CommunicationType::ChangeIotaData,
CommunicationType::DeleteUser,
CommunicationType::AttachUserBegin,
CommunicationType::AttachUserComplete,
CommunicationType::DeleteUserCredentialBegin,
CommunicationType::DeleteUserCredentialComplete,
CommunicationType::EraseHostedUserDataAck,
CommunicationType::ReleaseUserFromIota,
CommunicationType::DeleteIota,
CommunicationType::GetNotifications,
CommunicationType::ReadNotification,
CommunicationType::StateSubscribe,
] {
let value = CommunicationValue::new(message_type).with_id(1);
assert!(value.require_sender().is_err());
assert!(validate_dispatch_fields(&value).is_err());
}
}
#[test]
fn optional_client_changed_id_remains_optional() {
let value = CommunicationValue::new(CommunicationType::ClientChanged).without_id();
assert!(validate_dispatch_fields(&value).is_ok());
}
#[test]
fn push_notification_uses_its_logical_sender_field() {
let value = CommunicationValue::new(CommunicationType::PushNotification);
assert!(validate_dispatch_fields(&value).is_ok());
}
}

161
src/transport/handlers/account.rs
View file

@ -1,5 +1,5 @@
use super::super::connection::{
MtpValueCompat, OmikronConnection, OmikronResult, OptionalDataValueCompat,
OmikronConnection, OmikronResult, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::{
db::{iota_repo, user_repo},
@ -22,24 +22,30 @@ async fn delete(
Err(error) => CommunicationValue::new(CommunicationType::ErrorInternal)
.add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string())),
};
connection.send(&response.with_id(value.get_id())).await
connection
.send(&response.with_id(value.require_id()?))
.await
}
pub async fn user(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let user_id = UserId::from(value.get_sender() as i64);
value.require_id()?;
value.require_sender_i64()?;
let user_id = UserId::from(value.require_sender_i64()?);
complete_delete(connection, value, user_id).await
}
pub async fn iota(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
value.require_id()?;
value.require_sender_i64()?;
delete(
connection,
value.clone(),
iota_repo::delete_iota(IotaId::from(value.get_sender() as i64)),
iota_repo::delete_iota(IotaId::from(value.require_sender_i64()?)),
)
.await
}
@ -48,6 +54,8 @@ pub async fn release_from_iota(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
value.require_id()?;
value.require_sender_i64()?;
let Some(user_id) = value
.get_data(DataType::UserId)
.as_signed_number()
@ -55,18 +63,21 @@ pub async fn release_from_iota(
.filter(|id| *id > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidUserId)
.await;
};
let requester = IotaId::from(value.get_sender() as i64);
let requester = IotaId::from(value.require_sender_i64()?);
let Ok(user) = user_repo::get_by_user_id(UserId::from(user_id)).await else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(value.require_id()?, CommunicationType::ErrorNotFound)
.await;
};
if user.iota_id != Some(requester) {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotAuthenticated)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorNotAuthenticated,
)
.await;
}
let previous_iota = user.iota_id;
@ -76,14 +87,17 @@ pub async fn release_from_iota(
crate::transport::omikron_manager::publish_iota_user_snapshot(iota.0).await;
}
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.send(
&CommunicationValue::new(CommunicationType::Success)
.with_id(value.require_id()?),
)
.await
}
Err(error) => {
connection
.send(
&CommunicationValue::new(CommunicationType::ErrorInternal)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string())),
)
.await
@ -104,6 +118,8 @@ pub async fn attach_begin(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
value.require_id()?;
value.require_sender_i64()?;
let Some(user_id) = value
.get_data(DataType::UserId)
.as_signed_number()
@ -111,7 +127,7 @@ pub async fn attach_begin(
.filter(|v| *v > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidUserId)
.await;
};
if user_repo::get_by_user_id(UserId::from(user_id))
@ -119,10 +135,10 @@ pub async fn attach_begin(
.is_err()
{
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(value.require_id()?, CommunicationType::ErrorNotFound)
.await;
}
let requester = value.get_sender() as i64;
let requester = value.require_sender_i64()?;
let nonce =
connection
.state()
@ -130,7 +146,7 @@ pub async fn attach_begin(
connection
.send(
&CommunicationValue::new(CommunicationType::AttachUserChallenge)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()))
.add_typed_default(DataType::ServerNonce, DataValue::SignedNumber(nonce.into())),
)
@ -141,6 +157,8 @@ pub async fn attach_complete(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
value.require_id()?;
value.require_sender_i64()?;
let Some(user_id) = value
.get_data(DataType::UserId)
.as_signed_number()
@ -148,24 +166,30 @@ pub async fn attach_complete(
.filter(|v| *v > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidUserId)
.await;
};
let requester = value.get_sender() as i64;
let requester = value.require_sender_i64()?;
let Some(nonce) = value
.get_data(DataType::ServerNonce)
.as_signed_number()
.and_then(|v| u64::try_from(v).ok())
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorInvalidChallenge,
)
.await;
};
let signature = value.get_data(DataType::Signature).as_bytes();
let pq_signature = value.get_data(DataType::PqSignature).as_bytes();
let (Some(signature), Some(pq_signature)) = (signature, pq_signature) else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorInvalidChallenge,
)
.await;
};
if !connection.state().consume_challenge(
@ -175,12 +199,15 @@ pub async fn attach_complete(
nonce,
) {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorInvalidChallenge,
)
.await;
}
let Ok(user) = user_repo::get_by_user_id(UserId::from(user_id)).await else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(value.require_id()?, CommunicationType::ErrorNotFound)
.await;
};
let payload = lifecycle_payload(b"tensamin:user-attach:v1\0", user_id, requester, nonce);
@ -188,7 +215,10 @@ pub async fn attach_complete(
|| verify_ml_dsa(&user.public_key.sig_pq_public_key, &payload, &pq_signature).is_err()
{
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotAuthenticated)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorNotAuthenticated,
)
.await;
}
let previous_iota = user.iota_id;
@ -199,12 +229,15 @@ pub async fn attach_complete(
}
crate::transport::omikron_manager::publish_iota_user_snapshot(requester).await;
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.send(
&CommunicationValue::new(CommunicationType::Success)
.with_id(value.require_id()?),
)
.await
}
Err(_) => {
connection
.send_error_response(value.get_id(), CommunicationType::ErrorInternal)
.send_error_response(value.require_id()?, CommunicationType::ErrorInternal)
.await
}
}
@ -225,7 +258,7 @@ async fn complete_delete(
connection
.send(
&CommunicationValue::new(CommunicationType::Success)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(
DataType::CleanupPending,
DataValue::Bool(cleanup_pending),
@ -235,14 +268,14 @@ async fn complete_delete(
}
Err(crate::error::OmegaError::NotFound) => {
connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(value.require_id()?, CommunicationType::ErrorNotFound)
.await
}
Err(error) => {
connection
.send(
&CommunicationValue::new(CommunicationType::ErrorInternal)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string())),
)
.await
@ -254,6 +287,8 @@ pub async fn delete_credential_begin(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
value.require_id()?;
value.require_sender_i64()?;
let Some(user_id) = value
.get_data(DataType::UserId)
.as_signed_number()
@ -261,7 +296,7 @@ pub async fn delete_credential_begin(
.filter(|v| *v > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidUserId)
.await;
};
if user_repo::get_by_user_id(UserId::from(user_id))
@ -269,10 +304,10 @@ pub async fn delete_credential_begin(
.is_err()
{
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(value.require_id()?, CommunicationType::ErrorNotFound)
.await;
}
let requester = value.get_sender() as i64;
let requester = value.require_sender_i64()?;
let nonce =
connection
.state()
@ -280,7 +315,7 @@ pub async fn delete_credential_begin(
connection
.send(
&CommunicationValue::new(CommunicationType::DeleteUserCredentialChallenge)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()))
.add_typed_default(DataType::ServerNonce, DataValue::SignedNumber(nonce.into())),
)
@ -291,6 +326,8 @@ pub async fn delete_credential_complete(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
value.require_id()?;
value.require_sender_i64()?;
let Some(user_id) = value
.get_data(DataType::UserId)
.as_signed_number()
@ -298,17 +335,20 @@ pub async fn delete_credential_complete(
.filter(|v| *v > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidUserId)
.await;
};
let requester = value.get_sender() as i64;
let requester = value.require_sender_i64()?;
let Some(nonce) = value
.get_data(DataType::ServerNonce)
.as_signed_number()
.and_then(|v| u64::try_from(v).ok())
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorInvalidChallenge,
)
.await;
};
let (Some(signature), Some(pq_signature)) = (
@ -316,7 +356,10 @@ pub async fn delete_credential_complete(
value.get_data(DataType::PqSignature).as_bytes(),
) else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorInvalidChallenge,
)
.await;
};
if !connection.state().consume_challenge(
@ -326,12 +369,15 @@ pub async fn delete_credential_complete(
nonce,
) {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorInvalidChallenge,
)
.await;
}
let Ok(user) = user_repo::get_by_user_id(UserId::from(user_id)).await else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(value.require_id()?, CommunicationType::ErrorNotFound)
.await;
};
let payload = lifecycle_payload(b"tensamin:user-delete:v1\0", user_id, requester, nonce);
@ -339,7 +385,10 @@ pub async fn delete_credential_complete(
|| verify_ml_dsa(&user.public_key.sig_pq_public_key, &payload, &pq_signature).is_err()
{
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotAuthenticated)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorNotAuthenticated,
)
.await;
}
complete_delete(connection, value, user.id).await
@ -349,6 +398,8 @@ pub async fn erase_hosted_user_data_ack(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
value.require_id()?;
value.require_sender_i64()?;
let Some(user_id) = value
.get_data(DataType::UserId)
.as_signed_number()
@ -356,45 +407,31 @@ pub async fn erase_hosted_user_data_ack(
.filter(|v| *v > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidUserId)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidUserId)
.await;
};
let iota_id = IotaId::from(value.get_sender() as i64);
let iota_id = IotaId::from(value.require_sender_i64()?);
match user_repo::acknowledge_pending_erasure(UserId::from(user_id), iota_id).await {
Ok(true) => {
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.send(
&CommunicationValue::new(CommunicationType::Success)
.with_id(value.require_id()?),
)
.await
}
Ok(false) => {
connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotAuthenticated)
.send_error_response(
value.require_id()?,
CommunicationType::ErrorNotAuthenticated,
)
.await
}
Err(_) => {
connection
.send_error_response(value.get_id(), CommunicationType::ErrorInternal)
.send_error_response(value.require_id()?, CommunicationType::ErrorInternal)
.await
}
}
}
/// New lifecycle operation names are intentionally fail-closed until their
/// proof and durable-erasure handlers are enabled. This explicit dispatch
/// prevents either a bare Iota request or the legacy DeleteUser path from
/// acquiring account-deletion authority during a staged rollout.
pub async fn lifecycle_unavailable(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
connection
.send(
&CommunicationValue::new(CommunicationType::ErrorNotAuthenticated)
.with_id(value.get_id())
.add_typed_default(
DataType::ErrorType,
DataValue::Str("user lifecycle proof handler is not enabled".into()),
),
)
.await
}

4
src/transport/handlers/links.rs
View file

@ -1,5 +1,5 @@
use super::super::connection::{
MtpValueCompat, OmikronConnection, OmikronResult, OptionalDataValueCompat,
OmikronConnection, OmikronResult, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::server::short_link::add_short_link;
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
@ -17,7 +17,7 @@ pub async fn shorten(
.await
.map_err(|_| crate::error::OmegaError::Transport("short link error".to_string()))?;
let response = CommunicationValue::new(CommunicationType::ShortenLink)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(DataType::Link, DataValue::Str(short));
connection.send(&response).await
}

139
src/transport/handlers/notifications.rs
View file

@ -1,5 +1,5 @@
use super::super::connection::{
MtpValueCompat, OmikronConnection, OmikronResult, OptionalDataValueCompat,
OmikronConnection, OmikronResult, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::{db::notification_repo, log, models::UserId};
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
@ -10,35 +10,36 @@ pub async fn get(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let notifications =
match notification_repo::get_notifications(UserId::from(value.get_sender() as i64)).await {
Ok(items) => items
.into_iter()
.map(|item| {
let tm = TypeMap::latest();
let Some(sender) = DataType::SenderId.try_to_id(&tm) else {
return DataValue::Container(Vec::new());
};
let Some(amount) = DataType::Amount.try_to_id(&tm) else {
return DataValue::Container(Vec::new());
};
DataValue::Container(vec![
(sender, DataValue::SignedNumber(item.sender_id.0.into())),
(amount, DataValue::SignedNumber(item.amount.into())),
])
})
.collect(),
Err(error) => {
log!(
crate::util::logger::PrintType::General,
"SQL get_notifications error: {}",
error
);
Vec::new()
}
};
let request_id = value.require_id()?;
let sender = value.require_sender_i64()?;
let notifications = match notification_repo::get_notifications(UserId::from(sender)).await {
Ok(items) => items
.into_iter()
.map(|item| {
let tm = TypeMap::latest();
let Some(sender) = DataType::SenderId.try_to_id(&tm) else {
return DataValue::Container(Vec::new());
};
let Some(amount) = DataType::Amount.try_to_id(&tm) else {
return DataValue::Container(Vec::new());
};
DataValue::Container(vec![
(sender, DataValue::SignedNumber(item.sender_id.0.into())),
(amount, DataValue::SignedNumber(item.amount.into())),
])
})
.collect(),
Err(error) => {
log!(
crate::util::logger::PrintType::General,
"SQL get_notifications error: {}",
error
);
Vec::new()
}
};
let response = CommunicationValue::new(CommunicationType::GetNotifications)
.with_id(value.get_id())
.with_id(request_id)
.add_typed_default(DataType::Notifications, DataValue::Array(notifications));
connection.send(&response).await
}
@ -47,19 +48,17 @@ pub async fn read(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let receiver = match value.get_sender() {
sender if sender > 0 => sender as i64,
_ => match value.get_data(DataType::ReceiverId).as_number() {
Some(id) => id as i64,
None => return Ok(()),
},
};
let request_id = value.require_id()?;
let receiver = value.require_sender_i64()?;
let Some(other) = value
.get_data(DataType::SenderId)
.as_number()
.map(|id| id as i64)
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
else {
return Ok(());
return connection
.send_error_response(request_id, CommunicationType::ErrorInvalidData)
.await;
};
if let Err(error) =
notification_repo::read_notification(UserId::from(receiver), UserId::from(other)).await
@ -71,7 +70,7 @@ pub async fn read(
);
} else {
let response =
CommunicationValue::new(CommunicationType::ReadNotification).with_id(value.get_id());
CommunicationValue::new(CommunicationType::ReadNotification).with_id(request_id);
let _ = connection.send(&response).await;
let sync = CommunicationValue::new(CommunicationType::ReadNotification)
.with_receiver(receiver as u64)
@ -85,18 +84,54 @@ pub async fn push(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let receiver = match value.get_receiver() {
receiver if receiver > 0 => receiver as i64,
_ => match value.get_data(DataType::ReceiverId).as_number() {
Some(id) => id as i64,
None => return Ok(()),
},
let request_id = value.id().filter(|id| *id != 0);
let receiver = value.require_receiver_i64().ok().or_else(|| {
value
.get_data(DataType::ReceiverId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
});
let Some(receiver) = receiver else {
if let Some(request_id) = request_id {
return connection
.send_error_response(request_id, CommunicationType::ErrorInvalidData)
.await;
}
return Ok(());
};
let sender = value
let Some(sender) = value
.get_data(DataType::SenderId)
.as_number()
.map(|id| id as i64)
.unwrap_or(value.get_sender() as i64);
.and_then(|sender| i64::try_from(sender).ok())
.filter(|sender| *sender > 0)
else {
if let Some(request_id) = request_id {
return connection
.send_error_response(request_id, CommunicationType::ErrorInvalidData)
.await;
}
return Ok(());
};
let source_omikron = connection
.clone()
.get_omikron_id()
.await
.ok_or(crate::OmegaError::NotConnected)?;
if !connection
.state()
.presence
.routes_for_user(sender)
.iter()
.any(|route| route.omikron_id == source_omikron)
{
if let Some(request_id) = request_id {
return connection
.send_error_response(request_id, CommunicationType::ErrorNotAuthenticated)
.await;
}
return Ok(());
}
if let Err(error) =
notification_repo::add_notification(UserId::from(receiver), UserId::from(sender)).await
{
@ -106,9 +141,11 @@ pub async fn push(
error
);
} else {
let response =
CommunicationValue::new(CommunicationType::PushNotification).with_id(value.get_id());
let _ = connection.send(&response).await;
if let Some(request_id) = request_id {
let response =
CommunicationValue::new(CommunicationType::PushNotification).with_id(request_id);
let _ = connection.send(&response).await;
}
let push = CommunicationValue::new(CommunicationType::PushNotification)
.with_receiver(receiver as u64)
.add_typed_default(DataType::SenderId, DataValue::SignedNumber(sender.into()));

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

@ -1,5 +1,5 @@
use super::super::connection::{
MtpValueCompat, OmikronConnection, OmikronResult, OptionalDataValueCompat,
OmikronConnection, OmikronResult, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::{
db::user_repo, log_in, models::IotaId, sql::connection_status::UserStatus, state::OmegaState,
@ -11,7 +11,10 @@ use std::{
};
fn parse_subscription(value: &CommunicationValue) -> Result<(i64, i64, Vec<i64>), &'static str> {
let user_id = i64::try_from(value.get_sender())
let Some(sender) = value.sender() else {
return Err("user_id");
};
let user_id = i64::try_from(sender)
.ok()
.filter(|id| *id > 0)
.ok_or("user_id")?;
@ -51,13 +54,11 @@ fn states_for_users(state: &OmegaState, users: &[crate::models::User]) -> HashMa
users
.iter()
.map(|user| {
(
user.id.0,
state.presence.resolve_public_state(
user.id.0,
user.iota_id.map(|id| id.0).unwrap_or_default(),
),
)
let status = user
.iota_id
.map(|iota_id| state.presence.resolve_public_state(user.id.0, iota_id.0))
.unwrap_or(UserStatus::user_offline);
(user.id.0, status)
})
.collect()
}
@ -70,9 +71,10 @@ fn changed_states(
let mut changes = users
.iter()
.filter_map(|user| {
let after = state
.presence
.resolve_public_state(user.id.0, user.iota_id.map(|id| id.0).unwrap_or_default());
let after = user
.iota_id
.map(|iota_id| state.presence.resolve_public_state(user.id.0, iota_id.0))
.unwrap_or(UserStatus::user_offline);
(before.get(&user.id.0) != Some(&after)).then_some((user.id.0, after))
})
.collect::<Vec<_>>();
@ -96,21 +98,6 @@ fn state_notification(
.add_typed_default(DataType::UserState, DataValue::Str(user_state.to_string()))
}
fn private_state_notification(
user_id: i64,
session_id: i64,
user_state: &UserStatus,
) -> CommunicationValue {
CommunicationValue::new(CommunicationType::ClientChanged)
.with_receiver(user_id as u64)
.add_typed_default(
DataType::SessionId,
DataValue::SignedNumber(session_id.into()),
)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()))
.add_typed_default(DataType::UserState, DataValue::Str(user_state.to_string()))
}
async fn publish_state_changes(state: &OmegaState, changes: &[(i64, UserStatus)]) {
let mut grouped = BTreeMap::<i64, Vec<CommunicationValue>>::new();
for (user_id, user_state) in changes {
@ -143,28 +130,6 @@ async fn publish_changed_states(
publish_state_changes(state, &changed_states(state, before, users)).await;
}
async fn publish_private_state(state: &OmegaState, user_id: i64, user_state: &UserStatus) {
let mut grouped = BTreeMap::<i64, Vec<CommunicationValue>>::new();
for (session_id, route) in state.presence.sessions_for_user(user_id) {
grouped
.entry(route.omikron_id)
.or_default()
.push(private_state_notification(user_id, session_id, user_state));
}
for (omikron_id, notifications) in grouped {
if let Err(error) =
crate::transport::omikron_manager::send_state_batch(omikron_id, notifications).await
{
log_in!(
crate::util::logger::PrintType::General,
"Failed to deliver private presence state batch to Omikron {}: {}",
omikron_id,
error
);
}
}
}
pub async fn state_subscribe(
state: Arc<OmegaState>,
connection: Arc<OmikronConnection>,
@ -175,13 +140,13 @@ pub async fn state_subscribe(
Ok(subscription) => subscription,
Err("user_id") => {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNoUserId)
.send_error_response(value.require_id()?, CommunicationType::ErrorNoUserId)
.await;
}
Err(detail) => {
return connection
.send_error_response_with_detail(
value.get_id(),
value.require_id()?,
CommunicationType::ErrorInvalidData,
detail,
)
@ -190,14 +155,14 @@ pub async fn state_subscribe(
};
if !state.presence.owns_session(user_id, session_id, omikron_id) {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNoIota)
.send_error_response(value.require_id()?, CommunicationType::ErrorNoIota)
.await;
}
state
.presence
.replace_subscription(user_id, session_id, omikron_id, user_ids);
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.require_id()?))
.await
}
@ -242,7 +207,7 @@ pub async fn user_connected(
.filter(|id| *id > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let Some(session_id) = value
@ -252,7 +217,7 @@ pub async fn user_connected(
.filter(|id| *id > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let Some(iota_id) = value
@ -262,25 +227,25 @@ pub async fn user_connected(
.filter(|id| *id > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let user = match user_repo::get_by_user_id(user_id.into()).await {
Ok(user) => user,
Err(crate::error::OmegaError::NotFound) => {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(value.require_id()?, CommunicationType::ErrorNotFound)
.await;
}
Err(error) => return Err(error.into()),
Err(error) => return Err(error),
};
let preferences = match user_repo::get_presence_preferences(&[user_id]).await {
Ok(preferences) => preferences,
Err(error) => return Err(error.into()),
Err(error) => return Err(error),
};
if user.iota_id.map(|id| id.0) != Some(iota_id) || !state.presence.has_iota_route(iota_id) {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNoIota)
.send_error_response(value.require_id()?, CommunicationType::ErrorNoIota)
.await;
}
apply_preferences(&state, preferences);
@ -291,7 +256,7 @@ pub async fn user_connected(
.track_session(user_id, session_id, omikron_id, iota_id);
publish_changed_states(&state, &before, &users).await;
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.require_id()?))
.await
}
@ -309,7 +274,7 @@ pub async fn user_disconnected(
.filter(|id| *id > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let Some(session_id) = value
@ -319,7 +284,7 @@ pub async fn user_disconnected(
.filter(|id| *id > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
if let Ok(user) = user_repo::get_by_user_id(user_id.into()).await {
@ -337,137 +302,7 @@ pub async fn user_disconnected(
.remove_session(user_id, session_id, omikron_id);
}
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.await
}
pub async fn set_user_state(
state: Arc<OmegaState>,
connection: Arc<OmikronConnection>,
value: CommunicationValue,
omikron_id: i64,
) -> OmikronResult<()> {
let Some(user_id) = i64::try_from(value.get_sender()).ok().filter(|id| *id > 0) else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNoUserId)
.await;
};
if let Some(requested_user) = value.get_data(DataType::UserId) {
let Some(requested_user_id) = requested_user
.as_number()
.and_then(|id| i64::try_from(id).ok())
else {
return connection
.send_error_response_with_detail(
value.get_id(),
CommunicationType::ErrorInvalidData,
"user_id",
)
.await;
};
if requested_user_id != user_id {
return connection
.send_error_response_with_detail(
value.get_id(),
CommunicationType::ErrorInvalidData,
"user_id",
)
.await;
}
}
let Some(iota_id) = value
.get_data(DataType::IotaId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
else {
return connection
.send_error_response_with_detail(
value.get_id(),
CommunicationType::ErrorInvalidData,
"iota_id",
)
.await;
};
let Some(requested_state) = value
.get_data(DataType::UserState)
.as_str()
.and_then(UserStatus::from_client_preference)
else {
return connection
.send_error_response_with_detail(
value.get_id(),
CommunicationType::ErrorInvalidData,
"user_state",
)
.await;
};
if !state.presence.has_iota_route(iota_id) {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNoIota)
.await;
}
let Some(session_id) = value
.get_data(DataType::SessionId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
else {
return connection
.send_error_response_with_detail(
value.get_id(),
CommunicationType::ErrorInvalidData,
"session_id",
)
.await;
};
let Some(route) = state.presence.session_route(user_id, session_id) else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNoIota)
.await;
};
if route.omikron_id != omikron_id || route.iota_id != iota_id {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.await;
}
if !state.presence.has_active_session_for_iota(user_id, iota_id) {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNoIota)
.await;
}
let previous_preference = state.presence.preference(user_id);
let previous_state = state.presence.resolve_public_state(user_id, iota_id);
if let Err(error) =
user_repo::change_presence_preference(user_id.into(), requested_state.to_string()).await
{
log_in!(
crate::util::logger::PrintType::General,
"Failed to persist presence preference: {}",
error
);
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInternal)
.await;
}
state
.presence
.set_preference(user_id, requested_state.clone());
let new_state = state.presence.resolve_public_state(user_id, iota_id);
if requested_state != previous_preference {
publish_private_state(&state, user_id, &requested_state).await;
}
if requested_state != previous_preference && new_state != previous_state {
publish_state_changes(&state, &[(user_id, new_state)]).await;
}
connection
.send(
&CommunicationValue::new(CommunicationType::Success)
.with_id(value.get_id())
.add_typed_default(
DataType::UserState,
DataValue::Str(requested_state.to_string()),
),
)
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.require_id()?))
.await
}
@ -481,10 +316,11 @@ pub async fn iota_connected(
let Some(iota_id) = value
.get_data(DataType::IotaId)
.as_number()
.map(|id| id as i64)
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let users = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await?;
@ -503,7 +339,7 @@ pub async fn iota_connected(
crate::transport::omikron_manager::deliver_pending_erasures(iota_id).await;
publish_changed_states(&state, &before, &users).await;
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.require_id()?))
.await
}
@ -517,10 +353,11 @@ pub async fn iota_disconnected(
let Some(iota_id) = value
.get_data(DataType::IotaId)
.as_number()
.map(|id| id as i64)
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let users = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await?;
@ -530,7 +367,7 @@ pub async fn iota_disconnected(
state.presence.untrack_iota_connection(iota_id, omikron_id);
publish_changed_states(&state, &before, &users).await;
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.require_id()?))
.await
}
@ -540,7 +377,7 @@ pub async fn sync_status(
value: CommunicationValue,
omikron_id: i64,
) -> OmikronResult<()> {
let request_id = value.get_id();
let request_id = value.require_id()?;
let Some(DataValue::Array(iota_values)) = value.get_data(DataType::IotaIds) else {
return connection
.send_error_response(request_id, CommunicationType::ErrorInvalidData)

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

@ -1,5 +1,5 @@
use super::super::connection::{
MtpValueCompat, OmikronConnection, OmikronResult, OptionalDataValueCompat,
OmikronConnection, OmikronResult, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::{
db::{iota_repo, user_repo},
@ -22,13 +22,13 @@ pub async fn get_register(
.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)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let (register_id, registration_token) =
user_repo::allocate_registration(IotaId::from(iota_id), value.get_id()).await?;
user_repo::allocate_registration(IotaId::from(iota_id), value.require_id()?).await?;
let response = CommunicationValue::new(CommunicationType::GetRegister)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(
DataType::UserId,
DataValue::SignedNumber(register_id.0.into()),
@ -47,7 +47,7 @@ pub async fn complete_iota(
.and_then(|key| PublicKeyBundle::from_base64(key).ok());
let Some(public_key) = public_key else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
match iota_repo::create_new_iota(public_key).await {
@ -55,7 +55,7 @@ pub async fn complete_iota(
connection
.send(
&CommunicationValue::new(CommunicationType::CompleteRegisterIota)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(id.0.into())),
)
.await
@ -64,7 +64,7 @@ pub async fn complete_iota(
connection
.send(
&CommunicationValue::new(CommunicationType::ErrorInternal)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string())),
)
.await
@ -108,7 +108,7 @@ pub async fn complete_user(
})
else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
// Omikron supplies the authenticated Iota ID in the payload. The lease
@ -120,7 +120,7 @@ pub async fn complete_user(
.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)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
@ -136,14 +136,17 @@ pub async fn complete_user(
{
Ok(()) => {
connection
.send(&CommunicationValue::new(CommunicationType::Success).with_id(value.get_id()))
.send(
&CommunicationValue::new(CommunicationType::Success)
.with_id(value.require_id()?),
)
.await
}
Err(error) => {
connection
.send(
&CommunicationValue::new(CommunicationType::ErrorInternal)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string())),
)
.await

23
src/transport/handlers/states.rs
View file

@ -1,5 +1,5 @@
use super::super::connection::{
MtpValueCompat, OmikronConnection, OmikronResult, OptionalDataValueCompat,
OmikronConnection, OmikronResult, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::db::user_repo;
use mtp::{
@ -34,7 +34,7 @@ pub async fn get(
let Some(DataValue::Array(ids)) = value.get_data(DataType::UserIds) else {
return send_error(
connection,
value.get_id(),
value.require_id()?,
CommunicationType::ErrorInvalidData,
None,
)
@ -47,7 +47,7 @@ pub async fn get(
if session_id.is_none() && !legacy_peer {
return send_error(
connection,
value.get_id(),
value.require_id()?,
CommunicationType::ErrorInvalidData,
None,
)
@ -60,7 +60,7 @@ pub async fn get(
let DataValue::SignedNumber(id) = id else {
return send_error(
connection,
value.get_id(),
value.require_id()?,
CommunicationType::ErrorInvalidData,
session_id,
)
@ -69,7 +69,7 @@ pub async fn get(
let Ok(user_id) = i64::try_from(*id) else {
return send_error(
connection,
value.get_id(),
value.require_id()?,
CommunicationType::ErrorInvalidData,
session_id,
)
@ -78,7 +78,7 @@ pub async fn get(
if user_id <= 0 {
return send_error(
connection,
value.get_id(),
value.require_id()?,
CommunicationType::ErrorInvalidData,
session_id,
)
@ -95,7 +95,7 @@ pub async fn get(
Err(_) => {
return send_error(
connection,
value.get_id(),
value.require_id()?,
CommunicationType::ErrorInternal,
session_id,
)
@ -110,9 +110,10 @@ pub async fn get(
missing_user_ids.push(user_id);
continue;
};
let status = state
.presence
.resolve_public_state(user_id, user.iota_id.map(|id| id.0).unwrap_or_default())
let status = user
.iota_id
.map(|iota_id| state.presence.resolve_public_state(user_id, iota_id.0))
.unwrap_or(crate::sql::connection_status::UserStatus::user_offline)
.to_string();
let mut map = Vec::new();
if let Some(kind) = DataType::UserId.try_to_id(&tm) {
@ -124,7 +125,7 @@ pub async fn get(
states.push(DataValue::Container(map));
}
let response = CommunicationValue::new(CommunicationType::GetStates)
.with_id(value.get_id())
.with_id(value.require_id()?)
.add_typed_default(DataType::UserStates, DataValue::Array(states));
let response = if let Some(session_id) = session_id {
response.add_typed_default(DataType::SessionId, DataValue::SignedNumber(session_id))

134
src/transport/handlers/user_data.rs
View file

@ -1,5 +1,5 @@
use super::super::connection::{
MtpValueCompat, OmikronConnection, OmikronResult, OptionalDataValueCompat,
OmikronConnection, OmikronResult, OptionalDataValueCompat, RequiredMtpFields,
};
use crate::{
db::{iota_repo, user_repo},
@ -29,11 +29,16 @@ pub async fn get_user(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let request_id = value.require_id()?;
let sender = value.require_sender_i64()?;
let state = connection.state();
let user = if let Some(id) = value.get_data(DataType::UserId).as_number() {
user_repo::get_by_user_id(UserId::from(id as i64))
.await
.ok()
let user = if let Some(id) = value
.get_data(DataType::UserId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
{
user_repo::get_by_user_id(UserId::from(id)).await.ok()
} else if let Some(name) = value.get_data(DataType::Username).as_str() {
user_repo::get_by_username(name).await.ok()
} else {
@ -41,7 +46,7 @@ pub async fn get_user(
};
let Some(user) = user else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(request_id, CommunicationType::ErrorNotFound)
.await;
};
let id = user.id.0;
@ -52,11 +57,11 @@ pub async fn get_user(
.filter(|name| !name.is_empty())
.unwrap_or_else(|| username.clone());
let mut response = CommunicationValue::new(CommunicationType::GetUserData)
.with_id(value.get_id())
.with_id(request_id)
.add_typed_default(DataType::Username, DataValue::Str(username))
.add_typed_default(
DataType::PublicKey,
DataValue::Str(user.public_key.to_base64()),
DataValue::Str(user.public_key.try_to_base64()?),
)
.add_typed_default(DataType::UserId, DataValue::SignedNumber(id.into()))
.add_typed_default(DataType::Display, DataValue::Str(display))
@ -79,7 +84,7 @@ pub async fn get_user(
response.add_typed_default(DataType::Avatar, DataValue::Str(STANDARD.encode(avatar)));
}
let route = state.presence.user_route(id);
let private_request = value.get_sender() as i64 == id;
let private_request = sender == id;
let resolved_status = if private_request {
if !state
.presence
@ -125,13 +130,21 @@ pub async fn get_iota(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let found = if let Some(id) = value.get_data(DataType::IotaId).as_number() {
iota_repo::get_iota_by_id(IotaId::from(id as i64))
let request_id = value.require_id()?;
let found = if let Some(id) = value
.get_data(DataType::IotaId)
.as_number()
.and_then(|id| i64::try_from(id).ok())
.filter(|id| *id > 0)
{
iota_repo::get_iota_by_id(IotaId::from(id))
.await
.ok()
.map(|iota| (iota.id.0, iota.public_key, None, None))
} else if let Some(id) = value.get_data(DataType::UserId).as_number() {
if let Ok(user) = user_repo::get_by_user_id(UserId::from(id as i64)).await {
if let Some(id) = i64::try_from(id).ok().filter(|id| *id > 0)
&& let Ok(user) = user_repo::get_by_user_id(UserId::from(id)).await
{
match user.iota_id {
Some(iota_id) => iota_repo::get_iota_by_id(iota_id)
.await
@ -163,12 +176,12 @@ pub async fn get_iota(
};
let Some((id, key, user_id, username)) = found else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(request_id, CommunicationType::ErrorNotFound)
.await;
};
let mut response = CommunicationValue::new(CommunicationType::GetIotaData)
.with_id(value.get_id())
.add_typed_default(DataType::PublicKey, DataValue::Str(key.to_base64()))
.with_id(request_id)
.add_typed_default(DataType::PublicKey, DataValue::Str(key.try_to_base64()?))
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(id.into()))
.add_typed_default(DataType::OmikronConnections, connections(&connection, id));
if let Some(user_id) = user_id {
@ -185,7 +198,8 @@ async fn update_user(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let id = UserId::from(value.get_sender() as i64);
let request_id = value.require_id()?;
let id = UserId::from(value.require_sender_i64()?);
let mut error = None;
if let Some(name) = value.get_data(DataType::Username).as_str() {
error = user_repo::change_username(id, name.to_owned())
@ -193,56 +207,55 @@ async fn update_user(
.err()
.map(|e| e.to_string());
}
if error.is_none() {
if let Some(name) = value.get_data(DataType::Display).as_str() {
error = user_repo::change_display_name(id, name.to_owned())
.await
.err()
.map(|e| e.to_string());
}
if error.is_none()
&& let Some(name) = value.get_data(DataType::Display).as_str()
{
error = user_repo::change_display_name(id, name.to_owned())
.await
.err()
.map(|e| e.to_string());
}
if error.is_none() {
if let Some(avatar) = value.get_data(DataType::Avatar).as_str() {
error = user_repo::change_avatar(id, avatar.to_owned())
.await
.err()
.map(|e| e.to_string());
}
if error.is_none()
&& let Some(avatar) = value.get_data(DataType::Avatar).as_str()
{
error = user_repo::change_avatar(id, avatar.to_owned())
.await
.err()
.map(|e| e.to_string());
}
if error.is_none() {
if let Some(about) = value.get_data(DataType::About).as_str() {
error = user_repo::change_about(id, about.to_owned())
.await
.err()
.map(|e| e.to_string());
}
if error.is_none()
&& let Some(about) = value.get_data(DataType::About).as_str()
{
error = user_repo::change_about(id, about.to_owned())
.await
.err()
.map(|e| e.to_string());
}
if error.is_none() {
if let Some(status) = value.get_data(DataType::Status).as_str() {
error = user_repo::change_status(id, status.to_owned())
.await
.err()
.map(|e| e.to_string());
}
if error.is_none()
&& let Some(status) = value.get_data(DataType::Status).as_str()
{
error = user_repo::change_status(id, status.to_owned())
.await
.err()
.map(|e| e.to_string());
}
if error.is_none() {
if let Some(key) = value
if error.is_none()
&& let Some(key) = value
.get_data(DataType::PublicKey)
.as_str()
.and_then(|key| PublicKeyBundle::from_base64(key).ok())
{
error = user_repo::change_keys(id, key)
.await
.err()
.map(|e| e.to_string());
}
{
error = user_repo::change_keys(id, key)
.await
.err()
.map(|e| e.to_string());
}
let response = match error {
None => CommunicationValue::new(CommunicationType::Success),
Some(error) => CommunicationValue::new(CommunicationType::ErrorInternal)
.add_typed_default(DataType::ErrorType, DataValue::Str(error)),
};
connection.send(&response.with_id(value.get_id())).await
connection.send(&response.with_id(request_id)).await
}
pub async fn change_user(
@ -256,34 +269,35 @@ pub async fn change_iota(
connection: Arc<OmikronConnection>,
value: CommunicationValue,
) -> OmikronResult<()> {
let request_id = value.require_id()?;
let reset_data = value.get_data(DataType::ResetToken);
let Some(reset) = reset_data.as_str() else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let new_token_data = value.get_data(DataType::NewToken);
let Some(new_token) = new_token_data.as_str() else {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
.send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData)
.await;
};
let user_id = UserId::from(value.get_sender() as i64);
let user_id = UserId::from(value.require_sender_i64()?);
let user = match user_repo::get_by_user_id(user_id).await {
Ok(user) => user,
Err(_) => {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
.send_error_response(request_id, CommunicationType::ErrorNotFound)
.await;
}
};
if user.token != reset {
return connection
.send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge)
.send_error_response(request_id, CommunicationType::ErrorInvalidChallenge)
.await;
}
let result =
match user_repo::change_iota_id(user_id, Some(IotaId::from(value.get_sender() as i64)))
match user_repo::change_iota_id(user_id, Some(IotaId::from(value.require_sender_i64()?)))
.await
{
Ok(()) => user_repo::change_token(user_id, new_token.to_owned()).await,
@ -294,5 +308,5 @@ pub async fn change_iota(
Err(error) => CommunicationValue::new(CommunicationType::ErrorInternal)
.add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string())),
};
connection.send(&response.with_id(value.get_id())).await
connection.send(&response.with_id(request_id)).await
}

344
src/transport/omikron_connection.rs
View file

@ -1,12 +1,13 @@
use super::capabilities::{OmegaCapabilities, PeerCapabilities};
use crate::models::OmikronId;
use crate::{
load_keyring, log, log_cv_in, log_cv_out, log_err, log_in, server,
log, log_cv_in, log_cv_out, log_err, log_in, server,
state::OmegaState,
transport::connection::{RequiredMtpFields, validate_dispatch_fields},
transport::omikron_manager,
util::{file_util::load_file_vec, logger::PrintType},
};
use dashmap::DashMap;
use dashmap::{DashMap, mapref::entry::Entry};
use mtp::{
codec::{CommunicationType, CommunicationValue},
crypto::PublicKeyBundle,
@ -16,19 +17,21 @@ use mtp::{
use std::{
net::{IpAddr, Ipv4Addr},
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
Arc, Mutex as StdMutex,
atomic::{AtomicU32, AtomicUsize, Ordering},
},
time::{Duration, Instant},
};
use tokio::{
sync::{Mutex, mpsc},
sync::{Mutex, Semaphore, mpsc, oneshot},
time::interval,
};
use tokio_util::{sync::CancellationToken, task::TaskTracker};
const CLEANUP_INTERVAL: Duration = Duration::from_secs(30);
const MAX_WAITING_AGE: Duration = Duration::from_secs(60);
static ACTIVE_CONNECTIONS: AtomicUsize = AtomicUsize::new(0);
static NEXT_CORRELATION_ID: AtomicU32 = AtomicU32::new(1);
static ACTIVE_CONNECTIONS_BY_IP: once_cell::sync::Lazy<DashMap<IpAddr, usize>> =
once_cell::sync::Lazy::new(DashMap::new);
@ -47,16 +50,27 @@ impl Drop for ConnectionLimitGuard {
pub type OmikronResult<T> = crate::error::Result<T>;
pub struct WaitingTask {
pub task: Box<dyn Fn(Arc<OmikronConnection>, CommunicationValue) -> bool + Send + Sync>,
pub sender: oneshot::Sender<CommunicationValue>,
pub inserted_at: Instant,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum DispatchClass {
Concurrent,
Ordered,
}
pub struct OmikronConnection {
id: u64,
state: Arc<OmegaState>,
sender: Mutex<Option<WebMtpSender>>,
waiting_tasks: DashMap<u32, WaitingTask>,
cleanup_handle: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
handler_tasks: TaskTracker,
handler_cancel: CancellationToken,
handler_limit: Arc<Semaphore>,
ordered_sender: mpsc::Sender<CommunicationValue>,
ordered_receiver: StdMutex<Option<mpsc::Receiver<CommunicationValue>>>,
peer_capabilities: PeerCapabilities,
}
impl Drop for OmikronConnection {
@ -64,6 +78,9 @@ impl Drop for OmikronConnection {
if let Some(handle) = self.cleanup_handle.lock().unwrap().take() {
handle.abort();
}
self.handler_cancel.cancel();
self.handler_tasks.close();
self.waiting_tasks.clear();
}
}
@ -75,17 +92,24 @@ impl OmikronConnection {
state: Arc<OmegaState>,
authenticated: bool,
) -> Option<Arc<Self>> {
if !authenticated {
if !authenticated || id == 0 || i64::try_from(id).is_err() {
return None;
}
let peer_capabilities =
PeerCapabilities::from_identification_description(description).ok()?;
let handler_concurrency = state.omikron_handler_concurrency;
let (ordered_sender, ordered_receiver) = mpsc::channel(handler_concurrency);
Some(Arc::new(Self {
id,
state,
sender: Mutex::new(Some(sender)),
waiting_tasks: DashMap::new(),
cleanup_handle: std::sync::Mutex::new(None),
handler_tasks: TaskTracker::new(),
handler_cancel: CancellationToken::new(),
handler_limit: Arc::new(Semaphore::new(handler_concurrency)),
ordered_sender,
ordered_receiver: StdMutex::new(Some(ordered_receiver)),
peer_capabilities,
}))
}
@ -94,6 +118,115 @@ impl OmikronConnection {
&self.peer_capabilities
}
fn track_handler<F>(&self, task: F)
where
F: std::future::Future<Output = ()> + Send + 'static,
{
let handle = self.handler_tasks.spawn(task);
let id = self.id as i64;
tokio::spawn(async move {
if let Err(error) = handle.await {
log_err!(
id,
PrintType::Omega,
"Omikron handler task failed: {}",
error
);
}
});
}
fn start_ordered_worker(self: &Arc<Self>) {
let Some(mut receiver) = self.ordered_receiver.lock().unwrap().take() else {
return;
};
let connection = self.clone();
let cancellation = self.handler_cancel.clone();
self.track_handler(async move {
loop {
let value = tokio::select! {
_ = cancellation.cancelled() => break,
value = receiver.recv() => match value {
Some(value) => value,
None => break,
},
};
let global_permit = tokio::select! {
_ = cancellation.cancelled() => break,
permit = connection.state.global_handler_limit.clone().acquire_owned() => {
match permit {
Ok(permit) => permit,
Err(_) => break,
}
}
};
let permit = tokio::select! {
_ = cancellation.cancelled() => break,
permit = connection.handler_limit.clone().acquire_owned() => {
match permit {
Ok(permit) => permit,
Err(_) => break,
}
}
};
let result = tokio::select! {
_ = cancellation.cancelled() => break,
result = connection.clone().dispatch(value) => result,
};
drop(permit);
drop(global_permit);
if let Err(error) = result {
log_err!(
connection.id as i64,
PrintType::Omega,
"Error processing ordered Omikron message: {}",
error
);
}
}
});
}
async fn stop_handlers(&self) {
self.handler_cancel.cancel();
self.handler_tasks.close();
let _ = tokio::time::timeout(Duration::from_secs(1), self.handler_tasks.wait()).await;
self.waiting_tasks.clear();
}
fn correlation_response(value: &CommunicationValue) -> bool {
matches!(
value.get_comm_type_enum(),
Some(CommunicationType::Success)
| Some(CommunicationType::ErrorProtocol)
| Some(CommunicationType::ErrorAnonymous)
| Some(CommunicationType::ErrorInternal)
| Some(CommunicationType::ErrorInvalidData)
| Some(CommunicationType::ErrorInvalidUserId)
| Some(CommunicationType::ErrorInvalidOmikronId)
| Some(CommunicationType::ErrorNotFound)
| Some(CommunicationType::ErrorNotAuthenticated)
| Some(CommunicationType::ErrorNoIota)
| Some(CommunicationType::ErrorInvalidChallenge)
| Some(CommunicationType::ErrorInvalidSecret)
| Some(CommunicationType::ErrorInvalidPrivateKey)
| Some(CommunicationType::ErrorInvalidPublicKey)
| Some(CommunicationType::ErrorNoUserId)
| Some(CommunicationType::ErrorNoCallId)
| Some(CommunicationType::ErrorInvalidCallId)
)
}
fn next_correlation_id(&self) -> Option<u32> {
for _ in 0..u32::MAX {
let id = NEXT_CORRELATION_ID.fetch_add(1, Ordering::Relaxed);
if id != 0 && !self.waiting_tasks.contains_key(&id) {
return Some(id);
}
}
None
}
pub async fn handle(self: Arc<Self>, receiver: &mut WebMtpReceiver) {
log_in!(
self.id as i64,
@ -117,12 +250,16 @@ impl OmikronConnection {
self.clone().cleanup().await;
return;
}
let cleanup_conn = self.clone();
self.start_ordered_worker();
let cleanup_conn = Arc::downgrade(&self);
*self.cleanup_handle.lock().unwrap() = Some(tokio::spawn(async move {
let mut ticker = interval(CLEANUP_INTERVAL);
loop {
ticker.tick().await;
cleanup_conn
let Some(connection) = cleanup_conn.upgrade() else {
break;
};
connection
.waiting_tasks
.retain(|_, task| task.inserted_at.elapsed() < MAX_WAITING_AGE);
}
@ -164,16 +301,81 @@ impl OmikronConnection {
async fn process_message(self: Arc<Self>, value: CommunicationValue) -> OmikronResult<()> {
log_cv_in!(PrintType::Omikron, &value);
if value.is_type(CommunicationType::Relay) {
return self.dispatch(value).await;
}
if let Some(message_id) = value.id()
if Self::correlation_response(&value)
&& let Some(message_id) = value.id()
&& let Some((_, task)) = self.waiting_tasks.remove(&message_id)
{
let _ = (task.task)(self.clone(), value);
let _ = task.sender.send(value);
return Ok(());
}
self.dispatch(value).await
if !value.is_type(CommunicationType::Relay) {
validate_dispatch_fields(&value)?;
}
let dispatch_class = Self::dispatch_class(&value);
if dispatch_class == DispatchClass::Ordered {
tokio::select! {
_ = self.handler_cancel.cancelled() => {
return Err(crate::error::OmegaError::NotConnected);
}
result = self.ordered_sender.send(value) => {
result.map_err(|_| crate::error::OmegaError::NotConnected)?;
}
}
return Ok(());
}
let global_permit = tokio::select! {
_ = self.handler_cancel.cancelled() => {
return Err(crate::error::OmegaError::NotConnected);
}
permit = self.state.global_handler_limit.clone().acquire_owned() => {
permit.map_err(|_| crate::error::OmegaError::NotConnected)?
}
};
let permit = tokio::select! {
_ = self.handler_cancel.cancelled() => {
drop(global_permit);
return Err(crate::error::OmegaError::NotConnected);
}
permit = self.handler_limit.clone().acquire_owned() => {
permit.map_err(|_| crate::error::OmegaError::NotConnected)?
}
};
let connection = self.clone();
let cancellation = self.handler_cancel.clone();
self.track_handler(async move {
tokio::select! {
_ = cancellation.cancelled() => {}
result = connection.clone().dispatch(value) => {
if let Err(error) = result {
log_err!(
connection.id as i64,
PrintType::Omega,
"Error processing Omikron message: {}",
error
);
}
}
}
drop(permit);
drop(global_permit);
});
Ok(())
}
fn dispatch_class(value: &CommunicationValue) -> DispatchClass {
match value.get_comm_type_enum() {
Some(CommunicationType::UserConnected)
| Some(CommunicationType::UserDisconnected)
| Some(CommunicationType::IotaConnected)
| Some(CommunicationType::IotaDisconnected)
| Some(CommunicationType::SyncClientIotaStatus)
| Some(CommunicationType::StateSubscribe)
| Some(CommunicationType::ClientChanged) => DispatchClass::Ordered,
_ => DispatchClass::Concurrent,
}
}
async fn dispatch(self: Arc<Self>, value: CommunicationValue) -> OmikronResult<()> {
@ -341,33 +543,39 @@ impl OmikronConnection {
value: &CommunicationValue,
timeout: Duration,
) -> OmikronResult<CommunicationValue> {
let (tx, mut rx) = mpsc::channel(1);
let message_id = value.id().unwrap_or_default();
let task_tx = tx.clone();
self.waiting_tasks.insert(
message_id,
WaitingTask {
task: Box::new(move |_, response| {
let task_tx = task_tx.clone();
tokio::spawn(async move {
let _ = task_tx.send(response).await;
});
true
}),
inserted_at: Instant::now(),
},
);
value.require_id()?;
let (tx, rx) = oneshot::channel();
let message_id = self.next_correlation_id().ok_or_else(|| {
crate::OmegaError::Transport("no relay correlation id available".into())
})?;
match self.waiting_tasks.entry(message_id) {
Entry::Vacant(entry) => {
entry.insert(WaitingTask {
sender: tx,
inserted_at: Instant::now(),
});
}
Entry::Occupied(_) => {
return Err(crate::OmegaError::Transport(
"relay correlation id collision".into(),
));
}
}
let outbound = value.clone().with_id(message_id);
if let Err(error) = self.clone().send(value).await {
if let Err(error) = self.clone().send(&outbound).await {
self.waiting_tasks.remove(&message_id);
return Err(error);
}
match tokio::time::timeout(timeout, rx.recv()).await {
Ok(Some(response)) => Ok(response),
Ok(None) => Err(crate::error::OmegaError::Transport(
"Relay response channel closed".into(),
)),
match tokio::time::timeout(timeout, rx).await {
Ok(Ok(response)) => Ok(response),
Ok(Err(_)) => {
self.waiting_tasks.remove(&message_id);
Err(crate::error::OmegaError::Transport(
"Relay response channel closed".into(),
))
}
Err(_) => {
self.waiting_tasks.remove(&message_id);
Err(crate::error::OmegaError::Transport(
@ -411,6 +619,7 @@ impl OmikronConnection {
}
}
async fn cleanup(self: Arc<Self>) {
self.stop_handlers().await;
if self.id != 0 {
log_in!(self.id as i64, PrintType::Omega, "Omikron disconnected");
if omikron_manager::remove_omikron(self.id as i64, &self).await {
@ -441,7 +650,8 @@ pub async fn get_by_omikron_id(
description: Option<String>,
) -> Option<PublicKeyBundle> {
PeerCapabilities::from_identification_description(description.as_deref()).ok()?;
crate::db::omikron_repo::get_omikron_by_id(OmikronId::from(omikron_id as i64))
let omikron_id = i64::try_from(omikron_id).ok().filter(|id| *id > 0)?;
crate::db::omikron_repo::get_omikron_by_id(OmikronId::from(omikron_id))
.await
.ok()
.map(|omikron| omikron.public_key)
@ -453,7 +663,7 @@ pub async fn complete_register(_: PublicKeyBundle, _: Option<String>) -> u64 {
pub async fn start(port: u16, state: Arc<OmegaState>) -> 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 web_config = server::web::build_web_config(state.identity.clone())?
.serve_tcp_https(true)
.max_tcp_connections(256);
let ip = IpAddr::from(Ipv4Addr::new(0, 0, 0, 0));
@ -478,7 +688,7 @@ pub async fn start(port: u16, state: Arc<OmegaState>) -> Result<(), Box<dyn std:
max_frames_per_stream: None,
})
.with_authentication(
load_keyring(),
state.identity.clone_keyring()?,
Box::new(|id, description| Box::pin(get_by_omikron_id(id, description))),
Box::new(|key, description| Box::pin(complete_register(key, description))),
)
@ -499,7 +709,7 @@ pub async fn start(port: u16, state: Arc<OmegaState>) -> Result<(), Box<dyn std:
continue;
}
};
let config = crate::config::RateLimitConfig::from_env();
let config = &state.config.rate_limits;
let peer_ip = conn.remote_addr.map(|address| address.ip());
let active = ACTIVE_CONNECTIONS.fetch_add(1, Ordering::AcqRel) + 1;
let peer_active = peer_ip.map(|ip| {
@ -507,10 +717,11 @@ pub async fn start(port: u16, state: Arc<OmegaState>) -> Result<(), Box<dyn std:
*count += 1;
*count
});
let connection_limit_guard = ConnectionLimitGuard(peer_ip);
if active > config.transport_connections
|| peer_active.is_some_and(|count| count > config.transport_connections_per_ip)
{
drop(ConnectionLimitGuard(peer_ip));
drop(connection_limit_guard);
log_err!(
0,
PrintType::Omega,
@ -526,6 +737,7 @@ pub async fn start(port: u16, state: Arc<OmegaState>) -> Result<(), Box<dyn std:
state.clone(),
authenticated,
) else {
drop(connection_limit_guard);
log_err!(
0,
PrintType::Omega,
@ -534,10 +746,56 @@ pub async fn start(port: u16, state: Arc<OmegaState>) -> Result<(), Box<dyn std:
continue;
};
tokio::spawn(async move {
let _guard = ConnectionLimitGuard(peer_ip);
let _guard = connection_limit_guard;
omikron_manager::add_omikron(connection.clone()).await;
connection.handle(&mut conn.receiver).await;
});
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{DispatchClass, OmikronConnection};
use mtp::codec::{CommunicationType, CommunicationValue};
#[test]
fn relay_dispatch_is_not_on_the_ordered_state_lane() {
let value = CommunicationValue::new(CommunicationType::Relay);
assert_eq!(
OmikronConnection::dispatch_class(&value),
DispatchClass::Concurrent
);
}
#[test]
fn presence_lifecycle_dispatch_is_ordered() {
let value = CommunicationValue::new(CommunicationType::UserConnected);
assert_eq!(
OmikronConnection::dispatch_class(&value),
DispatchClass::Ordered
);
}
#[test]
fn client_presence_changes_are_ordered() {
let value = CommunicationValue::new(CommunicationType::ClientChanged);
assert_eq!(
OmikronConnection::dispatch_class(&value),
DispatchClass::Ordered
);
}
#[test]
fn only_protocol_responses_match_waiting_tasks() {
assert!(OmikronConnection::correlation_response(
&CommunicationValue::new(CommunicationType::Success).with_id(1)
));
assert!(OmikronConnection::correlation_response(
&CommunicationValue::new(CommunicationType::ErrorInvalidData).with_id(1)
));
assert!(!OmikronConnection::correlation_response(
&CommunicationValue::new(CommunicationType::GetUserData).with_id(1)
));
}
}

21
src/transport/omikron_manager.rs
View file

@ -9,7 +9,7 @@ use rand::prelude::IteratorRandom;
use std::sync::Arc;
pub static OMIKRON_CONNECTIONS: Lazy<DashMap<i64, Arc<OmikronConnection>>> =
Lazy::new(|| DashMap::new());
Lazy::new(DashMap::new);
pub async fn add_omikron(conn: Arc<OmikronConnection>) {
let id = match conn.clone().get_omikron_id().await {
@ -64,12 +64,11 @@ pub async fn get_all_connections()
.map_err(|_| ())?;
for user in users {
for route in state.presence.routes_for_user(user.id.0) {
if let Some(iotas) = result.get_mut(&route.omikron_id) {
if let Some(iota_id) = user.iota_id
&& let Some(users) = iotas.get_mut(&iota_id.0)
{
users.push(user.id.0);
}
if let Some(iotas) = result.get_mut(&route.omikron_id)
&& let Some(iota_id) = user.iota_id
&& let Some(users) = iotas.get_mut(&iota_id.0)
{
users.push(user.id.0);
}
}
}
@ -97,10 +96,10 @@ pub async fn send_state_batch(
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(connection) = get_connected_omikron(key) {
return Ok(connection);
}
if let Some(key) = keys.into_iter().choose(&mut rand::rng())
&& let Some(connection) = get_connected_omikron(key)
{
return Ok(connection);
}
Err(())

21
src/transport/relay_router.rs
View file

@ -21,14 +21,6 @@ pub enum RouteTarget {
}
impl RouteTarget {
pub fn wire_id(self) -> Option<u64> {
let (kind, id) = match self {
Self::User(id) => (USER_TARGET_KIND, id),
Self::Iota(id) => (IOTA_TARGET_KIND, id),
};
(id > 0 && id <= TARGET_ID_MASK).then_some(kind | id)
}
pub fn from_wire_id(value: u64) -> Option<Self> {
let id = value & TARGET_ID_MASK;
if id == 0 || value & !(TARGET_KIND_MASK | TARGET_ID_MASK) != 0 {
@ -40,12 +32,6 @@ impl RouteTarget {
_ => None,
}
}
pub const fn id(self) -> u64 {
match self {
Self::User(id) | Self::Iota(id) => id,
}
}
}
#[derive(Debug, Error)]
@ -183,10 +169,11 @@ mod tests {
use mtp::codec::DataValue;
fn wire(target: RouteTarget) -> u64 {
let Some(value) = target.wire_id() else {
panic!("valid route target was rejected");
let (kind, id) = match target {
RouteTarget::User(id) => (USER_TARGET_KIND, id),
RouteTarget::Iota(id) => (IOTA_TARGET_KIND, id),
};
value
kind | id
}
fn relay_frame() -> CommunicationValue {