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

[Add] Structure

Browse source
This commit is contained in:
Alex Emmet 2026-07-20 22:21:43 +02:00
commit 826fb9ce50
44 changed files with 2210 additions and 2721 deletions
Download patch file Download diff file Expand all files Collapse all files

720
Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

32
Cargo.toml
View file

@ -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.3"
rand = "0.8"
rand_core = { version = "0.6", features = ["getrandom", "std"] }
reqwest = { version = "0.13.2" }
rustls = { version = "0.23.37", default-features = false, features = [
once_cell = "1.21.4"
rand = "0.10.2"
reqwest = "0.13.4"
rustls = { version = "0.23.42", default-features = false, features = [
"std",
"tls12",
"aws-lc-rs",
"prefer-post-quantum",
] }
sha2 = "0.10.9"
sqlx = { version = "0.8.6", features = ["mysql", "runtime-async-std"] }
strum = "0.27.2"
strum_macros = "0.27.2"
sqlx = { version = "0.8.6", features = ["mysql", "runtime-tokio", "migrate"] }
strum = "0.28.0"
strum_macros = "0.28.0"
tokio = { version = "*", features = ["full"] }
uuid = { version = "1.19.0", features = ["v4"] }
x448 = "0.6.0"
zip = "6.0.0"
thiserror = "2.0.18"
uuid = { version = "1.24.0", features = ["v4", "v7"] }
zip = "8.6.0"
thiserror = "2.0.19"
serde = { version = "1.0.229", features = ["derive"] }
serde_json = "1.0.151"

44
migrations/001.sql Normal file
View file

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

1
src/api/mod.rs Normal file
View file

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

70
src/api/response.rs Normal file
View file

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

65
src/config.rs Normal file
View file

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

53
src/db/iota_repo.rs Normal file
View file

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

38
src/db/mod.rs Normal file
View file

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

38
src/db/notification_repo.rs Normal file
View file

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

29
src/db/omikron_repo.rs Normal file
View file

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

35
src/db/short_link_repo.rs Normal file
View file

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

194
src/db/user_repo.rs Normal file
View file

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

57
src/error.rs Normal file
View file

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

53
src/main.rs
View file

@ -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;
use crate::sql::sql::print_users;
pub use error::{OmegaError, Result};
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");
save_public_key_bundle(&kr.public_key_bundle(), KEYRING_PATH)
.expect("Failed to save generated public key bundle");
if let Err(error) = save_keyring_raw(&kr, KEYRING_PATH) {
eprintln!("Failed to save generated keyring: {error}");
}
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 {
log!("[ERROR] Failed to print users: {}", e);
} else {
log!(" Users");
}
let rate_limit_cleanup = crate::server::middleware::spawn_cleanup_task();
let short_link_cleanup = tokio::spawn(async {
let mut ticker = interval(Duration::from_secs(24 * 60 * 60));
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();
}

30
src/models/ids.rs Normal file
View file

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

9
src/models/iota.rs Normal file
View file

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

21
src/models/mod.rs Normal file
View file

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

9
src/models/notification.rs Normal file
View file

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

12
src/models/omikron.rs Normal file
View file

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

19
src/models/user.rs Normal file
View file

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

546
src/server/api.rs
View file

@ -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::sql::sql::{get_by_user_id, get_iota_by_id, get_omikron_by_id};
use crate::models::UserId;
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
.body
.map(|body| String::from_utf8_lossy(&body).to_string());
if method != Method::OPTIONS && !middleware::allow(request.remote_addr.ip(), &path) {
return response
.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(|s| !s.is_empty()).collect();
let _body: Option<JsonValue> = if let Some(ref bs) = body_string {
if let Ok(body_json) = json::parse(bs) {
Some(body_json)
} else {
None
}
} else {
None
};
let (status, body_text) = match path_parts.as_slice() {
// ==================================================
// DOWNLOAD IOTA FRONTEND
// ==================================================
["api", "download", "iota_frontend"] => {
let file_path = format!("{}/downloads/iota_frontend.zip", get_directory());
match std::fs::read(file_path) {
Ok(file_bytes) => {
return response
.status(StatusCode::OK)
.header("access-control-allow-origin", "*")
.header("content-type", "application/zip")
.header(
"content-disposition",
"attachment; filename=\"iota_frontend.zip\"",
)
.body(Bytes::from(file_bytes));
}
Err(_) => {
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())
}
};
let path_parts: Vec<&str> = path.split('/').filter(|part| !part.is_empty()).collect();
if let ["api", "download", "iota_frontend"] = path_parts.as_slice() {
let file_path = format!("{}/downloads/iota_frontend.zip", get_directory());
return match std::fs::read(file_path) {
Ok(bytes) => response
.status(StatusCode::OK)
.header("access-control-allow-origin", &crate::config::cors_origin())
.header("content-type", "application/zip")
.header(
"content-disposition",
"attachment; filename=\"iota_frontend.zip\"",
)
.body(Bytes::from(bytes)),
Err(_) => response
.status(StatusCode::NOT_FOUND)
.header("access-control-allow-origin", &crate::config::cors_origin())
.body(json(&StatusResponse {
status: "error_not_found",
})),
};
}
if let ["direct", short @ ..] = path_parts.as_slice() {
let short = short.join("");
let location = crate::server::short_link::get_short_link(&short)
.await
.unwrap_or_else(|_| "https://tensamin.net".to_string());
return response
.status(StatusCode::TEMPORARY_REDIRECT)
.header("location", &location);
}
let (status, body) = route(&path_parts)
.await
.unwrap_or_else(|error| (error.status_code(), error_body(&error)));
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(

84
src/server/middleware.rs Normal file
View file

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

2
src/server/mod.rs
View file

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

46
src/server/short_link.rs
View file

@ -1,37 +1,27 @@
use dashmap::DashMap;
use once_cell::sync::Lazy;
use rand::{Rng, thread_rng};
static LINKS: Lazy<DashMap<String, String>> = Lazy::new(DashMap::new);
use crate::db::short_link_repo;
use rand::RngExt;
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;
if !LINKS.contains_key(&short) {
return short;
let raw = generate_short_link().await;
if short_link_repo::insert(&raw, long).await.map_err(|_| ())? {
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(()) {
Ok(format!("{}{}", t, frag))
} else {
Err(())
}
let target = short_link_repo::get(&normalized)
.await
.map_err(|_| ())?
.ok_or(())?;
Ok(format!("{}{}", target, frag))
}
/* ---------------- helpers ---------------- */
fn short_length() -> usize {
let count = LINKS.len();
async fn short_length() -> usize {
let count = short_link_repo::count().await.unwrap_or(0);
match count {
0..=1_999 => 4,

42
src/server/validation.rs Normal file
View file

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

1
src/sql/mod.rs
View file

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

778
src/sql/sql.rs
View file

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

105
src/sql/user_online_tracker.rs
View file

@ -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 vec = r.value().clone();
vec.retain(|&id| id != omikron_id);
let empty = vec.is_empty();
drop(r);
IOTA_OMIKRON_CONNECTIONS.insert(iota_id, vec);
empty
} else {
false
};
if 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);
}
let mut replacement = None;
let mut connections_empty = false;
if let Some(mut entry) = IOTA_OMIKRON_CONNECTIONS.get_mut(&iota_id) {
entry.retain(|&id| id != omikron_id);
connections_empty = entry.is_empty();
replacement = entry.first().copied();
}
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 r in IOTA_OMIKRON_CONNECTIONS.iter() {
let iota_id = *r.key();
let mut connections = r.value().clone();
connections.retain(|&id| id != omikron_id);
if connections.is_empty() {
for mut entry in IOTA_OMIKRON_CONNECTIONS.iter_mut() {
let iota_id = *entry.key();
entry.retain(|&id| id != omikron_id);
if entry.is_empty() {
offline_iotas.push(iota_id);
}
if IOTA_PRIMARY_OMIKRON_CONNECTION
} else if IOTA_PRIMARY_OMIKRON_CONNECTION
.get(&iota_id)
.map(|p| *p == omikron_id)
.unwrap_or(false)
.is_some_and(|primary| *primary == omikron_id)
{
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 primary_to_remove {
IOTA_PRIMARY_OMIKRON_CONNECTION.remove(&iota_id);
for iota_id in &offline_iotas {
IOTA_OMIKRON_CONNECTIONS.remove_if(iota_id, |_, connections| connections.is_empty());
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);
}
}

1
src/transport/connection.rs Normal file
View file

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

43
src/transport/handlers/account.rs Normal file
View file

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

1
src/transport/handlers/calls.rs Normal file
View file

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

21
src/transport/handlers/links.rs Normal file
View file

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

1
src/transport/handlers/messaging.rs Normal file
View file

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

9
src/transport/handlers/mod.rs Normal file
View file

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

116
src/transport/handlers/notifications.rs Normal file
View file

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

128
src/transport/handlers/presence.rs Normal file
View file

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

151
src/transport/handlers/register.rs Normal file
View file

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

46
src/transport/handlers/states.rs Normal file
View file

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

271
src/transport/handlers/user_data.rs Normal file
View file

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

2
src/transport/mod.rs
View file

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

1020
src/transport/omikron_connection.rs
View file

File diff suppressed because it is too large Load diff

6
src/transport/omikron_manager.rs
View file

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

5
src/util/file_util.rs
View file

@ -8,9 +8,8 @@ use zip::ZipArchive;
use crate::log;
static WORKING_DIR: Lazy<PathBuf> = Lazy::new(|| {
std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
});
static WORKING_DIR: Lazy<PathBuf> =
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);

3
src/util/logger.rs
View file

@ -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 =
format_data_container(inner, version.clone());
let inner_formatted = format_data_container(inner, version.clone());
format!("{}={{ {} }}", key_str, inner_formatted)
}