Add rate limit for habit tick action

beaverhabits/frontend/components.py

@@ -13,7 +13,7 @@
 from beaverhabits.logging import logger
 from beaverhabits.storage.dict import DAY_MASK, MONTH_MASK
 from beaverhabits.storage.storage import CheckedRecord, Habit, HabitList, HabitStatus
-from beaverhabits.utils import WEEK_DAYS
+from beaverhabits.utils import WEEK_DAYS, ratelimiter
 
 strptime = datetime.datetime.strptime
 
@@ -92,6 +92,8 @@ async def note_tick(habit: Habit, day: datetime.date) -> bool | None:
     return record.done
 
 
+@ratelimiter(limit=30, window=30)
+@ratelimiter(limit=10, window=1)
 async def habit_tick(habit: Habit, day: datetime.date, value: bool):
     # Avoid duplicate tick
     record = habit.record_by(day)

beaverhabits/routes.py

@@ -1,11 +1,11 @@
 from typing import Optional
 
-from fastapi import Depends, FastAPI, Request
+from fastapi import Depends, FastAPI, HTTPException, Request, status
 from fastapi.responses import RedirectResponse
 from nicegui import app, ui
 
 from beaverhabits.frontend.import_page import import_ui_page
-from beaverhabits.frontend.layout import custom_header
+from beaverhabits.frontend.layout import custom_header, redirect
 from beaverhabits.frontend.order_page import order_page_ui
 
 from . import const, views
@@ -21,7 +21,6 @@
 from .frontend.cal_heatmap_page import heatmap_page
 from .frontend.habit_page import habit_page_ui
 from .frontend.index_page import index_page_ui
-from .logging import logger
 from .storage.meta import GUI_ROOT_PATH
 from .utils import dummy_days, get_user_today_date
 
@@ -53,6 +52,9 @@ async def demo_order_page() -> None:
 async def demo_habit_page(habit_id: str) -> None:
     today = await get_user_today_date()
     habit = await views.get_session_habit(habit_id)
+    if habit is None:
+        redirect("")
+        return
     habit_page_ui(today, habit)
 
 
@@ -61,6 +63,9 @@ async def demo_habit_page(habit_id: str) -> None:
 async def demo_habit_page_heatmap(habit_id: str) -> None:
     today = await get_user_today_date()
     habit = await views.get_session_habit(habit_id)
+    if habit is None:
+        redirect("")
+        return
     heatmap_page(today, habit)
 
 
@@ -194,6 +199,10 @@ async def try_register():
 
 
 def init_gui_routes(fastapi_app: FastAPI):
+    def handle_exception(exception: Exception):
+        if isinstance(exception, HTTPException):
+            if exception.status_code == status.HTTP_429_TOO_MANY_REQUESTS:
+                ui.notify(f"An error occurred: {exception}", type="negative")
 
     @app.middleware("http")
     async def AuthMiddleware(request: Request, call_next):
@@ -217,6 +226,7 @@ async def AuthMiddleware(request: Request, call_next):
         return response
 
     app.add_static_files("/statics", "statics")
+    app.on_exception(handle_exception)
     ui.run_with(
         fastapi_app,
         title=const.PAGE_TITLE,

beaverhabits/storage/dict.py

@@ -175,7 +175,7 @@ def __hash__(self) -> int:
         return hash(self.id)
 
     def __str__(self) -> str:
-        return f"{self.name}<{self.status.value}>"
+        return f"[{self.id}]{self.name}<{self.status.value}>"
 
     __repr__ = __str__
 

beaverhabits/utils.py

@@ -1,8 +1,12 @@
 import datetime
 import hashlib
+import time
 
 import pytz
+from cachetools import TTLCache
+from fastapi import HTTPException
 from nicegui import app, ui
+from starlette import status
 
 from beaverhabits.logging import logger
 
@@ -44,3 +48,37 @@ def generate_short_hash(name: str) -> str:
     h.update(name.encode())
     h.update(str(datetime.datetime.now()).encode())
     return h.hexdigest()[:6]
+
+
+def ratelimiter(limit: int, window: int):
+    if window <= 0 or window > 60 * 60:
+        raise ValueError("Window must be between 1 and 3600 seconds.")
+    cache = TTLCache(maxsize=128, ttl=60 * 60)
+
+    def decorator(func):
+        async def wrapper(*args, **kwargs):
+            current_time = time.time()
+            key = f"{args}_{kwargs}"
+
+            # Update timestamps
+            if key not in cache:
+                cache[key] = [current_time]
+            else:
+                cache[key].append(current_time)
+            cache[key] = [i for i in cache[key] if i >= current_time - window]
+
+            # Check with threshold
+            if len(cache[key]) > limit:
+                logger.warning(
+                    f"Rate limit exceeded for {func.__name__} with key {key}"
+                )
+                raise HTTPException(
+                    status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+                    detail=f"Rate limit exceeded. Try again later.",
+                )
+
+            return await func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

beaverhabits/views.py

@@ -2,7 +2,7 @@
 import json
 import random
 
-from fastapi import HTTPException, Request
+from fastapi import HTTPException
 from nicegui import app, ui
 
 from beaverhabits.app.auth import (
@@ -42,14 +42,14 @@ def get_session_habit_list() -> HabitList | None:
     return session_storage.get_user_habit_list()
 
 
-async def get_session_habit(habit_id: str) -> Habit:
+async def get_session_habit(habit_id: str) -> Habit | None:
     habit_list = get_session_habit_list()
     if habit_list is None:
-        raise HTTPException(status_code=404, detail="Habit list not found")
+        return None
 
     habit = await habit_list.get_habit_by(habit_id)
     if habit is None:
-        raise HTTPException(status_code=404, detail="Habit not found")
+        return None
 
     return habit
 

pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "sentry-sdk[fastapi]<3.0.0,>=2.15.0",
     "python-dateutil<3.0.0.0,>=2.9.0.post0",
     "loguru>=0.7.3",
+    "cachetools>=5.5.1",
 ]
 description = ""
 readme = "README.md"

tests/utils.py

@@ -1 +0,0 @@
-import datetime