branch: main
rate_limit.rs
2349 bytesRaw
use super::SomeSharedData;
use std::collections::HashMap;
use worker::{js_sys, Env, Request, Response, Result};

#[worker::send]
pub async fn handle_rate_limit_check(
    _req: Request,
    env: Env,
    _data: SomeSharedData,
) -> Result<Response> {
    let rate_limiter = env.rate_limiter("TEST_RATE_LIMITER")?;

    // Use a fixed key for testing
    let outcome = rate_limiter.limit("test-key".to_string()).await?;

    Response::from_json(&serde_json::json!({
        "success": outcome.success,
    }))
}

#[worker::send]
pub async fn handle_rate_limit_with_key(
    req: Request,
    env: Env,
    _data: SomeSharedData,
) -> Result<Response> {
    let uri = req.url()?;
    let segments = uri.path_segments().unwrap().collect::<Vec<_>>();
    let key = segments.get(2).unwrap_or(&"default-key");

    let rate_limiter = env.rate_limiter("TEST_RATE_LIMITER")?;
    let outcome = rate_limiter.limit(key.to_string()).await?;

    Response::from_json(&serde_json::json!({
        "success": outcome.success,
        "key": key,
    }))
}

#[worker::send]
pub async fn handle_rate_limit_bulk_test(
    _req: Request,
    env: Env,
    _data: SomeSharedData,
) -> Result<Response> {
    let rate_limiter = env.rate_limiter("TEST_RATE_LIMITER")?;

    // Test multiple requests to verify rate limiting behavior
    let mut results = Vec::new();
    for i in 0..15 {
        let key = format!("bulk-test-{}", i % 3); // Use 3 different keys
        let outcome = rate_limiter.limit(key.clone()).await?;
        results.push(serde_json::json!({
            "index": i,
            "key": key,
            "success": outcome.success,
        }));
    }

    Response::from_json(&serde_json::json!({
        "results": results,
    }))
}

#[worker::send]
pub async fn handle_rate_limit_reset(
    _req: Request,
    env: Env,
    _data: SomeSharedData,
) -> Result<Response> {
    let rate_limiter = env.rate_limiter("TEST_RATE_LIMITER")?;

    // Use a unique key to avoid interference with other tests
    let key = format!("reset-test-{}", js_sys::Date::now());

    // Make multiple requests with the same key
    let mut outcomes = HashMap::new();
    for i in 0..12 {
        let outcome = rate_limiter.limit(key.clone()).await?;
        outcomes.insert(format!("request_{}", i + 1), outcome.success);
    }

    Response::from_json(&outcomes)
}