complete
This commit is contained in:
CarterPerez-dev
parent
b34cc4a0db
commit
3c66c1818f
|
|
@ -0,0 +1,283 @@
|
|||
# fastapi-420
|
||||
|
||||
Production rate limiting for FastAPI. Uses HTTP code 420 "Enhance Your Calm" because 429 is boring.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install fastapi-420
|
||||
```
|
||||
|
||||
For Redis support:
|
||||
|
||||
```bash
|
||||
pip install fastapi-420[redis]
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Three ways to add rate limiting. Pick what fits your app.
|
||||
|
||||
### Middleware (global)
|
||||
|
||||
Limits all routes automatically.
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI
|
||||
from fastapi_420 import RateLimiter, RateLimitMiddleware
|
||||
|
||||
app = FastAPI()
|
||||
limiter = RateLimiter()
|
||||
|
||||
app.add_middleware(
|
||||
RateLimitMiddleware,
|
||||
limiter=limiter,
|
||||
default_limit="100/minute",
|
||||
)
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {"message": "hello"}
|
||||
```
|
||||
|
||||
### Decorator (per route)
|
||||
|
||||
Fine grained control on specific endpoints.
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi_420 import RateLimiter
|
||||
|
||||
app = FastAPI()
|
||||
limiter = RateLimiter()
|
||||
|
||||
@app.get("/search")
|
||||
@limiter.limit("30/minute")
|
||||
async def search(request: Request, q: str):
|
||||
return {"results": []}
|
||||
|
||||
@app.post("/upload")
|
||||
@limiter.limit("5/minute", "20/hour")
|
||||
async def upload(request: Request):
|
||||
return {"status": "ok"}
|
||||
```
|
||||
|
||||
### Dependency (FastAPI style)
|
||||
|
||||
Works with FastAPI's dependency injection.
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI, Depends
|
||||
from fastapi_420 import RateLimiter, RateLimitDep, set_global_limiter
|
||||
|
||||
app = FastAPI()
|
||||
limiter = RateLimiter()
|
||||
set_global_limiter(limiter)
|
||||
|
||||
@app.get("/api/data", dependencies=[Depends(RateLimitDep("50/minute"))])
|
||||
async def get_data():
|
||||
return {"data": []}
|
||||
```
|
||||
|
||||
## Common Patterns
|
||||
|
||||
### Different limits for different endpoints
|
||||
|
||||
Auth endpoints get strict limits. Public endpoints stay relaxed.
|
||||
|
||||
```python
|
||||
from fastapi_420 import ScopedRateLimiter
|
||||
|
||||
auth_limiter = ScopedRateLimiter(
|
||||
prefix="/auth",
|
||||
default_rules=["5/minute"],
|
||||
endpoint_rules={
|
||||
"POST:/auth/login": ["3/minute", "10/hour"],
|
||||
"POST:/auth/register": ["2/minute"],
|
||||
},
|
||||
)
|
||||
|
||||
@app.post("/auth/login", dependencies=[Depends(auth_limiter)])
|
||||
async def login():
|
||||
...
|
||||
```
|
||||
|
||||
### Using Redis
|
||||
|
||||
Memory storage works fine for single instances. Redis for distributed apps.
|
||||
|
||||
```python
|
||||
from fastapi_420 import RateLimiter, RateLimiterSettings, StorageSettings
|
||||
|
||||
settings = RateLimiterSettings(
|
||||
storage=StorageSettings(
|
||||
REDIS_URL="redis://localhost:6379/0",
|
||||
),
|
||||
)
|
||||
|
||||
limiter = RateLimiter(settings=settings)
|
||||
```
|
||||
|
||||
If Redis goes down, the limiter falls back to memory automatically.
|
||||
|
||||
### Trusting proxy headers
|
||||
|
||||
Behind nginx or a load balancer? Trust the forwarded headers.
|
||||
|
||||
```python
|
||||
from fastapi_420 import FingerprintSettings, RateLimiterSettings
|
||||
|
||||
settings = RateLimiterSettings(
|
||||
fingerprint=FingerprintSettings(
|
||||
TRUST_X_FORWARDED_FOR=True,
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
### Custom identification
|
||||
|
||||
Rate limit by user ID instead of IP.
|
||||
|
||||
```python
|
||||
def get_user_id(request):
|
||||
return request.state.user_id or request.client.host
|
||||
|
||||
@app.get("/api/resource")
|
||||
@limiter.limit("100/minute", key_func=get_user_id)
|
||||
async def resource(request: Request):
|
||||
...
|
||||
```
|
||||
|
||||
## Configuration Reference
|
||||
|
||||
All settings with their defaults.
|
||||
|
||||
### RateLimiterSettings
|
||||
|
||||
```python
|
||||
from fastapi_420 import RateLimiterSettings
|
||||
from fastapi_420.types import Algorithm
|
||||
|
||||
RateLimiterSettings(
|
||||
# Algorithm
|
||||
ALGORITHM=Algorithm.SLIDING_WINDOW, # SLIDING_WINDOW | TOKEN_BUCKET | FIXED_WINDOW
|
||||
|
||||
# Defaults applied when no rules specified
|
||||
DEFAULT_LIMIT="100/minute",
|
||||
DEFAULT_LIMITS=["100/minute"], # list form, multiple rules
|
||||
|
||||
# Storage key configuration
|
||||
KEY_PREFIX="rl",
|
||||
KEY_VERSION="v1",
|
||||
|
||||
# Response behavior
|
||||
INCLUDE_HEADERS=True, # add RateLimit-* headers
|
||||
HTTP_420_MESSAGE="Enhance Your Calm",
|
||||
HTTP_420_DETAIL={"error": "rate_limit_exceeded", "message": "Enhance Your Calm"},
|
||||
|
||||
# Failure handling
|
||||
FAIL_OPEN=True, # allow requests if storage fails
|
||||
LOG_VIOLATIONS=True, # log when limits exceeded
|
||||
|
||||
# Nested settings (see below)
|
||||
storage=StorageSettings(...),
|
||||
fingerprint=FingerprintSettings(...),
|
||||
)
|
||||
```
|
||||
|
||||
### StorageSettings
|
||||
|
||||
```python
|
||||
from fastapi_420 import StorageSettings
|
||||
|
||||
StorageSettings(
|
||||
# Redis (optional, falls back to memory if not set or unavailable)
|
||||
REDIS_URL=None, # "redis://localhost:6379/0"
|
||||
REDIS_KEY_PREFIX="rl",
|
||||
REDIS_SOCKET_TIMEOUT=5.0,
|
||||
REDIS_SOCKET_CONNECT_TIMEOUT=5.0,
|
||||
REDIS_MAX_CONNECTIONS=50,
|
||||
REDIS_RETRY_ON_TIMEOUT=True,
|
||||
REDIS_HEALTH_CHECK_INTERVAL=30,
|
||||
|
||||
# Memory storage
|
||||
MEMORY_MAX_KEYS=100_000, # max keys before LRU eviction
|
||||
MEMORY_CLEANUP_INTERVAL=60, # seconds between expired key cleanup
|
||||
)
|
||||
```
|
||||
|
||||
### FingerprintSettings
|
||||
|
||||
Controls how clients are identified. Higher levels are stricter but may cause issues with legitimate users behind proxies.
|
||||
|
||||
```python
|
||||
from fastapi_420 import FingerprintSettings
|
||||
from fastapi_420.types import FingerprintLevel
|
||||
|
||||
FingerprintSettings(
|
||||
LEVEL=FingerprintLevel.NORMAL, # RELAXED | NORMAL | STRICT
|
||||
|
||||
# What to trust
|
||||
TRUST_X_FORWARDED_FOR=False, # trust X-Forwarded-For header
|
||||
TRUSTED_PROXIES=[], # IPs that can set forwarded headers
|
||||
|
||||
# IPv6 handling
|
||||
IPV6_PREFIX_LENGTH=64, # normalize IPv6 to /64 prefix
|
||||
)
|
||||
```
|
||||
|
||||
**Fingerprint Levels:**
|
||||
|
||||
| Level | What it uses |
|
||||
|-------|-------------|
|
||||
| RELAXED | IP only |
|
||||
| NORMAL | IP + User-Agent |
|
||||
| STRICT | IP + User-Agent + Accept headers + Auth token hash |
|
||||
|
||||
### Algorithms
|
||||
|
||||
| Algorithm | Behavior | Best for |
|
||||
|-----------|----------|----------|
|
||||
| SLIDING_WINDOW | Smooth, accurate limits | Most cases (default) |
|
||||
| TOKEN_BUCKET | Allows short bursts | APIs with bursty traffic |
|
||||
| FIXED_WINDOW | Simple, less accurate at window edges | High performance needs |
|
||||
|
||||
## Rate Limit Format
|
||||
|
||||
Rules follow the pattern `{requests}/{period}`:
|
||||
|
||||
```
|
||||
100/minute
|
||||
50/hour
|
||||
1000/day
|
||||
10/second
|
||||
```
|
||||
|
||||
Multiple rules stack. The most restrictive one applies:
|
||||
|
||||
```python
|
||||
@limiter.limit("10/second", "100/minute", "1000/hour")
|
||||
async def endpoint(request: Request):
|
||||
...
|
||||
```
|
||||
|
||||
## Running the Example
|
||||
|
||||
```bash
|
||||
cd examples
|
||||
docker compose up -d
|
||||
pip install fastapi uvicorn
|
||||
python app.py
|
||||
```
|
||||
|
||||
Then hit http://localhost:8000/docs to see the API.
|
||||
|
||||
## Why 420?
|
||||
|
||||
Twitter used HTTP 420 "Enhance Your Calm" for rate limiting before switching to 429. It is more fun.
|
||||
|
||||
The exception is called `EnhanceYourCalm` and the response tells clients to chill out.
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
|
|
@ -0,0 +1,223 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
app.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uvicorn
|
||||
from typing import Annotated
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import (
|
||||
Depends,
|
||||
FastAPI,
|
||||
Request,
|
||||
)
|
||||
from fastapi_420 import (
|
||||
RateLimiter,
|
||||
RateLimiterSettings,
|
||||
RateLimitMiddleware,
|
||||
ScopedRateLimiter,
|
||||
FingerprintSettings,
|
||||
StorageSettings,
|
||||
set_global_limiter,
|
||||
RateLimitDep,
|
||||
)
|
||||
from fastapi_420.types import Algorithm, FingerprintLevel
|
||||
|
||||
|
||||
storage_settings = StorageSettings(
|
||||
REDIS_URL = "redis://localhost:6767/0",
|
||||
MEMORY_MAX_KEYS = 50_000,
|
||||
)
|
||||
|
||||
fingerprint_settings = FingerprintSettings(
|
||||
LEVEL = FingerprintLevel.NORMAL,
|
||||
TRUST_X_FORWARDED_FOR = True,
|
||||
)
|
||||
|
||||
settings = RateLimiterSettings(
|
||||
ALGORITHM = Algorithm.SLIDING_WINDOW,
|
||||
KEY_PREFIX = "myapp",
|
||||
INCLUDE_HEADERS = True,
|
||||
FAIL_OPEN = True,
|
||||
LOG_VIOLATIONS = True,
|
||||
storage = storage_settings,
|
||||
fingerprint = fingerprint_settings,
|
||||
)
|
||||
|
||||
limiter = RateLimiter(settings = settings)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
await limiter.init()
|
||||
set_global_limiter(limiter)
|
||||
yield
|
||||
await limiter.close()
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title = "Rate Limited API",
|
||||
description = "Example API demonstrating fastapi_420 rate limiting",
|
||||
lifespan = lifespan,
|
||||
)
|
||||
|
||||
auth_limiter = ScopedRateLimiter(
|
||||
prefix = "/auth",
|
||||
default_rules = ["5/minute",
|
||||
"20/hour"],
|
||||
endpoint_rules = {
|
||||
"POST:/auth/login": ["3/minute",
|
||||
"10/hour"],
|
||||
"POST:/auth/register": ["2/minute",
|
||||
"5/hour"],
|
||||
"POST:/auth/forgot-password": ["2/minute",
|
||||
"5/hour"],
|
||||
},
|
||||
)
|
||||
|
||||
public_limiter = ScopedRateLimiter(
|
||||
prefix = "/public",
|
||||
default_rules = ["100/minute",
|
||||
"1000/hour"],
|
||||
)
|
||||
|
||||
user_limiter = ScopedRateLimiter(
|
||||
prefix = "/api",
|
||||
default_rules = ["60/minute",
|
||||
"500/hour"],
|
||||
endpoint_rules = {
|
||||
"POST:/api/upload": ["10/minute"],
|
||||
"GET:/api/export": ["5/minute"],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@app.post("/auth/login", dependencies = [Depends(auth_limiter)])
|
||||
async def login(username: str, password: str):
|
||||
"""
|
||||
Strict rate limit: 3/minute, 10/hour
|
||||
Prevents brute force attacks
|
||||
"""
|
||||
return {"token": "fake_jwt_token", "user": username}
|
||||
|
||||
|
||||
@app.post("/auth/register", dependencies = [Depends(auth_limiter)])
|
||||
async def register(username: str, email: str, password: str):
|
||||
"""
|
||||
Very strict: 2/minute, 5/hour
|
||||
Prevents mass account creation
|
||||
"""
|
||||
return {"user_id": 123, "username": username}
|
||||
|
||||
|
||||
@app.post("/auth/forgot-password", dependencies = [Depends(auth_limiter)])
|
||||
async def forgot_password(email: str):
|
||||
"""
|
||||
Very strict: 2/minute, 5/hour
|
||||
Prevents email bombing
|
||||
"""
|
||||
return {"message": "If that email exists, we sent a reset link"}
|
||||
|
||||
|
||||
@app.get("/public/products")
|
||||
@limiter.limit("100/minute")
|
||||
async def list_products(request: Request):
|
||||
"""
|
||||
Relaxed limit for public browsing
|
||||
Uses decorator style
|
||||
"""
|
||||
return {"products": [{"id": 1, "name": "Widget"}]}
|
||||
|
||||
|
||||
@app.get("/public/search")
|
||||
@limiter.limit("50/minute")
|
||||
async def search(request: Request, q: str):
|
||||
"""
|
||||
Slightly stricter for search (more expensive)
|
||||
"""
|
||||
return {"results": [], "query": q}
|
||||
|
||||
|
||||
@app.get("/api/me", dependencies = [Depends(user_limiter)])
|
||||
async def get_current_user():
|
||||
"""
|
||||
Standard user endpoint: 60/minute
|
||||
"""
|
||||
return {"user_id": 123, "username": "johndoe"}
|
||||
|
||||
|
||||
@app.get(
|
||||
"/api/dashboard",
|
||||
dependencies = [Depends(user_limiter)],
|
||||
)
|
||||
async def dashboard():
|
||||
"""
|
||||
Standard user endpoint: 60/minute
|
||||
"""
|
||||
return {"stats": {"visits": 1000, "conversions": 50}}
|
||||
|
||||
|
||||
@app.post("/api/upload", dependencies = [Depends(user_limiter)])
|
||||
async def upload_file():
|
||||
"""
|
||||
Stricter for uploads: 10/minute
|
||||
"""
|
||||
return {"file_id": "abc123", "status": "uploaded"}
|
||||
|
||||
|
||||
@app.get("/api/export", dependencies = [Depends(user_limiter)])
|
||||
async def export_data():
|
||||
"""
|
||||
Very strict for exports: 5/minute
|
||||
Expensive operation
|
||||
"""
|
||||
return {"download_url": "/files/export_123.csv"}
|
||||
|
||||
|
||||
@app.get(
|
||||
"/api/settings",
|
||||
dependencies = [Depends(RateLimitDep("30/minute"))],
|
||||
)
|
||||
async def get_settings():
|
||||
"""
|
||||
Using RateLimitDep directly for one off limits
|
||||
"""
|
||||
return {"theme": "dark", "notifications": True}
|
||||
|
||||
|
||||
@app.get("/admin/users")
|
||||
async def admin_list_users():
|
||||
"""
|
||||
No rate limit for admin endpoints
|
||||
In production you would add auth middleware
|
||||
"""
|
||||
return {"users": [{"id": 1, "username": "admin"}]}
|
||||
|
||||
|
||||
@app.get("/admin/stats")
|
||||
async def admin_stats():
|
||||
"""
|
||||
No rate limit
|
||||
"""
|
||||
return {"total_users": 10000, "active_today": 500}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
"""
|
||||
Health check, excluded from rate limiting by default
|
||||
"""
|
||||
storage_healthy = await limiter._storage.health_check(
|
||||
) if limiter._storage else False
|
||||
return {
|
||||
"status": "healthy",
|
||||
"limiter_initialized": limiter.is_initialized,
|
||||
"storage_healthy": storage_healthy,
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(app, host = "0.0.0.0", port = 8000)
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
services:
|
||||
redis:
|
||||
image: redis:7-alpine
|
||||
ports:
|
||||
- "6767:6379"
|
||||
volumes:
|
||||
- redis_data:/data
|
||||
command: redis-server --appendonly yes
|
||||
healthcheck:
|
||||
test: ["CMD", "redis-cli", "ping"]
|
||||
interval: 5s
|
||||
timeout: 3s
|
||||
retries: 5
|
||||
|
||||
volumes:
|
||||
redis_data:
|
||||
|
|
@ -0,0 +1,141 @@
|
|||
# How the Rate Limiting Algorithms Actually Work
|
||||
|
||||
This package ships three algorithms. Most people just pick one without understanding what they actually do differently. Here is what is happening under the hood.
|
||||
|
||||
## The Short Version
|
||||
|
||||
| Algorithm | Good at | Bad at |
|
||||
|-----------|---------|--------|
|
||||
| Sliding Window | Accuracy, memory efficiency | Nothing really |
|
||||
| Token Bucket | Handling bursts gracefully | Can feel unpredictable |
|
||||
| Fixed Window | Simplicity, speed | Boundary bursts (see below) |
|
||||
|
||||
Sliding window is the default because Cloudflare tested it across 400 million requests and measured 99.997% accuracy. That is not marketing, that is just what the math works out to.
|
||||
|
||||
## Fixed Window (and why it has a problem)
|
||||
|
||||
Fixed window is the simplest approach. Divide time into chunks (say, 1 minute windows), count requests in each chunk, block when the count hits the limit.
|
||||
|
||||
```
|
||||
Window 1 (0:00-0:59) Window 2 (1:00-1:59)
|
||||
[--------------------] [--------------------]
|
||||
87 requests 12 requests
|
||||
```
|
||||
|
||||
The problem nobody thinks about: what happens at the boundary?
|
||||
|
||||
Say you have a 100 requests/minute limit. A client sends 100 requests at 0:59 (allowed, counter is 100). Then at 1:00, the window resets and they send another 100 requests (allowed, new window). That is 200 requests in 2 seconds. The limit is supposed to be 100/minute but they got 200.
|
||||
|
||||
This is called the boundary burst problem. It is not theoretical. Attackers know about it.
|
||||
|
||||
```
|
||||
0:00 0:59 1:00 1:59
|
||||
[---------------||--][--||---------------]
|
||||
100 100
|
||||
^
|
||||
200 requests in 2 seconds
|
||||
```
|
||||
|
||||
Fixed window is still useful when you need maximum speed and can tolerate this edge case. But for most APIs, you want something better.
|
||||
|
||||
## Sliding Window Counter (the default)
|
||||
|
||||
This fixes the boundary problem by looking at two windows and doing weighted interpolation.
|
||||
|
||||
The idea: instead of hard window boundaries, calculate a weighted average based on how far into the current window you are.
|
||||
|
||||
```python
|
||||
elapsed_ratio = (current_time % window_size) / window_size
|
||||
weighted_count = (previous_window_count * (1 - elapsed_ratio)) + current_window_count
|
||||
```
|
||||
|
||||
Say you are 40% into the current minute. You take 60% of the previous window count plus 100% of the current window count. This creates a "sliding" effect where the window smoothly moves forward rather than jumping.
|
||||
|
||||
```
|
||||
Previous window Current window
|
||||
[-------50-------] [---20---.........]
|
||||
^
|
||||
40% elapsed
|
||||
|
||||
weighted = (50 * 0.6) + 20 = 30 + 20 = 50
|
||||
```
|
||||
|
||||
Memory cost is just two integers per client (previous count, current count). The 99.997% accuracy comes from the fact that you are approximating a true sliding window with minimal storage.
|
||||
|
||||
The tradeoff is about 6% average error compared to a perfect sliding window log. In practice nobody notices.
|
||||
|
||||
## Token Bucket (for bursty traffic)
|
||||
|
||||
Token bucket thinks about rate limiting differently. Instead of counting requests, you have a bucket that fills with tokens over time. Each request takes a token. No tokens, no request.
|
||||
|
||||
```
|
||||
Bucket capacity: 10 tokens
|
||||
Refill rate: 1 token/second
|
||||
|
||||
[OOOOOOOOOO] <- full bucket (10 tokens)
|
||||
|
||||
5 requests come in at once:
|
||||
[OOOOO-----] <- 5 tokens left
|
||||
|
||||
Wait 3 seconds:
|
||||
[OOOOOOOO--] <- refilled to 8 tokens
|
||||
```
|
||||
|
||||
The key difference from sliding window: burst tolerance.
|
||||
|
||||
With sliding window at 60 requests/minute, a client can only ever make 60 requests in any 60 second period. Steady, predictable.
|
||||
|
||||
With token bucket at 60 tokens/minute capacity and 1 token/second refill, a client can burn all 60 tokens instantly if they saved up, then wait for refills. Bursty, but still respects the average rate.
|
||||
|
||||
The math:
|
||||
```python
|
||||
tokens = min(capacity, tokens + (elapsed_time * refill_rate))
|
||||
if tokens >= 1:
|
||||
tokens -= 1
|
||||
return allowed
|
||||
return denied
|
||||
```
|
||||
|
||||
When to use it: APIs where legitimate clients have bursty patterns. Mobile apps that sync on open. Batch processing endpoints. Anything where "60 per minute average" makes more sense than "never more than 60 in any minute."
|
||||
|
||||
## How We Store This Stuff
|
||||
|
||||
All three algorithms need to track state per client. The storage layer handles this.
|
||||
|
||||
For sliding window, we store:
|
||||
```
|
||||
key: "rl:v1:user:GET:/api/data:abc123:60"
|
||||
value: {count: 47, window_start: 1703894400}
|
||||
```
|
||||
|
||||
For token bucket:
|
||||
```
|
||||
key: "rl:v1:user:GET:/api/data:abc123:bucket"
|
||||
value: {tokens: 8.5, last_update: 1703894567.123}
|
||||
```
|
||||
|
||||
Memory storage uses Python dicts with LRU eviction when you hit max keys. Redis storage uses Lua scripts to make the read-modify-write atomic (more on why in the architecture doc).
|
||||
|
||||
## Picking an Algorithm
|
||||
|
||||
Most of the time: sliding window. It is accurate, memory efficient, and has no weird edge cases.
|
||||
|
||||
Use token bucket when:
|
||||
- Your clients legitimately need burst capacity
|
||||
- "Average rate" matters more than "instantaneous rate"
|
||||
- You are rate limiting something like file uploads or batch operations
|
||||
|
||||
Use fixed window when:
|
||||
- You need maximum performance and minimal complexity
|
||||
- The boundary burst is acceptable for your use case
|
||||
- You are doing something like daily API quotas where minute-level precision does not matter
|
||||
|
||||
## The Math Behind Sliding Window Accuracy
|
||||
|
||||
If you are curious why the 6% error claim holds up:
|
||||
|
||||
The worst case is when all requests from the previous window happened at the very end, and all requests from the current window happened at the very beginning. The weighted formula slightly miscounts in this scenario.
|
||||
|
||||
But in practice, requests are distributed somewhat evenly (especially under attack, which is when accuracy matters most). Real world testing shows the error averages around 0.003% under normal load. The 6% figure is the theoretical worst case with adversarial timing.
|
||||
|
||||
For rate limiting, this is more than good enough. You are not doing financial accounting here. Being off by one or two requests out of a hundred is fine.
|
||||
|
|
@ -0,0 +1,247 @@
|
|||
# Architecture and Design Decisions
|
||||
|
||||
This doc explains why the package is built the way it is. Not how to use it, but why certain decisions were made.
|
||||
|
||||
## The Three Layer Defense Model
|
||||
|
||||
The limiter uses three layers of protection that work together:
|
||||
|
||||
1. **Per user, per endpoint** - Each user gets their own limit for each endpoint
|
||||
2. **Per endpoint global** - Each endpoint has an overall limit across all users
|
||||
3. **Circuit breaker** - Global kill switch when things go really wrong
|
||||
|
||||
Why three layers? Because different attacks hit different surfaces.
|
||||
|
||||
A credential stuffing attack hammers one endpoint (`/login`) from many IPs. Per user limits do not help because each IP only tries a few times. You need a global endpoint limit.
|
||||
|
||||
A single abusive user with a valid account might scrape your entire API. Global limits do not help because they are lost in normal traffic. You need per user limits.
|
||||
|
||||
A massive DDoS might overwhelm everything. Individual limits cannot keep up with the volume. You need a circuit breaker that just shuts things down temporarily.
|
||||
|
||||
```
|
||||
Request
|
||||
│
|
||||
▼
|
||||
┌─────────────────────┐
|
||||
│ Per User/Endpoint │ ─── Stops individual abuse
|
||||
└─────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────┐
|
||||
│ Per Endpoint Global │ ─── Stops coordinated attacks
|
||||
└─────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────┐
|
||||
│ Circuit Breaker │ ─── Emergency stop
|
||||
└─────────────────────┘
|
||||
│
|
||||
▼
|
||||
Allowed
|
||||
```
|
||||
|
||||
## Why Fail Open is the Default
|
||||
|
||||
When Redis goes down, you have two choices:
|
||||
|
||||
1. **Fail closed** - Block all requests until Redis comes back
|
||||
2. **Fail open** - Allow all requests until Redis comes back
|
||||
|
||||
This package defaults to fail open. Here is why.
|
||||
|
||||
Rate limiting is a protective measure, not a core business function. If your rate limiter fails, the worst case is some extra load for a few minutes. If you fail closed, your entire API goes down.
|
||||
|
||||
Think about it from an attacker's perspective. If they know you fail closed, they just need to DoS your Redis instance and your whole API dies. You turned a rate limiting dependency into a single point of failure.
|
||||
|
||||
The fail open approach includes an in-memory fallback. When Redis is unreachable, the limiter switches to local memory storage. It is not perfect (each server instance counts separately) but it is better than nothing and way better than blocking everything.
|
||||
|
||||
```python
|
||||
# What happens internally
|
||||
try:
|
||||
return await redis_storage.check(key, limit)
|
||||
except RedisConnectionError:
|
||||
logger.warning("Redis down, using memory fallback")
|
||||
return await memory_storage.check(key, limit)
|
||||
```
|
||||
|
||||
You can change this behavior with `FAIL_OPEN=False` if your use case genuinely requires fail closed semantics. But think hard about whether you actually need that.
|
||||
|
||||
## Why Redis Needs Lua Scripts
|
||||
|
||||
Rate limiting seems simple: read counter, increment, check limit. But there is a race condition hiding in plain sight.
|
||||
|
||||
```python
|
||||
# This looks fine but it is broken
|
||||
count = await redis.get(key)
|
||||
if count < limit:
|
||||
await redis.incr(key)
|
||||
return allowed
|
||||
return denied
|
||||
```
|
||||
|
||||
The problem: between the GET and INCR, another request can sneak in.
|
||||
|
||||
```
|
||||
Request A: GET key -> 99
|
||||
Request B: GET key -> 99
|
||||
Request A: 99 < 100, so INCR -> 100
|
||||
Request B: 99 < 100, so INCR -> 101 # Limit bypassed!
|
||||
```
|
||||
|
||||
This is not theoretical. Under load, this happens constantly. Attackers can deliberately time requests to exploit it.
|
||||
|
||||
The fix is atomic operations. Redis Lua scripts run as a single atomic unit. Nothing can interleave.
|
||||
|
||||
```lua
|
||||
local count = tonumber(redis.call('GET', key)) or 0
|
||||
if count < limit then
|
||||
redis.call('INCR', key)
|
||||
return 1 -- allowed
|
||||
end
|
||||
return 0 -- denied
|
||||
```
|
||||
|
||||
We also use EVALSHA instead of EVAL. EVAL sends the entire script text every time. EVALSHA sends a 40 byte hash and Redis looks up the cached script. Saves bandwidth on every single request.
|
||||
|
||||
## Composite Fingerprinting
|
||||
|
||||
IP based rate limiting has an obvious problem: attackers can rotate IPs.
|
||||
|
||||
Botnets have millions of IPs. Cloud providers give you a new IP every few seconds. IPv6 lets you control billions of addresses from a single allocation.
|
||||
|
||||
Composite fingerprinting combines multiple signals to identify clients:
|
||||
|
||||
```
|
||||
IP Address
|
||||
+
|
||||
User-Agent
|
||||
+
|
||||
Accept Headers
|
||||
+
|
||||
Auth Token (hashed)
|
||||
=
|
||||
Fingerprint
|
||||
```
|
||||
|
||||
The idea is that while IPs are easy to rotate, the combination of browser characteristics is harder to fake consistently. A real browser has a specific User-Agent, accepts specific encodings and languages, and generally looks consistent across requests.
|
||||
|
||||
An attacker scripting requests often:
|
||||
- Uses a generic or missing User-Agent
|
||||
- Has inconsistent or minimal Accept headers
|
||||
- Changes these values between requests (which itself is a signal)
|
||||
|
||||
The fingerprint levels:
|
||||
|
||||
| Level | What it uses | When to use |
|
||||
|-------|--------------|-------------|
|
||||
| RELAXED | Just IP | High volume public APIs, CDN cached content |
|
||||
| NORMAL | IP + User-Agent | Most APIs (default) |
|
||||
| STRICT | IP + User-Agent + Accept + Auth | Sensitive endpoints, auth flows |
|
||||
|
||||
Stricter is not always better. STRICT mode might accidentally rate limit legitimate users behind corporate proxies who share an IP but have different auth tokens.
|
||||
|
||||
## IPv6 Requires Special Handling
|
||||
|
||||
Most rate limiters treat each IP as unique. This is fine for IPv4 where addresses are scarce.
|
||||
|
||||
IPv6 is different. ISPs typically allocate a /64 prefix to each customer, which means one user controls 18 quintillion addresses. They can rotate through a new IP for every single request and never repeat.
|
||||
|
||||
The fix is to normalize IPv6 addresses to their /64 prefix before using them as rate limit keys.
|
||||
|
||||
```python
|
||||
from ipaddress import ip_address, ip_network
|
||||
|
||||
def normalize_ip(ip: str) -> str:
|
||||
addr = ip_address(ip)
|
||||
if addr.version == 6:
|
||||
# Treat entire /64 as one "user"
|
||||
network = ip_network(f"{ip}/64", strict=False)
|
||||
return str(network.network_address)
|
||||
return ip
|
||||
```
|
||||
|
||||
Now `2001:db8::1` and `2001:db8::ffff:ffff:ffff:ffff` become the same key: `2001:db8::`.
|
||||
|
||||
The /64 choice is not arbitrary. It matches how ISPs actually allocate addresses to end users. Smaller prefixes risk grouping unrelated users. Larger prefixes miss the attack surface.
|
||||
|
||||
## Key Naming Strategy
|
||||
|
||||
Rate limit keys need to be:
|
||||
- Unique per client/endpoint/window combination
|
||||
- Predictable (so you can debug)
|
||||
- Namespaced (so they do not collide with other Redis usage)
|
||||
|
||||
The format:
|
||||
```
|
||||
{prefix}:{version}:{layer}:{endpoint}:{identifier}:{window}
|
||||
```
|
||||
|
||||
Example:
|
||||
```
|
||||
rl:v1:user:GET:/api/users:a1b2c3d4:60
|
||||
```
|
||||
|
||||
Breaking it down:
|
||||
- `rl` - Key prefix, configurable, avoids collision with other Redis keys
|
||||
- `v1` - Version, lets you migrate key formats without clearing everything
|
||||
- `user` - Layer (user, endpoint, or global)
|
||||
- `GET:/api/users` - Endpoint including method (GET and POST are separate limits)
|
||||
- `a1b2c3d4` - First 16 chars of fingerprint hash
|
||||
- `60` - Window size in seconds
|
||||
|
||||
The version field is important for migrations. If you change how keys are structured, bump the version and old keys will naturally expire while new keys use the new format.
|
||||
|
||||
## Response Headers
|
||||
|
||||
The package adds standard rate limit headers to responses:
|
||||
|
||||
```
|
||||
RateLimit-Limit: 100
|
||||
RateLimit-Remaining: 67
|
||||
RateLimit-Reset: 45
|
||||
```
|
||||
|
||||
These follow the IETF draft standard (draft-ietf-httpapi-ratelimit-headers). Using standard headers means clients do not need custom code to handle your API.
|
||||
|
||||
When a request is blocked:
|
||||
|
||||
```
|
||||
HTTP/1.1 420 Enhance Your Calm
|
||||
Retry-After: 45
|
||||
RateLimit-Limit: 100
|
||||
RateLimit-Remaining: 0
|
||||
RateLimit-Reset: 45
|
||||
```
|
||||
|
||||
The Retry-After header tells clients exactly how long to wait. Good clients respect this. Bad clients ignore it and get blocked again.
|
||||
|
||||
## Why HTTP 420
|
||||
|
||||
The spec says to use 429 Too Many Requests. We use 420 instead.
|
||||
|
||||
Twitter invented 420 "Enhance Your Calm" for their rate limiting before the 429 standard existed. It is technically non-standard now but it is more memorable and honestly more fun.
|
||||
|
||||
The code works the same either way. We just return a different number and a message telling you to calm down.
|
||||
|
||||
If you genuinely need 429 for compatibility, you can catch `EnhanceYourCalm` and return your own response. But where is the fun in that?
|
||||
|
||||
## Storage Abstraction
|
||||
|
||||
The storage layer is abstracted so you can swap backends:
|
||||
|
||||
```python
|
||||
class Storage(Protocol):
|
||||
async def increment(self, key: str, window: int, limit: int) -> Result
|
||||
async def get_window_state(self, key: str, window: int) -> State
|
||||
async def consume_token(self, key: str, capacity: int, rate: float) -> Result
|
||||
async def health_check(self) -> bool
|
||||
async def close(self) -> None
|
||||
```
|
||||
|
||||
Currently implemented:
|
||||
- **MemoryStorage** - Dict based, good for dev and single instance deployments
|
||||
- **RedisStorage** - Lua script based, required for distributed deployments
|
||||
|
||||
The abstraction means you can add new backends (Memcached, DynamoDB, whatever) without touching the limiter logic.
|
||||
|
||||
MemoryStorage includes LRU eviction when you hit max keys. This prevents unbounded memory growth from attackers creating millions of unique fingerprints. The default is 100k keys which is about 2.4MB for sliding window counters.
|
||||
|
|
@ -0,0 +1,325 @@
|
|||
# Security Considerations and Common Mistakes
|
||||
|
||||
Rate limiting seems simple until you realize how many ways it can fail. This doc covers real vulnerabilities, attack vectors, and the mistakes that show up in CVEs over and over again.
|
||||
|
||||
## Real CVEs: What Goes Wrong
|
||||
|
||||
### CVE-2023-2531 (AzuraCast) - X-Forwarded-For Spoofing
|
||||
|
||||
CVSS 7.5. The `getIp()` method trusted the X-Forwarded-For header without validation.
|
||||
|
||||
```python
|
||||
# What they did (broken)
|
||||
def get_ip(request):
|
||||
if 'X-Forwarded-For' in request.headers:
|
||||
return request.headers['X-Forwarded-For'].split(',')[0]
|
||||
return request.remote_addr
|
||||
```
|
||||
|
||||
An attacker just sends:
|
||||
```
|
||||
X-Forwarded-For: 1.2.3.4
|
||||
```
|
||||
|
||||
Now every request looks like it comes from a different IP. Rate limiting completely bypassed.
|
||||
|
||||
The fix is to only trust X-Forwarded-For from known proxy IPs:
|
||||
|
||||
```python
|
||||
TRUSTED_PROXIES = {'10.0.0.1', '10.0.0.2'}
|
||||
|
||||
def get_ip(request):
|
||||
if request.remote_addr in TRUSTED_PROXIES:
|
||||
forwarded = request.headers.get('X-Forwarded-For', '')
|
||||
if forwarded:
|
||||
return forwarded.split(',')[0].strip()
|
||||
return request.remote_addr
|
||||
```
|
||||
|
||||
This package has `TRUST_X_FORWARDED_FOR=False` by default. You have to explicitly enable it and should also configure `TRUSTED_PROXIES`.
|
||||
|
||||
### CVE-2023-46745 (LibreNMS) - Inconsistent Limit Application
|
||||
|
||||
Rate limiting was applied to POST `/login` but not GET `/login`. The login form could be submitted via GET with query parameters, bypassing limits entirely.
|
||||
|
||||
```
|
||||
POST /login <- rate limited
|
||||
GET /login?username=x&password=y <- not rate limited
|
||||
```
|
||||
|
||||
The lesson: rate limit the action, not just a specific method. This package includes the HTTP method in the endpoint key by default (`GET:/api/data` vs `POST:/api/data`), but you need to make sure all paths to sensitive operations are covered.
|
||||
|
||||
### GHSA-984p-xq9m-4rjw (express-brute) - Race Conditions
|
||||
|
||||
Non-atomic counter updates allowed concurrent requests to bypass limits. Classic TOCTOU (time of check to time of use) bug.
|
||||
|
||||
```javascript
|
||||
// Broken: gap between read and write
|
||||
const count = await store.get(key);
|
||||
if (count < limit) {
|
||||
await store.set(key, count + 1); // Race condition here
|
||||
next();
|
||||
}
|
||||
```
|
||||
|
||||
Between the get and set, dozens of requests can sneak through. Under load, this completely defeats rate limiting.
|
||||
|
||||
This is why we use Redis Lua scripts. The entire check-and-increment runs atomically.
|
||||
|
||||
## Attack Vectors to Understand
|
||||
|
||||
### Header Spoofing
|
||||
|
||||
Any header can be spoofed. X-Forwarded-For, X-Real-IP, CF-Connecting-IP. Attackers know which headers your reverse proxy sets.
|
||||
|
||||
The only safe approach:
|
||||
1. Know exactly which proxies sit in front of your app
|
||||
2. Only trust forwarded headers from those specific IPs
|
||||
3. Take the rightmost IP added by your trusted proxy, not the leftmost (which the client controls)
|
||||
|
||||
```
|
||||
X-Forwarded-For: attacker-spoofed, real-proxy-added
|
||||
^ ^
|
||||
client controls proxy controls
|
||||
```
|
||||
|
||||
### GraphQL Batching
|
||||
|
||||
This one catches people off guard. GraphQL lets you batch multiple operations in one HTTP request:
|
||||
|
||||
```json
|
||||
[
|
||||
{"query": "mutation { login(user: \"x\", pass: \"a\") }"},
|
||||
{"query": "mutation { login(user: \"x\", pass: \"b\") }"},
|
||||
{"query": "mutation { login(user: \"x\", pass: \"c\") }"}
|
||||
]
|
||||
```
|
||||
|
||||
That is one HTTP request but three login attempts. If you rate limit by request, you are letting 100 password guesses through per "request."
|
||||
|
||||
GraphQL endpoints need operation-level limiting:
|
||||
- Count mutations, not requests
|
||||
- Limit query depth and complexity
|
||||
- Cap batch sizes (5-10 operations max)
|
||||
|
||||
### Slowloris and Connection Exhaustion
|
||||
|
||||
Rate limiting usually focuses on completed requests. But what about requests that never complete?
|
||||
|
||||
Slowloris attacks open many connections and send data very slowly, tying up server resources without triggering request-based limits.
|
||||
|
||||
This is outside the scope of application-level rate limiting. You need:
|
||||
- Connection timeouts at the reverse proxy level
|
||||
- Per-IP connection limits in nginx/haproxy
|
||||
- Request body size limits
|
||||
|
||||
### IPv6 Address Rotation
|
||||
|
||||
Covered in the architecture doc, but worth repeating: a single residential IPv6 user controls 18 quintillion addresses in their /64 prefix.
|
||||
|
||||
```
|
||||
2001:db8:abcd:1234::1
|
||||
2001:db8:abcd:1234::2
|
||||
2001:db8:abcd:1234::ffff:ffff:ffff:ffff
|
||||
```
|
||||
|
||||
All controlled by the same person. If you rate limit per IP, they just rotate.
|
||||
|
||||
Normalize to /64 prefix before rate limiting:
|
||||
```python
|
||||
2001:db8:abcd:1234::1 -> 2001:db8:abcd:1234::
|
||||
```
|
||||
|
||||
## Redis Security
|
||||
|
||||
Redis is a common target. Default installations often have no authentication and listen on all interfaces.
|
||||
|
||||
### CVE-2025-49844 (RediShell) - CVSS 10.0
|
||||
|
||||
Remote code execution through Lua sandbox escape. Affected around 330,000 exposed instances.
|
||||
|
||||
Protection:
|
||||
1. Always require authentication (`requirepass`)
|
||||
2. Bind to localhost only unless you specifically need network access
|
||||
3. Use ACLs to restrict what the rate limiter can do
|
||||
4. Disable dangerous commands
|
||||
|
||||
```redis
|
||||
# Create a restricted user for rate limiting
|
||||
ACL SETUSER ratelimiter on >strongpassword ~rl:* +EVALSHA +GET +SET +INCR +EXPIRE +TTL
|
||||
```
|
||||
|
||||
This user can only:
|
||||
- Run EVALSHA (our Lua scripts)
|
||||
- Access keys starting with `rl:`
|
||||
- Do basic counter operations
|
||||
|
||||
They cannot FLUSHALL, CONFIG, DEBUG, or do anything administrative.
|
||||
|
||||
### Key Expiration
|
||||
|
||||
Every rate limit key must have a TTL. Without expiration:
|
||||
- Memory grows forever
|
||||
- Old keys from users who never return pile up
|
||||
- Attackers can create millions of keys by rotating identifiers
|
||||
|
||||
We set TTL to 2x the window size. A 60 second window gets a 120 second TTL. This ensures the key lives long enough for the sliding window to work, but expires after.
|
||||
|
||||
## Common Implementation Mistakes
|
||||
|
||||
### 1. Not Testing Under Load
|
||||
|
||||
Unit tests pass. Integration tests pass. Production gets slammed and requests slip through.
|
||||
|
||||
Race conditions only appear under concurrent load. You need to actually test with hundreds of simultaneous requests:
|
||||
|
||||
```python
|
||||
async def test_concurrent_limit():
|
||||
tasks = [make_request() for _ in range(200)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
allowed = sum(1 for r in results if r.status == 200)
|
||||
blocked = sum(1 for r in results if r.status == 420)
|
||||
|
||||
# With limit of 100, should see ~100 allowed, ~100 blocked
|
||||
assert 95 <= allowed <= 105
|
||||
```
|
||||
|
||||
### 2. Blocking Logging
|
||||
|
||||
```python
|
||||
# Bad: blocks the event loop
|
||||
logger.info(f"Rate limited {ip}") # Synchronous I/O
|
||||
```
|
||||
|
||||
Under attack, you might rate limit thousands of requests per second. Synchronous logging blocks the event loop and slows everything down.
|
||||
|
||||
Use async logging, batch writes, or background tasks:
|
||||
|
||||
```python
|
||||
# Better: non-blocking
|
||||
background_tasks.add_task(log_violation, ip, endpoint)
|
||||
```
|
||||
|
||||
### 3. Forgetting TTLs
|
||||
|
||||
```python
|
||||
# Bad: key lives forever
|
||||
await redis.incr(key)
|
||||
|
||||
# Good: key expires
|
||||
await redis.incr(key)
|
||||
await redis.expire(key, window_seconds * 2)
|
||||
```
|
||||
|
||||
Actually even better, do it atomically in Lua so there is no gap where the key exists without a TTL.
|
||||
|
||||
### 4. Integer Overflow
|
||||
|
||||
At extreme scale, counters can overflow. Python handles big integers fine, but Redis stores numbers as 64-bit signed integers. Max value is about 9.2 quintillion.
|
||||
|
||||
You probably will not hit this, but if you are doing high volume with long windows, consider it.
|
||||
|
||||
### 5. Clock Drift
|
||||
|
||||
Distributed systems with multiple app servers can have clock skew. Server A thinks it is 10:00:00, Server B thinks it is 10:00:02.
|
||||
|
||||
For sliding window, use Redis server time instead of local time:
|
||||
|
||||
```python
|
||||
server_time = await redis.time() # Returns (seconds, microseconds)
|
||||
```
|
||||
|
||||
Now all servers agree on what time it is.
|
||||
|
||||
### 6. Fail Closed Without Monitoring
|
||||
|
||||
If you choose to fail closed (block all requests when Redis is down), you absolutely need monitoring and alerting on Redis availability.
|
||||
|
||||
Otherwise Redis goes down at 3am, your API becomes completely unavailable, and you do not find out until customers complain.
|
||||
|
||||
### 7. Overly Strict Fingerprinting
|
||||
|
||||
STRICT fingerprinting mode includes Accept headers and auth tokens. But this can cause problems:
|
||||
|
||||
- Same user on different browsers gets different fingerprints
|
||||
- API clients with slightly different Accept headers look like different users
|
||||
- Token rotation creates new fingerprints
|
||||
|
||||
Start with NORMAL (IP + User-Agent) and only go stricter if you have a specific problem to solve.
|
||||
|
||||
## OWASP API4:2023 - Unrestricted Resource Consumption
|
||||
|
||||
Rate limiting directly addresses OWASP API4, which identifies these failure modes:
|
||||
|
||||
- Missing or inadequate rate limiting
|
||||
- Lack of payload size limits
|
||||
- No execution timeouts
|
||||
- Unlimited resource allocation per request
|
||||
|
||||
A complete defense includes:
|
||||
- Request rate limits (what this package does)
|
||||
- Payload size limits (configure in your web server)
|
||||
- Execution timeouts (configure in your ASGI server)
|
||||
- Memory limits (configure at the container/process level)
|
||||
|
||||
Rate limiting is one layer. You need the others too.
|
||||
|
||||
## Testing Your Implementation
|
||||
|
||||
### Verify limits actually work
|
||||
|
||||
```bash
|
||||
# Send 110 requests, expect ~100 to succeed
|
||||
for i in {1..110}; do
|
||||
curl -s -o /dev/null -w "%{http_code}\n" http://localhost:8000/api/test
|
||||
done | sort | uniq -c
|
||||
```
|
||||
|
||||
Should see roughly 100 responses of 200 and 10 responses of 420.
|
||||
|
||||
### Verify headers are present
|
||||
|
||||
```bash
|
||||
curl -i http://localhost:8000/api/test
|
||||
```
|
||||
|
||||
Look for:
|
||||
```
|
||||
RateLimit-Limit: 100
|
||||
RateLimit-Remaining: 99
|
||||
RateLimit-Reset: 58
|
||||
```
|
||||
|
||||
### Test concurrent requests
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
import httpx
|
||||
|
||||
async def blast():
|
||||
async with httpx.AsyncClient() as client:
|
||||
tasks = [client.get("http://localhost:8000/api/test") for _ in range(200)]
|
||||
responses = await asyncio.gather(*tasks)
|
||||
|
||||
codes = [r.status_code for r in responses]
|
||||
print(f"200: {codes.count(200)}, 420: {codes.count(420)}")
|
||||
|
||||
asyncio.run(blast())
|
||||
```
|
||||
|
||||
### Test failover
|
||||
|
||||
```bash
|
||||
# Stop Redis
|
||||
docker stop redis
|
||||
|
||||
# Requests should still work (fail open)
|
||||
curl http://localhost:8000/api/test
|
||||
|
||||
# Check logs for fallback warning
|
||||
```
|
||||
|
||||
## Reporting Security Issues
|
||||
|
||||
If you find a security vulnerability in this package, do not open a public issue. Email the maintainer directly so it can be fixed before disclosure.
|
||||
|
|
@ -3,7 +3,12 @@
|
|||
__init__.py
|
||||
"""
|
||||
|
||||
from fastapi_420.config import RateLimiterSettings, get_settings
|
||||
from fastapi_420.config import (
|
||||
FingerprintSettings,
|
||||
RateLimiterSettings,
|
||||
StorageSettings,
|
||||
get_settings,
|
||||
)
|
||||
from fastapi_420.defense import CircuitBreaker, LayeredDefense
|
||||
from fastapi_420.dependencies import (
|
||||
LimiterDep,
|
||||
|
|
@ -44,6 +49,7 @@ __all__ = [
|
|||
"EnhanceYourCalm",
|
||||
"FingerprintData",
|
||||
"FingerprintLevel",
|
||||
"FingerprintSettings",
|
||||
"Layer",
|
||||
"LayeredDefense",
|
||||
"LimiterDep",
|
||||
|
|
@ -58,6 +64,7 @@ __all__ = [
|
|||
"ScopedRateLimiter",
|
||||
"SlowDownMiddleware",
|
||||
"StorageError",
|
||||
"StorageSettings",
|
||||
"__version__",
|
||||
"create_rate_limit_dep",
|
||||
"get_limiter",
|
||||
|
|
|
|||
|
|
@ -180,7 +180,8 @@ class LayeredDefense:
|
|||
).build()
|
||||
|
||||
endpoint_rule = RateLimitRule(
|
||||
requests = rule.requests * self._settings.defense.ENDPOINT_LIMIT_MULTIPLIER,
|
||||
requests = rule.requests *
|
||||
self._settings.defense.ENDPOINT_LIMIT_MULTIPLIER,
|
||||
window_seconds = rule.window_seconds,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from fastapi_420.config import (
|
|||
)
|
||||
from fastapi_420.exceptions import (
|
||||
EnhanceYourCalm,
|
||||
StorageConnectionError,
|
||||
StorageError,
|
||||
)
|
||||
from fastapi_420.fingerprinting import (
|
||||
|
|
@ -91,18 +92,34 @@ class RateLimiter:
|
|||
if self._storage is None:
|
||||
self._storage = create_storage(self._settings.storage)
|
||||
|
||||
if isinstance(self._storage, RedisStorage):
|
||||
await self._storage.connect()
|
||||
|
||||
if isinstance(self._storage, MemoryStorage):
|
||||
await self._storage.start_cleanup_task()
|
||||
|
||||
if self._settings.storage.FALLBACK_TO_MEMORY:
|
||||
self._fallback_storage = MemoryStorage.from_settings(
|
||||
self._settings.storage
|
||||
)
|
||||
await self._fallback_storage.start_cleanup_task()
|
||||
|
||||
if isinstance(self._storage, RedisStorage):
|
||||
try:
|
||||
await self._storage.connect()
|
||||
except StorageConnectionError:
|
||||
if self._settings.FAIL_OPEN and self._fallback_storage:
|
||||
logger.warning(
|
||||
"Redis unavailable, using memory fallback",
|
||||
extra = {
|
||||
"redis_url":
|
||||
self._settings.storage.REDIS_URL
|
||||
},
|
||||
)
|
||||
self._storage = self._fallback_storage
|
||||
self._fallback_storage = None
|
||||
else:
|
||||
raise
|
||||
|
||||
if isinstance(self._storage,
|
||||
MemoryStorage
|
||||
) and self._storage != self._fallback_storage:
|
||||
await self._storage.start_cleanup_task()
|
||||
|
||||
self._algorithm = create_algorithm(self._settings.ALGORITHM)
|
||||
self._fingerprinter = CompositeFingerprinter.from_settings(
|
||||
self._settings.fingerprint
|
||||
|
|
|
|||
|
|
@ -87,32 +87,30 @@ class RateLimitMiddleware(BaseHTTPMiddleware):
|
|||
if not await self._should_limit(request):
|
||||
return await call_next(request) # type: ignore[no-any-return]
|
||||
|
||||
try:
|
||||
limit = self._get_limit_for_path(request.url.path)
|
||||
await self.limiter.check(
|
||||
request,
|
||||
limit,
|
||||
key_func = self.key_func,
|
||||
raise_on_limit = True,
|
||||
limit = self._get_limit_for_path(request.url.path)
|
||||
result = await self.limiter.check(
|
||||
request,
|
||||
limit,
|
||||
key_func = self.key_func,
|
||||
raise_on_limit = False,
|
||||
)
|
||||
|
||||
if not result.allowed:
|
||||
exc = EnhanceYourCalm(
|
||||
result = result,
|
||||
message = self.limiter.settings.HTTP_420_MESSAGE,
|
||||
detail = self.limiter.settings.HTTP_420_DETAIL,
|
||||
)
|
||||
|
||||
response = await call_next(request)
|
||||
|
||||
if self.limiter.settings.INCLUDE_HEADERS:
|
||||
result = await self.limiter.check(
|
||||
request,
|
||||
limit,
|
||||
key_func = self.key_func,
|
||||
raise_on_limit = False,
|
||||
)
|
||||
for header_name, header_value in result.headers.items():
|
||||
response.headers[header_name] = header_value
|
||||
|
||||
return response # type: ignore[no-any-return]
|
||||
|
||||
except EnhanceYourCalm as exc:
|
||||
return self._create_420_response(exc)
|
||||
|
||||
response = await call_next(request)
|
||||
|
||||
if self.limiter.settings.INCLUDE_HEADERS:
|
||||
for header_name, header_value in result.headers.items():
|
||||
response.headers[header_name] = header_value
|
||||
|
||||
return response # type: ignore[no-any-return]
|
||||
|
||||
async def _should_limit(self, request: Request) -> bool:
|
||||
"""
|
||||
Determine if request should be rate limited
|
||||
|
|
|
|||
|
|
@ -101,7 +101,7 @@ class MockScope:
|
|||
method: str = "GET"
|
||||
path: str = "/"
|
||||
query_string: bytes = b""
|
||||
headers: list[tuple[bytes, bytes]] = field(default_factory=list)
|
||||
headers: list[tuple[bytes, bytes]] = field(default_factory = list)
|
||||
client: tuple[str, int] | None = None
|
||||
route: Any = None
|
||||
|
||||
|
|
@ -124,9 +124,12 @@ class RequestFactory:
|
|||
path: str = TEST_ENDPOINT,
|
||||
client_ip: str = TEST_IP_V4,
|
||||
client_port: int = 12345,
|
||||
headers: dict[str, str] | None = None,
|
||||
query_params: dict[str, str] | None = None,
|
||||
cookies: dict[str, str] | None = None,
|
||||
headers: dict[str,
|
||||
str] | None = None,
|
||||
query_params: dict[str,
|
||||
str] | None = None,
|
||||
cookies: dict[str,
|
||||
str] | None = None,
|
||||
include_route: bool = True,
|
||||
) -> Request:
|
||||
"""
|
||||
|
|
@ -160,7 +163,8 @@ class RequestFactory:
|
|||
"path": path,
|
||||
"query_string": query_string,
|
||||
"headers": header_list,
|
||||
"client": (client_ip, client_port),
|
||||
"client": (client_ip,
|
||||
client_port),
|
||||
}
|
||||
|
||||
if include_route:
|
||||
|
|
@ -191,7 +195,7 @@ class RequestFactory:
|
|||
elif auth_type == "basic":
|
||||
headers["authorization"] = f"Basic {token}"
|
||||
|
||||
return RequestFactory.create(headers=headers, **kwargs)
|
||||
return RequestFactory.create(headers = headers, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def with_forwarded_for(
|
||||
|
|
@ -208,7 +212,7 @@ class RequestFactory:
|
|||
if real_ip:
|
||||
headers["x-real-ip"] = real_ip
|
||||
|
||||
return RequestFactory.create(headers=headers, **kwargs)
|
||||
return RequestFactory.create(headers = headers, **kwargs)
|
||||
|
||||
|
||||
class FingerprintFactory:
|
||||
|
|
@ -231,15 +235,15 @@ class FingerprintFactory:
|
|||
Create FingerprintData with sensible defaults
|
||||
"""
|
||||
return FingerprintData(
|
||||
ip=ip,
|
||||
ip_normalized=ip_normalized or ip,
|
||||
user_agent=user_agent,
|
||||
accept_language=accept_language,
|
||||
accept_encoding=accept_encoding,
|
||||
headers_hash=headers_hash,
|
||||
auth_identifier=auth_identifier,
|
||||
tls_fingerprint=tls_fingerprint,
|
||||
geo_asn=geo_asn,
|
||||
ip = ip,
|
||||
ip_normalized = ip_normalized or ip,
|
||||
user_agent = user_agent,
|
||||
accept_language = accept_language,
|
||||
accept_encoding = accept_encoding,
|
||||
headers_hash = headers_hash,
|
||||
auth_identifier = auth_identifier,
|
||||
tls_fingerprint = tls_fingerprint,
|
||||
geo_asn = geo_asn,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -253,16 +257,19 @@ class FingerprintFactory:
|
|||
"""
|
||||
identifier = auth_id
|
||||
if hash_id:
|
||||
identifier = hashlib.sha256(auth_id.encode()).hexdigest()[:16]
|
||||
identifier = hashlib.sha256(auth_id.encode()).hexdigest()[: 16]
|
||||
|
||||
return FingerprintFactory.create(auth_identifier=identifier, **kwargs)
|
||||
return FingerprintFactory.create(
|
||||
auth_identifier = identifier,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def anonymous(**kwargs: Any) -> FingerprintData:
|
||||
"""
|
||||
Create anonymous fingerprint (no auth)
|
||||
"""
|
||||
return FingerprintFactory.create(auth_identifier=None, **kwargs)
|
||||
return FingerprintFactory.create(auth_identifier = None, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def minimal(ip: str = TEST_IP_V4) -> FingerprintData:
|
||||
|
|
@ -270,8 +277,8 @@ class FingerprintFactory:
|
|||
Create minimal fingerprint (IP only)
|
||||
"""
|
||||
return FingerprintData(
|
||||
ip=ip,
|
||||
ip_normalized=ip,
|
||||
ip = ip,
|
||||
ip_normalized = ip,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -287,29 +294,44 @@ class RuleFactory:
|
|||
"""
|
||||
Create RateLimitRule with defaults
|
||||
"""
|
||||
return RateLimitRule(requests=requests, window_seconds=window_seconds)
|
||||
return RateLimitRule(
|
||||
requests = requests,
|
||||
window_seconds = window_seconds
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def per_second(requests: int = 10) -> RateLimitRule:
|
||||
return RateLimitRule(requests=requests, window_seconds=WINDOW_SECOND)
|
||||
return RateLimitRule(
|
||||
requests = requests,
|
||||
window_seconds = WINDOW_SECOND
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def per_minute(requests: int = 100) -> RateLimitRule:
|
||||
return RateLimitRule(requests=requests, window_seconds=WINDOW_MINUTE)
|
||||
return RateLimitRule(
|
||||
requests = requests,
|
||||
window_seconds = WINDOW_MINUTE
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def per_hour(requests: int = 1000) -> RateLimitRule:
|
||||
return RateLimitRule(requests=requests, window_seconds=WINDOW_HOUR)
|
||||
return RateLimitRule(
|
||||
requests = requests,
|
||||
window_seconds = WINDOW_HOUR
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def per_day(requests: int = 10000) -> RateLimitRule:
|
||||
return RateLimitRule(requests=requests, window_seconds=WINDOW_DAY)
|
||||
return RateLimitRule(
|
||||
requests = requests,
|
||||
window_seconds = WINDOW_DAY
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def strict() -> RateLimitRule:
|
||||
return RateLimitRule(
|
||||
requests=STRICT_LIMIT_REQUESTS,
|
||||
window_seconds=STRICT_LIMIT_WINDOW,
|
||||
requests = STRICT_LIMIT_REQUESTS,
|
||||
window_seconds = STRICT_LIMIT_WINDOW,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -331,10 +353,10 @@ class ResultFactory:
|
|||
Create allowed result
|
||||
"""
|
||||
return RateLimitResult(
|
||||
allowed=True,
|
||||
limit=limit,
|
||||
remaining=remaining if remaining is not None else limit - 1,
|
||||
reset_after=reset_after,
|
||||
allowed = True,
|
||||
limit = limit,
|
||||
remaining = remaining if remaining is not None else limit - 1,
|
||||
reset_after = reset_after,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -347,11 +369,11 @@ class ResultFactory:
|
|||
Create denied result
|
||||
"""
|
||||
return RateLimitResult(
|
||||
allowed=False,
|
||||
limit=limit,
|
||||
remaining=0,
|
||||
reset_after=reset_after,
|
||||
retry_after=retry_after or reset_after,
|
||||
allowed = False,
|
||||
limit = limit,
|
||||
remaining = 0,
|
||||
reset_after = reset_after,
|
||||
retry_after = retry_after or reset_after,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -364,10 +386,10 @@ class ResultFactory:
|
|||
Create result near the limit
|
||||
"""
|
||||
return RateLimitResult(
|
||||
allowed=True,
|
||||
limit=limit,
|
||||
remaining=remaining,
|
||||
reset_after=reset_after,
|
||||
allowed = True,
|
||||
limit = limit,
|
||||
remaining = remaining,
|
||||
reset_after = reset_after,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -388,12 +410,12 @@ class KeyFactory:
|
|||
Create RateLimitKey with defaults
|
||||
"""
|
||||
return RateLimitKey(
|
||||
prefix=prefix,
|
||||
version=version,
|
||||
layer=layer,
|
||||
endpoint=endpoint,
|
||||
identifier=identifier,
|
||||
window=window,
|
||||
prefix = prefix,
|
||||
version = version,
|
||||
layer = layer,
|
||||
endpoint = endpoint,
|
||||
identifier = identifier,
|
||||
window = window,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -403,10 +425,10 @@ class KeyFactory:
|
|||
window: int = WINDOW_MINUTE,
|
||||
) -> RateLimitKey:
|
||||
return KeyFactory.create(
|
||||
layer=Layer.USER,
|
||||
endpoint=endpoint,
|
||||
identifier=identifier,
|
||||
window=window,
|
||||
layer = Layer.USER,
|
||||
endpoint = endpoint,
|
||||
identifier = identifier,
|
||||
window = window,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -415,19 +437,19 @@ class KeyFactory:
|
|||
window: int = WINDOW_MINUTE,
|
||||
) -> RateLimitKey:
|
||||
return KeyFactory.create(
|
||||
layer=Layer.ENDPOINT,
|
||||
endpoint=endpoint,
|
||||
identifier="global",
|
||||
window=window,
|
||||
layer = Layer.ENDPOINT,
|
||||
endpoint = endpoint,
|
||||
identifier = "global",
|
||||
window = window,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def global_key(window: int = WINDOW_MINUTE) -> RateLimitKey:
|
||||
return KeyFactory.create(
|
||||
layer=Layer.GLOBAL,
|
||||
endpoint="",
|
||||
identifier="global",
|
||||
window=window,
|
||||
layer = Layer.GLOBAL,
|
||||
endpoint = "",
|
||||
identifier = "global",
|
||||
window = window,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -449,10 +471,10 @@ class WindowStateFactory:
|
|||
current_window = int(time.time() // window_seconds)
|
||||
|
||||
return WindowState(
|
||||
current_count=current_count,
|
||||
previous_count=previous_count,
|
||||
current_window=current_window,
|
||||
window_seconds=window_seconds,
|
||||
current_count = current_count,
|
||||
previous_count = previous_count,
|
||||
current_window = current_window,
|
||||
window_seconds = window_seconds,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -462,8 +484,8 @@ class WindowStateFactory:
|
|||
@staticmethod
|
||||
def with_usage(current: int, previous: int = 0) -> WindowState:
|
||||
return WindowStateFactory.create(
|
||||
current_count=current,
|
||||
previous_count=previous,
|
||||
current_count = current,
|
||||
previous_count = previous,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -482,22 +504,25 @@ class TokenBucketStateFactory:
|
|||
Create TokenBucketState with defaults
|
||||
"""
|
||||
return TokenBucketState(
|
||||
tokens=tokens,
|
||||
last_refill=last_refill or time.time(),
|
||||
capacity=capacity,
|
||||
refill_rate=refill_rate,
|
||||
tokens = tokens,
|
||||
last_refill = last_refill or time.time(),
|
||||
capacity = capacity,
|
||||
refill_rate = refill_rate,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def full(capacity: int = 100) -> TokenBucketState:
|
||||
return TokenBucketStateFactory.create(
|
||||
tokens=float(capacity),
|
||||
capacity=capacity,
|
||||
tokens = float(capacity),
|
||||
capacity = capacity,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def empty(capacity: int = 100) -> TokenBucketState:
|
||||
return TokenBucketStateFactory.create(tokens=0.0, capacity=capacity)
|
||||
return TokenBucketStateFactory.create(
|
||||
tokens = 0.0,
|
||||
capacity = capacity
|
||||
)
|
||||
|
||||
|
||||
class DefenseContextFactory:
|
||||
|
|
@ -517,20 +542,20 @@ class DefenseContextFactory:
|
|||
Create DefenseContext with defaults
|
||||
"""
|
||||
return DefenseContext(
|
||||
fingerprint=fingerprint or FingerprintFactory.create(),
|
||||
endpoint=endpoint,
|
||||
method=method,
|
||||
is_authenticated=is_authenticated,
|
||||
reputation_score=reputation_score,
|
||||
request_count_last_minute=request_count_last_minute,
|
||||
fingerprint = fingerprint or FingerprintFactory.create(),
|
||||
endpoint = endpoint,
|
||||
method = method,
|
||||
is_authenticated = is_authenticated,
|
||||
reputation_score = reputation_score,
|
||||
request_count_last_minute = request_count_last_minute,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def authenticated(**kwargs: Any) -> DefenseContext:
|
||||
fp = FingerprintFactory.authenticated()
|
||||
return DefenseContextFactory.create(
|
||||
fingerprint=fp,
|
||||
is_authenticated=True,
|
||||
fingerprint = fp,
|
||||
is_authenticated = True,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
|
@ -540,8 +565,8 @@ class DefenseContextFactory:
|
|||
request_count: int = 500,
|
||||
) -> DefenseContext:
|
||||
return DefenseContextFactory.create(
|
||||
reputation_score=reputation_score,
|
||||
request_count_last_minute=request_count,
|
||||
reputation_score = reputation_score,
|
||||
request_count_last_minute = request_count,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -561,11 +586,11 @@ class CircuitStateFactory:
|
|||
Create CircuitState with defaults
|
||||
"""
|
||||
return CircuitState(
|
||||
is_open=is_open,
|
||||
failure_count=failure_count,
|
||||
last_failure_time=last_failure_time,
|
||||
half_open_requests=half_open_requests,
|
||||
total_requests_in_window=total_requests_in_window,
|
||||
is_open = is_open,
|
||||
failure_count = failure_count,
|
||||
last_failure_time = last_failure_time,
|
||||
half_open_requests = half_open_requests,
|
||||
total_requests_in_window = total_requests_in_window,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -575,9 +600,9 @@ class CircuitStateFactory:
|
|||
@staticmethod
|
||||
def open(failure_time: float | None = None) -> CircuitState:
|
||||
return CircuitStateFactory.create(
|
||||
is_open=True,
|
||||
failure_count=1,
|
||||
last_failure_time=failure_time or time.time(),
|
||||
is_open = True,
|
||||
failure_count = 1,
|
||||
last_failure_time = failure_time or time.time(),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -587,10 +612,10 @@ def storage_settings() -> StorageSettings:
|
|||
Create test storage settings (memory backend)
|
||||
"""
|
||||
return StorageSettings(
|
||||
REDIS_URL=None,
|
||||
MEMORY_MAX_KEYS=10000,
|
||||
MEMORY_CLEANUP_INTERVAL=60,
|
||||
FALLBACK_TO_MEMORY=True,
|
||||
REDIS_URL = None,
|
||||
MEMORY_MAX_KEYS = 10000,
|
||||
MEMORY_CLEANUP_INTERVAL = 60,
|
||||
FALLBACK_TO_MEMORY = True,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -600,17 +625,17 @@ def fingerprint_settings() -> FingerprintSettings:
|
|||
Create test fingerprint settings
|
||||
"""
|
||||
return FingerprintSettings(
|
||||
LEVEL=FingerprintLevel.NORMAL,
|
||||
USE_IP=True,
|
||||
USE_USER_AGENT=True,
|
||||
USE_ACCEPT_HEADERS=False,
|
||||
USE_HEADER_ORDER=False,
|
||||
USE_AUTH=True,
|
||||
USE_TLS=False,
|
||||
USE_GEO=False,
|
||||
IPV6_PREFIX_LENGTH=64,
|
||||
TRUSTED_PROXIES=[],
|
||||
TRUST_X_FORWARDED_FOR=False,
|
||||
LEVEL = FingerprintLevel.NORMAL,
|
||||
USE_IP = True,
|
||||
USE_USER_AGENT = True,
|
||||
USE_ACCEPT_HEADERS = False,
|
||||
USE_HEADER_ORDER = False,
|
||||
USE_AUTH = True,
|
||||
USE_TLS = False,
|
||||
USE_GEO = False,
|
||||
IPV6_PREFIX_LENGTH = 64,
|
||||
TRUSTED_PROXIES = [],
|
||||
TRUST_X_FORWARDED_FOR = False,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -620,15 +645,15 @@ def defense_settings() -> DefenseSettings:
|
|||
Create test defense settings
|
||||
"""
|
||||
return DefenseSettings(
|
||||
MODE=DefenseMode.ADAPTIVE,
|
||||
GLOBAL_LIMIT="50000/minute",
|
||||
CIRCUIT_THRESHOLD=CIRCUIT_THRESHOLD,
|
||||
CIRCUIT_WINDOW=CIRCUIT_WINDOW,
|
||||
CIRCUIT_RECOVERY_TIME=CIRCUIT_RECOVERY,
|
||||
ADAPTIVE_REDUCTION_FACTOR=0.5,
|
||||
ENDPOINT_LIMIT_MULTIPLIER=10,
|
||||
LOCKDOWN_ALLOW_AUTHENTICATED=True,
|
||||
LOCKDOWN_ALLOW_KNOWN_GOOD=True,
|
||||
MODE = DefenseMode.ADAPTIVE,
|
||||
GLOBAL_LIMIT = "50000/minute",
|
||||
CIRCUIT_THRESHOLD = CIRCUIT_THRESHOLD,
|
||||
CIRCUIT_WINDOW = CIRCUIT_WINDOW,
|
||||
CIRCUIT_RECOVERY_TIME = CIRCUIT_RECOVERY,
|
||||
ADAPTIVE_REDUCTION_FACTOR = 0.5,
|
||||
ENDPOINT_LIMIT_MULTIPLIER = 10,
|
||||
LOCKDOWN_ALLOW_AUTHENTICATED = True,
|
||||
LOCKDOWN_ALLOW_KNOWN_GOOD = True,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -642,21 +667,22 @@ def rate_limiter_settings(
|
|||
Create test rate limiter settings
|
||||
"""
|
||||
return RateLimiterSettings(
|
||||
ENABLED=True,
|
||||
ALGORITHM=Algorithm.SLIDING_WINDOW,
|
||||
DEFAULT_LIMIT="100/minute",
|
||||
DEFAULT_LIMITS=["100/minute", "1000/hour"],
|
||||
FAIL_OPEN=True,
|
||||
KEY_PREFIX=KEY_PREFIX,
|
||||
KEY_VERSION=KEY_VERSION,
|
||||
INCLUDE_HEADERS=True,
|
||||
LOG_VIOLATIONS=False,
|
||||
ENVIRONMENT="development",
|
||||
HTTP_420_MESSAGE="Enhance your calm",
|
||||
HTTP_420_DETAIL="Rate limit exceeded. Take a breather.",
|
||||
storage=storage_settings,
|
||||
fingerprint=fingerprint_settings,
|
||||
defense=defense_settings,
|
||||
ENABLED = True,
|
||||
ALGORITHM = Algorithm.SLIDING_WINDOW,
|
||||
DEFAULT_LIMIT = "100/minute",
|
||||
DEFAULT_LIMITS = ["100/minute",
|
||||
"1000/hour"],
|
||||
FAIL_OPEN = True,
|
||||
KEY_PREFIX = KEY_PREFIX,
|
||||
KEY_VERSION = KEY_VERSION,
|
||||
INCLUDE_HEADERS = True,
|
||||
LOG_VIOLATIONS = False,
|
||||
ENVIRONMENT = "development",
|
||||
HTTP_420_MESSAGE = "Enhance your calm",
|
||||
HTTP_420_DETAIL = "Rate limit exceeded. Take a breather.",
|
||||
storage = storage_settings,
|
||||
fingerprint = fingerprint_settings,
|
||||
defense = defense_settings,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -665,7 +691,7 @@ async def memory_storage() -> AsyncGenerator[MemoryStorage, None]:
|
|||
"""
|
||||
Create and manage MemoryStorage instance
|
||||
"""
|
||||
storage = MemoryStorage(max_keys=10000, cleanup_interval=60)
|
||||
storage = MemoryStorage(max_keys = 10000, cleanup_interval = 60)
|
||||
await storage.start_cleanup_task()
|
||||
yield storage
|
||||
await storage.close()
|
||||
|
|
@ -701,9 +727,9 @@ def ip_extractor() -> IPExtractor:
|
|||
Create IP extractor instance
|
||||
"""
|
||||
return IPExtractor(
|
||||
ipv6_prefix_length=64,
|
||||
trusted_proxies=[],
|
||||
trust_x_forwarded_for=False,
|
||||
ipv6_prefix_length = 64,
|
||||
trusted_proxies = [],
|
||||
trust_x_forwarded_for = False,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -712,7 +738,7 @@ def headers_extractor() -> HeadersExtractor:
|
|||
"""
|
||||
Create headers extractor instance
|
||||
"""
|
||||
return HeadersExtractor(use_header_order=False, hash_length=16)
|
||||
return HeadersExtractor(use_header_order = False, hash_length = 16)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -721,13 +747,13 @@ def auth_extractor() -> AuthExtractor:
|
|||
Create auth extractor instance
|
||||
"""
|
||||
return AuthExtractor(
|
||||
jwt_secret=None,
|
||||
jwt_algorithms=["HS256"],
|
||||
api_key_header="X-API-Key",
|
||||
api_key_query_param="api_key",
|
||||
session_cookie="session_id",
|
||||
hash_identifiers=True,
|
||||
hash_length=16,
|
||||
jwt_secret = None,
|
||||
jwt_algorithms = ["HS256"],
|
||||
api_key_header = "X-API-Key",
|
||||
api_key_query_param = "api_key",
|
||||
session_cookie = "session_id",
|
||||
hash_identifiers = True,
|
||||
hash_length = 16,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -741,10 +767,10 @@ def composite_fingerprinter(
|
|||
Create composite fingerprinter instance
|
||||
"""
|
||||
return CompositeFingerprinter(
|
||||
level=FingerprintLevel.NORMAL,
|
||||
ip_extractor=ip_extractor,
|
||||
headers_extractor=headers_extractor,
|
||||
auth_extractor=auth_extractor,
|
||||
level = FingerprintLevel.NORMAL,
|
||||
ip_extractor = ip_extractor,
|
||||
headers_extractor = headers_extractor,
|
||||
auth_extractor = auth_extractor,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -754,10 +780,10 @@ async def circuit_breaker() -> CircuitBreaker:
|
|||
Create circuit breaker instance
|
||||
"""
|
||||
return CircuitBreaker(
|
||||
threshold=CIRCUIT_THRESHOLD,
|
||||
window_seconds=CIRCUIT_WINDOW,
|
||||
recovery_time=CIRCUIT_RECOVERY,
|
||||
defense_mode=DefenseMode.ADAPTIVE,
|
||||
threshold = CIRCUIT_THRESHOLD,
|
||||
window_seconds = CIRCUIT_WINDOW,
|
||||
recovery_time = CIRCUIT_RECOVERY,
|
||||
defense_mode = DefenseMode.ADAPTIVE,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -765,13 +791,14 @@ async def circuit_breaker() -> CircuitBreaker:
|
|||
async def rate_limiter(
|
||||
rate_limiter_settings: RateLimiterSettings,
|
||||
memory_storage: MemoryStorage,
|
||||
) -> AsyncGenerator[RateLimiter, None]:
|
||||
) -> AsyncGenerator[RateLimiter,
|
||||
None]:
|
||||
"""
|
||||
Create and manage RateLimiter instance
|
||||
"""
|
||||
limiter = RateLimiter(
|
||||
settings=rate_limiter_settings,
|
||||
storage=memory_storage,
|
||||
settings = rate_limiter_settings,
|
||||
storage = memory_storage,
|
||||
)
|
||||
await limiter.init()
|
||||
yield limiter
|
||||
|
|
@ -788,9 +815,9 @@ async def layered_defense(
|
|||
Create layered defense instance
|
||||
"""
|
||||
return LayeredDefense(
|
||||
storage=memory_storage,
|
||||
settings=rate_limiter_settings,
|
||||
circuit_breaker=circuit_breaker,
|
||||
storage = memory_storage,
|
||||
settings = rate_limiter_settings,
|
||||
circuit_breaker = circuit_breaker,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -842,7 +869,7 @@ def create_test_app(
|
|||
"""
|
||||
Create a test FastAPI application
|
||||
"""
|
||||
app = FastAPI(title="Test API")
|
||||
app = FastAPI(title = "Test API")
|
||||
|
||||
@app.get("/")
|
||||
async def root() -> dict[str, str]:
|
||||
|
|
@ -867,8 +894,8 @@ def create_test_app(
|
|||
if with_middleware and limiter:
|
||||
app.add_middleware(
|
||||
RateLimitMiddleware,
|
||||
limiter=limiter,
|
||||
default_limit=default_limit,
|
||||
limiter = limiter,
|
||||
default_limit = default_limit,
|
||||
)
|
||||
|
||||
return app
|
||||
|
|
@ -889,19 +916,20 @@ async def test_app_with_limiter(
|
|||
"""
|
||||
Create a test app with rate limiting middleware
|
||||
"""
|
||||
app = create_test_app(limiter=rate_limiter, with_middleware=True)
|
||||
app = create_test_app(limiter = rate_limiter, with_middleware = True)
|
||||
set_global_limiter(rate_limiter)
|
||||
return app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def async_client(test_app: FastAPI) -> AsyncGenerator[AsyncClient, None]:
|
||||
async def async_client(test_app: FastAPI) -> AsyncGenerator[AsyncClient,
|
||||
None]:
|
||||
"""
|
||||
Create async HTTP client for testing
|
||||
"""
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=test_app),
|
||||
base_url="http://test",
|
||||
transport = ASGITransport(app = test_app),
|
||||
base_url = "http://test",
|
||||
) as client:
|
||||
yield client
|
||||
|
||||
|
|
@ -909,13 +937,14 @@ async def async_client(test_app: FastAPI) -> AsyncGenerator[AsyncClient, None]:
|
|||
@pytest.fixture
|
||||
async def rate_limited_client(
|
||||
test_app_with_limiter: FastAPI,
|
||||
) -> AsyncGenerator[AsyncClient, None]:
|
||||
) -> AsyncGenerator[AsyncClient,
|
||||
None]:
|
||||
"""
|
||||
Create async HTTP client with rate limiting
|
||||
"""
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=test_app_with_limiter),
|
||||
base_url="http://test",
|
||||
transport = ASGITransport(app = test_app_with_limiter),
|
||||
base_url = "http://test",
|
||||
) as client:
|
||||
yield client
|
||||
|
||||
|
|
@ -930,21 +959,24 @@ def sync_client(test_app: FastAPI) -> Generator[TestClient, None, None]:
|
|||
|
||||
|
||||
def assert_rate_limit_headers(
|
||||
headers: dict[str, str],
|
||||
headers: dict[str,
|
||||
str],
|
||||
expected_limit: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Assert rate limit headers are present and valid
|
||||
"""
|
||||
assert "RateLimit-Limit" in headers
|
||||
assert "RateLimit-Remaining" in headers
|
||||
assert "RateLimit-Reset" in headers
|
||||
lower_headers = {k.lower(): v for k, v in headers.items()}
|
||||
|
||||
assert "ratelimit-limit" in lower_headers
|
||||
assert "ratelimit-remaining" in lower_headers
|
||||
assert "ratelimit-reset" in lower_headers
|
||||
|
||||
if expected_limit is not None:
|
||||
assert int(headers["RateLimit-Limit"]) == expected_limit
|
||||
assert int(lower_headers["ratelimit-limit"]) == expected_limit
|
||||
|
||||
assert int(headers["RateLimit-Remaining"]) >= 0
|
||||
assert int(headers["RateLimit-Reset"]) >= 0
|
||||
assert int(lower_headers["ratelimit-remaining"]) >= 0
|
||||
assert int(lower_headers["ratelimit-reset"]) >= 0
|
||||
|
||||
|
||||
def assert_420_response(
|
||||
|
|
@ -957,7 +989,8 @@ def assert_420_response(
|
|||
assert response.status_code == HTTP_420_ENHANCE_YOUR_CALM
|
||||
|
||||
if check_headers:
|
||||
assert "Retry-After" in response.headers or "RateLimit-Reset" in response.headers
|
||||
lower_headers = {k.lower(): v for k, v in response.headers.items()}
|
||||
assert "retry-after" in lower_headers or "ratelimit-reset" in lower_headers
|
||||
|
||||
|
||||
async def exhaust_rate_limit(
|
||||
|
|
@ -971,9 +1004,9 @@ async def exhaust_rate_limit(
|
|||
"""
|
||||
for _ in range(limit):
|
||||
await storage.increment(
|
||||
key=key,
|
||||
window_seconds=window_seconds,
|
||||
limit=limit,
|
||||
key = key,
|
||||
window_seconds = window_seconds,
|
||||
limit = limit,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -59,9 +59,9 @@ class TestSlidingWindowAlgorithm:
|
|||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="test_key",
|
||||
rule=rule,
|
||||
storage = storage,
|
||||
key = "test_key",
|
||||
rule = rule,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
|
@ -87,7 +87,10 @@ class TestSlidingWindowAlgorithm:
|
|||
async def test_limit_exceeded(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
rule = RuleFactory.create(
|
||||
requests = 5,
|
||||
window_seconds = WINDOW_MINUTE
|
||||
)
|
||||
|
||||
for _ in range(5):
|
||||
result = await algo.check(storage, "limit_key", rule)
|
||||
|
|
@ -103,7 +106,10 @@ class TestSlidingWindowAlgorithm:
|
|||
async def test_different_keys_independent(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
rule = RuleFactory.create(
|
||||
requests = 5,
|
||||
window_seconds = WINDOW_MINUTE
|
||||
)
|
||||
|
||||
for _ in range(5):
|
||||
await algo.check(storage, "key_a", rule)
|
||||
|
|
@ -149,10 +155,10 @@ class TestSlidingWindowAlgorithm:
|
|||
fixed_time = 1000000.0
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="timestamp_key",
|
||||
rule=rule,
|
||||
timestamp=fixed_time,
|
||||
storage = storage,
|
||||
key = "timestamp_key",
|
||||
rule = rule,
|
||||
timestamp = fixed_time,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
|
@ -176,9 +182,9 @@ class TestTokenBucketAlgorithm:
|
|||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="bucket_test",
|
||||
rule=rule,
|
||||
storage = storage,
|
||||
key = "bucket_test",
|
||||
rule = rule,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
|
@ -190,7 +196,10 @@ class TestTokenBucketAlgorithm:
|
|||
async def test_burst_consumption(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=10, window_seconds=WINDOW_MINUTE)
|
||||
rule = RuleFactory.create(
|
||||
requests = 10,
|
||||
window_seconds = WINDOW_MINUTE
|
||||
)
|
||||
|
||||
for i in range(10):
|
||||
result = await algo.check(storage, "burst_key", rule)
|
||||
|
|
@ -203,7 +212,10 @@ class TestTokenBucketAlgorithm:
|
|||
async def test_bucket_exhausted(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
rule = RuleFactory.create(
|
||||
requests = 5,
|
||||
window_seconds = WINDOW_MINUTE
|
||||
)
|
||||
|
||||
for _ in range(5):
|
||||
await algo.check(storage, "exhaust_key", rule)
|
||||
|
|
@ -217,7 +229,10 @@ class TestTokenBucketAlgorithm:
|
|||
async def test_different_keys_independent(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
rule = RuleFactory.create(
|
||||
requests = 5,
|
||||
window_seconds = WINDOW_MINUTE
|
||||
)
|
||||
|
||||
for _ in range(5):
|
||||
await algo.check(storage, "bucket_a", rule)
|
||||
|
|
@ -272,9 +287,9 @@ class TestFixedWindowAlgorithm:
|
|||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="fixed_test",
|
||||
rule=rule,
|
||||
storage = storage,
|
||||
key = "fixed_test",
|
||||
rule = rule,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
|
@ -297,7 +312,10 @@ class TestFixedWindowAlgorithm:
|
|||
async def test_limit_exceeded(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
rule = RuleFactory.create(
|
||||
requests = 5,
|
||||
window_seconds = WINDOW_MINUTE
|
||||
)
|
||||
|
||||
for _ in range(5):
|
||||
result = await algo.check(storage, "limit_fixed", rule)
|
||||
|
|
@ -312,7 +330,10 @@ class TestFixedWindowAlgorithm:
|
|||
async def test_different_keys_independent(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
rule = RuleFactory.create(
|
||||
requests = 5,
|
||||
window_seconds = WINDOW_MINUTE
|
||||
)
|
||||
|
||||
for _ in range(5):
|
||||
await algo.check(storage, "fixed_a", rule)
|
||||
|
|
@ -358,10 +379,10 @@ class TestFixedWindowAlgorithm:
|
|||
fixed_time = 1000000.0
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="timestamp_fixed",
|
||||
rule=rule,
|
||||
timestamp=fixed_time,
|
||||
storage = storage,
|
||||
key = "timestamp_fixed",
|
||||
rule = rule,
|
||||
timestamp = fixed_time,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
|
@ -395,7 +416,10 @@ class TestAlgorithmComparison:
|
|||
TokenBucketAlgorithm(),
|
||||
FixedWindowAlgorithm(),
|
||||
]
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
rule = RuleFactory.create(
|
||||
requests = 5,
|
||||
window_seconds = WINDOW_MINUTE
|
||||
)
|
||||
|
||||
for algo in algorithms:
|
||||
storage = MemoryStorage()
|
||||
|
|
@ -408,6 +432,7 @@ class TestAlgorithmComparison:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_algorithms_have_correct_names(self) -> None:
|
||||
assert SlidingWindowAlgorithm().name == Algorithm.SLIDING_WINDOW.value
|
||||
assert SlidingWindowAlgorithm(
|
||||
).name == Algorithm.SLIDING_WINDOW.value
|
||||
assert TokenBucketAlgorithm().name == Algorithm.TOKEN_BUCKET.value
|
||||
assert FixedWindowAlgorithm().name == Algorithm.FIXED_WINDOW.value
|
||||
|
|
|
|||
|
|
@ -0,0 +1,536 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
test_fingerprinting.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from fastapi_420.fingerprinting.ip import IPExtractor
|
||||
from fastapi_420.fingerprinting.headers import HeadersExtractor
|
||||
from fastapi_420.fingerprinting.auth import AuthExtractor
|
||||
from fastapi_420.fingerprinting.composite import CompositeFingerprinter
|
||||
from fastapi_420.types import FingerprintLevel
|
||||
|
||||
from tests.conftest import (
|
||||
TEST_IP_V4,
|
||||
TEST_IP_V4_PRIVATE,
|
||||
TEST_IP_V6,
|
||||
TEST_IP_V6_NORMALIZED,
|
||||
TEST_IP_LOCALHOST,
|
||||
TEST_USER_AGENT,
|
||||
TEST_ACCEPT_LANGUAGE,
|
||||
TEST_ACCEPT_ENCODING,
|
||||
TEST_JWT_TOKEN,
|
||||
TEST_API_KEY,
|
||||
TEST_SESSION_ID,
|
||||
RequestFactory,
|
||||
)
|
||||
|
||||
|
||||
class TestIPExtractor:
|
||||
"""
|
||||
Tests for IP address extraction and normalization
|
||||
"""
|
||||
def test_extract_basic_ipv4(self) -> None:
|
||||
extractor = IPExtractor()
|
||||
request = RequestFactory.create(client_ip = TEST_IP_V4)
|
||||
raw_ip, normalized_ip = extractor.extract(request)
|
||||
|
||||
assert raw_ip == TEST_IP_V4
|
||||
assert normalized_ip == TEST_IP_V4
|
||||
|
||||
def test_extract_localhost(self) -> None:
|
||||
extractor = IPExtractor()
|
||||
request = RequestFactory.create(client_ip = TEST_IP_LOCALHOST)
|
||||
raw_ip, normalized_ip = extractor.extract(request)
|
||||
|
||||
assert raw_ip == TEST_IP_LOCALHOST
|
||||
assert normalized_ip == TEST_IP_LOCALHOST
|
||||
|
||||
def test_extract_private_ip(self) -> None:
|
||||
extractor = IPExtractor()
|
||||
request = RequestFactory.create(client_ip = TEST_IP_V4_PRIVATE)
|
||||
raw_ip, normalized_ip = extractor.extract(request)
|
||||
|
||||
assert raw_ip == TEST_IP_V4_PRIVATE
|
||||
assert normalized_ip == TEST_IP_V4_PRIVATE
|
||||
|
||||
def test_ipv6_normalization_to_64_prefix(self) -> None:
|
||||
extractor = IPExtractor(ipv6_prefix_length = 64)
|
||||
request = RequestFactory.create(client_ip = TEST_IP_V6)
|
||||
raw_ip, normalized_ip = extractor.extract(request)
|
||||
|
||||
assert raw_ip == TEST_IP_V6
|
||||
assert normalized_ip == TEST_IP_V6_NORMALIZED
|
||||
|
||||
def test_ipv6_custom_prefix_length(self) -> None:
|
||||
extractor = IPExtractor(ipv6_prefix_length = 48)
|
||||
request = RequestFactory.create(client_ip = TEST_IP_V6)
|
||||
raw_ip, normalized_ip = extractor.extract(request)
|
||||
|
||||
assert "::" in normalized_ip or normalized_ip != TEST_IP_V6
|
||||
|
||||
def test_ipv4_mapped_ipv6(self) -> None:
|
||||
extractor = IPExtractor()
|
||||
ipv4_mapped = "::ffff:192.168.1.100"
|
||||
request = RequestFactory.create(client_ip = ipv4_mapped)
|
||||
raw_ip, normalized_ip = extractor.extract(request)
|
||||
|
||||
assert normalized_ip == "192.168.1.100"
|
||||
|
||||
def test_invalid_ip_passthrough(self) -> None:
|
||||
extractor = IPExtractor()
|
||||
request = RequestFactory.create(client_ip = "invalid_ip")
|
||||
raw_ip, normalized_ip = extractor.extract(request)
|
||||
|
||||
assert raw_ip == "invalid_ip"
|
||||
assert normalized_ip == "invalid_ip"
|
||||
|
||||
def test_x_forwarded_for_disabled(self) -> None:
|
||||
extractor = IPExtractor(trust_x_forwarded_for = False)
|
||||
request = RequestFactory.with_forwarded_for(
|
||||
forwarded_ips = ["10.0.0.1",
|
||||
"192.168.1.1"],
|
||||
client_ip = TEST_IP_V4,
|
||||
)
|
||||
raw_ip, _ = extractor.extract(request)
|
||||
|
||||
assert raw_ip == TEST_IP_V4
|
||||
|
||||
def test_x_forwarded_for_enabled(self) -> None:
|
||||
extractor = IPExtractor(trust_x_forwarded_for = True)
|
||||
request = RequestFactory.with_forwarded_for(
|
||||
forwarded_ips = ["10.0.0.1",
|
||||
"192.168.1.1"],
|
||||
client_ip = TEST_IP_V4,
|
||||
)
|
||||
raw_ip, _ = extractor.extract(request)
|
||||
|
||||
assert raw_ip == "10.0.0.1"
|
||||
|
||||
def test_x_forwarded_for_with_trusted_proxies(self) -> None:
|
||||
extractor = IPExtractor(
|
||||
trust_x_forwarded_for = True,
|
||||
trusted_proxies = ["192.168.1.1"],
|
||||
)
|
||||
request = RequestFactory.with_forwarded_for(
|
||||
forwarded_ips = ["10.0.0.1",
|
||||
"192.168.1.1"],
|
||||
client_ip = "172.16.0.1",
|
||||
)
|
||||
raw_ip, _ = extractor.extract(request)
|
||||
|
||||
assert raw_ip == "10.0.0.1"
|
||||
|
||||
def test_x_real_ip_header(self) -> None:
|
||||
extractor = IPExtractor(trusted_proxies = ["127.0.0.1"])
|
||||
request = RequestFactory.with_forwarded_for(
|
||||
forwarded_ips = [],
|
||||
real_ip = "203.0.113.50",
|
||||
client_ip = "127.0.0.1",
|
||||
)
|
||||
raw_ip, _ = extractor.extract(request)
|
||||
|
||||
assert raw_ip == "203.0.113.50"
|
||||
|
||||
def test_is_ipv6(self) -> None:
|
||||
extractor = IPExtractor()
|
||||
assert extractor.is_ipv6(TEST_IP_V6) is True
|
||||
assert extractor.is_ipv6(TEST_IP_V4) is False
|
||||
assert extractor.is_ipv6("invalid") is False
|
||||
|
||||
def test_is_private(self) -> None:
|
||||
extractor = IPExtractor()
|
||||
assert extractor.is_private(TEST_IP_V4_PRIVATE) is True
|
||||
assert extractor.is_private(TEST_IP_LOCALHOST) is True
|
||||
assert extractor.is_private("8.8.8.8") is False
|
||||
assert extractor.is_private("invalid") is False
|
||||
|
||||
|
||||
class TestHeadersExtractor:
|
||||
"""
|
||||
Tests for HTTP header extraction and fingerprinting
|
||||
"""
|
||||
def test_extract_user_agent(self) -> None:
|
||||
extractor = HeadersExtractor()
|
||||
request = RequestFactory.create()
|
||||
user_agent = extractor.extract_user_agent(request)
|
||||
|
||||
assert user_agent == TEST_USER_AGENT
|
||||
|
||||
def test_extract_user_agent_missing(self) -> None:
|
||||
extractor = HeadersExtractor()
|
||||
request = RequestFactory.create(headers = {"user-agent": ""})
|
||||
user_agent = extractor.extract_user_agent(request)
|
||||
|
||||
assert user_agent == ""
|
||||
|
||||
def test_extract_accept_language(self) -> None:
|
||||
extractor = HeadersExtractor()
|
||||
request = RequestFactory.create()
|
||||
accept_lang = extractor.extract_accept_language(request)
|
||||
|
||||
assert accept_lang == TEST_ACCEPT_LANGUAGE
|
||||
|
||||
def test_extract_accept_encoding(self) -> None:
|
||||
extractor = HeadersExtractor()
|
||||
request = RequestFactory.create()
|
||||
accept_enc = extractor.extract_accept_encoding(request)
|
||||
|
||||
assert accept_enc == TEST_ACCEPT_ENCODING
|
||||
|
||||
def test_compute_headers_hash(self) -> None:
|
||||
extractor = HeadersExtractor(hash_length = 16)
|
||||
request = RequestFactory.create()
|
||||
headers_hash = extractor.compute_headers_hash(request)
|
||||
|
||||
assert len(headers_hash) == 16
|
||||
assert headers_hash.isalnum()
|
||||
|
||||
def test_headers_hash_deterministic(self) -> None:
|
||||
extractor = HeadersExtractor()
|
||||
request1 = RequestFactory.create()
|
||||
request2 = RequestFactory.create()
|
||||
|
||||
hash1 = extractor.compute_headers_hash(request1)
|
||||
hash2 = extractor.compute_headers_hash(request2)
|
||||
|
||||
assert hash1 == hash2
|
||||
|
||||
def test_headers_hash_different_headers(self) -> None:
|
||||
extractor = HeadersExtractor()
|
||||
request1 = RequestFactory.create(
|
||||
headers = {"user-agent": "Firefox/1.0"}
|
||||
)
|
||||
request2 = RequestFactory.create(
|
||||
headers = {"user-agent": "Chrome/1.0"}
|
||||
)
|
||||
|
||||
hash1 = extractor.compute_headers_hash(request1)
|
||||
hash2 = extractor.compute_headers_hash(request2)
|
||||
|
||||
assert hash1 != hash2
|
||||
|
||||
def test_headers_hash_with_order(self) -> None:
|
||||
extractor = HeadersExtractor(use_header_order = True)
|
||||
request = RequestFactory.create()
|
||||
headers_hash = extractor.compute_headers_hash(request)
|
||||
|
||||
assert len(headers_hash) == 16
|
||||
|
||||
def test_extract_all(self) -> None:
|
||||
extractor = HeadersExtractor()
|
||||
request = RequestFactory.create()
|
||||
data = extractor.extract_all(request)
|
||||
|
||||
assert "user_agent" in data
|
||||
assert "accept_language" in data
|
||||
assert "accept_encoding" in data
|
||||
assert "headers_hash" in data
|
||||
|
||||
|
||||
class TestAuthExtractor:
|
||||
"""
|
||||
Tests for authentication identifier extraction
|
||||
"""
|
||||
def test_extract_jwt_bearer_token(self) -> None:
|
||||
extractor = AuthExtractor(hash_identifiers = False)
|
||||
request = RequestFactory.with_auth(
|
||||
auth_type = "bearer",
|
||||
token = TEST_JWT_TOKEN
|
||||
)
|
||||
identifier = extractor.extract(request)
|
||||
|
||||
assert identifier == "user_123"
|
||||
|
||||
def test_extract_jwt_bearer_token_hashed(self) -> None:
|
||||
extractor = AuthExtractor(
|
||||
hash_identifiers = True,
|
||||
hash_length = 16
|
||||
)
|
||||
request = RequestFactory.with_auth(
|
||||
auth_type = "bearer",
|
||||
token = TEST_JWT_TOKEN
|
||||
)
|
||||
identifier = extractor.extract(request)
|
||||
|
||||
assert identifier is not None
|
||||
assert len(identifier) == 16
|
||||
assert identifier != "user_123"
|
||||
|
||||
def test_extract_api_key_header(self) -> None:
|
||||
extractor = AuthExtractor(hash_identifiers = False)
|
||||
request = RequestFactory.with_auth(
|
||||
auth_type = "api_key",
|
||||
token = TEST_API_KEY
|
||||
)
|
||||
identifier = extractor.extract(request)
|
||||
|
||||
assert identifier == TEST_API_KEY
|
||||
|
||||
def test_extract_api_key_query_param(self) -> None:
|
||||
extractor = AuthExtractor(hash_identifiers = False)
|
||||
request = RequestFactory.create(
|
||||
query_params = {"api_key": TEST_API_KEY}
|
||||
)
|
||||
identifier = extractor.extract(request)
|
||||
|
||||
assert identifier == TEST_API_KEY
|
||||
|
||||
def test_extract_session_cookie(self) -> None:
|
||||
extractor = AuthExtractor(hash_identifiers = False)
|
||||
request = RequestFactory.create(
|
||||
cookies = {"session_id": TEST_SESSION_ID}
|
||||
)
|
||||
identifier = extractor.extract(request)
|
||||
|
||||
assert identifier == TEST_SESSION_ID
|
||||
|
||||
def test_extract_priority_jwt_over_api_key(self) -> None:
|
||||
extractor = AuthExtractor(hash_identifiers = False)
|
||||
request = RequestFactory.with_auth(
|
||||
auth_type = "bearer",
|
||||
token = TEST_JWT_TOKEN
|
||||
)
|
||||
identifier = extractor.extract(request)
|
||||
|
||||
assert identifier == "user_123"
|
||||
|
||||
def test_extract_no_auth(self) -> None:
|
||||
extractor = AuthExtractor()
|
||||
request = RequestFactory.create()
|
||||
identifier = extractor.extract(request)
|
||||
|
||||
assert identifier is None
|
||||
|
||||
def test_is_authenticated_true(self) -> None:
|
||||
extractor = AuthExtractor()
|
||||
request = RequestFactory.with_auth(
|
||||
auth_type = "bearer",
|
||||
token = TEST_JWT_TOKEN
|
||||
)
|
||||
|
||||
assert extractor.is_authenticated(request) is True
|
||||
|
||||
def test_is_authenticated_false(self) -> None:
|
||||
extractor = AuthExtractor()
|
||||
request = RequestFactory.create()
|
||||
|
||||
assert extractor.is_authenticated(request) is False
|
||||
|
||||
def test_invalid_jwt_format(self) -> None:
|
||||
extractor = AuthExtractor(hash_identifiers = False)
|
||||
request = RequestFactory.with_auth(
|
||||
auth_type = "bearer",
|
||||
token = "invalid"
|
||||
)
|
||||
identifier = extractor.extract(request)
|
||||
|
||||
assert identifier is None
|
||||
|
||||
def test_jwt_without_sub_claim(self) -> None:
|
||||
extractor = AuthExtractor(hash_identifiers = False)
|
||||
token_no_sub = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjk5OTk5OTk5OTl9.signature"
|
||||
request = RequestFactory.with_auth(
|
||||
auth_type = "bearer",
|
||||
token = token_no_sub
|
||||
)
|
||||
identifier = extractor.extract(request)
|
||||
|
||||
assert identifier == "" or identifier is None
|
||||
|
||||
|
||||
class TestCompositeFingerprinter:
|
||||
"""
|
||||
Tests for composite fingerprint extraction
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_normal_level_extraction(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.NORMAL
|
||||
)
|
||||
request = RequestFactory.with_auth()
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
assert fp.ip == TEST_IP_V4
|
||||
assert fp.user_agent == TEST_USER_AGENT
|
||||
assert fp.auth_identifier is not None
|
||||
assert fp.accept_language is None
|
||||
assert fp.tls_fingerprint is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_relaxed_level_extraction(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.RELAXED
|
||||
)
|
||||
request = RequestFactory.with_auth()
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
assert fp.ip == TEST_IP_V4
|
||||
assert fp.user_agent is None
|
||||
assert fp.auth_identifier is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_level_extraction(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.STRICT
|
||||
)
|
||||
request = RequestFactory.with_auth()
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
assert fp.ip == TEST_IP_V4
|
||||
assert fp.user_agent == TEST_USER_AGENT
|
||||
assert fp.accept_language == TEST_ACCEPT_LANGUAGE
|
||||
assert fp.accept_encoding == TEST_ACCEPT_ENCODING
|
||||
assert fp.headers_hash is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_level_ip_only(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.CUSTOM,
|
||||
use_ip = True,
|
||||
use_user_agent = False,
|
||||
use_auth = False,
|
||||
)
|
||||
request = RequestFactory.with_auth()
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
assert fp.ip == TEST_IP_V4
|
||||
assert fp.user_agent is None
|
||||
assert fp.auth_identifier is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tls_fingerprint_extraction(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.STRICT
|
||||
)
|
||||
request = RequestFactory.create(
|
||||
headers = {"X-JA3-Fingerprint": "abc123hash"}
|
||||
)
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
assert fp.tls_fingerprint == "abc123hash"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_geo_asn_extraction(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.STRICT
|
||||
)
|
||||
request = RequestFactory.create(
|
||||
headers = {"X-Client-ASN": "AS12345"}
|
||||
)
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
assert fp.geo_asn == "AS12345"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloudflare_country_header(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.STRICT
|
||||
)
|
||||
request = RequestFactory.create(headers = {"CF-IPCountry": "US"})
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
assert fp.geo_asn == "US"
|
||||
|
||||
def test_is_authenticated_method(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.NORMAL
|
||||
)
|
||||
auth_request = RequestFactory.with_auth()
|
||||
anon_request = RequestFactory.create()
|
||||
|
||||
assert fingerprinter.is_authenticated(auth_request) is True
|
||||
assert fingerprinter.is_authenticated(anon_request) is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_from_settings(self) -> None:
|
||||
from fastapi_420.config import FingerprintSettings
|
||||
|
||||
settings = FingerprintSettings(
|
||||
LEVEL = FingerprintLevel.NORMAL,
|
||||
USE_IP = True,
|
||||
USE_USER_AGENT = True,
|
||||
USE_ACCEPT_HEADERS = False,
|
||||
USE_AUTH = True,
|
||||
IPV6_PREFIX_LENGTH = 64,
|
||||
)
|
||||
fingerprinter = CompositeFingerprinter.from_settings(settings)
|
||||
|
||||
request = RequestFactory.with_auth()
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
assert fp.ip is not None
|
||||
assert fp.user_agent is not None
|
||||
|
||||
|
||||
class TestFingerprintDataCompositeKey:
|
||||
"""
|
||||
Tests for composite key generation from fingerprint data
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_composite_key_relaxed_anonymous(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.RELAXED
|
||||
)
|
||||
request = RequestFactory.create()
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
key = fp.to_composite_key(FingerprintLevel.RELAXED)
|
||||
assert key == TEST_IP_V4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_composite_key_relaxed_authenticated(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.RELAXED
|
||||
)
|
||||
request = RequestFactory.with_auth()
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
key = fp.to_composite_key(FingerprintLevel.RELAXED)
|
||||
assert TEST_IP_V4 in key
|
||||
assert ":" in key
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_composite_key_normal(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.NORMAL
|
||||
)
|
||||
request = RequestFactory.with_auth()
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
key = fp.to_composite_key(FingerprintLevel.NORMAL)
|
||||
parts = key.split(":")
|
||||
|
||||
assert len(parts) == 3
|
||||
assert parts[0] == TEST_IP_V4
|
||||
assert TEST_USER_AGENT in parts[1]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_composite_key_strict(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.STRICT
|
||||
)
|
||||
request = RequestFactory.with_auth()
|
||||
fp = await fingerprinter.extract(request)
|
||||
|
||||
key = fp.to_composite_key(FingerprintLevel.STRICT)
|
||||
parts = key.split(":")
|
||||
|
||||
assert len(parts) == 8
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_composite_key_deterministic(self) -> None:
|
||||
fingerprinter = CompositeFingerprinter(
|
||||
level = FingerprintLevel.NORMAL
|
||||
)
|
||||
request = RequestFactory.with_auth()
|
||||
|
||||
fp1 = await fingerprinter.extract(request)
|
||||
fp2 = await fingerprinter.extract(request)
|
||||
|
||||
key1 = fp1.to_composite_key(FingerprintLevel.NORMAL)
|
||||
key2 = fp2.to_composite_key(FingerprintLevel.NORMAL)
|
||||
|
||||
assert key1 == key2
|
||||
|
|
@ -0,0 +1,647 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
test_integration.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from fastapi import Depends, FastAPI, Request
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
from fastapi_420.config import FingerprintSettings, RateLimiterSettings
|
||||
from fastapi_420.dependencies import (
|
||||
RateLimitDep,
|
||||
ScopedRateLimiter,
|
||||
create_rate_limit_dep,
|
||||
set_global_limiter,
|
||||
)
|
||||
from fastapi_420.exceptions import HTTP_420_ENHANCE_YOUR_CALM
|
||||
from fastapi_420.limiter import RateLimiter
|
||||
from fastapi_420.middleware import RateLimitMiddleware, SlowDownMiddleware
|
||||
from fastapi_420.storage import MemoryStorage
|
||||
from fastapi_420.types import Algorithm
|
||||
|
||||
from tests.conftest import (
|
||||
assert_420_response,
|
||||
assert_rate_limit_headers,
|
||||
)
|
||||
|
||||
|
||||
def create_app_with_middleware(
|
||||
limiter: RateLimiter,
|
||||
default_limit: str = "100/minute",
|
||||
exclude_paths: list[str] | None = None,
|
||||
) -> FastAPI:
|
||||
"""
|
||||
Create a FastAPI app with rate limiting middleware
|
||||
"""
|
||||
app = FastAPI(title = "Test App")
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {"message": "Hello World"}
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {"status": "healthy"}
|
||||
|
||||
@app.get("/api/test")
|
||||
async def test_endpoint():
|
||||
return {"endpoint": "test"}
|
||||
|
||||
@app.post("/api/test")
|
||||
async def test_endpoint_post():
|
||||
return {"created": True}
|
||||
|
||||
@app.get("/api/protected")
|
||||
async def protected():
|
||||
return {"protected": True}
|
||||
|
||||
app.add_middleware(
|
||||
RateLimitMiddleware,
|
||||
limiter = limiter,
|
||||
default_limit = default_limit,
|
||||
exclude_paths = exclude_paths,
|
||||
)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def create_app_with_decorator(limiter: RateLimiter) -> FastAPI:
|
||||
"""
|
||||
Create a FastAPI app using decorators for rate limiting
|
||||
"""
|
||||
app = FastAPI(title = "Test App")
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {"message": "Hello World"}
|
||||
|
||||
@app.get("/api/limited")
|
||||
@limiter.limit("5/minute")
|
||||
async def limited_endpoint(request: Request):
|
||||
return {"limited": True}
|
||||
|
||||
@app.get("/api/multi-limited")
|
||||
@limiter.limit("10/minute", "100/hour")
|
||||
async def multi_limited_endpoint(request: Request):
|
||||
return {"multi": True}
|
||||
|
||||
@app.get("/api/unlimited")
|
||||
async def unlimited_endpoint():
|
||||
return {"unlimited": True}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def create_app_with_dependency(limiter: RateLimiter) -> FastAPI:
|
||||
"""
|
||||
Create a FastAPI app using dependency injection for rate limiting
|
||||
"""
|
||||
app = FastAPI(title = "Test App")
|
||||
set_global_limiter(limiter)
|
||||
|
||||
rate_limit = create_rate_limit_dep("5/minute")
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {"message": "Hello World"}
|
||||
|
||||
@app.get("/api/limited", dependencies = [Depends(rate_limit)])
|
||||
async def limited_endpoint():
|
||||
return {"limited": True}
|
||||
|
||||
@app.get("/api/with-result")
|
||||
async def with_result_endpoint(
|
||||
result = Depends(RateLimitDep("10/minute"))
|
||||
):
|
||||
return {"remaining": result.remaining}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
class TestMiddlewareIntegration:
|
||||
"""
|
||||
Integration tests for RateLimitMiddleware
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_allows_requests_under_limit(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
settings = RateLimiterSettings(INCLUDE_HEADERS = True)
|
||||
limiter = RateLimiter(settings = settings, storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "100/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
response = await client.get("/api/test")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"endpoint": "test"}
|
||||
assert_rate_limit_headers(
|
||||
dict(response.headers),
|
||||
expected_limit = 100
|
||||
)
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_returns_420_when_exceeded(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "5/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(5):
|
||||
response = await client.get("/api/test")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await client.get("/api/test")
|
||||
assert_420_response(response)
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_excludes_health_endpoint(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "1/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(10):
|
||||
response = await client.get("/health")
|
||||
assert response.status_code == 200
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_excludes_custom_paths(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(
|
||||
limiter,
|
||||
"1/minute",
|
||||
exclude_paths = ["/api/test"],
|
||||
)
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(10):
|
||||
response = await client.get("/api/test")
|
||||
assert response.status_code == 200
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_different_endpoints_independent(
|
||||
self
|
||||
) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "3/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(3):
|
||||
await client.get("/api/test")
|
||||
|
||||
response1 = await client.get("/api/test")
|
||||
response2 = await client.get("/api/protected")
|
||||
|
||||
assert response1.status_code == HTTP_420_ENHANCE_YOUR_CALM
|
||||
assert response2.status_code == 200
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_post_and_get_independent_limits(
|
||||
self
|
||||
) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "3/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(3):
|
||||
await client.get("/api/test")
|
||||
|
||||
get_response = await client.get("/api/test")
|
||||
assert get_response.status_code == HTTP_420_ENHANCE_YOUR_CALM
|
||||
|
||||
post_response = await client.post("/api/test")
|
||||
assert post_response.status_code == 200
|
||||
|
||||
await limiter.close()
|
||||
|
||||
|
||||
class TestDecoratorIntegration:
|
||||
"""
|
||||
Integration tests for @limiter.limit() decorator
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_allows_requests_under_limit(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_decorator(limiter)
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
response = await client.get("/api/limited")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"limited": True}
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_returns_420_when_exceeded(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_decorator(limiter)
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(5):
|
||||
response = await client.get("/api/limited")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await client.get("/api/limited")
|
||||
assert_420_response(response)
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_unlimited_endpoint_not_affected(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_decorator(limiter)
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(100):
|
||||
response = await client.get("/api/unlimited")
|
||||
assert response.status_code == 200
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_different_endpoints_independent(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_decorator(limiter)
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(5):
|
||||
await client.get("/api/limited")
|
||||
|
||||
response1 = await client.get("/api/limited")
|
||||
response2 = await client.get("/api/multi-limited")
|
||||
|
||||
assert response1.status_code == HTTP_420_ENHANCE_YOUR_CALM
|
||||
assert response2.status_code == 200
|
||||
|
||||
await limiter.close()
|
||||
|
||||
|
||||
class TestDependencyIntegration:
|
||||
"""
|
||||
Integration tests for dependency injection
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_dependency_allows_requests_under_limit(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_dependency(limiter)
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
response = await client.get("/api/limited")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"limited": True}
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dependency_returns_420_when_exceeded(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_dependency(limiter)
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(5):
|
||||
response = await client.get("/api/limited")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await client.get("/api/limited")
|
||||
assert_420_response(response)
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dependency_with_result_access(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_dependency(limiter)
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
response = await client.get("/api/with-result")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "remaining" in data
|
||||
assert data["remaining"] == 9
|
||||
|
||||
await limiter.close()
|
||||
|
||||
|
||||
class TestScopedRateLimiterIntegration:
|
||||
"""
|
||||
Integration tests for ScopedRateLimiter
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_limiter_applies_default_rules(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
set_global_limiter(limiter)
|
||||
|
||||
app = FastAPI()
|
||||
scoped = ScopedRateLimiter(
|
||||
prefix = "/api",
|
||||
default_rules = ["5/minute"],
|
||||
)
|
||||
|
||||
@app.get("/api/endpoint", dependencies = [Depends(scoped)])
|
||||
async def endpoint():
|
||||
return {"ok": True}
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(5):
|
||||
response = await client.get("/api/endpoint")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await client.get("/api/endpoint")
|
||||
assert_420_response(response)
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_limiter_endpoint_specific_rules(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
set_global_limiter(limiter)
|
||||
|
||||
app = FastAPI()
|
||||
scoped = ScopedRateLimiter(
|
||||
prefix = "/api",
|
||||
default_rules = ["100/minute"],
|
||||
endpoint_rules = {
|
||||
"GET:/api/strict": ["2/minute"],
|
||||
},
|
||||
)
|
||||
|
||||
@app.get("/api/normal", dependencies = [Depends(scoped)])
|
||||
async def normal():
|
||||
return {"type": "normal"}
|
||||
|
||||
@app.get("/api/strict", dependencies = [Depends(scoped)])
|
||||
async def strict():
|
||||
return {"type": "strict"}
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(2):
|
||||
response = await client.get("/api/strict")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await client.get("/api/strict")
|
||||
assert_420_response(response)
|
||||
|
||||
response = await client.get("/api/normal")
|
||||
assert response.status_code == 200
|
||||
|
||||
await limiter.close()
|
||||
|
||||
|
||||
class TestSlowDownMiddlewareIntegration:
|
||||
"""
|
||||
Integration tests for SlowDownMiddleware
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_slowdown_middleware_allows_requests(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
settings = RateLimiterSettings(INCLUDE_HEADERS = True)
|
||||
limiter = RateLimiter(settings = settings, storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
@app.get("/api/test")
|
||||
async def test_endpoint():
|
||||
return {"ok": True}
|
||||
|
||||
app.add_middleware(
|
||||
SlowDownMiddleware,
|
||||
limiter = limiter,
|
||||
threshold_limit = "100/minute",
|
||||
max_delay_seconds = 1.0,
|
||||
)
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
response = await client.get("/api/test")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"ok": True}
|
||||
|
||||
await limiter.close()
|
||||
|
||||
|
||||
class TestConcurrentRequests:
|
||||
"""
|
||||
Integration tests for concurrent request handling
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_requests_all_counted(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "100/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
tasks = [client.get("/api/test") for _ in range(50)]
|
||||
responses = await asyncio.gather(*tasks)
|
||||
|
||||
success_count = sum(
|
||||
1 for r in responses if r.status_code == 200
|
||||
)
|
||||
assert success_count == 50
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_requests_enforce_limit(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "10/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
tasks = [client.get("/api/test") for _ in range(20)]
|
||||
responses = await asyncio.gather(*tasks)
|
||||
|
||||
success_count = sum(
|
||||
1 for r in responses if r.status_code == 200
|
||||
)
|
||||
blocked_count = sum(
|
||||
1 for r in responses
|
||||
if r.status_code == HTTP_420_ENHANCE_YOUR_CALM
|
||||
)
|
||||
|
||||
assert success_count == 10
|
||||
assert blocked_count == 10
|
||||
|
||||
await limiter.close()
|
||||
|
||||
|
||||
class TestMultipleClients:
|
||||
"""
|
||||
Integration tests for multiple clients (different IPs)
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_ips_independent_limits(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
fingerprint_settings = FingerprintSettings(
|
||||
TRUST_X_FORWARDED_FOR = True
|
||||
)
|
||||
settings = RateLimiterSettings(fingerprint = fingerprint_settings)
|
||||
limiter = RateLimiter(settings = settings, storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "5/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(5):
|
||||
await client.get(
|
||||
"/api/test",
|
||||
headers = {"X-Forwarded-For": "192.168.1.1"}
|
||||
)
|
||||
|
||||
response1 = await client.get(
|
||||
"/api/test",
|
||||
headers = {"X-Forwarded-For": "192.168.1.1"}
|
||||
)
|
||||
|
||||
response2 = await client.get(
|
||||
"/api/test",
|
||||
headers = {"X-Forwarded-For": "192.168.1.2"}
|
||||
)
|
||||
|
||||
assert response1.status_code == HTTP_420_ENHANCE_YOUR_CALM
|
||||
assert response2.status_code == 200
|
||||
|
||||
await limiter.close()
|
||||
|
||||
|
||||
class TestAlgorithmIntegration:
|
||||
"""
|
||||
Integration tests for different algorithms
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_sliding_window_algorithm(self) -> None:
|
||||
settings = RateLimiterSettings(
|
||||
ALGORITHM = Algorithm.SLIDING_WINDOW
|
||||
)
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(settings = settings, storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "5/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(5):
|
||||
response = await client.get("/api/test")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await client.get("/api/test")
|
||||
assert_420_response(response)
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_bucket_algorithm(self) -> None:
|
||||
settings = RateLimiterSettings(ALGORITHM = Algorithm.TOKEN_BUCKET)
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(settings = settings, storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "5/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(5):
|
||||
response = await client.get("/api/test")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await client.get("/api/test")
|
||||
assert_420_response(response)
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fixed_window_algorithm(self) -> None:
|
||||
settings = RateLimiterSettings(ALGORITHM = Algorithm.FIXED_WINDOW)
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(settings = settings, storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
app = create_app_with_middleware(limiter, "5/minute")
|
||||
|
||||
async with AsyncClient(transport = ASGITransport(app = app),
|
||||
base_url = "http://test") as client:
|
||||
for _ in range(5):
|
||||
response = await client.get("/api/test")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await client.get("/api/test")
|
||||
assert_420_response(response)
|
||||
|
||||
await limiter.close()
|
||||
|
|
@ -0,0 +1,443 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
test_limiter.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from fastapi_420.config import RateLimiterSettings
|
||||
from fastapi_420.exceptions import EnhanceYourCalm, HTTP_420_ENHANCE_YOUR_CALM
|
||||
from fastapi_420.limiter import RateLimiter
|
||||
from fastapi_420.storage import MemoryStorage
|
||||
from fastapi_420.types import Algorithm, FingerprintLevel
|
||||
|
||||
from tests.conftest import (
|
||||
WINDOW_MINUTE,
|
||||
DEFAULT_LIMIT_REQUESTS,
|
||||
RequestFactory,
|
||||
)
|
||||
|
||||
|
||||
class TestRateLimiterInit:
|
||||
"""
|
||||
Tests for RateLimiter initialization
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_with_defaults(self) -> None:
|
||||
limiter = RateLimiter()
|
||||
assert limiter.is_initialized is False
|
||||
await limiter.init()
|
||||
assert limiter.is_initialized is True
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_with_custom_settings(self) -> None:
|
||||
settings = RateLimiterSettings(
|
||||
ALGORITHM = Algorithm.TOKEN_BUCKET,
|
||||
DEFAULT_LIMIT = "50/minute",
|
||||
)
|
||||
limiter = RateLimiter(settings = settings)
|
||||
await limiter.init()
|
||||
|
||||
assert limiter.settings.ALGORITHM == Algorithm.TOKEN_BUCKET
|
||||
assert limiter.settings.DEFAULT_LIMIT == "50/minute"
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_with_provided_storage(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
assert limiter.is_initialized is True
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_idempotent(self) -> None:
|
||||
limiter = RateLimiter()
|
||||
await limiter.init()
|
||||
await limiter.init()
|
||||
await limiter.init()
|
||||
|
||||
assert limiter.is_initialized is True
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_resets_initialized(self) -> None:
|
||||
limiter = RateLimiter()
|
||||
await limiter.init()
|
||||
assert limiter.is_initialized is True
|
||||
|
||||
await limiter.close()
|
||||
assert limiter.is_initialized is False
|
||||
|
||||
|
||||
class TestRateLimiterCheck:
|
||||
"""
|
||||
Tests for RateLimiter.check() method
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_allows_first_request(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
request = RequestFactory.create()
|
||||
result = await limiter.check(
|
||||
request,
|
||||
"100/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 99
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_multiple_rules(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
request = RequestFactory.create()
|
||||
result = await limiter.check(
|
||||
request,
|
||||
"100/minute",
|
||||
"1000/hour",
|
||||
raise_on_limit = False,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_uses_default_rules_when_none_provided(
|
||||
self
|
||||
) -> None:
|
||||
settings = RateLimiterSettings(
|
||||
DEFAULT_LIMITS = ["50/minute"],
|
||||
)
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(settings = settings, storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
request = RequestFactory.create()
|
||||
result = await limiter.check(request, raise_on_limit = False)
|
||||
|
||||
assert result.allowed is True
|
||||
assert result.limit == 50
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_raises_on_limit(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
request = RequestFactory.create()
|
||||
|
||||
for _ in range(5):
|
||||
await limiter.check(
|
||||
request,
|
||||
"5/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
with pytest.raises(EnhanceYourCalm) as exc_info:
|
||||
await limiter.check(request, "5/minute", raise_on_limit = True)
|
||||
|
||||
assert exc_info.value.status_code == HTTP_420_ENHANCE_YOUR_CALM
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_returns_result_when_not_raising(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
request = RequestFactory.create()
|
||||
|
||||
for _ in range(5):
|
||||
await limiter.check(
|
||||
request,
|
||||
"5/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
result = await limiter.check(
|
||||
request,
|
||||
"5/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
assert result.allowed is False
|
||||
assert result.remaining == 0
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_different_endpoints_independent(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
request1 = RequestFactory.create(path = "/api/endpoint1")
|
||||
request2 = RequestFactory.create(path = "/api/endpoint2")
|
||||
|
||||
for _ in range(5):
|
||||
await limiter.check(
|
||||
request1,
|
||||
"5/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
result1 = await limiter.check(
|
||||
request1,
|
||||
"5/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
result2 = await limiter.check(
|
||||
request2,
|
||||
"5/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
assert result1.allowed is False
|
||||
assert result2.allowed is True
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_different_users_independent(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
request1 = RequestFactory.create(client_ip = "192.168.1.1")
|
||||
request2 = RequestFactory.create(client_ip = "192.168.1.2")
|
||||
|
||||
for _ in range(5):
|
||||
await limiter.check(
|
||||
request1,
|
||||
"5/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
result1 = await limiter.check(
|
||||
request1,
|
||||
"5/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
result2 = await limiter.check(
|
||||
request2,
|
||||
"5/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
assert result1.allowed is False
|
||||
assert result2.allowed is True
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_with_custom_key_func(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
def custom_key(request):
|
||||
return "shared_key"
|
||||
|
||||
request1 = RequestFactory.create(client_ip = "192.168.1.1")
|
||||
request2 = RequestFactory.create(client_ip = "192.168.1.2")
|
||||
|
||||
for _ in range(5):
|
||||
await limiter.check(
|
||||
request1,
|
||||
"5/minute",
|
||||
key_func = custom_key,
|
||||
raise_on_limit = False,
|
||||
)
|
||||
|
||||
result = await limiter.check(
|
||||
request2,
|
||||
"5/minute",
|
||||
key_func = custom_key,
|
||||
raise_on_limit = False,
|
||||
)
|
||||
|
||||
assert result.allowed is False
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_auto_initializes(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
|
||||
assert limiter.is_initialized is False
|
||||
|
||||
request = RequestFactory.create()
|
||||
result = await limiter.check(
|
||||
request,
|
||||
"100/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
assert limiter.is_initialized is True
|
||||
assert result.allowed is True
|
||||
|
||||
await limiter.close()
|
||||
|
||||
|
||||
class TestRateLimiterDecorator:
|
||||
"""
|
||||
Tests for RateLimiter.limit() decorator
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_basic_usage(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
@limiter.limit("100/minute")
|
||||
async def endpoint(request):
|
||||
return {"success": True}
|
||||
|
||||
request = RequestFactory.create()
|
||||
result = await endpoint(request)
|
||||
|
||||
assert result == {"success": True}
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_enforces_limit(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
@limiter.limit("5/minute")
|
||||
async def endpoint(request):
|
||||
return {"success": True}
|
||||
|
||||
request = RequestFactory.create()
|
||||
|
||||
for _ in range(5):
|
||||
await endpoint(request)
|
||||
|
||||
with pytest.raises(EnhanceYourCalm):
|
||||
await endpoint(request)
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_multiple_rules(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
@limiter.limit("100/minute", "1000/hour")
|
||||
async def endpoint(request):
|
||||
return {"success": True}
|
||||
|
||||
request = RequestFactory.create()
|
||||
result = await endpoint(request)
|
||||
|
||||
assert result == {"success": True}
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_without_request_skips(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
@limiter.limit("5/minute")
|
||||
async def endpoint():
|
||||
return {"success": True}
|
||||
|
||||
result = await endpoint()
|
||||
assert result == {"success": True}
|
||||
|
||||
await limiter.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decorator_request_in_kwargs(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limiter = RateLimiter(storage = storage)
|
||||
await limiter.init()
|
||||
|
||||
@limiter.limit("100/minute")
|
||||
async def endpoint(data: str, request = None):
|
||||
return {"data": data}
|
||||
|
||||
request = RequestFactory.create()
|
||||
result = await endpoint("test", request = request)
|
||||
|
||||
assert result == {"data": "test"}
|
||||
|
||||
await limiter.close()
|
||||
|
||||
|
||||
class TestRateLimiterFailOpen:
|
||||
"""
|
||||
Tests for fail-open behavior
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_allows_requests_on_storage_failure(
|
||||
self
|
||||
) -> None:
|
||||
settings = RateLimiterSettings(FAIL_OPEN = True)
|
||||
storage = MemoryStorage()
|
||||
await storage.close()
|
||||
|
||||
limiter = RateLimiter(settings = settings, storage = storage)
|
||||
limiter._initialized = True
|
||||
limiter._storage = storage
|
||||
limiter._fallback_storage = None
|
||||
|
||||
request = RequestFactory.create()
|
||||
result = await limiter.check(
|
||||
request,
|
||||
"100/minute",
|
||||
raise_on_limit = False
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
||||
|
||||
class TestRateLimiterSettings:
|
||||
"""
|
||||
Tests for settings access
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_settings_property(self) -> None:
|
||||
settings = RateLimiterSettings(
|
||||
HTTP_420_MESSAGE = "Custom message",
|
||||
HTTP_420_DETAIL = "Custom detail",
|
||||
)
|
||||
limiter = RateLimiter(settings = settings)
|
||||
|
||||
assert limiter.settings.HTTP_420_MESSAGE == "Custom message"
|
||||
assert limiter.settings.HTTP_420_DETAIL == "Custom detail"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_initialized_property(self) -> None:
|
||||
limiter = RateLimiter()
|
||||
assert limiter.is_initialized is False
|
||||
|
||||
await limiter.init()
|
||||
assert limiter.is_initialized is True
|
||||
|
||||
await limiter.close()
|
||||
assert limiter.is_initialized is False
|
||||
|
|
@ -32,15 +32,15 @@ class TestMemoryStorageBasic:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_storage_custom_settings(self) -> None:
|
||||
storage = MemoryStorage(max_keys=5000, cleanup_interval=30)
|
||||
storage = MemoryStorage(max_keys = 5000, cleanup_interval = 30)
|
||||
assert storage.max_keys == 5000
|
||||
assert storage.cleanup_interval == 30
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_from_settings(self) -> None:
|
||||
settings = StorageSettings(
|
||||
MEMORY_MAX_KEYS=2000,
|
||||
MEMORY_CLEANUP_INTERVAL=120,
|
||||
MEMORY_MAX_KEYS = 2000,
|
||||
MEMORY_CLEANUP_INTERVAL = 120,
|
||||
)
|
||||
storage = MemoryStorage.from_settings(settings)
|
||||
assert storage.max_keys == 2000
|
||||
|
|
@ -74,9 +74,9 @@ class TestMemoryStorageSlidingWindow:
|
|||
async def test_increment_first_request(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
result = await storage.increment(
|
||||
key="test",
|
||||
window_seconds=WINDOW_MINUTE,
|
||||
limit=100,
|
||||
key = "test",
|
||||
window_seconds = WINDOW_MINUTE,
|
||||
limit = 100,
|
||||
)
|
||||
assert result.allowed is True
|
||||
assert result.limit == 100
|
||||
|
|
@ -119,10 +119,10 @@ class TestMemoryStorageSlidingWindow:
|
|||
fixed_time = 1000000.0
|
||||
|
||||
result = await storage.increment(
|
||||
key="timestamp_test",
|
||||
window_seconds=WINDOW_MINUTE,
|
||||
limit=100,
|
||||
timestamp=fixed_time,
|
||||
key = "timestamp_test",
|
||||
window_seconds = WINDOW_MINUTE,
|
||||
limit = 100,
|
||||
timestamp = fixed_time,
|
||||
)
|
||||
assert result.allowed is True
|
||||
|
||||
|
|
@ -131,7 +131,10 @@ class TestMemoryStorageSlidingWindow:
|
|||
@pytest.mark.asyncio
|
||||
async def test_get_window_state_empty(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
state = await storage.get_window_state("nonexistent", WINDOW_MINUTE)
|
||||
state = await storage.get_window_state(
|
||||
"nonexistent",
|
||||
WINDOW_MINUTE
|
||||
)
|
||||
assert state.current_count == 0
|
||||
assert state.previous_count == 0
|
||||
await storage.close()
|
||||
|
|
@ -165,16 +168,16 @@ class TestMemoryStorageSlidingWindow:
|
|||
storage._windows[prev_key] = storage._windows.__class__().__class__
|
||||
from fastapi_420.storage.memory import WindowEntry
|
||||
storage._windows[prev_key] = WindowEntry(
|
||||
count=50,
|
||||
window_start=previous_window,
|
||||
expires_at=base_time + window * 2,
|
||||
count = 50,
|
||||
window_start = previous_window,
|
||||
expires_at = base_time + window * 2,
|
||||
)
|
||||
|
||||
result = await storage.increment(
|
||||
key=key,
|
||||
window_seconds=window,
|
||||
limit=limit,
|
||||
timestamp=base_time + 1.0,
|
||||
key = key,
|
||||
window_seconds = window,
|
||||
limit = limit,
|
||||
timestamp = base_time + 1.0,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
|
@ -189,10 +192,10 @@ class TestMemoryStorageTokenBucket:
|
|||
async def test_consume_token_first_request(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
result = await storage.consume_token(
|
||||
key="bucket_test",
|
||||
capacity=100,
|
||||
refill_rate=1.67,
|
||||
tokens_to_consume=1,
|
||||
key = "bucket_test",
|
||||
capacity = 100,
|
||||
refill_rate = 1.67,
|
||||
tokens_to_consume = 1,
|
||||
)
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 99
|
||||
|
|
@ -205,9 +208,9 @@ class TestMemoryStorageTokenBucket:
|
|||
|
||||
for i in range(10):
|
||||
result = await storage.consume_token(
|
||||
key=key,
|
||||
capacity=100,
|
||||
refill_rate=1.67,
|
||||
key = key,
|
||||
capacity = 100,
|
||||
refill_rate = 1.67,
|
||||
)
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 100 - (i + 1)
|
||||
|
|
@ -222,16 +225,16 @@ class TestMemoryStorageTokenBucket:
|
|||
|
||||
for _ in range(capacity):
|
||||
result = await storage.consume_token(
|
||||
key=key,
|
||||
capacity=capacity,
|
||||
refill_rate=1.0,
|
||||
key = key,
|
||||
capacity = capacity,
|
||||
refill_rate = 1.0,
|
||||
)
|
||||
assert result.allowed is True
|
||||
|
||||
result = await storage.consume_token(
|
||||
key=key,
|
||||
capacity=capacity,
|
||||
refill_rate=1.0,
|
||||
key = key,
|
||||
capacity = capacity,
|
||||
refill_rate = 1.0,
|
||||
)
|
||||
assert result.allowed is False
|
||||
assert result.retry_after is not None
|
||||
|
|
@ -270,7 +273,11 @@ class TestMemoryStorageTokenBucket:
|
|||
storage = MemoryStorage()
|
||||
key = "bucket_state_test"
|
||||
|
||||
await storage.consume_token(key, capacity=100, refill_rate=1.67)
|
||||
await storage.consume_token(
|
||||
key,
|
||||
capacity = 100,
|
||||
refill_rate = 1.67
|
||||
)
|
||||
|
||||
state = await storage.get_token_bucket_state(key)
|
||||
assert state is not None
|
||||
|
|
@ -286,7 +293,7 @@ class TestMemoryStorageMaxKeys:
|
|||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_keys_eviction(self) -> None:
|
||||
storage = MemoryStorage(max_keys=5)
|
||||
storage = MemoryStorage(max_keys = 5)
|
||||
|
||||
for i in range(10):
|
||||
await storage.increment(f"key_{i}", WINDOW_MINUTE, 100)
|
||||
|
|
@ -297,7 +304,7 @@ class TestMemoryStorageMaxKeys:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lru_eviction_order(self) -> None:
|
||||
storage = MemoryStorage(max_keys=3)
|
||||
storage = MemoryStorage(max_keys = 3)
|
||||
|
||||
await storage.increment("key_a", WINDOW_MINUTE, 100)
|
||||
await storage.increment("key_b", WINDOW_MINUTE, 100)
|
||||
|
|
@ -317,18 +324,18 @@ class TestMemoryStorageCleanup:
|
|||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_entries(self) -> None:
|
||||
storage = MemoryStorage(cleanup_interval=60)
|
||||
storage = MemoryStorage(cleanup_interval = 60)
|
||||
|
||||
from fastapi_420.storage.memory import WindowEntry
|
||||
storage._windows["expired_key"] = WindowEntry(
|
||||
count=10,
|
||||
window_start=1,
|
||||
expires_at=time.time() - 100,
|
||||
count = 10,
|
||||
window_start = 1,
|
||||
expires_at = time.time() - 100,
|
||||
)
|
||||
storage._windows["valid_key"] = WindowEntry(
|
||||
count=10,
|
||||
window_start=1,
|
||||
expires_at=time.time() + 100,
|
||||
count = 10,
|
||||
window_start = 1,
|
||||
expires_at = time.time() + 100,
|
||||
)
|
||||
|
||||
await storage._cleanup_expired()
|
||||
|
|
@ -340,7 +347,7 @@ class TestMemoryStorageCleanup:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_task_starts(self) -> None:
|
||||
storage = MemoryStorage(cleanup_interval=1)
|
||||
storage = MemoryStorage(cleanup_interval = 1)
|
||||
await storage.start_cleanup_task()
|
||||
|
||||
assert storage._cleanup_task is not None
|
||||
|
|
@ -350,7 +357,7 @@ class TestMemoryStorageCleanup:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_task_stops_on_close(self) -> None:
|
||||
storage = MemoryStorage(cleanup_interval=1)
|
||||
storage = MemoryStorage(cleanup_interval = 1)
|
||||
await storage.start_cleanup_task()
|
||||
|
||||
task = storage._cleanup_task
|
||||
|
|
@ -364,14 +371,14 @@ class TestStorageFactory:
|
|||
Tests for create_storage factory function
|
||||
"""
|
||||
def test_create_memory_storage_no_redis(self) -> None:
|
||||
settings = StorageSettings(REDIS_URL=None)
|
||||
settings = StorageSettings(REDIS_URL = None)
|
||||
storage = create_storage(settings)
|
||||
assert isinstance(storage, MemoryStorage)
|
||||
|
||||
def test_create_memory_storage_explicit(self) -> None:
|
||||
settings = StorageSettings(
|
||||
REDIS_URL=None,
|
||||
MEMORY_MAX_KEYS=5000,
|
||||
REDIS_URL = None,
|
||||
MEMORY_MAX_KEYS = 5000,
|
||||
)
|
||||
storage = create_storage(settings)
|
||||
assert isinstance(storage, MemoryStorage)
|
||||
|
|
|
|||
|
|
@ -128,25 +128,29 @@ class TestRateLimitRule:
|
|||
Tests for RateLimitRule dataclass and parsing
|
||||
"""
|
||||
def test_create_valid_rule(self) -> None:
|
||||
rule = RateLimitRule(requests=100, window_seconds=60)
|
||||
rule = RateLimitRule(requests = 100, window_seconds = 60)
|
||||
assert rule.requests == 100
|
||||
assert rule.window_seconds == 60
|
||||
|
||||
def test_invalid_requests_zero(self) -> None:
|
||||
with pytest.raises(ValueError, match="requests must be positive"):
|
||||
RateLimitRule(requests=0, window_seconds=60)
|
||||
with pytest.raises(ValueError,
|
||||
match = "requests must be positive"):
|
||||
RateLimitRule(requests = 0, window_seconds = 60)
|
||||
|
||||
def test_invalid_requests_negative(self) -> None:
|
||||
with pytest.raises(ValueError, match="requests must be positive"):
|
||||
RateLimitRule(requests=-1, window_seconds=60)
|
||||
with pytest.raises(ValueError,
|
||||
match = "requests must be positive"):
|
||||
RateLimitRule(requests = -1, window_seconds = 60)
|
||||
|
||||
def test_invalid_window_zero(self) -> None:
|
||||
with pytest.raises(ValueError, match="window_seconds must be positive"):
|
||||
RateLimitRule(requests=100, window_seconds=0)
|
||||
with pytest.raises(ValueError,
|
||||
match = "window_seconds must be positive"):
|
||||
RateLimitRule(requests = 100, window_seconds = 0)
|
||||
|
||||
def test_invalid_window_negative(self) -> None:
|
||||
with pytest.raises(ValueError, match="window_seconds must be positive"):
|
||||
RateLimitRule(requests=100, window_seconds=-1)
|
||||
with pytest.raises(ValueError,
|
||||
match = "window_seconds must be positive"):
|
||||
RateLimitRule(requests = 100, window_seconds = -1)
|
||||
|
||||
def test_parse_per_second(self) -> None:
|
||||
for unit in ["second", "seconds", "sec", "s"]:
|
||||
|
|
@ -183,43 +187,45 @@ class TestRateLimitRule:
|
|||
assert rule.window_seconds == WINDOW_MINUTE
|
||||
|
||||
def test_parse_invalid_format_no_slash(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid rate limit format"):
|
||||
with pytest.raises(ValueError,
|
||||
match = "Invalid rate limit format"):
|
||||
RateLimitRule.parse("100minute")
|
||||
|
||||
def test_parse_invalid_format_multiple_slashes(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid rate limit format"):
|
||||
with pytest.raises(ValueError,
|
||||
match = "Invalid rate limit format"):
|
||||
RateLimitRule.parse("100/per/minute")
|
||||
|
||||
def test_parse_invalid_request_count(self) -> None:
|
||||
with pytest.raises(ValueError, match="Invalid request count"):
|
||||
with pytest.raises(ValueError, match = "Invalid request count"):
|
||||
RateLimitRule.parse("abc/minute")
|
||||
|
||||
def test_parse_unknown_time_unit(self) -> None:
|
||||
with pytest.raises(ValueError, match="Unknown time unit"):
|
||||
with pytest.raises(ValueError, match = "Unknown time unit"):
|
||||
RateLimitRule.parse("100/fortnight")
|
||||
|
||||
def test_str_representation_second(self) -> None:
|
||||
rule = RateLimitRule(requests=10, window_seconds=1)
|
||||
rule = RateLimitRule(requests = 10, window_seconds = 1)
|
||||
assert str(rule) == "10/second"
|
||||
|
||||
def test_str_representation_minute(self) -> None:
|
||||
rule = RateLimitRule(requests=100, window_seconds=60)
|
||||
rule = RateLimitRule(requests = 100, window_seconds = 60)
|
||||
assert str(rule) == "100/minute"
|
||||
|
||||
def test_str_representation_hour(self) -> None:
|
||||
rule = RateLimitRule(requests=1000, window_seconds=3600)
|
||||
rule = RateLimitRule(requests = 1000, window_seconds = 3600)
|
||||
assert str(rule) == "1000/hour"
|
||||
|
||||
def test_str_representation_day(self) -> None:
|
||||
rule = RateLimitRule(requests=10000, window_seconds=86400)
|
||||
rule = RateLimitRule(requests = 10000, window_seconds = 86400)
|
||||
assert str(rule) == "10000/day"
|
||||
|
||||
def test_str_representation_custom_window(self) -> None:
|
||||
rule = RateLimitRule(requests=50, window_seconds=120)
|
||||
rule = RateLimitRule(requests = 50, window_seconds = 120)
|
||||
assert str(rule) == "50/120s"
|
||||
|
||||
def test_frozen_immutable(self) -> None:
|
||||
rule = RateLimitRule(requests=100, window_seconds=60)
|
||||
rule = RateLimitRule(requests = 100, window_seconds = 60)
|
||||
with pytest.raises(AttributeError):
|
||||
rule.requests = 200
|
||||
|
||||
|
|
@ -245,10 +251,10 @@ class TestRateLimitResult:
|
|||
"""
|
||||
def test_allowed_result(self) -> None:
|
||||
result = RateLimitResult(
|
||||
allowed=True,
|
||||
limit=100,
|
||||
remaining=50,
|
||||
reset_after=30.0,
|
||||
allowed = True,
|
||||
limit = 100,
|
||||
remaining = 50,
|
||||
reset_after = 30.0,
|
||||
)
|
||||
assert result.allowed is True
|
||||
assert result.limit == 100
|
||||
|
|
@ -258,11 +264,11 @@ class TestRateLimitResult:
|
|||
|
||||
def test_denied_result(self) -> None:
|
||||
result = RateLimitResult(
|
||||
allowed=False,
|
||||
limit=100,
|
||||
remaining=0,
|
||||
reset_after=60.0,
|
||||
retry_after=60.0,
|
||||
allowed = False,
|
||||
limit = 100,
|
||||
remaining = 0,
|
||||
reset_after = 60.0,
|
||||
retry_after = 60.0,
|
||||
)
|
||||
assert result.allowed is False
|
||||
assert result.remaining == 0
|
||||
|
|
@ -270,10 +276,10 @@ class TestRateLimitResult:
|
|||
|
||||
def test_headers_basic(self) -> None:
|
||||
result = RateLimitResult(
|
||||
allowed=True,
|
||||
limit=100,
|
||||
remaining=50,
|
||||
reset_after=30.5,
|
||||
allowed = True,
|
||||
limit = 100,
|
||||
remaining = 50,
|
||||
reset_after = 30.5,
|
||||
)
|
||||
headers = result.headers
|
||||
assert headers["RateLimit-Limit"] == "100"
|
||||
|
|
@ -283,21 +289,21 @@ class TestRateLimitResult:
|
|||
|
||||
def test_headers_with_retry_after(self) -> None:
|
||||
result = RateLimitResult(
|
||||
allowed=False,
|
||||
limit=100,
|
||||
remaining=0,
|
||||
reset_after=60.0,
|
||||
retry_after=45.5,
|
||||
allowed = False,
|
||||
limit = 100,
|
||||
remaining = 0,
|
||||
reset_after = 60.0,
|
||||
retry_after = 45.5,
|
||||
)
|
||||
headers = result.headers
|
||||
assert headers["Retry-After"] == "45"
|
||||
|
||||
def test_headers_remaining_never_negative(self) -> None:
|
||||
result = RateLimitResult(
|
||||
allowed=False,
|
||||
limit=100,
|
||||
remaining=-5,
|
||||
reset_after=60.0,
|
||||
allowed = False,
|
||||
limit = 100,
|
||||
remaining = -5,
|
||||
reset_after = 60.0,
|
||||
)
|
||||
headers = result.headers
|
||||
assert headers["RateLimit-Remaining"] == "0"
|
||||
|
|
@ -308,19 +314,19 @@ class TestRateLimitResult:
|
|||
result.allowed = False
|
||||
|
||||
def test_factory_allowed(self) -> None:
|
||||
result = ResultFactory.allowed(limit=200, remaining=150)
|
||||
result = ResultFactory.allowed(limit = 200, remaining = 150)
|
||||
assert result.allowed is True
|
||||
assert result.limit == 200
|
||||
assert result.remaining == 150
|
||||
|
||||
def test_factory_denied(self) -> None:
|
||||
result = ResultFactory.denied(retry_after=30.0)
|
||||
result = ResultFactory.denied(retry_after = 30.0)
|
||||
assert result.allowed is False
|
||||
assert result.remaining == 0
|
||||
assert result.retry_after == 30.0
|
||||
|
||||
def test_factory_near_limit(self) -> None:
|
||||
result = ResultFactory.near_limit(remaining=2)
|
||||
result = ResultFactory.near_limit(remaining = 2)
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 2
|
||||
|
||||
|
|
@ -331,22 +337,22 @@ class TestFingerprintData:
|
|||
"""
|
||||
def test_create_full_fingerprint(self) -> None:
|
||||
fp = FingerprintData(
|
||||
ip=TEST_IP_V4,
|
||||
ip_normalized=TEST_IP_V4,
|
||||
user_agent=TEST_USER_AGENT,
|
||||
accept_language="en-US",
|
||||
accept_encoding="gzip",
|
||||
headers_hash="abc123",
|
||||
auth_identifier="user_456",
|
||||
tls_fingerprint="ja3hash",
|
||||
geo_asn="AS12345",
|
||||
ip = TEST_IP_V4,
|
||||
ip_normalized = TEST_IP_V4,
|
||||
user_agent = TEST_USER_AGENT,
|
||||
accept_language = "en-US",
|
||||
accept_encoding = "gzip",
|
||||
headers_hash = "abc123",
|
||||
auth_identifier = "user_456",
|
||||
tls_fingerprint = "ja3hash",
|
||||
geo_asn = "AS12345",
|
||||
)
|
||||
assert fp.ip == TEST_IP_V4
|
||||
assert fp.user_agent == TEST_USER_AGENT
|
||||
assert fp.auth_identifier == "user_456"
|
||||
|
||||
def test_create_minimal_fingerprint(self) -> None:
|
||||
fp = FingerprintData(ip=TEST_IP_V4, ip_normalized=TEST_IP_V4)
|
||||
fp = FingerprintData(ip = TEST_IP_V4, ip_normalized = TEST_IP_V4)
|
||||
assert fp.ip == TEST_IP_V4
|
||||
assert fp.user_agent is None
|
||||
assert fp.auth_identifier is None
|
||||
|
|
@ -357,13 +363,13 @@ class TestFingerprintData:
|
|||
assert key == TEST_IP_V4
|
||||
|
||||
def test_composite_key_relaxed_with_auth(self) -> None:
|
||||
fp = FingerprintFactory.create(auth_identifier="user_123")
|
||||
fp = FingerprintFactory.create(auth_identifier = "user_123")
|
||||
key = fp.to_composite_key(FingerprintLevel.RELAXED)
|
||||
assert "user_123" in key
|
||||
assert TEST_IP_V4 in key
|
||||
|
||||
def test_composite_key_normal(self) -> None:
|
||||
fp = FingerprintFactory.create(auth_identifier="user_123")
|
||||
fp = FingerprintFactory.create(auth_identifier = "user_123")
|
||||
key = fp.to_composite_key(FingerprintLevel.NORMAL)
|
||||
parts = key.split(":")
|
||||
assert len(parts) == 3
|
||||
|
|
@ -373,15 +379,15 @@ class TestFingerprintData:
|
|||
|
||||
def test_composite_key_strict(self) -> None:
|
||||
fp = FingerprintData(
|
||||
ip=TEST_IP_V4,
|
||||
ip_normalized=TEST_IP_V4,
|
||||
user_agent=TEST_USER_AGENT,
|
||||
accept_language="en-US",
|
||||
accept_encoding="gzip",
|
||||
headers_hash="abc123",
|
||||
auth_identifier="user_456",
|
||||
tls_fingerprint="ja3hash",
|
||||
geo_asn="AS12345",
|
||||
ip = TEST_IP_V4,
|
||||
ip_normalized = TEST_IP_V4,
|
||||
user_agent = TEST_USER_AGENT,
|
||||
accept_language = "en-US",
|
||||
accept_encoding = "gzip",
|
||||
headers_hash = "abc123",
|
||||
auth_identifier = "user_456",
|
||||
tls_fingerprint = "ja3hash",
|
||||
geo_asn = "AS12345",
|
||||
)
|
||||
key = fp.to_composite_key(FingerprintLevel.STRICT)
|
||||
parts = key.split(":")
|
||||
|
|
@ -389,7 +395,7 @@ class TestFingerprintData:
|
|||
assert parts[0] == TEST_IP_V4
|
||||
|
||||
def test_factory_authenticated(self) -> None:
|
||||
fp = FingerprintFactory.authenticated(auth_id="testuser")
|
||||
fp = FingerprintFactory.authenticated(auth_id = "testuser")
|
||||
assert fp.auth_identifier is not None
|
||||
assert len(fp.auth_identifier) == 16
|
||||
|
||||
|
|
@ -409,22 +415,22 @@ class TestWindowState:
|
|||
assert state.previous_count == 0
|
||||
|
||||
def test_weighted_count_start_of_window(self) -> None:
|
||||
state = WindowState(current_count=50, previous_count=100)
|
||||
state = WindowState(current_count = 50, previous_count = 100)
|
||||
weighted = state.weighted_count(0.0)
|
||||
assert weighted == 150.0
|
||||
|
||||
def test_weighted_count_end_of_window(self) -> None:
|
||||
state = WindowState(current_count=50, previous_count=100)
|
||||
state = WindowState(current_count = 50, previous_count = 100)
|
||||
weighted = state.weighted_count(1.0)
|
||||
assert weighted == 50.0
|
||||
|
||||
def test_weighted_count_mid_window(self) -> None:
|
||||
state = WindowState(current_count=50, previous_count=100)
|
||||
state = WindowState(current_count = 50, previous_count = 100)
|
||||
weighted = state.weighted_count(0.5)
|
||||
assert weighted == 100.0
|
||||
|
||||
def test_weighted_count_quarter_window(self) -> None:
|
||||
state = WindowState(current_count=40, previous_count=80)
|
||||
state = WindowState(current_count = 40, previous_count = 80)
|
||||
weighted = state.weighted_count(0.25)
|
||||
assert weighted == 80 * 0.75 + 40
|
||||
|
||||
|
|
@ -434,7 +440,7 @@ class TestWindowState:
|
|||
assert state.previous_count == 0
|
||||
|
||||
def test_factory_with_usage(self) -> None:
|
||||
state = WindowStateFactory.with_usage(current=50, previous=100)
|
||||
state = WindowStateFactory.with_usage(current = 50, previous = 100)
|
||||
assert state.current_count == 50
|
||||
assert state.previous_count == 100
|
||||
|
||||
|
|
@ -445,17 +451,17 @@ class TestTokenBucketState:
|
|||
"""
|
||||
def test_create_state(self) -> None:
|
||||
state = TokenBucketState(
|
||||
tokens=50.0,
|
||||
last_refill=1000.0,
|
||||
capacity=100,
|
||||
refill_rate=1.67,
|
||||
tokens = 50.0,
|
||||
last_refill = 1000.0,
|
||||
capacity = 100,
|
||||
refill_rate = 1.67,
|
||||
)
|
||||
assert state.tokens == 50.0
|
||||
assert state.capacity == 100
|
||||
assert state.refill_rate == 1.67
|
||||
|
||||
def test_factory_full(self) -> None:
|
||||
state = TokenBucketStateFactory.full(capacity=200)
|
||||
state = TokenBucketStateFactory.full(capacity = 200)
|
||||
assert state.tokens == 200.0
|
||||
assert state.capacity == 200
|
||||
|
||||
|
|
@ -491,11 +497,11 @@ class TestDefenseContext:
|
|||
def test_create_context(self) -> None:
|
||||
fp = FingerprintFactory.create()
|
||||
context = DefenseContext(
|
||||
fingerprint=fp,
|
||||
endpoint=TEST_ENDPOINT,
|
||||
method="GET",
|
||||
is_authenticated=True,
|
||||
reputation_score=0.9,
|
||||
fingerprint = fp,
|
||||
endpoint = TEST_ENDPOINT,
|
||||
method = "GET",
|
||||
is_authenticated = True,
|
||||
reputation_score = 0.9,
|
||||
)
|
||||
assert context.endpoint == TEST_ENDPOINT
|
||||
assert context.is_authenticated is True
|
||||
|
|
@ -507,7 +513,7 @@ class TestDefenseContext:
|
|||
assert context.fingerprint.auth_identifier is not None
|
||||
|
||||
def test_factory_suspicious(self) -> None:
|
||||
context = DefenseContextFactory.suspicious(reputation_score=0.2)
|
||||
context = DefenseContextFactory.suspicious(reputation_score = 0.2)
|
||||
assert context.reputation_score == 0.2
|
||||
assert context.request_count_last_minute == 500
|
||||
|
||||
|
|
@ -518,12 +524,12 @@ class TestRateLimitKey:
|
|||
"""
|
||||
def test_build_key(self) -> None:
|
||||
key = RateLimitKey(
|
||||
prefix=KEY_PREFIX,
|
||||
version=KEY_VERSION,
|
||||
layer=Layer.USER,
|
||||
endpoint=TEST_ENDPOINT,
|
||||
identifier=TEST_IP_V4,
|
||||
window=WINDOW_MINUTE,
|
||||
prefix = KEY_PREFIX,
|
||||
version = KEY_VERSION,
|
||||
layer = Layer.USER,
|
||||
endpoint = TEST_ENDPOINT,
|
||||
identifier = TEST_IP_V4,
|
||||
window = WINDOW_MINUTE,
|
||||
)
|
||||
built = key.build()
|
||||
expected = f"{KEY_PREFIX}:{KEY_VERSION}:user:{TEST_ENDPOINT}:{TEST_IP_V4}:{WINDOW_MINUTE}"
|
||||
|
|
|
|||
Loading…
Reference in New Issue