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

[Fix] Omikron Con no longer breaks on iota disconnect :D

Browse source
This commit is contained in:
Alex Emmet 2026-02-19 19:33:36 +01:00
commit defb4c0426
3 changed files with 64 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

25
src/server/omikron_connection.rs
View file

@ -63,7 +63,11 @@ impl OmikronConnection {
}
pub async fn get_omikron_id(&self) -> i64 {
*self.omikron_id.read().await
let omikron_id = {
let guard = self.omikron_id.read().await;
guard.clone()
};
omikron_id
}
pub async fn is_identified(&self) -> bool {
*self.identified.read().await && *self.challenged.read().await
@ -267,18 +271,23 @@ impl OmikronConnection {
}
if cv.is_type(CommunicationType::iota_disconnected) {
log_in!(PrintType::Omega, "IOTA disconnected");
if let Some(iota_id) = cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()) {
let iota_offline =
user_online_tracker::untrack_iota_connection(iota_id, omikron_id);
if iota_offline {
if let Ok(users) = sql::get_users_by_iota_id(iota_id).await {
let user_ids: Vec<i64> = users.iter().map(|u| u.0).collect();
user_online_tracker::untrack_many_users(&user_ids);
if let Some(v) = cv.get_data(DataTypes::iota_id) {
if let Some(iota_id) = v.as_i64() {
let iota_offline =
user_online_tracker::untrack_iota_connection(iota_id, omikron_id);
if iota_offline {
if let Ok(users) = sql::get_users_by_iota_id(iota_id).await {
let user_ids: Vec<i64> = users.iter().map(|u| u.0).collect();
user_online_tracker::untrack_many_users(&user_ids);
}
}
}
}
return;
}
if cv.is_type(CommunicationType::sync_client_iota_status) {
if let Some(json::JsonValue::Array(user_ids)) =
cv.get_data(DataTypes::user_ids).cloned()

4
src/sql/sql.rs
View file

@ -65,9 +65,9 @@ pub async fn initialize_db() -> Result<(), sqlx::Error> {
about VARCHAR(200) COLLATE utf8mb4_bin,
avatar MEDIUMBLOB,
sub_level INT(11) NOT NULL DEFAULT 0,
sub_end BIGINT(20) NOT NULL,
sub_end BIGINT(20) NOT NULL DEFAULT 0,
public_key TEXT NOT NULL COLLATE utf8mb4_bin,
private_key_hash TEXT NOT NULL COLLATE utf8mb4_bin,
private_key_hash TEXT NOT NULL COLLATE utf8mb4_bin DEFAULT '',
iota_id BIGINT UNSIGNED NOT NULL,
token VARCHAR(255) NOT NULL UNIQUE COLLATE utf8mb4_bin
)",

69
src/sql/user_online_tracker.rs
View file

@ -33,23 +33,30 @@ pub fn track_iota_connection(iota_id: i64, omikron_id: i64, primary: bool) {
}
pub fn untrack_iota_connection(iota_id: i64, omikron_id: i64) -> bool {
if let Some(mut connections) = IOTA_OMIKRON_CONNECTIONS.get_mut(&iota_id) {
connections.retain(|&id| id != omikron_id);
let connections_empty = if let Some(r) = IOTA_OMIKRON_CONNECTIONS.get(&iota_id) {
let mut vec = r.value().clone();
vec.retain(|&id| id != omikron_id);
let empty = vec.is_empty();
drop(r);
IOTA_OMIKRON_CONNECTIONS.insert(iota_id, vec);
empty
} else {
false
};
if IOTA_PRIMARY_OMIKRON_CONNECTION
.get(&iota_id)
.map(|p| *p == omikron_id)
.unwrap_or(false)
{
if let Some(primary_ref) = IOTA_PRIMARY_OMIKRON_CONNECTION.get(&iota_id) {
let primary_id = *primary_ref.value();
drop(primary_ref);
if primary_id == omikron_id {
IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
}
if connections.is_empty() {
IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
return true;
}
}
false
if connections_empty {
IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
}
connections_empty
}
pub fn get_iota_primary_omikron_connection(iota_id: i64) -> Option<i64> {
@ -81,34 +88,48 @@ pub fn untrack_many_users(user_ids: &[i64]) {
}
pub async fn untrack_omikron(omikron_id: i64) {
// Step 1: Collect all iota_ids where this omikron_id is connected
let mut offline_iotas = Vec::new();
let mut primary_to_remove = Vec::new();
IOTA_OMIKRON_CONNECTIONS.retain(|iota_id, connections| {
connections.retain(|id| *id != omikron_id);
// Collect iotas and primary info first
for r in IOTA_OMIKRON_CONNECTIONS.iter() {
let iota_id = *r.key();
let mut connections = r.value().clone();
connections.retain(|&id| id != omikron_id);
if connections.is_empty() {
offline_iotas.push(iota_id);
}
if IOTA_PRIMARY_OMIKRON_CONNECTION
.get(iota_id)
.get(&iota_id)
.map(|p| *p == omikron_id)
.unwrap_or(false)
{
IOTA_PRIMARY_OMIKRON_CONNECTION.remove(iota_id);
primary_to_remove.push(iota_id);
}
if connections.is_empty() {
offline_iotas.push(*iota_id);
false
} else {
true
}
});
// Update the connections vector after filtering
IOTA_OMIKRON_CONNECTIONS.insert(iota_id, connections);
}
// Step 2: Remove primary connections safely
for iota_id in primary_to_remove {
IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
}
// Step 3: Remove users that were on this omikron
USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);
// Step 4: For offline iotas, remove associated users from USER_STATUS_MAP
for iota_id in offline_iotas {
if let Ok(users) = sql::sql::get_users_by_iota_id(iota_id).await {
for user in users {
USER_STATUS_MAP.remove(&user.0);
}
}
// Finally remove the empty connections vector
IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
}
}