initial
This commit is contained in:
parent
43b2b5ca45
commit
b34cc4a0db
|
|
@ -0,0 +1,46 @@
|
|||
[style]
|
||||
based_on_style = pep8
|
||||
column_limit = 75
|
||||
indent_width = 4
|
||||
continuation_indent_width = 4
|
||||
indent_closing_brackets = false
|
||||
dedent_closing_brackets = true
|
||||
indent_blank_lines = false
|
||||
spaces_before_comment = 2
|
||||
spaces_around_power_operator = false
|
||||
spaces_around_default_or_named_assign = true
|
||||
space_between_ending_comma_and_closing_bracket = false
|
||||
space_inside_brackets = false
|
||||
spaces_around_subscript_colon = true
|
||||
blank_line_before_nested_class_or_def = false
|
||||
blank_line_before_class_docstring = false
|
||||
blank_lines_around_top_level_definition = 2
|
||||
blank_lines_between_top_level_imports_and_variables = 2
|
||||
blank_line_before_module_docstring = false
|
||||
split_before_logical_operator = true
|
||||
split_before_first_argument = true
|
||||
split_before_named_assigns = true
|
||||
split_complex_comprehension = true
|
||||
split_before_expression_after_opening_paren = false
|
||||
split_before_closing_bracket = true
|
||||
split_all_comma_separated_values = true
|
||||
split_all_top_level_comma_separated_values = false
|
||||
coalesce_brackets = false
|
||||
each_dict_entry_on_separate_line = true
|
||||
allow_multiline_lambdas = false
|
||||
allow_multiline_dictionary_keys = false
|
||||
split_penalty_import_names = 0
|
||||
join_multiple_lines = false
|
||||
align_closing_bracket_with_visual_indent = true
|
||||
arithmetic_precedence_indication = false
|
||||
split_penalty_for_added_line_split = 275
|
||||
use_tabs = false
|
||||
split_before_dot = false
|
||||
split_arguments_when_comma_terminated = true
|
||||
i18n_function_call = ['_', 'N_', 'gettext', 'ngettext']
|
||||
i18n_comment = ['# Translators:', '# i18n:']
|
||||
split_penalty_comprehension = 80
|
||||
split_penalty_after_opening_bracket = 280
|
||||
split_penalty_before_if_expr = 0
|
||||
split_penalty_bitwise_operator = 290
|
||||
split_penalty_logical_operator = 0
|
||||
|
|
@ -0,0 +1,242 @@
|
|||
[project]
|
||||
name = "fastapi-420"
|
||||
version = "0.1.0"
|
||||
description = "Enhance Your Calm - Advanced Rate Limiting & DDoS Protection for FastAPI"
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
requires-python = ">=3.12"
|
||||
keywords = ["fastapi", "rate-limiting", "ddos", "security", "api", "420"]
|
||||
classifiers = [
|
||||
"Development Status :: 4 - Beta",
|
||||
"Environment :: Web Environment",
|
||||
"Framework :: FastAPI",
|
||||
"Intended Audience :: Developers",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Operating System :: OS Independent",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Topic :: Internet :: WWW/HTTP",
|
||||
"Topic :: Security",
|
||||
"Typing :: Typed",
|
||||
]
|
||||
|
||||
dependencies = [
|
||||
"fastapi[standard]>=0.123.0,<1.0.0",
|
||||
"pydantic>=2.12.5,<3.0.0",
|
||||
"pydantic-settings>=2.12.0,<3.0.0",
|
||||
"redis>=7.1.0",
|
||||
"pyjwt>=2.10.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=9.0.2",
|
||||
"pytest-asyncio>=1.3.0",
|
||||
"pytest-cov>=6.0.0",
|
||||
"httpx>=0.28.1",
|
||||
"fakeredis>=2.26.0",
|
||||
"time-machine>=2.16.0",
|
||||
"asgi-lifespan>=2.1.0",
|
||||
"mypy>=1.19.0",
|
||||
"types-redis>=4.6.0",
|
||||
"ruff>=0.14.8",
|
||||
"pylint>=4.0.4",
|
||||
"pylint-pydantic>=0.4.1",
|
||||
"pylint-per-file-ignores>=3.2.0",
|
||||
"pre-commit>=4.2.0",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/CarterPerez-dev/fastapi-420"
|
||||
Documentation = "https://github.com/CarterPerez-dev/fastapi-420#readme"
|
||||
Repository = "https://github.com/CarterPerez-dev/fastapi-420"
|
||||
Issues = "https://github.com/CarterPerez-dev/fastapi-420/issues"
|
||||
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/fastapi_420"]
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py312"
|
||||
line-length = 88
|
||||
src = ["src"]
|
||||
exclude = ["alembic"]
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"E", # pycodestyle errors
|
||||
"W", # pycodestyle warnings
|
||||
"F", # pyflakes
|
||||
"B", # flake8-bugbear
|
||||
"C4", # flake8-comprehensions
|
||||
"UP", # pyupgrade
|
||||
"ARG", # flake8-unused-arguments
|
||||
"SIM", # flake8-simplify
|
||||
"PTH", # flake8-use-pathlib
|
||||
"RUF", # ruff-specific
|
||||
"ASYNC", # flake8-async
|
||||
"S", # flake8-bandit (security)
|
||||
"N", # pep8-naming
|
||||
]
|
||||
ignore = [
|
||||
"E501", # line too long (formatter handles this)
|
||||
"B008", # function call in default argument (FastAPI Depends)
|
||||
"S101", # assert usage (needed for tests)
|
||||
"S104", # 0.0.0.0 binding (intentional for Docker)
|
||||
"S105", # "bearer" token_type is not a password
|
||||
"ARG001", # unused function argument (common in FastAPI deps)
|
||||
"E712", # == False is REQUIRED for SQLAlchemy WHERE clauses
|
||||
"N999", # PascalCase module names (intentional: Base.py, User.py)
|
||||
"N818", # exception naming convention (style preference)
|
||||
"UP046", # Generic[T] syntax (keep for compatibility)
|
||||
"RUF005", # list concatenation style (preference)
|
||||
]
|
||||
|
||||
[tool.ruff.lint.per-file-ignores]
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.12"
|
||||
strict = true
|
||||
warn_return_any = true
|
||||
warn_unused_ignores = true
|
||||
disallow_untyped_defs = true
|
||||
disallow_incomplete_defs = true
|
||||
plugins = ["pydantic.mypy"]
|
||||
exclude = ["alembic", ".venv", "venv"]
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = ["tests.*", "conftest"]
|
||||
ignore_errors = true
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = ["core.logging"]
|
||||
disable_error_code = ["no-any-return"]
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = [
|
||||
"uuid6",
|
||||
"structlog",
|
||||
"structlog.*",
|
||||
"pwdlib",
|
||||
"slowapi",
|
||||
"slowapi.*",
|
||||
]
|
||||
ignore_missing_imports = true
|
||||
|
||||
|
||||
[tool.pydantic-mypy]
|
||||
init_forbid_extra = true
|
||||
init_typed = true
|
||||
warn_required_dynamic_aliases = true
|
||||
|
||||
|
||||
[tool.pylint.main]
|
||||
py-version = "3.11"
|
||||
jobs = 4
|
||||
load-plugins = [
|
||||
"pylint_pydantic",
|
||||
"pylint_per_file_ignores",
|
||||
]
|
||||
persistent = true
|
||||
ignore = [
|
||||
"alembic",
|
||||
"venv",
|
||||
".venv",
|
||||
"__pycache__",
|
||||
"build",
|
||||
"dist",
|
||||
".git",
|
||||
".pytest_cache",
|
||||
".mypy_cache",
|
||||
".ruff_cache",
|
||||
]
|
||||
ignore-paths = [
|
||||
"^alembic/.*",
|
||||
"^venv/.*",
|
||||
"^.venv/.*",
|
||||
"^build/.*",
|
||||
"^dist/.*",
|
||||
]
|
||||
|
||||
[tool.pylint.messages_control]
|
||||
disable = [
|
||||
"C0103", # invalid-name
|
||||
"C0116", # missing-function-docstring (we use minimal docs)
|
||||
"C0121", # singleton-comparison (== False required for SQLAlchemy)
|
||||
"C0301", # line-too-long
|
||||
"C0302", # too-many-lines
|
||||
"C0303", # trailing-whitespace
|
||||
"C0304", # final-newline-missing
|
||||
"C0305", # trailing-newlines
|
||||
"C0411", # wrong-import-order
|
||||
"C0412", # ungrouped-imports (style preference)
|
||||
"E0401", # import-error (uuid6/structlog/pwdlib not found by pylint)
|
||||
"E0611", # no-name-in-module (false positive for config re-exports)
|
||||
"E1102", # not-callable (false positive for SQLAlchemy func.now/count)
|
||||
"E1136", # unsubscriptable-object (false positive for generics)
|
||||
"R0801", # similar-lines
|
||||
"R0901", # too-many-ancestors (SQLAlchemy inheritance)
|
||||
"R0903", # too-few-public-methods
|
||||
"R0917", # too-many-positional-arguments (FastAPI patterns)
|
||||
"W0611", # unused-import (handled by ruff, config.py re-exports)
|
||||
"W0612", # unused-variable (handled by ruff)
|
||||
"W0613", # unused-argument (handled by ruff)
|
||||
"W0621", # redefined-outer-name (FastAPI route/param naming)
|
||||
"W0622", # redefined-builtin
|
||||
"W0718", # broad-exception-caught (intentional in health/rate-limit)
|
||||
]
|
||||
|
||||
[tool.pylint.format]
|
||||
max-line-length = 95
|
||||
|
||||
[tool.pylint.design]
|
||||
max-args = 12
|
||||
max-attributes = 10
|
||||
max-branches = 15
|
||||
max-locals = 20
|
||||
max-statements = 55
|
||||
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
testpaths = ["tests"]
|
||||
addopts = "-ra -q"
|
||||
filterwarnings = [
|
||||
"ignore::DeprecationWarning",
|
||||
]
|
||||
|
||||
[tool.coverage.run]
|
||||
branch = true
|
||||
source = ["src"]
|
||||
|
||||
[tool.coverage.report]
|
||||
exclude_lines = [
|
||||
"pragma: no cover",
|
||||
"if TYPE_CHECKING:",
|
||||
"raise NotImplementedError",
|
||||
]
|
||||
|
||||
|
||||
[tool.ty.src]
|
||||
include = ["src", "tests"]
|
||||
exclude = ["alembic/versions/**", ".venv/**"]
|
||||
respect-ignore-files = true
|
||||
|
||||
[tool.ty.environment]
|
||||
python-version = "3.12"
|
||||
root = ["./src"]
|
||||
python = "./.venv"
|
||||
|
||||
[tool.ty.rules]
|
||||
possibly-missing-attribute = "error"
|
||||
possibly-missing-import = "error"
|
||||
unused-ignore-comment = "warn"
|
||||
redundant-cast = "warn"
|
||||
undefined-reveal = "warn"
|
||||
|
||||
|
|
@ -0,0 +1,67 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
__init__.py
|
||||
"""
|
||||
|
||||
from fastapi_420.config import RateLimiterSettings, get_settings
|
||||
from fastapi_420.defense import CircuitBreaker, LayeredDefense
|
||||
from fastapi_420.dependencies import (
|
||||
LimiterDep,
|
||||
RateLimitDep,
|
||||
ScopedRateLimiter,
|
||||
create_rate_limit_dep,
|
||||
get_limiter,
|
||||
require_rate_limit,
|
||||
set_global_limiter,
|
||||
)
|
||||
from fastapi_420.exceptions import (
|
||||
EnhanceYourCalm,
|
||||
HTTP_420_ENHANCE_YOUR_CALM,
|
||||
RateLimitError,
|
||||
RateLimitExceeded,
|
||||
StorageError,
|
||||
)
|
||||
from fastapi_420.limiter import RateLimiter
|
||||
from fastapi_420.middleware import RateLimitMiddleware, SlowDownMiddleware
|
||||
from fastapi_420.types import (
|
||||
Algorithm,
|
||||
DefenseMode,
|
||||
FingerprintData,
|
||||
FingerprintLevel,
|
||||
Layer,
|
||||
RateLimitResult,
|
||||
RateLimitRule,
|
||||
)
|
||||
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
||||
__all__ = [
|
||||
"HTTP_420_ENHANCE_YOUR_CALM",
|
||||
"Algorithm",
|
||||
"CircuitBreaker",
|
||||
"DefenseMode",
|
||||
"EnhanceYourCalm",
|
||||
"FingerprintData",
|
||||
"FingerprintLevel",
|
||||
"Layer",
|
||||
"LayeredDefense",
|
||||
"LimiterDep",
|
||||
"RateLimitDep",
|
||||
"RateLimitError",
|
||||
"RateLimitExceeded",
|
||||
"RateLimitMiddleware",
|
||||
"RateLimitResult",
|
||||
"RateLimitRule",
|
||||
"RateLimiter",
|
||||
"RateLimiterSettings",
|
||||
"ScopedRateLimiter",
|
||||
"SlowDownMiddleware",
|
||||
"StorageError",
|
||||
"__version__",
|
||||
"create_rate_limit_dep",
|
||||
"get_limiter",
|
||||
"get_settings",
|
||||
"require_rate_limit",
|
||||
"set_global_limiter",
|
||||
]
|
||||
|
|
@ -0,0 +1,39 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
__init__.py
|
||||
"""
|
||||
|
||||
from fastapi_420.algorithms.base import BaseAlgorithm
|
||||
from fastapi_420.algorithms.fixed_window import FixedWindowAlgorithm
|
||||
from fastapi_420.algorithms.sliding_window import SlidingWindowAlgorithm
|
||||
from fastapi_420.algorithms.token_bucket import TokenBucketAlgorithm
|
||||
from fastapi_420.types import Algorithm
|
||||
|
||||
|
||||
def create_algorithm(algorithm_type: Algorithm) -> BaseAlgorithm:
|
||||
"""
|
||||
Factory function to create appropriate algorithm instance
|
||||
"""
|
||||
algorithm_map: dict[Algorithm,
|
||||
type[BaseAlgorithm]] = {
|
||||
Algorithm.SLIDING_WINDOW:
|
||||
SlidingWindowAlgorithm,
|
||||
Algorithm.TOKEN_BUCKET: TokenBucketAlgorithm,
|
||||
Algorithm.FIXED_WINDOW: FixedWindowAlgorithm,
|
||||
Algorithm.LEAKY_BUCKET: SlidingWindowAlgorithm,
|
||||
}
|
||||
|
||||
algorithm_class = algorithm_map.get(
|
||||
algorithm_type,
|
||||
SlidingWindowAlgorithm
|
||||
)
|
||||
return algorithm_class()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BaseAlgorithm",
|
||||
"FixedWindowAlgorithm",
|
||||
"SlidingWindowAlgorithm",
|
||||
"TokenBucketAlgorithm",
|
||||
"create_algorithm",
|
||||
]
|
||||
|
|
@ -0,0 +1,52 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
base.py
|
||||
"""
|
||||
# pylint: disable=unnecessary-ellipsis
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.storage import Storage
|
||||
from fastapi_420.types import RateLimitResult, RateLimitRule
|
||||
|
||||
|
||||
class BaseAlgorithm(ABC):
|
||||
"""
|
||||
Abstract base class for rate limiting algorithms
|
||||
"""
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""
|
||||
Algorithm name for logging and debugging
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def check(
|
||||
self,
|
||||
storage: Storage,
|
||||
key: str,
|
||||
rule: RateLimitRule,
|
||||
timestamp: float | None = None,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Check if request is allowed under rate limit
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_current_usage(
|
||||
self,
|
||||
storage: Storage,
|
||||
key: str,
|
||||
rule: RateLimitRule,
|
||||
) -> int:
|
||||
"""
|
||||
Get current usage count without incrementing
|
||||
"""
|
||||
...
|
||||
|
|
@ -0,0 +1,76 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
fixed_window.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi_420.algorithms.base import BaseAlgorithm
|
||||
from fastapi_420.storage.redis_backend import RedisStorage
|
||||
from fastapi_420.types import Algorithm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.storage import Storage
|
||||
from fastapi_420.types import RateLimitResult, RateLimitRule
|
||||
|
||||
|
||||
class FixedWindowAlgorithm(BaseAlgorithm):
|
||||
"""
|
||||
Fixed window counter algorithm
|
||||
|
||||
Simple implementation but suffers from boundary burst problem:
|
||||
clients can make 2x the limit by timing requests at window edges.
|
||||
|
||||
Use sliding_window for production unless simplicity is required.
|
||||
"""
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return Algorithm.FIXED_WINDOW.value
|
||||
|
||||
async def check(
|
||||
self,
|
||||
storage: Storage,
|
||||
key: str,
|
||||
rule: RateLimitRule,
|
||||
timestamp: float | None = None,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Check and increment counter using fixed window algorithm
|
||||
"""
|
||||
if isinstance(storage, RedisStorage):
|
||||
return await storage.increment_fixed_window(
|
||||
key = key,
|
||||
window_seconds = rule.window_seconds,
|
||||
limit = rule.requests,
|
||||
timestamp = timestamp,
|
||||
)
|
||||
|
||||
return await storage.increment(
|
||||
key = key,
|
||||
window_seconds = rule.window_seconds,
|
||||
limit = rule.requests,
|
||||
timestamp = timestamp,
|
||||
)
|
||||
|
||||
async def get_current_usage(
|
||||
self,
|
||||
storage: Storage,
|
||||
key: str,
|
||||
rule: RateLimitRule,
|
||||
) -> int:
|
||||
"""
|
||||
Get current window count without incrementing
|
||||
"""
|
||||
now = time.time()
|
||||
current_window = int(now // rule.window_seconds)
|
||||
window_key = f"{key}:{current_window}"
|
||||
|
||||
state = await storage.get_window_state(
|
||||
key = window_key,
|
||||
window_seconds = rule.window_seconds,
|
||||
)
|
||||
|
||||
return state.current_count
|
||||
|
|
@ -0,0 +1,66 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
sliding_window.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi_420.algorithms.base import BaseAlgorithm
|
||||
from fastapi_420.types import Algorithm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.storage import Storage
|
||||
from fastapi_420.types import RateLimitResult, RateLimitRule
|
||||
|
||||
|
||||
class SlidingWindowAlgorithm(BaseAlgorithm):
|
||||
"""
|
||||
Sliding window counter algorithm
|
||||
|
||||
The recommended default for production rate limiting.
|
||||
Achieves ~99.997% accuracy with O(1) memory per client.
|
||||
Uses weighted interpolation between two fixed windows to
|
||||
approximate a true sliding window.
|
||||
"""
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return Algorithm.SLIDING_WINDOW.value
|
||||
|
||||
async def check(
|
||||
self,
|
||||
storage: Storage,
|
||||
key: str,
|
||||
rule: RateLimitRule,
|
||||
timestamp: float | None = None,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Check and increment counter using sliding window algorithm
|
||||
"""
|
||||
return await storage.increment(
|
||||
key = key,
|
||||
window_seconds = rule.window_seconds,
|
||||
limit = rule.requests,
|
||||
timestamp = timestamp,
|
||||
)
|
||||
|
||||
async def get_current_usage(
|
||||
self,
|
||||
storage: Storage,
|
||||
key: str,
|
||||
rule: RateLimitRule,
|
||||
) -> int:
|
||||
"""
|
||||
Get current weighted usage count without incrementing
|
||||
"""
|
||||
now = time.time()
|
||||
elapsed_ratio = (now % rule.window_seconds) / rule.window_seconds
|
||||
|
||||
state = await storage.get_window_state(
|
||||
key = key,
|
||||
window_seconds = rule.window_seconds,
|
||||
)
|
||||
|
||||
return int(state.weighted_count(elapsed_ratio))
|
||||
|
|
@ -0,0 +1,66 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
token_bucket.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi_420.algorithms.base import BaseAlgorithm
|
||||
from fastapi_420.types import Algorithm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.storage import Storage
|
||||
from fastapi_420.types import RateLimitResult, RateLimitRule
|
||||
|
||||
|
||||
class TokenBucketAlgorithm(BaseAlgorithm):
|
||||
"""
|
||||
Token bucket algorithm
|
||||
|
||||
Allows controlled bursting up to bucket capacity while
|
||||
enforcing average rate limits. Tokens refill at a constant
|
||||
rate and are consumed per request.
|
||||
|
||||
Best for APIs that need burst tolerance with eventual rate enforcement.
|
||||
"""
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return Algorithm.TOKEN_BUCKET.value
|
||||
|
||||
async def check(
|
||||
self,
|
||||
storage: Storage,
|
||||
key: str,
|
||||
rule: RateLimitRule,
|
||||
timestamp: float | None = None, # noqa: ARG002
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Check and consume token using token bucket algorithm
|
||||
"""
|
||||
capacity = rule.requests
|
||||
refill_rate = rule.requests / rule.window_seconds
|
||||
|
||||
return await storage.consume_token(
|
||||
key = key,
|
||||
capacity = capacity,
|
||||
refill_rate = refill_rate,
|
||||
tokens_to_consume = 1,
|
||||
)
|
||||
|
||||
async def get_current_usage(
|
||||
self,
|
||||
storage: Storage,
|
||||
key: str,
|
||||
rule: RateLimitRule,
|
||||
) -> int:
|
||||
"""
|
||||
Get current token count (inverted as usage)
|
||||
"""
|
||||
state = await storage.get_token_bucket_state(key = key)
|
||||
|
||||
if state is None:
|
||||
return 0
|
||||
|
||||
return rule.requests - int(state.tokens)
|
||||
|
|
@ -0,0 +1,188 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
config.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import (
|
||||
Field,
|
||||
RedisDsn,
|
||||
model_validator,
|
||||
)
|
||||
from pydantic_settings import (
|
||||
BaseSettings,
|
||||
SettingsConfigDict,
|
||||
)
|
||||
from fastapi_420.types import (
|
||||
Algorithm,
|
||||
DefenseMode,
|
||||
FingerprintLevel,
|
||||
RateLimitRule,
|
||||
)
|
||||
|
||||
|
||||
class StorageSettings(BaseSettings):
|
||||
"""
|
||||
Storage backend configuration
|
||||
"""
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix = "RATELIMIT_",
|
||||
env_file = ".env",
|
||||
env_file_encoding = "utf-8",
|
||||
extra = "ignore",
|
||||
)
|
||||
|
||||
REDIS_URL: RedisDsn | None = None
|
||||
REDIS_MAX_CONNECTIONS: Annotated[int, Field(ge = 1, le = 1000)] = 100
|
||||
REDIS_SOCKET_TIMEOUT: Annotated[float,
|
||||
Field(ge = 0.1,
|
||||
le = 60.0)] = 5.0
|
||||
REDIS_RETRY_ON_TIMEOUT: bool = True
|
||||
REDIS_DECODE_RESPONSES: bool = True
|
||||
|
||||
MEMORY_MAX_KEYS: Annotated[int,
|
||||
Field(ge = 100,
|
||||
le = 10_000_000)] = 100_000
|
||||
MEMORY_CLEANUP_INTERVAL: Annotated[int, Field(ge = 1, le = 3600)] = 60
|
||||
|
||||
FALLBACK_TO_MEMORY: bool = True
|
||||
|
||||
|
||||
class FingerprintSettings(BaseSettings):
|
||||
"""
|
||||
Fingerprinting configuration
|
||||
"""
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix = "RATELIMIT_FP_",
|
||||
env_file = ".env",
|
||||
env_file_encoding = "utf-8",
|
||||
extra = "ignore",
|
||||
)
|
||||
|
||||
LEVEL: FingerprintLevel = FingerprintLevel.NORMAL
|
||||
USE_IP: bool = True
|
||||
USE_USER_AGENT: bool = True
|
||||
USE_ACCEPT_HEADERS: bool = False
|
||||
USE_HEADER_ORDER: bool = False
|
||||
USE_AUTH: bool = True
|
||||
USE_TLS: bool = False
|
||||
USE_GEO: bool = False
|
||||
|
||||
IPV6_PREFIX_LENGTH: Annotated[int, Field(ge = 32, le = 128)] = 64
|
||||
TRUSTED_PROXIES: list[str] = []
|
||||
TRUST_X_FORWARDED_FOR: bool = False
|
||||
|
||||
|
||||
class DefenseSettings(BaseSettings):
|
||||
"""
|
||||
DDoS defense layer configuration
|
||||
"""
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix = "RATELIMIT_DEFENSE_",
|
||||
env_file = ".env",
|
||||
env_file_encoding = "utf-8",
|
||||
extra = "ignore",
|
||||
)
|
||||
|
||||
MODE: DefenseMode = DefenseMode.ADAPTIVE
|
||||
GLOBAL_LIMIT: str = "50000/minute"
|
||||
CIRCUIT_THRESHOLD: Annotated[int, Field(ge = 1)] = 10000
|
||||
CIRCUIT_WINDOW: Annotated[int, Field(ge = 1, le = 3600)] = 60
|
||||
CIRCUIT_RECOVERY_TIME: Annotated[int, Field(ge = 1, le = 3600)] = 30
|
||||
|
||||
ADAPTIVE_REDUCTION_FACTOR: Annotated[float,
|
||||
Field(ge = 0.1,
|
||||
le = 1.0)] = 0.5
|
||||
ENDPOINT_LIMIT_MULTIPLIER: Annotated[int, Field(ge = 1, le = 100)] = 10
|
||||
LOCKDOWN_ALLOW_AUTHENTICATED: bool = True
|
||||
LOCKDOWN_ALLOW_KNOWN_GOOD: bool = True
|
||||
|
||||
@model_validator(mode = "after")
|
||||
def validate_global_limit(self) -> DefenseSettings:
|
||||
"""
|
||||
Validate global limit can be parsed.
|
||||
"""
|
||||
RateLimitRule.parse(self.GLOBAL_LIMIT)
|
||||
return self
|
||||
|
||||
|
||||
class RateLimiterSettings(BaseSettings):
|
||||
"""
|
||||
Main rate limiter settings with environment variable support
|
||||
"""
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix = "RATELIMIT_",
|
||||
env_file = ".env",
|
||||
env_file_encoding = "utf-8",
|
||||
extra = "ignore",
|
||||
)
|
||||
|
||||
ENABLED: bool = True
|
||||
ALGORITHM: Algorithm = Algorithm.SLIDING_WINDOW
|
||||
DEFAULT_LIMIT: str = "100/minute"
|
||||
DEFAULT_LIMITS: list[str] = ["100/minute", "1000/hour"]
|
||||
FAIL_OPEN: bool = True
|
||||
KEY_PREFIX: str = "ratelimit"
|
||||
KEY_VERSION: str = "v1"
|
||||
INCLUDE_HEADERS: bool = True
|
||||
LOG_VIOLATIONS: bool = True
|
||||
ENVIRONMENT: Literal["development",
|
||||
"staging",
|
||||
"production"] = "development"
|
||||
|
||||
HTTP_420_MESSAGE: str = "Enhance your calm"
|
||||
HTTP_420_DETAIL: str = "Rate limit exceeded. Take a breather."
|
||||
|
||||
endpoint_limits: dict[str, list[RateLimitRule]] = {}
|
||||
|
||||
storage: StorageSettings = StorageSettings()
|
||||
fingerprint: FingerprintSettings = FingerprintSettings()
|
||||
defense: DefenseSettings = DefenseSettings()
|
||||
|
||||
@model_validator(mode = "after")
|
||||
def validate_limits(self) -> RateLimiterSettings:
|
||||
"""
|
||||
Validate all limit strings can be parsed
|
||||
"""
|
||||
RateLimitRule.parse(self.DEFAULT_LIMIT)
|
||||
for limit in self.DEFAULT_LIMITS:
|
||||
RateLimitRule.parse(limit)
|
||||
return self
|
||||
|
||||
@model_validator(mode = "after")
|
||||
def validate_production_settings(self) -> RateLimiterSettings:
|
||||
"""
|
||||
Enforce stricter settings in production
|
||||
"""
|
||||
if self.ENVIRONMENT == "production": # noqa: SIM102
|
||||
if self.storage.REDIS_URL is None and not self.storage.FALLBACK_TO_MEMORY:
|
||||
raise ValueError(
|
||||
"Production requires Redis URL or FALLBACK_TO_MEMORY=True"
|
||||
)
|
||||
return self
|
||||
|
||||
def get_default_rules(self) -> list[RateLimitRule]:
|
||||
"""
|
||||
Parse and return default rate limit rules
|
||||
"""
|
||||
return [
|
||||
RateLimitRule.parse(limit) for limit in self.DEFAULT_LIMITS
|
||||
]
|
||||
|
||||
def get_global_limit_rule(self) -> RateLimitRule:
|
||||
"""
|
||||
Parse and return global defense limit rule
|
||||
"""
|
||||
return RateLimitRule.parse(self.defense.GLOBAL_LIMIT)
|
||||
|
||||
|
||||
@lru_cache
|
||||
def get_settings() -> RateLimiterSettings:
|
||||
"""
|
||||
Cached settings instance
|
||||
"""
|
||||
return RateLimiterSettings()
|
||||
|
|
@ -0,0 +1,14 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
__init__.py
|
||||
"""
|
||||
|
||||
from fastapi_420.defense.circuit_breaker import CircuitBreaker
|
||||
from fastapi_420.defense.layers import LayeredDefense, LayerResult
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CircuitBreaker",
|
||||
"LayerResult",
|
||||
"LayeredDefense",
|
||||
]
|
||||
|
|
@ -0,0 +1,142 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
circuit_breaker.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import asyncio
|
||||
import logging
|
||||
from dataclasses import (
|
||||
field,
|
||||
dataclass,
|
||||
)
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi_420.types import CircuitState, DefenseMode
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.storage import Storage
|
||||
|
||||
|
||||
logger = logging.getLogger("fastapi_420")
|
||||
|
||||
|
||||
@dataclass
|
||||
class CircuitBreaker:
|
||||
"""
|
||||
Global circuit breaker for API-wide protection
|
||||
|
||||
When request volume exceeds threshold, the circuit opens
|
||||
and applies configured defense strategy.
|
||||
"""
|
||||
threshold: int = 10000
|
||||
window_seconds: int = 60
|
||||
recovery_time: int = 30
|
||||
defense_mode: DefenseMode = DefenseMode.ADAPTIVE
|
||||
|
||||
_state: CircuitState = field(default_factory = CircuitState)
|
||||
_lock: asyncio.Lock = field(default_factory = asyncio.Lock)
|
||||
_counter_key: str = "circuit:global:requests"
|
||||
|
||||
async def check(self, storage: Storage) -> bool:
|
||||
"""
|
||||
Check if circuit is allowing requests
|
||||
|
||||
Returns True if requests should be allowed
|
||||
"""
|
||||
async with self._lock:
|
||||
now = time.time()
|
||||
|
||||
if self._state.is_open:
|
||||
if now - self._state.last_failure_time >= self.recovery_time:
|
||||
await self._enter_half_open()
|
||||
return True
|
||||
return False
|
||||
|
||||
request_count = await self._get_request_count(storage)
|
||||
self._state.total_requests_in_window = request_count
|
||||
|
||||
if request_count >= self.threshold:
|
||||
await self._trip(now)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
async def record_request(self, storage: Storage) -> None:
|
||||
"""
|
||||
Record a request in the circuit breaker counter
|
||||
"""
|
||||
now = time.time()
|
||||
window = int(now // self.window_seconds)
|
||||
key = f"{self._counter_key}:{window}"
|
||||
|
||||
await storage.increment(
|
||||
key = key,
|
||||
window_seconds = self.window_seconds,
|
||||
limit = self.threshold * 10,
|
||||
timestamp = now,
|
||||
)
|
||||
|
||||
async def _get_request_count(self, storage: Storage) -> int:
|
||||
"""
|
||||
Get current request count in window
|
||||
"""
|
||||
now = time.time()
|
||||
window = int(now // self.window_seconds)
|
||||
key = f"{self._counter_key}:{window}"
|
||||
|
||||
state = await storage.get_window_state(
|
||||
key = key,
|
||||
window_seconds = self.window_seconds,
|
||||
)
|
||||
|
||||
return state.current_count
|
||||
|
||||
async def _trip(self, now: float) -> None:
|
||||
"""
|
||||
Trip the circuit breaker
|
||||
"""
|
||||
self._state.is_open = True
|
||||
self._state.last_failure_time = now
|
||||
self._state.failure_count += 1
|
||||
self._state.half_open_requests = 0
|
||||
|
||||
logger.warning(
|
||||
"Circuit breaker tripped",
|
||||
extra = {
|
||||
"threshold": self.threshold,
|
||||
"total_requests": self._state.total_requests_in_window,
|
||||
"defense_mode": self.defense_mode.value,
|
||||
},
|
||||
)
|
||||
|
||||
async def _enter_half_open(self) -> None:
|
||||
"""
|
||||
Enter half-open state for recovery testing
|
||||
"""
|
||||
self._state.half_open_requests = 0
|
||||
logger.info("Circuit breaker entering half-open state")
|
||||
|
||||
async def reset(self) -> None:
|
||||
"""
|
||||
Reset circuit breaker to closed state
|
||||
"""
|
||||
async with self._lock:
|
||||
self._state = CircuitState()
|
||||
logger.info("Circuit breaker reset to closed state")
|
||||
|
||||
@property
|
||||
def is_open(self) -> bool:
|
||||
"""
|
||||
Check if circuit is open
|
||||
"""
|
||||
return self._state.is_open
|
||||
|
||||
@property
|
||||
def current_state(self) -> CircuitState:
|
||||
"""
|
||||
Get current circuit state
|
||||
"""
|
||||
return self._state
|
||||
|
|
@ -0,0 +1,314 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
layers.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi_420.algorithms import create_algorithm
|
||||
from fastapi_420.exceptions import EnhanceYourCalm
|
||||
from fastapi_420.types import (
|
||||
DefenseContext,
|
||||
DefenseMode,
|
||||
FingerprintData,
|
||||
Layer,
|
||||
RateLimitKey,
|
||||
RateLimitResult,
|
||||
RateLimitRule,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
|
||||
from fastapi_420.config import RateLimiterSettings
|
||||
from fastapi_420.defense.circuit_breaker import CircuitBreaker
|
||||
from fastapi_420.storage import Storage
|
||||
|
||||
|
||||
logger = logging.getLogger("fastapi_420")
|
||||
|
||||
|
||||
@dataclass
|
||||
class LayerResult:
|
||||
"""
|
||||
Result from a defense layer check
|
||||
"""
|
||||
layer: Layer
|
||||
allowed: bool
|
||||
result: RateLimitResult
|
||||
should_continue: bool = True
|
||||
|
||||
|
||||
class LayeredDefense:
|
||||
"""
|
||||
Three-layer defense system for comprehensive rate limiting
|
||||
|
||||
Layer 1: Per-User Per-Endpoint (granular)
|
||||
Layer 2: Per-Endpoint Global (endpoint protection)
|
||||
Layer 3: Global Circuit Breaker (DDoS protection)
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
storage: Storage,
|
||||
settings: RateLimiterSettings,
|
||||
circuit_breaker: CircuitBreaker | None = None,
|
||||
) -> None:
|
||||
self._storage = storage
|
||||
self._settings = settings
|
||||
self._circuit_breaker = circuit_breaker
|
||||
self._algorithm = create_algorithm(settings.ALGORITHM)
|
||||
|
||||
async def check_all_layers(
|
||||
self,
|
||||
request: Request,
|
||||
fingerprint: FingerprintData,
|
||||
endpoint: str,
|
||||
rules: list[RateLimitRule],
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Check all defense layers in order
|
||||
"""
|
||||
context = DefenseContext(
|
||||
fingerprint = fingerprint,
|
||||
endpoint = endpoint,
|
||||
method = request.method,
|
||||
is_authenticated = fingerprint.auth_identifier is not None,
|
||||
)
|
||||
|
||||
layer3_result = await self._check_layer3_global(context)
|
||||
if not layer3_result.allowed:
|
||||
self._log_violation(layer3_result, context)
|
||||
raise EnhanceYourCalm(
|
||||
result = layer3_result.result,
|
||||
message = self._settings.HTTP_420_MESSAGE,
|
||||
detail = "API is under heavy load. Please try again later.",
|
||||
)
|
||||
|
||||
layer2_result = await self._check_layer2_endpoint(context, rules)
|
||||
if not layer2_result.allowed:
|
||||
self._log_violation(layer2_result, context)
|
||||
raise EnhanceYourCalm(
|
||||
result = layer2_result.result,
|
||||
message = self._settings.HTTP_420_MESSAGE,
|
||||
detail = self._settings.HTTP_420_DETAIL,
|
||||
)
|
||||
|
||||
layer1_result = await self._check_layer1_user(context, rules)
|
||||
if not layer1_result.allowed:
|
||||
self._log_violation(layer1_result, context)
|
||||
raise EnhanceYourCalm(
|
||||
result = layer1_result.result,
|
||||
message = self._settings.HTTP_420_MESSAGE,
|
||||
detail = self._settings.HTTP_420_DETAIL,
|
||||
)
|
||||
|
||||
return layer1_result.result
|
||||
|
||||
async def _check_layer1_user(
|
||||
self,
|
||||
context: DefenseContext,
|
||||
rules: list[RateLimitRule],
|
||||
) -> LayerResult:
|
||||
"""
|
||||
Layer 1: Per-user per-endpoint rate limiting
|
||||
"""
|
||||
identifier = context.fingerprint.to_composite_key(
|
||||
self._settings.fingerprint.LEVEL
|
||||
)
|
||||
|
||||
worst_result: RateLimitResult | None = None
|
||||
|
||||
for rule in rules:
|
||||
key = RateLimitKey(
|
||||
prefix = self._settings.KEY_PREFIX,
|
||||
version = self._settings.KEY_VERSION,
|
||||
layer = Layer.USER,
|
||||
endpoint = context.endpoint,
|
||||
identifier = identifier,
|
||||
window = rule.window_seconds,
|
||||
).build()
|
||||
|
||||
result = await self._algorithm.check(
|
||||
storage = self._storage,
|
||||
key = key,
|
||||
rule = rule,
|
||||
)
|
||||
|
||||
if not result.allowed: # noqa: SIM102
|
||||
if worst_result is None or (result.retry_after or 0) > (
|
||||
worst_result.retry_after or 0):
|
||||
worst_result = result
|
||||
|
||||
if worst_result:
|
||||
return LayerResult(
|
||||
layer = Layer.USER,
|
||||
allowed = False,
|
||||
result = worst_result,
|
||||
)
|
||||
|
||||
return LayerResult(
|
||||
layer = Layer.USER,
|
||||
allowed = True,
|
||||
result = result,
|
||||
)
|
||||
|
||||
async def _check_layer2_endpoint(
|
||||
self,
|
||||
context: DefenseContext,
|
||||
rules: list[RateLimitRule], # noqa: ARG002
|
||||
) -> LayerResult:
|
||||
"""
|
||||
Layer 2: Per-endpoint global rate limiting
|
||||
"""
|
||||
endpoint_rules = self._settings.endpoint_limits.get(
|
||||
context.endpoint,
|
||||
self._settings.get_default_rules()
|
||||
)
|
||||
|
||||
for rule in endpoint_rules:
|
||||
key = RateLimitKey(
|
||||
prefix = self._settings.KEY_PREFIX,
|
||||
version = self._settings.KEY_VERSION,
|
||||
layer = Layer.ENDPOINT,
|
||||
endpoint = context.endpoint,
|
||||
identifier = "global",
|
||||
window = rule.window_seconds,
|
||||
).build()
|
||||
|
||||
endpoint_rule = RateLimitRule(
|
||||
requests = rule.requests * self._settings.defense.ENDPOINT_LIMIT_MULTIPLIER,
|
||||
window_seconds = rule.window_seconds,
|
||||
)
|
||||
|
||||
result = await self._algorithm.check(
|
||||
storage = self._storage,
|
||||
key = key,
|
||||
rule = endpoint_rule,
|
||||
)
|
||||
|
||||
if not result.allowed:
|
||||
return LayerResult(
|
||||
layer = Layer.ENDPOINT,
|
||||
allowed = False,
|
||||
result = result,
|
||||
)
|
||||
|
||||
return LayerResult(
|
||||
layer = Layer.ENDPOINT,
|
||||
allowed = True,
|
||||
result = result,
|
||||
)
|
||||
|
||||
async def _check_layer3_global(
|
||||
self,
|
||||
context: DefenseContext,
|
||||
) -> LayerResult:
|
||||
"""
|
||||
Layer 3: Global circuit breaker
|
||||
"""
|
||||
if self._circuit_breaker is None:
|
||||
return LayerResult(
|
||||
layer = Layer.GLOBAL,
|
||||
allowed = True,
|
||||
result = RateLimitResult(
|
||||
allowed = True,
|
||||
limit = 0,
|
||||
remaining = 0,
|
||||
reset_after = 0,
|
||||
),
|
||||
)
|
||||
|
||||
is_allowed = await self._circuit_breaker.check(self._storage)
|
||||
|
||||
if not is_allowed:
|
||||
if self._should_bypass_circuit(context):
|
||||
return LayerResult(
|
||||
layer = Layer.GLOBAL,
|
||||
allowed = True,
|
||||
result = RateLimitResult(
|
||||
allowed = True,
|
||||
limit = 0,
|
||||
remaining = 0,
|
||||
reset_after = 0,
|
||||
),
|
||||
)
|
||||
|
||||
return LayerResult(
|
||||
layer = Layer.GLOBAL,
|
||||
allowed = False,
|
||||
result = RateLimitResult(
|
||||
allowed = False,
|
||||
limit = self._circuit_breaker.threshold,
|
||||
remaining = 0,
|
||||
reset_after = float(
|
||||
self._circuit_breaker.recovery_time
|
||||
),
|
||||
retry_after = float(
|
||||
self._circuit_breaker.recovery_time
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
await self._circuit_breaker.record_request(self._storage)
|
||||
|
||||
return LayerResult(
|
||||
layer = Layer.GLOBAL,
|
||||
allowed = True,
|
||||
result = RateLimitResult(
|
||||
allowed = True,
|
||||
limit = self._circuit_breaker.threshold,
|
||||
remaining = max(
|
||||
0,
|
||||
self._circuit_breaker.threshold - self._circuit_breaker
|
||||
.current_state.total_requests_in_window,
|
||||
),
|
||||
reset_after = float(self._circuit_breaker.window_seconds),
|
||||
),
|
||||
)
|
||||
|
||||
def _should_bypass_circuit(self, context: DefenseContext) -> bool:
|
||||
"""
|
||||
Determine if request should bypass open circuit
|
||||
"""
|
||||
mode = self._settings.defense.MODE
|
||||
|
||||
if mode == DefenseMode.DISABLED:
|
||||
return True
|
||||
|
||||
if mode == DefenseMode.LOCKDOWN:
|
||||
if self._settings.defense.LOCKDOWN_ALLOW_AUTHENTICATED and context.is_authenticated:
|
||||
return True
|
||||
|
||||
return self._settings.defense.LOCKDOWN_ALLOW_KNOWN_GOOD and context.reputation_score >= 0.9
|
||||
|
||||
if mode == DefenseMode.ADAPTIVE:
|
||||
return bool(context.is_authenticated)
|
||||
|
||||
return False
|
||||
|
||||
def _log_violation(
|
||||
self,
|
||||
layer_result: LayerResult,
|
||||
context: DefenseContext,
|
||||
) -> None:
|
||||
"""
|
||||
Log rate limit violation
|
||||
"""
|
||||
if not self._settings.LOG_VIOLATIONS:
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
"Rate limit exceeded",
|
||||
extra = {
|
||||
"layer": layer_result.layer.value,
|
||||
"endpoint": context.endpoint,
|
||||
"method": context.method,
|
||||
"is_authenticated": context.is_authenticated,
|
||||
"remaining": layer_result.result.remaining,
|
||||
"reset_after": layer_result.result.reset_after,
|
||||
},
|
||||
)
|
||||
|
|
@ -0,0 +1,177 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
dependencies.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
from collections.abc import Callable
|
||||
|
||||
from fastapi import Depends, Request
|
||||
|
||||
from fastapi_420.limiter import RateLimiter
|
||||
from fastapi_420.types import RateLimitResult, RateLimitRule
|
||||
|
||||
|
||||
_global_limiter: RateLimiter | None = None
|
||||
|
||||
|
||||
def set_global_limiter(limiter: RateLimiter) -> None:
|
||||
"""
|
||||
Set the global rate limiter instance for dependency injection
|
||||
"""
|
||||
global _global_limiter # pylint: disable=global-statement
|
||||
_global_limiter = limiter
|
||||
|
||||
|
||||
def get_limiter() -> RateLimiter:
|
||||
"""
|
||||
Get the global rate limiter instance
|
||||
"""
|
||||
if _global_limiter is None:
|
||||
raise RuntimeError(
|
||||
"Rate limiter not initialized. "
|
||||
"Call set_global_limiter() or use RateLimiterDep with explicit limiter."
|
||||
)
|
||||
return _global_limiter
|
||||
|
||||
|
||||
class RateLimitDep:
|
||||
"""
|
||||
FastAPI dependency for rate limiting
|
||||
|
||||
Usage:
|
||||
@app.get("/api/data", dependencies=[Depends(RateLimitDep("100/minute"))])
|
||||
async def get_data():
|
||||
return {"data": "value"}
|
||||
|
||||
# Or with result access:
|
||||
@app.get("/api/data")
|
||||
async def get_data(limit_result: Annotated[RateLimitResult, Depends(RateLimitDep("100/minute"))]):
|
||||
return {"remaining": limit_result.remaining}
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
*rules: str,
|
||||
limiter: RateLimiter | None = None,
|
||||
key_func: Callable[[Request],
|
||||
str] | None = None,
|
||||
) -> None:
|
||||
self.rules = [RateLimitRule.parse(rule) for rule in rules]
|
||||
self._limiter = limiter
|
||||
self.key_func = key_func
|
||||
|
||||
@property
|
||||
def limiter(self) -> RateLimiter:
|
||||
"""
|
||||
Get limiter instance
|
||||
"""
|
||||
if self._limiter is not None:
|
||||
return self._limiter
|
||||
return get_limiter()
|
||||
|
||||
async def __call__(self, request: Request) -> RateLimitResult:
|
||||
"""
|
||||
Check rate limit and return result
|
||||
"""
|
||||
rule_strings = [str(rule) for rule in self.rules]
|
||||
return await self.limiter.check(
|
||||
request,
|
||||
*rule_strings,
|
||||
key_func = self.key_func,
|
||||
raise_on_limit = True,
|
||||
)
|
||||
|
||||
|
||||
def create_rate_limit_dep(
|
||||
*rules: str,
|
||||
limiter: RateLimiter | None = None,
|
||||
key_func: Callable[[Request],
|
||||
str] | None = None,
|
||||
) -> RateLimitDep:
|
||||
"""
|
||||
Factory function to create rate limit dependency
|
||||
|
||||
Usage:
|
||||
rate_limit = create_rate_limit_dep("100/minute", "1000/hour")
|
||||
|
||||
@app.get("/api/data", dependencies=[Depends(rate_limit)])
|
||||
async def get_data():
|
||||
return {"data": "value"}
|
||||
"""
|
||||
return RateLimitDep(*rules, limiter = limiter, key_func = key_func)
|
||||
|
||||
|
||||
LimiterDep = Annotated[RateLimiter, Depends(get_limiter)]
|
||||
|
||||
|
||||
async def require_rate_limit(
|
||||
request: Request,
|
||||
limiter: LimiterDep,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Dependency that applies default rate limits
|
||||
|
||||
Usage:
|
||||
@app.get("/api/data")
|
||||
async def get_data(
|
||||
limit_result: Annotated[RateLimitResult, Depends(require_rate_limit)]
|
||||
):
|
||||
return {"remaining": limit_result.remaining}
|
||||
"""
|
||||
return await limiter.check(request, raise_on_limit = True)
|
||||
|
||||
|
||||
class ScopedRateLimiter:
|
||||
"""
|
||||
Rate limiter scoped to specific endpoints or route groups
|
||||
|
||||
Usage:
|
||||
api_limiter = ScopedRateLimiter(
|
||||
prefix="/api/v1",
|
||||
default_rules=["100/minute"],
|
||||
endpoint_rules={
|
||||
"POST:/api/v1/upload": ["10/minute"],
|
||||
"POST:/api/v1/login": ["5/minute"],
|
||||
}
|
||||
)
|
||||
|
||||
@app.post("/api/v1/upload", dependencies=[Depends(api_limiter)])
|
||||
async def upload():
|
||||
return {"status": "ok"}
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
prefix: str = "",
|
||||
default_rules: list[str] | None = None,
|
||||
endpoint_rules: dict[str,
|
||||
list[str]] | None = None,
|
||||
limiter: RateLimiter | None = None,
|
||||
) -> None:
|
||||
self.prefix = prefix
|
||||
self.default_rules = default_rules or ["100/minute"]
|
||||
self.endpoint_rules = endpoint_rules or {}
|
||||
self._limiter = limiter
|
||||
|
||||
@property
|
||||
def limiter(self) -> RateLimiter:
|
||||
"""
|
||||
Get limiter instance
|
||||
"""
|
||||
if self._limiter is not None:
|
||||
return self._limiter
|
||||
return get_limiter()
|
||||
|
||||
async def __call__(self, request: Request) -> RateLimitResult:
|
||||
"""
|
||||
Apply appropriate rate limit based on endpoint
|
||||
"""
|
||||
endpoint = f"{request.method}:{request.url.path}"
|
||||
|
||||
rules = self.endpoint_rules.get(endpoint, self.default_rules)
|
||||
|
||||
return await self.limiter.check(
|
||||
request,
|
||||
*rules,
|
||||
raise_on_limit = True,
|
||||
)
|
||||
|
|
@ -0,0 +1,211 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
exceptions.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from fastapi_420.types import Layer, StorageType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.types import RateLimitResult
|
||||
|
||||
|
||||
HTTP_420_ENHANCE_YOUR_CALM = 420
|
||||
|
||||
|
||||
class RateLimitError(Exception):
|
||||
"""
|
||||
Base exception for rate limiting errors
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
details: dict[str,
|
||||
Any] | None = None
|
||||
) -> None:
|
||||
self.message = message
|
||||
self.details = details or {}
|
||||
super().__init__(self.message)
|
||||
|
||||
|
||||
class EnhanceYourCalm(HTTPException):
|
||||
"""
|
||||
HTTP 420 response - the signature rate limit response
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
result: RateLimitResult | None = None,
|
||||
message: str = "Enhance your calm",
|
||||
detail: str = "Rate limit exceeded. Take a breather.",
|
||||
headers: dict[str,
|
||||
str] | None = None,
|
||||
) -> None:
|
||||
final_headers = {}
|
||||
if result:
|
||||
final_headers.update(result.headers)
|
||||
if headers:
|
||||
final_headers.update(headers)
|
||||
|
||||
super().__init__(
|
||||
status_code = HTTP_420_ENHANCE_YOUR_CALM,
|
||||
detail = {
|
||||
"message": message,
|
||||
"detail": detail,
|
||||
"limit_info": result.headers if result else {}
|
||||
},
|
||||
headers = final_headers if final_headers else None,
|
||||
)
|
||||
self.result = result
|
||||
|
||||
|
||||
class RateLimitExceeded(RateLimitError):
|
||||
"""
|
||||
Raised when a rate limit is exceeded at any layer
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
result: RateLimitResult,
|
||||
layer: Layer = Layer.USER,
|
||||
endpoint: str = "",
|
||||
identifier: str = "",
|
||||
) -> None:
|
||||
self.result = result
|
||||
self.layer = layer
|
||||
self.endpoint = endpoint
|
||||
self.identifier = identifier
|
||||
super().__init__(
|
||||
message = f"Rate limit exceeded on {layer.value} layer",
|
||||
details = {
|
||||
"layer": layer.value,
|
||||
"endpoint": endpoint,
|
||||
"remaining": result.remaining,
|
||||
"reset_after": result.reset_after,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class StorageError(RateLimitError):
|
||||
"""
|
||||
Raised when storage backend operations fail
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
operation: str,
|
||||
backend: StorageType | None = None,
|
||||
original_error: Exception | None = None,
|
||||
) -> None:
|
||||
self.operation = operation
|
||||
self.backend = backend
|
||||
self.original_error = original_error
|
||||
backend_name = backend.value if backend else "unknown"
|
||||
super().__init__(
|
||||
message =
|
||||
f"Storage operation '{operation}' failed on {backend_name}",
|
||||
details = {
|
||||
"operation": operation,
|
||||
"backend": backend_name,
|
||||
"error": str(original_error) if original_error else None,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class StorageConnectionError(StorageError):
|
||||
"""
|
||||
Raised when unable to connect to storage backend
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
backend: StorageType,
|
||||
original_error: Exception | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
operation = "connect",
|
||||
backend = backend,
|
||||
original_error = original_error,
|
||||
)
|
||||
|
||||
|
||||
class StorageUnavailable(StorageError):
|
||||
"""
|
||||
Raised when storage backend is temporarily unavailable
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
backend: StorageType,
|
||||
original_error: Exception | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
operation = "health_check",
|
||||
backend = backend,
|
||||
original_error = original_error,
|
||||
)
|
||||
|
||||
|
||||
class FingerprintError(RateLimitError):
|
||||
"""
|
||||
Raised when fingerprint extraction fails
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
reason: str,
|
||||
original_error: Exception | None = None,
|
||||
) -> None:
|
||||
self.reason = reason
|
||||
self.original_error = original_error
|
||||
super().__init__(
|
||||
message = f"Fingerprint extraction failed: {reason}",
|
||||
details = {
|
||||
"reason": reason,
|
||||
"error": str(original_error) if original_error else None,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class CircuitBreakerOpen(RateLimitError):
|
||||
"""
|
||||
Raised when global circuit breaker is open
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
recovery_time: float,
|
||||
total_requests: int,
|
||||
threshold: int,
|
||||
) -> None:
|
||||
self.recovery_time = recovery_time
|
||||
self.total_requests = total_requests
|
||||
self.threshold = threshold
|
||||
super().__init__(
|
||||
message = "Circuit breaker is open - API is in defense mode",
|
||||
details = {
|
||||
"recovery_time": recovery_time,
|
||||
"total_requests": total_requests,
|
||||
"threshold": threshold,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class ConfigurationError(RateLimitError):
|
||||
"""
|
||||
Raised when rate limiter configuration is invalid
|
||||
"""
|
||||
def __init__(self, reason: str) -> None:
|
||||
super().__init__(
|
||||
message = f"Invalid configuration: {reason}",
|
||||
details = {"reason": reason},
|
||||
)
|
||||
|
||||
|
||||
class InvalidRuleError(ConfigurationError):
|
||||
"""
|
||||
Raised when a rate limit rule string is invalid
|
||||
"""
|
||||
def __init__(self, rule_string: str, reason: str) -> None:
|
||||
self.rule_string = rule_string
|
||||
super().__init__(
|
||||
reason = f"Invalid rule '{rule_string}': {reason}"
|
||||
)
|
||||
|
|
@ -0,0 +1,17 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
__init__.py
|
||||
"""
|
||||
|
||||
from fastapi_420.fingerprinting.auth import AuthExtractor
|
||||
from fastapi_420.fingerprinting.composite import CompositeFingerprinter
|
||||
from fastapi_420.fingerprinting.headers import HeadersExtractor
|
||||
from fastapi_420.fingerprinting.ip import IPExtractor
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AuthExtractor",
|
||||
"CompositeFingerprinter",
|
||||
"HeadersExtractor",
|
||||
"IPExtractor",
|
||||
]
|
||||
|
|
@ -0,0 +1,138 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
auth.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import jwt
|
||||
import json
|
||||
import base64
|
||||
import hashlib
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
|
||||
|
||||
class AuthExtractor:
|
||||
"""
|
||||
Extract authentication identifiers from requests
|
||||
|
||||
Supports JWT tokens, API keys in headers/query params,
|
||||
and session-based authentication.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
jwt_secret: str | None = None,
|
||||
jwt_algorithms: list[str] | None = None,
|
||||
api_key_header: str = "X-API-Key",
|
||||
api_key_query_param: str = "api_key",
|
||||
session_cookie: str = "session_id",
|
||||
hash_identifiers: bool = True,
|
||||
hash_length: int = 16,
|
||||
) -> None:
|
||||
self.jwt_secret = jwt_secret
|
||||
self.jwt_algorithms = jwt_algorithms or ["HS256"]
|
||||
self.api_key_header = api_key_header
|
||||
self.api_key_query_param = api_key_query_param
|
||||
self.session_cookie = session_cookie
|
||||
self.hash_identifiers = hash_identifiers
|
||||
self.hash_length = hash_length
|
||||
|
||||
def extract(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract authentication identifier using fallback chain
|
||||
|
||||
Order: JWT -> API Key (header) -> API Key (query) -> Session -> None
|
||||
"""
|
||||
identifier = (
|
||||
self._extract_jwt_subject(request)
|
||||
or self._extract_api_key_header(request)
|
||||
or self._extract_api_key_query(request)
|
||||
or self._extract_session(request)
|
||||
)
|
||||
|
||||
if identifier and self.hash_identifiers:
|
||||
return self._hash_identifier(identifier)
|
||||
|
||||
return identifier
|
||||
|
||||
def _extract_jwt_subject(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract subject claim from JWT Bearer token
|
||||
"""
|
||||
auth_header = request.headers.get("Authorization", "")
|
||||
|
||||
if not auth_header.startswith("Bearer "):
|
||||
return None
|
||||
|
||||
token = auth_header[7 :]
|
||||
|
||||
if not self.jwt_secret:
|
||||
return self._extract_jwt_subject_unsafe(token)
|
||||
|
||||
try:
|
||||
payload = jwt.decode(
|
||||
token,
|
||||
self.jwt_secret,
|
||||
algorithms = self.jwt_algorithms,
|
||||
)
|
||||
return str(payload.get("sub", ""))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _extract_jwt_subject_unsafe(self, token: str) -> str | None:
|
||||
"""
|
||||
Extract subject from JWT without verification
|
||||
|
||||
Used when jwt_secret is not configured - only extracts
|
||||
the claim for rate limiting purposes, does NOT validate.
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) != 3:
|
||||
return None
|
||||
|
||||
payload_b64 = parts[1]
|
||||
padding = 4 - len(payload_b64) % 4
|
||||
if padding != 4:
|
||||
payload_b64 += "=" * padding
|
||||
|
||||
payload_bytes = base64.urlsafe_b64decode(payload_b64)
|
||||
payload = json.loads(payload_bytes)
|
||||
|
||||
return str(payload.get("sub", ""))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _extract_api_key_header(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract API key from header
|
||||
"""
|
||||
return request.headers.get(self.api_key_header)
|
||||
|
||||
def _extract_api_key_query(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract API key from query parameter
|
||||
"""
|
||||
return request.query_params.get(self.api_key_query_param)
|
||||
|
||||
def _extract_session(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract session ID from cookie
|
||||
"""
|
||||
return request.cookies.get(self.session_cookie)
|
||||
|
||||
def _hash_identifier(self, identifier: str) -> str:
|
||||
"""
|
||||
Hash identifier for privacy
|
||||
"""
|
||||
hash_bytes = hashlib.sha256(identifier.encode()).hexdigest()
|
||||
return hash_bytes[: self.hash_length]
|
||||
|
||||
def is_authenticated(self, request: Request) -> bool:
|
||||
"""
|
||||
Check if request has any authentication
|
||||
"""
|
||||
return self.extract(request) is not None
|
||||
|
|
@ -0,0 +1,198 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
composite.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi_420.fingerprinting.auth import AuthExtractor
|
||||
from fastapi_420.fingerprinting.headers import HeadersExtractor
|
||||
from fastapi_420.fingerprinting.ip import IPExtractor
|
||||
from fastapi_420.types import FingerprintData, FingerprintLevel
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
|
||||
from fastapi_420.config import FingerprintSettings
|
||||
|
||||
|
||||
class CompositeFingerprinter: # pylint: disable=too-many-instance-attributes
|
||||
"""
|
||||
Combines multiple fingerprinting methods based on configuration
|
||||
|
||||
Preset levels determine which extractors are used:
|
||||
- strict: All methods (IP, headers, auth, TLS, geo)
|
||||
- normal: IP + User-Agent + Auth (default)
|
||||
- relaxed: IP + Auth only
|
||||
- custom: Configured via settings
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
level: FingerprintLevel = FingerprintLevel.NORMAL,
|
||||
ip_extractor: IPExtractor | None = None,
|
||||
headers_extractor: HeadersExtractor | None = None,
|
||||
auth_extractor: AuthExtractor | None = None,
|
||||
use_ip: bool = True,
|
||||
use_user_agent: bool = True,
|
||||
use_accept_headers: bool = False,
|
||||
use_header_order: bool = False,
|
||||
use_auth: bool = True,
|
||||
use_tls: bool = False,
|
||||
use_geo: bool = False,
|
||||
) -> None:
|
||||
self.level = level
|
||||
self._ip_extractor = ip_extractor or IPExtractor()
|
||||
self._headers_extractor = headers_extractor or HeadersExtractor(
|
||||
use_header_order = use_header_order,
|
||||
)
|
||||
self._auth_extractor = auth_extractor or AuthExtractor()
|
||||
|
||||
if level == FingerprintLevel.STRICT:
|
||||
self.use_ip = True
|
||||
self.use_user_agent = True
|
||||
self.use_accept_headers = True
|
||||
self.use_header_order = True
|
||||
self.use_auth = True
|
||||
self.use_tls = True
|
||||
self.use_geo = True
|
||||
elif level == FingerprintLevel.RELAXED:
|
||||
self.use_ip = True
|
||||
self.use_user_agent = False
|
||||
self.use_accept_headers = False
|
||||
self.use_header_order = False
|
||||
self.use_auth = True
|
||||
self.use_tls = False
|
||||
self.use_geo = False
|
||||
elif level == FingerprintLevel.CUSTOM:
|
||||
self.use_ip = use_ip
|
||||
self.use_user_agent = use_user_agent
|
||||
self.use_accept_headers = use_accept_headers
|
||||
self.use_header_order = use_header_order
|
||||
self.use_auth = use_auth
|
||||
self.use_tls = use_tls
|
||||
self.use_geo = use_geo
|
||||
else:
|
||||
self.use_ip = True
|
||||
self.use_user_agent = True
|
||||
self.use_accept_headers = False
|
||||
self.use_header_order = False
|
||||
self.use_auth = True
|
||||
self.use_tls = False
|
||||
self.use_geo = False
|
||||
|
||||
@classmethod
|
||||
def from_settings(
|
||||
cls,
|
||||
settings: FingerprintSettings
|
||||
) -> CompositeFingerprinter:
|
||||
"""
|
||||
Create fingerprinter from settings
|
||||
"""
|
||||
ip_extractor = IPExtractor(
|
||||
ipv6_prefix_length = settings.IPV6_PREFIX_LENGTH,
|
||||
trusted_proxies = settings.TRUSTED_PROXIES,
|
||||
trust_x_forwarded_for = settings.TRUST_X_FORWARDED_FOR,
|
||||
)
|
||||
|
||||
headers_extractor = HeadersExtractor(
|
||||
use_header_order = settings.USE_HEADER_ORDER,
|
||||
)
|
||||
|
||||
return cls(
|
||||
level = settings.LEVEL,
|
||||
ip_extractor = ip_extractor,
|
||||
headers_extractor = headers_extractor,
|
||||
use_ip = settings.USE_IP,
|
||||
use_user_agent = settings.USE_USER_AGENT,
|
||||
use_accept_headers = settings.USE_ACCEPT_HEADERS,
|
||||
use_header_order = settings.USE_HEADER_ORDER,
|
||||
use_auth = settings.USE_AUTH,
|
||||
use_tls = settings.USE_TLS,
|
||||
use_geo = settings.USE_GEO,
|
||||
)
|
||||
|
||||
async def extract(self, request: Request) -> FingerprintData:
|
||||
"""
|
||||
Extract fingerprint data from request
|
||||
"""
|
||||
raw_ip = ""
|
||||
normalized_ip = ""
|
||||
user_agent = None
|
||||
accept_language = None
|
||||
accept_encoding = None
|
||||
headers_hash = None
|
||||
auth_identifier = None
|
||||
tls_fingerprint = None
|
||||
geo_asn = None
|
||||
|
||||
if self.use_ip:
|
||||
raw_ip, normalized_ip = self._ip_extractor.extract(request)
|
||||
|
||||
if self.use_user_agent:
|
||||
user_agent = self._headers_extractor.extract_user_agent(
|
||||
request
|
||||
)
|
||||
|
||||
if self.use_accept_headers:
|
||||
accept_language = self._headers_extractor.extract_accept_language(
|
||||
request
|
||||
)
|
||||
accept_encoding = self._headers_extractor.extract_accept_encoding(
|
||||
request
|
||||
)
|
||||
|
||||
if self.use_header_order:
|
||||
headers_hash = self._headers_extractor.compute_headers_hash(
|
||||
request
|
||||
)
|
||||
|
||||
if self.use_auth:
|
||||
auth_identifier = self._auth_extractor.extract(request)
|
||||
|
||||
if self.use_tls:
|
||||
tls_fingerprint = self._extract_tls_fingerprint(request)
|
||||
|
||||
if self.use_geo:
|
||||
geo_asn = self._extract_geo_asn(request)
|
||||
|
||||
return FingerprintData(
|
||||
ip = raw_ip,
|
||||
ip_normalized = normalized_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,
|
||||
)
|
||||
|
||||
def _extract_tls_fingerprint(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract TLS/JA3 fingerprint if available
|
||||
|
||||
Requires proxy to pass fingerprint in header
|
||||
"""
|
||||
return (
|
||||
request.headers.get("X-JA3-Fingerprint")
|
||||
or request.headers.get("X-TLS-Fingerprint")
|
||||
)
|
||||
|
||||
def _extract_geo_asn(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract geographic ASN if available
|
||||
|
||||
Requires geo lookup service to populate header
|
||||
"""
|
||||
return (
|
||||
request.headers.get("X-Client-ASN")
|
||||
or request.headers.get("CF-IPCountry")
|
||||
)
|
||||
|
||||
def is_authenticated(self, request: Request) -> bool:
|
||||
"""
|
||||
Check if request has authentication
|
||||
"""
|
||||
return self._auth_extractor.is_authenticated(request)
|
||||
|
|
@ -0,0 +1,95 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
headers.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
|
||||
|
||||
class HeadersExtractor:
|
||||
"""
|
||||
Extract fingerprint data from HTTP headers
|
||||
|
||||
Header ordering is browser-specific and not user-configurable,
|
||||
making it useful for fingerprinting even when other headers are spoofed.
|
||||
"""
|
||||
FINGERPRINT_HEADERS: ClassVar[list[str]] = [
|
||||
"user-agent",
|
||||
"accept",
|
||||
"accept-language",
|
||||
"accept-encoding",
|
||||
"connection",
|
||||
"upgrade-insecure-requests",
|
||||
"sec-fetch-site",
|
||||
"sec-fetch-mode",
|
||||
"sec-fetch-user",
|
||||
"sec-fetch-dest",
|
||||
"sec-ch-ua",
|
||||
"sec-ch-ua-mobile",
|
||||
"sec-ch-ua-platform",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
use_header_order: bool = False,
|
||||
hash_length: int = 16,
|
||||
) -> None:
|
||||
self.use_header_order = use_header_order
|
||||
self.hash_length = hash_length
|
||||
|
||||
def extract_user_agent(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract User-Agent header
|
||||
"""
|
||||
return request.headers.get("user-agent")
|
||||
|
||||
def extract_accept_language(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract Accept-Language header
|
||||
"""
|
||||
return request.headers.get("accept-language")
|
||||
|
||||
def extract_accept_encoding(self, request: Request) -> str | None:
|
||||
"""
|
||||
Extract Accept-Encoding header
|
||||
"""
|
||||
return request.headers.get("accept-encoding")
|
||||
|
||||
def compute_headers_hash(self, request: Request) -> str:
|
||||
"""
|
||||
Compute hash of fingerprint-relevant headers
|
||||
|
||||
Includes header ordering if configured, which is
|
||||
browser-specific and harder to spoof.
|
||||
"""
|
||||
components: list[str] = []
|
||||
|
||||
if self.use_header_order:
|
||||
header_keys = [k.lower() for k in request.headers]
|
||||
components.append("|".join(header_keys))
|
||||
|
||||
for header_name in self.FINGERPRINT_HEADERS:
|
||||
value = request.headers.get(header_name, "")
|
||||
components.append(f"{header_name}={value}")
|
||||
|
||||
fingerprint_string = "\n".join(components)
|
||||
hash_bytes = hashlib.sha256(fingerprint_string.encode()
|
||||
).hexdigest()
|
||||
|
||||
return hash_bytes[: self.hash_length]
|
||||
|
||||
def extract_all(self, request: Request) -> dict[str, str | None]:
|
||||
"""
|
||||
Extract all header-based fingerprint data
|
||||
"""
|
||||
return {
|
||||
"user_agent": self.extract_user_agent(request),
|
||||
"accept_language": self.extract_accept_language(request),
|
||||
"accept_encoding": self.extract_accept_encoding(request),
|
||||
"headers_hash": self.compute_headers_hash(request),
|
||||
}
|
||||
|
|
@ -0,0 +1,134 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
ip.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ipaddress import (
|
||||
IPv6Address,
|
||||
ip_address,
|
||||
ip_network,
|
||||
)
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
|
||||
|
||||
class IPExtractor:
|
||||
"""
|
||||
Extract and normalize client IP addresses
|
||||
|
||||
Handles IPv6 /64 prefix normalization since users control
|
||||
entire /64 prefixes (~18 quintillion addresses), making
|
||||
per-IP limits trivially bypassable without normalization.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
ipv6_prefix_length: int = 64,
|
||||
trusted_proxies: list[str] | None = None,
|
||||
trust_x_forwarded_for: bool = False,
|
||||
) -> None:
|
||||
self.ipv6_prefix_length = ipv6_prefix_length
|
||||
self.trusted_proxies = set(trusted_proxies or [])
|
||||
self.trust_x_forwarded_for = trust_x_forwarded_for
|
||||
|
||||
def extract(self, request: Request) -> tuple[str, str]:
|
||||
"""
|
||||
Extract raw IP and normalized IP from request
|
||||
|
||||
Returns tuple of (raw_ip, normalized_ip)
|
||||
"""
|
||||
raw_ip = self._get_client_ip(request)
|
||||
normalized_ip = self._normalize_ip(raw_ip)
|
||||
return raw_ip, normalized_ip
|
||||
|
||||
def _get_client_ip(self, request: Request) -> str:
|
||||
"""
|
||||
Get client IP address, handling proxy headers if configured
|
||||
"""
|
||||
if self.trust_x_forwarded_for:
|
||||
forwarded_for = request.headers.get("X-Forwarded-For")
|
||||
if forwarded_for:
|
||||
return self._parse_x_forwarded_for(forwarded_for, request)
|
||||
|
||||
x_real_ip = request.headers.get("X-Real-IP")
|
||||
if x_real_ip and self._is_trusted_proxy(
|
||||
request.client.host if request.client else ""):
|
||||
return x_real_ip.strip()
|
||||
|
||||
if request.client:
|
||||
return request.client.host
|
||||
|
||||
return "0.0.0.0"
|
||||
|
||||
def _parse_x_forwarded_for(self, header: str, request: Request) -> str:
|
||||
"""
|
||||
Parse X-Forwarded-For header safely
|
||||
|
||||
Uses rightmost-trusted approach: walk backwards through
|
||||
the chain and return the first IP not in trusted proxies
|
||||
"""
|
||||
ips = [ip.strip() for ip in header.split(",")]
|
||||
|
||||
if not self.trusted_proxies:
|
||||
return ips[0]
|
||||
|
||||
for ip in reversed(ips):
|
||||
if ip not in self.trusted_proxies:
|
||||
return ip
|
||||
|
||||
return ips[0] if ips else (
|
||||
request.client.host if request.client else "0.0.0.0"
|
||||
)
|
||||
|
||||
def _is_trusted_proxy(self, ip: str) -> bool:
|
||||
"""
|
||||
Check if IP is in trusted proxies list
|
||||
"""
|
||||
return ip in self.trusted_proxies
|
||||
|
||||
def _normalize_ip(self, ip_str: str) -> str:
|
||||
"""
|
||||
Normalize IP address for rate limiting
|
||||
|
||||
IPv6 addresses are normalized to their /64 network prefix
|
||||
since users typically control entire /64 blocks.
|
||||
"""
|
||||
try:
|
||||
addr = ip_address(ip_str)
|
||||
except ValueError:
|
||||
return ip_str
|
||||
|
||||
if isinstance(addr, IPv6Address):
|
||||
if addr.ipv4_mapped:
|
||||
return str(addr.ipv4_mapped)
|
||||
|
||||
network = ip_network(
|
||||
f"{ip_str}/{self.ipv6_prefix_length}",
|
||||
strict = False,
|
||||
)
|
||||
return str(network.network_address)
|
||||
|
||||
return str(addr)
|
||||
|
||||
def is_ipv6(self, ip_str: str) -> bool:
|
||||
"""
|
||||
Check if IP string is IPv6
|
||||
"""
|
||||
try:
|
||||
addr = ip_address(ip_str)
|
||||
return isinstance(addr, IPv6Address)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
def is_private(self, ip_str: str) -> bool:
|
||||
"""
|
||||
Check if IP is a private/internal address
|
||||
"""
|
||||
try:
|
||||
addr = ip_address(ip_str)
|
||||
return addr.is_private
|
||||
except ValueError:
|
||||
return False
|
||||
|
|
@ -0,0 +1,347 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
limiter.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import functools
|
||||
from typing import (
|
||||
Any,
|
||||
ParamSpec,
|
||||
TypeVar,
|
||||
TYPE_CHECKING,
|
||||
)
|
||||
from collections.abc import Callable
|
||||
|
||||
from starlette.requests import Request
|
||||
|
||||
from fastapi_420.algorithms import (
|
||||
create_algorithm,
|
||||
)
|
||||
from fastapi_420.config import (
|
||||
RateLimiterSettings,
|
||||
get_settings,
|
||||
)
|
||||
from fastapi_420.exceptions import (
|
||||
EnhanceYourCalm,
|
||||
StorageError,
|
||||
)
|
||||
from fastapi_420.fingerprinting import (
|
||||
CompositeFingerprinter,
|
||||
)
|
||||
from fastapi_420.storage import (
|
||||
MemoryStorage,
|
||||
RedisStorage,
|
||||
create_storage,
|
||||
)
|
||||
from fastapi_420.types import (
|
||||
Layer,
|
||||
RateLimitKey,
|
||||
RateLimitResult,
|
||||
RateLimitRule,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.algorithms.base import BaseAlgorithm
|
||||
from fastapi_420.storage import Storage
|
||||
|
||||
|
||||
logger = logging.getLogger("fastapi_420")
|
||||
|
||||
P = ParamSpec("P")
|
||||
R = TypeVar("R")
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""
|
||||
Main rate limiter class for FastAPI applications.
|
||||
|
||||
Usage:
|
||||
limiter = RateLimiter()
|
||||
|
||||
@app.get("/api/data")
|
||||
@limiter.limit("100/minute", "1000/hour")
|
||||
async def get_data(request: Request):
|
||||
return {"data": "value"}
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
settings: RateLimiterSettings | None = None,
|
||||
storage: Storage | None = None,
|
||||
) -> None:
|
||||
self._settings = settings or get_settings()
|
||||
self._storage = storage
|
||||
self._fallback_storage: MemoryStorage | None = None
|
||||
self._algorithm: BaseAlgorithm | None = None
|
||||
self._fingerprinter: CompositeFingerprinter | None = None
|
||||
self._initialized = False
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def init(self) -> None:
|
||||
"""
|
||||
Initialize storage, algorithm, and fingerprinter
|
||||
"""
|
||||
async with self._lock:
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
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()
|
||||
|
||||
self._algorithm = create_algorithm(self._settings.ALGORITHM)
|
||||
self._fingerprinter = CompositeFingerprinter.from_settings(
|
||||
self._settings.fingerprint
|
||||
)
|
||||
|
||||
self._initialized = True
|
||||
logger.info(
|
||||
"Rate limiter initialized",
|
||||
extra = {
|
||||
"algorithm":
|
||||
self._settings.ALGORITHM.value,
|
||||
"storage":
|
||||
self._storage.storage_type.value,
|
||||
"fingerprint_level":
|
||||
self._settings.fingerprint.LEVEL.value,
|
||||
},
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""
|
||||
Close storage connections
|
||||
"""
|
||||
if self._storage:
|
||||
await self._storage.close()
|
||||
|
||||
if self._fallback_storage:
|
||||
await self._fallback_storage.close()
|
||||
|
||||
self._initialized = False
|
||||
|
||||
def limit(
|
||||
self,
|
||||
*rules: str,
|
||||
key_func: Callable[[Request],
|
||||
str] | None = None,
|
||||
) -> Callable[[Callable[P,
|
||||
R]],
|
||||
Callable[P,
|
||||
R]]:
|
||||
"""
|
||||
Decorator to apply rate limits to an endpoint
|
||||
|
||||
Args:
|
||||
rules: Rate limit strings like "100/minute", "1000/hour"
|
||||
key_func: Optional custom function to generate rate limit key
|
||||
"""
|
||||
parsed_rules = [RateLimitRule.parse(rule) for rule in rules]
|
||||
|
||||
def decorator(func: Callable[P, R]) -> Callable[P, R]:
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
|
||||
request = self._extract_request(args, kwargs)
|
||||
|
||||
if request is None:
|
||||
return await func(*args, **kwargs) # type: ignore[misc, no-any-return]
|
||||
|
||||
await self._check_rate_limits(
|
||||
request = request,
|
||||
rules = parsed_rules,
|
||||
key_func = key_func,
|
||||
)
|
||||
|
||||
return await func(*args, **kwargs) # type: ignore[misc, no-any-return]
|
||||
|
||||
return wrapper # type: ignore[return-value]
|
||||
|
||||
return decorator
|
||||
|
||||
async def check(
|
||||
self,
|
||||
request: Request,
|
||||
*rules: str,
|
||||
key_func: Callable[[Request],
|
||||
str] | None = None,
|
||||
raise_on_limit: bool = True,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Manually check rate limit without decorator
|
||||
|
||||
Returns the result of the strictest (most restrictive) rule
|
||||
"""
|
||||
parsed_rules = [RateLimitRule.parse(rule) for rule in rules]
|
||||
|
||||
if not parsed_rules:
|
||||
parsed_rules = self._settings.get_default_rules()
|
||||
|
||||
return await self._check_rate_limits(
|
||||
request = request,
|
||||
rules = parsed_rules,
|
||||
key_func = key_func,
|
||||
raise_on_limit = raise_on_limit,
|
||||
)
|
||||
|
||||
async def _check_rate_limits(
|
||||
self,
|
||||
request: Request,
|
||||
rules: list[RateLimitRule],
|
||||
key_func: Callable[[Request],
|
||||
str] | None = None,
|
||||
raise_on_limit: bool = True,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Check all rules and return/raise for the most restrictive failure
|
||||
"""
|
||||
if not self._initialized:
|
||||
await self.init()
|
||||
|
||||
storage = await self._get_active_storage()
|
||||
if storage is None:
|
||||
if self._settings.FAIL_OPEN:
|
||||
return RateLimitResult(
|
||||
allowed = True,
|
||||
limit = 0,
|
||||
remaining = 0,
|
||||
reset_after = 0,
|
||||
)
|
||||
raise StorageError(operation = "check", backend = None)
|
||||
|
||||
fingerprint = await self._fingerprinter.extract(request) # type: ignore[union-attr]
|
||||
endpoint = self._get_endpoint(request)
|
||||
|
||||
if key_func:
|
||||
identifier = key_func(request)
|
||||
else:
|
||||
identifier = fingerprint.to_composite_key(
|
||||
self._settings.fingerprint.LEVEL
|
||||
)
|
||||
|
||||
worst_result: RateLimitResult | None = None
|
||||
|
||||
for rule in rules:
|
||||
key = RateLimitKey(
|
||||
prefix = self._settings.KEY_PREFIX,
|
||||
version = self._settings.KEY_VERSION,
|
||||
layer = Layer.USER,
|
||||
endpoint = endpoint,
|
||||
identifier = identifier,
|
||||
window = rule.window_seconds,
|
||||
).build()
|
||||
|
||||
result = await self._algorithm.check( # type: ignore[union-attr]
|
||||
storage = storage,
|
||||
key = key,
|
||||
rule = rule,
|
||||
)
|
||||
|
||||
if not result.allowed: # noqa: SIM102
|
||||
if worst_result is None or result.retry_after > (worst_result.retry_after or 0): # type: ignore[operator]
|
||||
worst_result = result
|
||||
|
||||
if worst_result is not None:
|
||||
if self._settings.LOG_VIOLATIONS:
|
||||
logger.warning(
|
||||
"Rate limit exceeded",
|
||||
extra = {
|
||||
"endpoint": endpoint,
|
||||
"identifier": identifier[: 16],
|
||||
"remaining": worst_result.remaining,
|
||||
"reset_after": worst_result.reset_after,
|
||||
},
|
||||
)
|
||||
|
||||
if raise_on_limit:
|
||||
raise EnhanceYourCalm(
|
||||
result = worst_result,
|
||||
message = self._settings.HTTP_420_MESSAGE,
|
||||
detail = self._settings.HTTP_420_DETAIL,
|
||||
)
|
||||
|
||||
return worst_result
|
||||
|
||||
best_result = result
|
||||
return best_result
|
||||
|
||||
async def _get_active_storage(self) -> Storage | None:
|
||||
"""
|
||||
Get active storage, falling back to memory if primary fails
|
||||
"""
|
||||
if self._storage is None:
|
||||
return self._fallback_storage
|
||||
|
||||
try:
|
||||
is_healthy = await self._storage.health_check()
|
||||
if is_healthy:
|
||||
return self._storage
|
||||
except Exception: # noqa: S110
|
||||
pass
|
||||
|
||||
if self._fallback_storage:
|
||||
logger.warning(
|
||||
"Primary storage unavailable, using memory fallback",
|
||||
extra = {
|
||||
"primary_storage": self._storage.storage_type.value
|
||||
},
|
||||
)
|
||||
return self._fallback_storage
|
||||
|
||||
return None
|
||||
|
||||
def _extract_request(
|
||||
self,
|
||||
args: tuple[Any,
|
||||
...],
|
||||
kwargs: dict[str,
|
||||
Any],
|
||||
) -> Request | None:
|
||||
"""
|
||||
Extract Request object from function arguments
|
||||
"""
|
||||
for arg in args:
|
||||
if isinstance(arg, Request):
|
||||
return arg
|
||||
|
||||
for value in kwargs.values():
|
||||
if isinstance(value, Request):
|
||||
return value
|
||||
|
||||
return None
|
||||
|
||||
def _get_endpoint(self, request: Request) -> str:
|
||||
"""
|
||||
Get endpoint identifier from request
|
||||
"""
|
||||
route = request.scope.get("route")
|
||||
if route:
|
||||
return f"{request.method}:{route.path}"
|
||||
|
||||
return f"{request.method}:{request.url.path}"
|
||||
|
||||
@property
|
||||
def settings(self) -> RateLimiterSettings:
|
||||
"""
|
||||
Get current settings
|
||||
"""
|
||||
return self._settings
|
||||
|
||||
@property
|
||||
def is_initialized(self) -> bool:
|
||||
"""
|
||||
Check if limiter is initialized
|
||||
"""
|
||||
return self._initialized
|
||||
|
|
@ -0,0 +1,210 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
middleware.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
)
|
||||
from collections.abc import Callable
|
||||
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import (
|
||||
JSONResponse,
|
||||
Response,
|
||||
)
|
||||
|
||||
from fastapi_420.exceptions import (
|
||||
HTTP_420_ENHANCE_YOUR_CALM,
|
||||
EnhanceYourCalm,
|
||||
)
|
||||
from fastapi_420.limiter import RateLimiter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.types import ASGIApp
|
||||
|
||||
|
||||
logger = logging.getLogger("fastapi_420")
|
||||
|
||||
|
||||
class RateLimitMiddleware(BaseHTTPMiddleware):
|
||||
"""
|
||||
ASGI middleware for automatic rate limiting on all routes
|
||||
|
||||
Usage:
|
||||
from fastapi import FastAPI
|
||||
from fastapi_420.middleware import RateLimitMiddleware
|
||||
from fastapi_420.limiter import RateLimiter
|
||||
|
||||
app = FastAPI()
|
||||
limiter = RateLimiter()
|
||||
|
||||
app.add_middleware(
|
||||
RateLimitMiddleware,
|
||||
limiter=limiter,
|
||||
default_limit="100/minute",
|
||||
)
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
app: ASGIApp,
|
||||
limiter: RateLimiter,
|
||||
default_limit: str = "100/minute",
|
||||
exclude_paths: list[str] | None = None,
|
||||
exclude_patterns: list[str] | None = None,
|
||||
include_paths: list[str] | None = None,
|
||||
path_limits: dict[str,
|
||||
str] | None = None,
|
||||
key_func: Callable[[Request],
|
||||
str] | None = None,
|
||||
) -> None:
|
||||
super().__init__(app)
|
||||
self.limiter = limiter
|
||||
self.default_limit = default_limit
|
||||
self.exclude_paths = set(exclude_paths or [])
|
||||
self.exclude_patterns = [
|
||||
re.compile(p) for p in (exclude_patterns or [])
|
||||
]
|
||||
self.include_paths = set(include_paths) if include_paths else None
|
||||
self.path_limits = path_limits or {}
|
||||
self.key_func = key_func
|
||||
|
||||
self.exclude_paths.add("/health")
|
||||
self.exclude_paths.add("/healthz")
|
||||
self.exclude_paths.add("/ready")
|
||||
self.exclude_paths.add("/metrics")
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Callable) -> Response: # type: ignore[type-arg]
|
||||
"""
|
||||
Process request and apply rate limiting
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
async def _should_limit(self, request: Request) -> bool:
|
||||
"""
|
||||
Determine if request should be rate limited
|
||||
"""
|
||||
path = request.url.path
|
||||
|
||||
if path in self.exclude_paths:
|
||||
return False
|
||||
|
||||
for pattern in self.exclude_patterns:
|
||||
if pattern.match(path):
|
||||
return False
|
||||
|
||||
if self.include_paths is not None: # noqa: SIM102
|
||||
if path not in self.include_paths:
|
||||
for include_path in self.include_paths:
|
||||
if path.startswith(include_path):
|
||||
break
|
||||
else:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _get_limit_for_path(self, path: str) -> str:
|
||||
"""
|
||||
Get rate limit for specific path
|
||||
"""
|
||||
if path in self.path_limits:
|
||||
return self.path_limits[path]
|
||||
|
||||
for pattern_path, limit in self.path_limits.items():
|
||||
if path.startswith(pattern_path):
|
||||
return limit
|
||||
|
||||
return self.default_limit
|
||||
|
||||
def _create_420_response(self, exc: EnhanceYourCalm) -> JSONResponse:
|
||||
"""
|
||||
Create HTTP 420 response
|
||||
"""
|
||||
headers = {}
|
||||
if exc.result:
|
||||
headers.update(exc.result.headers)
|
||||
|
||||
return JSONResponse(
|
||||
status_code = HTTP_420_ENHANCE_YOUR_CALM,
|
||||
content = exc.detail,
|
||||
headers = headers,
|
||||
)
|
||||
|
||||
|
||||
class SlowDownMiddleware(BaseHTTPMiddleware):
|
||||
"""
|
||||
Alternative middleware that adds delays instead of blocking
|
||||
|
||||
Useful for gradual throttling rather than hard limits
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
app: ASGIApp,
|
||||
limiter: RateLimiter,
|
||||
threshold_limit: str = "50/minute",
|
||||
max_delay_seconds: float = 5.0,
|
||||
delay_increment: float = 0.5,
|
||||
) -> None:
|
||||
super().__init__(app)
|
||||
self.limiter = limiter
|
||||
self.threshold_limit = threshold_limit
|
||||
self.max_delay_seconds = max_delay_seconds
|
||||
self.delay_increment = delay_increment
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Callable) -> Response: # type: ignore[type-arg]
|
||||
"""
|
||||
Process request with potential delays
|
||||
"""
|
||||
result = await self.limiter.check(
|
||||
request,
|
||||
self.threshold_limit,
|
||||
raise_on_limit = False,
|
||||
)
|
||||
|
||||
if result.remaining <= 0:
|
||||
delay = min(
|
||||
self.max_delay_seconds,
|
||||
result.retry_after or self.delay_increment,
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
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]
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
__init__.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, TypeAlias
|
||||
|
||||
from fastapi_420.storage.memory import MemoryStorage
|
||||
from fastapi_420.storage.redis_backend import RedisStorage
|
||||
from fastapi_420.types import StorageType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.config import StorageSettings
|
||||
|
||||
|
||||
Storage: TypeAlias = MemoryStorage | RedisStorage # noqa: UP040
|
||||
|
||||
|
||||
def create_storage(settings: StorageSettings) -> Storage:
|
||||
"""
|
||||
Factory function to create appropriate storage backend.
|
||||
"""
|
||||
if settings.REDIS_URL is not None:
|
||||
return RedisStorage.from_settings(settings)
|
||||
return MemoryStorage.from_settings(settings)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MemoryStorage",
|
||||
"RedisStorage",
|
||||
"Storage",
|
||||
"StorageType",
|
||||
"create_storage",
|
||||
]
|
||||
|
|
@ -0,0 +1,32 @@
|
|||
--[[
|
||||
ⒸAngelaMos | 2025
|
||||
fixed_window.lua
|
||||
|
||||
Atomic fixed window counter rate limiting.
|
||||
Simple but has boundary burst problem - use sliding_window for production.
|
||||
Returns: {allowed (0/1), remaining, reset_after, retry_after}
|
||||
--]]
|
||||
|
||||
local key = KEYS[1]
|
||||
local window_seconds = tonumber(ARGV[1])
|
||||
local limit = tonumber(ARGV[2])
|
||||
local now = tonumber(ARGV[3])
|
||||
|
||||
local current_window = math.floor(now / window_seconds)
|
||||
local window_key = key .. ":" .. current_window
|
||||
|
||||
local count = tonumber(redis.call('GET', window_key)) or 0
|
||||
local reset_after = window_seconds - (now % window_seconds)
|
||||
|
||||
if count >= limit then
|
||||
return {0, 0, reset_after, reset_after}
|
||||
end
|
||||
|
||||
local new_count = redis.call('INCR', window_key)
|
||||
if new_count == 1 then
|
||||
redis.call('EXPIRE', window_key, window_seconds)
|
||||
end
|
||||
|
||||
local remaining = math.max(0, limit - new_count)
|
||||
|
||||
return {1, remaining, reset_after, 0}
|
||||
|
|
@ -0,0 +1,37 @@
|
|||
--[[
|
||||
ⒸAngelaMos | 2025
|
||||
sliding_window.lua
|
||||
|
||||
Atomic sliding window counter rate limiting.
|
||||
Returns: {allowed (0/1), remaining, reset_after}
|
||||
--]]
|
||||
|
||||
local key = KEYS[1]
|
||||
local window_seconds = tonumber(ARGV[1])
|
||||
local limit = tonumber(ARGV[2])
|
||||
local now = tonumber(ARGV[3])
|
||||
|
||||
local current_window = math.floor(now / window_seconds)
|
||||
local previous_window = current_window - 1
|
||||
local elapsed_ratio = (now % window_seconds) / window_seconds
|
||||
|
||||
local current_key = key .. ":" .. current_window
|
||||
local previous_key = key .. ":" .. previous_window
|
||||
|
||||
local current_count = tonumber(redis.call('GET', current_key)) or 0
|
||||
local previous_count = tonumber(redis.call('GET', previous_key)) or 0
|
||||
|
||||
local weighted_count = math.floor(previous_count * (1 - elapsed_ratio) + current_count)
|
||||
local reset_after = window_seconds - (now % window_seconds)
|
||||
|
||||
if weighted_count >= limit then
|
||||
return {0, 0, reset_after, reset_after}
|
||||
end
|
||||
|
||||
redis.call('INCR', current_key)
|
||||
redis.call('EXPIRE', current_key, window_seconds * 2)
|
||||
|
||||
local new_weighted = math.floor(previous_count * (1 - elapsed_ratio) + current_count + 1)
|
||||
local remaining = math.max(0, limit - new_weighted)
|
||||
|
||||
return {1, remaining, reset_after, 0}
|
||||
|
|
@ -0,0 +1,49 @@
|
|||
--[[
|
||||
ⒸAngelaMos | 2025
|
||||
token_bucket.lua
|
||||
|
||||
Atomic token bucket rate limiting.
|
||||
Returns: {allowed (0/1), remaining, reset_after, retry_after}
|
||||
--]]
|
||||
|
||||
local key = KEYS[1]
|
||||
local capacity = tonumber(ARGV[1])
|
||||
local refill_rate = tonumber(ARGV[2])
|
||||
local tokens_to_consume = tonumber(ARGV[3])
|
||||
local now = tonumber(ARGV[4])
|
||||
|
||||
local bucket_key = key .. ":bucket"
|
||||
local data = redis.call('HMGET', bucket_key, 'tokens', 'last_refill')
|
||||
|
||||
local tokens = tonumber(data[1])
|
||||
local last_refill = tonumber(data[2])
|
||||
|
||||
if tokens == nil then
|
||||
tokens = capacity
|
||||
last_refill = now
|
||||
end
|
||||
|
||||
local elapsed = now - last_refill
|
||||
local tokens_to_add = elapsed * refill_rate
|
||||
tokens = math.min(capacity, tokens + tokens_to_add)
|
||||
|
||||
if tokens >= tokens_to_consume then
|
||||
tokens = tokens - tokens_to_consume
|
||||
redis.call('HMSET', bucket_key, 'tokens', tokens, 'last_refill', now)
|
||||
redis.call('EXPIRE', bucket_key, 3600)
|
||||
|
||||
local time_to_full = 0
|
||||
if refill_rate > 0 then
|
||||
time_to_full = (capacity - tokens) / refill_rate
|
||||
end
|
||||
|
||||
return {1, math.floor(tokens), time_to_full, 0}
|
||||
end
|
||||
|
||||
redis.call('HMSET', bucket_key, 'tokens', tokens, 'last_refill', now)
|
||||
redis.call('EXPIRE', bucket_key, 3600)
|
||||
|
||||
local tokens_needed = tokens_to_consume - tokens
|
||||
local wait_time = tokens_needed / refill_rate
|
||||
|
||||
return {0, 0, wait_time, wait_time}
|
||||
|
|
@ -0,0 +1,294 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
memory.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi_420.types import (
|
||||
RateLimitResult,
|
||||
StorageType,
|
||||
TokenBucketState,
|
||||
WindowState,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.config import StorageSettings
|
||||
|
||||
|
||||
@dataclass
|
||||
class WindowEntry:
|
||||
"""
|
||||
Storage entry for sliding window counter.
|
||||
"""
|
||||
count: int = 0
|
||||
window_start: int = 0
|
||||
expires_at: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class MemoryStorage:
|
||||
"""
|
||||
In memory storage backend for rate limiting.
|
||||
|
||||
Thread safe through asyncio locks. Suitable for single instance
|
||||
deployments, development, and as a fallback when Redis is unavailable
|
||||
"""
|
||||
max_keys: int = 100_000
|
||||
cleanup_interval: int = 60
|
||||
_windows: OrderedDict[str,
|
||||
WindowEntry] = field(
|
||||
default_factory = OrderedDict
|
||||
)
|
||||
_buckets: dict[str, TokenBucketState] = field(default_factory = dict)
|
||||
_lock: asyncio.Lock = field(default_factory = asyncio.Lock)
|
||||
_cleanup_task: asyncio.Task[None] | None = field(
|
||||
default = None,
|
||||
repr = False
|
||||
)
|
||||
_closed: bool = field(default = False)
|
||||
|
||||
@classmethod
|
||||
def from_settings(cls, settings: StorageSettings) -> MemoryStorage:
|
||||
"""
|
||||
Create storage instance from settings
|
||||
"""
|
||||
return cls(
|
||||
max_keys = settings.MEMORY_MAX_KEYS,
|
||||
cleanup_interval = settings.MEMORY_CLEANUP_INTERVAL,
|
||||
)
|
||||
|
||||
async def start_cleanup_task(self) -> None:
|
||||
"""
|
||||
Start background cleanup task for expired entries
|
||||
"""
|
||||
if self._cleanup_task is None:
|
||||
self._cleanup_task = asyncio.create_task(self._cleanup_loop())
|
||||
|
||||
async def _cleanup_loop(self) -> None:
|
||||
"""
|
||||
Background loop to clean up expired entries
|
||||
"""
|
||||
while not self._closed:
|
||||
await asyncio.sleep(self.cleanup_interval)
|
||||
await self._cleanup_expired()
|
||||
|
||||
async def _cleanup_expired(self) -> None:
|
||||
"""
|
||||
Remove expired entries from storage
|
||||
"""
|
||||
now = time.time()
|
||||
async with self._lock:
|
||||
expired_keys = [
|
||||
key for key, entry in self._windows.items()
|
||||
if entry.expires_at < now
|
||||
]
|
||||
for key in expired_keys:
|
||||
del self._windows[key]
|
||||
|
||||
expired_buckets = [
|
||||
key for key, state in self._buckets.items()
|
||||
if state.last_refill + 3600 < now
|
||||
]
|
||||
for key in expired_buckets:
|
||||
del self._buckets[key]
|
||||
|
||||
async def _enforce_max_keys(self) -> None:
|
||||
"""
|
||||
Evict oldest entries when max_keys is exceeded
|
||||
"""
|
||||
while len(self._windows) > self.max_keys:
|
||||
self._windows.popitem(last = False)
|
||||
|
||||
async def get_window_state(
|
||||
self,
|
||||
key: str,
|
||||
window_seconds: int,
|
||||
) -> WindowState:
|
||||
"""
|
||||
Get current sliding window state
|
||||
"""
|
||||
now = time.time()
|
||||
current_window = int(now // window_seconds)
|
||||
previous_window = current_window - 1
|
||||
|
||||
current_key = f"{key}:{current_window}"
|
||||
previous_key = f"{key}:{previous_window}"
|
||||
|
||||
async with self._lock:
|
||||
current_entry = self._windows.get(current_key)
|
||||
previous_entry = self._windows.get(previous_key)
|
||||
|
||||
return WindowState(
|
||||
current_count = current_entry.count
|
||||
if current_entry else 0,
|
||||
previous_count = previous_entry.count
|
||||
if previous_entry else 0,
|
||||
current_window = current_window,
|
||||
window_seconds = window_seconds,
|
||||
)
|
||||
|
||||
async def increment(
|
||||
self,
|
||||
key: str,
|
||||
window_seconds: int,
|
||||
limit: int,
|
||||
timestamp: float | None = None,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Atomically check and increment counter using sliding window algorithm
|
||||
"""
|
||||
now = timestamp if timestamp is not None else time.time()
|
||||
current_window = int(now // window_seconds)
|
||||
previous_window = current_window - 1
|
||||
elapsed_ratio = (now % window_seconds) / window_seconds
|
||||
|
||||
current_key = f"{key}:{current_window}"
|
||||
previous_key = f"{key}:{previous_window}"
|
||||
|
||||
async with self._lock:
|
||||
current_entry = self._windows.get(current_key)
|
||||
previous_entry = self._windows.get(previous_key)
|
||||
|
||||
current_count = current_entry.count if current_entry else 0
|
||||
previous_count = previous_entry.count if previous_entry else 0
|
||||
|
||||
weighted_count = int(
|
||||
previous_count * (1 - elapsed_ratio) + current_count
|
||||
)
|
||||
|
||||
if weighted_count >= limit:
|
||||
reset_after = window_seconds - (now % window_seconds)
|
||||
return RateLimitResult(
|
||||
allowed = False,
|
||||
limit = limit,
|
||||
remaining = 0,
|
||||
reset_after = reset_after,
|
||||
retry_after = reset_after,
|
||||
)
|
||||
|
||||
if current_entry:
|
||||
current_entry.count += 1
|
||||
else:
|
||||
self._windows[current_key] = WindowEntry(
|
||||
count = 1,
|
||||
window_start = current_window,
|
||||
expires_at = now + (window_seconds * 2),
|
||||
)
|
||||
self._windows.move_to_end(current_key)
|
||||
await self._enforce_max_keys()
|
||||
|
||||
new_weighted = int(
|
||||
previous_count * (1 - elapsed_ratio) + current_count + 1
|
||||
)
|
||||
remaining = max(0, limit - new_weighted)
|
||||
reset_after = window_seconds - (now % window_seconds)
|
||||
|
||||
return RateLimitResult(
|
||||
allowed = True,
|
||||
limit = limit,
|
||||
remaining = remaining,
|
||||
reset_after = reset_after,
|
||||
)
|
||||
|
||||
async def get_token_bucket_state(
|
||||
self,
|
||||
key: str,
|
||||
) -> TokenBucketState | None:
|
||||
"""
|
||||
Get token bucket state if it exists
|
||||
"""
|
||||
async with self._lock:
|
||||
return self._buckets.get(key)
|
||||
|
||||
async def consume_token(
|
||||
self,
|
||||
key: str,
|
||||
capacity: int,
|
||||
refill_rate: float,
|
||||
tokens_to_consume: int = 1,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Attempt to consume tokens from bucket
|
||||
"""
|
||||
now = time.time()
|
||||
|
||||
async with self._lock:
|
||||
state = self._buckets.get(key)
|
||||
|
||||
if state is None:
|
||||
state = TokenBucketState(
|
||||
tokens = float(capacity),
|
||||
last_refill = now,
|
||||
capacity = capacity,
|
||||
refill_rate = refill_rate,
|
||||
)
|
||||
self._buckets[key] = state
|
||||
|
||||
elapsed = now - state.last_refill
|
||||
tokens_to_add = elapsed * refill_rate
|
||||
state.tokens = min(
|
||||
float(capacity),
|
||||
state.tokens + tokens_to_add
|
||||
)
|
||||
state.last_refill = now
|
||||
|
||||
if state.tokens >= tokens_to_consume:
|
||||
state.tokens -= tokens_to_consume
|
||||
time_to_full = (
|
||||
capacity - state.tokens
|
||||
) / refill_rate if refill_rate > 0 else 0
|
||||
|
||||
return RateLimitResult(
|
||||
allowed = True,
|
||||
limit = capacity,
|
||||
remaining = int(state.tokens),
|
||||
reset_after = time_to_full,
|
||||
)
|
||||
|
||||
tokens_needed = tokens_to_consume - state.tokens
|
||||
wait_time = tokens_needed / refill_rate if refill_rate > 0 else float(
|
||||
"inf"
|
||||
)
|
||||
|
||||
return RateLimitResult(
|
||||
allowed = False,
|
||||
limit = capacity,
|
||||
remaining = 0,
|
||||
reset_after = wait_time,
|
||||
retry_after = wait_time,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""
|
||||
Close storage and cleanup resources
|
||||
"""
|
||||
self._closed = True
|
||||
if self._cleanup_task:
|
||||
self._cleanup_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._cleanup_task
|
||||
self._cleanup_task = None
|
||||
|
||||
async with self._lock:
|
||||
self._windows.clear()
|
||||
self._buckets.clear()
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""
|
||||
Check if storage is healthy
|
||||
"""
|
||||
return not self._closed
|
||||
|
||||
@property
|
||||
def storage_type(self) -> StorageType:
|
||||
"""
|
||||
Return storage type identifier
|
||||
"""
|
||||
return StorageType.MEMORY
|
||||
|
|
@ -0,0 +1,408 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
redis_backend.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import (
|
||||
dataclass,
|
||||
field,
|
||||
)
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import redis.asyncio as redis
|
||||
from redis.asyncio.connection import ConnectionPool
|
||||
from redis.exceptions import (
|
||||
ConnectionError as RedisConnectionError,
|
||||
)
|
||||
from redis.exceptions import (
|
||||
ResponseError,
|
||||
TimeoutError,
|
||||
)
|
||||
from fastapi_420.exceptions import (
|
||||
StorageConnectionError,
|
||||
StorageError,
|
||||
)
|
||||
from fastapi_420.types import (
|
||||
RateLimitResult,
|
||||
StorageType,
|
||||
TokenBucketState,
|
||||
WindowState,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi_420.config import StorageSettings
|
||||
|
||||
|
||||
LUA_SCRIPTS_DIR = Path(__file__).parent / "lua"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RedisStorage:
|
||||
"""
|
||||
Redis storage backend with atomic Lua script operations
|
||||
|
||||
Uses EVALSHA with cached script hashes for optimal performance
|
||||
All rate limiting operations are atomic to prevent race conditions
|
||||
"""
|
||||
|
||||
url: str
|
||||
max_connections: int = 100
|
||||
socket_timeout: float = 5.0
|
||||
retry_on_timeout: bool = True
|
||||
decode_responses: bool = True
|
||||
|
||||
_client: redis.Redis[Any] | None = field(default = None, repr = False)
|
||||
_pool: ConnectionPool[Any] | None = field(default = None, repr = False)
|
||||
_script_shas: dict[str,
|
||||
str] = field(
|
||||
default_factory = dict,
|
||||
repr = False
|
||||
)
|
||||
_scripts_loaded: bool = field(default = False, repr = False)
|
||||
|
||||
@classmethod
|
||||
def from_settings(cls, settings: StorageSettings) -> RedisStorage:
|
||||
"""
|
||||
Create storage instance from settings
|
||||
"""
|
||||
if settings.REDIS_URL is None:
|
||||
raise StorageConnectionError(
|
||||
backend = StorageType.REDIS,
|
||||
original_error = ValueError("REDIS_URL is required"),
|
||||
)
|
||||
|
||||
return cls(
|
||||
url = str(settings.REDIS_URL),
|
||||
max_connections = settings.REDIS_MAX_CONNECTIONS,
|
||||
socket_timeout = settings.REDIS_SOCKET_TIMEOUT,
|
||||
retry_on_timeout = settings.REDIS_RETRY_ON_TIMEOUT,
|
||||
decode_responses = settings.REDIS_DECODE_RESPONSES,
|
||||
)
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""
|
||||
Establish Redis connection and load Lua scripts.
|
||||
"""
|
||||
try:
|
||||
self._pool = ConnectionPool.from_url(
|
||||
self.url,
|
||||
max_connections = self.max_connections,
|
||||
socket_timeout = self.socket_timeout,
|
||||
retry_on_timeout = self.retry_on_timeout,
|
||||
decode_responses = self.decode_responses,
|
||||
)
|
||||
self._client = redis.Redis(connection_pool = self._pool)
|
||||
|
||||
await self._client.ping()
|
||||
await self._load_scripts()
|
||||
|
||||
except (RedisConnectionError, TimeoutError, OSError) as e:
|
||||
raise StorageConnectionError(
|
||||
backend = StorageType.REDIS,
|
||||
original_error = e,
|
||||
) from e
|
||||
|
||||
async def _load_scripts(self) -> None:
|
||||
"""
|
||||
Load Lua scripts into Redis and cache their SHA1 hashes
|
||||
"""
|
||||
if self._client is None:
|
||||
raise StorageError(
|
||||
operation = "load_scripts",
|
||||
backend = StorageType.REDIS,
|
||||
original_error = RuntimeError("Client not connected"),
|
||||
)
|
||||
|
||||
script_files = {
|
||||
"sliding_window": LUA_SCRIPTS_DIR / "sliding_window.lua",
|
||||
"token_bucket": LUA_SCRIPTS_DIR / "token_bucket.lua",
|
||||
"fixed_window": LUA_SCRIPTS_DIR / "fixed_window.lua",
|
||||
}
|
||||
|
||||
for name, path in script_files.items():
|
||||
script_content = path.read_text()
|
||||
sha = await self._client.script_load(script_content) # type: ignore[no-untyped-call]
|
||||
self._script_shas[name] = sha
|
||||
|
||||
self._scripts_loaded = True
|
||||
|
||||
async def _ensure_connected(self) -> redis.Redis[Any]:
|
||||
"""
|
||||
Ensure client is connected and scripts are loaded
|
||||
"""
|
||||
if self._client is None:
|
||||
await self.connect()
|
||||
|
||||
if self._client is None:
|
||||
raise StorageError(
|
||||
operation = "ensure_connected",
|
||||
backend = StorageType.REDIS,
|
||||
original_error = RuntimeError(
|
||||
"Failed to establish connection"
|
||||
),
|
||||
)
|
||||
|
||||
return self._client
|
||||
|
||||
async def _execute_script(
|
||||
self,
|
||||
script_name: str,
|
||||
keys: list[str],
|
||||
args: list[str | int | float],
|
||||
) -> list[int | float]:
|
||||
"""
|
||||
Execute a Lua script using EVALSHA
|
||||
"""
|
||||
client = await self._ensure_connected()
|
||||
|
||||
if script_name not in self._script_shas:
|
||||
raise StorageError(
|
||||
operation = "execute_script",
|
||||
backend = StorageType.REDIS,
|
||||
original_error = ValueError(
|
||||
f"Unknown script: {script_name}"
|
||||
),
|
||||
)
|
||||
|
||||
sha = self._script_shas[script_name]
|
||||
|
||||
try:
|
||||
result = await client.evalsha(sha, len(keys), *keys, *args) # type: ignore[no-untyped-call]
|
||||
return result # type: ignore[no-any-return]
|
||||
|
||||
except ResponseError as e:
|
||||
if "NOSCRIPT" in str(e):
|
||||
await self._load_scripts()
|
||||
sha = self._script_shas[script_name]
|
||||
result = await client.evalsha(sha, len(keys), *keys, *args) # type: ignore[no-untyped-call]
|
||||
return result # type: ignore[no-any-return]
|
||||
raise StorageError(
|
||||
operation = "execute_script",
|
||||
backend = StorageType.REDIS,
|
||||
original_error = e,
|
||||
) from e
|
||||
|
||||
except (RedisConnectionError, TimeoutError) as e:
|
||||
raise StorageError(
|
||||
operation = "execute_script",
|
||||
backend = StorageType.REDIS,
|
||||
original_error = e,
|
||||
) from e
|
||||
|
||||
async def get_window_state(
|
||||
self,
|
||||
key: str,
|
||||
window_seconds: int,
|
||||
) -> WindowState:
|
||||
"""
|
||||
Get current sliding window state
|
||||
"""
|
||||
client = await self._ensure_connected()
|
||||
now = time.time()
|
||||
current_window = int(now // window_seconds)
|
||||
previous_window = current_window - 1
|
||||
|
||||
current_key = f"{key}:{current_window}"
|
||||
previous_key = f"{key}:{previous_window}"
|
||||
|
||||
try:
|
||||
pipeline = client.pipeline()
|
||||
pipeline.get(current_key)
|
||||
pipeline.get(previous_key)
|
||||
results = await pipeline.execute()
|
||||
|
||||
current_count = int(results[0]) if results[0] else 0
|
||||
previous_count = int(results[1]) if results[1] else 0
|
||||
|
||||
return WindowState(
|
||||
current_count = current_count,
|
||||
previous_count = previous_count,
|
||||
current_window = current_window,
|
||||
window_seconds = window_seconds,
|
||||
)
|
||||
|
||||
except (RedisConnectionError, TimeoutError) as e:
|
||||
raise StorageError(
|
||||
operation = "get_window_state",
|
||||
backend = StorageType.REDIS,
|
||||
original_error = e,
|
||||
) from e
|
||||
|
||||
async def increment(
|
||||
self,
|
||||
key: str,
|
||||
window_seconds: int,
|
||||
limit: int,
|
||||
timestamp: float | None = None,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Atomically check and increment counter using sliding window algorithm
|
||||
"""
|
||||
now = timestamp if timestamp is not None else time.time()
|
||||
|
||||
result = await self._execute_script(
|
||||
"sliding_window",
|
||||
keys = [key],
|
||||
args = [window_seconds,
|
||||
limit,
|
||||
now],
|
||||
)
|
||||
|
||||
allowed = bool(result[0])
|
||||
remaining = int(result[1])
|
||||
reset_after = float(result[2])
|
||||
retry_after = float(result[3]) if result[3] else None
|
||||
|
||||
return RateLimitResult(
|
||||
allowed = allowed,
|
||||
limit = limit,
|
||||
remaining = remaining,
|
||||
reset_after = reset_after,
|
||||
retry_after = retry_after if not allowed else None,
|
||||
)
|
||||
|
||||
async def increment_fixed_window(
|
||||
self,
|
||||
key: str,
|
||||
window_seconds: int,
|
||||
limit: int,
|
||||
timestamp: float | None = None,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Atomically check and increment counter using fixed window algorithm
|
||||
"""
|
||||
now = timestamp if timestamp is not None else time.time()
|
||||
|
||||
result = await self._execute_script(
|
||||
"fixed_window",
|
||||
keys = [key],
|
||||
args = [window_seconds,
|
||||
limit,
|
||||
now],
|
||||
)
|
||||
|
||||
allowed = bool(result[0])
|
||||
remaining = int(result[1])
|
||||
reset_after = float(result[2])
|
||||
retry_after = float(result[3]) if result[3] else None
|
||||
|
||||
return RateLimitResult(
|
||||
allowed = allowed,
|
||||
limit = limit,
|
||||
remaining = remaining,
|
||||
reset_after = reset_after,
|
||||
retry_after = retry_after if not allowed else None,
|
||||
)
|
||||
|
||||
async def get_token_bucket_state(
|
||||
self,
|
||||
key: str,
|
||||
) -> TokenBucketState | None:
|
||||
"""
|
||||
Get token bucket state if it exists
|
||||
"""
|
||||
client = await self._ensure_connected()
|
||||
bucket_key = f"{key}:bucket"
|
||||
|
||||
try:
|
||||
data = await client.hmget(
|
||||
bucket_key,
|
||||
"tokens",
|
||||
"last_refill",
|
||||
"capacity",
|
||||
"refill_rate"
|
||||
)
|
||||
|
||||
if data[0] is None:
|
||||
return None
|
||||
|
||||
return TokenBucketState(
|
||||
tokens = float(data[0]),
|
||||
last_refill = float(data[1]), # type: ignore[arg-type]
|
||||
capacity = int(data[2]), # type: ignore[arg-type]
|
||||
refill_rate = float(data[3]), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
except (RedisConnectionError, TimeoutError) as e:
|
||||
raise StorageError(
|
||||
operation = "get_token_bucket_state",
|
||||
backend = StorageType.REDIS,
|
||||
original_error = e,
|
||||
) from e
|
||||
|
||||
async def consume_token(
|
||||
self,
|
||||
key: str,
|
||||
capacity: int,
|
||||
refill_rate: float,
|
||||
tokens_to_consume: int = 1,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Attempt to consume tokens from bucket atomically
|
||||
"""
|
||||
now = time.time()
|
||||
|
||||
result = await self._execute_script(
|
||||
"token_bucket",
|
||||
keys = [key],
|
||||
args = [capacity,
|
||||
refill_rate,
|
||||
tokens_to_consume,
|
||||
now],
|
||||
)
|
||||
|
||||
allowed = bool(result[0])
|
||||
remaining = int(result[1])
|
||||
reset_after = float(result[2])
|
||||
retry_after = float(result[3]) if result[3] else None
|
||||
|
||||
return RateLimitResult(
|
||||
allowed = allowed,
|
||||
limit = capacity,
|
||||
remaining = remaining,
|
||||
reset_after = reset_after,
|
||||
retry_after = retry_after if not allowed else None,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""
|
||||
Close Redis connection.
|
||||
"""
|
||||
if self._client:
|
||||
await self._client.close()
|
||||
self._client = None
|
||||
|
||||
if self._pool:
|
||||
await self._pool.disconnect()
|
||||
self._pool = None
|
||||
|
||||
self._scripts_loaded = False
|
||||
self._script_shas.clear()
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""
|
||||
Check if Redis connection is healthy
|
||||
"""
|
||||
try:
|
||||
client = await self._ensure_connected()
|
||||
await client.ping()
|
||||
return True
|
||||
except (StorageError, RedisConnectionError, TimeoutError):
|
||||
return False
|
||||
|
||||
@property
|
||||
def storage_type(self) -> StorageType:
|
||||
"""
|
||||
Return storage type identifier
|
||||
"""
|
||||
return StorageType.REDIS
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""
|
||||
Check if client is connected
|
||||
"""
|
||||
return self._client is not None and self._scripts_loaded
|
||||
|
|
@ -0,0 +1,401 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
types.py
|
||||
"""
|
||||
# pylint: disable=unnecessary-ellipsis
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Protocol,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from starlette.requests import Request
|
||||
|
||||
|
||||
class Algorithm(StrEnum):
|
||||
"""
|
||||
Rate limiting algorithm selection
|
||||
"""
|
||||
SLIDING_WINDOW = "sliding_window"
|
||||
TOKEN_BUCKET = "token_bucket"
|
||||
FIXED_WINDOW = "fixed_window"
|
||||
LEAKY_BUCKET = "leaky_bucket"
|
||||
|
||||
|
||||
class FingerprintLevel(StrEnum):
|
||||
"""
|
||||
Preset fingerprinting intensity levels
|
||||
"""
|
||||
STRICT = "strict"
|
||||
NORMAL = "normal"
|
||||
RELAXED = "relaxed"
|
||||
CUSTOM = "custom"
|
||||
|
||||
|
||||
class DefenseMode(StrEnum):
|
||||
"""
|
||||
Global circuit breaker defense strategies
|
||||
"""
|
||||
ADAPTIVE = "adaptive"
|
||||
LOCKDOWN = "lockdown"
|
||||
CHALLENGE = "challenge"
|
||||
DISABLED = "disabled"
|
||||
|
||||
|
||||
class StorageType(StrEnum):
|
||||
"""
|
||||
Backend storage type selection
|
||||
"""
|
||||
REDIS = "redis"
|
||||
MEMORY = "memory"
|
||||
|
||||
|
||||
class Layer(StrEnum):
|
||||
"""
|
||||
Rate limiting defense layers
|
||||
"""
|
||||
USER = "user"
|
||||
ENDPOINT = "endpoint"
|
||||
GLOBAL = "global"
|
||||
|
||||
|
||||
@dataclass(frozen = True, slots = True)
|
||||
class RateLimitResult:
|
||||
"""
|
||||
Result of a rate limit check operation
|
||||
"""
|
||||
allowed: bool
|
||||
limit: int
|
||||
remaining: int
|
||||
reset_after: float
|
||||
retry_after: float | None = None
|
||||
|
||||
@property
|
||||
def headers(self) -> dict[str, str]:
|
||||
"""
|
||||
Generate IETF draft-compliant rate limit headers
|
||||
"""
|
||||
hdrs = {
|
||||
"RateLimit-Limit": str(self.limit),
|
||||
"RateLimit-Remaining": str(max(0,
|
||||
self.remaining)),
|
||||
"RateLimit-Reset": str(int(self.reset_after)),
|
||||
}
|
||||
if self.retry_after is not None:
|
||||
hdrs["Retry-After"] = str(int(self.retry_after))
|
||||
return hdrs
|
||||
|
||||
|
||||
@dataclass(frozen = True, slots = True)
|
||||
class RateLimitRule:
|
||||
"""
|
||||
A single rate limit rule defining requests allowed per time window
|
||||
"""
|
||||
requests: int
|
||||
window_seconds: int
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.requests <= 0:
|
||||
raise ValueError("requests must be positive")
|
||||
if self.window_seconds <= 0:
|
||||
raise ValueError("window_seconds must be positive")
|
||||
|
||||
@classmethod
|
||||
def parse(cls, rule_string: str) -> RateLimitRule:
|
||||
"""
|
||||
Parse rate limit string like '100/minute' or '1000/hour'
|
||||
"""
|
||||
parts = rule_string.strip().lower().split("/")
|
||||
if len(parts) != 2:
|
||||
raise ValueError(f"Invalid rate limit format: {rule_string}")
|
||||
|
||||
try:
|
||||
requests = int(parts[0])
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Invalid request count: {parts[0]}") from e
|
||||
|
||||
window_map = {
|
||||
"second": 1,
|
||||
"seconds": 1,
|
||||
"sec": 1,
|
||||
"s": 1,
|
||||
"minute": 60,
|
||||
"minutes": 60,
|
||||
"min": 60,
|
||||
"m": 60,
|
||||
"hour": 3600,
|
||||
"hours": 3600,
|
||||
"hr": 3600,
|
||||
"h": 3600,
|
||||
"day": 86400,
|
||||
"days": 86400,
|
||||
"d": 86400,
|
||||
}
|
||||
|
||||
window_str = parts[1].strip()
|
||||
if window_str not in window_map:
|
||||
raise ValueError(f"Unknown time unit: {window_str}")
|
||||
|
||||
return cls(
|
||||
requests = requests,
|
||||
window_seconds = window_map[window_str]
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
if self.window_seconds == 1:
|
||||
unit = "second"
|
||||
elif self.window_seconds == 60:
|
||||
unit = "minute"
|
||||
elif self.window_seconds == 3600:
|
||||
unit = "hour"
|
||||
elif self.window_seconds == 86400:
|
||||
unit = "day"
|
||||
else:
|
||||
unit = f"{self.window_seconds}s"
|
||||
return f"{self.requests}/{unit}"
|
||||
|
||||
|
||||
@dataclass(slots = True)
|
||||
class FingerprintData:
|
||||
"""
|
||||
Collected fingerprint data from a request
|
||||
"""
|
||||
ip: str
|
||||
ip_normalized: str
|
||||
user_agent: str | None = None
|
||||
accept_language: str | None = None
|
||||
accept_encoding: str | None = None
|
||||
headers_hash: str | None = None
|
||||
auth_identifier: str | None = None
|
||||
tls_fingerprint: str | None = None
|
||||
geo_asn: str | None = None
|
||||
|
||||
def to_composite_key(self, level: FingerprintLevel) -> str:
|
||||
"""
|
||||
Generate composite fingerprint key based on level.
|
||||
"""
|
||||
if level == FingerprintLevel.RELAXED:
|
||||
components = [self.ip_normalized]
|
||||
if self.auth_identifier:
|
||||
components.append(self.auth_identifier)
|
||||
return ":".join(filter(None, components))
|
||||
|
||||
if level == FingerprintLevel.NORMAL:
|
||||
components = [
|
||||
self.ip_normalized,
|
||||
self.user_agent or "",
|
||||
self.auth_identifier or "",
|
||||
]
|
||||
return ":".join(components)
|
||||
|
||||
components = [
|
||||
self.ip_normalized,
|
||||
self.user_agent or "",
|
||||
self.accept_language or "",
|
||||
self.accept_encoding or "",
|
||||
self.headers_hash or "",
|
||||
self.auth_identifier or "",
|
||||
self.tls_fingerprint or "",
|
||||
self.geo_asn or "",
|
||||
]
|
||||
return ":".join(components)
|
||||
|
||||
|
||||
@dataclass(slots = True)
|
||||
class WindowState:
|
||||
"""
|
||||
State for sliding window counter algorithm
|
||||
"""
|
||||
current_count: int = 0
|
||||
previous_count: int = 0
|
||||
current_window: int = 0
|
||||
window_seconds: int = 60
|
||||
|
||||
def weighted_count(self, elapsed_ratio: float) -> float:
|
||||
"""
|
||||
Calculate weighted count using sliding window interpolation.
|
||||
"""
|
||||
return self.previous_count * (
|
||||
1 - elapsed_ratio
|
||||
) + self.current_count
|
||||
|
||||
|
||||
@dataclass(slots = True)
|
||||
class TokenBucketState:
|
||||
"""
|
||||
State for token bucket algorithm
|
||||
"""
|
||||
tokens: float
|
||||
last_refill: float
|
||||
capacity: int
|
||||
refill_rate: float
|
||||
|
||||
|
||||
@dataclass(slots = True)
|
||||
class CircuitState:
|
||||
"""
|
||||
Global circuit breaker state
|
||||
"""
|
||||
is_open: bool = False
|
||||
failure_count: int = 0
|
||||
last_failure_time: float = 0.0
|
||||
half_open_requests: int = 0
|
||||
total_requests_in_window: int = 0
|
||||
|
||||
|
||||
@dataclass(slots = True)
|
||||
class DefenseContext:
|
||||
"""
|
||||
Context passed to defense strategies
|
||||
"""
|
||||
fingerprint: FingerprintData
|
||||
endpoint: str
|
||||
method: str
|
||||
is_authenticated: bool = False
|
||||
reputation_score: float = 1.0
|
||||
request_count_last_minute: int = 0
|
||||
|
||||
|
||||
@dataclass(slots = True)
|
||||
class RateLimitKey:
|
||||
"""
|
||||
Structured rate limit key components
|
||||
"""
|
||||
prefix: str = "ratelimit"
|
||||
version: str = "v1"
|
||||
layer: Layer = Layer.USER
|
||||
endpoint: str = ""
|
||||
identifier: str = ""
|
||||
window: int = 0
|
||||
|
||||
def build(self) -> str:
|
||||
"""
|
||||
Build the full Redis key string.
|
||||
"""
|
||||
return f"{self.prefix}:{self.version}:{self.layer.value}:{self.endpoint}:{self.identifier}:{self.window}"
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class StorageBackend(Protocol):
|
||||
"""
|
||||
Protocol for rate limit storage backends
|
||||
"""
|
||||
async def get_window_state(
|
||||
self,
|
||||
key: str,
|
||||
window_seconds: int,
|
||||
) -> WindowState:
|
||||
"""
|
||||
Get current sliding window state
|
||||
"""
|
||||
...
|
||||
|
||||
async def increment(
|
||||
self,
|
||||
key: str,
|
||||
window_seconds: int,
|
||||
limit: int,
|
||||
timestamp: float | None = None,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Atomically check and increment counter, returning result
|
||||
"""
|
||||
...
|
||||
|
||||
async def get_token_bucket_state(
|
||||
self,
|
||||
key: str,
|
||||
) -> TokenBucketState | None:
|
||||
"""
|
||||
Get token bucket state if it exists
|
||||
"""
|
||||
...
|
||||
|
||||
async def consume_token(
|
||||
self,
|
||||
key: str,
|
||||
capacity: int,
|
||||
refill_rate: float,
|
||||
tokens_to_consume: int = 1,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Attempt to consume tokens from bucket
|
||||
"""
|
||||
...
|
||||
|
||||
async def close(self) -> None:
|
||||
"""
|
||||
Close storage connections.
|
||||
"""
|
||||
...
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""
|
||||
Check if storage backend is healthy
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Fingerprinter(Protocol):
|
||||
"""
|
||||
Protocol for request fingerprinting strategies.
|
||||
"""
|
||||
async def extract(self, request: Request) -> FingerprintData:
|
||||
"""
|
||||
Extract fingerprint data from request.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RateLimitAlgorithm(Protocol):
|
||||
"""
|
||||
Protocol for rate limiting algorithms
|
||||
"""
|
||||
async def check(
|
||||
self,
|
||||
storage: StorageBackend,
|
||||
key: str,
|
||||
rule: RateLimitRule,
|
||||
timestamp: float | None = None,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Check if request is allowed under rate limit
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@dataclass
|
||||
class LimiterConfig:
|
||||
"""
|
||||
Main configuration for the rate limiter.
|
||||
"""
|
||||
default_rules: list[RateLimitRule] = field(default_factory = list)
|
||||
algorithm: Algorithm = Algorithm.SLIDING_WINDOW
|
||||
fingerprint_level: FingerprintLevel = FingerprintLevel.NORMAL
|
||||
defense_mode: DefenseMode = DefenseMode.ADAPTIVE
|
||||
fail_open: bool = True
|
||||
key_prefix: str = "ratelimit"
|
||||
include_headers: bool = True
|
||||
log_violations: bool = True
|
||||
|
||||
global_limit: RateLimitRule | None = None
|
||||
endpoint_limits: dict[str,
|
||||
list[RateLimitRule]] = field(
|
||||
default_factory = dict
|
||||
)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not self.default_rules:
|
||||
self.default_rules = [
|
||||
RateLimitRule.parse("100/minute"),
|
||||
RateLimitRule.parse("1000/hour"),
|
||||
]
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
__init__.py
|
||||
"""
|
||||
|
|
@ -0,0 +1,984 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
conftest.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Generator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, Request
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from fastapi_420.algorithms import create_algorithm
|
||||
from fastapi_420.algorithms.base import BaseAlgorithm
|
||||
from fastapi_420.algorithms.fixed_window import FixedWindowAlgorithm
|
||||
from fastapi_420.algorithms.sliding_window import SlidingWindowAlgorithm
|
||||
from fastapi_420.algorithms.token_bucket import TokenBucketAlgorithm
|
||||
from fastapi_420.config import (
|
||||
DefenseSettings,
|
||||
FingerprintSettings,
|
||||
RateLimiterSettings,
|
||||
StorageSettings,
|
||||
)
|
||||
from fastapi_420.defense.circuit_breaker import CircuitBreaker
|
||||
from fastapi_420.defense.layers import LayeredDefense
|
||||
from fastapi_420.dependencies import set_global_limiter
|
||||
from fastapi_420.exceptions import HTTP_420_ENHANCE_YOUR_CALM
|
||||
from fastapi_420.fingerprinting.auth import AuthExtractor
|
||||
from fastapi_420.fingerprinting.composite import CompositeFingerprinter
|
||||
from fastapi_420.fingerprinting.headers import HeadersExtractor
|
||||
from fastapi_420.fingerprinting.ip import IPExtractor
|
||||
from fastapi_420.limiter import RateLimiter
|
||||
from fastapi_420.middleware import RateLimitMiddleware
|
||||
from fastapi_420.storage import MemoryStorage
|
||||
from fastapi_420.types import (
|
||||
Algorithm,
|
||||
CircuitState,
|
||||
DefenseContext,
|
||||
DefenseMode,
|
||||
FingerprintData,
|
||||
FingerprintLevel,
|
||||
Layer,
|
||||
RateLimitKey,
|
||||
RateLimitResult,
|
||||
RateLimitRule,
|
||||
StorageType,
|
||||
TokenBucketState,
|
||||
WindowState,
|
||||
)
|
||||
|
||||
|
||||
WINDOW_SECOND = 1
|
||||
WINDOW_MINUTE = 60
|
||||
WINDOW_HOUR = 3600
|
||||
WINDOW_DAY = 86400
|
||||
|
||||
DEFAULT_LIMIT_REQUESTS = 100
|
||||
DEFAULT_LIMIT_WINDOW = WINDOW_MINUTE
|
||||
STRICT_LIMIT_REQUESTS = 10
|
||||
STRICT_LIMIT_WINDOW = WINDOW_MINUTE
|
||||
|
||||
TEST_IP_V4 = "192.168.1.100"
|
||||
TEST_IP_V4_PRIVATE = "10.0.0.1"
|
||||
TEST_IP_V6 = "2001:0db8:85a3:0000:0000:8a2e:0370:7334"
|
||||
TEST_IP_V6_NORMALIZED = "2001:db8:85a3::"
|
||||
TEST_IP_LOCALHOST = "127.0.0.1"
|
||||
|
||||
TEST_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
|
||||
TEST_ACCEPT_LANGUAGE = "en-US,en;q=0.9"
|
||||
TEST_ACCEPT_ENCODING = "gzip, deflate, br"
|
||||
|
||||
TEST_JWT_TOKEN = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyXzEyMyIsImV4cCI6OTk5OTk5OTk5OX0.signature"
|
||||
TEST_JWT_SUBJECT = "user_123"
|
||||
TEST_API_KEY = "sk-test-api-key-12345"
|
||||
TEST_SESSION_ID = "sess_abc123def456"
|
||||
|
||||
TEST_ENDPOINT = "/api/v1/test"
|
||||
TEST_METHOD = "GET"
|
||||
|
||||
CIRCUIT_THRESHOLD = 1000
|
||||
CIRCUIT_WINDOW = WINDOW_MINUTE
|
||||
CIRCUIT_RECOVERY = 30
|
||||
|
||||
KEY_PREFIX = "ratelimit"
|
||||
KEY_VERSION = "v1"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockScope:
|
||||
"""
|
||||
Mock ASGI scope for request creation
|
||||
"""
|
||||
type: str = "http"
|
||||
method: str = "GET"
|
||||
path: str = "/"
|
||||
query_string: bytes = b""
|
||||
headers: list[tuple[bytes, bytes]] = field(default_factory=list)
|
||||
client: tuple[str, int] | None = None
|
||||
route: Any = None
|
||||
|
||||
|
||||
class MockRoute:
|
||||
"""
|
||||
Mock route object for endpoint extraction
|
||||
"""
|
||||
def __init__(self, path: str = TEST_ENDPOINT) -> None:
|
||||
self.path = path
|
||||
|
||||
|
||||
class RequestFactory:
|
||||
"""
|
||||
Factory for creating mock Starlette Request objects
|
||||
"""
|
||||
@staticmethod
|
||||
def create(
|
||||
method: str = TEST_METHOD,
|
||||
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,
|
||||
include_route: bool = True,
|
||||
) -> Request:
|
||||
"""
|
||||
Create a mock Request with configurable parameters
|
||||
"""
|
||||
header_list: list[tuple[bytes, bytes]] = []
|
||||
|
||||
default_headers = {
|
||||
"host": "localhost",
|
||||
"user-agent": TEST_USER_AGENT,
|
||||
"accept": "*/*",
|
||||
"accept-language": TEST_ACCEPT_LANGUAGE,
|
||||
"accept-encoding": TEST_ACCEPT_ENCODING,
|
||||
}
|
||||
|
||||
if headers:
|
||||
default_headers.update(headers)
|
||||
|
||||
for key, value in default_headers.items():
|
||||
header_list.append((key.lower().encode(), value.encode()))
|
||||
|
||||
query_string = b""
|
||||
if query_params:
|
||||
query_string = "&".join(
|
||||
f"{k}={v}" for k, v in query_params.items()
|
||||
).encode()
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": method,
|
||||
"path": path,
|
||||
"query_string": query_string,
|
||||
"headers": header_list,
|
||||
"client": (client_ip, client_port),
|
||||
}
|
||||
|
||||
if include_route:
|
||||
scope["route"] = MockRoute(path)
|
||||
|
||||
request = Request(scope)
|
||||
|
||||
if cookies:
|
||||
request._cookies = cookies
|
||||
|
||||
return request
|
||||
|
||||
@staticmethod
|
||||
def with_auth(
|
||||
auth_type: str = "bearer",
|
||||
token: str = TEST_JWT_TOKEN,
|
||||
**kwargs: Any,
|
||||
) -> Request:
|
||||
"""
|
||||
Create request with authentication header
|
||||
"""
|
||||
headers = kwargs.pop("headers", {}) or {}
|
||||
|
||||
if auth_type == "bearer":
|
||||
headers["authorization"] = f"Bearer {token}"
|
||||
elif auth_type == "api_key":
|
||||
headers["x-api-key"] = token
|
||||
elif auth_type == "basic":
|
||||
headers["authorization"] = f"Basic {token}"
|
||||
|
||||
return RequestFactory.create(headers=headers, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def with_forwarded_for(
|
||||
forwarded_ips: list[str],
|
||||
real_ip: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Request:
|
||||
"""
|
||||
Create request with proxy headers
|
||||
"""
|
||||
headers = kwargs.pop("headers", {}) or {}
|
||||
headers["x-forwarded-for"] = ", ".join(forwarded_ips)
|
||||
|
||||
if real_ip:
|
||||
headers["x-real-ip"] = real_ip
|
||||
|
||||
return RequestFactory.create(headers=headers, **kwargs)
|
||||
|
||||
|
||||
class FingerprintFactory:
|
||||
"""
|
||||
Factory for creating FingerprintData instances
|
||||
"""
|
||||
@staticmethod
|
||||
def create(
|
||||
ip: str = TEST_IP_V4,
|
||||
ip_normalized: str | None = None,
|
||||
user_agent: str | None = TEST_USER_AGENT,
|
||||
accept_language: str | None = TEST_ACCEPT_LANGUAGE,
|
||||
accept_encoding: str | None = TEST_ACCEPT_ENCODING,
|
||||
headers_hash: str | None = None,
|
||||
auth_identifier: str | None = None,
|
||||
tls_fingerprint: str | None = None,
|
||||
geo_asn: str | None = None,
|
||||
) -> FingerprintData:
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def authenticated(
|
||||
auth_id: str = "user_123",
|
||||
hash_id: bool = True,
|
||||
**kwargs: Any,
|
||||
) -> FingerprintData:
|
||||
"""
|
||||
Create authenticated fingerprint
|
||||
"""
|
||||
identifier = auth_id
|
||||
if hash_id:
|
||||
identifier = hashlib.sha256(auth_id.encode()).hexdigest()[:16]
|
||||
|
||||
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)
|
||||
|
||||
@staticmethod
|
||||
def minimal(ip: str = TEST_IP_V4) -> FingerprintData:
|
||||
"""
|
||||
Create minimal fingerprint (IP only)
|
||||
"""
|
||||
return FingerprintData(
|
||||
ip=ip,
|
||||
ip_normalized=ip,
|
||||
)
|
||||
|
||||
|
||||
class RuleFactory:
|
||||
"""
|
||||
Factory for creating RateLimitRule instances
|
||||
"""
|
||||
@staticmethod
|
||||
def create(
|
||||
requests: int = DEFAULT_LIMIT_REQUESTS,
|
||||
window_seconds: int = DEFAULT_LIMIT_WINDOW,
|
||||
) -> RateLimitRule:
|
||||
"""
|
||||
Create RateLimitRule with defaults
|
||||
"""
|
||||
return RateLimitRule(requests=requests, window_seconds=window_seconds)
|
||||
|
||||
@staticmethod
|
||||
def per_second(requests: int = 10) -> RateLimitRule:
|
||||
return RateLimitRule(requests=requests, window_seconds=WINDOW_SECOND)
|
||||
|
||||
@staticmethod
|
||||
def per_minute(requests: int = 100) -> RateLimitRule:
|
||||
return RateLimitRule(requests=requests, window_seconds=WINDOW_MINUTE)
|
||||
|
||||
@staticmethod
|
||||
def per_hour(requests: int = 1000) -> RateLimitRule:
|
||||
return RateLimitRule(requests=requests, window_seconds=WINDOW_HOUR)
|
||||
|
||||
@staticmethod
|
||||
def per_day(requests: int = 10000) -> RateLimitRule:
|
||||
return RateLimitRule(requests=requests, window_seconds=WINDOW_DAY)
|
||||
|
||||
@staticmethod
|
||||
def strict() -> RateLimitRule:
|
||||
return RateLimitRule(
|
||||
requests=STRICT_LIMIT_REQUESTS,
|
||||
window_seconds=STRICT_LIMIT_WINDOW,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def parse(rule_string: str) -> RateLimitRule:
|
||||
return RateLimitRule.parse(rule_string)
|
||||
|
||||
|
||||
class ResultFactory:
|
||||
"""
|
||||
Factory for creating RateLimitResult instances
|
||||
"""
|
||||
@staticmethod
|
||||
def allowed(
|
||||
limit: int = DEFAULT_LIMIT_REQUESTS,
|
||||
remaining: int | None = None,
|
||||
reset_after: float = 60.0,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Create allowed result
|
||||
"""
|
||||
return RateLimitResult(
|
||||
allowed=True,
|
||||
limit=limit,
|
||||
remaining=remaining if remaining is not None else limit - 1,
|
||||
reset_after=reset_after,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def denied(
|
||||
limit: int = DEFAULT_LIMIT_REQUESTS,
|
||||
reset_after: float = 60.0,
|
||||
retry_after: float | None = None,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Create denied result
|
||||
"""
|
||||
return RateLimitResult(
|
||||
allowed=False,
|
||||
limit=limit,
|
||||
remaining=0,
|
||||
reset_after=reset_after,
|
||||
retry_after=retry_after or reset_after,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def near_limit(
|
||||
limit: int = DEFAULT_LIMIT_REQUESTS,
|
||||
remaining: int = 1,
|
||||
reset_after: float = 30.0,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
Create result near the limit
|
||||
"""
|
||||
return RateLimitResult(
|
||||
allowed=True,
|
||||
limit=limit,
|
||||
remaining=remaining,
|
||||
reset_after=reset_after,
|
||||
)
|
||||
|
||||
|
||||
class KeyFactory:
|
||||
"""
|
||||
Factory for creating RateLimitKey instances
|
||||
"""
|
||||
@staticmethod
|
||||
def create(
|
||||
prefix: str = KEY_PREFIX,
|
||||
version: str = KEY_VERSION,
|
||||
layer: Layer = Layer.USER,
|
||||
endpoint: str = TEST_ENDPOINT,
|
||||
identifier: str = TEST_IP_V4,
|
||||
window: int = WINDOW_MINUTE,
|
||||
) -> RateLimitKey:
|
||||
"""
|
||||
Create RateLimitKey with defaults
|
||||
"""
|
||||
return RateLimitKey(
|
||||
prefix=prefix,
|
||||
version=version,
|
||||
layer=layer,
|
||||
endpoint=endpoint,
|
||||
identifier=identifier,
|
||||
window=window,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def user_key(
|
||||
endpoint: str = TEST_ENDPOINT,
|
||||
identifier: str = TEST_IP_V4,
|
||||
window: int = WINDOW_MINUTE,
|
||||
) -> RateLimitKey:
|
||||
return KeyFactory.create(
|
||||
layer=Layer.USER,
|
||||
endpoint=endpoint,
|
||||
identifier=identifier,
|
||||
window=window,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def endpoint_key(
|
||||
endpoint: str = TEST_ENDPOINT,
|
||||
window: int = WINDOW_MINUTE,
|
||||
) -> RateLimitKey:
|
||||
return KeyFactory.create(
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class WindowStateFactory:
|
||||
"""
|
||||
Factory for creating WindowState instances
|
||||
"""
|
||||
@staticmethod
|
||||
def create(
|
||||
current_count: int = 0,
|
||||
previous_count: int = 0,
|
||||
current_window: int | None = None,
|
||||
window_seconds: int = WINDOW_MINUTE,
|
||||
) -> WindowState:
|
||||
"""
|
||||
Create WindowState with defaults
|
||||
"""
|
||||
if current_window is None:
|
||||
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,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def empty() -> WindowState:
|
||||
return WindowStateFactory.create()
|
||||
|
||||
@staticmethod
|
||||
def with_usage(current: int, previous: int = 0) -> WindowState:
|
||||
return WindowStateFactory.create(
|
||||
current_count=current,
|
||||
previous_count=previous,
|
||||
)
|
||||
|
||||
|
||||
class TokenBucketStateFactory:
|
||||
"""
|
||||
Factory for creating TokenBucketState instances
|
||||
"""
|
||||
@staticmethod
|
||||
def create(
|
||||
tokens: float = 100.0,
|
||||
last_refill: float | None = None,
|
||||
capacity: int = 100,
|
||||
refill_rate: float = 1.67,
|
||||
) -> TokenBucketState:
|
||||
"""
|
||||
Create TokenBucketState with defaults
|
||||
"""
|
||||
return TokenBucketState(
|
||||
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,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def empty(capacity: int = 100) -> TokenBucketState:
|
||||
return TokenBucketStateFactory.create(tokens=0.0, capacity=capacity)
|
||||
|
||||
|
||||
class DefenseContextFactory:
|
||||
"""
|
||||
Factory for creating DefenseContext instances
|
||||
"""
|
||||
@staticmethod
|
||||
def create(
|
||||
fingerprint: FingerprintData | None = None,
|
||||
endpoint: str = TEST_ENDPOINT,
|
||||
method: str = TEST_METHOD,
|
||||
is_authenticated: bool = False,
|
||||
reputation_score: float = 1.0,
|
||||
request_count_last_minute: int = 0,
|
||||
) -> DefenseContext:
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def authenticated(**kwargs: Any) -> DefenseContext:
|
||||
fp = FingerprintFactory.authenticated()
|
||||
return DefenseContextFactory.create(
|
||||
fingerprint=fp,
|
||||
is_authenticated=True,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def suspicious(
|
||||
reputation_score: float = 0.3,
|
||||
request_count: int = 500,
|
||||
) -> DefenseContext:
|
||||
return DefenseContextFactory.create(
|
||||
reputation_score=reputation_score,
|
||||
request_count_last_minute=request_count,
|
||||
)
|
||||
|
||||
|
||||
class CircuitStateFactory:
|
||||
"""
|
||||
Factory for creating CircuitState instances
|
||||
"""
|
||||
@staticmethod
|
||||
def create(
|
||||
is_open: bool = False,
|
||||
failure_count: int = 0,
|
||||
last_failure_time: float = 0.0,
|
||||
half_open_requests: int = 0,
|
||||
total_requests_in_window: int = 0,
|
||||
) -> CircuitState:
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def closed() -> CircuitState:
|
||||
return CircuitStateFactory.create()
|
||||
|
||||
@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(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rate_limiter_settings(
|
||||
storage_settings: StorageSettings,
|
||||
fingerprint_settings: FingerprintSettings,
|
||||
defense_settings: DefenseSettings,
|
||||
) -> RateLimiterSettings:
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def memory_storage() -> AsyncGenerator[MemoryStorage, None]:
|
||||
"""
|
||||
Create and manage MemoryStorage instance
|
||||
"""
|
||||
storage = MemoryStorage(max_keys=10000, cleanup_interval=60)
|
||||
await storage.start_cleanup_task()
|
||||
yield storage
|
||||
await storage.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sliding_window_algorithm() -> SlidingWindowAlgorithm:
|
||||
"""
|
||||
Create sliding window algorithm instance
|
||||
"""
|
||||
return SlidingWindowAlgorithm()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def token_bucket_algorithm() -> TokenBucketAlgorithm:
|
||||
"""
|
||||
Create token bucket algorithm instance
|
||||
"""
|
||||
return TokenBucketAlgorithm()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fixed_window_algorithm() -> FixedWindowAlgorithm:
|
||||
"""
|
||||
Create fixed window algorithm instance
|
||||
"""
|
||||
return FixedWindowAlgorithm()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ip_extractor() -> IPExtractor:
|
||||
"""
|
||||
Create IP extractor instance
|
||||
"""
|
||||
return IPExtractor(
|
||||
ipv6_prefix_length=64,
|
||||
trusted_proxies=[],
|
||||
trust_x_forwarded_for=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def headers_extractor() -> HeadersExtractor:
|
||||
"""
|
||||
Create headers extractor instance
|
||||
"""
|
||||
return HeadersExtractor(use_header_order=False, hash_length=16)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def composite_fingerprinter(
|
||||
ip_extractor: IPExtractor,
|
||||
headers_extractor: HeadersExtractor,
|
||||
auth_extractor: AuthExtractor,
|
||||
) -> CompositeFingerprinter:
|
||||
"""
|
||||
Create composite fingerprinter instance
|
||||
"""
|
||||
return CompositeFingerprinter(
|
||||
level=FingerprintLevel.NORMAL,
|
||||
ip_extractor=ip_extractor,
|
||||
headers_extractor=headers_extractor,
|
||||
auth_extractor=auth_extractor,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def rate_limiter(
|
||||
rate_limiter_settings: RateLimiterSettings,
|
||||
memory_storage: MemoryStorage,
|
||||
) -> AsyncGenerator[RateLimiter, None]:
|
||||
"""
|
||||
Create and manage RateLimiter instance
|
||||
"""
|
||||
limiter = RateLimiter(
|
||||
settings=rate_limiter_settings,
|
||||
storage=memory_storage,
|
||||
)
|
||||
await limiter.init()
|
||||
yield limiter
|
||||
await limiter.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def layered_defense(
|
||||
memory_storage: MemoryStorage,
|
||||
rate_limiter_settings: RateLimiterSettings,
|
||||
circuit_breaker: CircuitBreaker,
|
||||
) -> LayeredDefense:
|
||||
"""
|
||||
Create layered defense instance
|
||||
"""
|
||||
return LayeredDefense(
|
||||
storage=memory_storage,
|
||||
settings=rate_limiter_settings,
|
||||
circuit_breaker=circuit_breaker,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_request() -> Request:
|
||||
"""
|
||||
Create a default test request
|
||||
"""
|
||||
return RequestFactory.create()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def authenticated_request() -> Request:
|
||||
"""
|
||||
Create an authenticated test request
|
||||
"""
|
||||
return RequestFactory.with_auth()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_fingerprint() -> FingerprintData:
|
||||
"""
|
||||
Create a default test fingerprint
|
||||
"""
|
||||
return FingerprintFactory.create()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_rule() -> RateLimitRule:
|
||||
"""
|
||||
Create a default test rule
|
||||
"""
|
||||
return RuleFactory.per_minute()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def strict_rule() -> RateLimitRule:
|
||||
"""
|
||||
Create a strict test rule
|
||||
"""
|
||||
return RuleFactory.strict()
|
||||
|
||||
|
||||
def create_test_app(
|
||||
limiter: RateLimiter | None = None,
|
||||
with_middleware: bool = False,
|
||||
default_limit: str = "100/minute",
|
||||
) -> FastAPI:
|
||||
"""
|
||||
Create a test FastAPI application
|
||||
"""
|
||||
app = FastAPI(title="Test API")
|
||||
|
||||
@app.get("/")
|
||||
async def root() -> dict[str, str]:
|
||||
return {"message": "Hello World"}
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> dict[str, str]:
|
||||
return {"status": "healthy"}
|
||||
|
||||
@app.get(TEST_ENDPOINT)
|
||||
async def test_endpoint() -> dict[str, str]:
|
||||
return {"endpoint": "test"}
|
||||
|
||||
@app.post(TEST_ENDPOINT)
|
||||
async def test_endpoint_post() -> dict[str, str]:
|
||||
return {"created": "true"}
|
||||
|
||||
@app.get("/api/v1/protected")
|
||||
async def protected_endpoint() -> dict[str, str]:
|
||||
return {"protected": "true"}
|
||||
|
||||
if with_middleware and limiter:
|
||||
app.add_middleware(
|
||||
RateLimitMiddleware,
|
||||
limiter=limiter,
|
||||
default_limit=default_limit,
|
||||
)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_app() -> FastAPI:
|
||||
"""
|
||||
Create a basic test app without middleware
|
||||
"""
|
||||
return create_test_app()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def test_app_with_limiter(
|
||||
rate_limiter: RateLimiter,
|
||||
) -> FastAPI:
|
||||
"""
|
||||
Create a test app with rate limiting middleware
|
||||
"""
|
||||
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]:
|
||||
"""
|
||||
Create async HTTP client for testing
|
||||
"""
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=test_app),
|
||||
base_url="http://test",
|
||||
) as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def rate_limited_client(
|
||||
test_app_with_limiter: FastAPI,
|
||||
) -> AsyncGenerator[AsyncClient, None]:
|
||||
"""
|
||||
Create async HTTP client with rate limiting
|
||||
"""
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=test_app_with_limiter),
|
||||
base_url="http://test",
|
||||
) as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sync_client(test_app: FastAPI) -> Generator[TestClient, None, None]:
|
||||
"""
|
||||
Create sync test client
|
||||
"""
|
||||
with TestClient(test_app) as client:
|
||||
yield client
|
||||
|
||||
|
||||
def assert_rate_limit_headers(
|
||||
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
|
||||
|
||||
if expected_limit is not None:
|
||||
assert int(headers["RateLimit-Limit"]) == expected_limit
|
||||
|
||||
assert int(headers["RateLimit-Remaining"]) >= 0
|
||||
assert int(headers["RateLimit-Reset"]) >= 0
|
||||
|
||||
|
||||
def assert_420_response(
|
||||
response: Any,
|
||||
check_headers: bool = True,
|
||||
) -> None:
|
||||
"""
|
||||
Assert response is HTTP 420 with proper structure
|
||||
"""
|
||||
assert response.status_code == HTTP_420_ENHANCE_YOUR_CALM
|
||||
|
||||
if check_headers:
|
||||
assert "Retry-After" in response.headers or "RateLimit-Reset" in response.headers
|
||||
|
||||
|
||||
async def exhaust_rate_limit(
|
||||
storage: MemoryStorage,
|
||||
key: str,
|
||||
limit: int,
|
||||
window_seconds: int = WINDOW_MINUTE,
|
||||
) -> None:
|
||||
"""
|
||||
Helper to exhaust a rate limit by making requests
|
||||
"""
|
||||
for _ in range(limit):
|
||||
await storage.increment(
|
||||
key=key,
|
||||
window_seconds=window_seconds,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
|
||||
async def wait_for_window_reset(window_seconds: int = 1) -> None:
|
||||
"""
|
||||
Helper to wait for a window to reset
|
||||
"""
|
||||
await asyncio.sleep(window_seconds + 0.1)
|
||||
|
|
@ -0,0 +1,413 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
test_algorithms.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from fastapi_420.algorithms import create_algorithm
|
||||
from fastapi_420.algorithms.sliding_window import SlidingWindowAlgorithm
|
||||
from fastapi_420.algorithms.token_bucket import TokenBucketAlgorithm
|
||||
from fastapi_420.algorithms.fixed_window import FixedWindowAlgorithm
|
||||
from fastapi_420.storage import MemoryStorage
|
||||
from fastapi_420.types import Algorithm
|
||||
|
||||
from tests.conftest import (
|
||||
WINDOW_MINUTE,
|
||||
WINDOW_SECOND,
|
||||
RuleFactory,
|
||||
)
|
||||
|
||||
|
||||
class TestAlgorithmFactory:
|
||||
"""
|
||||
Tests for algorithm factory function
|
||||
"""
|
||||
def test_create_sliding_window(self) -> None:
|
||||
algo = create_algorithm(Algorithm.SLIDING_WINDOW)
|
||||
assert isinstance(algo, SlidingWindowAlgorithm)
|
||||
|
||||
def test_create_token_bucket(self) -> None:
|
||||
algo = create_algorithm(Algorithm.TOKEN_BUCKET)
|
||||
assert isinstance(algo, TokenBucketAlgorithm)
|
||||
|
||||
def test_create_fixed_window(self) -> None:
|
||||
algo = create_algorithm(Algorithm.FIXED_WINDOW)
|
||||
assert isinstance(algo, FixedWindowAlgorithm)
|
||||
|
||||
def test_create_from_string(self) -> None:
|
||||
algo = create_algorithm(Algorithm.SLIDING_WINDOW)
|
||||
assert algo.name == Algorithm.SLIDING_WINDOW.value
|
||||
|
||||
|
||||
class TestSlidingWindowAlgorithm:
|
||||
"""
|
||||
Tests for sliding window counter algorithm
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_algorithm_name(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
assert algo.name == "sliding_window"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_request_allowed(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="test_key",
|
||||
rule=rule,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 99
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_requests(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
for i in range(50):
|
||||
result = await algo.check(storage, f"multi_key", rule)
|
||||
assert result.allowed is True
|
||||
|
||||
assert result.remaining == 50
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limit_exceeded(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
|
||||
for _ in range(5):
|
||||
result = await algo.check(storage, "limit_key", rule)
|
||||
assert result.allowed is True
|
||||
|
||||
result = await algo.check(storage, "limit_key", rule)
|
||||
assert result.allowed is False
|
||||
assert result.retry_after is not None
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_keys_independent(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
|
||||
for _ in range(5):
|
||||
await algo.check(storage, "key_a", rule)
|
||||
|
||||
result_a = await algo.check(storage, "key_a", rule)
|
||||
result_b = await algo.check(storage, "key_b", rule)
|
||||
|
||||
assert result_a.allowed is False
|
||||
assert result_b.allowed is True
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_usage_empty(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute()
|
||||
|
||||
usage = await algo.get_current_usage(storage, "empty_key", rule)
|
||||
assert usage == 0
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_usage_after_requests(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
for _ in range(25):
|
||||
await algo.check(storage, "usage_key", rule)
|
||||
|
||||
usage = await algo.get_current_usage(storage, "usage_key", rule)
|
||||
assert usage >= 24
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_explicit_timestamp(self) -> None:
|
||||
algo = SlidingWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
fixed_time = 1000000.0
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="timestamp_key",
|
||||
rule=rule,
|
||||
timestamp=fixed_time,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
||||
await storage.close()
|
||||
|
||||
|
||||
class TestTokenBucketAlgorithm:
|
||||
"""
|
||||
Tests for token bucket algorithm
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_algorithm_name(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
assert algo.name == "token_bucket"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_request_allowed(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="bucket_test",
|
||||
rule=rule,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 99
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_burst_consumption(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=10, window_seconds=WINDOW_MINUTE)
|
||||
|
||||
for i in range(10):
|
||||
result = await algo.check(storage, "burst_key", rule)
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 10 - (i + 1)
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bucket_exhausted(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
|
||||
for _ in range(5):
|
||||
await algo.check(storage, "exhaust_key", rule)
|
||||
|
||||
result = await algo.check(storage, "exhaust_key", rule)
|
||||
assert result.allowed is False
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_keys_independent(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
|
||||
for _ in range(5):
|
||||
await algo.check(storage, "bucket_a", rule)
|
||||
|
||||
result_a = await algo.check(storage, "bucket_a", rule)
|
||||
result_b = await algo.check(storage, "bucket_b", rule)
|
||||
|
||||
assert result_a.allowed is False
|
||||
assert result_b.allowed is True
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_usage_empty(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
usage = await algo.get_current_usage(storage, "empty_bucket", rule)
|
||||
assert usage == 0
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_usage_after_consumption(self) -> None:
|
||||
algo = TokenBucketAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
for _ in range(25):
|
||||
await algo.check(storage, "usage_bucket", rule)
|
||||
|
||||
usage = await algo.get_current_usage(storage, "usage_bucket", rule)
|
||||
assert usage >= 24
|
||||
|
||||
await storage.close()
|
||||
|
||||
|
||||
class TestFixedWindowAlgorithm:
|
||||
"""
|
||||
Tests for fixed window counter algorithm
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_algorithm_name(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
assert algo.name == "fixed_window"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_request_allowed(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="fixed_test",
|
||||
rule=rule,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_requests(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
for _ in range(50):
|
||||
result = await algo.check(storage, "multi_fixed", rule)
|
||||
assert result.allowed is True
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limit_exceeded(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
|
||||
for _ in range(5):
|
||||
result = await algo.check(storage, "limit_fixed", rule)
|
||||
assert result.allowed is True
|
||||
|
||||
result = await algo.check(storage, "limit_fixed", rule)
|
||||
assert result.allowed is False
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_keys_independent(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
|
||||
for _ in range(5):
|
||||
await algo.check(storage, "fixed_a", rule)
|
||||
|
||||
result_a = await algo.check(storage, "fixed_a", rule)
|
||||
result_b = await algo.check(storage, "fixed_b", rule)
|
||||
|
||||
assert result_a.allowed is False
|
||||
assert result_b.allowed is True
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_usage_empty(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
usage = await algo.get_current_usage(storage, "empty_fixed", rule)
|
||||
assert usage == 0
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_usage_after_requests(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
for _ in range(25):
|
||||
await algo.check(storage, "usage_fixed", rule)
|
||||
|
||||
usage = await algo.get_current_usage(storage, "usage_fixed", rule)
|
||||
assert usage >= 0
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_explicit_timestamp(self) -> None:
|
||||
algo = FixedWindowAlgorithm()
|
||||
storage = MemoryStorage()
|
||||
rule = RuleFactory.per_minute(100)
|
||||
fixed_time = 1000000.0
|
||||
|
||||
result = await algo.check(
|
||||
storage=storage,
|
||||
key="timestamp_fixed",
|
||||
rule=rule,
|
||||
timestamp=fixed_time,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
||||
await storage.close()
|
||||
|
||||
|
||||
class TestAlgorithmComparison:
|
||||
"""
|
||||
Comparative tests between algorithms
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_algorithms_allow_first_request(self) -> None:
|
||||
algorithms = [
|
||||
SlidingWindowAlgorithm(),
|
||||
TokenBucketAlgorithm(),
|
||||
FixedWindowAlgorithm(),
|
||||
]
|
||||
rule = RuleFactory.per_minute(100)
|
||||
|
||||
for algo in algorithms:
|
||||
storage = MemoryStorage()
|
||||
result = await algo.check(storage, "compare_key", rule)
|
||||
assert result.allowed is True, f"{algo.name} failed first request"
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_algorithms_enforce_limit(self) -> None:
|
||||
algorithms = [
|
||||
SlidingWindowAlgorithm(),
|
||||
TokenBucketAlgorithm(),
|
||||
FixedWindowAlgorithm(),
|
||||
]
|
||||
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
|
||||
|
||||
for algo in algorithms:
|
||||
storage = MemoryStorage()
|
||||
for _ in range(5):
|
||||
await algo.check(storage, "enforce_key", rule)
|
||||
|
||||
result = await algo.check(storage, "enforce_key", rule)
|
||||
assert result.allowed is False, f"{algo.name} didn't enforce limit"
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_algorithms_have_correct_names(self) -> None:
|
||||
assert SlidingWindowAlgorithm().name == Algorithm.SLIDING_WINDOW.value
|
||||
assert TokenBucketAlgorithm().name == Algorithm.TOKEN_BUCKET.value
|
||||
assert FixedWindowAlgorithm().name == Algorithm.FIXED_WINDOW.value
|
||||
|
|
@ -0,0 +1,421 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
test_storage.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from fastapi_420.storage import MemoryStorage, create_storage
|
||||
from fastapi_420.config import StorageSettings
|
||||
from fastapi_420.types import StorageType
|
||||
|
||||
from tests.conftest import (
|
||||
WINDOW_MINUTE,
|
||||
WINDOW_SECOND,
|
||||
DEFAULT_LIMIT_REQUESTS,
|
||||
)
|
||||
|
||||
|
||||
class TestMemoryStorageBasic:
|
||||
"""
|
||||
Basic MemoryStorage creation and lifecycle tests
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_storage(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
assert storage.storage_type == StorageType.MEMORY
|
||||
assert storage.max_keys == 100_000
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_storage_custom_settings(self) -> None:
|
||||
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,
|
||||
)
|
||||
storage = MemoryStorage.from_settings(settings)
|
||||
assert storage.max_keys == 2000
|
||||
assert storage.cleanup_interval == 120
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_healthy(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
assert await storage.health_check() is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_after_close(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
await storage.close()
|
||||
assert await storage.health_check() is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_clears_data(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
await storage.increment("test_key", WINDOW_MINUTE, 100)
|
||||
await storage.close()
|
||||
assert len(storage._windows) == 0
|
||||
assert len(storage._buckets) == 0
|
||||
|
||||
|
||||
class TestMemoryStorageSlidingWindow:
|
||||
"""
|
||||
Tests for sliding window counter in MemoryStorage
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_first_request(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
result = await storage.increment(
|
||||
key="test",
|
||||
window_seconds=WINDOW_MINUTE,
|
||||
limit=100,
|
||||
)
|
||||
assert result.allowed is True
|
||||
assert result.limit == 100
|
||||
assert result.remaining == 99
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_multiple_requests(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
key = "multi_test"
|
||||
|
||||
for i in range(10):
|
||||
result = await storage.increment(key, WINDOW_MINUTE, 100)
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 100 - (i + 1)
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_reaches_limit(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
key = "limit_test"
|
||||
limit = 5
|
||||
|
||||
for _ in range(limit):
|
||||
result = await storage.increment(key, WINDOW_MINUTE, limit)
|
||||
assert result.allowed is True
|
||||
|
||||
result = await storage.increment(key, WINDOW_MINUTE, limit)
|
||||
assert result.allowed is False
|
||||
assert result.remaining == 0
|
||||
assert result.retry_after is not None
|
||||
assert result.retry_after > 0
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_with_explicit_timestamp(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
fixed_time = 1000000.0
|
||||
|
||||
result = await storage.increment(
|
||||
key="timestamp_test",
|
||||
window_seconds=WINDOW_MINUTE,
|
||||
limit=100,
|
||||
timestamp=fixed_time,
|
||||
)
|
||||
assert result.allowed is True
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_window_state_empty(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
state = await storage.get_window_state("nonexistent", WINDOW_MINUTE)
|
||||
assert state.current_count == 0
|
||||
assert state.previous_count == 0
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_window_state_with_data(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
key = "state_test"
|
||||
|
||||
for _ in range(5):
|
||||
await storage.increment(key, WINDOW_MINUTE, 100)
|
||||
|
||||
state = await storage.get_window_state(key, WINDOW_MINUTE)
|
||||
assert state.current_count == 5
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sliding_window_weighted_count(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
key = "weighted_test"
|
||||
window = 2
|
||||
limit = 100
|
||||
|
||||
base_time = 1000.0
|
||||
current_window = int(base_time // window)
|
||||
previous_window = current_window - 1
|
||||
|
||||
prev_key = f"{key}:{previous_window}"
|
||||
curr_key = f"{key}:{current_window}"
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
result = await storage.increment(
|
||||
key=key,
|
||||
window_seconds=window,
|
||||
limit=limit,
|
||||
timestamp=base_time + 1.0,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
await storage.close()
|
||||
|
||||
|
||||
class TestMemoryStorageTokenBucket:
|
||||
"""
|
||||
Tests for token bucket algorithm in MemoryStorage
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
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,
|
||||
)
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 99
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consume_token_multiple(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
key = "multi_bucket"
|
||||
|
||||
for i in range(10):
|
||||
result = await storage.consume_token(
|
||||
key=key,
|
||||
capacity=100,
|
||||
refill_rate=1.67,
|
||||
)
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 100 - (i + 1)
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consume_token_exhausted(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
key = "exhaust_bucket"
|
||||
capacity = 5
|
||||
|
||||
for _ in range(capacity):
|
||||
result = await storage.consume_token(
|
||||
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,
|
||||
)
|
||||
assert result.allowed is False
|
||||
assert result.retry_after is not None
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consume_token_refill(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
key = "refill_test"
|
||||
capacity = 10
|
||||
refill_rate = 10.0
|
||||
|
||||
for _ in range(capacity):
|
||||
await storage.consume_token(key, capacity, refill_rate)
|
||||
|
||||
result = await storage.consume_token(key, capacity, refill_rate)
|
||||
assert result.allowed is False
|
||||
|
||||
await asyncio.sleep(0.15)
|
||||
|
||||
result = await storage.consume_token(key, capacity, refill_rate)
|
||||
assert result.allowed is True
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_token_bucket_state_empty(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
state = await storage.get_token_bucket_state("nonexistent")
|
||||
assert state is None
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_token_bucket_state_exists(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
key = "bucket_state_test"
|
||||
|
||||
await storage.consume_token(key, capacity=100, refill_rate=1.67)
|
||||
|
||||
state = await storage.get_token_bucket_state(key)
|
||||
assert state is not None
|
||||
assert state.tokens == 99.0
|
||||
assert state.capacity == 100
|
||||
|
||||
await storage.close()
|
||||
|
||||
|
||||
class TestMemoryStorageMaxKeys:
|
||||
"""
|
||||
Tests for key eviction when max_keys is exceeded
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_keys_eviction(self) -> None:
|
||||
storage = MemoryStorage(max_keys=5)
|
||||
|
||||
for i in range(10):
|
||||
await storage.increment(f"key_{i}", WINDOW_MINUTE, 100)
|
||||
|
||||
assert len(storage._windows) <= 6
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lru_eviction_order(self) -> None:
|
||||
storage = MemoryStorage(max_keys=3)
|
||||
|
||||
await storage.increment("key_a", WINDOW_MINUTE, 100)
|
||||
await storage.increment("key_b", WINDOW_MINUTE, 100)
|
||||
await storage.increment("key_c", WINDOW_MINUTE, 100)
|
||||
|
||||
await storage.increment("key_d", WINDOW_MINUTE, 100)
|
||||
|
||||
keys = list(storage._windows.keys())
|
||||
assert not any("key_a" in k for k in keys)
|
||||
|
||||
await storage.close()
|
||||
|
||||
|
||||
class TestMemoryStorageCleanup:
|
||||
"""
|
||||
Tests for automatic cleanup of expired entries
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_entries(self) -> None:
|
||||
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,
|
||||
)
|
||||
storage._windows["valid_key"] = WindowEntry(
|
||||
count=10,
|
||||
window_start=1,
|
||||
expires_at=time.time() + 100,
|
||||
)
|
||||
|
||||
await storage._cleanup_expired()
|
||||
|
||||
assert "expired_key" not in storage._windows
|
||||
assert "valid_key" in storage._windows
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_task_starts(self) -> None:
|
||||
storage = MemoryStorage(cleanup_interval=1)
|
||||
await storage.start_cleanup_task()
|
||||
|
||||
assert storage._cleanup_task is not None
|
||||
assert not storage._cleanup_task.done()
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_task_stops_on_close(self) -> None:
|
||||
storage = MemoryStorage(cleanup_interval=1)
|
||||
await storage.start_cleanup_task()
|
||||
|
||||
task = storage._cleanup_task
|
||||
await storage.close()
|
||||
|
||||
assert storage._cleanup_task is None
|
||||
|
||||
|
||||
class TestStorageFactory:
|
||||
"""
|
||||
Tests for create_storage factory function
|
||||
"""
|
||||
def test_create_memory_storage_no_redis(self) -> 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,
|
||||
)
|
||||
storage = create_storage(settings)
|
||||
assert isinstance(storage, MemoryStorage)
|
||||
assert storage.max_keys == 5000
|
||||
|
||||
|
||||
class TestMemoryStorageConcurrency:
|
||||
"""
|
||||
Tests for concurrent access to MemoryStorage
|
||||
"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_increments(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
key = "concurrent_test"
|
||||
limit = 1000
|
||||
|
||||
async def increment() -> bool:
|
||||
result = await storage.increment(key, WINDOW_MINUTE, limit)
|
||||
return result.allowed
|
||||
|
||||
tasks = [increment() for _ in range(100)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
assert all(results)
|
||||
|
||||
state = await storage.get_window_state(key, WINDOW_MINUTE)
|
||||
assert state.current_count == 100
|
||||
|
||||
await storage.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_different_keys(self) -> None:
|
||||
storage = MemoryStorage()
|
||||
limit = 100
|
||||
|
||||
async def increment(key: str) -> int:
|
||||
for _ in range(10):
|
||||
await storage.increment(key, WINDOW_MINUTE, limit)
|
||||
state = await storage.get_window_state(key, WINDOW_MINUTE)
|
||||
return state.current_count
|
||||
|
||||
tasks = [increment(f"key_{i}") for i in range(10)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
assert all(r == 10 for r in results)
|
||||
|
||||
await storage.close()
|
||||
|
|
@ -0,0 +1,547 @@
|
|||
"""
|
||||
ⒸAngelaMos | 2025
|
||||
test_types.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from fastapi_420.types import (
|
||||
Algorithm,
|
||||
CircuitState,
|
||||
DefenseContext,
|
||||
DefenseMode,
|
||||
FingerprintData,
|
||||
FingerprintLevel,
|
||||
Layer,
|
||||
RateLimitKey,
|
||||
RateLimitResult,
|
||||
RateLimitRule,
|
||||
StorageType,
|
||||
TokenBucketState,
|
||||
WindowState,
|
||||
)
|
||||
|
||||
from tests.conftest import (
|
||||
KEY_PREFIX,
|
||||
KEY_VERSION,
|
||||
TEST_ENDPOINT,
|
||||
TEST_IP_V4,
|
||||
TEST_USER_AGENT,
|
||||
WINDOW_DAY,
|
||||
WINDOW_HOUR,
|
||||
WINDOW_MINUTE,
|
||||
WINDOW_SECOND,
|
||||
CircuitStateFactory,
|
||||
DefenseContextFactory,
|
||||
FingerprintFactory,
|
||||
KeyFactory,
|
||||
ResultFactory,
|
||||
RuleFactory,
|
||||
TokenBucketStateFactory,
|
||||
WindowStateFactory,
|
||||
)
|
||||
|
||||
|
||||
class TestAlgorithmEnum:
|
||||
"""
|
||||
Tests for Algorithm enum values
|
||||
"""
|
||||
def test_sliding_window_value(self) -> None:
|
||||
assert Algorithm.SLIDING_WINDOW.value == "sliding_window"
|
||||
|
||||
def test_token_bucket_value(self) -> None:
|
||||
assert Algorithm.TOKEN_BUCKET.value == "token_bucket"
|
||||
|
||||
def test_fixed_window_value(self) -> None:
|
||||
assert Algorithm.FIXED_WINDOW.value == "fixed_window"
|
||||
|
||||
def test_leaky_bucket_value(self) -> None:
|
||||
assert Algorithm.LEAKY_BUCKET.value == "leaky_bucket"
|
||||
|
||||
def test_enum_is_str_enum(self) -> None:
|
||||
assert isinstance(Algorithm.SLIDING_WINDOW, str)
|
||||
assert Algorithm.SLIDING_WINDOW == "sliding_window"
|
||||
|
||||
|
||||
class TestFingerprintLevelEnum:
|
||||
"""
|
||||
Tests for FingerprintLevel enum values
|
||||
"""
|
||||
def test_strict_value(self) -> None:
|
||||
assert FingerprintLevel.STRICT.value == "strict"
|
||||
|
||||
def test_normal_value(self) -> None:
|
||||
assert FingerprintLevel.NORMAL.value == "normal"
|
||||
|
||||
def test_relaxed_value(self) -> None:
|
||||
assert FingerprintLevel.RELAXED.value == "relaxed"
|
||||
|
||||
def test_custom_value(self) -> None:
|
||||
assert FingerprintLevel.CUSTOM.value == "custom"
|
||||
|
||||
|
||||
class TestDefenseModeEnum:
|
||||
"""
|
||||
Tests for DefenseMode enum values
|
||||
"""
|
||||
def test_adaptive_value(self) -> None:
|
||||
assert DefenseMode.ADAPTIVE.value == "adaptive"
|
||||
|
||||
def test_lockdown_value(self) -> None:
|
||||
assert DefenseMode.LOCKDOWN.value == "lockdown"
|
||||
|
||||
def test_challenge_value(self) -> None:
|
||||
assert DefenseMode.CHALLENGE.value == "challenge"
|
||||
|
||||
def test_disabled_value(self) -> None:
|
||||
assert DefenseMode.DISABLED.value == "disabled"
|
||||
|
||||
|
||||
class TestStorageTypeEnum:
|
||||
"""
|
||||
Tests for StorageType enum values
|
||||
"""
|
||||
def test_redis_value(self) -> None:
|
||||
assert StorageType.REDIS.value == "redis"
|
||||
|
||||
def test_memory_value(self) -> None:
|
||||
assert StorageType.MEMORY.value == "memory"
|
||||
|
||||
|
||||
class TestLayerEnum:
|
||||
"""
|
||||
Tests for Layer enum values
|
||||
"""
|
||||
def test_user_value(self) -> None:
|
||||
assert Layer.USER.value == "user"
|
||||
|
||||
def test_endpoint_value(self) -> None:
|
||||
assert Layer.ENDPOINT.value == "endpoint"
|
||||
|
||||
def test_global_value(self) -> None:
|
||||
assert Layer.GLOBAL.value == "global"
|
||||
|
||||
|
||||
class TestRateLimitRule:
|
||||
"""
|
||||
Tests for RateLimitRule dataclass and parsing
|
||||
"""
|
||||
def test_create_valid_rule(self) -> None:
|
||||
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)
|
||||
|
||||
def test_invalid_requests_negative(self) -> None:
|
||||
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)
|
||||
|
||||
def test_invalid_window_negative(self) -> None:
|
||||
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"]:
|
||||
rule = RateLimitRule.parse(f"10/{unit}")
|
||||
assert rule.requests == 10
|
||||
assert rule.window_seconds == WINDOW_SECOND
|
||||
|
||||
def test_parse_per_minute(self) -> None:
|
||||
for unit in ["minute", "minutes", "min", "m"]:
|
||||
rule = RateLimitRule.parse(f"100/{unit}")
|
||||
assert rule.requests == 100
|
||||
assert rule.window_seconds == WINDOW_MINUTE
|
||||
|
||||
def test_parse_per_hour(self) -> None:
|
||||
for unit in ["hour", "hours", "hr", "h"]:
|
||||
rule = RateLimitRule.parse(f"1000/{unit}")
|
||||
assert rule.requests == 1000
|
||||
assert rule.window_seconds == WINDOW_HOUR
|
||||
|
||||
def test_parse_per_day(self) -> None:
|
||||
for unit in ["day", "days", "d"]:
|
||||
rule = RateLimitRule.parse(f"10000/{unit}")
|
||||
assert rule.requests == 10000
|
||||
assert rule.window_seconds == WINDOW_DAY
|
||||
|
||||
def test_parse_with_whitespace(self) -> None:
|
||||
rule = RateLimitRule.parse(" 100 / minute ")
|
||||
assert rule.requests == 100
|
||||
assert rule.window_seconds == WINDOW_MINUTE
|
||||
|
||||
def test_parse_case_insensitive(self) -> None:
|
||||
rule = RateLimitRule.parse("100/MINUTE")
|
||||
assert rule.requests == 100
|
||||
assert rule.window_seconds == WINDOW_MINUTE
|
||||
|
||||
def test_parse_invalid_format_no_slash(self) -> None:
|
||||
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"):
|
||||
RateLimitRule.parse("100/per/minute")
|
||||
|
||||
def test_parse_invalid_request_count(self) -> None:
|
||||
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"):
|
||||
RateLimitRule.parse("100/fortnight")
|
||||
|
||||
def test_str_representation_second(self) -> None:
|
||||
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)
|
||||
assert str(rule) == "100/minute"
|
||||
|
||||
def test_str_representation_hour(self) -> None:
|
||||
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)
|
||||
assert str(rule) == "10000/day"
|
||||
|
||||
def test_str_representation_custom_window(self) -> None:
|
||||
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)
|
||||
with pytest.raises(AttributeError):
|
||||
rule.requests = 200
|
||||
|
||||
def test_factory_create(self) -> None:
|
||||
rule = RuleFactory.create()
|
||||
assert rule.requests == 100
|
||||
assert rule.window_seconds == WINDOW_MINUTE
|
||||
|
||||
def test_factory_per_second(self) -> None:
|
||||
rule = RuleFactory.per_second(10)
|
||||
assert rule.requests == 10
|
||||
assert rule.window_seconds == WINDOW_SECOND
|
||||
|
||||
def test_factory_strict(self) -> None:
|
||||
rule = RuleFactory.strict()
|
||||
assert rule.requests == 10
|
||||
assert rule.window_seconds == WINDOW_MINUTE
|
||||
|
||||
|
||||
class TestRateLimitResult:
|
||||
"""
|
||||
Tests for RateLimitResult dataclass and headers
|
||||
"""
|
||||
def test_allowed_result(self) -> None:
|
||||
result = RateLimitResult(
|
||||
allowed=True,
|
||||
limit=100,
|
||||
remaining=50,
|
||||
reset_after=30.0,
|
||||
)
|
||||
assert result.allowed is True
|
||||
assert result.limit == 100
|
||||
assert result.remaining == 50
|
||||
assert result.reset_after == 30.0
|
||||
assert result.retry_after is None
|
||||
|
||||
def test_denied_result(self) -> None:
|
||||
result = RateLimitResult(
|
||||
allowed=False,
|
||||
limit=100,
|
||||
remaining=0,
|
||||
reset_after=60.0,
|
||||
retry_after=60.0,
|
||||
)
|
||||
assert result.allowed is False
|
||||
assert result.remaining == 0
|
||||
assert result.retry_after == 60.0
|
||||
|
||||
def test_headers_basic(self) -> None:
|
||||
result = RateLimitResult(
|
||||
allowed=True,
|
||||
limit=100,
|
||||
remaining=50,
|
||||
reset_after=30.5,
|
||||
)
|
||||
headers = result.headers
|
||||
assert headers["RateLimit-Limit"] == "100"
|
||||
assert headers["RateLimit-Remaining"] == "50"
|
||||
assert headers["RateLimit-Reset"] == "30"
|
||||
assert "Retry-After" not in headers
|
||||
|
||||
def test_headers_with_retry_after(self) -> None:
|
||||
result = RateLimitResult(
|
||||
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,
|
||||
)
|
||||
headers = result.headers
|
||||
assert headers["RateLimit-Remaining"] == "0"
|
||||
|
||||
def test_frozen_immutable(self) -> None:
|
||||
result = ResultFactory.allowed()
|
||||
with pytest.raises(AttributeError):
|
||||
result.allowed = False
|
||||
|
||||
def test_factory_allowed(self) -> None:
|
||||
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)
|
||||
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)
|
||||
assert result.allowed is True
|
||||
assert result.remaining == 2
|
||||
|
||||
|
||||
class TestFingerprintData:
|
||||
"""
|
||||
Tests for FingerprintData and composite key generation
|
||||
"""
|
||||
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",
|
||||
)
|
||||
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)
|
||||
assert fp.ip == TEST_IP_V4
|
||||
assert fp.user_agent is None
|
||||
assert fp.auth_identifier is None
|
||||
|
||||
def test_composite_key_relaxed_ip_only(self) -> None:
|
||||
fp = FingerprintFactory.anonymous()
|
||||
key = fp.to_composite_key(FingerprintLevel.RELAXED)
|
||||
assert key == TEST_IP_V4
|
||||
|
||||
def test_composite_key_relaxed_with_auth(self) -> None:
|
||||
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")
|
||||
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]
|
||||
assert parts[2] == "user_123"
|
||||
|
||||
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",
|
||||
)
|
||||
key = fp.to_composite_key(FingerprintLevel.STRICT)
|
||||
parts = key.split(":")
|
||||
assert len(parts) == 8
|
||||
assert parts[0] == TEST_IP_V4
|
||||
|
||||
def test_factory_authenticated(self) -> None:
|
||||
fp = FingerprintFactory.authenticated(auth_id="testuser")
|
||||
assert fp.auth_identifier is not None
|
||||
assert len(fp.auth_identifier) == 16
|
||||
|
||||
def test_factory_minimal(self) -> None:
|
||||
fp = FingerprintFactory.minimal()
|
||||
assert fp.ip == TEST_IP_V4
|
||||
assert fp.user_agent is None
|
||||
|
||||
|
||||
class TestWindowState:
|
||||
"""
|
||||
Tests for WindowState and weighted count calculation
|
||||
"""
|
||||
def test_create_empty_state(self) -> None:
|
||||
state = WindowState()
|
||||
assert state.current_count == 0
|
||||
assert state.previous_count == 0
|
||||
|
||||
def test_weighted_count_start_of_window(self) -> None:
|
||||
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)
|
||||
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)
|
||||
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)
|
||||
weighted = state.weighted_count(0.25)
|
||||
assert weighted == 80 * 0.75 + 40
|
||||
|
||||
def test_factory_empty(self) -> None:
|
||||
state = WindowStateFactory.empty()
|
||||
assert state.current_count == 0
|
||||
assert state.previous_count == 0
|
||||
|
||||
def test_factory_with_usage(self) -> None:
|
||||
state = WindowStateFactory.with_usage(current=50, previous=100)
|
||||
assert state.current_count == 50
|
||||
assert state.previous_count == 100
|
||||
|
||||
|
||||
class TestTokenBucketState:
|
||||
"""
|
||||
Tests for TokenBucketState
|
||||
"""
|
||||
def test_create_state(self) -> None:
|
||||
state = TokenBucketState(
|
||||
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)
|
||||
assert state.tokens == 200.0
|
||||
assert state.capacity == 200
|
||||
|
||||
def test_factory_empty(self) -> None:
|
||||
state = TokenBucketStateFactory.empty()
|
||||
assert state.tokens == 0.0
|
||||
|
||||
|
||||
class TestCircuitState:
|
||||
"""
|
||||
Tests for CircuitState
|
||||
"""
|
||||
def test_default_closed(self) -> None:
|
||||
state = CircuitState()
|
||||
assert state.is_open is False
|
||||
assert state.failure_count == 0
|
||||
|
||||
def test_factory_closed(self) -> None:
|
||||
state = CircuitStateFactory.closed()
|
||||
assert state.is_open is False
|
||||
|
||||
def test_factory_open(self) -> None:
|
||||
state = CircuitStateFactory.open()
|
||||
assert state.is_open is True
|
||||
assert state.failure_count == 1
|
||||
assert state.last_failure_time > 0
|
||||
|
||||
|
||||
class TestDefenseContext:
|
||||
"""
|
||||
Tests for DefenseContext
|
||||
"""
|
||||
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,
|
||||
)
|
||||
assert context.endpoint == TEST_ENDPOINT
|
||||
assert context.is_authenticated is True
|
||||
assert context.reputation_score == 0.9
|
||||
|
||||
def test_factory_authenticated(self) -> None:
|
||||
context = DefenseContextFactory.authenticated()
|
||||
assert context.is_authenticated is True
|
||||
assert context.fingerprint.auth_identifier is not None
|
||||
|
||||
def test_factory_suspicious(self) -> None:
|
||||
context = DefenseContextFactory.suspicious(reputation_score=0.2)
|
||||
assert context.reputation_score == 0.2
|
||||
assert context.request_count_last_minute == 500
|
||||
|
||||
|
||||
class TestRateLimitKey:
|
||||
"""
|
||||
Tests for RateLimitKey and key building
|
||||
"""
|
||||
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,
|
||||
)
|
||||
built = key.build()
|
||||
expected = f"{KEY_PREFIX}:{KEY_VERSION}:user:{TEST_ENDPOINT}:{TEST_IP_V4}:{WINDOW_MINUTE}"
|
||||
assert built == expected
|
||||
|
||||
def test_factory_user_key(self) -> None:
|
||||
key = KeyFactory.user_key()
|
||||
built = key.build()
|
||||
assert ":user:" in built
|
||||
assert TEST_ENDPOINT in built
|
||||
|
||||
def test_factory_endpoint_key(self) -> None:
|
||||
key = KeyFactory.endpoint_key()
|
||||
built = key.build()
|
||||
assert ":endpoint:" in built
|
||||
assert ":global:" in built
|
||||
|
||||
def test_factory_global_key(self) -> None:
|
||||
key = KeyFactory.global_key()
|
||||
built = key.build()
|
||||
assert ":global:" in built
|
||||
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue