From 6dadeeb4882a61a1ab9559113ffa93ac690fcc21 Mon Sep 17 00:00:00 2001
From: Rui Alves <alvesrui951@gmail.com>
Date: Sat, 21 Sep 2024 12:00:10 +0100
Subject: [PATCH] Use slowapi as the limiter

---
 docker/main/requirements-wheels.txt |  1 +
 frigate/api/auth.py                 | 26 ++++++++++++--------------
 frigate/api/fastapi_app.py          | 14 ++++++++------
 frigate/api/media.py                |  7 ++++++-
 4 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/docker/main/requirements-wheels.txt b/docker/main/requirements-wheels.txt
index 712fb0878..af0073fe6 100644
--- a/docker/main/requirements-wheels.txt
+++ b/docker/main/requirements-wheels.txt
@@ -2,6 +2,7 @@ click == 8.1.*
 # FastAPI
 starlette-context == 0.3.6
 fastapi == 0.115.0
+slowapi == 0.1.9
 imutils == 0.5.*
 joserfc == 1.0.*
 markupsafe == 2.1.*
diff --git a/frigate/api/auth.py b/frigate/api/auth.py
index 60713051c..efe2a5530 100644
--- a/frigate/api/auth.py
+++ b/frigate/api/auth.py
@@ -16,6 +16,7 @@ from fastapi import APIRouter, Request, Response
 from fastapi.responses import JSONResponse, RedirectResponse
 from joserfc import jwt
 from peewee import DoesNotExist
+from slowapi import Limiter
 
 from frigate.api.defs.app_body import (
     AppPostLoginBody,
@@ -74,17 +75,6 @@ def get_remote_addr(request: Request):
     return request.remote_addr or "127.0.0.1"
 
 
-# TODO: Rui
-# limiter = Limiter(
-#    get_remote_addr,
-#    storage_uri="memory://",
-# )
-
-
-def get_rate_limit(request: Request):
-    return request.app.frigate_config.auth.failed_login_rate_limit
-
-
 def get_jwt_secret() -> str:
     jwt_secret = None
     # check env var
@@ -306,9 +296,17 @@ def logout(request: Request):
     return response
 
 
+limiter = Limiter(key_func=get_remote_addr)
+
+
+def get_rate_limit(request: Request):
+    return request.app.frigate_config.auth.failed_login_rate_limit
+
+
 @router.post("/login")
-# TODO: Rui Implement limiter for FastAPI
-# @limiter.limit(get_rate_limit, deduct_when=lambda response: response.status_code == 400)
+# Ideally, this would be a decorator @limiter.limit(limit_value=get_rate_limit) but that way the request object is not passed to the method
+# See: https://github.com/laurentS/slowapi/issues/41
+# @limiter.limit(limit_value=get_rate_limit)
 def login(request: Request, body: AppPostLoginBody):
     JWT_COOKIE_NAME = request.app.frigate_config.auth.cookie_name
     JWT_COOKIE_SECURE = request.app.frigate_config.auth.cookie_secure
@@ -325,7 +323,7 @@ def login(request: Request, body: AppPostLoginBody):
     if verify_password(password, password_hash):
         expiration = int(time.time()) + JWT_SESSION_LENGTH
         encoded_jwt = create_encoded_jwt(user, expiration, request.app.jwt_token)
-        response = Response({}, 200)
+        response = Response("", 200)
         set_jwt_cookie(
             response, JWT_COOKIE_NAME, encoded_jwt, expiration, JWT_COOKIE_SECURE
         )
diff --git a/frigate/api/fastapi_app.py b/frigate/api/fastapi_app.py
index eaacb9326..20e1b5fba 100644
--- a/frigate/api/fastapi_app.py
+++ b/frigate/api/fastapi_app.py
@@ -4,12 +4,15 @@ from typing import Optional
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 from playhouse.sqliteq import SqliteQueueDatabase
+from slowapi import _rate_limit_exceeded_handler
+from slowapi.errors import RateLimitExceeded
+from slowapi.middleware import SlowAPIMiddleware
 from starlette_context import middleware, plugins
 from starlette_context.plugins import Plugin
 
 from frigate.api import app as main_app
 from frigate.api import auth, event, export, media, notification, preview, review
-from frigate.api.auth import get_jwt_secret
+from frigate.api.auth import get_jwt_secret, limiter
 from frigate.embeddings import EmbeddingsContext
 from frigate.events.external import ExternalEventProcessor
 from frigate.plus import PlusApi
@@ -53,6 +56,7 @@ def create_fastapi_app(
     )
 
     # update the request_address with the x-forwarded-for header from nginx
+    # https://starlette-context.readthedocs.io/en/latest/plugins.html#forwarded-for
     app.add_middleware(
         middleware.ContextMiddleware,
         plugins=(plugins.ForwardedForPlugin(),),
@@ -75,11 +79,9 @@ def create_fastapi_app(
             database.close()
         return response
 
-    # TODO: Rui
-    # initialize the rate limiter for the login endpoint
-    # limiter.init_app(app)
-    # if frigate_config.auth.failed_login_rate_limit is None:
-    #    limiter.enabled = False
+    app.state.limiter = limiter
+    app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+    app.add_middleware(SlowAPIMiddleware)
 
     # Routes
     app.include_router(main_app.router)
diff --git a/frigate/api/media.py b/frigate/api/media.py
index ffda91cec..4d11bd409 100644
--- a/frigate/api/media.py
+++ b/frigate/api/media.py
@@ -18,8 +18,8 @@ from fastapi import APIRouter, Path, Query, Request, Response
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 from peewee import DoesNotExist, fn
 from tzlocal import get_localzone_name
-from werkzeug.utils import secure_filename
 
+from frigate.api.defs.media_query_parameters import MediaLatestFrameQueryParams
 from frigate.api.defs.tags import Tags
 from frigate.config import FrigateConfig
 from frigate.const import (
@@ -39,6 +39,11 @@ logger = logging.getLogger(__name__)
 router = APIRouter(tags=[Tags.media])
 
 
+# TODO: Rui Implement or get from existing 3rd party
+def secure_filename(file_name: str):
+    return file_name
+
+
 @router.get("/media/camera/{camera_name}")
 def mjpeg_feed(
     request: Request,