initial commit

src/modules/token/mod.rs

@@ -0,0 +1,357 @@
+//
+// Copyright (c) 2025 rustmailer.com (https://rustmailer.com)
+//
+// This file is part of the Bichon Email Archiving Project
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+
+use crate::modules::account::migration::AccountModel;
+use crate::modules::database::delete_impl;
+use crate::modules::database::manager::DB_MANAGER;
+use crate::modules::database::{insert_impl, list_all_impl, update_impl};
+use crate::modules::token::payload::AccessTokenUpdateRequest;
+use crate::raise_error;
+use crate::{
+    generate_token, modules::error::BichonResult,
+    modules::token::payload::AccessTokenCreateRequest, utc_now,
+};
+use native_db::*;
+use native_model::{native_model, Model};
+use poem_openapi::Object;
+use serde::{Deserialize, Serialize};
+use std::collections::BTreeSet;
+use std::net::IpAddr;
+
+use super::error::code::ErrorCode;
+
+pub mod payload;
+pub mod root;
+
+#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize, Object)]
+#[native_model(id = 1, version = 1)]
+#[native_db]
+pub struct AccessToken {
+    /// The unique token string used for authentication
+    #[primary_key]
+    pub token: String,
+    /// A set of account information associated with the token.
+    pub accounts: BTreeSet<AccountInfo>,
+    /// The timestamp (in milliseconds since epoch) when the token was created.
+    pub created_at: i64,
+    /// The timestamp (in milliseconds since epoch) when the token was last updated.
+    pub updated_at: i64,
+    /// An optional description of the token's purpose or usage.
+    pub description: Option<String>,
+    /// The timestamp (in milliseconds since epoch) when the token was last used.
+    pub last_access_at: i64,
+    /// Optional access control settings
+    pub acl: Option<AccessControl>,
+}
+
+#[derive(Clone, Debug, Hash, PartialEq, Eq, Deserialize, Serialize, Object)]
+pub struct AccountInfo {
+    /// The unique identifier for the account.
+    pub id: u64,
+    /// The email address associated with the account.
+    pub email: String,
+}
+
+impl Ord for AccountInfo {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.id.cmp(&other.id)
+    }
+}
+
+impl PartialOrd for AccountInfo {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize, Object)]
+pub struct AccessControl {
+    /// An optional set of valid IPv4 or IPv6 addresses allowed to use the access token.
+    pub ip_whitelist: Option<BTreeSet<String>>,
+    /// An optional rate limit configuration for the access token.
+    pub rate_limit: Option<RateLimit>,
+}
+
+impl AccessControl {
+    pub fn validate(&self) -> BichonResult<()> {
+        if let Some(ip_whitelist) = &self.ip_whitelist {
+            for ip in ip_whitelist {
+                if ip.parse::<IpAddr>().is_err() {
+                    return Err(raise_error!(
+                        format!("Invalid IP address: {}", ip),
+                        ErrorCode::InvalidParameter
+                    ));
+                }
+            }
+        }
+
+        // Validate rate limit
+        if let Some(rate_limit) = &self.rate_limit {
+            if rate_limit.interval < 1 {
+                return Err(raise_error!(
+                    "Rate limit interval must be at least 1 second".into(),
+                    ErrorCode::InvalidParameter
+                ));
+            }
+            if rate_limit.quota < 1 {
+                return Err(raise_error!(
+                    "Rate limit quota must be at least 1".into(),
+                    ErrorCode::InvalidParameter
+                ));
+            }
+        }
+
+        Ok(())
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize, Object)]
+pub struct RateLimit {
+    /// The time window in seconds for the rate limit.
+    pub interval: u64,
+    /// The maximum number of allowed requests within the time window.
+    pub quota: u32,
+}
+
+impl AccessToken {
+    pub fn new(
+        token: String,
+        accounts: BTreeSet<AccountInfo>,
+        description: Option<String>,
+        acl: Option<AccessControl>,
+    ) -> Self {
+        Self {
+            token,
+            accounts,
+            created_at: utc_now!(),
+            updated_at: utc_now!(),
+            description,
+            last_access_at: Default::default(),
+            acl,
+        }
+    }
+
+    pub async fn try_update_access_timestamp(token: &str) -> BichonResult<AccessToken> {
+        let token = token.to_string();
+        update_impl(
+            DB_MANAGER.meta_db(),
+            |rw| {
+                rw.get()
+                    .primary::<AccessToken>(token)
+                    .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?
+                    .ok_or_else(|| {
+                        raise_error!("Token not exist.".into(), ErrorCode::ResourceNotFound)
+                    })
+            },
+            |current| {
+                let mut updated = current.clone();
+                updated.last_access_at = utc_now!();
+                Ok(updated)
+            },
+        )
+        .await
+    }
+
+    pub async fn grant_account_access(token: &str, account: AccountInfo) -> BichonResult<()> {
+        let token = token.to_string();
+        update_impl(
+            DB_MANAGER.meta_db(),
+            move |rw| {
+                rw.get()
+                    .primary::<AccessToken>(token.clone())
+                    .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?
+                    .ok_or_else(|| {
+                        raise_error!(
+                            format!(
+                            "The access token with token={} that you want to modify was not found.",
+                            token
+                        ),
+                            ErrorCode::ResourceNotFound
+                        )
+                    })
+            },
+            |current| {
+                let mut updated = current.clone();
+                updated.accounts.insert(account);
+                updated.updated_at = utc_now!();
+                Ok(updated)
+            },
+        )
+        .await?;
+        Ok(())
+    }
+
+    pub async fn update(token: &str, request: AccessTokenUpdateRequest) -> BichonResult<()> {
+        if request.should_skip_update() {
+            return Err(raise_error!(
+                "No changes detected in access scopes, description, or accounts. \
+                 Please modify at least one of these fields to perform an update."
+                    .into(),
+                ErrorCode::InvalidParameter
+            ));
+        }
+        request.validate().await?;
+
+        let account_infos = if let Some(accounts) = &request.accounts {
+            let mut account_infos = BTreeSet::new();
+            for account_id in accounts {
+                let account = AccountModel::get(*account_id).await?;
+                account_infos.insert(AccountInfo {
+                    id: *account_id,
+                    email: account.email,
+                });
+            }
+            account_infos
+        } else {
+            BTreeSet::new()
+        };
+
+        let token = token.to_string();
+        update_impl(
+            DB_MANAGER.meta_db(),
+            move |rw| {
+                rw.get()
+                    .primary::<AccessToken>(token.clone())
+                    .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?
+                    .ok_or_else(|| {
+                        raise_error!(
+                            format!(
+                            "The access token with token={} that you want to modify was not found.",
+                            token
+                        ),
+                            ErrorCode::ResourceNotFound
+                        )
+                    })
+            },
+            move |current| {
+                let mut updated = current.clone();
+                if let Some(description) = request.description {
+                    updated.description = Some(description);
+                }
+
+                if request.accounts.is_some() {
+                    updated.accounts = account_infos;
+                }
+
+                if let Some(acl) = request.acl {
+                    updated.acl = Some(acl);
+                }
+
+                updated.updated_at = utc_now!();
+                Ok(updated)
+            },
+        )
+        .await?;
+        Ok(())
+    }
+
+    pub async fn create(request: AccessTokenCreateRequest) -> BichonResult<String> {
+        // Validate request parameters first
+        request.validate().await?;
+
+        let AccessTokenCreateRequest {
+            accounts,
+            description,
+            acl,
+        } = request;
+
+        let mut account_infos = BTreeSet::new();
+        for &account_id in &accounts {
+            let account = AccountModel::get(account_id).await?;
+            account_infos.insert(AccountInfo {
+                id: account_id,
+                email: account.email,
+            });
+        }
+
+        let token = generate_token!(128);
+        let access_token = AccessToken::new(token.clone(), account_infos, description, acl);
+        insert_impl(DB_MANAGER.meta_db(), access_token).await?;
+        Ok(token)
+    }
+
+    pub async fn delete(token: &str) -> BichonResult<()> {
+        let token = token.to_string();
+        delete_impl(DB_MANAGER.meta_db(), move |rw| {
+            rw.get()
+                .primary::<AccessToken>(token.clone())
+                .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?
+                .ok_or_else(|| {
+                    raise_error!(
+                        format!("Token '{}' not found during deletion process.", token),
+                        ErrorCode::ResourceNotFound
+                    )
+                })
+        })
+        .await
+    }
+
+    pub async fn list_all() -> BichonResult<Vec<AccessToken>> {
+        list_all_impl(DB_MANAGER.meta_db()).await
+    }
+
+    pub async fn list_account_tokens(account_id: u64) -> BichonResult<Vec<AccessToken>> {
+        let all = AccessToken::list_all().await?;
+        let result: Vec<AccessToken> = all
+            .into_iter()
+            .filter(|e| {
+                e.accounts
+                    .iter()
+                    .any(|account_info| account_info.id == account_id)
+            })
+            .collect();
+        Ok(result)
+    }
+
+    pub async fn cleanup_account(account_id: u64) -> BichonResult<()> {
+        let tokens = Self::list_account_tokens(account_id).await?;
+        if tokens.is_empty() {
+            return Ok(());
+        }
+
+        for token in tokens {
+            update_impl(
+                DB_MANAGER.meta_db(),
+                move |rw| {
+                    rw.get()
+                        .primary::<AccessToken>(token.token.clone())
+                        .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?
+                        .ok_or_else(|| {
+                            raise_error!(
+                                format!("Cannot find access token, {}", token.token),
+                                ErrorCode::ResourceNotFound
+                            )
+                        })
+                },
+                move |current| {
+                    let mut updated = current.clone();
+                    updated.updated_at = utc_now!();
+                    updated.accounts.retain(|account| account.id != account_id);
+                    Ok(updated)
+                },
+            )
+            .await?;
+        }
+        Ok(())
+    }
+
+    pub fn can_access_account(&self, account_id: u64) -> bool {
+        self.accounts.iter().any(|account| account.id == account_id)
+    }
+}