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 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::PostgresRegistry;
use crank_runtime::SecretCrypto;
use crank_schema::{Schema, SchemaKind};
use serde_json::{Value, json};
use serial_test::serial;
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";

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)
    }
}

#[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(),
        ),
        trust_forwarded_headers: false,
    });
    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");
}