feat: add multi-user support and role-based access control #31

src/modules/token/mod.rs

@@ -16,12 +16,17 @@
 // 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 std::collections::HashMap;
 
-use crate::modules::account::migration::AccountModel;
-use crate::modules::database::delete_impl;
+use super::error::code::ErrorCode;
 use crate::modules::database::manager::DB_MANAGER;
+use crate::modules::database::{
+    async_find_impl, delete_impl, filter_by_secondary_key_impl, with_transaction,
+};
 use crate::modules::database::{insert_impl, list_all_impl, update_impl};
-use crate::modules::token::payload::AccessTokenUpdateRequest;
+use crate::modules::settings::cli::SETTINGS;
+use crate::modules::token::view::AccessTokenResp;
+use crate::modules::users::BichonUser;
 use crate::raise_error;
 use crate::{
     generate_token, modules::error::BichonResult,
@@ -29,259 +34,227 @@ use crate::{
 };
 use native_db::*;
 use native_model::{native_model, Model};
-use poem_openapi::Object;
+use poem_openapi::{Enum, Object};
 use serde::{Deserialize, Serialize};
-use std::collections::BTreeSet;
-use std::net::IpAddr;
-
-use super::error::code::ErrorCode;
 
 pub mod payload;
 pub mod root;
+pub mod view;
+
+// Starting from version 0.2.0, this model is deprecated/no longer used
+// #[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, PartialEq, Eq, Deserialize, Serialize, Enum)]
+pub enum TokenType {
+    WebUI,
+    Api,
+}
 
 #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize, Object)]
-#[native_model(id = 1, version = 1)]
+#[native_model(id = 11, version = 1)]
 #[native_db]
-pub struct AccessToken {
+pub struct AccessTokenModel {
+    /// The ID of the user who owns this token
+    #[secondary_key]
+    pub user_id: u64,
     /// 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>,
+    /// An optional name of the token.
+    pub name: Option<String>,
+    /// Token type: WebUI or API
+    pub token_type: TokenType,
     /// 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 expires.
+    /// None means the token does not expire (this applies only to API tokens).
+    pub expire_at: Option<i64>,
     /// 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(
+impl AccessTokenModel {
+    pub fn new_api_token(
         token: String,
-        accounts: BTreeSet<AccountInfo>,
-        description: Option<String>,
-        acl: Option<AccessControl>,
+        user_id: u64,
+        name: Option<String>,
+        expire_at: Option<i64>,
     ) -> Self {
         Self {
             token,
-            accounts,
             created_at: utc_now!(),
             updated_at: utc_now!(),
-            description,
             last_access_at: Default::default(),
-            acl,
+            name,
+            user_id,
+            token_type: TokenType::Api,
+            expire_at,
         }
     }
 
-    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 fn new_webui_token(user_id: u64) -> AccessTokenModel {
+        let now = utc_now!();
+        AccessTokenModel {
+            token: generate_token!(128),
+            created_at: now,
+            updated_at: now,
+            last_access_at: Default::default(),
+            name: None,
+            user_id,
+            token_type: TokenType::WebUI,
+            expire_at: None,
+        }
     }
 
-    pub async fn grant_account_access(token: &str, account: AccountInfo) -> BichonResult<()> {
-        let token = token.to_string();
-        update_impl(
+    pub async fn reset_webui_token(user_id: u64) -> BichonResult<String> {
+        let old_token = Self::get_user_webui_token(user_id).await?;
+        let new_token = Self::new_webui_token(user_id);
+        let new_token_str = new_token.token.clone();
+
+        match old_token {
+            Some(old) => {
+                with_transaction(DB_MANAGER.meta_db(), move |rw| {
+                    rw.remove(old)
+                        .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
+
+                    rw.insert(new_token)
+                        .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?;
+
+                    Ok(())
+                })
+                .await?;
+            }
+            None => {
+                insert_impl(DB_MANAGER.meta_db(), new_token).await?;
+            }
+        }
+
+        Ok(new_token_str)
+    }
+
+    pub async fn get_user_webui_token(user_id: u64) -> BichonResult<Option<AccessTokenModel>> {
+        let tokens = filter_by_secondary_key_impl::<AccessTokenModel>(
             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)
-            },
+            AccessTokenModelKey::user_id,
+            user_id,
         )
         .await?;
-        Ok(())
+
+        Ok(tokens
+            .into_iter()
+            .find(|t| t.token_type == TokenType::WebUI))
     }
 
-    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?;
+    pub async fn get_user_api_tokens(user_id: u64) -> BichonResult<Vec<AccessTokenModel>> {
+        let tokens = filter_by_secondary_key_impl::<AccessTokenModel>(
+            DB_MANAGER.meta_db(),
+            AccessTokenModelKey::user_id,
+            user_id,
+        )
+        .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,
-                });
+        Ok(tokens
+            .into_iter()
+            .filter(|t| t.token_type == TokenType::Api)
+            .collect())
+    }
+
+    pub async fn resolve_user_from_token(token: &str) -> BichonResult<BichonUser> {
+        let token = token.to_string();
+        let token_option = async_find_impl::<AccessTokenModel>(DB_MANAGER.meta_db(), token).await?;
+        let token = match token_option {
+            Some(token) => token,
+            None => {
+                return Err(raise_error!(
+                    "Permission denied: no valid access token provided.".into(),
+                    ErrorCode::PermissionDenied
+                ))
             }
-            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 matches!(token.token_type, TokenType::WebUI) {
+            let life = utc_now!() - token.created_at;
+            let max_life = SETTINGS.bichon_webui_token_expiration_hours * 60 * 60 * 1000;
 
-                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,
-            });
+            if life > (max_life as i64) {
+                return Err(raise_error!(
+                    "Permission denied: the WebUI token has expired.".into(),
+                    ErrorCode::PermissionDenied
+                ));
+            }
         }
 
+        if matches!(token.token_type, TokenType::Api) {
+            if let Some(expire_at) = token.expire_at {
+                if utc_now!() > expire_at {
+                    return Err(raise_error!(
+                        "Your API token has expired and is no longer valid.".into(),
+                        ErrorCode::PermissionDenied
+                    ));
+                }
+            }
+            let token = token.token.clone();
+            update_impl(
+                DB_MANAGER.meta_db(),
+                |rw| {
+                    rw.get()
+                        .primary::<AccessTokenModel>(token)
+                        .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?
+                        .ok_or_else(|| {
+                            raise_error!(
+                                "The access token does not exist or has been reset.".into(),
+                                ErrorCode::ResourceNotFound
+                            )
+                        })
+                },
+                |current| {
+                    let mut updated = current.clone();
+                    updated.last_access_at = utc_now!();
+                    Ok(updated)
+                },
+            )
+            .await?;
+        }
+
+        let user = BichonUser::find(token.user_id)
+            .await?
+            .ok_or_else(|| raise_error!("The user associated with this access token does not exist or may have been deleted.".into(), ErrorCode::ResourceNotFound))?;
+        Ok(user)
+    }
+
+    pub async fn create_api_token(
+        user_id: u64,
+        request: AccessTokenCreateRequest,
+    ) -> BichonResult<String> {
+        // Validate request parameters first
+        request.validate().await?;
+        let expire_at = request
+            .expire_in
+            .map(|hours| utc_now!() + (hours as i64) * 60 * 60 * 1000);
         let token = generate_token!(128);
-        let access_token = AccessToken::new(token.clone(), account_infos, description, acl);
+        let access_token =
+            AccessTokenModel::new_api_token(token.clone(), user_id, request.name, expire_at);
         insert_impl(DB_MANAGER.meta_db(), access_token).await?;
         Ok(token)
     }
@@ -290,7 +263,7 @@ impl AccessToken {
         let token = token.to_string();
         delete_impl(DB_MANAGER.meta_db(), move |rw| {
             rw.get()
-                .primary::<AccessToken>(token.clone())
+                .primary::<AccessTokenModel>(token.clone())
                 .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?
                 .ok_or_else(|| {
                     raise_error!(
@@ -302,56 +275,47 @@ impl AccessToken {
         .await
     }
 
-    pub async fn list_all() -> BichonResult<Vec<AccessToken>> {
-        list_all_impl(DB_MANAGER.meta_db()).await
+    pub async fn get_token(token: &str) -> BichonResult<AccessTokenModel> {
+        async_find_impl(DB_MANAGER.meta_db(), token.to_string())
+            .await?
+            .ok_or_else(|| {
+                raise_error!(
+                    format!("Access token '{}' not found", token),
+                    ErrorCode::ResourceNotFound
+                )
+            })
     }
 
-    pub async fn list_account_tokens(account_id: u64) -> BichonResult<Vec<AccessToken>> {
-        let all = AccessToken::list_all().await?;
-        let result: Vec<AccessToken> = all
+    pub async fn list_all_api_tokens() -> BichonResult<Vec<AccessTokenResp>> {
+        let users = BichonUser::list_all().await?;
+        let mut all = list_all_impl::<AccessTokenModel>(DB_MANAGER.meta_db()).await?;
+
+        all.retain(|t| t.token_type == TokenType::Api);
+        let user_map: HashMap<u64, BichonUser> = users.into_iter().map(|u| (u.id, u)).collect();
+
+        let resp = all
             .into_iter()
-            .filter(|e| {
-                e.accounts
-                    .iter()
-                    .any(|account_info| account_info.id == account_id)
+            .map(|token| {
+                let user = user_map.get(&token.user_id);
+                AccessTokenResp {
+                    user_name: user
+                        .map(|u| u.username.clone())
+                        .unwrap_or_else(|| "Unknown".to_string()),
+                    user_email: user
+                        .map(|u| u.email.clone())
+                        .unwrap_or_else(|| "N/A".to_string()),
+                    user_id: token.user_id,
+                    name: token.name,
+                    token: token.token,
+                    token_type: token.token_type,
+                    created_at: token.created_at,
+                    updated_at: token.updated_at,
+                    expire_at: token.expire_at,
+                    last_access_at: token.last_access_at,
+                }
             })
             .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)
+        Ok(resp)
     }
 }