pds-gatekeeper / src / gate.rs

use crate::AppState;
use crate::helpers::{generate_gate_token, json_error_response};
use axum::Form;
use axum::extract::{Query, State};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Redirect, Response};
use axum_template::RenderHtml;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::env;
use tracing::log;

#[derive(Deserialize)]
pub struct GateQuery {
    handle: String,
    state: String,
    #[serde(default)]
    error: Option<String>,
    #[serde(default)]
    redirect_url: Option<String>,
}

#[derive(Deserialize, Serialize)]
pub struct CaptchaPage {
    handle: String,
    state: String,
    captcha_site_key: String,
    error_message: Option<String>,
    pds: String,
    redirect_url: Option<String>,
}

#[derive(Deserialize)]
pub struct CaptchaForm {
    #[serde(rename = "h-captcha-response")]
    h_captcha_response: String,
    #[serde(default)]
    redirect_url: Option<String>,
}

/// GET /gate - Display the captcha page
pub async fn get_gate(
    Query(params): Query<GateQuery>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    let hcaptcha_site_key = match env::var("PDS_HCAPTCHA_SITE_KEY") {
        Ok(key) => key,
        Err(_) => {
            return json_error_response(
                StatusCode::INTERNAL_SERVER_ERROR,
                "ServerError",
                "hCaptcha is not configured",
            )
            .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response());
        }
    };

    let error_message = match params.error {
        None => None,
        Some(error) => Some(html_escape::encode_safe(&error).to_string()),
    };

    RenderHtml(
        "captcha.hbs",
        state.template_engine,
        CaptchaPage {
            handle: params.handle,
            state: params.state,
            captcha_site_key: hcaptcha_site_key,
            error_message,
            pds: state.app_config.pds_service_did.replace("did:web:", ""),
            redirect_url: params.redirect_url,
        },
    )
    .into_response()
}

/// POST /gate - Verify captcha and redirect
pub async fn post_gate(
    State(state): State<AppState>,
    Query(params): Query<GateQuery>,
    Form(form): Form<CaptchaForm>,
) -> Response {
    // Verify hCaptcha response
    let hcaptcha_secret = match env::var("PDS_HCAPTCHA_SECRET_KEY") {
        Ok(secret) => secret,
        Err(_) => {
            return json_error_response(
                StatusCode::INTERNAL_SERVER_ERROR,
                "ServerError",
                "hCaptcha is not configured",
            )
            .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response());
        }
    };

    let client = match reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(10))
        .build()
    {
        Ok(c) => c,
        Err(e) => {
            log::error!("Failed to create HTTP client: {}", e);
            return json_error_response(
                StatusCode::INTERNAL_SERVER_ERROR,
                "ServerError",
                "Failed to verify captcha",
            )
            .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response());
        }
    };

    #[derive(Deserialize, Serialize)]
    struct HCaptchaResponse {
        success: bool,
        challenge_ts: DateTime<Utc>,
        hostname: String,
        #[serde(rename = "error-codes", default)]
        error_codes: Vec<String>,
    }

    let verification_result = client
        .post("https://api.hcaptcha.com/siteverify")
        .form(&[
            ("secret", hcaptcha_secret.as_str()),
            ("response", form.h_captcha_response.as_str()),
        ])
        .send()
        .await;

    let verification_response = match verification_result {
        Ok(resp) => resp,
        Err(e) => {
            log::error!("Failed to verify hCaptcha: {}", e);

            return Redirect::to(&format!(
                "/gate?handle={}&state={}&error={}",
                url_encode(&params.handle),
                url_encode(&params.state),
                url_encode("Verification failed. Please try again.")
            ))
            .into_response();
        }
    };

    let captcha_result: HCaptchaResponse = match verification_response.json().await {
        Ok(result) => result,
        Err(e) => {
            log::error!("Failed to parse hCaptcha response: {}", e);

            return Redirect::to(&format!(
                "/gate?handle={}&state={}&error={}",
                url_encode(&params.handle),
                url_encode(&params.state),
                url_encode("Verification failed. Please try again.")
            ))
            .into_response();
        }
    };

    if !captcha_result.success {
        log::warn!(
            "hCaptcha verification failed for handle {}: {:?}",
            params.handle,
            captcha_result.error_codes
        );
        return Redirect::to(&format!(
            "/gate?handle={}&state={}&error={}",
            url_encode(&params.handle),
            url_encode(&params.state),
            url_encode("Verification failed. Please try again.")
        ))
        .into_response();
    }

    // Generate secure JWE verification token
    let code = match generate_gate_token(&params.handle, &state.app_config.gate_jwe_key) {
        Ok(token) => token,
        Err(e) => {
            log::error!("Failed to generate gate token: {}", e);
            return json_error_response(
                StatusCode::INTERNAL_SERVER_ERROR,
                "ServerError",
                "Failed to create verification code",
            )
            .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response());
        }
    };

    let now = Utc::now();

    // Store the encrypted token in the database
    let result = sqlx::query(
        "INSERT INTO gate_codes (code, handle, created_at)
         VALUES (?, ?, ?)",
    )
    .bind(&code)
    .bind(&params.handle)
    .bind(now)
    .execute(&state.pds_gatekeeper_pool)
    .await;

    if let Err(e) = result {
        log::error!("Failed to store gate code: {}", e);
        return json_error_response(
            StatusCode::INTERNAL_SERVER_ERROR,
            "ServerError",
            "Failed to create verification code",
        )
        .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response());
    }

    // Redirects by origin if it's found. If not redirect to the configured URL.
    let mut base_redirect = state.app_config.default_successful_redirect_url.clone();
    if let Some(ref redirect_url) = form.redirect_url {
        let trimmed = redirect_url.trim();
        if !trimmed.is_empty()
            && (trimmed.starts_with("https://") || trimmed.starts_with("http://"))
        {
            base_redirect = trimmed.trim_end_matches('/').to_string();
        }
    }

    let base_redirect = match state
        .app_config
        .captcha_success_redirects
        .contains(&base_redirect)
    {
        true => base_redirect,
        false => state.app_config.default_successful_redirect_url.clone(),
    };

    // Redirect to client app with code and state
    let redirect_url = format!(
        "{}/?code={}&state={}",
        base_redirect,
        url_encode(&code),
        url_encode(&params.state)
    );

    Redirect::to(&redirect_url).into_response()
}

/// Simple URL encode function
fn url_encode(s: &str) -> String {
    urlencoding::encode(s).to_string()
}