feat(auth): Enhance Google OAuth flow with async email extraction and caching

This commit is contained in:
MythEclipse
2025-08-12 19:25:10 +07:00
parent 44b1e09551
commit b40a430c49
11 changed files with 204 additions and 34 deletions
Generated
+1
View File
@@ -2292,6 +2292,7 @@ dependencies = [
"imphnen-entities",
"imphnen-libs",
"rand 0.9.1",
"reqwest 0.11.27",
"serde",
"serde_json",
"sha2",
@@ -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<Arc<Self>>, Query(auth_request): Query<AuthRequest>| async move {
move |State(controller): State<Arc<Self>>, Extension(app_state): Extension<AppState>, Query(auth_request): Query<AuthRequest>| 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<Json<AuthLoginResponsetDto>, 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<Json<AuthLoginResponsetDto>, Error> {
let (user, token) = self.google_oauth_service.google_oauth_callback(auth_request, app_state).await?;
let auth_response = AuthLoginResponsetDto {
user,
token,
@@ -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<A: AuthServiceTrait + Send + Sync + 'static, U: Use
fn with_services(auth_service: A, users_service: U, env: &'static Env) -> 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))
}
@@ -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<PermissionsEnum>,
) -> 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
+21 -3
View File
@@ -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"),
+14 -5
View File
@@ -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<Response, Infallible> {
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<UsersDetailQueryDto> =
match repository.query_user_by_email(email).await {
@@ -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,
+1
View File
@@ -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"] }
+109 -4
View File
@@ -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<String> {
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<String> {
}
};
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<String> {
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<String> {
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::<Value>().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<String> {
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<String> {
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<String> {
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
}
}
}
+1 -1
View File
@@ -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::*;
@@ -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<String, ErrorResponse>;
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>;
}
}