Merge remote-tracking branch 'refs/remotes/origin/main'

Cargo.toml

@@ -10,33 +10,27 @@ mtp = { git = "https://git.methanium.net/methanium/mtp", features = [
     "web-server",
 ] }
 
-aes-gcm = "*"
 ansi_term = "0.12.1"
-anyhow = "1.0.101"
 base64 = "0.22.1"
 bytes = "1"
-dashmap = "6.1.0"
+dashmap = "6.2.1"
 dotenv = "0.15.0"
-hex = "0.4.3"
-hkdf = "0.12.4"
 http = "1"
-json = "0.12.4"
+once_cell = "1.21.4"
-once_cell = "1.21.3"
+rand = "0.10.2"
-rand = "0.8"
+reqwest = "0.13.4"
-rand_core = { version = "0.6", features = ["getrandom", "std"] }
+rustls = { version = "0.23.42", default-features = false, features = [
-reqwest = { version = "0.13.2" }
-rustls = { version = "0.23.37", default-features = false, features = [
     "std",
     "tls12",
     "aws-lc-rs",
     "prefer-post-quantum",
 ] }
-sha2 = "0.10.9"
+sqlx = { version = "0.8.6", features = ["mysql", "runtime-tokio", "migrate"] }
-sqlx = { version = "0.8.6", features = ["mysql", "runtime-async-std"] }
+strum = "0.28.0"
-strum = "0.27.2"
+strum_macros = "0.28.0"
-strum_macros = "0.27.2"
 tokio = { version = "*", features = ["full"] }
-uuid = { version = "1.19.0", features = ["v4"] }
+uuid = { version = "1.24.0", features = ["v4", "v7"] }
-x448 = "0.6.0"
+zip = "8.6.0"
-zip = "6.0.0"
+thiserror = "2.0.19"
-thiserror = "2.0.18"
+serde = { version = "1.0.229", features = ["derive"] }
+serde_json = "1.0.151"

migrations/001.sql

@@ -0,0 +1,44 @@
+CREATE TABLE IF NOT EXISTS users (
+    id BIGINT NOT NULL PRIMARY KEY,
+    iota_id BIGINT NOT NULL,
+    username VARBINARY(255) NOT NULL,
+    display VARBINARY(255),
+    status VARBINARY(255),
+    about VARBINARY(1000),
+    avatar BLOB,
+    sub_level INT NOT NULL DEFAULT 0,
+    sub_end BIGINT NOT NULL DEFAULT 0,
+    public_key BLOB NOT NULL,
+    token BLOB NOT NULL,
+    UNIQUE KEY uk_users_username (username),
+    UNIQUE KEY uk_users_iota_id (iota_id)
+);
+
+CREATE TABLE IF NOT EXISTS iotas (
+    id BIGINT NOT NULL PRIMARY KEY,
+    public_key BLOB NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS omikrons (
+    id BIGINT NOT NULL PRIMARY KEY,
+    public_key BLOB NOT NULL,
+    location VARBINARY(255) NOT NULL,
+    ip_address VARBINARY(45) NOT NULL,
+    port INT NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS notifications (
+    id BIGINT NOT NULL AUTO_INCREMENT PRIMARY KEY,
+    sender_id BIGINT NOT NULL,
+    receiver_id BIGINT NOT NULL,
+    amount BIGINT NOT NULL DEFAULT 1,
+    UNIQUE KEY uk_notifications_sender_receiver (sender_id, receiver_id),
+    INDEX idx_notifications_receiver (receiver_id)
+);
+
+CREATE TABLE IF NOT EXISTS short_links (
+    short_key VARCHAR(12) CHARACTER SET ascii COLLATE ascii_bin NOT NULL PRIMARY KEY,
+    long_url TEXT NOT NULL,
+    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    INDEX idx_short_links_created_at (created_at)
+);

src/api/mod.rs

@@ -0,0 +1 @@
+pub mod response;

src/api/response.rs

@@ -0,0 +1,70 @@
+use serde::Serialize;
+use std::collections::BTreeMap;
+
+#[derive(Serialize)]
+pub struct StatusResponse {
+    pub status: &'static str,
+}
+
+#[derive(Serialize)]
+pub struct OmikronResponse {
+    pub status: &'static str,
+    pub id: i64,
+    pub public_key: String,
+    pub ip_address: String,
+    pub port: u16,
+}
+
+#[derive(Serialize)]
+pub struct IotaResponse {
+    pub status: &'static str,
+    pub iota_id: i64,
+    pub public_key: String,
+}
+
+#[derive(Serialize)]
+pub struct UserResponse {
+    pub status: &'static str,
+    pub username: String,
+    pub public_key: String,
+    pub user_id: i64,
+    pub iota_id: i64,
+    pub sub_level: i32,
+    pub sub_end: i64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub display: Option<String>,
+    #[serde(rename = "status_message", skip_serializing_if = "Option::is_none")]
+    pub status_message: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub about: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub avatar: Option<String>,
+}
+
+#[derive(Serialize)]
+pub struct UsernameResponse {
+    pub status: &'static str,
+    pub username: String,
+    pub public_key: String,
+    pub user_id: i64,
+    pub iota_id: i64,
+    pub sub_level: i32,
+    pub sub_end: i64,
+}
+
+#[derive(Serialize)]
+pub struct ConnectionsResponse {
+    pub status: &'static str,
+    #[serde(flatten)]
+    pub connections: BTreeMap<String, BTreeMap<String, Vec<i64>>>,
+}
+
+#[derive(Serialize)]
+pub struct PublicKeyResponse {
+    pub status: &'static str,
+    pub public_key: String,
+}
+
+pub fn json<T: Serialize>(value: &T) -> String {
+    serde_json::to_string(value).unwrap_or_else(|_| "{\"status\":\"error\"}".to_string())
+}

src/config.rs

@@ -0,0 +1,65 @@
+use std::{env, time::Duration};
+
+#[derive(Clone, Debug)]
+pub struct RateLimitConfig {
+    pub window: Duration,
+    pub general_requests: usize,
+    pub registration_requests: usize,
+    pub transport_connections: usize,
+    pub transport_connections_per_ip: usize,
+}
+
+pub fn cors_origin() -> String {
+    env::var("CORS_ORIGIN").unwrap_or_else(|_| "https://tensamin.net".to_string())
+}
+
+impl Default for RateLimitConfig {
+    fn default() -> Self {
+        Self {
+            window: Duration::from_secs(60),
+            general_requests: 120,
+            registration_requests: 20,
+            transport_connections: 512,
+            transport_connections_per_ip: 32,
+        }
+    }
+}
+
+impl RateLimitConfig {
+    pub fn from_env() -> Self {
+        let defaults = Self::default();
+        Self {
+            window: env_duration("RATE_LIMIT_WINDOW_SECONDS", defaults.window),
+            general_requests: env_usize("RATE_LIMIT_GENERAL_REQUESTS", defaults.general_requests),
+            registration_requests: env_usize(
+                "RATE_LIMIT_REGISTRATION_REQUESTS",
+                defaults.registration_requests,
+            ),
+            transport_connections: env_usize(
+                "RATE_LIMIT_TRANSPORT_CONNECTIONS",
+                defaults.transport_connections,
+            ),
+            transport_connections_per_ip: env_usize(
+                "RATE_LIMIT_TRANSPORT_CONNECTIONS_PER_IP",
+                defaults.transport_connections_per_ip,
+            ),
+        }
+    }
+}
+
+fn env_usize(name: &str, fallback: usize) -> usize {
+    env::var(name)
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .filter(|value| *value > 0)
+        .unwrap_or(fallback)
+}
+
+fn env_duration(name: &str, fallback: Duration) -> Duration {
+    env::var(name)
+        .ok()
+        .and_then(|value| value.parse::<u64>().ok())
+        .filter(|value| *value > 0)
+        .map(Duration::from_secs)
+        .unwrap_or(fallback)
+}

src/db/iota_repo.rs

@@ -0,0 +1,53 @@
+use crate::{
+    db::pool,
+    error::{OmegaError, Result},
+    models::{Iota, IotaId},
+};
+use mtp::crypto::PublicKeyBundle;
+use sqlx::Row;
+
+pub async fn get_iota_by_id(id: IotaId) -> Result<Iota> {
+    let row = sqlx::query("SELECT id, public_key FROM iotas WHERE id = ?")
+        .bind(id.0)
+        .fetch_optional(&pool().await?)
+        .await?
+        .ok_or(OmegaError::NotFound)?;
+    let key = PublicKeyBundle::from_bytes(&row.get::<Vec<u8>, _>("public_key"))
+        .map_err(|error| OmegaError::Validation(error.to_string()))?;
+    Ok(Iota {
+        id: row.get::<i64, _>("id").into(),
+        public_key: key,
+    })
+}
+
+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)
+}
+
+pub async fn register_complete_iota(id: IotaId, public_key: PublicKeyBundle) -> Result<()> {
+    sqlx::query("INSERT INTO iotas (id, public_key) VALUES (?, ?)")
+        .bind(id.0)
+        .bind(public_key.as_bytes())
+        .execute(&pool().await?)
+        .await?;
+    Ok(())
+}
+
+pub async fn change_iota_key(id: IotaId, key: PublicKeyBundle) -> Result<()> {
+    sqlx::query("UPDATE iotas SET public_key = ? WHERE id = ?")
+        .bind(key.as_bytes())
+        .bind(id.0)
+        .execute(&pool().await?)
+        .await?;
+    Ok(())
+}
+pub async fn delete_iota(id: IotaId) -> Result<()> {
+    sqlx::query("DELETE FROM iotas WHERE id = ?")
+        .bind(id.0)
+        .execute(&pool().await?)
+        .await?;
+    Ok(())
+}

src/db/mod.rs

@@ -0,0 +1,38 @@
+use crate::error::{OmegaError, Result};
+use once_cell::sync::Lazy;
+use sqlx::{MySql, Pool, mysql::MySqlPoolOptions};
+use std::{env, sync::Arc};
+use tokio::sync::RwLock;
+
+pub mod iota_repo;
+pub mod notification_repo;
+pub mod omikron_repo;
+pub mod short_link_repo;
+pub mod user_repo;
+
+pub type DbPool = Pool<MySql>;
+
+static POOL: Lazy<Arc<RwLock<Option<DbPool>>>> = Lazy::new(|| Arc::new(RwLock::new(None)));
+
+pub async fn pool() -> Result<DbPool> {
+    POOL.read()
+        .await
+        .as_ref()
+        .cloned()
+        .ok_or(OmegaError::DatabaseNotInitialized)
+}
+
+pub async fn initialize() -> Result<()> {
+    let url =
+        env::var("DB_URL").map_err(|_| OmegaError::Validation("DB_URL is not set".to_string()))?;
+    let pool = MySqlPoolOptions::new()
+        .max_connections(200)
+        .connect(&url)
+        .await?;
+    sqlx::migrate!()
+        .run(&pool)
+        .await
+        .map_err(|error| OmegaError::Validation(format!("database migration failed: {error}")))?;
+    *POOL.write().await = Some(pool);
+    Ok(())
+}

src/db/notification_repo.rs

@@ -0,0 +1,38 @@
+use crate::{
+    db::pool,
+    error::Result,
+    models::{Notification, UserId},
+};
+use sqlx::Row;
+
+pub async fn add_notification(sender_id: UserId, receiver_id: UserId) -> Result<()> {
+    sqlx::query("INSERT INTO notifications (sender_id, receiver_id, amount) VALUES (?, ?, 1) ON DUPLICATE KEY UPDATE amount = amount + 1").bind(sender_id.0).bind(receiver_id.0).execute(&pool().await?).await?;
+    Ok(())
+}
+
+pub async fn read_notification(sender_id: UserId, receiver_id: UserId) -> Result<()> {
+    sqlx::query("DELETE FROM notifications WHERE sender_id = ? AND receiver_id = ?")
+        .bind(sender_id.0)
+        .bind(receiver_id.0)
+        .execute(&pool().await?)
+        .await?;
+    Ok(())
+}
+
+pub async fn get_notifications(receiver_id: UserId) -> Result<Vec<Notification>> {
+    let rows = sqlx::query(
+        "SELECT id, sender_id, receiver_id, amount FROM notifications WHERE receiver_id = ?",
+    )
+    .bind(receiver_id.0)
+    .fetch_all(&pool().await?)
+    .await?;
+    Ok(rows
+        .into_iter()
+        .map(|row| Notification {
+            id: row.get("id"),
+            sender_id: row.get::<i64, _>("sender_id").into(),
+            receiver_id: row.get::<i64, _>("receiver_id").into(),
+            amount: row.get("amount"),
+        })
+        .collect())
+}

src/db/omikron_repo.rs

@@ -0,0 +1,29 @@
+use crate::{
+    db::pool,
+    error::{OmegaError, Result},
+    models::{Omikron, OmikronId},
+};
+use mtp::crypto::PublicKeyBundle;
+use sqlx::Row;
+
+pub async fn get_omikron_by_id(id: OmikronId) -> Result<Omikron> {
+    let row =
+        sqlx::query("SELECT id, public_key, location, ip_address, port FROM omikrons WHERE id = ?")
+            .bind(id.0)
+            .fetch_optional(&pool().await?)
+            .await?
+            .ok_or(OmegaError::NotFound)?;
+    let key = PublicKeyBundle::from_bytes(&row.get::<Vec<u8>, _>("public_key"))
+        .map_err(|error| OmegaError::Validation(error.to_string()))?;
+    let location = String::from_utf8(row.get("location"))
+        .map_err(|error| OmegaError::Validation(format!("invalid location UTF-8: {error}")))?;
+    let ip_address = String::from_utf8(row.get("ip_address"))
+        .map_err(|error| OmegaError::Validation(format!("invalid IP address UTF-8: {error}")))?;
+    Ok(Omikron {
+        id: row.get::<i64, _>("id").into(),
+        public_key: key,
+        location,
+        ip_address,
+        port: row.get::<i32, _>("port") as u16,
+    })
+}

src/db/short_link_repo.rs

@@ -0,0 +1,35 @@
+use crate::{db::pool, error::Result};
+
+pub async fn count() -> Result<u64> {
+    let count = sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM short_links")
+        .fetch_one(&pool().await?)
+        .await?;
+    Ok(count.max(0) as u64)
+}
+
+pub async fn insert(short_key: &str, long_url: &str) -> Result<bool> {
+    let result = sqlx::query("INSERT IGNORE INTO short_links (short_key, long_url) VALUES (?, ?)")
+        .bind(short_key)
+        .bind(long_url)
+        .execute(&pool().await?)
+        .await?;
+    Ok(result.rows_affected() == 1)
+}
+
+pub async fn get(short_key: &str) -> Result<Option<String>> {
+    Ok(
+        sqlx::query_scalar::<_, String>("SELECT long_url FROM short_links WHERE short_key = ?")
+            .bind(short_key)
+            .fetch_optional(&pool().await?)
+            .await?,
+    )
+}
+
+pub async fn delete_expired() -> Result<u64> {
+    let result = sqlx::query(
+        "DELETE FROM short_links WHERE created_at < CURRENT_TIMESTAMP - INTERVAL 7 DAY",
+    )
+    .execute(&pool().await?)
+    .await?;
+    Ok(result.rows_affected())
+}

src/db/user_repo.rs

@@ -0,0 +1,194 @@
+use crate::{
+    db::pool,
+    error::{OmegaError, Result},
+    models::{IotaId, User, UserId},
+};
+use mtp::crypto::PublicKeyBundle;
+use sqlx::FromRow;
+
+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;
+    Ok(UserId::from(id.max(1)))
+}
+
+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 = ?";
+
+#[derive(FromRow)]
+struct UserRow {
+    id: i64,
+    iota_id: i64,
+    username: Vec<u8>,
+    display: Option<Vec<u8>>,
+    status: Option<Vec<u8>>,
+    about: Option<Vec<u8>>,
+    avatar: Option<Vec<u8>>,
+    sub_level: i32,
+    sub_end: i64,
+    public_key: Vec<u8>,
+    token: Vec<u8>,
+}
+
+impl TryFrom<UserRow> for User {
+    type Error = sqlx::Error;
+
+    fn try_from(row: UserRow) -> std::result::Result<Self, Self::Error> {
+        let public_key = PublicKeyBundle::from_bytes(&row.public_key)
+            .map_err(|error| sqlx::Error::Decode(Box::new(error)))?;
+        let decode =
+            |value| String::from_utf8(value).map_err(|error| sqlx::Error::Decode(Box::new(error)));
+        Ok(User {
+            id: row.id.into(),
+            iota_id: row.iota_id.into(),
+            username: decode(row.username)?,
+            display: row.display.map(decode).transpose()?,
+            status: row.status.map(decode).transpose()?,
+            about: row.about.map(decode).transpose()?,
+            avatar: row.avatar,
+            sub_level: row.sub_level,
+            sub_end: row.sub_end,
+            public_key,
+            token: decode(row.token)?,
+        })
+    }
+}
+
+pub async fn get_by_username(username: &str) -> Result<User> {
+    let row = sqlx::query_as::<_, UserRow>(USER_BY_USERNAME_QUERY)
+        .bind(username)
+        .fetch_optional(&pool().await?)
+        .await?
+        .ok_or(OmegaError::NotFound)?;
+    row.try_into().map_err(OmegaError::from)
+}
+
+pub async fn get_by_user_id(id: UserId) -> Result<User> {
+    let row = sqlx::query_as::<_, UserRow>(USER_BY_ID_QUERY)
+        .bind(id.0)
+        .fetch_optional(&pool().await?)
+        .await?
+        .ok_or(OmegaError::NotFound)?;
+    row.try_into().map_err(OmegaError::from)
+}
+
+pub async fn get_users_by_iota_id(id: IotaId) -> Result<Vec<User>> {
+    let rows = sqlx::query_as::<_, UserRow>(USERS_BY_IOTA_ID_QUERY)
+        .bind(id.0)
+        .fetch_all(&pool().await?)
+        .await?;
+    rows.into_iter()
+        .map(|row| row.try_into().map_err(OmegaError::from))
+        .collect()
+}
+
+async fn update(
+    id: UserId,
+    query: &'static str,
+    value: impl Send + sqlx::Encode<'static, sqlx::MySql> + sqlx::Type<sqlx::MySql> + 'static,
+) -> Result<()> {
+    sqlx::query(query)
+        .bind(value)
+        .bind(id.0)
+        .execute(&pool().await?)
+        .await?;
+    Ok(())
+}
+
+pub async fn change_username(id: UserId, value: String) -> Result<()> {
+    update(
+        id,
+        "UPDATE users SET username = ? WHERE id = ?",
+        value.into_bytes(),
+    )
+    .await
+}
+pub async fn change_display_name(id: UserId, value: String) -> Result<()> {
+    update(
+        id,
+        "UPDATE users SET display = ? WHERE id = ?",
+        value.into_bytes(),
+    )
+    .await
+}
+pub async fn change_avatar(id: UserId, value: String) -> Result<()> {
+    update(
+        id,
+        "UPDATE users SET avatar = ? WHERE id = ?",
+        value.into_bytes(),
+    )
+    .await
+}
+pub async fn change_about(id: UserId, value: String) -> Result<()> {
+    update(
+        id,
+        "UPDATE users SET about = ? WHERE id = ?",
+        value.into_bytes(),
+    )
+    .await
+}
+pub async fn change_status(id: UserId, value: String) -> Result<()> {
+    update(
+        id,
+        "UPDATE users SET status = ? WHERE id = ?",
+        value.into_bytes(),
+    )
+    .await
+}
+
+pub async fn change_iota_id(id: UserId, value: IotaId) -> Result<()> {
+    sqlx::query("UPDATE users SET iota_id = ? WHERE id = ?")
+        .bind(value.0)
+        .bind(id.0)
+        .execute(&pool().await?)
+        .await?;
+    Ok(())
+}
+pub async fn change_token(id: UserId, value: String) -> Result<()> {
+    update(
+        id,
+        "UPDATE users SET token = ? WHERE id = ?",
+        value.into_bytes(),
+    )
+    .await
+}
+pub async fn delete_user(id: UserId) -> Result<()> {
+    sqlx::query("DELETE FROM users WHERE id = ?")
+        .bind(id.0)
+        .execute(&pool().await?)
+        .await?;
+    Ok(())
+}
+
+pub async fn change_keys(id: UserId, public_key: PublicKeyBundle) -> Result<()> {
+    sqlx::query("UPDATE users SET public_key = ? WHERE id = ?")
+        .bind(public_key.as_bytes())
+        .bind(id.0)
+        .execute(&pool().await?)
+        .await?;
+    Ok(())
+}
+
+pub async fn register_complete_user(
+    id: UserId,
+    username: String,
+    public_key: PublicKeyBundle,
+    iota_id: IotaId,
+    token: String,
+) -> Result<()> {
+    sqlx::query(
+        "INSERT INTO users (id, username, public_key, iota_id, token) VALUES (?, ?, ?, ?, ?)",
+    )
+    .bind(id.0)
+    .bind(username.into_bytes())
+    .bind(public_key.as_bytes())
+    .bind(iota_id.0)
+    .bind(token.into_bytes())
+    .execute(&pool().await?)
+    .await?;
+    Ok(())
+}

src/error.rs

@@ -0,0 +1,57 @@
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum OmegaError {
+    #[error("database pool is not initialized")]
+    DatabaseNotInitialized,
+    #[error("database error: {0}")]
+    Database(sqlx::Error),
+    #[error("invalid input: {0}")]
+    Validation(String),
+    #[error("resource not found")]
+    NotFound,
+    #[error("transport error: {0}")]
+    Transport(String),
+    #[error("not connected")]
+    NotConnected,
+    #[error("not authenticated")]
+    NotAuthenticated,
+    #[error("invalid response")]
+    InvalidResponse,
+    #[error("authentication failed")]
+    AuthenticationFailed,
+    #[error("send error: {0}")]
+    SendError(String),
+    #[error("I/O error: {0}")]
+    Io(#[from] std::io::Error),
+}
+
+pub type Result<T> = std::result::Result<T, OmegaError>;
+
+impl From<sqlx::Error> for OmegaError {
+    fn from(error: sqlx::Error) -> Self {
+        if matches!(error, sqlx::Error::RowNotFound) {
+            Self::NotFound
+        } else {
+            Self::Database(error)
+        }
+    }
+}
+
+impl OmegaError {
+    pub fn status_code(&self) -> http::StatusCode {
+        match self {
+            Self::Validation(_) => http::StatusCode::BAD_REQUEST,
+            Self::NotFound => http::StatusCode::NOT_FOUND,
+            Self::DatabaseNotInitialized
+            | Self::Database(_)
+            | Self::Transport(_)
+            | Self::NotConnected
+            | Self::NotAuthenticated
+            | Self::InvalidResponse
+            | Self::AuthenticationFailed
+            | Self::SendError(_)
+            | Self::Io(_) => http::StatusCode::INTERNAL_SERVER_ERROR,
+        }
+    }
+}

src/main.rs

@@ -1,30 +1,41 @@
+mod api;
+mod config;
+mod db;
+pub mod error;
+mod models;
 mod server;
 mod sql;
 mod transport;
 mod util;
 
-use crate::sql::sql::initialize_db;
+pub use error::{OmegaError, Result};
-use crate::sql::sql::print_users;
+
+use crate::db::initialize;
 use crate::transport::omikron_connection;
 use crate::util::file_util::get_directory;
 use crate::util::logger::PrintType;
 use crate::util::logger::startup;
 use dotenv::from_path;
-use mtp::files::{load_keyring_raw, save_keyring_raw, save_public_key_bundle};
 use mtp::crypto::Keyring;
+use mtp::files::{load_keyring_raw, save_keyring_raw, save_public_key_bundle};
 use once_cell::sync::Lazy;
 use rustls::crypto::aws_lc_rs::default_provider;
 use std::env;
 use std::path::Path;
+use std::time::Duration;
+use tokio::time::interval;
 
 const KEYRING_PATH: &str = "./omega.mk";
 
 static KEYRING: Lazy<Keyring> = Lazy::new(|| {
     load_keyring_raw(KEYRING_PATH).unwrap_or_else(|_| {
         let kr = Keyring::generate();
-        save_keyring_raw(&kr, KEYRING_PATH).expect("Failed to save generated keyring");
+        if let Err(error) = save_keyring_raw(&kr, KEYRING_PATH) {
-        save_public_key_bundle(&kr.public_key_bundle(), KEYRING_PATH)
+            eprintln!("Failed to save generated keyring: {error}");
-            .expect("Failed to save generated public key bundle");
+        }
+        if let Err(error) = save_public_key_bundle(&kr.public_key_bundle(), KEYRING_PATH) {
+            eprintln!("Failed to save generated public key bundle: {error}");
+        }
         eprintln!("Generated new keyring at {}", KEYRING_PATH);
         kr
     })
@@ -34,7 +45,10 @@ pub fn get_keyring() -> &'static Keyring {
     &KEYRING
 }
 pub fn load_keyring() -> Keyring {
-    Keyring::from_bytes(&KEYRING.to_bytes()).unwrap()
+    Keyring::from_bytes(&KEYRING.to_bytes()).unwrap_or_else(|error| {
+        eprintln!("Failed to clone keyring: {error}");
+        Keyring::generate()
+    })
 }
 
 #[tokio::main]
@@ -50,7 +64,7 @@ async fn main() {
 
     log!("Started");
     log!("  .env");
-    if let Err(e) = initialize_db().await {
+    if let Err(e) = initialize().await {
         log!("[FATAL] Database initialization failed: {}", e);
         log!(
             "[FATAL] Please ensure the database is running and the .env file is configured correctly."
@@ -59,12 +73,21 @@ async fn main() {
     } else {
         log!("  DB");
     }
-    if let Err(e) = print_users().await {
+    let rate_limit_cleanup = crate::server::middleware::spawn_cleanup_task();
-        log!("[ERROR] Failed to print users: {}", e);
+    let short_link_cleanup = tokio::spawn(async {
-    } else {
+        let mut ticker = interval(Duration::from_secs(24 * 60 * 60));
-        log!("  Users");
+        loop {
-    }
+            ticker.tick().await;
-
+            if let Err(error) = crate::db::short_link_repo::delete_expired().await {
+                log_err!(
+                    0,
+                    PrintType::General,
+                    "Short-link cleanup failed: {}",
+                    error
+                );
+            }
+        }
+    });
     let port: u16 = env::var("PORT")
         .ok()
         .and_then(|s| s.parse().ok())
@@ -80,4 +103,6 @@ async fn main() {
             log!("Shutting down on signal...");
         }
     }
+    rate_limit_cleanup.abort();
+    short_link_cleanup.abort();
 }

src/models/ids.rs

@@ -0,0 +1,30 @@
+use std::fmt::{Display, Formatter};
+
+macro_rules! id_type {
+    ($name:ident) => {
+        #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, serde::Serialize)]
+        pub struct $name(pub i64);
+
+        impl From<i64> for $name {
+            fn from(value: i64) -> Self {
+                Self(value)
+            }
+        }
+
+        impl From<$name> for i64 {
+            fn from(value: $name) -> Self {
+                value.0
+            }
+        }
+
+        impl Display for $name {
+            fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
+                self.0.fmt(formatter)
+            }
+        }
+    };
+}
+
+id_type!(UserId);
+id_type!(IotaId);
+id_type!(OmikronId);

src/models/iota.rs

@@ -0,0 +1,9 @@
+use super::IotaId;
+use mtp::crypto::PublicKeyBundle;
+
+#[derive(Clone, Debug, serde::Serialize)]
+pub struct Iota {
+    pub id: IotaId,
+    #[serde(serialize_with = "crate::models::serialize_public_key")]
+    pub public_key: PublicKeyBundle,
+}

src/models/mod.rs

@@ -0,0 +1,21 @@
+mod ids;
+mod iota;
+mod notification;
+mod omikron;
+mod user;
+
+fn serialize_public_key<S>(
+    key: &mtp::crypto::PublicKeyBundle,
+    serializer: S,
+) -> std::result::Result<S::Ok, S::Error>
+where
+    S: serde::Serializer,
+{
+    serializer.serialize_str(&key.to_base64())
+}
+
+pub use ids::{IotaId, OmikronId, UserId};
+pub use iota::Iota;
+pub use notification::Notification;
+pub use omikron::Omikron;
+pub use user::User;

src/models/notification.rs

@@ -0,0 +1,9 @@
+use super::UserId;
+
+#[derive(Clone, Debug, serde::Serialize)]
+pub struct Notification {
+    pub id: i64,
+    pub sender_id: UserId,
+    pub receiver_id: UserId,
+    pub amount: i64,
+}

src/models/omikron.rs

@@ -0,0 +1,12 @@
+use super::OmikronId;
+use mtp::crypto::PublicKeyBundle;
+
+#[derive(Clone, Debug, serde::Serialize)]
+pub struct Omikron {
+    pub id: OmikronId,
+    #[serde(serialize_with = "crate::models::serialize_public_key")]
+    pub public_key: PublicKeyBundle,
+    pub location: String,
+    pub ip_address: String,
+    pub port: u16,
+}

src/models/user.rs

@@ -0,0 +1,19 @@
+use super::{IotaId, UserId};
+use mtp::crypto::PublicKeyBundle;
+
+#[derive(Clone, Debug, serde::Serialize)]
+pub struct User {
+    pub id: UserId,
+    pub iota_id: IotaId,
+    pub username: String,
+    pub display: Option<String>,
+    pub status: Option<String>,
+    pub about: Option<String>,
+    pub avatar: Option<Vec<u8>>,
+    pub sub_level: i32,
+    pub sub_end: i64,
+    #[serde(serialize_with = "crate::models::serialize_public_key")]
+    pub public_key: PublicKeyBundle,
+    #[serde(skip_serializing)]
+    pub token: String,
+}

src/server/api.rs

@@ -1,356 +1,240 @@
+use crate::api::response::{
+    ConnectionsResponse, IotaResponse, OmikronResponse, PublicKeyResponse, StatusResponse,
+    UserResponse, UsernameResponse, json,
+};
+use crate::db::{
+    iota_repo::get_iota_by_id,
+    omikron_repo::get_omikron_by_id,
+    user_repo::{get_by_user_id, get_by_username},
+};
+use crate::error::{OmegaError, Result};
 use crate::load_keyring;
-use crate::sql::sql;
+use crate::models::UserId;
-use crate::sql::sql::{get_by_user_id, get_iota_by_id, get_omikron_by_id};
+use crate::server::{
+    middleware,
+    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::util::file_util::get_directory;
 use base64::Engine as _;
 use bytes::Bytes;
 use http::{Method, StatusCode};
-use json::JsonValue;
 use mtp::webserver::{Http3Request, Http3Response, RouteParams};
+use std::collections::BTreeMap;
+
+fn error_body(error: &OmegaError) -> String {
+    json(&StatusResponse {
+        status: match error {
+            OmegaError::Validation(_) => "error_bad_request",
+            OmegaError::NotFound => "error_not_found",
+            _ => "error",
+        },
+    })
+}
+
+fn user_response(user: crate::models::User) -> UserResponse {
+    UserResponse {
+        status: "success",
+        username: user.username,
+        public_key: user.public_key.to_base64(),
+        user_id: user.id.0,
+        iota_id: user.iota_id.0,
+        sub_level: user.sub_level,
+        sub_end: user.sub_end,
+        display: user.display,
+        status_message: user.status,
+        about: user.about,
+        avatar: user
+            .avatar
+            .map(|value| base64::engine::general_purpose::STANDARD.encode(value)),
+    }
+}
+
+async fn route(path_parts: &[&str]) -> Result<(StatusCode, String)> {
+    match path_parts {
+        ["api", "get", "omikron"] => {
+            let connection = get_random_omikron()
+                .await
+                .map_err(|_| OmegaError::NotFound)?;
+            let id = connection
+                .get_omikron_id()
+                .await
+                .ok_or(OmegaError::NotFound)?;
+            let omikron = get_omikron_by_id(id.into()).await?;
+            Ok((
+                StatusCode::OK,
+                json(&OmikronResponse {
+                    status: "success",
+                    id,
+                    public_key: omikron.public_key.to_base64(),
+                    ip_address: omikron.ip_address,
+                    port: omikron.port,
+                }),
+            ))
+        }
+        ["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?
+                }
+            };
+            Ok((
+                StatusCode::OK,
+                json(&OmikronResponse {
+                    status: "success",
+                    id: omikron.id.0,
+                    public_key: omikron.public_key.to_base64(),
+                    ip_address: omikron.ip_address,
+                    port: omikron.port,
+                }),
+            ))
+        }
+        ["api", "get", "connections"] => {
+            let connections = get_all_connections()
+                .await
+                .map_err(|_| OmegaError::Transport("failed to load connections".to_string()))?;
+            let connections = connections
+                .into_iter()
+                .map(|(omikron_id, iotas)| {
+                    let iotas = iotas
+                        .into_iter()
+                        .map(|(iota_id, users)| {
+                            (
+                                iota_id.to_string(),
+                                users.into_iter().map(i64::from).collect(),
+                            )
+                        })
+                        .collect();
+                    (omikron_id.to_string(), iotas)
+                })
+                .collect::<BTreeMap<_, _>>();
+            Ok((
+                StatusCode::OK,
+                json(&ConnectionsResponse {
+                    status: "success",
+                    connections,
+                }),
+            ))
+        }
+        ["api", "get", "iota", id] => {
+            let id = parse_positive_id(id)?;
+            let iota = get_iota_by_id(id.into()).await?;
+            Ok((
+                StatusCode::OK,
+                json(&IotaResponse {
+                    status: "success",
+                    iota_id: iota.id.0,
+                    public_key: iota.public_key.to_base64(),
+                }),
+            ))
+        }
+        ["api", "get", "id", username] => {
+            validate_non_empty(username, "username", 15)?;
+            let user = get_by_username(username).await?;
+            Ok((
+                StatusCode::OK,
+                json(&UsernameResponse {
+                    status: "success",
+                    username: user.username,
+                    public_key: user.public_key.to_base64(),
+                    user_id: user.id.0,
+                    iota_id: user.iota_id.0,
+                    sub_level: user.sub_level,
+                    sub_end: user.sub_end,
+                }),
+            ))
+        }
+        ["api", "get", "public_key"] => {
+            let public_key = base64::engine::general_purpose::STANDARD
+                .encode(load_keyring().public_key_bundle().as_bytes());
+            Ok((
+                StatusCode::OK,
+                json(&PublicKeyResponse {
+                    status: "success",
+                    public_key,
+                }),
+            ))
+        }
+        ["api", "get", "user", id] => {
+            let id = parse_positive_id(id)?;
+            let user = get_by_user_id(UserId::from(id)).await?;
+            Ok((StatusCode::OK, json(&user_response(user))))
+        }
+        _ => Ok((
+            StatusCode::INTERNAL_SERVER_ERROR,
+            json(&StatusResponse { status: "error" }),
+        )),
+    }
+}
 
 pub async fn handle(request: Http3Request, response: Http3Response) -> Http3Response {
     let method = request.method;
     let path = request.uri.path().to_string();
-    let body_string = request
+    if method != Method::OPTIONS && !middleware::allow(request.remote_addr.ip(), &path) {
-        .body
+        return response
-        .map(|body| String::from_utf8_lossy(&body).to_string());
+            .status(StatusCode::TOO_MANY_REQUESTS)
-
+            .header("access-control-allow-origin", &crate::config::cors_origin())
+            .body(json(&StatusResponse {
+                status: "error_rate_limited",
+            }));
+    }
     if method == Method::OPTIONS {
         return response
             .status(StatusCode::OK)
-            .header("access-control-allow-origin", "*")
+            .header("access-control-allow-origin", &crate::config::cors_origin())
             .header("access-control-allow-methods", "GET, POST, OPTIONS")
             .header("access-control-allow-headers", "*");
     }
-
+    let path_parts: Vec<&str> = path.split('/').filter(|part| !part.is_empty()).collect();
-    let path_parts: Vec<&str> = path.split('/').filter(|s| !s.is_empty()).collect();
+    if let ["api", "download", "iota_frontend"] = path_parts.as_slice() {
-
+        let file_path = format!("{}/downloads/iota_frontend.zip", get_directory());
-    let _body: Option<JsonValue> = if let Some(ref bs) = body_string {
+        return match std::fs::read(file_path) {
-        if let Ok(body_json) = json::parse(bs) {
+            Ok(bytes) => response
-            Some(body_json)
+                .status(StatusCode::OK)
-        } else {
+                .header("access-control-allow-origin", &crate::config::cors_origin())
-            None
+                .header("content-type", "application/zip")
-        }
+                .header(
-    } else {
+                    "content-disposition",
-        None
+                    "attachment; filename=\"iota_frontend.zip\"",
-    };
+                )
-
+                .body(Bytes::from(bytes)),
-    let (status, body_text) = match path_parts.as_slice() {
+            Err(_) => response
-        // ==================================================
+                .status(StatusCode::NOT_FOUND)
-        // DOWNLOAD IOTA FRONTEND
+                .header("access-control-allow-origin", &crate::config::cors_origin())
-        // ==================================================
+                .body(json(&StatusResponse {
-        ["api", "download", "iota_frontend"] => {
+                    status: "error_not_found",
-            let file_path = format!("{}/downloads/iota_frontend.zip", get_directory());
+                })),
-
+        };
-            match std::fs::read(file_path) {
+    }
-                Ok(file_bytes) => {
+    if let ["direct", short @ ..] = path_parts.as_slice() {
-                    return response
+        let short = short.join("");
-                        .status(StatusCode::OK)
+        let location = crate::server::short_link::get_short_link(&short)
-                        .header("access-control-allow-origin", "*")
+            .await
-                        .header("content-type", "application/zip")
+            .unwrap_or_else(|_| "https://tensamin.net".to_string());
-                        .header(
+        return response
-                            "content-disposition",
+            .status(StatusCode::TEMPORARY_REDIRECT)
-                            "attachment; filename=\"iota_frontend.zip\"",
+            .header("location", &location);
-                        )
+    }
-                        .body(Bytes::from(file_bytes));
+    let (status, body) = route(&path_parts)
-                }
+        .await
-                Err(_) => {
+        .unwrap_or_else(|error| (error.status_code(), error_body(&error)));
-                    let mut res = JsonValue::new_object();
-                    res["status"] = "error_not_found".into();
-                    return response
-                        .status(StatusCode::NOT_FOUND)
-                        .header("access-control-allow-origin", "*")
-                        .body(res.dump());
-                }
-            }
-        }
-
-        // ==================================================
-        // GET RANDOM OMIKRON
-        // ==================================================
-        ["api", "get", "omikron"] => {
-            if let Ok(omikron_conn) = get_random_omikron().await {
-                if let Some(id) = omikron_conn.get_omikron_id().await {
-                    if let Ok((public_key, ip_address, port)) = sql::get_omikron_by_id(id).await {
-                        let mut res = JsonValue::new_object();
-                        res["status"] = "success".into();
-                        res["id"] = id.into();
-                        res["public_key"] = public_key.to_base64().into();
-                        res["ip_address"] = ip_address.into();
-                        res["port"] = port.into();
-
-                        (StatusCode::OK, res.dump())
-                    } else {
-                        let mut res = JsonValue::new_object();
-                        res["status"] = "error".into();
-                        (StatusCode::INTERNAL_SERVER_ERROR, res.dump())
-                    }
-                } else {
-                    let mut res = JsonValue::new_object();
-                    res["status"] = "error".into();
-                    (StatusCode::INTERNAL_SERVER_ERROR, res.dump())
-                }
-            } else {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error_not_found".into();
-                (StatusCode::NOT_FOUND, res.dump())
-            }
-        }
-
-        // ==================================================
-        // GET OMIKRON BY ID
-        // ==================================================
-        ["api", "get", "omikron", id] => {
-            let id = id.parse::<i64>().unwrap_or(0);
-
-            if id == 0 {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error_bad_request".into();
-                (StatusCode::BAD_REQUEST, res.dump())
-            } else if let Ok((public_key, ip_address, port)) = get_omikron_by_id(id).await {
-                let mut res = JsonValue::new_object();
-                res["status"] = "success".into();
-                res["id"] = id.into();
-                res["public_key"] = public_key.to_base64().into();
-                res["ip_address"] = ip_address.into();
-                res["port"] = port.into();
-                (StatusCode::OK, res.dump())
-            } else if let Some(omikron_id) = get_iota_primary_omikron_connection(id) {
-                if let Ok((public_key, ip_address, port)) = get_omikron_by_id(omikron_id).await {
-                    let mut res = JsonValue::new_object();
-                    res["status"] = "success".into();
-                    res["id"] = omikron_id.into();
-                    res["public_key"] = public_key.to_base64().into();
-                    res["ip_address"] = ip_address.into();
-                    res["port"] = port.into();
-                    (StatusCode::OK, res.dump())
-                } else {
-                    let mut res = JsonValue::new_object();
-                    res["status"] = "error_not_found".into();
-                    (StatusCode::NOT_FOUND, res.dump())
-                }
-            } else if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) = get_by_user_id(id).await
-            {
-                if let Some(omikron_id) = get_iota_primary_omikron_connection(iota_id) {
-                    if let Ok((public_key, ip_address, port)) = get_omikron_by_id(omikron_id).await
-                    {
-                        let mut res = JsonValue::new_object();
-                        res["status"] = "success".into();
-                        res["id"] = omikron_id.into();
-                        res["public_key"] = public_key.to_base64().into();
-                        res["ip_address"] = ip_address.into();
-                        res["port"] = port.into();
-                        (StatusCode::OK, res.dump())
-                    } else {
-                        let mut res = JsonValue::new_object();
-                        res["status"] = "error_not_found".into();
-                        (StatusCode::NOT_FOUND, res.dump())
-                    }
-                } else {
-                    let mut res = JsonValue::new_object();
-                    res["status"] = "error_not_found".into();
-                    (StatusCode::NOT_FOUND, res.dump())
-                }
-            } else {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error_not_found".into();
-                (StatusCode::NOT_FOUND, res.dump())
-            }
-        }
-
-        ["api", "get", "connections"] => {
-            if let Ok(connections) = get_all_connections().await {
-                let mut res = JsonValue::new_object();
-                res["status"] = "success".into();
-
-                for (omikron_id, iota_map) in connections {
-                    let mut omikron_obj = JsonValue::new_object();
-                    for (iota_id, user_ids) in iota_map {
-                        let mut user_arr = JsonValue::new_array();
-                        for user_id in user_ids {
-                            let _ = user_arr.push(user_id);
-                        }
-                        omikron_obj[&iota_id.to_string()] = user_arr;
-                    }
-                    res[&omikron_id.to_string()] = omikron_obj;
-                }
-
-                (StatusCode::OK, res.dump())
-            } else {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error".into();
-                (StatusCode::INTERNAL_SERVER_ERROR, res.dump())
-            }
-        }
-
-        // ==================================================
-        // GET IOTA BY ID
-        // ==================================================
-        ["api", "get", "iota", id] => {
-            let id: i64 = id.parse().unwrap_or(0);
-
-            if id == 0 {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error_bad_request".into();
-                (StatusCode::BAD_REQUEST, res.dump())
-            } else if let Ok((id, public_key)) = get_iota_by_id(id).await {
-                let mut res = JsonValue::new_object();
-                res["status"] = "success".into();
-                res["iota_id"] = id.into();
-                res["public_key"] = public_key.to_base64().into();
-                (StatusCode::OK, res.dump())
-            } else {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error_not_found".into();
-                (StatusCode::NOT_FOUND, res.dump())
-            }
-        }
-
-        // ==================================================
-        // GET ID BY USERNAME
-        // ==================================================
-        ["api", "get", "id", username] => {
-            if username.is_empty() {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error_bad_request".into();
-                (StatusCode::BAD_REQUEST, res.dump())
-            } else if let Ok((
-                id,
-                iota_id,
-                username,
-                _,
-                _,
-                _,
-                _,
-                sub_level,
-                sub_end,
-                public_key,
-                _,
-                _,
-            )) = sql::get_by_username(username).await
-            {
-                let mut res = JsonValue::new_object();
-                res["status"] = "success".into();
-                res["username"] = username.into();
-                res["public_key"] = public_key.to_base64().into();
-                res["user_id"] = id.into();
-                res["iota_id"] = iota_id.into();
-                res["sub_level"] = sub_level.into();
-                res["sub_end"] = sub_end.into();
-
-                (StatusCode::OK, res.dump())
-            } else {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error_not_found".into();
-                (StatusCode::OK, res.dump())
-            }
-        }
-
-        // ==================================================
-        // GET SERVER PUBLIC KEY
-        // ==================================================
-        ["api", "get", "public_key"] => {
-            let mut res = JsonValue::new_object();
-            res["status"] = "success".into();
-            let bundle = load_keyring().public_key_bundle();
-            res["public_key"] = base64::engine::general_purpose::STANDARD
-                .encode(bundle.as_bytes())
-                .into();
-            (StatusCode::OK, res.dump())
-        }
-
-        // ==================================================
-        // GET USER BY ID
-        // ==================================================
-        ["api", "get", "user", id] => {
-            let id: i64 = id.parse().unwrap_or(0);
-
-            if id == 0 {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error_bad_request".into();
-                (StatusCode::BAD_REQUEST, res.dump())
-            } else if let Ok((
-                id,
-                iota_id,
-                username,
-                display,
-                status_msg,
-                about,
-                avatar,
-                sub_level,
-                sub_end,
-                public_key,
-                _,
-                _,
-            )) = sql::get_by_user_id(id).await
-            {
-                let mut res = JsonValue::new_object();
-                res["status"] = "success".into();
-                res["username"] = username.into();
-                res["public_key"] = public_key.to_base64().into();
-                res["user_id"] = id.into();
-                res["iota_id"] = iota_id.into();
-                res["sub_level"] = sub_level.into();
-                res["sub_end"] = sub_end.into();
-
-                if let Some(display) = display {
-                    res["display"] = display.into();
-                }
-                if let Some(status_msg) = status_msg {
-                    res["status_message"] = status_msg.into();
-                }
-                if let Some(about) = about {
-                    res["about"] = about.into();
-                }
-                if let Some(avatar) = avatar {
-                    res["avatar"] = base64::engine::general_purpose::STANDARD
-                        .encode(avatar)
-                        .into();
-                }
-
-                (StatusCode::OK, res.dump())
-            } else {
-                let mut res = JsonValue::new_object();
-                res["status"] = "error_not_found".into();
-                (StatusCode::OK, res.dump())
-            }
-        }
-
-        // ==================================================
-        // DIRECT - SHORT LINK RESOLUTION
-        // ==================================================
-        ["direct", short @ ..] => {
-            let short_str = short.join("/");
-            let short = short_str.replace("/", "");
-            if let Ok(long) = crate::server::short_link::get_short_link(&short).await {
-                return response
-                    .status(StatusCode::TEMPORARY_REDIRECT)
-                    .header("location", &long);
-            } else {
-                return response
-                    .status(StatusCode::TEMPORARY_REDIRECT)
-                    .header("location", "https://tensamin.net");
-            }
-        }
-
-        // ==================================================
-        // DEFAULT
-        // ==================================================
-        _ => {
-            let mut res = JsonValue::new_object();
-            res["status"] = "error".into();
-            (StatusCode::INTERNAL_SERVER_ERROR, res.dump())
-        }
-    };
-
     response
         .status(status)
-        .header("access-control-allow-origin", "*")
+        .header("access-control-allow-origin", &crate::config::cors_origin())
         .header("access-control-allow-headers", "*")
         .header("access-control-allow-methods", "GET, POST, OPTIONS")
-        .body(body_text)
+        .body(body)
 }
 
 pub async fn handle_pattern(

src/server/middleware.rs

@@ -0,0 +1,84 @@
+use dashmap::DashMap;
+use once_cell::sync::Lazy;
+use std::net::IpAddr;
+use std::{collections::VecDeque, time::Instant};
+use tokio::time::interval;
+
+static REQUESTS: Lazy<DashMap<(IpAddr, String), VecDeque<Instant>>> = Lazy::new(DashMap::new);
+static CONFIG: Lazy<crate::config::RateLimitConfig> =
+    Lazy::new(crate::config::RateLimitConfig::from_env);
+const MAX_TRACKED_CLIENT_BUCKETS: usize = 100_000;
+
+pub fn allow(remote_addr: IpAddr, path: &str) -> bool {
+    let key = if path.contains("register") {
+        "registration"
+    } else {
+        "general"
+    };
+    let map_key = (remote_addr, key.to_string());
+    if REQUESTS.len() >= MAX_TRACKED_CLIENT_BUCKETS {
+        cleanup_expired();
+        if REQUESTS.len() >= MAX_TRACKED_CLIENT_BUCKETS && !REQUESTS.contains_key(&map_key) {
+            return false;
+        }
+    }
+    let limit = if key == "registration" {
+        CONFIG.registration_requests
+    } else {
+        CONFIG.general_requests
+    };
+    let now = Instant::now();
+    let mut entries = REQUESTS.entry(map_key).or_default();
+    while entries
+        .front()
+        .is_some_and(|time| now.duration_since(*time) >= CONFIG.window)
+    {
+        entries.pop_front();
+    }
+    if entries.len() >= limit {
+        return false;
+    }
+    entries.push_back(now);
+    true
+}
+
+pub fn spawn_cleanup_task() -> tokio::task::JoinHandle<()> {
+    tokio::spawn(async {
+        let mut ticker = interval(CONFIG.window);
+        loop {
+            ticker.tick().await;
+            cleanup_expired();
+        }
+    })
+}
+
+fn cleanup_expired() {
+    let now = Instant::now();
+    REQUESTS.retain(|_, entries| {
+        while entries
+            .front()
+            .is_some_and(|time| now.duration_since(*time) >= CONFIG.window)
+        {
+            entries.pop_front();
+        }
+        !entries.is_empty()
+    });
+}
+
+#[cfg(test)]
+mod tests {
+    use super::allow;
+    use std::net::{IpAddr, Ipv4Addr};
+
+    #[test]
+    fn tracks_clients_independently() {
+        let first: IpAddr = "192.0.2.10"
+            .parse()
+            .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
+        let second: IpAddr = "192.0.2.11"
+            .parse()
+            .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
+        assert!(allow(first, "/api/get/user/1"));
+        assert!(allow(second, "/api/get/user/1"));
+    }
+}

src/server/mod.rs

@@ -1,4 +1,6 @@
 pub mod api;
 pub mod index;
+pub mod middleware;
 pub mod server;
 pub mod short_link;
+pub mod validation;

src/server/short_link.rs

@@ -1,37 +1,27 @@
-use dashmap::DashMap;
+use crate::db::short_link_repo;
-use once_cell::sync::Lazy;
+use rand::RngExt;
-use rand::{Rng, thread_rng};
-
-static LINKS: Lazy<DashMap<String, String>> = Lazy::new(DashMap::new);
 
 const CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHJKLMNPRSTUVWXYZ1234567890";
 
 pub async fn add_short_link(long: &str) -> Result<String, ()> {
-    let raw = generate_unique_short_link().await;
-    LINKS.insert(raw.clone(), long.to_string());
-
-    Ok(format!(
-        "hmtps://omega.tensamin.net/direct/{}",
-        format_with_dashes(&raw)
-    ))
-}
-
-async fn generate_unique_short_link() -> String {
     loop {
-        let short = generate_short_link().await;
+        let raw = generate_short_link().await;
-        if !LINKS.contains_key(&short) {
+        if short_link_repo::insert(&raw, long).await.map_err(|_| ())? {
-            return short;
+            return Ok(format!(
+                "https://omega.tensamin.net/direct/{}",
+                format_with_dashes(&raw)
+            ));
         }
     }
 }
 
 pub async fn generate_short_link() -> String {
-    let len = short_length();
+    let len = short_length().await;
 
-    let mut rng = thread_rng();
+    let mut rng = rand::rng();
     (0..len)
         .map(|_| {
-            let idx = rng.gen_range(0..CHARSET.len());
+            let idx = rng.random_range(0..CHARSET.len());
             CHARSET[idx] as char
         })
         .collect()
@@ -46,17 +36,17 @@ pub async fn get_short_link(short: &str) -> Result<String, ()> {
     let frag = short.replace(key, "");
     let normalized = normalize_short(&key);
 
-    if let Ok(t) = LINKS.get(&normalized).map(|v| v.value().clone()).ok_or(()) {
+    let target = short_link_repo::get(&normalized)
-        Ok(format!("{}{}", t, frag))
+        .await
-    } else {
+        .map_err(|_| ())?
-        Err(())
+        .ok_or(())?;
-    }
+    Ok(format!("{}{}", target, frag))
 }
 
 /* ---------------- helpers ---------------- */
 
-fn short_length() -> usize {
+async fn short_length() -> usize {
-    let count = LINKS.len();
+    let count = short_link_repo::count().await.unwrap_or(0);
 
     match count {
         0..=1_999 => 4,

src/server/validation.rs

@@ -0,0 +1,42 @@
+use crate::error::{OmegaError, Result};
+
+pub fn parse_positive_id(value: &str) -> Result<i64> {
+    let id = value
+        .parse::<i64>()
+        .map_err(|_| OmegaError::Validation("ID must be a positive integer".to_string()))?;
+    if id <= 0 {
+        return Err(OmegaError::Validation(
+            "ID must be a positive integer".to_string(),
+        ));
+    }
+    Ok(id)
+}
+
+pub fn validate_non_empty(value: &str, field: &str, max_len: usize) -> Result<()> {
+    let length = value.chars().count();
+    if length == 0 || length > max_len {
+        return Err(OmegaError::Validation(format!(
+            "{field} must contain 1 to {max_len} characters"
+        )));
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{parse_positive_id, validate_non_empty};
+
+    #[test]
+    fn rejects_invalid_ids() {
+        assert!(parse_positive_id("bad").is_err());
+        assert!(parse_positive_id("0").is_err());
+        assert_eq!(parse_positive_id("7").expect("valid test ID"), 7);
+    }
+
+    #[test]
+    fn enforces_string_bounds() {
+        assert!(validate_non_empty("", "username", 15).is_err());
+        assert!(validate_non_empty("abcdefghijklmnop", "username", 15).is_err());
+        assert!(validate_non_empty("alice", "username", 15).is_ok());
+    }
+}

src/sql/mod.rs

@@ -1,3 +1,2 @@
 pub mod connection_status;
-pub mod sql;
 pub mod user_online_tracker;

src/sql/sql.rs

@@ -1,778 +0,0 @@
-use crate::log;
-use mtp::crypto::PublicKeyBundle;
-use once_cell::sync::Lazy;
-use sqlx::{MySql, Pool, Row, mysql::MySqlPoolOptions};
-use std::sync::atomic::{AtomicU64, Ordering};
-use std::{
-    env,
-    sync::Arc,
-    time::{SystemTime, UNIX_EPOCH},
-};
-use tokio::sync::RwLock;
-
-/*
-use crate::sql::{
-    iota_omikron_tracker::{get_omikron_for_iota, track_iota_omikron, untrack_iota},
-    sql::{
-        change_about, change_avatar, change_display_name, change_iota_id, change_iota_key,
-        change_keys, change_status, change_username, get_by_id, get_by_username, get_iota_by_id,
-        get_register_id, register_complete_iota, register_complete_user,
-    },
-};
- */
-
-static SQL_DB: Lazy<Arc<RwLock<Option<Pool<MySql>>>>> = Lazy::new(|| Arc::new(RwLock::new(None)));
-
-pub async fn connect() -> Result<Pool<MySql>, sqlx::Error> {
-    let url = env::var("DB_URL").expect("DB_URL is not set");
-
-    MySqlPoolOptions::new()
-        .max_connections(200)
-        .connect(&url)
-        .await
-}
-// Omega
-// - Omikron
-//   - Iota
-//     - User
-//     - User
-//   - Iota
-//     - User
-//     - User
-// - Omikron
-//   - Iota
-//     - User
-//     - User
-//   - Iota
-//     - User
-//     - User
-pub async fn initialize_db() -> Result<(), sqlx::Error> {
-    let pool = connect().await?;
-    let mut db_lock = SQL_DB.write().await;
-    // create tables
-    // with indexes
-    let _ = sqlx::query(
-        "CREATE TABLE IF NOT EXISTS
-        users (
-        id BIGINT NOT NULL PRIMARY KEY,
-        username VARCHAR(15) NOT NULL UNIQUE COLLATE utf8mb4_bin,
-        display VARCHAR(15) COLLATE utf8mb4_bin,
-        status VARCHAR(15) COLLATE utf8mb4_bin,
-        about VARCHAR(200) COLLATE utf8mb4_bin,
-        avatar MEDIUMBLOB,
-        sub_level INT(11) NOT NULL DEFAULT 0,
-        sub_end BIGINT(20) NOT NULL DEFAULT 0,
-        public_key BLOB NOT NULL,
-        private_key_hash TEXT NOT NULL COLLATE utf8mb4_bin DEFAULT '',
-        iota_id BIGINT NOT NULL,
-        token VARCHAR(256) NOT NULL UNIQUE COLLATE utf8mb4_bin
-        )",
-    )
-    .execute(&pool)
-    .await;
-    let _ = sqlx::query(
-        "CREATE TABLE IF NOT EXISTS
-        iotas (
-        id BIGINT NOT NULL PRIMARY KEY,
-        public_key BLOB NOT NULL
-        )",
-    )
-    .execute(&pool)
-    .await;
-    let _ = sqlx::query(
-        "CREATE TABLE IF NOT EXISTS
-        omikrons (
-        id BIGINT NOT NULL PRIMARY KEY,
-        public_key BLOB NOT NULL,
-        location VARCHAR(255) NOT NULL COLLATE utf8mb4_bin,
-        ip_address VARCHAR(255) NOT NULL COLLATE utf8mb4_bin,
-        port INT(11) NOT NULL DEFAULT 959
-        )",
-    )
-    .execute(&pool)
-    .await;
-    // Retrofits `port` onto omikrons tables created before this column existed;
-    // `CREATE TABLE IF NOT EXISTS` above is a no-op against an already-existing table.
-    let _ = sqlx::query(
-        "ALTER TABLE omikrons ADD COLUMN IF NOT EXISTS port INT(11) NOT NULL DEFAULT 959",
-    )
-    .execute(&pool)
-    .await;
-    let _ = sqlx::query(
-        "CREATE TABLE IF NOT EXISTS
-        notifications (
-        id BIGINT NOT NULL PRIMARY KEY AUTO_INCREMENT,
-        sender_id BIGINT NOT NULL,
-        receiver_id BIGINT NOT NULL,
-        amount BIGINT NOT NULL DEFAULT 0
-        )",
-    )
-    .execute(&pool)
-    .await;
-    *db_lock = Some(pool);
-    Ok(())
-}
-
-// ==========================================================================================
-//                                           REGISTER
-// ==========================================================================================
-pub static CURRENT_MILLI_USED: Lazy<Arc<AtomicU64>> = Lazy::new(|| Arc::new(AtomicU64::new(0)));
-pub static CURRENT_REGISTER_PROCESS: Lazy<Arc<RwLock<Vec<u64>>>> =
-    Lazy::new(|| Arc::new(RwLock::new(Vec::new())));
-
-pub async fn get_register_id() -> u64 {
-    let mut current_time = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap_or_default()
-        .as_millis() as u64;
-
-    loop {
-        let current_locked = CURRENT_MILLI_USED.load(Ordering::SeqCst);
-
-        if current_locked < current_time {
-            let result = CURRENT_MILLI_USED.compare_exchange(
-                current_locked,   // expected value
-                current_time,     // new value
-                Ordering::SeqCst, // acquire/release ordering
-                Ordering::SeqCst, // failure ordering
-            );
-
-            match result {
-                Ok(_) => {
-                    CURRENT_REGISTER_PROCESS.write().await.push(current_time);
-                    return current_time;
-                }
-                Err(_) => {
-                    continue;
-                }
-            }
-        } else {
-            current_time = current_locked + 1;
-        }
-    }
-}
-
-// ==========================================================================================
-//                                           USERS
-// ==========================================================================================
-
-pub async fn get_by_username(
-    username: &str,
-) -> Result<
-    (
-        i64,
-        i64,
-        String,
-        Option<String>,
-        Option<String>,
-        Option<String>,
-        Option<Vec<u8>>,
-        i32,
-        i64,
-        PublicKeyBundle,
-        String,
-        String,
-    ),
-    sqlx::Error,
-> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    let row = sqlx::query(
-        "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE username = ?",
-    )
-    .bind(username)
-    .fetch_optional(&pool)
-    .await?;
-
-    match row {
-        Some(row) => {
-            let id: i64 = row.get("id");
-            let iota_id: i64 = row.get("iota_id");
-            let username: Vec<u8> = row.get("username");
-            let display: Option<Vec<u8>> = row.get("display");
-            let status: Option<Vec<u8>> = row.get("status");
-            let about: Option<Vec<u8>> = row.get("about");
-            let avatar: Option<Vec<u8>> = row.get("avatar");
-            let sub_level: i32 = row.get("sub_level");
-            let sub_end: i64 = row.get("sub_end");
-            let public_key = PublicKeyBundle::from_bytes(&row.get::<Vec<u8>, _>("public_key"))
-                .map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
-            let private_key_hash: Vec<u8> = row.get("private_key_hash");
-            let token: Vec<u8> = row.get("token");
-
-            Ok((
-                id,
-                iota_id,
-                String::from_utf8_lossy(&username).to_string(),
-                display.map(|d| String::from_utf8_lossy(&d).to_string()),
-                status.map(|s| String::from_utf8_lossy(&s).to_string()),
-                about.map(|a| String::from_utf8_lossy(&a).to_string()),
-                avatar,
-                sub_level,
-                sub_end,
-                public_key,
-                String::from_utf8_lossy(&private_key_hash).to_string(),
-                String::from_utf8_lossy(&token).to_string(),
-            ))
-        }
-        _ => Err(sqlx::Error::RowNotFound),
-    }
-}
-
-pub async fn get_by_user_id(
-    id: i64,
-) -> Result<
-    (
-        i64,
-        i64,
-        String,
-        Option<String>,
-        Option<String>,
-        Option<String>,
-        Option<Vec<u8>>,
-        i32,
-        i64,
-        PublicKeyBundle,
-        String,
-        String,
-    ),
-    sqlx::Error,
-> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    let row = sqlx::query(
-        "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE id = ?",
-    )
-    .bind(id)
-    .fetch_optional(&pool)
-    .await?;
-
-    match row {
-        Some(row) => {
-            let id: i64 = row.get("id");
-            let iota_id: i64 = row.get("iota_id");
-            let username: Vec<u8> = row.get("username");
-            let display: Option<Vec<u8>> = row.get("display");
-            let status: Option<Vec<u8>> = row.get("status");
-            let about: Option<Vec<u8>> = row.get("about");
-            let avatar: Option<Vec<u8>> = row.get("avatar");
-            let sub_level: i32 = row.get("sub_level");
-            let sub_end: i64 = row.get("sub_end");
-            let public_key = PublicKeyBundle::from_bytes(&row.get::<Vec<u8>, _>("public_key"))
-                .map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
-            let private_key_hash: Vec<u8> = row.get("private_key_hash");
-            let token: Vec<u8> = row.get("token");
-
-            Ok((
-                id,
-                iota_id,
-                String::from_utf8_lossy(&username).to_string(),
-                display.map(|d| String::from_utf8_lossy(&d).to_string()),
-                status.map(|s| String::from_utf8_lossy(&s).to_string()),
-                about.map(|a| String::from_utf8_lossy(&a).to_string()),
-                avatar,
-                sub_level,
-                sub_end,
-                public_key,
-                String::from_utf8_lossy(&private_key_hash).to_string(),
-                String::from_utf8_lossy(&token).to_string(),
-            ))
-        }
-        _ => Err(sqlx::Error::RowNotFound),
-    }
-}
-
-pub async fn get_users_by_iota_id(
-    iota_id_param: i64,
-) -> Result<
-    Vec<(
-        i64,
-        i64,
-        String,
-        Option<String>,
-        Option<String>,
-        Option<String>,
-        Option<Vec<u8>>,
-        i32,
-        i64,
-        PublicKeyBundle,
-        String,
-        String,
-    )>,
-    sqlx::Error,
-> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    let rows = sqlx::query(
-        "SELECT id, iota_id, username, display, status, about, avatar, sub_level, sub_end, public_key, private_key_hash, token FROM users WHERE iota_id = ?",
-    )
-    .bind(iota_id_param)
-    .fetch_all(&pool)
-    .await?;
-
-    let mut users = Vec::new();
-    for row in rows {
-        let id: i64 = row.get("id");
-        let iota_id: i64 = row.get("iota_id");
-        let username: Vec<u8> = row.get("username");
-        let display: Option<Vec<u8>> = row.get("display");
-        let status: Option<Vec<u8>> = row.get("status");
-        let about: Option<Vec<u8>> = row.get("about");
-        let avatar: Option<Vec<u8>> = row.get("avatar");
-        let sub_level: i32 = row.get("sub_level");
-        let sub_end: i64 = row.get("sub_end");
-        let public_key = PublicKeyBundle::from_bytes(&row.get::<Vec<u8>, _>("public_key"))
-            .map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
-        let private_key_hash: Vec<u8> = row.get("private_key_hash");
-        let token: Vec<u8> = row.get("token");
-
-        users.push((
-            id,
-            iota_id,
-            String::from_utf8_lossy(&username).to_string(),
-            display.map(|d| String::from_utf8_lossy(&d).to_string()),
-            status.map(|s| String::from_utf8_lossy(&s).to_string()),
-            about.map(|a| String::from_utf8_lossy(&a).to_string()),
-            avatar,
-            sub_level,
-            sub_end,
-            public_key,
-            String::from_utf8_lossy(&private_key_hash).to_string(),
-            String::from_utf8_lossy(&token).to_string(),
-        ));
-    }
-
-    Ok(users)
-}
-
-pub async fn change_username(id: i64, new_username: String) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("UPDATE users SET username = ? WHERE id = ?")
-        .bind(new_username.as_bytes().to_vec())
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-
-pub async fn change_display_name(id: i64, new_display: String) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("UPDATE users SET display = ? WHERE id = ?")
-        .bind(new_display.as_bytes().to_vec())
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-
-pub async fn change_avatar(id: i64, new_avatar: String) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("UPDATE users SET avatar = ? WHERE id = ?")
-        .bind(new_avatar.as_bytes().to_vec())
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-
-pub async fn change_about(id: i64, new_about: String) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("UPDATE users SET about = ? WHERE id = ?")
-        .bind(new_about.as_bytes().to_vec())
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-
-pub async fn change_status(id: i64, new_status: String) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("UPDATE users SET status = ? WHERE id = ?")
-        .bind(new_status.as_bytes().to_vec())
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-
-pub async fn delete_user(id: i64) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("DELETE FROM users WHERE id = ?")
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-pub async fn change_iota_id(id: i64, new_iota_id: i64) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("UPDATE users SET iota_id = ? WHERE id = ?")
-        .bind(new_iota_id)
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-
-pub async fn change_keys(
-    id: i64,
-    new_public_key: PublicKeyBundle,
-    new_private_key_hash: String,
-) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("UPDATE users SET public_key = ?, private_key_hash = ? WHERE id = ?")
-        .bind(new_public_key.as_bytes())
-        .bind(new_private_key_hash.as_bytes().to_vec())
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-pub async fn change_token(id: i64, new_token: String) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("UPDATE users SET token = ? WHERE id = ?")
-        .bind(new_token.as_bytes().to_vec())
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-pub async fn register_complete_user(
-    id: i64,
-    username: String,
-    public_key: PublicKeyBundle,
-    iota_id: i64,
-    token: String,
-) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query(
-        "INSERT INTO users (id, username, public_key, iota_id, token) VALUES (?, ?, ?, ?, ?)",
-    )
-    .bind(id)
-    .bind(username.as_bytes().to_vec())
-    .bind(public_key.as_bytes())
-    .bind(iota_id)
-    .bind(token.as_bytes().to_vec())
-    .execute(&pool)
-    .await?;
-
-    Ok(())
-}
-pub async fn print_users() -> Result<(), Box<dyn std::error::Error>> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    log!("Printing users...");
-    for row in sqlx::query(
-        "SELECT id, iota_id, username, display, status, about, sub_level, sub_end, public_key, private_key_hash, token FROM users",
-    )
-    .fetch_all(&pool)
-    .await?
-    .iter()
-    {
-        let id: i64 = row.get("id");
-        let iota_id: i64 = row.get("iota_id");
-        let username: Vec<u8> = row.get("username");
-        let display: Option<Vec<u8>> = row.get("display");
-        let status: Option<Vec<u8>> = row.get("status");
-        let about: Option<Vec<u8>> = row.get("about");
-        let sub_level: i32 = row.get("sub_level");
-        let sub_end: i64 = row.get("sub_end");
-
-        log!(
-            "User: {:?}",
-            (
-                id,
-                iota_id,
-                String::from_utf8_lossy(&username).to_string(),
-                display.map_or("".to_string(), |d| String::from_utf8_lossy(&d).to_string()),
-                status.map_or("".to_string(), |s| String::from_utf8_lossy(&s).to_string()),
-                about.map_or("".to_string(), |a| String::from_utf8_lossy(&a).to_string()),
-                sub_level,
-                sub_end
-            )
-        );
-    }
-
-    Ok(())
-}
-
-// ==========================================================================================
-//                                                IOTA
-// ==========================================================================================
-pub async fn create_new_iota(public_key: PublicKeyBundle) -> Result<i64, sqlx::Error> {
-    let new_id = get_register_id().await as i64;
-    register_complete_iota(new_id, public_key).await?;
-    Ok(new_id)
-}
-
-pub async fn register_complete_iota(
-    id: i64,
-    public_key: PublicKeyBundle,
-) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("INSERT INTO iotas (id, public_key) VALUES (?, ?)")
-        .bind(id)
-        .bind(public_key.as_bytes())
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-
-pub async fn get_iota_by_id(id: i64) -> Result<(i64, PublicKeyBundle), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    let result =
-        sqlx::query_as::<_, (i64, Vec<u8>)>("SELECT id, public_key FROM iotas WHERE id = ?")
-            .bind(id)
-            .fetch_optional(&pool)
-            .await;
-
-    match result {
-        Ok(optional_row) => match optional_row {
-            Some((id_i64, public_key)) => {
-                let bundle = PublicKeyBundle::from_bytes(&public_key)
-                    .map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
-                Ok((id_i64, bundle))
-            }
-            _ => Err(sqlx::Error::RowNotFound),
-        },
-        Err(e) => Err(e),
-    }
-}
-
-pub async fn change_iota_key(id: i64, new_key: PublicKeyBundle) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("UPDATE iotas SET public_key = ? WHERE id = ?")
-        .bind(new_key.as_bytes())
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-
-pub async fn delete_iota(id: i64) -> Result<(), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    sqlx::query("DELETE FROM iotas WHERE id = ?")
-        .bind(id)
-        .execute(&pool)
-        .await?;
-
-    Ok(())
-}
-
-// ==========================================================================================
-//                                         OMIKRONS
-// ==========================================================================================
-
-pub async fn get_omikron_by_id(id: i64) -> Result<(PublicKeyBundle, String, u16), sqlx::Error> {
-    let pool = {
-        let db_lock = SQL_DB.read().await;
-        db_lock
-            .as_ref()
-            .cloned()
-            .expect("Database pool not initialized")
-    };
-
-    let row = sqlx::query_as::<_, (Vec<u8>, Vec<u8>, i32)>(
-        "SELECT public_key, ip_address, port FROM omikrons WHERE id = ?",
-    )
-    .bind(id)
-    .fetch_optional(&pool)
-    .await?;
-
-    match row {
-        Some((public_key, ip_address, port)) => {
-            let bundle = PublicKeyBundle::from_bytes(&public_key)
-                .map_err(|e| sqlx::Error::Decode(Box::new(e)))?;
-            Ok((
-                bundle,
-                String::from_utf8_lossy(&ip_address).to_string(),
-                port as u16,
-            ))
-        }
-        _ => Err(sqlx::Error::RowNotFound),
-    }
-}
-
-// ==========================================================================================
-//                                         PHI
-// ==========================================================================================
-
-pub async fn add_notification(sender_id: i64, receiver_id: i64) -> Result<(), sqlx::Error> {
-    let db_lock = SQL_DB.read().await;
-    let pool = db_lock.as_ref().expect("Database pool is not initialized");
-
-    sqlx::query(
-        r#"
-        INSERT INTO notifications (sender_id, receiver_id, amount)
-        VALUES (?, ?, 1)
-        ON DUPLICATE KEY UPDATE amount = amount + 1
-        "#,
-    )
-    .bind(sender_id)
-    .bind(receiver_id)
-    .execute(pool)
-    .await?;
-
-    Ok(())
-}
-pub async fn read_notification(sender_id: i64, receiver_id: i64) -> Result<(), sqlx::Error> {
-    let db_lock = SQL_DB.read().await;
-    let pool = db_lock.as_ref().expect("Database pool is not initialized");
-
-    sqlx::query(
-        r#"
-        DELETE FROM notifications WHERE sender_id = ? AND receiver_id = ?
-        "#,
-    )
-    .bind(sender_id)
-    .bind(receiver_id)
-    .execute(pool)
-    .await?;
-
-    Ok(())
-}
-pub async fn get_notifications(user_id: i64) -> Result<Vec<(i64, i64)>, sqlx::Error> {
-    let db_lock = SQL_DB.read().await;
-    let pool = db_lock.as_ref().expect("Database pool is not initialized");
-
-    sqlx::query_as::<_, (i64, i64)>(
-        r#"
-        SELECT sender_id, amount FROM notifications WHERE receiver_id = ?
-        "#,
-    )
-    .bind(user_id)
-    .fetch_all(pool)
-    .await
-}

src/sql/user_online_tracker.rs

@@ -1,4 +1,5 @@
-use crate::sql;
+use crate::db::user_repo;
+use crate::models::IotaId;
 use crate::sql::connection_status::UserStatus;
 use dashmap::DashMap;
 use once_cell::sync::Lazy;
@@ -33,27 +34,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 {
-    let connections_empty = if let Some(r) = IOTA_OMIKRON_CONNECTIONS.get(&iota_id) {
+    let mut replacement = None;
-        let mut vec = r.value().clone();
+    let mut connections_empty = false;
-        vec.retain(|&id| id != omikron_id);
+    if let Some(mut entry) = IOTA_OMIKRON_CONNECTIONS.get_mut(&iota_id) {
-        let empty = vec.is_empty();
+        entry.retain(|&id| id != omikron_id);
-        drop(r);
+        connections_empty = entry.is_empty();
-        IOTA_OMIKRON_CONNECTIONS.insert(iota_id, vec);
+        replacement = entry.first().copied();
-        empty
-    } else {
-        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_empty {
-        IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
+        IOTA_OMIKRON_CONNECTIONS.remove_if(&iota_id, |_, connections| connections.is_empty());
+    }
+
+    if IOTA_PRIMARY_OMIKRON_CONNECTION
+        .get(&iota_id)
+        .is_some_and(|primary| *primary == omikron_id)
+    {
+        match replacement {
+            Some(omikron_id) => {
+                IOTA_PRIMARY_OMIKRON_CONNECTION.insert(iota_id, omikron_id);
+            }
+            None => {
+                IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
+            }
+        }
     }
 
     connections_empty
@@ -87,9 +91,9 @@ pub async fn get_all_connections()
         let iota_id = *entry.key();
         let omikron_ids = entry.value().clone();
 
-        if let Ok(users) = sql::sql::get_users_by_iota_id(iota_id.try_into().unwrap()).await {
+        if let Ok(users) = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await {
             for user in users {
-                let user_id = user.0 as i64;
+                let user_id = user.id.0;
                 if let Some(conn) = USER_STATUS_MAP.get(&user_id) {
                     let user_omikron_id = conn.omikron_id;
                     if omikron_ids.contains(&user_omikron_id) {
@@ -117,6 +121,12 @@ pub fn track_user_status(user_id: i64, status: UserStatus, omikron_id: i64) {
     );
 }
 
+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 get_user_status(user_id: i64) -> Option<UserConnection> {
     USER_STATUS_MAP.get(&user_id).map(|v| v.clone())
 }
@@ -128,57 +138,38 @@ pub fn untrack_many_users(user_ids: &[i64]) {
 }
 
 pub async fn untrack_omikron(omikron_id: i64) {
-    let primary_keys_to_remove: Vec<i64> = IOTA_PRIMARY_OMIKRON_CONNECTION
-        .iter()
-        .filter(|entry| *entry.value() == omikron_id)
-        .map(|entry| *entry.key())
-        .collect();
-
-    for key in primary_keys_to_remove {
-        IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&key);
-    }
-
     let mut offline_iotas = Vec::new();
-    let mut primary_to_remove = Vec::new();
+    let mut primary_replacements = Vec::new();
 
-    // Collect iotas and primary info first
+    for mut entry in IOTA_OMIKRON_CONNECTIONS.iter_mut() {
-    for r in IOTA_OMIKRON_CONNECTIONS.iter() {
+        let iota_id = *entry.key();
-        let iota_id = *r.key();
+        entry.retain(|&id| id != omikron_id);
-        let mut connections = r.value().clone();
+        if entry.is_empty() {
-        connections.retain(|&id| id != omikron_id);
-
-        if connections.is_empty() {
             offline_iotas.push(iota_id);
-        }
+        } else if IOTA_PRIMARY_OMIKRON_CONNECTION
-
-        if IOTA_PRIMARY_OMIKRON_CONNECTION
             .get(&iota_id)
-            .map(|p| *p == omikron_id)
+            .is_some_and(|primary| *primary == omikron_id)
-            .unwrap_or(false)
         {
-            primary_to_remove.push(iota_id);
+            primary_replacements.push((iota_id, entry[0]));
         }
-
-        // Update the connections vector after filtering
-        IOTA_OMIKRON_CONNECTIONS.insert(iota_id, connections);
     }
 
-    // Step 2: Remove primary connections safely
+    for iota_id in &offline_iotas {
-    for iota_id in primary_to_remove {
+        IOTA_OMIKRON_CONNECTIONS.remove_if(iota_id, |_, connections| connections.is_empty());
-        IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
+        IOTA_PRIMARY_OMIKRON_CONNECTION.remove(iota_id);
+    }
+
+    for (iota_id, replacement) in primary_replacements {
+        IOTA_PRIMARY_OMIKRON_CONNECTION.insert(iota_id, replacement);
     }
 
-    // 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.try_into().unwrap()).await {
+        if let Ok(users) = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await {
             for user in users {
-                USER_STATUS_MAP.remove(&(user.0 as i64));
+                USER_STATUS_MAP.remove(&user.id.0);
             }
         }
-        // Finally remove the empty connections vector
-        IOTA_OMIKRON_CONNECTIONS.remove(&iota_id);
     }
 }

src/transport/connection.rs

@@ -0,0 +1 @@
+pub(crate) use super::omikron_connection::OmikronConnection;

src/transport/handlers/account.rs

@@ -0,0 +1,43 @@
+use super::super::omikron_connection::{OmikronConnection, OmikronResult};
+use crate::{
+    db::{iota_repo, user_repo},
+    models::{IotaId, UserId},
+};
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
+use std::sync::Arc;
+
+async fn delete(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+    result: impl std::future::Future<Output = crate::error::Result<()>>,
+) -> OmikronResult<()> {
+    let response = match result.await {
+        Ok(()) => CommunicationValue::new(CommunicationType::Success),
+        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
+}
+
+pub async fn user(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    delete(
+        connection,
+        value.clone(),
+        user_repo::delete_user(UserId::from(value.get_sender() as i64)),
+    )
+    .await
+}
+pub async fn iota(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    delete(
+        connection,
+        value.clone(),
+        iota_repo::delete_iota(IotaId::from(value.get_sender() as i64)),
+    )
+    .await
+}

src/transport/handlers/calls.rs

@@ -0,0 +1 @@
+/* Call and WebRTC commands are reserved for the transport extensions that define those protocol values. */

src/transport/handlers/links.rs

@@ -0,0 +1,21 @@
+use super::super::omikron_connection::{OmikronConnection, OmikronResult};
+use crate::server::short_link::add_short_link;
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
+use std::sync::Arc;
+
+pub async fn shorten(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    let link = value
+        .get_data(DataType::Link)
+        .as_str()
+        .ok_or(crate::error::OmegaError::InvalidResponse)?;
+    let short = add_short_link(link)
+        .await
+        .map_err(|_| crate::error::OmegaError::Transport("short link error".to_string()))?;
+    let response = CommunicationValue::new(CommunicationType::ShortenLink)
+        .with_id(value.get_id())
+        .add_typed_default(DataType::Link, DataValue::Str(short));
+    connection.send(&response).await
+}

src/transport/handlers/messaging.rs

@@ -0,0 +1 @@
+/* Message delivery remains in the connection dispatcher until the protocol exposes a separate message handler contract. */

src/transport/handlers/mod.rs

@@ -0,0 +1,9 @@
+pub mod account;
+pub mod calls;
+pub mod links;
+pub mod messaging;
+pub mod notifications;
+pub mod presence;
+pub mod register;
+pub mod states;
+pub mod user_data;

src/transport/handlers/notifications.rs

@@ -0,0 +1,116 @@
+use super::super::omikron_connection::{OmikronConnection, OmikronResult};
+use crate::{db::notification_repo, log, models::UserId};
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
+use mtp::type_map::TypeMap;
+use std::sync::Arc;
+
+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 response = CommunicationValue::new(CommunicationType::GetNotifications)
+        .with_id(value.get_id())
+        .add_typed_default(DataType::Notifications, DataValue::Array(notifications));
+    connection.send(&response).await
+}
+
+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 Some(other) = value
+        .get_data(DataType::SenderId)
+        .as_number()
+        .map(|id| id as i64)
+    else {
+        return Ok(());
+    };
+    if let Err(error) =
+        notification_repo::read_notification(UserId::from(receiver), UserId::from(other)).await
+    {
+        log!(
+            crate::util::logger::PrintType::General,
+            "SQL read_notification error: {}",
+            error
+        );
+    } else {
+        let response =
+            CommunicationValue::new(CommunicationType::ReadNotification).with_id(value.get_id());
+        let _ = connection.send(&response).await;
+        let sync = CommunicationValue::new(CommunicationType::ReadNotification)
+            .with_receiver(receiver as u64)
+            .add_typed_default(DataType::SenderId, DataValue::SignedNumber(other.into()));
+        crate::transport::omikron_manager::send_to_user(receiver, &sync).await;
+    }
+    Ok(())
+}
+
+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 sender = value
+        .get_data(DataType::SenderId)
+        .as_number()
+        .map(|id| id as i64)
+        .unwrap_or(value.get_sender() as i64);
+    if let Err(error) =
+        notification_repo::add_notification(UserId::from(receiver), UserId::from(sender)).await
+    {
+        log!(
+            crate::util::logger::PrintType::General,
+            "SQL add_notification error: {}",
+            error
+        );
+    } else {
+        let response =
+            CommunicationValue::new(CommunicationType::PushNotification).with_id(value.get_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()));
+        crate::transport::omikron_manager::send_to_user(receiver, &push).await;
+    }
+    Ok(())
+}

src/transport/handlers/presence.rs

@@ -0,0 +1,128 @@
+use super::super::omikron_connection::{OmikronConnection, OmikronResult};
+use crate::{
+    db::user_repo,
+    log_in,
+    models::IotaId,
+    sql::{connection_status::UserStatus, user_online_tracker},
+};
+use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
+use std::sync::Arc;
+
+pub async fn user_connected(
+    _connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+    omikron_id: i64,
+) -> OmikronResult<()> {
+    log_in!(crate::util::logger::PrintType::Omega, "User connected");
+    if let Some(user_id) = value.get_data(DataType::UserId).as_number() {
+        let status = value
+            .get_data(DataType::UserState)
+            .as_str()
+            .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);
+        }
+    }
+    Ok(())
+}
+
+pub async fn user_disconnected(
+    _: Arc<OmikronConnection>,
+    value: CommunicationValue,
+    omikron_id: i64,
+) -> 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);
+    }
+    Ok(())
+}
+
+pub async fn iota_connected(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+    omikron_id: i64,
+) -> OmikronResult<()> {
+    log_in!(crate::util::logger::PrintType::Omega, "IOTA connected");
+    let Some(iota_id) = value
+        .get_data(DataType::IotaId)
+        .as_number()
+        .map(|id| id as i64)
+    else {
+        return Ok(());
+    };
+    user_online_tracker::track_iota_connection(iota_id, omikron_id, true);
+    let mut user_ids = Vec::new();
+    match user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await {
+        Ok(users) => {
+            for user in users {
+                user_ids.push(DataValue::SignedNumber(user.id.0.into()));
+                user_online_tracker::track_user_status(
+                    user.id.0,
+                    UserStatus::user_offline,
+                    omikron_id,
+                );
+            }
+        }
+        Err(_) => log_in!(
+            crate::util::logger::PrintType::General,
+            "SQL error loading users for IOTA"
+        ),
+    }
+    let response = CommunicationValue::new(CommunicationType::IotaUserData)
+        .with_id(value.get_id())
+        .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into()))
+        .add_typed_default(DataType::UserIds, DataValue::Array(user_ids));
+    let _ = connection.send(&response).await;
+    Ok(())
+}
+
+pub async fn iota_disconnected(
+    _: Arc<OmikronConnection>,
+    value: CommunicationValue,
+    omikron_id: i64,
+) -> OmikronResult<()> {
+    log_in!(crate::util::logger::PrintType::Omega, "IOTA disconnected");
+    let Some(iota_id) = value
+        .get_data(DataType::IotaId)
+        .as_number()
+        .map(|id| id as i64)
+    else {
+        return Ok(());
+    };
+    if user_online_tracker::untrack_iota_connection(iota_id, omikron_id) {
+        if let Ok(users) = user_repo::get_users_by_iota_id(IotaId::from(iota_id)).await {
+            user_online_tracker::untrack_many_users(
+                &users.iter().map(|user| user.id.0).collect::<Vec<_>>(),
+            );
+        }
+    }
+    Ok(())
+}
+
+pub async fn sync_status(
+    _: Arc<OmikronConnection>,
+    value: CommunicationValue,
+    omikron_id: i64,
+) -> OmikronResult<()> {
+    if let DataValue::Array(ids) = value.get_data(DataType::UserIds) {
+        for id in ids {
+            if let DataValue::SignedNumber(id) = id {
+                user_online_tracker::track_user_status(
+                    *id as i64,
+                    UserStatus::user_offline,
+                    omikron_id,
+                );
+            }
+        }
+    }
+    if let DataValue::Array(ids) = value.get_data(DataType::IotaIds) {
+        for id in ids {
+            if let DataValue::SignedNumber(id) = id {
+                user_online_tracker::track_iota_connection(*id as i64, omikron_id, true);
+            }
+        }
+    }
+    Ok(())
+}

src/transport/handlers/register.rs

@@ -0,0 +1,151 @@
+use super::super::omikron_connection::{OmikronConnection, OmikronResult};
+use crate::{
+    db::{iota_repo, user_repo},
+    models::{IotaId, UserId},
+};
+use mtp::{
+    codec::{CommunicationType, CommunicationValue, DataType, DataValue},
+    crypto::PublicKeyBundle,
+};
+use std::sync::Arc;
+
+pub async fn get_register(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    let register_id = user_repo::get_register_id().await?;
+    let response = CommunicationValue::new(CommunicationType::GetRegister)
+        .with_id(value.get_id())
+        .add_typed_default(
+            DataType::UserId,
+            DataValue::SignedNumber(register_id.0.into()),
+        );
+    connection.send(&response).await
+}
+
+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()
+        .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)
+            .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
+                }
+            }
+        }
+        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
+            }
+        },
+    }
+}
+
+pub async fn complete_user(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    let user_id = value
+        .get_data(DataType::UserId)
+        .as_number()
+        .map(|id| id as i64);
+    let username = value
+        .get_data(DataType::Username)
+        .as_str()
+        .map(str::to_owned);
+    let public_key = value
+        .get_data(DataType::PublicKey)
+        .as_str()
+        .and_then(|key| PublicKeyBundle::from_base64(key).ok());
+    let reset_token = value
+        .get_data(DataType::ResetToken)
+        .as_str()
+        .map(str::to_owned);
+    let Some((user_id, username, public_key, reset_token)) = user_id
+        .zip(username)
+        .zip(public_key)
+        .zip(reset_token)
+        .map(|(((id, name), key), token)| (id, name, key, token))
+    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),
+        reset_token,
+    )
+    .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
+        }
+    }
+}

src/transport/handlers/states.rs

@@ -0,0 +1,46 @@
+use super::super::omikron_connection::{OmikronConnection, OmikronResult};
+use crate::sql::{connection_status::UserStatus, user_online_tracker};
+use mtp::{
+    codec::{CommunicationType, CommunicationValue, DataType, DataValue},
+    type_map::TypeMap,
+};
+use std::sync::Arc;
+
+pub async fn get(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    let DataValue::Array(ids) = value.get_data(DataType::UserIds) else {
+        return Ok(());
+    };
+    let tm = TypeMap::latest();
+    let states = ids
+        .iter()
+        .filter_map(|id| {
+            let DataValue::SignedNumber(id) = id else {
+                return None;
+            };
+            let status = user_online_tracker::get_user_status(*id as i64)
+                .map(|status| {
+                    if status.connection_type == UserStatus::user_invisible {
+                        UserStatus::user_offline.to_string()
+                    } else {
+                        status.connection_type.to_string()
+                    }
+                })
+                .unwrap_or_else(|| UserStatus::iota_offline.to_string());
+            let mut map = Vec::new();
+            if let Some(kind) = DataType::UserId.try_to_id(&tm) {
+                map.push((kind, DataValue::SignedNumber((*id as i64).into())));
+            }
+            if let Some(kind) = DataType::UserState.try_to_id(&tm) {
+                map.push((kind, DataValue::Str(status)));
+            }
+            Some(DataValue::Container(map))
+        })
+        .collect();
+    let response = CommunicationValue::new(CommunicationType::GetStates)
+        .with_id(value.get_id())
+        .add_typed_default(DataType::UserStates, DataValue::Array(states));
+    connection.send(&response).await
+}

src/transport/handlers/user_data.rs

@@ -0,0 +1,271 @@
+use super::super::omikron_connection::{OmikronConnection, OmikronResult};
+use crate::{
+    db::{iota_repo, user_repo},
+    models::{IotaId, UserId},
+    sql::{connection_status::UserStatus, user_online_tracker},
+};
+use base64::{Engine as _, engine::general_purpose::STANDARD};
+use mtp::{
+    codec::{CommunicationType, CommunicationValue, DataType, DataValue},
+    crypto::PublicKeyBundle,
+};
+use std::sync::Arc;
+
+fn connections(iota_id: i64) -> DataValue {
+    DataValue::Array(
+        user_online_tracker::get_iota_omikron_connections(iota_id)
+            .unwrap_or_default()
+            .into_iter()
+            .map(|id| DataValue::SignedNumber(id.into()))
+            .collect(),
+    )
+}
+
+pub async fn get_user(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    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()
+    } else if let Some(name) = value.get_data(DataType::Username).as_str() {
+        user_repo::get_by_username(name).await.ok()
+    } else {
+        None
+    };
+    let Some(user) = user else {
+        return connection
+            .send_error_response(value.get_id(), CommunicationType::ErrorNotFound)
+            .await;
+    };
+    let id = user.id.0;
+    let iota_id = user.iota_id.0;
+    let username = user.username.clone();
+    let display = user
+        .display
+        .filter(|name| !name.is_empty())
+        .unwrap_or_else(|| username.clone());
+    let mut response = CommunicationValue::new(CommunicationType::GetUserData)
+        .with_id(value.get_id())
+        .add_typed_default(DataType::Username, DataValue::Str(username))
+        .add_typed_default(
+            DataType::PublicKey,
+            DataValue::Str(user.public_key.to_base64()),
+        )
+        .add_typed_default(DataType::UserId, DataValue::SignedNumber(id.into()))
+        .add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into()))
+        .add_typed_default(DataType::Display, DataValue::Str(display))
+        .add_typed_default(
+            DataType::SubLevel,
+            DataValue::SignedNumber(user.sub_level as i128),
+        )
+        .add_typed_default(
+            DataType::SubEnd,
+            DataValue::SignedNumber(user.sub_end.into()),
+        );
+    if let Some(status) = user.status.filter(|value| !value.is_empty()) {
+        response = response.add_typed_default(DataType::Status, DataValue::Str(status));
+    }
+    if let Some(about) = user.about.filter(|value| !value.is_empty()) {
+        response = response.add_typed_default(DataType::About, DataValue::Str(about));
+    }
+    if let Some(avatar) = user.avatar {
+        response =
+            response.add_typed_default(DataType::Avatar, DataValue::Str(STANDARD.encode(avatar)));
+    }
+    let online = user_online_tracker::get_user_status(id);
+    response = response
+        .add_typed_default(
+            DataType::OnlineStatus,
+            DataValue::Str(
+                online
+                    .as_ref()
+                    .map(|status| {
+                        if status.connection_type == UserStatus::user_invisible {
+                            UserStatus::user_offline.to_string()
+                        } else {
+                            status.connection_type.to_string()
+                        }
+                    })
+                    .unwrap_or_else(|| UserStatus::iota_offline.to_string()),
+            ),
+        )
+        .add_typed_default(DataType::OmikronConnections, connections(iota_id));
+    if let Some(status) = online {
+        response = response.add_typed_default(
+            DataType::OmikronId,
+            DataValue::SignedNumber(status.omikron_id.into()),
+        );
+    }
+    connection.send(&response).await
+}
+
+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))
+            .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 {
+            iota_repo::get_iota_by_id(user.iota_id)
+                .await
+                .ok()
+                .map(|iota| (iota.id.0, iota.public_key, Some(user.id.0), None))
+        } else {
+            None
+        }
+    } else if let Some(name) = value.get_data(DataType::Username).as_str() {
+        if let Ok(user) = user_repo::get_by_username(name).await {
+            iota_repo::get_iota_by_id(user.iota_id)
+                .await
+                .ok()
+                .map(|iota| {
+                    (
+                        iota.id.0,
+                        iota.public_key,
+                        Some(user.id.0),
+                        Some(name.to_owned()),
+                    )
+                })
+        } else {
+            None
+        }
+    } else {
+        None
+    };
+    let Some((id, key, user_id, username)) = found else {
+        return connection
+            .send_error_response(value.get_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()))
+        .add_typed_default(DataType::IotaId, DataValue::SignedNumber(id.into()))
+        .add_typed_default(DataType::OmikronConnections, connections(id));
+    if let Some(user_id) = user_id {
+        response =
+            response.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()));
+    }
+    if let Some(username) = username {
+        response = response.add_typed_default(DataType::Username, DataValue::Str(username));
+    }
+    connection.send(&response).await
+}
+
+async fn update_user(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    let id = UserId::from(value.get_sender() as 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())
+            .await
+            .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() {
+        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() {
+        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() {
+        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() {
+        if 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());
+        }
+    }
+    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
+}
+
+pub async fn change_user(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    update_user(connection, value).await
+}
+
+pub async fn change_iota(
+    connection: Arc<OmikronConnection>,
+    value: CommunicationValue,
+) -> OmikronResult<()> {
+    let Some(reset) = value.get_data(DataType::ResetToken).as_str() else {
+        return connection
+            .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
+            .await;
+    };
+    let Some(new_token) = value.get_data(DataType::NewToken).as_str() else {
+        return connection
+            .send_error_response(value.get_id(), CommunicationType::ErrorInvalidData)
+            .await;
+    };
+    let user_id = UserId::from(value.get_sender() as 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)
+                .await;
+        }
+    };
+    if user.token != reset {
+        return connection
+            .send_error_response(value.get_id(), CommunicationType::ErrorInvalidChallenge)
+            .await;
+    }
+    let result =
+        match user_repo::change_iota_id(user_id, IotaId::from(value.get_sender() as i64)).await {
+            Ok(()) => user_repo::change_token(user_id, new_token.to_owned()).await,
+            Err(error) => Err(error),
+        };
+    let response = match result {
+        Ok(()) => CommunicationValue::new(CommunicationType::Success),
+        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
+}

src/transport/mod.rs

@@ -1,2 +1,4 @@
+pub mod connection;
+pub mod handlers;
 pub mod omikron_connection;
 pub mod omikron_manager;

src/transport/omikron_manager.rs

@@ -1,4 +1,4 @@
-use crate::transport::omikron_connection::OmikronConnection;
+use crate::transport::connection::OmikronConnection;
 use dashmap::DashMap;
 use mtp::codec::CommunicationValue;
 use once_cell::sync::Lazy;
@@ -27,11 +27,9 @@ pub async fn remove_omikron(omikron_id: i64) {
 }
 
 pub async fn get_random_omikron() -> Result<Arc<OmikronConnection>, ()> {
-    let mut rng = rand::thread_rng();
-
     let keys: Vec<_> = OMIKRON_CONNECTIONS.iter().map(|e| *e.key()).collect();
 
-    if let Some(key) = keys.into_iter().choose(&mut rng) {
+    if let Some(key) = keys.into_iter().choose(&mut rand::rng()) {
         if let Some(entry) = OMIKRON_CONNECTIONS.get(&key) {
             return Ok(entry.clone());
         }

src/util/file_util.rs

@@ -8,9 +8,8 @@ use zip::ZipArchive;
 
 use crate::log;
 
-static WORKING_DIR: Lazy<PathBuf> = Lazy::new(|| {
+static WORKING_DIR: Lazy<PathBuf> =
-    std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
+    Lazy::new(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
-});
 
 pub fn delete_file(path: &str, name: &str) -> bool {
     let dir = Path::new(&get_directory()).join(path);

src/util/logger.rs

@@ -269,8 +269,7 @@ fn format_data_container(data: Vec<(DataTypeId, DataValue)>, version: Version) -
                 DataValue::Str(s) => format!("{}=\"{}\"", key_str, s),
 
                 DataValue::Container(inner) => {
-                    let inner_formatted =
+                    let inner_formatted = format_data_container(inner, version.clone());
-                        format_data_container(inner, version.clone());
                     format!("{}={{ {} }}", key_str, inner_formatted)
                 }