From 4e9a3ee292322ad311a21208e1142672140e2eda Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Mon, 2 Jan 2023 08:02:35 +0800
Subject: [PATCH 01/11] Add fastapi-limiter to requirements.txt

---
 backend/requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/backend/requirements.txt b/backend/requirements.txt
index dd11aa18..fedf8ee3 100644
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -1,5 +1,6 @@
 alembic==1.8.1
 fastapi==0.88.0
+fastapi-limiter==0.1.5
 loguru==0.6.0
 numpy==1.22.4
 psycopg2-binary==2.9.5

From 144a5bc424679c5cff50365d764dfd7b2c4f98a3 Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Mon, 2 Jan 2023 08:04:59 +0800
Subject: [PATCH 02/11] Initialise redis connection for Fastapi limiter

---
 backend/main.py                 | 26 ++++++++++++++++++++++++++
 backend/oasst_backend/config.py |  4 ++++
 2 files changed, 30 insertions(+)

diff --git a/backend/main.py b/backend/main.py
index 9cf43701..81adcccd 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 from http import HTTPStatus
+from math import ceil
 from pathlib import Path
 from typing import Optional
 
@@ -7,6 +8,8 @@ import alembic.command
 import alembic.config
 import fastapi
 import pydantic
+import redis.asyncio as redis
+from fastapi_limiter import FastAPILimiter
 from loguru import logger
 from oasst_backend.api.deps import get_dummy_api_client
 from oasst_backend.api.v1.api import api_router
@@ -63,6 +66,29 @@ if settings.UPDATE_ALEMBIC:
             logger.exception("Alembic upgrade failed on startup")
 
 
+if settings.RATE_LIMIT:
+
+    @app.on_event("startup")
+    async def connect_redis():
+        async def http_callback(request: fastapi.Request, response: fastapi.Response, pexpire: int):
+            """Error callback function when too many requests"""
+            expire = ceil(pexpire / 1000)
+            raise OasstError(
+                f"Too Many Requests. Retry After {expire} seconds.",
+                OasstErrorCode.TOO_MANY_REQUESTS,
+                HTTPStatus.TOO_MANY_REQUESTS,
+            )
+
+        try:
+            redis_client = redis.from_url(
+                f"redis://{settings.REDIS_HOST}:{settings.REDIS_PORT}/0", encoding="utf-8", decode_responses=True
+            )
+            logger.info(f"Connected to {redis_client=}")
+            await FastAPILimiter.init(redis_client, http_callback=http_callback)
+        except Exception:
+            logger.exception("Failed to establish Redis connection")
+
+
 if settings.DEBUG_USE_SEED_DATA:
 
     @app.on_event("startup")
diff --git a/backend/oasst_backend/config.py b/backend/oasst_backend/config.py
index 602780be..cfe0d4c7 100644
--- a/backend/oasst_backend/config.py
+++ b/backend/oasst_backend/config.py
@@ -15,6 +15,10 @@ class Settings(BaseSettings):
     POSTGRES_DB: str = "postgres"
     DATABASE_URI: Optional[PostgresDsn] = None
 
+    RATE_LIMIT: bool = True
+    REDIS_HOST: str = "localhost"
+    REDIS_PORT: str = "6379"
+
     DEBUG_ALLOW_ANY_API_KEY: bool = False
     DEBUG_SKIP_API_KEY_CHECK: bool = False
     DEBUG_USE_SEED_DATA: bool = False

From 5daadd697789f0be11b7ae8f4d321b041ec993ac Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Mon, 2 Jan 2023 08:09:35 +0800
Subject: [PATCH 03/11] Add TOO_MANY_REQUEST to OasstErrorCode

---
 backend/oasst_backend/exceptions.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/backend/oasst_backend/exceptions.py b/backend/oasst_backend/exceptions.py
index f431b05b..7d4da6ef 100644
--- a/backend/oasst_backend/exceptions.py
+++ b/backend/oasst_backend/exceptions.py
@@ -18,6 +18,7 @@ class OasstErrorCode(IntEnum):
     DATABASE_URI_NOT_SET = 1
     API_CLIENT_NOT_AUTHORIZED = 2
     SERVER_ERROR = 3
+    TOO_MANY_REQUESTS = 429
 
     # 1000-2000: tasks endpoint
     TASK_INVALID_REQUEST_TYPE = 1000

From 95089d1ccf5b4c21082778e56863511597062e88 Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Tue, 3 Jan 2023 00:19:55 +0800
Subject: [PATCH 04/11] Add rate limiter for task request

---
 backend/oasst_backend/api/deps.py     | 60 ++++++++++++++++++++++++++-
 backend/oasst_backend/api/v1/tasks.py |  9 +++-
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/backend/oasst_backend/api/deps.py b/backend/oasst_backend/api/deps.py
index e0286ba3..cc5e08ae 100644
--- a/backend/oasst_backend/api/deps.py
+++ b/backend/oasst_backend/api/deps.py
@@ -4,8 +4,9 @@ from secrets import token_hex
 from typing import Generator
 from uuid import UUID
 
-from fastapi import Depends, Security
+from fastapi import Depends, Request, Response, Security
 from fastapi.security.api_key import APIKey, APIKeyHeader, APIKeyQuery
+from fastapi_limiter.depends import RateLimiter
 from loguru import logger
 from oasst_backend.config import settings
 from oasst_backend.database import engine
@@ -85,3 +86,60 @@ def get_trusted_api_client(
             http_status_code=HTTPStatus.FORBIDDEN,
         )
     return client
+
+
+class UserRateLimiter(RateLimiter):
+    def __init__(
+        self, times: int = 100, milliseconds: int = 0, seconds: int = 0, minutes: int = 1, hours: int = 0
+    ) -> None:
+        async def identifier(request: Request) -> str:
+            """Identify a request based on api_key and user.id"""
+            api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+            user = (await request.json()).get("user")
+            return f"{api_key}:{user.get('id')}"
+
+        super().__init__(times, milliseconds, seconds, minutes, hours, identifier)
+
+    async def __call__(self, request: Request, response: Response, api_key: str = Depends(get_api_key)) -> None:
+        # Skip if rate limiting is disabled
+        if not settings.RATE_LIMIT:
+            return
+
+        # Attempt to retrieve api_key and user information
+        api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+        user = (await request.json()).get("user")
+
+        # Skip when api_key and user information are not available
+        # (such that it will be handled by `APIClientRateLimiter`)
+        if not api_key or not user or not user.get("id"):
+            return
+
+        return await super().__call__(request, response)
+
+
+class APIClientRateLimiter(RateLimiter):
+    def __init__(
+        self, times: int = 10_000, milliseconds: int = 0, seconds: int = 0, minutes: int = 1, hours: int = 0
+    ) -> None:
+        async def identifier(request: Request) -> str:
+            """Identify a request based on api_key and user.id"""
+            api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+            return f"{api_key}"
+
+        super().__init__(times, milliseconds, seconds, minutes, hours, identifier)
+
+    async def __call__(self, request: Request, response: Response, api_key: str = Depends(get_api_key)) -> None:
+        # Skip if rate limiting is disabled
+        if not settings.RATE_LIMIT:
+            return
+
+        # Attempt to retrieve api_key and user information
+        api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+        user = (await request.json()).get("user")
+
+        # Skip if user information is available
+        # (such that it will be handled by `UserRateLimiter`)
+        if not api_key or user:
+            return
+
+        return await super().__call__(request, response)
diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py
index 636f0feb..c0614ce0 100644
--- a/backend/oasst_backend/api/v1/tasks.py
+++ b/backend/oasst_backend/api/v1/tasks.py
@@ -128,7 +128,14 @@ def generate_task(
     return task, message_tree_id, parent_message_id
 
 
-@router.post("/", response_model=protocol_schema.AnyTask)  # work with Union once more types are added
+@router.post(
+    "/",
+    response_model=protocol_schema.AnyTask,
+    dependencies=[
+        Depends(deps.UserRateLimiter(times=100, minutes=5)),
+        Depends(deps.APIClientRateLimiter(times=10_000, minutes=1)),
+    ],
+)  # work with Union once more types are added
 def request_task(
     *,
     db: Session = Depends(deps.get_db),

From 4e82991b2d6ecd22e7c49cca980a3d3fd4004cc8 Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Tue, 3 Jan 2023 00:22:37 +0800
Subject: [PATCH 05/11] Update docker compose to start redis on backend-dev and
 update REDIS_HOST in compose environment

---
 docker-compose.yaml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docker-compose.yaml b/docker-compose.yaml
index ed72c820..4fe1c2e7 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -4,7 +4,7 @@ services:
   # Use `docker compose up backend-dev --build --attach-dependencies` to start a database and work and the backend.
   backend-dev:
     image: sverrirab/sleep
-    depends_on: [db, adminer]
+    depends_on: [db, adminer, redis, redis-insights]
 
   # Use `docker compose up frontend-dev --build --attach-dependencies` to start all services needed to work on the frontend.
   frontend-dev:
@@ -90,6 +90,7 @@ services:
     image: oasst-backend
     environment:
       - POSTGRES_HOST=db
+      - REDIS_HOST=redis
       - DEBUG_SKIP_API_KEY_CHECK=True
       - DEBUG_USE_SEED_DATA=True
       - MAX_WORKERS=1

From af0313ad5d42ee60c144bb2b1bd8b6a0fc0d5d02 Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Mon, 2 Jan 2023 08:02:35 +0800
Subject: [PATCH 06/11] Add fastapi-limiter to requirements.txt

---
 backend/requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/backend/requirements.txt b/backend/requirements.txt
index dd11aa18..fedf8ee3 100644
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -1,5 +1,6 @@
 alembic==1.8.1
 fastapi==0.88.0
+fastapi-limiter==0.1.5
 loguru==0.6.0
 numpy==1.22.4
 psycopg2-binary==2.9.5

From 2a2968f6d41993ee54d43bb78837d1aec596c777 Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Mon, 2 Jan 2023 08:04:59 +0800
Subject: [PATCH 07/11] Initialise redis connection for Fastapi limiter

---
 backend/main.py                 | 26 ++++++++++++++++++++++++++
 backend/oasst_backend/config.py |  4 ++++
 2 files changed, 30 insertions(+)

diff --git a/backend/main.py b/backend/main.py
index d9c35095..4eb4a436 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -1,4 +1,5 @@
 from http import HTTPStatus
+from math import ceil
 from pathlib import Path
 from typing import Optional
 
@@ -6,6 +7,8 @@ import alembic.command
 import alembic.config
 import fastapi
 import pydantic
+import redis.asyncio as redis
+from fastapi_limiter import FastAPILimiter
 from loguru import logger
 from oasst_backend.api.deps import get_dummy_api_client
 from oasst_backend.api.v1.api import api_router
@@ -62,6 +65,29 @@ if settings.UPDATE_ALEMBIC:
             logger.exception("Alembic upgrade failed on startup")
 
 
+if settings.RATE_LIMIT:
+
+    @app.on_event("startup")
+    async def connect_redis():
+        async def http_callback(request: fastapi.Request, response: fastapi.Response, pexpire: int):
+            """Error callback function when too many requests"""
+            expire = ceil(pexpire / 1000)
+            raise OasstError(
+                f"Too Many Requests. Retry After {expire} seconds.",
+                OasstErrorCode.TOO_MANY_REQUESTS,
+                HTTPStatus.TOO_MANY_REQUESTS,
+            )
+
+        try:
+            redis_client = redis.from_url(
+                f"redis://{settings.REDIS_HOST}:{settings.REDIS_PORT}/0", encoding="utf-8", decode_responses=True
+            )
+            logger.info(f"Connected to {redis_client=}")
+            await FastAPILimiter.init(redis_client, http_callback=http_callback)
+        except Exception:
+            logger.exception("Failed to establish Redis connection")
+
+
 if settings.DEBUG_USE_SEED_DATA:
 
     @app.on_event("startup")
diff --git a/backend/oasst_backend/config.py b/backend/oasst_backend/config.py
index c675c148..fef59832 100644
--- a/backend/oasst_backend/config.py
+++ b/backend/oasst_backend/config.py
@@ -14,6 +14,10 @@ class Settings(BaseSettings):
     POSTGRES_DB: str = "postgres"
     DATABASE_URI: Optional[PostgresDsn] = None
 
+    RATE_LIMIT: bool = True
+    REDIS_HOST: str = "localhost"
+    REDIS_PORT: str = "6379"
+
     DEBUG_ALLOW_ANY_API_KEY: bool = False
     DEBUG_SKIP_API_KEY_CHECK: bool = False
     DEBUG_USE_SEED_DATA: bool = False

From bbf385bc076366c3e340a23bf64f53500838d071 Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Mon, 2 Jan 2023 08:09:35 +0800
Subject: [PATCH 08/11] Add TOO_MANY_REQUEST to OasstErrorCode

---
 backend/oasst_backend/exceptions.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/backend/oasst_backend/exceptions.py b/backend/oasst_backend/exceptions.py
index b6fb2d7f..49eeb088 100644
--- a/backend/oasst_backend/exceptions.py
+++ b/backend/oasst_backend/exceptions.py
@@ -17,6 +17,7 @@ class OasstErrorCode(IntEnum):
     DATABASE_URI_NOT_SET = 1
     API_CLIENT_NOT_AUTHORIZED = 2
     SERVER_ERROR = 3
+    TOO_MANY_REQUESTS = 429
 
     # 1000-2000: tasks endpoint
     TASK_INVALID_REQUEST_TYPE = 1000

From aa6d5786c610614164c8f8422c4d02266ae024ab Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Tue, 3 Jan 2023 00:19:55 +0800
Subject: [PATCH 09/11] Add rate limiter for task request

---
 backend/oasst_backend/api/deps.py     | 60 ++++++++++++++++++++++++++-
 backend/oasst_backend/api/v1/tasks.py |  9 +++-
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/backend/oasst_backend/api/deps.py b/backend/oasst_backend/api/deps.py
index fe26f0a6..b347da48 100644
--- a/backend/oasst_backend/api/deps.py
+++ b/backend/oasst_backend/api/deps.py
@@ -3,8 +3,9 @@ from secrets import token_hex
 from typing import Generator
 from uuid import UUID
 
-from fastapi import Depends, Security
+from fastapi import Depends, Request, Response, Security
 from fastapi.security.api_key import APIKey, APIKeyHeader, APIKeyQuery
+from fastapi_limiter.depends import RateLimiter
 from loguru import logger
 from oasst_backend.config import settings
 from oasst_backend.database import engine
@@ -84,3 +85,60 @@ def get_trusted_api_client(
             http_status_code=HTTPStatus.FORBIDDEN,
         )
     return client
+
+
+class UserRateLimiter(RateLimiter):
+    def __init__(
+        self, times: int = 100, milliseconds: int = 0, seconds: int = 0, minutes: int = 1, hours: int = 0
+    ) -> None:
+        async def identifier(request: Request) -> str:
+            """Identify a request based on api_key and user.id"""
+            api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+            user = (await request.json()).get("user")
+            return f"{api_key}:{user.get('id')}"
+
+        super().__init__(times, milliseconds, seconds, minutes, hours, identifier)
+
+    async def __call__(self, request: Request, response: Response, api_key: str = Depends(get_api_key)) -> None:
+        # Skip if rate limiting is disabled
+        if not settings.RATE_LIMIT:
+            return
+
+        # Attempt to retrieve api_key and user information
+        api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+        user = (await request.json()).get("user")
+
+        # Skip when api_key and user information are not available
+        # (such that it will be handled by `APIClientRateLimiter`)
+        if not api_key or not user or not user.get("id"):
+            return
+
+        return await super().__call__(request, response)
+
+
+class APIClientRateLimiter(RateLimiter):
+    def __init__(
+        self, times: int = 10_000, milliseconds: int = 0, seconds: int = 0, minutes: int = 1, hours: int = 0
+    ) -> None:
+        async def identifier(request: Request) -> str:
+            """Identify a request based on api_key and user.id"""
+            api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+            return f"{api_key}"
+
+        super().__init__(times, milliseconds, seconds, minutes, hours, identifier)
+
+    async def __call__(self, request: Request, response: Response, api_key: str = Depends(get_api_key)) -> None:
+        # Skip if rate limiting is disabled
+        if not settings.RATE_LIMIT:
+            return
+
+        # Attempt to retrieve api_key and user information
+        api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+        user = (await request.json()).get("user")
+
+        # Skip if user information is available
+        # (such that it will be handled by `UserRateLimiter`)
+        if not api_key or user:
+            return
+
+        return await super().__call__(request, response)
diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py
index aaa4a8c1..6571a5a9 100644
--- a/backend/oasst_backend/api/v1/tasks.py
+++ b/backend/oasst_backend/api/v1/tasks.py
@@ -127,7 +127,14 @@ def generate_task(
     return task, message_tree_id, parent_message_id
 
 
-@router.post("/", response_model=protocol_schema.AnyTask)  # work with Union once more types are added
+@router.post(
+    "/",
+    response_model=protocol_schema.AnyTask,
+    dependencies=[
+        Depends(deps.UserRateLimiter(times=100, minutes=5)),
+        Depends(deps.APIClientRateLimiter(times=10_000, minutes=1)),
+    ],
+)  # work with Union once more types are added
 def request_task(
     *,
     db: Session = Depends(deps.get_db),

From 0c1eaeb06671138ea2d234d32a3d40bba2e1d089 Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Tue, 3 Jan 2023 00:22:37 +0800
Subject: [PATCH 10/11] Update docker compose to start redis on backend-dev and
 update REDIS_HOST in compose environment

---
 docker-compose.yaml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docker-compose.yaml b/docker-compose.yaml
index c8d1377e..dc147c73 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -4,7 +4,7 @@ services:
   # Use `docker compose up backend-dev --build --attach-dependencies` to start a database and work and the backend.
   backend-dev:
     image: sverrirab/sleep
-    depends_on: [db, adminer]
+    depends_on: [db, adminer, redis, redis-insights]
 
   # Use `docker compose up frontend-dev --build --attach-dependencies` to start all services needed to work on the frontend.
   frontend-dev:
@@ -91,6 +91,7 @@ services:
     image: oasst-backend
     environment:
       - POSTGRES_HOST=db
+      - REDIS_HOST=redis
       - DEBUG_SKIP_API_KEY_CHECK=True
       - DEBUG_USE_SEED_DATA=True
       - MAX_WORKERS=1

From 4796146aa78b5ad66a9ffb23022db880d99fa262 Mon Sep 17 00:00:00 2001
From: kiritowu <wongzhaowu@gmail.com>
Date: Tue, 3 Jan 2023 01:17:36 +0800
Subject: [PATCH 11/11] Fix getting api_key using dependencies

---
 backend/oasst_backend/api/deps.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/backend/oasst_backend/api/deps.py b/backend/oasst_backend/api/deps.py
index b347da48..103de4fd 100644
--- a/backend/oasst_backend/api/deps.py
+++ b/backend/oasst_backend/api/deps.py
@@ -105,7 +105,6 @@ class UserRateLimiter(RateLimiter):
             return
 
         # Attempt to retrieve api_key and user information
-        api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
         user = (await request.json()).get("user")
 
         # Skip when api_key and user information are not available
@@ -133,7 +132,6 @@ class APIClientRateLimiter(RateLimiter):
             return
 
         # Attempt to retrieve api_key and user information
-        api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
         user = (await request.json()).get("user")
 
         # Skip if user information is available