Implement ratelimits

config-example.yml

@@ -14,6 +14,11 @@ kanae:
   # Port that the server binds to.
   # Defaults to 8000
   port: 8000
+
+  # Global ratelimit for Kanae
+  # Syntax is amount/minute. Ex: 1/minute
+  ratelimits:
+    - "10/second"
 
 
   # Prometheus exporter for Kanae. The following keys are used in order to control
@@ -34,4 +39,8 @@ kanae:
 # The PostgreSQL connection URI that is used to connect to the database
 # The URI must be valid, and components will need to be quoted.
 # See https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING
-postgres_uri: "postgresql://user:password@localhost:5432/user"
+postgres_uri: "postgresql://user:password@localhost:5432/user"
+
+# The Redis connection URI that is used
+# See https://github.com/redis/lettuce/wiki/Redis-URI-and-connection-details#uri-syntax for URI syntax
+redis_uri: "redis://localhost:6379/0"

docker/docker-compose.dev.yml

@@ -14,5 +14,13 @@ services:
     ports:
       - 5432:5432
 
+  redis:
+    container_name: kanae_redis
+    image: redis:7-alpine
+    healthcheck:
+      test: redis-cli ping || exit 1
+    ports:
+      - 6379:6379
+
 volumes:
-  database:
+  database:

docker/docker-compose.prod.yml

@@ -8,11 +8,12 @@ services:
       # Do not edit the next line. If you want to change the path of the configuration file, please edit the CONFIG_LOCATION variable
       - ${CONFIG_LOCATION}:/kanae/server/config.yml
     ports:
-      - 9619:9619
+      - 8000:8000
     depends_on:
       - database
+      - redis
     # Safety script to fully wait until PostgreSQL is up
-    command: sh -c '/kanae/wait-for database:5432 -- echo "[Wait-for] PostgreSQL is fully up. Starting Kanae." && /kanae/start.sh'
+    command: sh -c '/kanae/wait-for database:5432 -- echo "[Wait-for] PostgreSQL is fully up. Waiting for Redis" && /kanae/wait-for redis:6379 -- echo "[Wait-for] Both PostgreSQL and Redis are fully ready. Starting up Kanae" && /kanae/start.sh'
     restart: always
 
   database:
@@ -34,6 +35,13 @@ services:
       start_period: 5m
     restart: always
 
+  redis:
+    container_name: kanae_redis
+    image: redis:7-alpine
+    healthcheck:
+      test: redis-cli ping || exit 1
+    restart: always
+
   kanae-prometheus:
     container_name: kanae_prometheus
     ports:
@@ -55,4 +63,4 @@ services:
       - grafana-data:/var/lib/grafana
 
 volumes:
-  database:
+  database:

docker/docker-compose.yml

@@ -8,11 +8,11 @@ services:
       # Do not edit the next line. If you want to change the path of the configuration file, please edit the CONFIG_LOCATION variable
       - ${CONFIG_LOCATION}:/kanae/server/config.yml
     ports:
-      - 9619:9619
+      - 8000:8000
     depends_on:
       - database
     # Safety script to fully wait until PostgreSQL is up
-    command: sh -c '/kanae/wait-for database:5432 -- echo "[Wait-for] PostgreSQL is fully up. Starting Kanae." && /kanae/start.sh'
+    command: sh -c '/kanae/wait-for database:5432 -- echo "[Wait-for] PostgreSQL is fully up. Waiting for Redis" && /kanae/wait-for redis:6379 -- echo "[Wait-for] Both PostgreSQL and Redis are fully ready. Starting up Kanae" && /kanae/start.sh'
     restart: always
 
   database:
@@ -34,5 +34,11 @@ services:
       start_period: 5m
     restart: always
 
+  redis:
+    container_name: kanae_redis
+    image: redis:7-alpine
+    healthcheck:
+      test: redis-cli ping || exit 1
+
 volumes:
   database:

requirements.txt

@@ -2,4 +2,6 @@ fastapi>=0.111.0,<1
 uvicorn[standard]>=0.12.0,<1
 asyncpg>=0.29.0,<1
 click>=8.1.7,<9
-typing-extensions>=4.12.2,<5
+typing-extensions>=4.12.2,<5
+slowapi>=0.1.9,<1
+redis>=5.0.8,<6.0.0

server/launcher.py

@@ -6,6 +6,8 @@
 import uvicorn
 from core import Kanae
 from routes import router
+from slowapi import _rate_limit_exceeded_handler
+from slowapi.errors import RateLimitExceeded
 from utils.config import KanaeConfig
 from uvicorn.supervisors import Multiprocess
 
@@ -14,6 +16,9 @@
 
 app = Kanae(config=config)
 app.include_router(router)
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)  # type: ignore
+app.state.limiter = router.limiter
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()

server/routes/__init__.py

@@ -1,9 +1,9 @@
 import importlib
 from pkgutil import iter_modules
 
-from fastapi import APIRouter
+from utils.router import KanaeRouter
 
-router = APIRouter()
+router = KanaeRouter()
 route_modules = [module.name for module in iter_modules(__path__, f"{__package__}.")]
 
 for route in route_modules:

server/routes/user.py

@@ -1,6 +1,6 @@
-from fastapi import APIRouter
 from pydantic import BaseModel
 from utils.request import RouteRequest
+from utils.router import KanaeRouter
 
 
 class NotFound(BaseModel):
@@ -11,7 +11,7 @@ class GetUser(BaseModel):
     user: str
 
 
-router = APIRouter(prefix="/users", tags=["Users"])
+router = KanaeRouter(prefix="/users", tags=["Users"])
 
 
 @router.get(
@@ -20,6 +20,7 @@ class GetUser(BaseModel):
     responses={200: {"model": GetUser}, 404: {"model": NotFound}},
     name="Get users",
 )
+@router.limiter.limit("1/minute")
 async def get_users(request: RouteRequest) -> GetUser:
     query = "SELECT 1;"
     status = await request.app.pool.execute(query)

server/utils/router.py

@@ -0,0 +1,40 @@
+from pathlib import Path
+
+from fastapi import APIRouter
+from pydantic import BaseModel
+from slowapi import Limiter
+from slowapi.util import get_remote_address
+
+from .config import KanaeConfig
+
+CONFIG_PATH = Path(__file__).parents[1] / "config.yml"
+
+
+class PartialConfig(BaseModel, frozen=True):
+    redis_uri: str
+    ratelimits: list[str]
+    dev_mode: bool = False
+
+
+class KanaeRouter(APIRouter):
+    limiter: Limiter
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        # This isn't my favorite implementation, but will do for now - Noelle
+        self._config = self._load_config()
+        self.limiter = Limiter(
+            key_func=get_remote_address,
+            storage_uri=self._config.redis_uri,
+            default_limits=self._config.ratelimits,  # type: ignore
+            enabled=not self._config.dev_mode,
+        )
+
+    def _load_config(self) -> PartialConfig:
+        config = KanaeConfig(CONFIG_PATH)
+        return PartialConfig(
+            redis_uri=config["redis_uri"],
+            ratelimits=config["kanae"]["ratelimits"],
+            dev_mode=config["kanae"]["dev_mode"],
+        )