Rate limiting

loon/web/__init__.py

@@ -1,6 +1,8 @@
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.errors import RateLimitExceeded
 from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.middleware.cors import CORSMiddleware
 
@@ -14,10 +16,15 @@ async def lifespan(app: FastAPI):
     app.state.mqtt_manager = MQTTManager.setup_and_start()
     yield
 
-docs_url = "docs" if settings.web.show_docs else None
+limiter = Limiter(key_func=lambda request: request.client.host)
+
+docs_url = "/docs" if settings.web.show_docs else None
 
 app = FastAPI(lifespan=lifespan, docs_url=docs_url)
 
+app.state.limiter = limiter
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+
 app.add_middleware(
     AuthenticationMiddleware,
     backend=JWTAuthBackend(),

loon/web/auth/router.py

@@ -1,6 +1,7 @@
-from fastapi import APIRouter, HTTPException, status
+from fastapi import APIRouter, HTTPException, Request, status
 from sqlmodel import Session, select
 
+from loon.web import limiter
 from loon.web.auth.schema import JWTResponse, JWTRequest
 from loon.web.db import engine
 from loon.web.users.models import User
@@ -11,7 +12,8 @@ router = APIRouter(prefix="/auth", tags=["auth"])
 
 
 @router.post("/jwt/", response_model=JWTResponse)
-async def get_jwt(body: JWTRequest) -> JWTResponse:
+@limiter.limit("5/minute")
+async def get_jwt(request: Request, body: JWTRequest) -> JWTResponse:
     with Session(engine) as session:
         statement = select(User).where(User.internal_username == body.username)
         user = session.exec(statement).first()