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_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": "-"})
def get_application() -> FastAPI: application = FastAPI(debug=True, title="Smart Server", description="MP session and game handler", version="0.0.1", docs_url="/api/docs", redoc_url="/api/redocs") application.add_middleware( CORSMiddleware, allow_origins=ALLOWED_HOSTS or ["*"], allow_credentials=True, allow_methods=["GET", "POST"], allow_headers=["*"], ) application.add_middleware(SessionMiddleware, secret_key=SECRET_KEY, same_site='None') ''' Note: If the API is a proxy to base play url (eg: play.localhost/cjs/...), it's better to use from_session else, for generic purpose, you can rate limit using IP address too. ''' #AUTH_FUNCTION = client_ip or from_session application.add_middleware( RateLimitMiddleware, authenticate=AUTH_FUNCTION, backend=RedisBackend(), config=router_rate_limits #{ # r"^/second_limit": [Rule(second=1), Rule(group="admin")], # r"^/minute_limit": [Rule(minute=1), Rule(group="admin")], #}, ) application.add_event_handler("startup", create_start_app_handler(application)) application.add_event_handler("shutdown", create_stop_app_handler(application)) application.add_exception_handler(HTTPException, http_error_handler) application.add_exception_handler(RequestValidationError, http422_error_handler) application.include_router(api_router) return application
app.mount("/static", StaticFiles(directory=os.path.join(DIR, "static")), name="static") if get_settings().ratelimit_redis_host is not None: 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)
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
""" Resolve the user's unique identifier and the user's group from ASGI SCOPE. 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()
# Get Redis credentials 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.*":