From 23f9647e90fdc074d956b7ee34edbf16f03052dc Mon Sep 17 00:00:00 2001 From: Martin Stefcek Date: Thu, 5 Dec 2024 12:56:22 +0400 Subject: [PATCH 1/4] feat: add registration, login, api tokens endpoints --- Cargo.lock | 18 +++- atoma-auth/src/auth.rs | 90 +++++++++++++++- atoma-proxy-service/Cargo.toml | 2 + atoma-proxy-service/src/proxy_service.rs | 132 ++++++++++++++++++++++- atoma-proxy/src/main.rs | 4 +- atoma-state/src/handlers.rs | 88 +++++++++++++++ atoma-state/src/state_manager.rs | 42 +++++++- atoma-state/src/types.rs | 27 ++++- 8 files changed, 387 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f8a3485f..8a2ca6fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -609,12 +609,14 @@ name = "atoma-proxy-service" version = "0.1.0" dependencies = [ "anyhow", + "atoma-auth", "atoma-state", "atoma-sui", "axum", "config", "serde", "tokio", + "tower-http 0.6.2", "tracing", "tracing-subscriber", "utoipa", @@ -4399,7 +4401,7 @@ dependencies = [ "tonic", "tonic-health", "tower 0.4.13", - "tower-http", + "tower-http 0.5.2", "tracing", ] @@ -8030,6 +8032,20 @@ dependencies = [ "uuid", ] +[[package]] +name = "tower-http" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "403fa3b783d4b626a8ad51d766ab03cb6d2dbfc46b1c5d4448395e6628dc9697" +dependencies = [ + "bitflags 2.6.0", + "bytes", + "http 1.2.0", + "pin-project-lite", + "tower-layer", + "tower-service", +] + [[package]] name = "tower-layer" version = "0.3.3" diff --git a/atoma-auth/src/auth.rs b/atoma-auth/src/auth.rs index bc723aa5..62d7ce76 100644 --- a/atoma-auth/src/auth.rs +++ b/atoma-auth/src/auth.rs @@ -28,6 +28,7 @@ pub struct Claims { } /// The Auth struct +#[derive(Clone)] pub struct Auth { /// The secret key for JWT authentication. secret_key: String, @@ -67,8 +68,7 @@ impl Auth { /// # Errors /// /// * If the token generation fails - #[instrument(level = "trace", skip(self))] - fn generate_refresh_token(&self, user_id: i64) -> Result { + async fn generate_refresh_token(&self, user_id: i64) -> Result { let expiration = Utc::now() + Duration::days(self.refresh_token_lifetime as i64); let claims = Claims { user_id, @@ -80,6 +80,11 @@ impl Auth { &claims, &EncodingKey::from_secret(self.secret_key.as_ref()), )?; + self.state_manager_sender + .send(AtomaAtomaStateManagerEvent::StoreRefreshToken { + user_id, + refresh_token_hash: self.hash_string(&token), + })?; Ok(token) } @@ -185,13 +190,34 @@ impl Auth { /// # Returns /// /// * `String` - The hashed password - fn hash_string(&self, text: &str) -> String { + pub fn hash_string(&self, text: &str) -> String { let mut hasher = Blake2b::new(); hasher.update(text); let hash_result: GenericArray = hasher.finalize(); hex::encode(hash_result) } + /// Register user with username/password. + /// This method will register a new user with a username and password + /// The password is hashed and stored in the DB + /// The method will generate a new refresh and access token + pub async fn register(&self, username: &str, password: &str) -> Result<(String, String)> { + let (result_sender, result_receiver) = oneshot::channel(); + self.state_manager_sender + .send(AtomaAtomaStateManagerEvent::RegisterUserWithPassword { + username: username.to_string(), + password: self.hash_string(password), + result_sender, + })?; + let user_id = result_receiver + .await?? + .map(|user_id| user_id as u64) + .ok_or_else(|| anyhow::anyhow!("User already registred"))?; + let refresh_token = self.generate_refresh_token(user_id as i64).await?; + let access_token = self.generate_access_token(&refresh_token).await?; + Ok((refresh_token, access_token)) + } + /// Check the user password /// This method will check if the user password is correct /// The password is hashed and compared with the hashed password in the DB @@ -214,7 +240,7 @@ impl Auth { .await?? .map(|user_id| user_id as u64) .ok_or_else(|| anyhow::anyhow!("User not found"))?; - let refresh_token = self.generate_refresh_token(user_id as i64)?; + let refresh_token = self.generate_refresh_token(user_id as i64).await?; let access_token = self.generate_access_token(&refresh_token).await?; Ok((refresh_token, access_token)) } @@ -256,6 +282,50 @@ impl Auth { })?; Ok(api_token) } + + pub async fn revoke_api_token(&self, jwt: &str, api_token: &str) -> Result<()> { + let claims = self.validate_token(jwt, false)?; + if !self + .check_refresh_token_validity( + claims.user_id, + &claims + .refresh_token_hash + .expect("Access token should have refresh token hash"), + ) + .await? + { + return Err(anyhow::anyhow!("Access token was revoked")); + } + self.state_manager_sender + .send(AtomaAtomaStateManagerEvent::RevokeApiToken { + user_id: claims.user_id, + api_token: api_token.to_string(), + })?; + Ok(()) + } + + pub async fn get_all_api_tokens(&self, jwt: &str) -> Result> { + let claims = self.validate_token(jwt, false)?; + if !self + .check_refresh_token_validity( + claims.user_id, + &claims + .refresh_token_hash + .expect("Access token should have refresh token hash"), + ) + .await? + { + return Err(anyhow::anyhow!("Access token was revoked")); + } + + let (result_sender, result_receiver) = oneshot::channel(); + self.state_manager_sender + .send(AtomaAtomaStateManagerEvent::GetApiTokensForUser { + user_id: claims.user_id, + result_sender, + })?; + Ok(result_receiver.await??) + } } // TODO: Add more comprehensive tests, for now test the happy path only @@ -279,9 +349,14 @@ mod test { async fn test_access_token_regenerate() { let (auth, receiver) = setup_test(); let user_id = 123; - let refresh_token = auth.generate_refresh_token(user_id).unwrap(); + let refresh_token = auth.generate_refresh_token(user_id).await.unwrap(); let refresh_token_hash = auth.hash_string(&refresh_token); let mock_handle = tokio::task::spawn(async move { + let event = receiver.recv_async().await.unwrap(); + match event { + AtomaAtomaStateManagerEvent::StoreRefreshToken { .. } => {} + _ => panic!("Unexpected event"), + } let event = receiver.recv_async().await.unwrap(); match event { AtomaAtomaStateManagerEvent::IsRefreshTokenValid { @@ -330,6 +405,11 @@ mod test { } _ => panic!("Unexpected event"), } + let event = receiver.recv_async().await.unwrap(); + match event { + AtomaAtomaStateManagerEvent::StoreRefreshToken { .. } => {} + _ => panic!("Unexpected event"), + } for _ in 0..2 { // During the token generation, the refresh token is checked for validity // 1) when the user logs in diff --git a/atoma-proxy-service/Cargo.toml b/atoma-proxy-service/Cargo.toml index 8b817bcd..cb3fa60f 100644 --- a/atoma-proxy-service/Cargo.toml +++ b/atoma-proxy-service/Cargo.toml @@ -6,12 +6,14 @@ license.workspace = true [dependencies] anyhow.workspace = true +atoma-auth.workspace = true atoma-state.workspace = true atoma-sui.workspace = true axum.workspace = true config.workspace = true serde = { workspace = true, features = ["derive"] } tokio = { workspace = true, features = ["full"] } +tower-http = { version = "0.6.2", features = ["cors"] } tracing.workspace = true tracing-subscriber.workspace = true utoipa = { workspace = true, features = ["axum_extras"] } diff --git a/atoma-proxy-service/src/proxy_service.rs b/atoma-proxy-service/src/proxy_service.rs index 5154476b..d25bb951 100644 --- a/atoma-proxy-service/src/proxy_service.rs +++ b/atoma-proxy-service/src/proxy_service.rs @@ -1,15 +1,17 @@ +use atoma_auth::Auth; use atoma_state::{ - types::{NodeSubscription, Stack, Task}, + types::{AuthRequest, AuthResponse, NodeSubscription, RevokeApiTokenRequest, Stack, Task}, AtomaState, }; use axum::{ extract::{Path, State}, - http::StatusCode, - routing::get, + http::{HeaderMap, Method, StatusCode}, + routing::{get, post}, Json, Router, }; use tokio::{net::TcpListener, sync::watch::Receiver}; +use tower_http::cors::{Any, CorsLayer}; use tracing::{error, instrument}; type Result = std::result::Result; @@ -46,6 +48,9 @@ pub struct ProxyServiceState { /// Manages the persistent state of nodes, tasks, and other system components. /// Handles database operations and state synchronization. pub atoma_state: AtomaState, + + /// The authentication manager for the proxy service. + pub auth: Auth, } /// Starts and runs the Atoma proxy service service, handling HTTP requests and graceful shutdown. @@ -153,16 +158,137 @@ pub async fn run_proxy_service( /// .await?; /// ``` pub fn create_proxy_service_router(proxy_service_state: ProxyServiceState) -> Router { + let cors = CorsLayer::new() + .allow_origin(Any) + .allow_methods(vec![Method::GET, Method::POST]) + .allow_headers(Any); Router::new() .route("/subscriptions", get(get_all_subscriptions)) .route("/tasks", get(get_all_tasks)) .route("/task/:id", get(get_nodes_for_tasks)) .route("/stacks/:id", get(get_node_stacks)) .route("/get_stacks", get(get_current_stacks)) + .route("/register", post(register)) + .route("/login", post(login)) + .route("/api_tokens", get(get_all_api_tokens)) + .route("/generate_api_token", get(generate_api_token)) + .route("/revoke_api_token", post(revoke_api_token)) + .layer(cors) .with_state(proxy_service_state) .route("/health", get(health)) } +async fn get_all_api_tokens( + State(proxy_service_state): State, + headers: HeaderMap, +) -> Result>> { + let auth_header = headers + .get("Authorization") + .ok_or(StatusCode::UNAUTHORIZED)? + .to_str() + .map_err(|_| StatusCode::UNAUTHORIZED)?; + + let jwt = auth_header + .strip_prefix("Bearer ") + .ok_or(StatusCode::UNAUTHORIZED)?; + Ok(Json( + proxy_service_state + .auth + .get_all_api_tokens(jwt) + .await + .map_err(|e| { + error!("Failed to get all api tokens: {:?}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?, + )) +} + +async fn generate_api_token( + State(proxy_service_state): State, + headers: HeaderMap, +) -> Result> { + let auth_header = headers + .get("Authorization") + .ok_or(StatusCode::UNAUTHORIZED)? + .to_str() + .map_err(|_| StatusCode::UNAUTHORIZED)?; + + let jwt = auth_header + .strip_prefix("Bearer ") + .ok_or(StatusCode::UNAUTHORIZED)?; + Ok(Json( + proxy_service_state + .auth + .generate_api_token(jwt) + .await + .map_err(|e| { + error!("Failed to generate api token: {:?}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?, + )) +} + +async fn revoke_api_token( + State(proxy_service_state): State, + headers: HeaderMap, + body: Json, +) -> Result> { + let auth_header = headers + .get("Authorization") + .ok_or(StatusCode::UNAUTHORIZED)? + .to_str() + .map_err(|_| StatusCode::UNAUTHORIZED)?; + + let jwt = auth_header + .strip_prefix("Bearer ") + .ok_or(StatusCode::UNAUTHORIZED)?; + proxy_service_state + .auth + .revoke_api_token(jwt, &body.api_token) + .await + .map_err(|e| { + error!("Failed to revoke api token: {:?}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + Ok(Json(())) +} + +async fn register( + State(proxy_service_state): State, + body: Json, +) -> Result> { + let (refresh_token, access_token) = proxy_service_state + .auth + .register(&body.username, &body.password) + .await + .map_err(|e| { + error!("Failed to register user: {:?}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + Ok(Json(AuthResponse { + access_token, + refresh_token, + })) +} + +async fn login( + State(proxy_service_state): State, + body: Json, +) -> Result> { + let (refresh_token, access_token) = proxy_service_state + .auth + .check_user_password(&body.username, &body.password) + .await + .map_err(|e| { + error!("Failed to register user: {:?}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + Ok(Json(AuthResponse { + access_token, + refresh_token, + })) +} + /// Retrieves all stacks that are not settled. /// /// # Arguments diff --git a/atoma-proxy/src/main.rs b/atoma-proxy/src/main.rs index ed0038e6..d0bb7295 100644 --- a/atoma-proxy/src/main.rs +++ b/atoma-proxy/src/main.rs @@ -131,8 +131,7 @@ async fn main() -> Result<()> { let (confidential_compute_service_sender, _confidential_compute_service_receiver) = tokio::sync::mpsc::unbounded_channel(); - // TODO: Use this in the proxy service - let _auth = Auth::new(config.auth, state_manager_sender.clone()); + let auth = Auth::new(config.auth, state_manager_sender.clone()); let (_stack_retrieve_sender, stack_retrieve_receiver) = tokio::sync::mpsc::unbounded_channel(); let sui_subscriber = atoma_sui::SuiEventSubscriber::new( @@ -182,6 +181,7 @@ async fn main() -> Result<()> { let proxy_service_state = ProxyServiceState { atoma_state: AtomaState::new_from_url(&config.state.database_url).await?, + auth, }; let proxy_service_handle = spawn_with_shutdown( diff --git a/atoma-state/src/handlers.rs b/atoma-state/src/handlers.rs index a2e2528c..e26063d7 100644 --- a/atoma-state/src/handlers.rs +++ b/atoma-state/src/handlers.rs @@ -915,6 +915,94 @@ pub(crate) async fn handle_state_manager_event( .send(user_id) .map_err(|_| AtomaStateManagerError::ChannelSendError)?; } + AtomaAtomaStateManagerEvent::RegisterUserWithPassword { + username, + password, + result_sender, + } => { + let user_id = state_manager.state.register(&username, &password).await; + result_sender + .send(user_id) + .map_err(|_| AtomaStateManagerError::ChannelSendError)?; + } + AtomaAtomaStateManagerEvent::IsRefreshTokenValid { + user_id, + refresh_token_hash, + result_sender, + } => { + let is_valid = state_manager + .state + .is_refresh_token_valid(user_id, &refresh_token_hash) + .await; + result_sender + .send(is_valid) + .map_err(|_| AtomaStateManagerError::ChannelSendError)?; + } + AtomaAtomaStateManagerEvent::StoreRefreshToken { + user_id, + refresh_token_hash, + } => { + state_manager + .state + .store_refresh_token(user_id, &refresh_token_hash) + .await?; + } + AtomaAtomaStateManagerEvent::RevokeRefreshToken { + user_id, + refresh_token_hash, + } => { + state_manager + .state + .delete_refresh_token(user_id, &refresh_token_hash) + .await?; + } + AtomaAtomaStateManagerEvent::IsApiTokenValid { + user_id, + api_token, + result_sender, + } => { + let is_valid = state_manager + .state + .is_api_token_valid(user_id, &api_token) + .await; + result_sender + .send(is_valid) + .map_err(|_| AtomaStateManagerError::ChannelSendError)?; + } + AtomaAtomaStateManagerEvent::StoreNewApiToken { user_id, api_token } => { + state_manager + .state + .store_api_token(user_id, &api_token) + .await?; + } + AtomaAtomaStateManagerEvent::RevokeApiToken { user_id, api_token } => { + state_manager + .state + .delete_api_token(user_id, &api_token) + .await?; + } + AtomaAtomaStateManagerEvent::GetApiTokensForUser { + user_id, + result_sender, + } => { + let api_tokens = state_manager.state.get_api_tokens_for_user(user_id).await; + result_sender + .send(api_tokens) + .map_err(|_| AtomaStateManagerError::ChannelSendError)?; + } + AtomaAtomaStateManagerEvent::GetUserIdByUsernamePassword { + username, + password, + result_sender, + } => { + let user_id = state_manager + .state + .get_user_id_by_username_password(&username, &password) + .await; + result_sender + .send(user_id) + .map_err(|_| AtomaStateManagerError::ChannelSendError)?; + } AtomaAtomaStateManagerEvent::IsRefreshTokenValid { user_id, refresh_token_hash, diff --git a/atoma-state/src/state_manager.rs b/atoma-state/src/state_manager.rs index 55046faf..e13c77b9 100644 --- a/atoma-state/src/state_manager.rs +++ b/atoma-state/src/state_manager.rs @@ -616,6 +616,44 @@ impl AtomaState { .collect() } + /// Register user with password. + /// + /// This method inserts a new entry into the `users` table to register a new user. In case the user already exists, it returns None. + /// + /// # Arguments + /// + /// * `username` - The username of the user. + /// * `password_hash` - The password hash of the user. + /// + /// # Returns + /// + /// - `Result>`: A result containing either: + /// - `Ok(Some(i64))`: The ID of the user if the user was successfully registered. + /// - `Ok(None)`: If the user already exists. + /// - `Err(AtomaStateManagerError)`: An error if the database query fails. + /// + /// # Errors + /// + /// This function will return an error if the database query fails to execute. + /// + /// # Example + /// + /// ```rust,ignore + /// use atoma_node::atoma_state::AtomaStateManager; + /// + /// async fn register_user(state_manager: &AtomaStateManager, username: &str, password_hash: &str) -> Result, AtomaStateManagerError> { + /// state_manager.register(username, password_hash).await + /// } + /// ``` + pub async fn register(&self, username: &str, password_hash: &str) -> Result> { + let result = sqlx::query("INSERT INTO users (username, password_hash) VALUES ($1, $2) ON CONFLICT (username) DO NOTHING RETURNING id") + .bind(username) + .bind(password_hash) + .fetch_optional(&self.db) + .await?; + Ok(result.map(|record| record.get("id"))) + } + /// Checks if a node is subscribed to a specific task. /// /// This method queries the `node_subscriptions` table to determine if there's @@ -2676,7 +2714,7 @@ impl AtomaState { username: &str, hashed_password: &str, ) -> Result> { - let user = sqlx::query("SELECT id FROM users WHERE username = $1 AND password = $2") + let user = sqlx::query("SELECT id FROM users WHERE username = $1 AND password_hash = $2") .bind(username) .bind(hashed_password) .fetch_optional(&self.db) @@ -3428,7 +3466,7 @@ pub(crate) mod queries { "CREATE TABLE IF NOT EXISTS users ( id BIGSERIAL PRIMARY KEY, username VARCHAR(50) UNIQUE NOT NULL, - password_hash VARCHAR(50) NOT NULL + password_hash VARCHAR(64) NOT NULL )", ) .execute(db) diff --git a/atoma-state/src/types.rs b/atoma-state/src/types.rs index c9353146..249d6634 100644 --- a/atoma-state/src/types.rs +++ b/atoma-state/src/types.rs @@ -7,6 +7,22 @@ use tokio::sync::oneshot; use crate::state_manager::Result; +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, FromRow)] +pub struct RevokeApiTokenRequest { + pub api_token: String, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, FromRow)] +pub struct AuthRequest { + pub username: String, + pub password: String, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, FromRow)] +pub struct AuthResponse { + pub access_token: String, + pub refresh_token: String, +} /// Represents a task in the system #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, FromRow)] pub struct Task { @@ -287,6 +303,11 @@ pub enum AtomaAtomaStateManagerEvent { selected_node_id: i64, result_sender: oneshot::Sender>>>, }, + RegisterUserWithPassword { + username: String, + password: String, + result_sender: oneshot::Sender>>, + }, GetUserIdByUsernamePassword { username: String, password: String, @@ -297,13 +318,13 @@ pub enum AtomaAtomaStateManagerEvent { refresh_token_hash: String, result_sender: oneshot::Sender>, }, - StoresRefreshToken { + StoreRefreshToken { user_id: i64, - refresh_token: String, + refresh_token_hash: String, }, RevokeRefreshToken { user_id: i64, - refresh_token: String, + refresh_token_hash: String, }, IsApiTokenValid { user_id: i64, From 665662a978c3acc4dc9fb84447427ebd4fef48ee Mon Sep 17 00:00:00 2001 From: Martin Stefcek Date: Tue, 10 Dec 2024 09:26:17 +0100 Subject: [PATCH 2/4] fix clippy --- atoma-state/src/handlers.rs | 78 ------------------------------------- 1 file changed, 78 deletions(-) diff --git a/atoma-state/src/handlers.rs b/atoma-state/src/handlers.rs index e26063d7..ab91682e 100644 --- a/atoma-state/src/handlers.rs +++ b/atoma-state/src/handlers.rs @@ -990,84 +990,6 @@ pub(crate) async fn handle_state_manager_event( .send(api_tokens) .map_err(|_| AtomaStateManagerError::ChannelSendError)?; } - AtomaAtomaStateManagerEvent::GetUserIdByUsernamePassword { - username, - password, - result_sender, - } => { - let user_id = state_manager - .state - .get_user_id_by_username_password(&username, &password) - .await; - result_sender - .send(user_id) - .map_err(|_| AtomaStateManagerError::ChannelSendError)?; - } - AtomaAtomaStateManagerEvent::IsRefreshTokenValid { - user_id, - refresh_token_hash, - result_sender, - } => { - let is_valid = state_manager - .state - .is_refresh_token_valid(user_id, &refresh_token_hash) - .await; - result_sender - .send(is_valid) - .map_err(|_| AtomaStateManagerError::ChannelSendError)?; - } - AtomaAtomaStateManagerEvent::StoresRefreshToken { - user_id, - refresh_token, - } => { - state_manager - .state - .store_refresh_token(user_id, &refresh_token) - .await?; - } - AtomaAtomaStateManagerEvent::RevokeRefreshToken { - user_id, - refresh_token, - } => { - state_manager - .state - .delete_refresh_token(user_id, &refresh_token) - .await?; - } - AtomaAtomaStateManagerEvent::IsApiTokenValid { - user_id, - api_token, - result_sender, - } => { - let is_valid = state_manager - .state - .is_api_token_valid(user_id, &api_token) - .await; - result_sender - .send(is_valid) - .map_err(|_| AtomaStateManagerError::ChannelSendError)?; - } - AtomaAtomaStateManagerEvent::StoreNewApiToken { user_id, api_token } => { - state_manager - .state - .store_api_token(user_id, &api_token) - .await?; - } - AtomaAtomaStateManagerEvent::RevokeApiToken { user_id, api_token } => { - state_manager - .state - .delete_api_token(user_id, &api_token) - .await?; - } - AtomaAtomaStateManagerEvent::GetApiTokensForUser { - user_id, - result_sender, - } => { - let api_tokens = state_manager.state.get_api_tokens_for_user(user_id).await; - result_sender - .send(api_tokens) - .map_err(|_| AtomaStateManagerError::ChannelSendError)?; - } } Ok(()) } From f9132565601ef0e77201f720de7a6fb4aabf53ce Mon Sep 17 00:00:00 2001 From: Martin Stefcek Date: Tue, 10 Dec 2024 10:07:59 +0100 Subject: [PATCH 3/4] address comments --- Cargo.lock | 1 + atoma-auth/src/auth.rs | 25 ++++++ atoma-proxy-service/src/proxy_service.rs | 98 ++++++++++++++++++++++++ atoma-state/Cargo.toml | 1 + atoma-state/src/types.rs | 5 +- 5 files changed, 128 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8a2ca6fc..022ec9f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -635,6 +635,7 @@ dependencies = [ "thiserror 1.0.69", "tokio", "tracing", + "utoipa", ] [[package]] diff --git a/atoma-auth/src/auth.rs b/atoma-auth/src/auth.rs index 62d7ce76..8481d5f9 100644 --- a/atoma-auth/src/auth.rs +++ b/atoma-auth/src/auth.rs @@ -283,6 +283,19 @@ impl Auth { Ok(api_token) } + /// Revoke an API token + /// This method will revoke an API token for the user + /// The method will check if the access token and its corresponding refresh token is valid and revoke the API token in the state manager + /// + /// # Arguments + /// + /// * `jwt` - The access token to be used to revoke the API token + /// * `api_token` - The API token to be revoked + /// + /// # Returns + /// + /// * `Result<()>` - If the API token was revoked + #[instrument(level = "info", skip(self))] pub async fn revoke_api_token(&self, jwt: &str, api_token: &str) -> Result<()> { let claims = self.validate_token(jwt, false)?; if !self @@ -304,6 +317,18 @@ impl Auth { Ok(()) } + /// Get all API tokens for a user + /// This method will get all API tokens for a user + /// The method will check if the access token and its corresponding refresh token is valid + /// + /// # Arguments + /// + /// * `jwt` - The access token to be used to get the API tokens + /// + /// # Returns + /// + /// * `Result>` - The list of API tokens + #[instrument(level = "info", skip(self))] pub async fn get_all_api_tokens(&self, jwt: &str) -> Result> { let claims = self.validate_token(jwt, false)?; if !self diff --git a/atoma-proxy-service/src/proxy_service.rs b/atoma-proxy-service/src/proxy_service.rs index d25bb951..ad71c84a 100644 --- a/atoma-proxy-service/src/proxy_service.rs +++ b/atoma-proxy-service/src/proxy_service.rs @@ -178,6 +178,25 @@ pub fn create_proxy_service_router(proxy_service_state: ProxyServiceState) -> Ro .route("/health", get(health)) } +/// Retrieves all API tokens for the user. +/// +/// # Arguments +/// * `proxy_service_state` - The shared state containing the state manager +/// * `headers` - The headers of the request +/// +/// # Returns +/// +/// * `Result>>` - A JSON response containing a list of API tokens +#[utoipa::path( + get, + path = "", + responses( + (status = OK, description = "Retrieves all API tokens for the user", body = Value), + (status = UNAUTHORIZED, description = "Unauthorized request"), + (status = INTERNAL_SERVER_ERROR, description = "Failed to get all api tokens") + ) +)] +#[instrument(level = "info", skip_all)] async fn get_all_api_tokens( State(proxy_service_state): State, headers: HeaderMap, @@ -203,6 +222,26 @@ async fn get_all_api_tokens( )) } +/// Generates an API token for the user. +/// +/// # Arguments +/// +/// * `proxy_service_state` - The shared state containing the state manager +/// * `headers` - The headers of the request +/// +/// # Returns +/// +/// * `Result>` - A JSON response containing the generated API token +#[utoipa::path( + get, + path = "", + responses( + (status = OK, description = "Generates an API token for the user", body = Value), + (status = UNAUTHORIZED, description = "Unauthorized request"), + (status = INTERNAL_SERVER_ERROR, description = "Failed to generate api token") + ) +)] +#[instrument(level = "info", skip_all)] async fn generate_api_token( State(proxy_service_state): State, headers: HeaderMap, @@ -228,6 +267,27 @@ async fn generate_api_token( )) } +/// Revokes an API token for the user. +/// +/// # Arguments +/// +/// * `proxy_service_state` - The shared state containing the state manager +/// * `headers` - The headers of the request +/// * `body` - The request body containing the API token to revoke +/// +/// # Returns +/// +/// * `Result>` - A JSON response indicating the success of the operation +#[utoipa::path( + post, + path = "", + responses( + (status = OK, description = "Revokes an API token for the user", body = Value), + (status = UNAUTHORIZED, description = "Unauthorized request"), + (status = INTERNAL_SERVER_ERROR, description = "Failed to revoke api token") + ) +)] +#[instrument(level = "info", skip_all)] async fn revoke_api_token( State(proxy_service_state): State, headers: HeaderMap, @@ -253,6 +313,25 @@ async fn revoke_api_token( Ok(Json(())) } +/// Registers a new user with the proxy service. +/// +/// # Arguments +/// +/// * `proxy_service_state` - The shared state containing the state manager +/// * `body` - The request body containing the username and password of the new user +/// +/// # Returns +/// +/// * `Result>` - A JSON response containing the access and refresh tokens +#[utoipa::path( + post, + path = "", + responses( + (status = OK, description = "Registers a new user", body = Value), + (status = INTERNAL_SERVER_ERROR, description = "Failed to register user") + ) +)] +#[instrument(level = "trace", skip_all)] async fn register( State(proxy_service_state): State, body: Json, @@ -271,6 +350,25 @@ async fn register( })) } +/// Logs in a user with the proxy service. +/// +/// # Arguments +/// +/// * `proxy_service_state` - The shared state containing the state manager +/// * `body` - The request body containing the username and password of the user +/// +/// # Returns +/// +/// * `Result>` - A JSON response containing the access and refresh tokens +#[utoipa::path( + post, + path = "", + responses( + (status = OK, description = "Logs in a user", body = Value), + (status = INTERNAL_SERVER_ERROR, description = "Failed to login user") + ) +)] +#[instrument(level = "trace", skip_all)] async fn login( State(proxy_service_state): State, body: Json, diff --git a/atoma-state/Cargo.toml b/atoma-state/Cargo.toml index 5bc2f26d..35d962da 100644 --- a/atoma-state/Cargo.toml +++ b/atoma-state/Cargo.toml @@ -14,3 +14,4 @@ sqlx = { workspace = true, features = ["runtime-tokio-native-tls", "sqlite"] } thiserror = { workspace = true } tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } +utoipa.workspace = true diff --git a/atoma-state/src/types.rs b/atoma-state/src/types.rs index 249d6634..ec4fe2bc 100644 --- a/atoma-state/src/types.rs +++ b/atoma-state/src/types.rs @@ -4,15 +4,16 @@ use atoma_sui::events::{ use serde::{Deserialize, Serialize}; use sqlx::FromRow; use tokio::sync::oneshot; +use utoipa::ToSchema; use crate::state_manager::Result; -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, FromRow)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, FromRow, ToSchema)] pub struct RevokeApiTokenRequest { pub api_token: String, } -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, FromRow)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, FromRow, ToSchema)] pub struct AuthRequest { pub username: String, pub password: String, From 9b581b341095225c3693ad97a8e3f0292db3dc6a Mon Sep 17 00:00:00 2001 From: Martin Stefcek Date: Tue, 10 Dec 2024 10:11:45 +0100 Subject: [PATCH 4/4] update --- atoma-auth/src/auth.rs | 3 ++- atoma-state/src/state_manager.rs | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/atoma-auth/src/auth.rs b/atoma-auth/src/auth.rs index 8481d5f9..01d4cfa9 100644 --- a/atoma-auth/src/auth.rs +++ b/atoma-auth/src/auth.rs @@ -10,7 +10,7 @@ use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation} use rand::Rng; use serde::{Deserialize, Serialize}; use tokio::sync::oneshot; -use tracing::instrument; +use tracing::{error, instrument}; use crate::AtomaAuthConfig; @@ -340,6 +340,7 @@ impl Auth { ) .await? { + error!("Access token was revoked"); return Err(anyhow::anyhow!("Access token was revoked")); } diff --git a/atoma-state/src/state_manager.rs b/atoma-state/src/state_manager.rs index e13c77b9..4a0e3cd6 100644 --- a/atoma-state/src/state_manager.rs +++ b/atoma-state/src/state_manager.rs @@ -645,6 +645,7 @@ impl AtomaState { /// state_manager.register(username, password_hash).await /// } /// ``` + #[instrument(level = "trace", skip(self))] pub async fn register(&self, username: &str, password_hash: &str) -> Result> { let result = sqlx::query("INSERT INTO users (username, password_hash) VALUES ($1, $2) ON CONFLICT (username) DO NOTHING RETURNING id") .bind(username)