crank/apps/admin-api/tests/integration/auth_rate_limit.rs

#![allow(dead_code, unused_imports)]

use super::common::*;

use std::{
    collections::BTreeMap,
    env, fmt,
    sync::Arc,
    time::{SystemTime, UNIX_EPOCH},
};

use async_trait::async_trait;
use axum::{Json, Router, routing::post};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use crank_core::{
    ExecutionConfig, HttpMethod, MembershipRole, OperationSecurityLevel, Protocol,
    ResponseCachePolicy, RestTarget, SecretKind, Target, ToolDescription, WorkspaceId,
};
use crank_core::{IdentityError, IdentityProvider, IdentityProviderKind, LoginOutcome};
use crank_mapping::{MappingRule, MappingSet};
use crank_registry::{
    CreateAdminBootstrapContractRequest, MASTER_KEY_CIPHER_CONTRACT, MasterKeyIdentityCandidate,
    MigrationAuthority, PostgresRegistry,
};
use crank_runtime::SecretCrypto;
use crank_schema::{Schema, SchemaKind};
use serde_json::{Value, json};
use serial_test::serial;
use sha2::{Digest, Sha256};
use tokio::net::TcpListener;

use admin_api::{
    app::build_app,
    auth::{AuthSettings, BootstrapAdminConfig, hash_password},
    service::{AdminService, AdminServiceBuilder, OperationPayload},
    state::AppState,
};

const DEFAULT_WORKSPACE_ID: &str = "ws_default";
const TEST_AUTH_EMAIL: &str = "owner@crank.local";
const TEST_AUTH_PASSWORD: &str = "test-password";
const TEST_PASSWORD_PEPPER: &str = "test-password-pepper";
const TEST_SESSION_SECRET: &str = "test-session-secret";
const TEST_MASTER_KEY: &str = "test-master-key-00000000000000000000000000000000";

struct TestServer {
    base_url: String,
    shutdown: Option<tokio::sync::oneshot::Sender<()>>,
    handle: Option<tokio::task::JoinHandle<()>>,
}

impl Drop for TestServer {
    fn drop(&mut self) {
        let shutdown = self.shutdown.take();
        let handle = self.handle.take();

        tokio::task::block_in_place(|| {
            if let Some(shutdown) = shutdown {
                let _ = shutdown.send(());
            }
            if let Some(handle) = handle {
                let _ = tokio::runtime::Handle::current().block_on(handle);
            }
        });
    }
}

impl fmt::Display for TestServer {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.base_url)
    }
}

impl AsRef<str> for TestServer {
    fn as_ref(&self) -> &str {
        &self.base_url
    }
}

struct RejectingIdentityProvider;

#[async_trait]
impl IdentityProvider for RejectingIdentityProvider {
    fn id(&self) -> &str {
        "rejecting-test-provider"
    }

    fn kind(&self) -> IdentityProviderKind {
        IdentityProviderKind::Password
    }

    async fn login_password(
        &self,
        _payload: crank_core::LoginPayload,
    ) -> Result<LoginOutcome, IdentityError> {
        Err(IdentityError::BadCredentials)
    }
}

async fn empty_registry(name: &str) -> PostgresRegistry {
    let database_url = crank_test_support::postgres_schema_url(name).await;
    let pool = sqlx::PgPool::connect(&database_url).await.unwrap();
    MigrationAuthority::apply(&pool).await.unwrap();
    let registry = PostgresRegistry::connect(&database_url).await.unwrap();
    let secret_crypto = test_secret_crypto();
    registry
        .verify_or_register_master_key_identity(MasterKeyIdentityCandidate {
            epoch: secret_crypto.master_key_epoch(),
            fingerprint: secret_crypto.master_key_fingerprint(),
            cipher_contract: MASTER_KEY_CIPHER_CONTRACT,
            observed_at: &time::OffsetDateTime::now_utc(),
        })
        .await
        .unwrap();
    registry
}

fn bootstrap_token_hash(token: &str) -> String {
    let digest = Sha256::digest(format!("bootstrap:{token}").as_bytes());
    URL_SAFE_NO_PAD.encode(digest)
}

async fn login_client_without_csrf(root_url: &str) -> reqwest::Client {
    let client = reqwest::Client::builder()
        .cookie_store(true)
        .build()
        .unwrap();
    client
        .post(format!("{root_url}/api/auth/login"))
        .json(&json!({
            "email": TEST_AUTH_EMAIL,
            "password": TEST_AUTH_PASSWORD,
        }))
        .send()
        .await
        .unwrap()
        .error_for_status()
        .unwrap();
    client
}

#[tokio::test(flavor = "multi_thread")]
#[serial]
async fn rejects_rapid_login_requests_with_429() {
    let registry = test_registry().await;
    let storage_root = test_storage_root("login_rate_limit");
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();
    let app = build_app(AppState {
        service: test_service(
            registry,
            storage_root,
            test_auth_settings(),
            test_secret_crypto(),
        ),
        api_rate_limiter: crank_runtime::RequestRateLimiter::new(
            crank_runtime::RequestRateLimitConfig::new(1, 1).unwrap(),
        ),
        trusted_proxy_ips: Vec::new(),
    });
    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
    let handle = tokio::spawn(async move {
        axum::serve(
            listener,
            app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
        )
        .with_graceful_shutdown(async move {
            let _ = shutdown_rx.await;
        })
        .await
        .unwrap();
    });

    let client = reqwest::Client::new();
    let root_url = format!("http://{address}");
    let first_response = client
        .post(format!("{root_url}/api/auth/login"))
        .header("x-request-id", "req_admin_rate_01")
        .json(&json!({
            "email": TEST_AUTH_EMAIL,
            "password": TEST_AUTH_PASSWORD,
        }))
        .send()
        .await
        .unwrap();
    assert_eq!(first_response.status(), reqwest::StatusCode::OK);

    let second_response = client
        .post(format!("{root_url}/api/auth/login"))
        .header("x-request-id", "req_admin_rate_02")
        .json(&json!({
            "email": TEST_AUTH_EMAIL,
            "password": TEST_AUTH_PASSWORD,
        }))
        .send()
        .await
        .unwrap();

    assert_eq!(
        second_response.status(),
        reqwest::StatusCode::TOO_MANY_REQUESTS
    );
    assert_eq!(
        second_response.headers()["x-request-id"].to_str().unwrap(),
        "req_admin_rate_02"
    );
    let payload = second_response.json::<Value>().await.unwrap();
    assert_eq!(payload["error"]["code"], "rate_limited");
    let retry_after_ms = payload["error"]["context"]["retry_after_ms"]
        .as_u64()
        .unwrap();
    assert!((1..=1000).contains(&retry_after_ms));

    let _ = shutdown_tx.send(());
    let _ = handle.await;
}

#[tokio::test(flavor = "multi_thread")]
#[serial]
async fn login_uses_identity_provider_when_configured() {
    let registry = test_registry().await;
    let storage_root = test_storage_root("identity_provider_login");
    let service = AdminServiceBuilder::new(
        registry,
        storage_root,
        test_auth_settings(),
        test_secret_crypto(),
        crank_runtime::community_default().build(),
    )
    .with_identity_provider(Arc::new(RejectingIdentityProvider))
    .build();

    let error = match service
        .login(admin_api::service::LoginPayload {
            email: TEST_AUTH_EMAIL.to_owned(),
            password: TEST_AUTH_PASSWORD.to_owned(),
        })
        .await
    {
        Ok(_) => panic!("login should delegate to identity provider"),
        Err(error) => error,
    };
    assert_eq!(error.to_string(), "invalid email or password");
}

#[tokio::test(flavor = "multi_thread")]
#[serial]
async fn bootstrap_token_creates_first_admin_once_and_never_replays() {
    let registry = empty_registry("test_admin_api_bootstrap_token").await;
    let token = "bootstrap-token-story-1-12-local-operator-only";
    registry
        .create_admin_bootstrap_contract(CreateAdminBootstrapContractRequest {
            id: "boot_test_story_1_12",
            token_hash: &bootstrap_token_hash(token),
            email: "first-owner@crank.local",
            display_name: "First Owner",
            expires_at: &(time::OffsetDateTime::now_utc() + time::Duration::minutes(15)),
        })
        .await
        .unwrap();

    let storage_root = test_storage_root("bootstrap_token");
    let base_url = spawn_admin_api(build_test_app(registry, storage_root)).await;
    let root_url = base_url
        .as_ref()
        .split("/api/admin/workspaces/")
        .next()
        .unwrap()
        .to_owned();
    let client = reqwest::Client::builder()
        .cookie_store(true)
        .build()
        .unwrap();

    let status = assert_success_json(
        client
            .get(format!("{root_url}/api/auth/bootstrap/status"))
            .send()
            .await
            .unwrap(),
    )
    .await;
    assert_eq!(status["bootstrap_required"], true);

    let completed = assert_success_json(
        client
            .post(format!("{root_url}/api/auth/bootstrap/complete"))
            .json(&json!({
                "token": token,
                "password": "first-owner-password-123"
            }))
            .send()
            .await
            .unwrap(),
    )
    .await;
    assert_eq!(completed["user"]["email"], "first-owner@crank.local");
    assert!(completed["csrf_token"].as_str().unwrap().len() >= 32);

    let replay = client
        .post(format!("{root_url}/api/auth/bootstrap/complete"))
        .json(&json!({
            "token": token,
            "password": "second-password-should-not-work"
        }))
        .send()
        .await
        .unwrap();
    assert_eq!(replay.status(), reqwest::StatusCode::UNAUTHORIZED);

    let status_after = assert_success_json(
        client
            .get(format!("{root_url}/api/auth/bootstrap/status"))
            .send()
            .await
            .unwrap(),
    )
    .await;
    assert_eq!(status_after["bootstrap_required"], false);
}

#[tokio::test(flavor = "multi_thread")]
#[serial]
async fn browser_mutations_require_csrf_and_reject_cross_origin_requests() {
    let registry = test_registry().await;
    let storage_root = test_storage_root("csrf_and_origin");
    let base_url = spawn_admin_api(build_test_app(registry, storage_root)).await;
    let root_url = base_url
        .as_ref()
        .split("/api/admin/workspaces/")
        .next()
        .unwrap()
        .to_owned();
    let client = login_client_without_csrf(&root_url).await;

    let missing_csrf = client
        .patch(format!("{root_url}/api/auth/profile"))
        .json(&json!({
            "display_name": "Blocked Without CSRF",
            "email": TEST_AUTH_EMAIL
        }))
        .send()
        .await
        .unwrap();
    assert_eq!(missing_csrf.status(), reqwest::StatusCode::FORBIDDEN);
    let missing_body = missing_csrf.text().await.unwrap();
    assert!(!missing_body.contains(TEST_SESSION_SECRET));
    assert!(!missing_body.contains(TEST_AUTH_PASSWORD));

    let cross_origin = client
        .get(format!("{root_url}/api/auth/profile"))
        .header("origin", "https://attacker.example")
        .send()
        .await
        .unwrap();
    assert_eq!(cross_origin.status(), reqwest::StatusCode::FORBIDDEN);
    let cross_origin_body = cross_origin.text().await.unwrap();
    assert!(!cross_origin_body.contains(TEST_SESSION_SECRET));
    assert!(!cross_origin_body.contains(TEST_AUTH_PASSWORD));
}