[Add] Structure

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