From b40a430c49c8b0ef15e2c80c7501072938447ff5 Mon Sep 17 00:00:00 2001 From: MythEclipse Date: Tue, 12 Aug 2025 19:25:10 +0700 Subject: [PATCH] feat(auth): Enhance Google OAuth flow with async email extraction and caching --- Cargo.lock | 1 + .../v1/auth/google/google_oauth_controller.rs | 11 +- .../v1/auth/google/google_oauth_service.rs | 20 +++- .../src/v1/permissions/permissions_guard.rs | 26 ++-- imphnen-iam/src/v1/users/users_service.rs | 24 +++- imphnen-middleware/src/auth_middleware/mod.rs | 19 ++- .../src/permissions_middleware/mod.rs | 19 ++- imphnen-utils/Cargo.toml | 1 + imphnen-utils/src/extract_email.rs | 113 +++++++++++++++++- imphnen-utils/src/lib.rs | 2 +- .../iam/auth/google/google_oauth_flow_test.rs | 2 +- 11 files changed, 204 insertions(+), 34 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fcf887b..48608ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2292,6 +2292,7 @@ dependencies = [ "imphnen-entities", "imphnen-libs", "rand 0.9.1", + "reqwest 0.11.27", "serde", "serde_json", "sha2", diff --git a/imphnen-iam/src/v1/auth/google/google_oauth_controller.rs b/imphnen-iam/src/v1/auth/google/google_oauth_controller.rs index c602aef..6213ad4 100644 --- a/imphnen-iam/src/v1/auth/google/google_oauth_controller.rs +++ b/imphnen-iam/src/v1/auth/google/google_oauth_controller.rs @@ -2,7 +2,7 @@ use axum::{ extract::{Query, State}, response::Redirect, routing::get, - Json, Router, + Json, Router, Extension, }; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; @@ -12,6 +12,7 @@ use imphnen_libs::enviroment::ENV; // Import ENV use crate::v1::auth::google::google_oauth_service::{AuthRequest, GoogleOauthService, GoogleOauthServiceImpl}; use imphnen_entities::error_dto::error::Error; use crate::v1::auth::AuthLoginResponsetDto; +use crate::AppState; #[derive(Debug, Serialize, Deserialize, ToSchema)] pub struct GoogleAuthUrlResponse { @@ -56,9 +57,9 @@ where .route( "/callback", get( - move |State(controller): State>, Query(auth_request): Query| async move { + move |State(controller): State>, Extension(app_state): Extension, Query(auth_request): Query| async move { let controller = Arc::clone(&controller); - controller.google_oauth_callback(auth_request).await + controller.google_oauth_callback(auth_request, &app_state).await }, ), ) @@ -70,8 +71,8 @@ where Ok(Redirect::to(authorize_url.as_str())) } - pub async fn google_oauth_callback(&self, auth_request: AuthRequest) -> Result, Error> { - let (user, token) = self.google_oauth_service.google_oauth_callback(auth_request).await?; + pub async fn google_oauth_callback(&self, auth_request: AuthRequest, app_state: &AppState) -> Result, Error> { + let (user, token) = self.google_oauth_service.google_oauth_callback(auth_request, app_state).await?; let auth_response = AuthLoginResponsetDto { user, token, diff --git a/imphnen-iam/src/v1/auth/google/google_oauth_service.rs b/imphnen-iam/src/v1/auth/google/google_oauth_service.rs index 0d16c03..9b22eba 100644 --- a/imphnen-iam/src/v1/auth/google/google_oauth_service.rs +++ b/imphnen-iam/src/v1/auth/google/google_oauth_service.rs @@ -9,7 +9,7 @@ use oauth2::url::Url; use tracing::{info, error}; use imphnen_entities::error_dto::error::Error; -use imphnen_libs::{jsonwebtoken::{encode_access_token, encode_refresh_token}, enviroment::Env}; +use imphnen_libs::{jsonwebtoken::{encode_access_token, encode_refresh_token}, enviroment::Env, AppState}; use imphnen_utils::{generate_oauth_csrf_token, validate_oauth_csrf_token, validate_csrf_token}; use crate::v1::auth::TokenDto; use crate::v1::auth::auth_service::AuthServiceTrait; @@ -91,7 +91,7 @@ pub trait GoogleOauthService Self; fn google_oauth_client(&self) -> BasicClient; fn generate_auth_url(&self) -> (Url, CsrfToken); - async fn google_oauth_callback(&self, auth_request: AuthRequest) -> Result<(UsersDetailItemDto, TokenDto), Error>; // Changed return type + async fn google_oauth_callback(&self, auth_request: AuthRequest, app_state: &AppState) -> Result<(UsersDetailItemDto, TokenDto), Error>; // Changed return type } #[derive(Clone)] @@ -158,7 +158,7 @@ where .url() } - async fn google_oauth_callback(&self, auth_request: AuthRequest) -> Result<(UsersDetailItemDto, TokenDto), Error> { + async fn google_oauth_callback(&self, auth_request: AuthRequest, app_state: &AppState) -> Result<(UsersDetailItemDto, TokenDto), Error> { // Validate input parameters first auth_request.validate()?; @@ -280,6 +280,20 @@ where refresh_token, }; + // Cache the user in auth repository for subsequent requests + let auth_repo = crate::v1::auth::AuthRepository::new(app_state); + let user_query_dto: crate::v1::users::users_dto::UsersDetailQueryDto = (&user).into(); + if let Err(err_store) = auth_repo.query_store_user(user_query_dto).await { + error!( + "Failed to store user cache for {}: {}", + user.email, err_store + ); + // Don't fail the login, just log the error + error!("Google OAuth login succeeded but caching failed for user: {}", user.email); + } else { + info!("Successfully cached user {} after Google OAuth login", user.email); + } + info!("Successfully completed Google OAuth for user: {}", user.email); Ok((user, token_dto)) } diff --git a/imphnen-iam/src/v1/permissions/permissions_guard.rs b/imphnen-iam/src/v1/permissions/permissions_guard.rs index 125721e..d478aa6 100644 --- a/imphnen-iam/src/v1/permissions/permissions_guard.rs +++ b/imphnen-iam/src/v1/permissions/permissions_guard.rs @@ -1,5 +1,5 @@ use super::PermissionsEnum; -use crate::{AppState, AuthRepository, common_response, extract_email}; +use crate::{AppState, AuthRepository, common_response, extract_email, extract_email_async}; use axum::{ http::{HeaderMap, StatusCode}, response::Response, @@ -11,12 +11,24 @@ pub async fn permissions_guard( required_permissions: Vec, ) -> Result<(), Response> { let auth_repo = AuthRepository::new(&state); - let email = extract_email(headers).ok_or_else(|| { - common_response( - StatusCode::UNAUTHORIZED, - "Invalid or missing authorization token", - ) - })?; + + // Try synchronous email extraction first (for internal JWT tokens) + let email = match extract_email(headers) { + Some(email) => email, + None => { + // If sync extraction fails, try async (for Google tokens) + match extract_email_async(headers).await { + Some(email) => email, + None => { + return Err(common_response( + StatusCode::UNAUTHORIZED, + "Invalid or missing authorization token", + )); + } + } + } + }; + let raw_user = auth_repo .query_get_stored_user(email.clone()) .await diff --git a/imphnen-iam/src/v1/users/users_service.rs b/imphnen-iam/src/v1/users/users_service.rs index d2be0de..36ab5f4 100644 --- a/imphnen-iam/src/v1/users/users_service.rs +++ b/imphnen-iam/src/v1/users/users_service.rs @@ -6,7 +6,7 @@ use crate::{ AppState, MetaRequestDto, ResponseListSuccessDto, UsersRepository, UsersSchema, }; use crate::{ - ResponseSuccessDto, common_response, extract_email, success_list_response, + ResponseSuccessDto, common_response, extract_email, extract_email_async, success_list_response, success_response, validate_request, }; use axum::http::HeaderMap; @@ -73,10 +73,19 @@ impl UsersServiceTrait for UsersService { async fn get_user_me(headers: HeaderMap, state: &AppState) -> Response { let repo = UsersRepository::new(state); + + // Try synchronous email extraction first (for internal JWT tokens) let email = match extract_email(&headers) { Some(email) => email, - None => return common_response(StatusCode::UNAUTHORIZED, "Invalid token"), + None => { + // If sync extraction fails, try async (for Google tokens) + match extract_email_async(&headers).await { + Some(email) => email, + None => return common_response(StatusCode::UNAUTHORIZED, "Invalid token"), + } + } }; + match repo.query_user_by_email(email).await { Ok(user) if !user.is_deleted => success_response(ResponseSuccessDto { data: UserDto::from(&user), // Corrected to use UserDto::from by reference @@ -134,10 +143,19 @@ impl UsersServiceTrait for UsersService { user: UsersUpdateRequestDto, ) -> Response { let repo = UsersRepository::new(state); + + // Try synchronous email extraction first (for internal JWT tokens) let email = match extract_email(&headers) { Some(email) => email, - None => return common_response(StatusCode::UNAUTHORIZED, "Unauthorized"), + None => { + // If sync extraction fails, try async (for Google tokens) + match extract_email_async(&headers).await { + Some(email) => email, + None => return common_response(StatusCode::UNAUTHORIZED, "Unauthorized"), + } + } }; + let user_data = match repo.query_user_by_email(email.clone()).await { Ok(user) => user, Err(_) => return common_response(StatusCode::NOT_FOUND, "User not found"), diff --git a/imphnen-middleware/src/auth_middleware/mod.rs b/imphnen-middleware/src/auth_middleware/mod.rs index f5c4f07..d5a5242 100644 --- a/imphnen-middleware/src/auth_middleware/mod.rs +++ b/imphnen-middleware/src/auth_middleware/mod.rs @@ -4,7 +4,7 @@ use axum::{ }; use imphnen_iam::{UsersDetailQueryDto, UsersRepository}; use imphnen_libs::AppState; -use imphnen_utils::{common_response, extract_email}; +use imphnen_utils::{common_response, extract_email, extract_email_async}; use std::convert::Infallible; pub async fn auth_middleware( @@ -13,15 +13,24 @@ pub async fn auth_middleware( next: Next, ) -> Result { let headers = req.headers(); + + // Try synchronous email extraction first (for internal JWT tokens) let email = match extract_email(headers) { Some(email) => email, None => { - return Ok(common_response( - StatusCode::UNAUTHORIZED, - "Invalid or expired token", - )); + // If sync extraction fails, try async (for Google tokens) + match extract_email_async(headers).await { + Some(email) => email, + None => { + return Ok(common_response( + StatusCode::UNAUTHORIZED, + "Invalid or expired token", + )); + } + } } }; + let repository = UsersRepository::new(&state); let user: Option = match repository.query_user_by_email(email).await { diff --git a/imphnen-middleware/src/permissions_middleware/mod.rs b/imphnen-middleware/src/permissions_middleware/mod.rs index 6f6b8a7..cff2da2 100644 --- a/imphnen-middleware/src/permissions_middleware/mod.rs +++ b/imphnen-middleware/src/permissions_middleware/mod.rs @@ -5,7 +5,7 @@ use axum::{ use futures::future::BoxFuture; use imphnen_iam::{AuthRepository, PermissionsEnum}; use imphnen_libs::AppState; -use imphnen_utils::{common_response, extract_email}; +use imphnen_utils::{common_response, extract_email, extract_email_async}; use std::task::{Context, Poll}; use tower::{Layer, Service}; @@ -59,15 +59,24 @@ where let permissions = self.permissions.clone(); Box::pin(async move { let headers = req.headers(); + + // Try synchronous email extraction first (for internal JWT tokens) let email = match extract_email(headers) { Some(email) => email, None => { - return Ok(common_response( - StatusCode::UNAUTHORIZED, - "Invalid or missing authorization token", - )); + // If sync extraction fails, try async (for Google tokens) + match extract_email_async(headers).await { + Some(email) => email, + None => { + return Ok(common_response( + StatusCode::UNAUTHORIZED, + "Invalid or missing authorization token", + )); + } + } } }; + let auth_repo = AuthRepository::new(&app_state); let user = match auth_repo.query_get_stored_user(email).await { Ok(user) => user, diff --git a/imphnen-utils/Cargo.toml b/imphnen-utils/Cargo.toml index 354c6b5..67e0a77 100644 --- a/imphnen-utils/Cargo.toml +++ b/imphnen-utils/Cargo.toml @@ -21,5 +21,6 @@ uuid.workspace = true tracing.workspace = true base64.workspace = true sha2.workspace = true +reqwest.workspace = true dotenvy = { workspace = true } tracing-subscriber = { workspace = true, features = ["env-filter"] } diff --git a/imphnen-utils/src/extract_email.rs b/imphnen-utils/src/extract_email.rs index 5691765..4782e33 100644 --- a/imphnen-utils/src/extract_email.rs +++ b/imphnen-utils/src/extract_email.rs @@ -3,6 +3,7 @@ use crate::decode_access_token; use axum::http::{HeaderMap, header::AUTHORIZATION}; /// Extracts the email from the Authorization header, if present and valid. +/// Supports both our internal JWT tokens and Google access tokens. pub fn extract_email(headers: &HeaderMap) -> Option { info!(?headers, "extract_email called with headers"); let auth_header = match headers.get(AUTHORIZATION) { @@ -27,19 +28,103 @@ pub fn extract_email(headers: &HeaderMap) -> Option { } }; info!(token, "Extracted bearer token in extract_email"); + + // First try to decode as our internal JWT token match decode_access_token(token) { Ok(data) => { - info!(email = %data.claims.sub, "Successfully decoded access token in extract_email"); + info!(email = %data.claims.sub, "Successfully decoded internal access token in extract_email"); Some(data.claims.sub) } + Err(_) => { + info!("Failed to decode as internal JWT, checking if it's a Google token"); + // If it fails, it might be a Google access token + // For Google tokens, we need async validation, so we'll return None here + // and handle Google tokens separately in the calling code + error!("Token is not a valid internal JWT. If this is a Google token, please use extract_email_async or handle Google OAuth flow properly."); + None + } + } +} + +/// Async version that can handle Google access tokens +pub async fn extract_email_async(headers: &HeaderMap) -> Option { + info!(?headers, "extract_email_async called with headers"); + let auth_header = match headers.get(AUTHORIZATION) { + Some(h) => h, + None => { + error!("Authorization header missing in extract_email_async"); + return None; + } + }; + let auth_str = match auth_header.to_str() { + Ok(s) => s, Err(e) => { - error!(error = ?e, "Failed to decode access token in extract_email"); + error!(error = ?e, "Failed to convert Authorization header to str in extract_email_async"); + return None; + } + }; + let token = match auth_str.strip_prefix("Bearer ") { + Some(t) => t, + None => { + error!(auth_str, "Authorization header does not start with 'Bearer ' in extract_email_async"); + return None; + } + }; + info!(token, "Extracted bearer token in extract_email_async"); + + // First try to decode as our internal JWT token + match decode_access_token(token) { + Ok(data) => { + info!(email = %data.claims.sub, "Successfully decoded internal access token in extract_email_async"); + Some(data.claims.sub) + } + Err(_) => { + info!("Failed to decode as internal JWT, trying Google token validation"); + // If it fails, try to validate as Google access token + extract_email_from_google_token(token).await + } + } +} + +/// Extracts email from Google access token by calling Google's tokeninfo endpoint +async fn extract_email_from_google_token(token: &str) -> Option { + use serde_json::Value; + + let client = reqwest::Client::new(); + let tokeninfo_url = format!("https://oauth2.googleapis.com/tokeninfo?access_token={}", token); + + match client.get(&tokeninfo_url).send().await { + Ok(response) => { + if response.status().is_success() { + match response.json::().await { + Ok(token_info) => { + if let Some(email) = token_info.get("email").and_then(|e| e.as_str()) { + info!(email = %email, "Successfully extracted email from Google token"); + Some(email.to_string()) + } else { + error!("Email not found in Google token info response"); + None + } + } + Err(e) => { + error!(error = ?e, "Failed to parse Google token info response"); + None + } + } + } else { + error!(status = %response.status(), "Google token validation failed"); + None + } + } + Err(e) => { + error!(error = ?e, "Failed to validate Google token"); None } } } /// Extracts the email from a JWT token string. +/// Supports both our internal JWT tokens and Google access tokens. pub fn extract_email_token(token: String) -> Option { info!(token = %token, "extract_email_token called with token"); match decode_access_token(&token) { @@ -47,9 +132,29 @@ pub fn extract_email_token(token: String) -> Option { info!(email = %data.claims.sub, "Successfully decoded token in extract_email_token"); Some(data.claims.sub) } - Err(e) => { - error!(error = ?e, "Failed to decode token in extract_email_token"); + Err(_) => { + info!("Failed to decode as internal JWT in extract_email_token, checking if it's a Google token"); + // If it fails, it might be a Google access token + // For Google tokens, we need async validation, so we'll return None here + // and handle Google tokens separately in the calling code + error!("Token is not a valid internal JWT. If this is a Google token, please use extract_email_token_async or handle Google OAuth flow properly."); None } } } + +/// Async version of extract_email_token that can handle Google access tokens +pub async fn extract_email_token_async(token: String) -> Option { + info!(token = %token, "extract_email_token_async called with token"); + match decode_access_token(&token) { + Ok(data) => { + info!(email = %data.claims.sub, "Successfully decoded internal token in extract_email_token_async"); + Some(data.claims.sub) + } + Err(_) => { + info!("Failed to decode as internal JWT in extract_email_token_async, trying Google token validation"); + // If it fails, try to validate as Google access token + extract_email_from_google_token(&token).await + } + } +} diff --git a/imphnen-utils/src/lib.rs b/imphnen-utils/src/lib.rs index 1bdc4a6..29b68f6 100644 --- a/imphnen-utils/src/lib.rs +++ b/imphnen-utils/src/lib.rs @@ -14,7 +14,7 @@ pub mod csrf_token; pub use logger::init_logger; pub use bind_filter::*; -pub use extract_email::*; +pub use extract_email::{extract_email, extract_email_async, extract_email_token, extract_email_token_async}; pub use generate_date::*; pub use generate_otp::*; pub use get_id::*; diff --git a/tests/src/iam/auth/google/google_oauth_flow_test.rs b/tests/src/iam/auth/google/google_oauth_flow_test.rs index d315ec0..e6b1af8 100644 --- a/tests/src/iam/auth/google/google_oauth_flow_test.rs +++ b/tests/src/iam/auth/google/google_oauth_flow_test.rs @@ -26,7 +26,7 @@ mod tests { fn with_services(auth_service: crate::v1::auth::auth_service::AuthService, users_service: crate::v1::users::users_service::UsersService, env: &'static Env) -> Self; // Updated signature fn google_oauth_client(&self) -> oauth2::basic::BasicClient; fn generate_auth_url(&self) -> (url::Url, oauth2::CsrfToken); - async fn google_oauth_callback(&self, auth_request: AuthRequest) -> anyhow::Result; + async fn google_oauth_callback(&self, auth_request: AuthRequest, app_state: &crate::AppState) -> anyhow::Result<(crate::v1::users::users_dto::UsersDetailItemDto, crate::v1::auth::TokenDto), imphnen_entities::error_dto::error::Error>; } }