pds-gatekeeper / src / admin / oauth.rs

use super::session;
use crate::AppState;
use crate::admin::store::SqlAuthStore;
use axum::{
    extract::{Query, State},
    http::StatusCode,
    response::{Html, IntoResponse, Redirect, Response},
};
use axum_extra::extract::cookie::{Cookie, SameSite, SignedCookieJar};
use jacquard_common::CowStr;
use jacquard_identity::JacquardResolver;
use jacquard_identity::resolver::IdentityResolver;
use jacquard_oauth::session::ClientSessionData;
use jacquard_oauth::{
    atproto::{AtprotoClientMetadata, GrantType},
    client::OAuthClient,
    session::ClientData,
    types::{AuthorizeOptions, CallbackParams},
};
use serde::Deserialize;
use sqlx::SqlitePool;
use tracing::log;

/// Type alias for the concrete OAuthClient we use.
pub type AdminOAuthClient = OAuthClient<JacquardResolver, SqlAuthStore>;

/// Initialize the OAuth client for admin portal authentication.
pub fn init_oauth_client(
    pds_hostname: &str,
    pool: SqlitePool,
) -> Result<AdminOAuthClient, anyhow::Error> {
    // Build client metadata
    let client_id = format!("https://{}/admin/client-metadata.json", pds_hostname)
        .parse()
        .map_err(|_: url::ParseError| anyhow::anyhow!("Invalid client_id URL"))?;
    let redirect_uri = format!("https://{}/admin/oauth/callback", pds_hostname)
        .parse()
        .map_err(|_: url::ParseError| anyhow::anyhow!("Invalid redirect_uri URL"))?;
    let client_uri = format!("https://{}/admin/", pds_hostname)
        .parse()
        .map_err(|_: url::ParseError| anyhow::anyhow!("Invalid client_uri URL"))?;

    let config = AtprotoClientMetadata::new(
        client_id,
        Some(client_uri),
        vec![redirect_uri],
        vec![GrantType::AuthorizationCode],
        vec![jacquard_oauth::scopes::Scope::parse("atproto").expect("valid scope")],
        None,
    );

    let client_data = ClientData::new(None, config);
    let store = SqlAuthStore::new(pool);
    let client = OAuthClient::new(store, client_data);

    Ok(client)
}

/// GET /admin/client-metadata.json — Serves the OAuth client metadata.
pub async fn client_metadata_json(State(state): State<AppState>) -> Response {
    let pds_hostname = &state.app_config.pds_hostname;
    let client_id = format!("https://{}/admin/client-metadata.json", pds_hostname);
    let redirect_uri = format!("https://{}/admin/oauth/callback", pds_hostname);
    let client_uri = format!("https://{}/admin/", pds_hostname);

    let metadata = serde_json::json!({
        "client_id": client_id,
        "client_uri": client_uri,
        "redirect_uris": [redirect_uri],
        "grant_types": ["authorization_code"],
        "response_types": ["code"],
        "scope": "atproto",
        "token_endpoint_auth_method": "none",
        "application_type": "web",
        "dpop_bound_access_tokens": true,

    });

    (
        StatusCode::OK,
        [(axum::http::header::CONTENT_TYPE, "application/json")],
        serde_json::to_string_pretty(&metadata).unwrap_or_default(),
    )
        .into_response()
}

/// GET /admin/login — Renders the login page.
pub async fn get_login(
    State(state): State<AppState>,
    Query(params): Query<LoginQueryParams>,
) -> Response {
    let mut data = serde_json::json!({
        "pds_hostname": state.app_config.pds_hostname,
    });

    if let Some(error) = params.error {
        data["error"] = serde_json::Value::String(error);
    }

    use axum_template::TemplateEngine;
    match state.template_engine.render("admin/login.hbs", data) {
        Ok(html) => Html(html).into_response(),
        Err(e) => {
            tracing::error!("Failed to render login template: {}", e);
            StatusCode::INTERNAL_SERVER_ERROR.into_response()
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct LoginQueryParams {
    pub error: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct LoginForm {
    pub handle: String,
}

/// POST /admin/login — Initiates the OAuth flow.
pub async fn post_login(
    State(state): State<AppState>,
    axum::extract::Form(form): axum::extract::Form<LoginForm>,
) -> Response {
    let oauth_client: &AdminOAuthClient = match &state.admin_oauth_client {
        Some(client) => client,
        None => return StatusCode::NOT_FOUND.into_response(),
    };

    let pds_hostname = &state.app_config.pds_hostname;
    let redirect_uri: url::Url = match format!("https://{}/admin/oauth/callback", pds_hostname)
        .parse()
    {
        Ok(u) => u,
        Err(_) => {
            return Redirect::to("/admin/login?error=Invalid+server+configuration").into_response();
        }
    };

    let options = AuthorizeOptions {
        redirect_uri: Some(redirect_uri),
        scopes: vec![
            jacquard_oauth::scopes::Scope::parse("atproto").expect("valid scope"),
            jacquard_oauth::scopes::Scope::parse("transition:generic").expect("valid scope"),
        ],
        prompt: None,
        state: None,
    };

    match oauth_client.start_auth(&form.handle, options).await {
        Ok(auth_url) => Redirect::to(&auth_url).into_response(),
        Err(e) => {
            tracing::error!("OAuth start_auth failed: {}", e);
            let msg = format!("Login failed: {}", e);
            let error_msg = urlencoding::encode(&msg);
            Redirect::to(&format!("/admin/login?error={}", error_msg)).into_response()
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct OAuthCallbackParams {
    pub code: String,
    pub state: Option<String>,
    pub iss: Option<String>,
}

/// GET /admin/oauth/callback — Handles the OAuth callback.
pub async fn oauth_callback(
    State(state): State<AppState>,
    Query(params): Query<OAuthCallbackParams>,
    jar: SignedCookieJar,
) -> Response {
    let oauth_client: &AdminOAuthClient = match &state.admin_oauth_client {
        Some(client) => client,
        None => return StatusCode::NOT_FOUND.into_response(),
    };

    let rbac = match &state.admin_rbac_config {
        Some(rbac) => rbac,
        None => return StatusCode::NOT_FOUND.into_response(),
    };

    let callback_params = CallbackParams {
        code: params.code.as_str().into(),
        state: params.state.as_deref().map(Into::into),
        iss: params.iss.as_deref().map(Into::into),
    };

    // Exchange authorization code for session
    let oauth_session = match oauth_client.callback(callback_params).await {
        Ok(session) => session,
        Err(e) => {
            tracing::error!("OAuth callback failed: {}", e);
            let msg = format!("Authentication failed: {}", e);
            let error_msg = urlencoding::encode(&msg);
            return Redirect::to(&format!("/admin/login?error={}", error_msg)).into_response();
        }
    };

    // Extract DID and handle from the OAuth session
    let (did, session_id) = oauth_session.session_info().await;
    let session_id_string = session_id.to_string();
    log::info!("Authenticated as DID {}", did);
    let handle = match state.resolver.resolve_did_doc(&did).await {
        Ok(did_doc) => match did_doc.parse() {
            Ok(parsed_did_doc) => {
                let handles = parsed_did_doc.handles();
                if handles.len() > 0 {
                    handles[0].to_string()
                } else {
                    "Not found".to_string()
                }
            }
            Err(err) => {
                tracing::error!("Failed to parse DID document for {}: {}", did, err);
                "Not found".to_string()
            }
        },
        Err(err) => {
            tracing::error!("Failed to resolve DID document for {}: {}", did, err);
            //just default to the did
            "Not found".to_string()
        }
    };
    let did_str = did.to_string();

    // Check if this DID is a member in the RBAC config
    if !rbac.is_member(&did_str) {
        tracing::warn!("Access denied for DID {} (not in RBAC config)", did_str);
        return render_error(
            &state,
            "Access Denied",
            &format!(
                "Your identity ({}) is not authorized to access the admin portal. Contact your PDS administrator.",
                did_str
            ),
        );
    }

    // Create admin session
    let ttl_hours = state.app_config.admin_session_ttl_hours;
    let session_id = match session::create_session(
        &state.pds_gatekeeper_pool,
        &did_str,
        &handle,
        &session_id_string,
        ttl_hours,
    )
    .await
    {
        Ok(id) => id,
        Err(e) => {
            tracing::error!("Failed to create admin session: {}", e);
            return Redirect::to("/admin/login?error=Session+creation+failed").into_response();
        }
    };

    // Set signed cookie
    let mut cookie = Cookie::new("__gatekeeper_admin_session", session_id);
    cookie.set_http_only(true);
    cookie.set_secure(true);
    cookie.set_same_site(SameSite::Lax);
    cookie.set_path("/admin/");

    let updated_jar = jar.add(cookie);

    (updated_jar, Redirect::to("/admin/dashboard")).into_response()
}

fn render_error(state: &AppState, title: &str, message: &str) -> Response {
    let data = serde_json::json!({
        "error_title": title,
        "error_message": message,
        "pds_hostname": state.app_config.pds_hostname,
    });

    use axum_template::TemplateEngine;
    match state.template_engine.render("admin/error.hbs", data) {
        Ok(html) => Html(html).into_response(),
        Err(_) => (StatusCode::FORBIDDEN, format!("{}: {}", title, message)).into_response(),
    }
}