bot/src/server/auth.rs

use axum::{
    Json, RequestPartsExt,
    extract::{FromRef, FromRequestParts, Query, State},
    http::request::Parts,
    response::IntoResponse,
};
use axum_extra::{
    TypedHeader,
    extract::cookie::{Cookie, CookieJar, SameSite},
    headers::{Authorization, authorization::Bearer},
};
use std::sync::Arc;

use super::AppState;
use crate::error::{AppError, AppResult};

pub struct AuthenticatedUser(pub crate::domain::auth::Claims);

#[axum::async_trait]
impl<S> FromRequestParts<S> for AuthenticatedUser
where
    Arc<AppState>: axum::extract::FromRef<S>,
    S: Send + Sync,
{
    type Rejection = AppError;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let app_state = Arc::<AppState>::from_ref(state);

        let token = if let Ok(TypedHeader(Authorization(bearer))) =
            parts.extract::<TypedHeader<Authorization<Bearer>>>().await
        {
            Some(bearer.token().to_string())
        } else {
            let jar = parts
                .extract::<CookieJar>()
                .await
                .map_err(|e| AppError::Internal(e.to_string()))?;
            jar.get("access_token")
                .map(|cookie| cookie.value().to_string())
        };

        let token = token
            .ok_or_else(|| AppError::Unauthorized("Missing or invalid access token".into()))?;

        let claims = app_state
            .verifier
            .verify(&token)
            .await
            .map_err(|e| AppError::Unauthorized(format!("Token verification failed: {}", e)))?;

        Ok(AuthenticatedUser(claims))
    }
}

#[derive(serde::Deserialize)]
pub struct AuthCallbackQuery {
    pub code: String,
    pub redirect_uri: String,
}

#[derive(serde::Deserialize)]
pub struct RefreshRequest {
    pub refresh_token: Option<String>,
}

pub async fn auth_refresh(
    State(state): State<Arc<AppState>>,
    jar: CookieJar,
    Json(payload): Json<RefreshRequest>,
) -> AppResult<impl IntoResponse> {
    let refresh_token = payload
        .refresh_token
        .filter(|token| !token.is_empty())
        .or_else(|| {
            jar.get("refresh_token")
                .map(|cookie| cookie.value().to_string())
        })
        .ok_or_else(|| AppError::Unauthorized("Missing refresh token".into()))?;

    let data = state
        .authenticator
        .refresh_token(refresh_token)
        .await
        .map_err(|e| AppError::Unauthorized(e.to_string()))?;

    let jar = update_auth_cookies(jar, &data, &state.config);
    Ok((jar, Json(data)))
}

pub async fn auth_callback(
    State(state): State<Arc<AppState>>,
    jar: CookieJar,
    Query(query): Query<AuthCallbackQuery>,
) -> AppResult<impl IntoResponse> {
    let data = state
        .authenticator
        .exchange_code(query.code, query.redirect_uri)
        .await
        .map_err(|e| AppError::Internal(format!("Token exchange failed: {}", e)))?;

    let jar = update_auth_cookies(jar, &data, &state.config);
    Ok((jar, Json(data)))
}

pub async fn auth_logout(State(state): State<Arc<AppState>>, jar: CookieJar) -> impl IntoResponse {
    let jar = clear_auth_cookies(jar, &state.config);
    (jar, axum::http::StatusCode::NO_CONTENT)
}

pub async fn auth_session(user: AuthenticatedUser) -> Json<serde_json::Value> {
    Json(serde_json::json!({
        "authenticated": true,
        "user": user.0
    }))
}

fn secure(config: &crate::config::Config) -> bool {
    config.cookie_secure
}

pub fn update_auth_cookies(
    jar: CookieJar,
    data: &serde_json::Value,
    config: &crate::config::Config,
) -> CookieJar {
    let access_token = data.get("access_token");
    let refresh_token = data.get("refresh_token");

    let mut jar = jar;

    if let Some(token) = access_token.and_then(|t| t.as_str()) {
        let cookie = Cookie::build(("access_token", token.to_owned()))
            .path("/")
            .http_only(true)
            .same_site(SameSite::Lax)
            .secure(secure(config))
            .build();
        jar = jar.add(cookie);
    }

    if let Some(token) = refresh_token.and_then(|t| t.as_str()) {
        let cookie = Cookie::build(("refresh_token", token.to_owned()))
            .path("/")
            .http_only(true)
            .same_site(SameSite::Lax)
            .secure(secure(config))
            .build();
        jar = jar.add(cookie);
    }

    jar
}

pub fn clear_auth_cookies(jar: CookieJar, config: &crate::config::Config) -> CookieJar {
    let mut jar = jar;
    for name in ["access_token", "refresh_token"] {
        let cookie = Cookie::build((name, ""))
            .path("/")
            .http_only(true)
            .same_site(SameSite::Lax)
            .secure(secure(config))
            .max_age(cookie::time::Duration::seconds(0))
            .build();
        jar = jar.add(cookie);
    }

    jar
}