sqlpage--sqlpage/tests/oidc/mod.rs

use actix_web::{
    App, HttpResponse, HttpServer, Responder,
    cookie::Cookie,
    http::{StatusCode, header},
    test,
    web::{self, Data},
};
use base64::Engine;
use openidconnect::url::Url;
use serde::{Deserialize, Serialize};
use serde_json::json;
use sqlpage::webserver::http::create_app;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio_util::sync::{CancellationToken, DropGuard};

fn base64url_encode(data: &[u8]) -> String {
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
}

pub fn make_jwt(claims: &serde_json::Value, secret: &str) -> String {
    use hmac::{Hmac, KeyInit, Mac};
    use sha2::Sha256;

    let header = json!({
        "alg": "HS256",
        "typ": "JWT",
        "kid": "test"
    });

    let header_b64 = base64url_encode(header.to_string().as_bytes());
    let payload_b64 = base64url_encode(claims.to_string().as_bytes());

    let message = format!("{}.{}", header_b64, payload_b64);

    let mut mac =
        Hmac::<Sha256>::new_from_slice(secret.as_bytes()).expect("HMAC accepts any key size");
    mac.update(message.as_bytes());
    let signature = mac.finalize().into_bytes();
    let signature_b64 = base64url_encode(&signature);

    format!("{}.{}.{}", header_b64, payload_b64, signature_b64)
}

type JwtCustomizer<'a> = dyn Fn(serde_json::Value, &str) -> String + Send + Sync + 'a;

struct ProviderState<'a> {
    secret: String,
    issuer_url: String,
    client_id: String,
    auth_codes: HashMap<String, String>, // code -> nonce
    jwt_customizer: Option<Box<JwtCustomizer<'a>>>,
    token_endpoint_delay: Duration,
    discovery_count: usize,
}

type ProviderStateWithLifetime<'a> = ProviderState<'a>;
type SharedProviderState = Arc<Mutex<ProviderStateWithLifetime<'static>>>;

#[derive(Serialize, Deserialize)]
struct DiscoveryResponse {
    issuer: String,
    authorization_endpoint: String,
    token_endpoint: String,
    jwks_uri: String,
    response_types_supported: Vec<String>,
    subject_types_supported: Vec<String>,
    id_token_signing_alg_values_supported: Vec<String>,
    end_session_endpoint: String,
}

#[derive(Serialize)]
struct JwksResponse {
    keys: Vec<serde_json::Value>,
}

#[derive(Serialize)]
struct TokenResponse {
    access_token: String,
    token_type: String,
    id_token: String,
    expires_in: i64,
}

async fn discovery_endpoint(state: Data<SharedProviderState>) -> impl Responder {
    let mut state = state.lock().unwrap();
    state.discovery_count += 1;
    let discovery = DiscoveryResponse {
        issuer: state.issuer_url.clone(),
        authorization_endpoint: format!("{}/auth", state.issuer_url),
        token_endpoint: format!("{}/token", state.issuer_url),
        jwks_uri: format!("{}/jwks", state.issuer_url),
        response_types_supported: vec!["code".to_string()],
        subject_types_supported: vec!["public".to_string()],
        id_token_signing_alg_values_supported: vec!["HS256".to_string()],
        end_session_endpoint: format!("{}/logout", state.issuer_url),
    };
    drop(state);
    HttpResponse::Ok()
        .insert_header((header::CONTENT_TYPE, "application/json"))
        .json(discovery)
}

async fn jwks_endpoint(state: Data<SharedProviderState>) -> impl Responder {
    let state = state.lock().unwrap();
    let jwks = JwksResponse {
        keys: vec![json!({
            "kty": "oct",
            "kid": "test",
            "use": "sig",
            "alg": "HS256",
            "k": base64url_encode(state.secret.as_bytes())
        })],
    };
    HttpResponse::Ok()
        .insert_header((header::CONTENT_TYPE, "application/json"))
        .json(jwks)
}

async fn token_endpoint(
    state: Data<SharedProviderState>,
    req: web::Form<HashMap<String, String>>,
) -> impl Responder {
    let mut state = state.lock().unwrap();
    let Some(code) = req.get("code") else {
        return HttpResponse::BadRequest().body("Missing code");
    };
    let nonce = state.auth_codes.get(code).cloned().unwrap_or_default();
    if nonce.is_empty() {
        return HttpResponse::BadRequest().body("Unknown code");
    }

    let now = chrono::Utc::now().timestamp();
    let claims = json!({
        "iss": state.issuer_url.as_str(),
        "sub": "test_user",
        "aud": state.client_id.as_str(),
        "exp": now + 3600,
        "iat": now,
        "nonce": nonce,
    });

    let id_token = state
        .jwt_customizer
        .take()
        .map(|customizer| customizer(claims.clone(), &state.secret))
        .unwrap_or_else(|| make_jwt(&claims, &state.secret));

    let delay = state.token_endpoint_delay;
    drop(state);

    let response = TokenResponse {
        access_token: "test_access_token".to_string(),
        token_type: "Bearer".to_string(),
        id_token,
        expires_in: 3600,
    };

    let json_bytes = serde_json::to_vec(&response).unwrap();
    let body = futures_util::stream::once(async move {
        tokio::time::sleep(delay).await;
        Ok::<web::Bytes, actix_web::Error>(web::Bytes::from(json_bytes))
    });
    HttpResponse::Ok()
        .insert_header((header::CONTENT_TYPE, "application/json"))
        .streaming(body)
}

pub struct FakeOidcProvider {
    pub issuer_url: String,
    pub client_id: String,
    pub client_secret: String,
    state: SharedProviderState,
    _stop_on_drop: DropGuard,
}

fn extract_set_cookies(headers: &header::HeaderMap) -> Vec<Cookie<'static>> {
    headers
        .get_all(header::SET_COOKIE)
        .filter_map(|h| h.to_str().ok())
        .filter_map(|s| Cookie::parse(s).ok())
        .map(Cookie::into_owned)
        .collect()
}

impl FakeOidcProvider {
    pub async fn new() -> Self {
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let port = listener.local_addr().unwrap().port();
        let issuer_url = format!("http://127.0.0.1:{}", port);
        let client_id = "test_client".to_string();
        let client_secret = "test_secret".to_string();

        let state: SharedProviderState = Arc::new(Mutex::new(ProviderState {
            secret: client_secret.clone(),
            issuer_url: issuer_url.clone(),
            client_id: client_id.clone(),
            auth_codes: HashMap::new(),
            jwt_customizer: None,
            token_endpoint_delay: Duration::ZERO,
            discovery_count: 0,
        }));

        let state_for_server = Arc::clone(&state);

        let server_stop = CancellationToken::new();
        let stop_on_drop = server_stop.clone().drop_guard();

        let server = HttpServer::new(move || {
            let state = Data::new(Arc::clone(&state_for_server));
            App::new()
                .app_data(state.clone())
                .route(
                    "/.well-known/openid-configuration",
                    web::get().to(discovery_endpoint),
                )
                .route("/jwks", web::get().to(jwks_endpoint))
                .route("/token", web::post().to(token_endpoint))
        })
        .workers(1)
        .listen(listener)
        .unwrap()
        .shutdown_timeout(1)
        .shutdown_signal(server_stop.cancelled_owned())
        .run();

        tokio::spawn(server);

        Self {
            issuer_url,
            client_id,
            client_secret,
            state,
            _stop_on_drop: stop_on_drop,
        }
    }

    fn with_state_mut<R>(&self, f: impl FnOnce(&mut ProviderState) -> R) -> R {
        let mut state = self.state.lock().unwrap();
        f(&mut state)
    }

    pub fn set_token_endpoint_delay(&self, delay: Duration) {
        self.with_state_mut(|s| s.token_endpoint_delay = delay);
    }

    pub fn discovery_count(&self) -> usize {
        self.state.lock().unwrap().discovery_count
    }

    pub fn store_auth_code(&self, code: String, nonce: String) {
        self.with_state_mut(|s| {
            s.auth_codes.insert(code, nonce);
        });
    }
}

fn get_query_param(url: &Url, name: &str) -> String {
    url.query_pairs()
        .find(|(k, _)| k == name)
        .unwrap()
        .1
        .to_string()
}

macro_rules! request_with_cookies {
    ($app:expr, $req:expr, $cookies:expr) => {{
        let mut req = $req;
        for cookie in $cookies.iter() {
            req = req.cookie(cookie.clone());
        }
        let resp = test::call_service(&$app, req.to_request()).await;
        for new_cookie in extract_set_cookies(resp.headers()) {
            $cookies.retain(|c: &Cookie| c.name() != new_cookie.name());
            if !new_cookie.value().is_empty() {
                $cookies.push(new_cookie);
            }
        }
        resp
    }};
}

async fn setup_oidc_test(
    provider_mutator: impl FnOnce(&mut ProviderState),
) -> (
    impl actix_web::dev::Service<
        actix_http::Request,
        Response = actix_web::dev::ServiceResponse<impl actix_web::body::MessageBody>,
        Error = actix_web::Error,
    >,
    FakeOidcProvider,
) {
    use sqlpage::{
        AppState,
        app_config::{AppConfig, test_database_url},
    };
    crate::common::init_log();
    let provider = FakeOidcProvider::new().await;
    provider.with_state_mut(provider_mutator);

    let db_url = test_database_url();
    let config_json = format!(
        r#"{{
        "database_url": "{db_url}",
        "max_database_pool_connections": 1,
        "database_connection_retries": 3,
        "database_connection_acquire_timeout_seconds": 15,
        "allow_exec": true,
        "max_uploaded_file_size": 123456,
        "listen_on": "127.0.0.1:0",
        "system_root_ca_certificates": false,
        "oidc_issuer_url": "{}",
        "oidc_client_id": "{}",
        "oidc_client_secret": "{}",
        "oidc_protected_paths": ["/"],
        "host": "localhost:1"
    }}"#,
        provider.issuer_url, provider.client_id, provider.client_secret
    );

    let config: AppConfig = serde_json::from_str(&config_json).unwrap();
    let app_state = AppState::init(&config).await.unwrap();
    let app = test::init_service(create_app(Data::new(app_state))).await;
    (app, provider)
}

#[actix_web::test]
async fn test_oidc_happy_path() {
    let (app, provider) = setup_oidc_test(|_| {}).await;
    let mut cookies: Vec<Cookie<'static>> = Vec::new();

    let resp = request_with_cookies!(app, test::TestRequest::get().uri("/"), cookies);
    assert_eq!(resp.status(), StatusCode::SEE_OTHER);
    let auth_url = Url::parse(resp.headers().get("location").unwrap().to_str().unwrap()).unwrap();

    let state = get_query_param(&auth_url, "state");
    let nonce = get_query_param(&auth_url, "nonce");
    let redirect_uri = get_query_param(&auth_url, "redirect_uri");
    provider.store_auth_code("test_auth_code".to_string(), nonce);

    let callback_uri = format!(
        "{}?code=test_auth_code&state={}",
        Url::parse(&redirect_uri).unwrap().path(),
        state
    );
    let callback_resp =
        request_with_cookies!(app, test::TestRequest::get().uri(&callback_uri), cookies);
    assert_eq!(callback_resp.status(), StatusCode::SEE_OTHER);

    let final_resp = request_with_cookies!(app, test::TestRequest::get().uri("/"), cookies);
    assert_eq!(final_resp.status(), StatusCode::OK);
}

async fn assert_oidc_login_fails(
    provider_mutator: impl FnOnce(&mut ProviderState),
    state_override: Option<String>,
) {
    let (app, provider) = setup_oidc_test(provider_mutator).await;
    let mut cookies: Vec<Cookie<'static>> = Vec::new();

    let resp = request_with_cookies!(app, test::TestRequest::get().uri("/"), cookies);
    assert_eq!(resp.status(), StatusCode::SEE_OTHER);
    let auth_url = Url::parse(resp.headers().get("location").unwrap().to_str().unwrap()).unwrap();

    let state = get_query_param(&auth_url, "state");
    let nonce = get_query_param(&auth_url, "nonce");
    let redirect_uri = get_query_param(&auth_url, "redirect_uri");
    provider.store_auth_code("test_auth_code".to_string(), nonce);

    let callback_state = state_override.unwrap_or(state);
    let callback_uri = format!(
        "{}?code=test_auth_code&state={}",
        Url::parse(&redirect_uri).unwrap().path(),
        callback_state
    );
    let callback_resp =
        request_with_cookies!(app, test::TestRequest::get().uri(&callback_uri), cookies);

    assert_eq!(callback_resp.status(), StatusCode::SEE_OTHER);
    let location = callback_resp
        .headers()
        .get("location")
        .unwrap()
        .to_str()
        .unwrap();
    assert!(
        location.starts_with(&provider.issuer_url),
        "Expected redirect to OIDC provider, but got {location}"
    );

    let auth_cookie_present = cookies.iter().any(|c| c.name() == "sqlpage_auth");
    assert!(
        !auth_cookie_present,
        "Authentication cookie should not be set on failure"
    );
}

async fn assert_oidc_callback_fails_with_bad_jwt(
    mutate_jwt_claims: impl FnMut(&mut serde_json::Value) + Send + Sync + 'static,
) {
    let mutate_jwt_claims = Mutex::new(mutate_jwt_claims);
    assert_oidc_login_fails(
        |state| {
            state.jwt_customizer = Some(Box::new(move |mut claims, secret| {
                mutate_jwt_claims.lock().unwrap()(&mut claims);
                make_jwt(&claims, secret)
            }));
        },
        None,
    )
    .await;
}

#[actix_web::test]
async fn test_oidc_csrf_state_mismatch_is_rejected() {
    assert_oidc_login_fails(|_| {}, Some("wrong_state".to_string())).await;
}

#[actix_web::test]
async fn test_oidc_nonce_mismatch_is_rejected() {
    assert_oidc_callback_fails_with_bad_jwt(|claims| {
        claims["nonce"] = json!("wrong_nonce");
    })
    .await;
}

#[actix_web::test]
async fn test_oidc_bad_signature_is_rejected() {
    assert_oidc_login_fails(
        |state| {
            state.jwt_customizer = Some(Box::new(|claims, _| make_jwt(&claims, "wrong_secret")));
        },
        None,
    )
    .await;
}

#[actix_web::test]
async fn test_oidc_wrong_audience_is_rejected() {
    assert_oidc_callback_fails_with_bad_jwt(|claims| {
        claims["aud"] = json!("wrong_client");
    })
    .await;
}

#[actix_web::test]
async fn test_oidc_wrong_issuer_is_rejected() {
    assert_oidc_callback_fails_with_bad_jwt(|claims| {
        claims["iss"] = json!("https://wrong-issuer.com");
    })
    .await;
}

#[actix_web::test]
async fn test_oidc_expired_token_is_rejected() {
    assert_oidc_callback_fails_with_bad_jwt(|claims| {
        let current_exp = claims["exp"].as_i64().unwrap();
        claims["exp"] = json!(current_exp - 7200);
    })
    .await;
}

async fn setup_oidc_test_with_prefix(
    provider_mutator: impl FnOnce(&mut ProviderState),
    site_prefix: &str,
) -> (
    impl actix_web::dev::Service<
        actix_http::Request,
        Response = actix_web::dev::ServiceResponse<impl actix_web::body::MessageBody>,
        Error = actix_web::Error,
    >,
    FakeOidcProvider,
) {
    use sqlpage::{
        AppState,
        app_config::{AppConfig, test_database_url},
    };
    crate::common::init_log();
    let provider = FakeOidcProvider::new().await;
    provider.with_state_mut(provider_mutator);

    let db_url = test_database_url();
    let config_json = format!(
        r#"{{
        "database_url": "{db_url}",
        "oidc_issuer_url": "{}",
        "oidc_client_id": "{}",
        "oidc_client_secret": "{}",
        "oidc_protected_paths": ["/"],
        "site_prefix": "{site_prefix}"
    }}"#,
        provider.issuer_url, provider.client_id, provider.client_secret
    );

    let config: AppConfig = serde_json::from_str(&config_json).unwrap();
    let app_state = AppState::init(&config).await.unwrap();
    let app = test::init_service(create_app(Data::new(app_state))).await;
    (app, provider)
}

#[actix_web::test]
async fn test_oidc_with_site_prefix() {
    let (app, _provider) = setup_oidc_test_with_prefix(|_| {}, "/my-app/").await;
    let mut cookies: Vec<Cookie<'static>> = Vec::new();

    // Access the app with the prefix
    let resp = request_with_cookies!(app, test::TestRequest::get().uri("/my-app/"), cookies);
    assert_eq!(resp.status(), StatusCode::SEE_OTHER);
    let auth_url = Url::parse(resp.headers().get("location").unwrap().to_str().unwrap()).unwrap();

    // Check if the redirect_uri parameter in the auth URL contains the site prefix
    let redirect_uri = get_query_param(&auth_url, "redirect_uri");
    assert!(
        redirect_uri.contains("/my-app/sqlpage/oidc_callback"),
        "Redirect URI should contain site prefix. Got: {}",
        redirect_uri
    );
}

#[actix_web::test]
async fn test_oidc_logout_uses_correct_scheme() {
    use sqlpage::{
        AppState,
        app_config::{AppConfig, test_database_url},
    };

    crate::common::init_log();
    let provider = FakeOidcProvider::new().await;

    let db_url = test_database_url();
    let config_json = format!(
        r#"{{
        "database_url": "{db_url}",
        "oidc_issuer_url": "{}",
        "oidc_client_id": "{}",
        "oidc_client_secret": "{}",
        "https_domain": "example.com"
    }}"#,
        provider.issuer_url, provider.client_id, provider.client_secret
    );

    let config: AppConfig = serde_json::from_str(&config_json).unwrap();
    let app_state = AppState::init(&config).await.unwrap();
    let logout_path = app_state
        .oidc_state
        .as_ref()
        .unwrap()
        .config
        .create_logout_url("/logged_out");
    let app = test::init_service(create_app(Data::new(app_state))).await;
    // make sure the logout path includes the configured domain
    assert!(logout_path.starts_with("/sqlpage/oidc_logout"));

    let req = test::TestRequest::get().uri(&logout_path).to_request();
    let resp = test::call_service(&app, req).await;

    assert_eq!(resp.status(), StatusCode::SEE_OTHER);
    let location = resp.headers().get("location").unwrap().to_str().unwrap();
    let location_url = Url::parse(location).unwrap();
    assert_eq!(location_url.path(), "/logout");
    let params: HashMap<String, String> = location_url.query_pairs().into_owned().collect();
    let post_logout = params.get("post_logout_redirect_uri").unwrap();
    assert_eq!(post_logout, "https://example.com/logged_out");
}

/// An OIDC provider metadata refresh must not block authenticated requests.
/// The refresh should happen in the background while existing requests are
/// served using the current (possibly stale) OIDC client.
#[actix_web::test]
async fn test_slow_discovery_does_not_block_authenticated_requests() {
    let (app, provider) = setup_oidc_test(|_| {}).await;
    let mut cookies: Vec<Cookie<'static>> = Vec::new();

    // Complete a full login to get auth cookies
    let resp = request_with_cookies!(app, test::TestRequest::get().uri("/"), cookies);
    assert_eq!(resp.status(), StatusCode::SEE_OTHER);
    let auth_url = Url::parse(resp.headers().get("location").unwrap().to_str().unwrap()).unwrap();
    let state_param = get_query_param(&auth_url, "state");
    let nonce = get_query_param(&auth_url, "nonce");
    let redirect_uri = get_query_param(&auth_url, "redirect_uri");
    provider.store_auth_code("test_auth_code".to_string(), nonce);
    let callback_uri = format!(
        "{}?code=test_auth_code&state={}",
        Url::parse(&redirect_uri).unwrap().path(),
        state_param
    );
    let callback_resp =
        request_with_cookies!(app, test::TestRequest::get().uri(&callback_uri), cookies);
    assert_eq!(callback_resp.status(), StatusCode::SEE_OTHER);

    // Advance time so the OIDC snapshot appears stale.
    // The next request triggers a background refresh.
    let count_before = provider.discovery_count();
    tokio::time::pause();
    tokio::time::advance(Duration::from_secs(3601)).await;
    // Resume real time so the DB pool and background refresh work normally.
    tokio::time::resume();

    // An authenticated request must succeed immediately, even though
    // it triggers a background refresh.
    let resp = request_with_cookies!(app, test::TestRequest::get().uri("/"), cookies);
    assert_eq!(resp.status(), StatusCode::OK);

    // Let the background refresh task complete.
    tokio::task::yield_now().await;
    assert!(
        provider.discovery_count() > count_before,
        "OIDC provider metadata was not refreshed"
    );
}

/// A slow OIDC token endpoint must not freeze the server.
/// The body-read timeout fires and the request completes with a redirect.
#[actix_web::test]
async fn test_slow_token_endpoint_does_not_freeze_server() {
    let (app, provider) = setup_oidc_test(|_| {}).await;
    let mut cookies: Vec<Cookie<'static>> = Vec::new();

    let resp = request_with_cookies!(app, test::TestRequest::get().uri("/"), cookies);
    assert_eq!(resp.status(), StatusCode::SEE_OTHER);
    let auth_url = Url::parse(resp.headers().get("location").unwrap().to_str().unwrap()).unwrap();
    let state_param = get_query_param(&auth_url, "state");
    let nonce = get_query_param(&auth_url, "nonce");
    let redirect_uri = get_query_param(&auth_url, "redirect_uri");
    provider.store_auth_code("test_auth_code".to_string(), nonce);

    provider.set_token_endpoint_delay(Duration::from_secs(999));

    let callback_uri = format!(
        "{}?code=test_auth_code&state={}",
        Url::parse(&redirect_uri).unwrap().path(),
        state_param
    );

    let handle = tokio::task::spawn_local(async move {
        let mut req = test::TestRequest::get().uri(&callback_uri);
        for cookie in cookies.iter() {
            req = req.cookie(cookie.clone());
        }
        test::call_service(&app, req.to_request()).await
    });

    // Let the TCP round-trip complete so awc reads HTTP headers,
    // then advance past the body-read timeout.
    tokio::task::yield_now().await;
    tokio::time::pause();
    tokio::time::advance(Duration::from_secs(60)).await;

    let resp = tokio::time::timeout(Duration::from_secs(1), handle)
        .await
        .expect("OIDC callback hung on a slow token endpoint")
        .unwrap();
    assert_eq!(resp.status(), StatusCode::SEE_OTHER);
}