[Fix] Stability

src/config.rs

@@ -9,8 +9,8 @@ pub struct RateLimitConfig {
     pub transport_connections_per_ip: usize,
 }
 
-pub fn cors_origin() -> String {
-    env::var("CORS_ORIGIN").unwrap_or_else(|_| "https://tensamin.net".to_string())
+pub const fn cors_origin() -> &'static str {
+    "*"
 }
 
 impl Default for RateLimitConfig {

src/db/iota_repo.rs

@@ -21,13 +21,31 @@ pub async fn get_iota_by_id(id: IotaId) -> Result<Iota> {
 }
 
 pub async fn create_new_iota(public_key: PublicKeyBundle) -> Result<IotaId> {
-    let id = crate::db::user_repo::get_register_id().await?;
-    let iota_id = IotaId::from(id.0);
-    register_complete_iota(iota_id, public_key).await?;
-    Ok(iota_id)
+    for _ in 0..16 {
+        let id = crate::db::user_repo::get_register_id().await?;
+        let iota_id = IotaId::from(id.0);
+        match register_complete_iota(iota_id, public_key.clone()).await {
+            Ok(()) => return Ok(iota_id),
+            Err(OmegaError::Database(error)) => {
+                if crate::db::user_repo::is_duplicate_key(&error) {
+                    continue;
+                }
+                return Err(OmegaError::Database(error));
+            }
+            Err(error) => return Err(error),
+        }
+    }
+    Err(OmegaError::Validation(
+        "could not allocate a unique Iota ID".into(),
+    ))
 }
 
 pub async fn register_complete_iota(id: IotaId, public_key: PublicKeyBundle) -> Result<()> {
+    if !crate::db::user_repo::valid_protocol_id(id.0) {
+        return Err(OmegaError::Validation(
+            "Iota ID is outside the 48-bit protocol range".into(),
+        ));
+    }
     sqlx::query("INSERT INTO iotas (id, public_key) VALUES (?, ?)")
         .bind(id.0)
         .bind(public_key.as_bytes())

src/db/user_repo.rs

@@ -4,17 +4,94 @@ use crate::{
     models::{IotaId, User, UserId},
 };
 use mtp::crypto::PublicKeyBundle;
-use sqlx::FromRow;
+use sqlx::{FromRow, Row};
+
+pub const MAX_PROTOCOL_ID: i64 = (1_i64 << 48) - 1;
+const ID_ALLOCATION_ATTEMPTS: usize = 16;
 
 pub async fn get_register_id() -> Result<UserId> {
-    let bytes = *uuid::Uuid::now_v7().as_bytes();
-    let id =
-        i64::from_be_bytes(bytes[8..].try_into().map_err(|_| {
-            OmegaError::Validation("generated ID has an invalid length".to_string())
-        })?) & i64::MAX;
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    let timestamp = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_secs())
+        .unwrap_or(0);
+
+    let ts = timestamp as i64;
+    if ts >= 1 && ts <= MAX_PROTOCOL_ID {
+        return Ok(UserId::from(ts));
+    }
+
+    // Fall back to random if the timestamp is outside the 48-bit range.
+    let id = (rand::random::<u64>() & ((1_u64 << 48) - 1)) as i64;
     Ok(UserId::from(id.max(1)))
 }
 
+pub fn valid_protocol_id(id: i64) -> bool {
+    (1..=MAX_PROTOCOL_ID).contains(&id)
+}
+
+/// Allocate an ID that is durable, short-lived, and bound to the connected
+/// Iota.  `token` is presented again when completing the registration.
+pub async fn allocate_registration(iota_id: IotaId, request_id: u32) -> Result<(UserId, String)> {
+    if !valid_protocol_id(iota_id.0) {
+        return Err(OmegaError::Validation(
+            "Iota ID is outside the 48-bit protocol range".into(),
+        ));
+    }
+
+    for _ in 0..ID_ALLOCATION_ATTEMPTS {
+        let id = get_register_id().await?;
+        let token = uuid::Uuid::new_v4().to_string();
+        let result = sqlx::query(
+            "INSERT INTO registration_leases (token, user_id, iota_id, request_id, expires_at) \
+             VALUES (?, ?, ?, ?, DATE_ADD(UTC_TIMESTAMP(), INTERVAL 10 MINUTE))",
+        )
+        .bind(&token)
+        .bind(id.0)
+        .bind(iota_id.0)
+        .bind(request_id)
+        .execute(&pool().await?)
+        .await;
+        match result {
+            Ok(_) => return Ok((id, token)),
+            Err(error) if is_duplicate_key(&error) => {
+                let existing = sqlx::query(
+                    "SELECT user_id, token, expires_at >= UTC_TIMESTAMP() AS current \
+                     FROM registration_leases WHERE iota_id = ? AND request_id = ?",
+                )
+                .bind(iota_id.0)
+                .bind(request_id)
+                .fetch_optional(&pool().await?)
+                .await?;
+                if let Some(existing) = existing {
+                    let current: i8 = existing.get("current");
+                    if current != 0 {
+                        return Ok((
+                            UserId::from(existing.get::<i64, _>("user_id")),
+                            existing.get::<String, _>("token"),
+                        ));
+                    }
+                }
+                continue;
+            }
+            Err(error) => return Err(error.into()),
+        }
+    }
+    Err(OmegaError::Validation(
+        "could not allocate a unique registration ID".into(),
+    ))
+}
+
+pub(crate) fn is_duplicate_key(error: &sqlx::Error) -> bool {
+    error.as_database_error().is_some_and(|database| {
+        // MySQL's generic database-error API exposes SQLSTATE (23000) as
+        // `code()`.  The driver-specific duplicate-key number is retained in
+        // the diagnostic message.
+        database.code().as_deref() == Some("23000") && database.message().contains("1062")
+    })
+}
+
 const USER_BY_USERNAME_QUERY: &str = "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE username = ?";
 const USER_BY_ID_QUERY: &str = "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE id = ?";
 const USERS_BY_IOTA_ID_QUERY: &str = "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE iota_id = ?";
@@ -179,16 +256,99 @@ pub async fn register_complete_user(
     public_key: PublicKeyBundle,
     iota_id: IotaId,
     token: String,
+    registration_token: String,
 ) -> Result<()> {
-    sqlx::query(
+    if !valid_protocol_id(id.0) || !valid_protocol_id(iota_id.0) {
+        return Err(OmegaError::Validation(
+            "user or Iota ID is outside the 48-bit protocol range".into(),
+        ));
+    }
+    if !valid_username(&username) {
+        return Err(OmegaError::Validation("invalid username".into()));
+    }
+
+    let mut transaction = pool().await?.begin().await?;
+    let lease = sqlx::query(
+        "SELECT iota_id, completed_at IS NOT NULL AS completed, \
+         expires_at >= UTC_TIMESTAMP() AS current FROM registration_leases \
+         WHERE token = ? AND user_id = ? FOR UPDATE",
+    )
+    .bind(&registration_token)
+    .bind(id.0)
+    .fetch_optional(&mut *transaction)
+    .await?
+    .ok_or_else(|| OmegaError::Validation("unknown registration lease".into()))?;
+    let lease_iota_id: i64 = lease.get("iota_id");
+    let completed: i8 = lease.get("completed");
+    let current: i8 = lease.get("current");
+    if lease_iota_id != iota_id.0 || (completed == 0 && current == 0) {
+        return Err(OmegaError::Validation(
+            "expired or mismatched registration lease".into(),
+        ));
+    }
+
+    let insert_result = sqlx::query(
         "INSERT INTO users (id, username, public_key, iota_id, token) VALUES (?, ?, ?, ?, ?)",
     )
     .bind(id.0)
-    .bind(username.into_bytes())
+    .bind(username.as_bytes())
     .bind(public_key.as_bytes())
     .bind(iota_id.0)
-    .bind(token.into_bytes())
-    .execute(&pool().await?)
-    .await?;
+    .bind(token.as_bytes())
+    .execute(&mut *transaction)
+    .await;
+
+    let result: Result<()> = match insert_result {
+        Ok(_) => Ok(()),
+        Err(insert_error) => {
+            let existing = sqlx::query_as::<_, UserRow>(USER_BY_ID_QUERY)
+                .bind(id.0)
+                .fetch_optional(&mut *transaction)
+                .await?;
+            match existing
+                .map(User::try_from)
+                .transpose()
+                .map_err(OmegaError::from)?
+            {
+                Some(existing)
+                    if existing.iota_id == iota_id
+                        && existing.username == username
+                        && existing.public_key.as_bytes() == public_key.as_bytes()
+                        && existing.token == token =>
+                {
+                    Ok(())
+                }
+                _ => Err(insert_error.into()),
+            }
+        }
+    };
+    result?;
+    sqlx::query("UPDATE registration_leases SET completed_at = UTC_TIMESTAMP() WHERE token = ?")
+        .bind(&registration_token)
+        .execute(&mut *transaction)
+        .await?;
+    transaction.commit().await?;
     Ok(())
 }
+
+fn valid_username(username: &str) -> bool {
+    !username.is_empty()
+        && username.chars().count() <= 15
+        && !username.chars().any(char::is_control)
+        && !username.contains(['/', '\\'])
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{MAX_PROTOCOL_ID, get_register_id, valid_protocol_id};
+
+    #[tokio::test]
+    async fn generated_registration_ids_fit_the_mtp_wire_range() {
+        for _ in 0..128 {
+            assert!(valid_protocol_id(get_register_id().await.unwrap().0));
+        }
+        assert!(!valid_protocol_id(0));
+        assert!(valid_protocol_id(MAX_PROTOCOL_ID));
+        assert!(!valid_protocol_id(MAX_PROTOCOL_ID + 1));
+    }
+}

src/server/api.rs

@@ -15,12 +15,12 @@ use crate::server::{
     validation::{parse_positive_id, validate_non_empty},
 };
 use crate::sql::user_online_tracker::{get_all_connections, get_iota_primary_omikron_connection};
-use crate::transport::omikron_manager::get_random_omikron;
+use crate::transport::omikron_manager::{get_connected_omikron, get_random_omikron};
 use crate::util::file_util::get_directory;
 use base64::Engine as _;
 use bytes::Bytes;
 use http::{Method, StatusCode};
-use mtp::webserver::{Http3Request, Http3Response, RouteParams};
+use mtp::webserver::{HttpRequest, HttpResponse, RouteParams};
 use std::collections::BTreeMap;
 
 fn error_body(error: &OmegaError) -> String {
@@ -75,20 +75,19 @@ async fn route(path_parts: &[&str]) -> Result<(StatusCode, String)> {
         }
         ["api", "get", "omikron", id] => {
             let id = parse_positive_id(id)?;
-            let omikron = match get_omikron_by_id(id.into()).await {
-                Ok(value) => value,
-                Err(_) => {
-                    let fallback_id =
-                        if let Some(fallback_id) = get_iota_primary_omikron_connection(id) {
-                            fallback_id
-                        } else {
-                            let user = get_by_user_id(UserId::from(id)).await?;
-                            get_iota_primary_omikron_connection(user.iota_id.0)
-                                .ok_or(OmegaError::NotFound)?
-                        };
-                    get_omikron_by_id(fallback_id.into()).await?
-                }
+            let omikron_id = if get_connected_omikron(id).is_some() {
+                id
+            } else if let Some(omikron_id) = get_iota_primary_omikron_connection(id) {
+                omikron_id
+            } else {
+                let user = get_by_user_id(UserId::from(id)).await?;
+                get_iota_primary_omikron_connection(user.iota_id.0).ok_or(OmegaError::NotFound)?
             };
+
+            // Database rows describe registered Omikrons. The public discovery
+            // API must expose only routes backed by a currently live transport.
+            get_connected_omikron(omikron_id).ok_or(OmegaError::NotFound)?;
+            let omikron = get_omikron_by_id(omikron_id.into()).await?;
             Ok((
                 StatusCode::OK,
                 json(&OmikronResponse {
@@ -178,7 +177,7 @@ async fn route(path_parts: &[&str]) -> Result<(StatusCode, String)> {
     }
 }
 
-pub async fn handle(request: Http3Request, response: Http3Response) -> Http3Response {
+pub async fn handle(request: HttpRequest, response: HttpResponse) -> HttpResponse {
     let method = request.method;
     let path = request.uri.path().to_string();
     if method != Method::OPTIONS && !middleware::allow(request.remote_addr.ip(), &path) {
@@ -238,9 +237,9 @@ pub async fn handle(request: Http3Request, response: Http3Response) -> Http3Resp
 }
 
 pub async fn handle_pattern(
-    request: Http3Request,
-    response: Http3Response,
+    request: HttpRequest,
+    response: HttpResponse,
     _params: RouteParams,
-) -> Http3Response {
+) -> HttpResponse {
     handle(request, response).await
 }

src/server/index.rs

@@ -1,13 +1,34 @@
 use http::StatusCode;
-use mtp::webserver::Http3Response;
+use mtp::webserver::HttpResponse;
 
-pub fn index_handler(response: Http3Response) -> Http3Response {
+pub fn index_handler(response: HttpResponse) -> HttpResponse {
     let documentation = r#"
 Omega API Server
 
 Available Routes:
-- /api/*    : API endpoints for the Omega server.
-- /direct/* : Resolution for shortened links.
+- /api/get/omikron
+    Returns a randomly selected connected Omikron.
+- /api/get/omikron/{id}
+    Returns an Omikron by ID. If {id} is a connected Iota ID or user ID,
+    returns that account's primary connected Omikron instead.
+- /api/get/connections
+    Returns the current live connection map as
+    { "status": "success", "connections": { omikron_id: { iota_id: [user_id] } } }.
+- /api/get/iota/{id}
+    Returns the Iota's public identity data.
+- /api/get/user/{id}
+    Returns user profile and public identity data.
+- /api/get/id/{username}
+    Resolves a username to its user and Iota IDs.
+- /api/get/public_key
+    Returns Omega's public key.
+- /api/download/iota_frontend
+    Downloads the Iota frontend archive.
+- /direct/{short_key}
+    Resolves a shortened link with a temporary redirect.
+
+All IDs must be positive decimal integers. The connections endpoint reports
+live in-memory state; it is empty after an Omega restart until Omikrons sync.
 
 All other routes will return this documentation.
 "#;

src/server/server.rs

@@ -3,10 +3,6 @@ use mtp::webserver::WebServerConfig;
 
 pub fn build_web_config() -> Result<WebServerConfig, mtp::webserver::RouterError> {
     WebServerConfig::new()
-        .route(
-            "/",
-            |_request, response| async move { index_handler(response) },
-        )?
         .route("/api/download/iota_frontend", api::handle)?
         .route("/api/get/omikron", api::handle)?
         .route("/api/get/connections", api::handle)?
@@ -15,5 +11,6 @@ pub fn build_web_config() -> Result<WebServerConfig, mtp::webserver::RouterError
         .route_pattern("/api/get/iota/{id}", api::handle_pattern)?
         .route_pattern("/api/get/id/{username}", api::handle_pattern)?
         .route_pattern("/api/get/user/{id}", api::handle_pattern)?
-        .route_pattern("/direct/{short}", api::handle_pattern)
+        .route_pattern("/direct/{short}", api::handle_pattern)?
+        .fallback(|_request, response| async move { index_handler(response) })
 }

src/sql/user_online_tracker.rs

@@ -19,6 +19,11 @@ static IOTA_OMIKRON_CONNECTIONS: Lazy<DashMap<i64, Vec<i64>>> = Lazy::new(DashMa
 // UserID -> UserStatus
 static USER_STATUS_MAP: Lazy<DashMap<i64, UserConnection>> = Lazy::new(DashMap::new);
 
+// The legacy account map is kept for Iota-only/offline compatibility.  Client
+// transports are tracked independently: two devices must never overwrite each
+// other's route merely because they authenticate as the same account.
+static USER_SESSION_STATUS_MAP: Lazy<DashMap<(i64, i64), UserConnection>> = Lazy::new(DashMap::new);
+
 pub fn track_iota_connection(iota_id: i64, omikron_id: i64, primary: bool) {
     let mut entry = IOTA_OMIKRON_CONNECTIONS
         .entry(iota_id)
@@ -121,13 +126,63 @@ pub fn track_user_status(user_id: i64, status: UserStatus, omikron_id: i64) {
     );
 }
 
+pub fn track_user_session_status(
+    user_id: i64,
+    session_id: i64,
+    status: UserStatus,
+    omikron_id: i64,
+) {
+    USER_SESSION_STATUS_MAP.insert(
+        (user_id, session_id),
+        UserConnection {
+            connection_type: status,
+            omikron_id,
+        },
+    );
+}
+
 pub fn untrack_user_status(user_id: i64, omikron_id: i64) {
     USER_STATUS_MAP.remove_if(&user_id, |_, connection| {
         connection.omikron_id == omikron_id
     });
 }
 
+pub fn untrack_user_session_status(user_id: i64, session_id: i64, omikron_id: i64) {
+    USER_SESSION_STATUS_MAP.remove_if(&(user_id, session_id), |_, connection| {
+        connection.omikron_id == omikron_id
+    });
+}
+
+pub fn update_user_session_status(user_id: i64, status: UserStatus) {
+    for mut entry in USER_SESSION_STATUS_MAP.iter_mut() {
+        if entry.key().0 == user_id {
+            entry.connection_type = status.clone();
+        }
+    }
+    // Preserve the account preference for legacy routes as well.
+    if let Some(mut entry) = USER_STATUS_MAP.get_mut(&user_id) {
+        entry.connection_type = status;
+    }
+}
+
 pub fn get_user_status(user_id: i64) -> Option<UserConnection> {
+    // A connected visible session is preferred.  Invisible sessions remain
+    // routable but are intentionally presented as offline when they are the
+    // only active routes.
+    let sessions: Vec<UserConnection> = USER_SESSION_STATUS_MAP
+        .iter()
+        .filter(|entry| entry.key().0 == user_id)
+        .map(|entry| entry.value().clone())
+        .collect();
+    if let Some(status) = sessions
+        .iter()
+        .find(|status| status.connection_type != UserStatus::user_invisible)
+    {
+        return Some(status.clone());
+    }
+    if let Some(status) = sessions.first() {
+        return Some(status.clone());
+    }
     USER_STATUS_MAP.get(&user_id).map(|v| v.clone())
 }
 
@@ -164,6 +219,7 @@ pub async fn untrack_omikron(omikron_id: i64) {
     }
 
     USER_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);
+    USER_SESSION_STATUS_MAP.retain(|_, status| status.omikron_id != omikron_id);
 
     for iota_id in offline_iotas {
         if let Ok(users) = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await {

src/transport/handlers/presence.rs

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

src/transport/handlers/register.rs

@@ -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
     {

src/transport/omikron_connection.rs

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

src/transport/omikron_manager.rs

@@ -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);
         }
     }
 

src/util/logger.rs

@@ -244,8 +244,11 @@ pub fn format_cv(cv: &CommunicationValue) -> String {
         parts.push(format!("> {}", receiver));
     }
 
-    let comm_type = cv.get_type().to_string();
-    parts.push(format!("{}", comm_type));
+    let comm_type = cv
+        .get_comm_type_enum()
+        .map(|kind| kind.to_string())
+        .unwrap_or_else(|| cv.get_type().to_string());
+    parts.push(format!("{} (id={})", comm_type, cv.get_id()));
 
     let data = cv.data();