tests.rs

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

use async_trait::async_trait;
use axum::body::Body;
use axum::extract::ConnectInfo;
use axum::http::{HeaderMap, HeaderValue, Method, Request, StatusCode, Uri, header};
use axum::response::{IntoResponse as _, Response};
use axum::{Router, middleware, routing::get};
use peryx_core::Ecosystem;
use peryx_identity::IndexAcl;
use peryx_index::{Index, IndexKind};
use peryx_policy::Policy;
use peryx_search::default_indexer;
use peryx_storage::blob::BlobStore;
use peryx_storage::meta::MetaStore;
use tower::ServiceExt as _;

use super::{
    ActorKey, ForwardedClient, RateLimitConfig, RateLimiter, RouteClass, RouteLimit, UpstreamLimited, UpstreamLimits,
    UpstreamPermit, ecosystem_route_class, limited_response, malformed_forwarded_response, real_ip,
};
use crate::serving::{
    AbsoluteProtocolDriver, EcosystemDriver, IndexedProtocolDriver, ProtocolDriver, RateLimitPrincipal, ServiceDriver,
};
use crate::state::{AppState, ServingState};
use crate::{RouteDescriptor, RouteMethod, RoutePosture, RouteRateLimit, client_address};

struct IndexedDriver;

impl EcosystemDriver for IndexedDriver {
    fn ecosystem(&self) -> Ecosystem {
        Ecosystem::new("example")
    }
}

impl RateLimitPrincipal for IndexedDriver {
    fn resolve(
        &self,
        _state: &crate::ServingState,
        _position: Option<usize>,
        _headers: &HeaderMap,
    ) -> peryx_identity::Principal {
        peryx_identity::Principal::Named {
            subject: "reader".to_owned(),
        }
    }
}

#[async_trait]
impl IndexedProtocolDriver for IndexedDriver {
    fn classify_route(&self, _path: &str) -> RouteClass {
        RouteClass::Artifact
    }

    async fn get(
        &self,
        _state: Arc<ServingState>,
        _position: usize,
        _rest: String,
        _uri: Uri,
        _headers: HeaderMap,
        _method: Method,
    ) -> Response {
        StatusCode::NO_CONTENT.into_response()
    }

    async fn post(&self, _state: Arc<ServingState>, _path: String, _request: axum::extract::Request) -> Response {
        StatusCode::NO_CONTENT.into_response()
    }

    async fn put(&self, _state: Arc<ServingState>, _request: axum::extract::Request) -> Response {
        StatusCode::NO_CONTENT.into_response()
    }

    async fn delete(&self, _state: Arc<ServingState>, _request: axum::extract::Request) -> Response {
        StatusCode::NO_CONTENT.into_response()
    }
}

#[async_trait]
impl ServiceDriver for IndexedDriver {
    fn classify_service_post(&self, path: &str, _headers: &HeaderMap) -> Option<RouteClass> {
        (path == "+special").then_some(RouteClass::Admin)
    }

    async fn service_post(&self, _state: Arc<crate::ServingState>, _request: axum::extract::Request) -> Response {
        StatusCode::NO_CONTENT.into_response()
    }
}

struct AbsoluteDriver;

impl EcosystemDriver for AbsoluteDriver {
    fn ecosystem(&self) -> Ecosystem {
        Ecosystem::new("absolute")
    }
}

#[async_trait]
impl AbsoluteProtocolDriver for AbsoluteDriver {
    fn prefixes(&self) -> &'static [&'static str] {
        &["/artifacts"]
    }

    fn classify_route(&self, _path: &str) -> RouteClass {
        RouteClass::Artifact
    }

    async fn serve(&self, _state: Arc<ServingState>, _request: axum::extract::Request) -> Response {
        StatusCode::NO_CONTENT.into_response()
    }
}

#[tokio::test]
async fn test_protocol_fixtures_serve_supported_requests() {
    let (_dir, state) = app(RateLimitConfig::default());
    let serving = Arc::clone(&state.serving);
    let indexed = IndexedDriver;

    assert_eq!(
        indexed
            .get(
                Arc::clone(&serving),
                0,
                "resource".to_owned(),
                Uri::from_static("/items/resource"),
                HeaderMap::new(),
                Method::GET,
            )
            .await
            .status(),
        StatusCode::NO_CONTENT
    );
    assert_eq!(
        indexed
            .post(
                Arc::clone(&serving),
                "/items".to_owned(),
                Request::builder().body(Body::from("post body")).unwrap(),
            )
            .await
            .status(),
        StatusCode::NO_CONTENT
    );
    assert_eq!(
        indexed
            .put(
                Arc::clone(&serving),
                Request::builder()
                    .method(Method::PUT)
                    .uri("/items/resource")
                    .body(Body::from("put body"))
                    .unwrap(),
            )
            .await
            .status(),
        StatusCode::NO_CONTENT
    );
    assert_eq!(
        indexed
            .delete(Arc::clone(&serving), Request::new(Body::empty()))
            .await
            .status(),
        StatusCode::NO_CONTENT
    );
    assert_eq!(
        indexed
            .service_post(Arc::clone(&serving), Request::new(Body::empty()))
            .await
            .status(),
        StatusCode::NO_CONTENT
    );
    assert_eq!(
        AbsoluteDriver
            .serve(serving, Request::get("/artifacts/item").body(Body::empty()).unwrap())
            .await
            .status(),
        StatusCode::NO_CONTENT
    );
}

#[test]
fn test_check_client_allows_within_limit_then_denies_per_client() {
    let limiter = RateLimiter::new(RateLimitConfig {
        listing: RouteLimit::new(2, 60),
        ..RateLimitConfig::enabled_defaults()
    });
    let client = IpAddr::V4(Ipv4Addr::new(198, 51, 100, 7));
    assert!(limiter.check_client(RouteClass::Listing, client));
    assert!(limiter.check_client(RouteClass::Listing, client));
    assert!(!limiter.check_client(RouteClass::Listing, client));
    assert!(limiter.check_client(RouteClass::Listing, IpAddr::V4(Ipv4Addr::new(203, 0, 113, 9))));
}

#[test]
fn test_the_window_resets_and_readmits_once_time_advances_past_it() {
    let millis = Arc::new(AtomicU64::new(0));
    let handle = Arc::clone(&millis);
    let limiter = RateLimiter::with_clock(
        RateLimitConfig {
            listing: RouteLimit::new(1, 1),
            ..RateLimitConfig::enabled_defaults()
        },
        Arc::new(move || Duration::from_millis(handle.load(Ordering::SeqCst))),
    );
    let client = IpAddr::V4(Ipv4Addr::new(198, 51, 100, 7));
    assert!(
        limiter.check_client(RouteClass::Listing, client),
        "the first request in the window is admitted"
    );
    assert!(
        !limiter.check_client(RouteClass::Listing, client),
        "the second exhausts the one-per-window budget"
    );

    millis.store(1_001, Ordering::SeqCst);
    assert!(
        limiter.check_client(RouteClass::Listing, client),
        "a window whose reset time has passed readmits the client"
    );
}

#[test]
fn test_ecosystem_route_class_handles_writes_and_common_reads() {
    assert_eq!(
        ecosystem_route_class(&Method::POST, "/alpha/items"),
        Some(RouteClass::Upload)
    );
    assert_eq!(
        ecosystem_route_class(&Method::GET, "/alpha/hosted/+api"),
        Some(RouteClass::Admin)
    );
    assert_eq!(
        ecosystem_route_class(&Method::GET, "/alpha/resources/widget/details"),
        None
    );
}

#[test]
fn test_ecosystem_route_class_treats_head_and_options_as_reads() {
    assert_eq!(ecosystem_route_class(&Method::HEAD, "/service/resources/current"), None);
    assert_eq!(ecosystem_route_class(&Method::OPTIONS, "/alpha/items/current"), None);
    for method in [Method::PUT, Method::PATCH, Method::DELETE] {
        assert_eq!(
            ecosystem_route_class(&method, "/service/resources/1"),
            Some(RouteClass::Upload)
        );
    }
    assert_eq!(
        ecosystem_route_class(&Method::TRACE, "/alpha/items"),
        Some(RouteClass::Upload)
    );
}

#[test]
fn test_route_classes_expose_stable_names_and_limits() {
    let config = RateLimitConfig::enabled_defaults();
    let expected = [
        (RouteClass::Listing, "listing", config.listing),
        (RouteClass::Metadata, "metadata", config.metadata),
        (RouteClass::Artifact, "artifact", config.artifact),
        (RouteClass::Upload, "upload", config.upload),
        (RouteClass::Admin, "admin", config.admin),
        (RouteClass::Authentication, "authentication", config.authentication),
    ];

    assert_eq!(RouteClass::all(), expected.map(|(class, _, _)| class));
    for (class, name, limit) in expected {
        assert_eq!(class.as_str(), name);
        assert_eq!(config.limit(class), limit);
    }
}

#[test]
fn test_default_limiter_is_disabled_and_counts_each_class() {
    let limiter = RateLimiter::default();

    assert!(!limiter.enabled());
    for class in RouteClass::all() {
        assert!(limiter.check_client(class, IpAddr::V4(Ipv4Addr::LOCALHOST)));
    }
    assert_eq!(
        limiter
            .counters()
            .into_iter()
            .map(|snapshot| (snapshot.class, snapshot.allowed, snapshot.denied))
            .collect::<Vec<_>>(),
        vec![
            ("listing", 1, 0),
            ("metadata", 1, 0),
            ("artifact", 1, 0),
            ("upload", 1, 0),
            ("admin", 1, 0),
            ("authentication", 1, 0),
        ]
    );
}

#[test]
fn test_zero_limit_is_unbounded() {
    let limiter = RateLimiter::new(RateLimitConfig {
        listing: RouteLimit::new(0, 60),
        ..RateLimitConfig::enabled_defaults()
    });

    assert!(limiter.check_client(RouteClass::Listing, IpAddr::V4(Ipv4Addr::LOCALHOST)));
    assert!(limiter.check_client(RouteClass::Listing, IpAddr::V4(Ipv4Addr::LOCALHOST)));
}

/// The hot path asks this before touching forwarded headers at all, so a limiter that trusted
/// nobody yet answered "some proxy" would parse headers it then ignores, and one that answered
/// "none" over a configured proxy would bucket every proxied client by the proxy's address.
#[test]
fn test_a_limiter_reports_whether_any_proxy_is_trusted() {
    let proxied = RateLimiter::new(RateLimitConfig {
        trusted_proxies: vec!["10.0.0.0/8".parse().unwrap()],
        ..RateLimitConfig::enabled_defaults()
    });

    assert_eq!(
        (RateLimiter::default().trusts_any_proxy(), proxied.trusts_any_proxy()),
        (false, true)
    );
}

#[test]
fn test_proxy_trust_canonicalizes_addresses() {
    let limiter = RateLimiter::new(RateLimitConfig {
        trusted_proxies: vec!["127.0.0.0/8".parse().unwrap()],
        ..RateLimitConfig::enabled_defaults()
    });

    assert!(limiter.trusts_proxy("::ffff:127.0.0.1".parse().unwrap()));
    assert!(!limiter.trusts_proxy("198.51.100.1".parse().unwrap()));
}

#[test]
fn test_actor_key_uses_subject_or_resolved_client() {
    let limiter = RateLimiter::default();
    let request = Request::new(Body::empty());

    assert!(matches!(
        limiter
            .actor_key(
                peryx_identity::Principal::Named {
                    subject: "user".to_owned(),
                },
                &request,
            )
            .unwrap(),
        ActorKey::Token(_)
    ));
    assert_eq!(
        limiter
            .actor_key(peryx_identity::Principal::Anonymous, &request)
            .unwrap(),
        ActorKey::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST))
    );
}

fn proxied_limiter() -> RateLimiter {
    RateLimiter::new(RateLimitConfig {
        trusted_proxies: vec!["10.0.0.0/8".parse().unwrap()],
        ..RateLimitConfig::enabled_defaults()
    })
}

fn proxied_request() -> Request<Body> {
    let mut request = Request::new(Body::empty());
    request
        .extensions_mut()
        .insert(ConnectInfo(SocketAddr::from(([10, 0, 0, 1], 8080))));
    request
}

#[test]
fn test_client_ip_ignores_forwarded_headers_from_untrusted_peer() {
    let limiter = proxied_limiter();
    let mut request = Request::new(Body::empty());
    request
        .extensions_mut()
        .insert(ConnectInfo(SocketAddr::from(([192, 0, 2, 1], 8080))));
    request
        .headers_mut()
        .insert("x-forwarded-for", HeaderValue::from_static("198.51.100.1"));

    assert_eq!(
        limiter.client_ip(&request).unwrap(),
        Some(IpAddr::V4(Ipv4Addr::new(192, 0, 2, 1)))
    );
}

#[test]
fn test_actor_key_uses_the_attached_client_address() {
    let limiter = proxied_limiter();
    let mut request = proxied_request();
    request
        .headers_mut()
        .insert("x-forwarded-for", HeaderValue::from_static("198.51.100.1"));
    let mut request = client_address::attach(&limiter, request);
    request
        .headers_mut()
        .insert("x-forwarded-for", HeaderValue::from_static("198.51.100.2"));

    assert_eq!(
        limiter
            .actor_key(peryx_identity::Principal::Anonymous, &request)
            .unwrap(),
        ActorKey::Ip("198.51.100.1".parse().unwrap())
    );
}

#[test]
fn test_forwarded_chain_uses_rightmost_untrusted_client() {
    let limiter = proxied_limiter();
    let mut headers = HeaderMap::new();
    headers.append("x-forwarded-for", HeaderValue::from_static("192.0.2.1, 10.0.0.2"));
    headers.append("x-forwarded-for", HeaderValue::from_static("198.51.100.2"));

    assert!(matches!(
        limiter.forwarded_client_ip(&headers),
        ForwardedClient::Resolved(IpAddr::V4(address)) if address == Ipv4Addr::new(198, 51, 100, 2)
    ));
}

#[test]
fn test_forwarded_chain_rejects_malformed_suffix() {
    let limiter = proxied_limiter();
    let mut headers = HeaderMap::new();
    headers.insert("x-forwarded-for", HeaderValue::from_static("192.0.2.1, malformed"));

    assert!(matches!(
        limiter.forwarded_client_ip(&headers),
        ForwardedClient::Malformed
    ));
}

#[test]
fn test_fully_trusted_chain_uses_peer() {
    let limiter = proxied_limiter();
    let mut request = proxied_request();
    request
        .headers_mut()
        .insert("x-forwarded-for", HeaderValue::from_static("10.0.0.2"));

    assert_eq!(
        limiter.client_ip(&request).unwrap(),
        Some(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)))
    );
}

#[test]
fn test_real_ip_requires_one_valid_address() {
    assert!(matches!(real_ip(&HeaderMap::new()), ForwardedClient::TrustedChain));

    let mut headers = HeaderMap::new();
    headers.insert("x-real-ip", HeaderValue::from_static("198.51.100.2"));
    assert!(matches!(real_ip(&headers), ForwardedClient::Resolved(_)));

    headers.append("x-real-ip", HeaderValue::from_static("198.51.100.3"));
    assert!(matches!(real_ip(&headers), ForwardedClient::Malformed));

    let mut invalid = HeaderMap::new();
    invalid.insert("x-real-ip", HeaderValue::from_bytes(b"\xff").unwrap());
    assert!(matches!(real_ip(&invalid), ForwardedClient::Malformed));
}

#[tokio::test]
async fn test_upstream_limits_handle_unconfigured_unbounded_and_bounded_indexes() {
    let limits = UpstreamLimits::new([("unbounded".to_owned(), 0), ("bounded".to_owned(), 1)]);

    limits.acquire("missing").await.unwrap();
    limits.acquire("unbounded").await.unwrap();
    let permit = limits.acquire("bounded").await.unwrap();
    assert_eq!(
        limits
            .snapshots()
            .into_iter()
            .map(|snapshot| (
                snapshot.index,
                snapshot.max_concurrent,
                snapshot.max_waiting,
                snapshot.in_flight,
                snapshot.waiting,
                snapshot.denied,
                snapshot.admission_denied,
            ))
            .collect::<Vec<_>>(),
        vec![
            ("bounded".to_owned(), 1, 4, 1, 0, 0, 0),
            ("unbounded".to_owned(), 0, 0, 0, 0, 0, 0)
        ]
    );
    assert_eq!(limits.totals().in_flight, 1);
    drop(permit);
    assert_eq!(limits.totals().in_flight, 0);
}

#[tokio::test(start_paused = true)]
async fn test_upstream_limit_times_out_with_retry_horizon() {
    let limits = Arc::new(UpstreamLimits::new([("bounded".to_owned(), 1)]));
    let _permit = limits.acquire("bounded").await.unwrap();
    let waiting_limits = Arc::clone(&limits);
    let waiting = tokio::spawn(async move { waiting_limits.acquire("bounded").await });
    tokio::time::advance(Duration::from_secs(30)).await;

    let error = waiting.await.unwrap().unwrap_err();

    assert_eq!(error.retry_after, 30);
    assert_eq!(limits.snapshots()[0].denied, 1);
    assert_eq!(limits.totals().denied, 1);
    assert_eq!(limits.totals().admission_denied, 0);
}

type Acquisition<'a> =
    std::pin::Pin<Box<dyn std::future::Future<Output = Result<UpstreamPermit, UpstreamLimited>> + Send + 'a>>;

/// Parks `count` acquisitions of `index` on the semaphore queue, returning them still pending so
/// their admission slots stay held for as long as the caller keeps them alive.
async fn queued<'a>(limits: &'a UpstreamLimits, index: &'static str, count: usize) -> Vec<Acquisition<'a>> {
    let mut waiters: Vec<Acquisition<'a>> = (0..count)
        .map(|_| Box::pin(limits.acquire(index)) as Acquisition<'a>)
        .collect();
    for waiter in &mut waiters {
        assert!(futures_util::poll!(waiter.as_mut()).is_pending());
    }
    waiters
}

#[tokio::test(start_paused = true)]
async fn test_upstream_index_allowance_rejects_extra_waiters_before_queueing() {
    let limits = UpstreamLimits::new([("bounded".to_owned(), 1)]);
    let _active = limits.acquire("bounded").await.unwrap();

    let _waiters = queued(&limits, "bounded", 4).await;
    let rejected = limits.acquire("bounded").await.unwrap_err();

    assert_eq!(rejected.retry_after, 30);
    assert_eq!(limits.totals().waiting, 4);
    assert_eq!(limits.totals().admission_denied, 1);
    // Nothing waited, so the wait horizon never expired.
    assert_eq!(limits.totals().denied, 0);
}

#[tokio::test(start_paused = true)]
async fn test_upstream_full_index_queue_leaves_other_index_admitted() {
    let limits = UpstreamLimits::new([("busy".to_owned(), 1), ("quiet".to_owned(), 1)]);
    let _active = limits.acquire("busy").await.unwrap();
    let _waiters = queued(&limits, "busy", 4).await;

    limits.acquire("busy").await.unwrap_err();
    let _quiet = limits.acquire("quiet").await.unwrap();

    assert_eq!(limits.totals().in_flight, 2);
    assert_eq!(limits.totals().admission_denied, 1);
}

#[tokio::test(start_paused = true)]
async fn test_upstream_process_admission_caps_active_and_waiting_work() {
    let limits = UpstreamLimits::sharing(
        &Arc::new(tokio::sync::Semaphore::new(2)),
        [("busy".to_owned(), 1), ("quiet".to_owned(), 1)],
    );
    let _active = limits.acquire("busy").await.unwrap();
    let _waiter = queued(&limits, "busy", 1).await;

    let rejected = limits.acquire("quiet").await.unwrap_err();

    assert_eq!(rejected.retry_after, 30);
    assert_eq!(limits.totals().in_flight + limits.totals().waiting, 2);
    assert_eq!(limits.totals().admission_denied, 1);
}

#[tokio::test(start_paused = true)]
async fn test_upstream_cancelled_waiter_returns_its_admission_slot() {
    let limits = UpstreamLimits::sharing(
        &Arc::new(tokio::sync::Semaphore::new(2)),
        [("busy".to_owned(), 1), ("quiet".to_owned(), 1)],
    );
    let _active = limits.acquire("busy").await.unwrap();
    let waiter = queued(&limits, "busy", 1).await;

    drop(waiter);

    assert_eq!(limits.totals().waiting, 0);
    limits.acquire("quiet").await.unwrap();
}

#[tokio::test(start_paused = true)]
async fn test_upstream_sibling_gates_spend_one_process_admission() {
    let artifacts = UpstreamLimits::sharing(&Arc::new(tokio::sync::Semaphore::new(1)), [("bounded".to_owned(), 1)]);
    let metadata = artifacts.sibling([("bounded".to_owned(), 1)]);
    let _active = artifacts.acquire("bounded").await.unwrap();

    let rejected = metadata.acquire("bounded").await.unwrap_err();

    assert_eq!(rejected.retry_after, 30);
    assert_eq!(metadata.totals().admission_denied, 1);
    assert_eq!(metadata.totals().in_flight, 0);
}

fn app(config: RateLimitConfig) -> (tempfile::TempDir, AppState) {
    let dir = tempfile::tempdir().unwrap();
    let meta = MetaStore::open(dir.path().join("peryx.redb")).unwrap();
    let blobs = BlobStore::new(dir.path().join("blobs"));
    (dir, AppState::with_rate_limits(meta, blobs, 60, Vec::new(), config, []))
}

fn router(state: AppState) -> Router {
    Router::new()
        .fallback(get(|| async { StatusCode::NO_CONTENT }))
        .layer(middleware::from_fn_with_state(Arc::new(state), super::enforce))
}

fn process_request(path: &'static str, rate_limit: RouteRateLimit) -> Request<Body> {
    let mut request = Request::get(path).body(Body::empty()).unwrap();
    request.extensions_mut().insert(RouteDescriptor::new(
        RouteMethod::Get,
        path,
        RoutePosture::Read,
        rate_limit,
    ));
    request
}

#[tokio::test]
async fn test_enforce_bypasses_health_and_limits_admin_requests() {
    let config = RateLimitConfig {
        admin: RouteLimit::new(1, 60),
        ..RateLimitConfig::enabled_defaults()
    };
    let (_dir, state) = app(config);
    let router = router(state);

    assert_eq!(
        router
            .clone()
            .oneshot(process_request("/+health", RouteRateLimit::Exempt))
            .await
            .unwrap()
            .status(),
        StatusCode::NO_CONTENT
    );
    assert_eq!(
        router
            .clone()
            .oneshot(process_request("/+status", RouteRateLimit::Class(RouteClass::Admin)))
            .await
            .unwrap()
            .status(),
        StatusCode::NO_CONTENT
    );
    let response = router
        .oneshot(process_request("/+status", RouteRateLimit::Class(RouteClass::Admin)))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    assert!(
        (1..=60).contains(
            &response.headers()[header::RETRY_AFTER]
                .to_str()
                .unwrap()
                .parse::<u64>()
                .unwrap()
        )
    );
}

#[tokio::test]
async fn test_enforce_rejects_malformed_forwarded_identity() {
    let config = RateLimitConfig {
        trusted_proxies: vec!["10.0.0.0/8".parse().unwrap()],
        ..RateLimitConfig::enabled_defaults()
    };
    let (_dir, state) = app(config);
    let mut request = Request::get("/+status").body(Body::empty()).unwrap();
    request
        .extensions_mut()
        .insert(ConnectInfo(SocketAddr::from(([10, 0, 0, 1], 8080))));
    request
        .headers_mut()
        .insert("x-forwarded-for", HeaderValue::from_static("malformed"));

    assert_eq!(
        router(state).oneshot(request).await.unwrap().status(),
        StatusCode::BAD_REQUEST
    );
}

#[tokio::test]
async fn test_enforce_uses_the_indexed_drivers_route_class() {
    let dir = tempfile::tempdir().unwrap();
    let meta = MetaStore::open(dir.path().join("peryx.redb")).unwrap();
    let blobs = BlobStore::new(dir.path().join("blobs"));
    let mut state = AppState::with_rate_limits(
        meta,
        blobs,
        60,
        vec![Index {
            name: "items".to_owned(),
            route: "items".to_owned(),
            ecosystem: Ecosystem::new("example"),
            kind: IndexKind::Hosted { volatile: true },
            policy: Policy::default(),
            acl: IndexAcl::default(),
        }],
        RateLimitConfig {
            artifact: RouteLimit::new(1, 60),
            ..RateLimitConfig::enabled_defaults()
        },
        [],
    );
    state.register_rate_limit_principal(Ecosystem::new("example"), &IndexedDriver);
    state
        .register_protocol(ProtocolDriver::Indexed(Arc::new(IndexedDriver)), default_indexer())
        .unwrap();
    let router = router(state);

    let request = || {
        Request::get("/items/resource")
            .header(header::AUTHORIZATION, "opaque")
            .body(Body::empty())
            .unwrap()
    };
    assert_eq!(
        router.clone().oneshot(request()).await.unwrap().status(),
        StatusCode::NO_CONTENT
    );
    assert_eq!(
        router.oneshot(request()).await.unwrap().status(),
        StatusCode::TOO_MANY_REQUESTS
    );
}

#[tokio::test]
async fn test_enforce_uses_an_ecosystem_service_post_class() {
    let config = RateLimitConfig {
        admin: RouteLimit::new(1, 60),
        ..RateLimitConfig::enabled_defaults()
    };
    let (_dir, mut state) = app(config);
    state.register_capabilities(|registrar| {
        registrar.register_service(Ecosystem::new("example"), Arc::new(IndexedDriver));
    });
    let router = router(state);
    let request = || Request::post("/+special").body(Body::empty()).unwrap();

    assert_eq!(
        router.clone().oneshot(request()).await.unwrap().status(),
        StatusCode::METHOD_NOT_ALLOWED
    );
    assert_eq!(
        router.oneshot(request()).await.unwrap().status(),
        StatusCode::TOO_MANY_REQUESTS
    );
}

#[tokio::test]
async fn test_enforce_uses_an_absolute_drivers_route_class() {
    let config = RateLimitConfig {
        artifact: RouteLimit::new(1, 60),
        ..RateLimitConfig::enabled_defaults()
    };
    let (_dir, mut state) = app(config);
    state
        .register_protocol(ProtocolDriver::Absolute(Arc::new(AbsoluteDriver)), default_indexer())
        .unwrap();
    let router = router(state);
    let request = || Request::get("/artifacts/item").body(Body::empty()).unwrap();

    assert_eq!(
        router.clone().oneshot(request()).await.unwrap().status(),
        StatusCode::NO_CONTENT
    );
    assert_eq!(
        router.oneshot(request()).await.unwrap().status(),
        StatusCode::TOO_MANY_REQUESTS
    );
}

#[tokio::test]
async fn test_enforce_falls_back_to_the_listing_class() {
    let config = RateLimitConfig {
        listing: RouteLimit::new(1, 60),
        ..RateLimitConfig::enabled_defaults()
    };
    let (_dir, state) = app(config);
    let router = router(state);
    let request = || Request::get("/unknown").body(Body::empty()).unwrap();

    assert_eq!(
        router.clone().oneshot(request()).await.unwrap().status(),
        StatusCode::NO_CONTENT
    );
    assert_eq!(
        router.oneshot(request()).await.unwrap().status(),
        StatusCode::TOO_MANY_REQUESTS
    );
}

#[tokio::test]
async fn test_enforce_skips_drivers_without_service_posts() {
    let config = RateLimitConfig {
        upload: RouteLimit::new(1, 60),
        ..RateLimitConfig::enabled_defaults()
    };
    let (_dir, mut state) = app(config);
    state.register_driver(Arc::new(AbsoluteDriver));
    let router = router(state);
    let request = || Request::post("/upload").body(Body::empty()).unwrap();

    assert_eq!(
        router.clone().oneshot(request()).await.unwrap().status(),
        StatusCode::METHOD_NOT_ALLOWED
    );
    assert_eq!(
        router.oneshot(request()).await.unwrap().status(),
        StatusCode::TOO_MANY_REQUESTS
    );
}

#[tokio::test]
async fn test_enforce_falls_back_when_service_posts_decline() {
    let config = RateLimitConfig {
        upload: RouteLimit::new(1, 60),
        ..RateLimitConfig::enabled_defaults()
    };
    let (_dir, mut state) = app(config);
    state.register_capabilities(|registrar| {
        registrar.register_service(Ecosystem::new("example"), Arc::new(IndexedDriver));
    });
    let router = router(state);
    let request = || Request::post("/upload").body(Body::empty()).unwrap();

    assert_eq!(
        router.clone().oneshot(request()).await.unwrap().status(),
        StatusCode::METHOD_NOT_ALLOWED
    );
    assert_eq!(
        router.oneshot(request()).await.unwrap().status(),
        StatusCode::TOO_MANY_REQUESTS
    );
}

#[test]
fn test_client_ip_resolves_forwarded_and_real_ip_headers() {
    let limiter = proxied_limiter();
    let mut forwarded = proxied_request();
    forwarded
        .headers_mut()
        .insert("x-forwarded-for", HeaderValue::from_static("198.51.100.1"));
    let mut real = proxied_request();
    real.headers_mut()
        .insert("x-real-ip", HeaderValue::from_static("203.0.113.2"));

    assert_eq!(
        limiter.client_ip(&forwarded).unwrap(),
        Some("198.51.100.1".parse().unwrap())
    );
    assert_eq!(limiter.client_ip(&real).unwrap(), Some("203.0.113.2".parse().unwrap()));
}

#[test]
fn test_forwarded_chain_rejects_non_utf8_header_values() {
    let limiter = proxied_limiter();
    let mut headers = HeaderMap::new();
    headers.insert("x-forwarded-for", HeaderValue::from_bytes(&[0xff]).unwrap());

    assert!(matches!(
        limiter.forwarded_client_ip(&headers),
        ForwardedClient::Malformed
    ));
}

#[test]
fn test_rate_limit_errors_keep_status_and_retry_header() {
    assert_eq!(malformed_forwarded_response().status(), StatusCode::BAD_REQUEST);
    let response = limited_response(41);
    assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    assert_eq!(response.headers()[header::RETRY_AFTER], "41");
}

/// The classes a request actually spent, so a test can say which bucket took the charge rather than
/// only that some bucket refused. Two mutations of the classifier produce the same status by
/// charging the wrong class, and a status assertion cannot tell them apart.
fn charged(serving: &ServingState) -> Vec<(&'static str, u64, u64)> {
    serving
        .rate_limits
        .counters()
        .into_iter()
        .filter(|snapshot| snapshot.allowed > 0 || snapshot.denied > 0)
        .map(|snapshot| (snapshot.class, snapshot.allowed, snapshot.denied))
        .collect()
}

/// A service only classifies its own POST. The same path arriving as a GET is classified by route,
/// so a guard that stopped checking the method would charge a read to whatever class the service
/// claims for its write - here the admin bucket, which is far scarcer than listing.
#[tokio::test]
async fn test_enforce_keeps_a_service_post_class_off_a_get() {
    let (_dir, mut state) = app(RateLimitConfig::enabled_defaults());
    state.register_capabilities(|registrar| {
        registrar.register_service(Ecosystem::new("example"), Arc::new(IndexedDriver));
    });
    let serving = state.serving.clone();
    let router = router(state);

    let status = router
        .oneshot(Request::get("/+special").body(Body::empty()).unwrap())
        .await
        .unwrap()
        .status();

    assert_eq!(
        (status, charged(&serving)),
        (StatusCode::NO_CONTENT, vec![("listing", 1, 0)])
    );
}

fn routable_state(config: RateLimitConfig) -> (tempfile::TempDir, AppState) {
    let dir = tempfile::tempdir().unwrap();
    let meta = MetaStore::open(dir.path().join("peryx.redb")).unwrap();
    let blobs = BlobStore::new(dir.path().join("blobs"));
    let mut state = AppState::with_rate_limits(
        meta,
        blobs,
        60,
        vec![Index {
            name: "items".to_owned(),
            route: "items".to_owned(),
            ecosystem: Ecosystem::new("example"),
            kind: IndexKind::Hosted { volatile: true },
            policy: Policy::default(),
            acl: IndexAcl::default(),
        }],
        config,
        [],
    );
    state.register_rate_limit_principal(Ecosystem::new("example"), &IndexedDriver);
    state
        .register_protocol(ProtocolDriver::Indexed(Arc::new(IndexedDriver)), default_indexer())
        .unwrap();
    (dir, state)
}

fn declared_listing(credential: Option<&str>) -> Request<Body> {
    let mut request = process_request("/items/resource", RouteRateLimit::Class(RouteClass::Listing));
    if let Some(credential) = credential {
        request
            .headers_mut()
            .insert(header::AUTHORIZATION, HeaderValue::from_str(credential).unwrap());
    }
    request
}

/// A declared class settles the classification on its own, so the middleware skips the driver lookup
/// and every request from one address shares that address's bucket, credential or not.
///
/// Resolving a driver anyway would bucket the credentialed request by its subject instead, and the
/// pair would stop sharing - so one address could spend the class twice over. The status alone does
/// not show it: the first request succeeds either way, and only the second reveals whether the two
/// were charged to the same bucket.
#[tokio::test]
async fn test_enforce_shares_one_address_bucket_when_the_class_is_declared() {
    let (_dir, state) = routable_state(RateLimitConfig {
        listing: RouteLimit::new(1, 60),
        ..RateLimitConfig::enabled_defaults()
    });
    let serving = state.serving.clone();
    let router = router(state);

    let credentialed = router
        .clone()
        .oneshot(declared_listing(Some("opaque")))
        .await
        .unwrap()
        .status();
    let anonymous = router.oneshot(declared_listing(None)).await.unwrap().status();

    assert_eq!(
        (credentialed, anonymous, charged(&serving)),
        (
            StatusCode::NO_CONTENT,
            StatusCode::TOO_MANY_REQUESTS,
            vec![("listing", 1, 1)]
        )
    );
}

/// A configured limiter reports itself enabled, which is what every caller checks before spending
/// work on limiting at all. The sibling test covers the default, and a limiter that always answered
/// "disabled" would pass that one while turning limiting off everywhere.
#[test]
fn test_a_configured_limiter_reports_itself_enabled() {
    assert!(RateLimiter::new(RateLimitConfig::enabled_defaults()).enabled());
}

/// A refusal names the kind of client it charged, so an operator reading the security log can tell a
/// shared address apart from one credential spending its own bucket. Two kinds run here rather than
/// one, because a constant answer satisfies either on its own.
///
/// The credentialed pair is asserted first. It has to appear before the anonymous pair's absence of
/// `client=ip` would mean anything, since a capture that never reached the callsite reads the same as
/// one whose event carried a different value.
#[tokio::test]
async fn test_a_refusal_names_the_client_kind_it_charged() {
    let captured = crate::capture::Captured::install();
    let (_dir, state) = routable_state(RateLimitConfig {
        listing: RouteLimit::new(1, 60),
        artifact: RouteLimit::new(1, 60),
        ..RateLimitConfig::enabled_defaults()
    });
    let router = router(state);
    let credentialed = || {
        Request::get("/items/resource")
            .header(header::AUTHORIZATION, "Bearer opaque")
            .body(Body::empty())
            .unwrap()
    };

    router.clone().oneshot(credentialed()).await.unwrap();
    let by_credential = router.clone().oneshot(credentialed()).await.unwrap().status();
    router.clone().oneshot(declared_listing(None)).await.unwrap();
    let by_address = router.oneshot(declared_listing(None)).await.unwrap().status();

    assert_eq!(
        (by_credential, by_address),
        (StatusCode::TOO_MANY_REQUESTS, StatusCode::TOO_MANY_REQUESTS)
    );
    let output = captured.output();
    assert_eq!(output.matches(r#"client="token""#).count(), 1, "{output}");
    assert_eq!(output.matches(r#"client="ip""#).count(), 1, "{output}");
}