use axum::{ extract::{Path, Query, State}, http::StatusCode, response::{IntoResponse, Redirect}, }; use axum_cookie::CookieManager; use jsonwebtoken::{encode, Algorithm, Header}; use serde::Serialize; use tracing::{debug, debug_span, info, instrument, trace, warn}; use crate::data::user as user_repo; use crate::{app::AppState, errors::ErrorResponse, session}; #[derive(Debug, serde::Deserialize)] pub struct OAuthCallbackParams { pub code: Option, pub state: Option, pub error: Option, pub error_description: Option, } #[derive(Debug, serde::Deserialize)] pub struct AuthorizeQuery { pub link: Option, } /// Handles the beginning of the OAuth authorization flow. /// /// Requires the `provider` path parameter, which determines the OAuth provider to use. #[instrument(skip_all, fields(provider = %provider))] pub async fn oauth_authorize_handler( State(app_state): State, Path(provider): Path, Query(aq): Query, cookie: CookieManager, ) -> axum::response::Response { let Some(prov) = app_state.auth.get(&provider) else { warn!(%provider, "Unknown OAuth provider"); return ErrorResponse::bad_request("invalid_provider", Some(provider)).into_response(); }; let is_linking = aq.link == Some(true); // Persist link intent using a short-lived cookie; callbacks won't carry our query params. if is_linking { cookie.add( axum_cookie::cookie::Cookie::builder("link", "1") .http_only(true) .same_site(axum_cookie::prelude::SameSite::Lax) .path("/") // TODO: Pick a reasonable max age that aligns with how long OAuth providers can successfully complete the flow. .max_age(std::time::Duration::from_secs(60 * 60)) .build(), ); } trace!(linking = %is_linking, "Starting OAuth authorization"); // Try to acquire the existing session (PKCE session is ignored) let existing_session = match session::get_session_token(&cookie) { Some(token) => match session::decode_jwt(&token, &app_state.jwt_decoding_key) { Some(claims) if !session::is_pkce_session(&claims) => Some(claims), Some(_) => { debug!("existing session ignored; it is a PKCE session"); None } None => { debug!("invalid session token"); return ErrorResponse::unauthorized("invalid session token").into_response(); } }, None => { debug!("missing session cookie"); return ErrorResponse::unauthorized("missing session cookie").into_response(); } }; // If linking is enabled, error if the session doesn't exist or is a PKCE session if is_linking && existing_session.is_none() { warn!("missing session cookie during linking flow, refusing"); return ErrorResponse::unauthorized("missing session cookie").into_response(); } let auth_info = match prov.authorize(&app_state.jwt_encoding_key).await { Ok(info) => info, Err(e) => return e.into_response(), }; let final_token = if let Some(mut claims) = existing_session { // We have a user session and are linking. Merge PKCE info into it. if let Some(pkce_claims) = session::decode_jwt(&auth_info.session_token, &app_state.jwt_decoding_key) { claims.pkce_verifier = pkce_claims.pkce_verifier; claims.csrf_state = pkce_claims.csrf_state; // re-encode encode(&Header::new(Algorithm::HS256), &claims, &app_state.jwt_encoding_key).expect("jwt sign") } else { warn!("Failed to decode PKCE session token during linking flow"); // Fallback to just using the PKCE token, which will break linking but not panic. auth_info.session_token } } else { // Not linking or no existing session, just use the new token. auth_info.session_token }; session::set_session_cookie(&cookie, &final_token); trace!("Redirecting to provider authorization page"); Redirect::to(auth_info.authorize_url.as_str()).into_response() } /// Handles the callback from the OAuth provider after the user has authorized the app. /// /// Requires the `provider` path parameter, which determines the OAuth provider to use for finishing the OAuth flow. /// Requires the `code` and `state` query parameters, which are returned by the OAuth provider after the user has authorized the app. pub async fn oauth_callback_handler( State(app_state): State, Path(provider): Path, Query(params): Query, cookie: CookieManager, ) -> axum::response::Response { // Validate provider let Some(prov) = app_state.auth.get(&provider) else { warn!(%provider, "Unknown OAuth provider"); return ErrorResponse::bad_request("invalid_provider", Some(provider)).into_response(); }; // Process callback-returned errors from provider if let Some(error) = params.error { warn!(%provider, error = %error, desc = ?params.error_description, "OAuth callback returned an error"); return ErrorResponse::bad_request(error, params.error_description).into_response(); } // Acquire required parameters let Some(code) = params.code.as_deref() else { return ErrorResponse::bad_request("invalid_request", Some("missing code".into())).into_response(); }; let Some(state) = params.state.as_deref() else { return ErrorResponse::bad_request("invalid_request", Some("missing state".into())).into_response(); }; debug_span!("oauth_callback_handler", provider = %provider, code = %code, state = %state); // Handle callback from provider let user = match prov.handle_callback(code, state, &cookie, &app_state.jwt_decoding_key).await { Ok(u) => u, Err(e) => { warn!(%provider, "OAuth callback handling failed"); return e.into_response(); } }; debug!(cookies = ?cookie.cookie().iter().collect::>(), "Cookies"); // Linking or sign-in flow. Determine link intent from cookie (set at authorize time) let link_cookie = cookie.get("link").map(|c| c.value().to_string()); if link_cookie.is_some() { cookie.remove("link"); } let email = user.email.as_deref(); // Determine linking intent with a valid session if link_cookie.as_deref() == Some("1") { debug!("Link intent present"); if let Some(claims) = session::get_session_token(&cookie).and_then(|t| session::decode_jwt(&t, &app_state.jwt_decoding_key)) { // Perform linking with current session user let (cur_prov, cur_id) = claims.subject.split_once(':').unwrap_or(("", "")); let current_user = match user_repo::find_user_by_provider_id(&app_state.db, cur_prov, cur_id).await { Ok(Some(u)) => u, Ok(None) => { warn!("Current session user not found; proceeding as normal sign-in"); return ErrorResponse::bad_request("invalid_request", Some("current session user not found".into())) .into_response(); } Err(_) => { return ErrorResponse::with_status(StatusCode::INTERNAL_SERVER_ERROR, "database_error", None).into_response(); } }; if let Err(e) = user_repo::link_oauth_account( &app_state.db, current_user.id, &provider, &user.id, email, Some(&user.username), user.name.as_deref(), user.avatar_url.as_deref(), ) .await { warn!(error = %e, %provider, "Failed to link OAuth account"); return ErrorResponse::with_status(StatusCode::INTERNAL_SERVER_ERROR, "database_error", None).into_response(); } return (StatusCode::FOUND, Redirect::to("/profile")).into_response(); } else { warn!(%provider, "Link intent present but session missing/invalid; proceeding as normal sign-in"); } } // Normal sign-in: do NOT auto-link by email (security). If email exists, require linking flow. if let Some(e) = email { if let Ok(Some(existing)) = user_repo::find_user_by_email(&app_state.db, e).await { // Only block if the user already has at least one linked provider. // NOTE: We do not check whether providers are currently active. If a user has exactly one provider and it is inactive, // this may lock them out until the provider is reactivated or a manual admin link is performed. match user_repo::get_oauth_account_count_for_user(&app_state.db, existing.id).await { Ok(count) if count > 0 => { // Check if the "new" provider is already linked to the user match user_repo::find_user_by_provider_id(&app_state.db, &provider, &user.id).await { Ok(Some(_)) => { debug!( %provider, %existing.id, "Provider already linked to user, signing in normally"); } Ok(None) => { debug!( %provider, %existing.id, "Provider not linked to user, failing" ); return ErrorResponse::bad_request( "account_exists", Some(format!( "An account already exists for {}. Sign in with your existing provider, then visit /auth/{}?link=true to add this provider.", e, provider )), ) .into_response(); } Err(e) => { warn!(error = %e, %provider, "Failed to find user by provider ID"); return ErrorResponse::with_status(StatusCode::INTERNAL_SERVER_ERROR, "database_error", None) .into_response(); } } } Ok(_) => { // No providers linked yet: safe to associate this provider if let Err(e) = user_repo::link_oauth_account( &app_state.db, existing.id, &provider, &user.id, email, Some(&user.username), user.name.as_deref(), user.avatar_url.as_deref(), ) .await { warn!(error = %e, %provider, "Failed to link OAuth account to existing user with no providers"); return ErrorResponse::with_status(StatusCode::INTERNAL_SERVER_ERROR, "database_error", None) .into_response(); } } Err(e) => { warn!(error = %e, "Failed to count oauth accounts for user"); return ErrorResponse::with_status(StatusCode::INTERNAL_SERVER_ERROR, "database_error", None).into_response(); } } } else { // Create new user with email match user_repo::create_user( &app_state.db, &user.username, user.name.as_deref(), email, user.avatar_url.as_deref(), &provider, &user.id, ) .await { Ok(u) => u, Err(e) => { warn!(error = %e, %provider, "Failed to create user"); return ErrorResponse::with_status(StatusCode::INTERNAL_SERVER_ERROR, "database_error", None).into_response(); } }; } } else { // No email available: disallow sign-in for safety return ErrorResponse::bad_request( "invalid_request", Some("account has no email; sign in with a different provider".into()), ) .into_response(); } // Create session token let session_token = session::create_jwt_for_user(&provider, &user, &app_state.jwt_encoding_key); session::set_session_cookie(&cookie, &session_token); info!(%provider, "Signed in successfully"); // Process avatar asynchronously (don't block the response) if let Some(avatar_url) = user.avatar_url.as_deref() { let image_storage = app_state.image_storage.clone(); let user_public_id = user.id.clone(); let avatar_url = avatar_url.to_string(); debug!(%user_public_id, %avatar_url, "Processing avatar"); tokio::spawn(async move { match image_storage.process_avatar(&user_public_id, &avatar_url).await { Ok(avatar_urls) => { info!( user_id = %user_public_id, original_url = %avatar_urls.original_url, mini_url = %avatar_urls.mini_url, "Avatar processed successfully" ); } Err(e) => { warn!( user_id = %user_public_id, avatar_url = %avatar_url, error = %e, "Failed to process avatar" ); } } }); } (StatusCode::FOUND, Redirect::to("/profile")).into_response() } /// Handles the request to the profile endpoint. /// /// Requires the `session` cookie to be present. pub async fn profile_handler(State(app_state): State, cookie: CookieManager) -> axum::response::Response { let Some(token_str) = session::get_session_token(&cookie) else { debug!("Missing session cookie"); return ErrorResponse::unauthorized("missing session cookie").into_response(); }; let Some(claims) = session::decode_jwt(&token_str, &app_state.jwt_decoding_key) else { debug!("Invalid session token"); return ErrorResponse::unauthorized("invalid session token").into_response(); }; // sub format: provider:provider_user_id let (prov, prov_user_id) = match claims.subject.split_once(':') { Some((p, id)) => (p, id), None => { debug!("Malformed session token subject"); return ErrorResponse::unauthorized("invalid session token").into_response(); } }; match user_repo::find_user_by_provider_id(&app_state.db, prov, prov_user_id).await { Ok(Some(db_user)) => { // Include linked providers in the profile payload match user_repo::list_user_providers(&app_state.db, db_user.id).await { Ok(providers) => { #[derive(Serialize)] struct ProfilePayload { id: i64, email: Option, providers: Vec, created_at: chrono::DateTime, updated_at: chrono::DateTime, } let body = ProfilePayload { id: db_user.id, email: db_user.email.clone(), providers, created_at: db_user.created_at, updated_at: db_user.updated_at, }; axum::Json(body).into_response() } Err(e) => { warn!(error = %e, "Failed to list user providers"); ErrorResponse::with_status( StatusCode::INTERNAL_SERVER_ERROR, "database_error", Some("could not fetch providers".into()), ) .into_response() } } } Ok(None) => { debug!("User not found for session"); ErrorResponse::unauthorized("session not found").into_response() } Err(e) => { warn!(error = %e, "Failed to fetch user for session"); ErrorResponse::with_status( StatusCode::INTERNAL_SERVER_ERROR, "database_error", Some("could not fetch user".into()), ) .into_response() } } } pub async fn logout_handler(State(app_state): State, cookie: CookieManager) -> axum::response::Response { if let Some(token_str) = session::get_session_token(&cookie) { // Remove from in-memory sessions if present app_state.sessions.remove(&token_str); } session::clear_session_cookie(&cookie); info!("Signed out successfully"); (StatusCode::FOUND, Redirect::to("/")).into_response() } #[derive(Serialize)] struct ProviderInfo { id: &'static str, name: &'static str, active: bool, } pub async fn list_providers_handler(State(app_state): State) -> axum::response::Response { let providers: Vec = app_state .auth .values() .map(|provider| ProviderInfo { id: provider.id(), name: provider.label(), active: provider.active(), }) .collect(); axum::Json(providers).into_response() } pub async fn health_handler( State(app_state): State, Query(params): Query>, ) -> axum::response::Response { // Force health check in debug mode #[cfg(debug_assertions)] if params.contains_key("force") { app_state.check_health().await; } let ok = app_state.health.read().await.ok(); let status = if ok { StatusCode::OK } else { StatusCode::SERVICE_UNAVAILABLE }; let body = serde_json::json!({ "ok": ok }); (status, axum::Json(body)).into_response() }