async def test_on_auth_error_default(): rate_limit = RateLimitMiddleware( hello_world, auth_func, RedisBackend(), { r"/": [Rule(group="admin")], }, ) async with httpx.AsyncClient( app=rate_limit, base_url="http://testserver") as client: # type: httpx.AsyncClient response = await client.get("/", headers={ "user": "******", "group": "default" }) assert response.status_code == 200 assert response.text == "Hello world!" # No headers result in EmptyInformation with pytest.raises(EmptyInformation): await client.get("/", headers=None) # Raise the right exception with pytest.raises(AssertionError): await client.get("/", headers={"user": "******", "group": "-"})
async def test_on_auth_error_with_handler(): rate_limit = RateLimitMiddleware( hello_world, auth_func, RedisBackend(), { r"/": [Rule(group="admin")], }, on_auth_error=handle_auth_error, ) async with httpx.AsyncClient( app=rate_limit, base_url="http://testserver") as client: # type: httpx.AsyncClient response = await client.get("/", headers={ "user": "******", "group": "default" }) assert response.status_code == 200 assert response.text == "Hello world!" response = await client.get("/", headers=None) assert response.status_code == 401 assert response.text == ""
def setup(app): app.add_middleware( RateLimitMiddleware, authenticate=ratelimit, backend=RedisBackend(host="redis"), config={r'^/ts/v1/rainmeter': [Rule(minute=59, block_time=60)]}, on_blocked=handle_429)
async def test_multiple(redisbackend): await StrictRedis().flushdb() rate_limit = RateLimitMiddleware( hello_world, auth_func, redisbackend(), {r"/multiple": [Rule(second=1, minute=3)]}, ) async with httpx.AsyncClient( app=rate_limit, base_url="http://testserver") as client: # type: httpx.AsyncClient # multiple 1/s and 3/min # 1 3 response = await client.get("/multiple", headers={ "user": "******", "group": "default" }) assert response.status_code == 200 # 1-1 3-1 = 0 2 response = await client.get("/multiple", headers={ "user": "******", "group": "default" }) assert response.status_code == 429 await asyncio.sleep(1) # 0+1 2-1 = 1 1 response = await client.get("/multiple", headers={ "user": "******", "group": "default" }) assert response.status_code == 200 # 1-1 1-1 = 0 0 response = await client.get("/multiple", headers={ "user": "******", "group": "default" }) assert response.status_code == 429 await asyncio.sleep(2) # 0+1 0+0 = 1 0 response = await client.get("/multiple", headers={ "user": "******", "group": "default" }) assert response.status_code == 429
async def auth_function(scope: Any) -> Tuple[str, str]: try: return client_ip(scope) except Exception: return "default", "default" app.add_middleware( RateLimitMiddleware, authenticate=auth_function, backend=RedisBackend( host=get_settings().ratelimit_redis_host, port=get_settings().ratelimit_redis_port, ), config={ r"^/login/check": [Rule(minute=20)], }, ) templates = Jinja2Templates(directory=os.path.join(DIR, "templates")) mflog.set_config() @cache def get_auth_controller() -> AuthenticationController: auth_backend = make_auth_backend() return AuthenticationController(auth_backend) def check_redirect_url(url: str): parsed_url = urlparse(url) if parsed_url.netloc != "":
async def test_redis(): await StrictRedis().flushdb() rate_limit = RateLimitMiddleware( hello_world, auth_func, RedisBackend(), { r"/second_limit": [Rule(second=1), Rule(group="admin")], r"/minute.*": [Rule(minute=1), Rule(group="admin")], r"/block": [Rule(second=1, block_time=5)], }, ) async with httpx.AsyncClient( app=rate_limit, base_url="http://testserver") as client: # type: httpx.AsyncClient response = await client.get("/") assert response.status_code == 200 response = await client.get("/second_limit", headers={ "user": "******", "group": "default" }) assert response.status_code == 200 response = await client.get("/second_limit", headers={ "user": "******", "group": "default" }) assert response.status_code == 429 response = await client.get("/second_limit", headers={ "user": "******", "group": "admin" }) assert response.status_code == 200 await asyncio.sleep(1) response = await client.get("/second_limit", headers={ "user": "******", "group": "default" }) assert response.status_code == 200 response = await client.get("/minute_limit", headers={ "user": "******", "group": "default" }) assert response.status_code == 200 response = await client.get("/minute_limit", headers={ "user": "******", "group": "default" }) assert response.status_code == 429 response = await client.get("/minute_limit", headers={ "user": "******", "group": "admin" }) assert response.status_code == 200 response = await client.get("/block", headers={ "user": "******", "group": "default" }) assert response.status_code == 200 response = await client.get("/block", headers={ "user": "******", "group": "default" }) assert response.status_code == 429 await asyncio.sleep(1) response = await client.get("/block", headers={ "user": "******", "group": "default" }) assert response.status_code == 429 response = await client.get("/block", headers={ "user": "******", "group": "default" }) assert response.status_code == 429 await asyncio.sleep(4) response = await client.get("/block", headers={ "user": "******", "group": "default" }) assert response.status_code == 200
from fastapi import APIRouter, Response from fastapi.responses import PlainTextResponse from ratelimit import Rule from api.routes import snfgenerator, apiv2 router_rate_limits = { r'^/snfgenerator/': [Rule(second=1)], r'^/api/v0.2/xxx/game/': [Rule(second=0.2)] } router = APIRouter() router.include_router(snfgenerator.router, tags=["SNFGenerator"], prefix="/snfgenerator") router.include_router(apiv2.router, tags=["API v2.0"], prefix="/api/v0.2/xxx/game") @router.get("/crossdomain.xml") @router.get("/{path:path}/crossdomain.xml") def get_cross_domain(path=None): return Response( content= '<cross-domain-policy><allow-access-from domain="*"/></cross-domain-policy>', media_type="application/xml") @router.get("/api/ping", response_class=PlainTextResponse) def ping():
If there is no user information, it should raise `EmptyInformation`. If there is no group information, it should return "default". """ return "1", "default" app = FastAPI(exception_handlers={404: not_found}) if os.getenv("ENVIRONMENT") != "testing": app.add_middleware( RateLimitMiddleware, authenticate=AUTH_FUNCTION, backend=RedisBackend(), config={ r"^/person": [Rule(second=1), Rule(group="default")], r"^/": [Rule(minute=5), Rule(group="default")], }, ) def get_db(): db = SessionLocal() try: yield db finally: db.close() @app.post("/person/")
redis_password, redis_socket = REDIS_URL.replace("redis://:", "").split("@") redis_host, redis_port = redis_socket.split(":") redis_port = int(redis_port) app = FastAPI() app.openapi = custom_openapi # type: ignore # Add rate limiting middleware app.add_middleware( CustomRateLimitMiddleware, authenticate=create_jwt_auth(key=SECRET_KEY, algorithms=[HASH_ALGORITHM]), backend=RedisBackend(host=redis_host, port=redis_port, password=redis_password), config={ r"^/$": [Rule(second=5, group="default"), Rule(group="unlimited")], r"^/genres.*": [Rule(second=5, group="default"), Rule(group="unlimited")], r"^/artists$": [Rule(second=2, group="default"), Rule(group="unlimited")], r"^/artists/[0-9]+$": [ Rule(second=2, group="default"), Rule(group="unlimited"), ], r"^/preferences.*": [Rule(second=1, group="default"), Rule(group="unlimited")], r"^/recommendations.*": [