Phase 1: harden rate limits — JWT-aware keys, per-endpoint caps, daily limits on all job types

bioai-platform/backend/app/deps.py

@@ -1,4 +1,34 @@
+import base64
+import json
+import logging
+
+from fastapi import Request
 from slowapi import Limiter
 from slowapi.util import get_remote_address
 
-limiter = Limiter(key_func=get_remote_address, default_limits=["30/minute"])
+logger = logging.getLogger(__name__)
+
+
+def _rate_limit_key(request: Request) -> str:
+    """Use user ID from JWT for authenticated requests, fall back to IP."""
+    auth = request.headers.get("Authorization", "")
+    if auth.startswith("Bearer "):
+        try:
+            token = auth[7:]
+            parts = token.split(".")
+            if len(parts) == 3:
+                payload = parts[1]
+                padding = 4 - len(payload) % 4
+                if padding != 4:
+                    payload += "=" * padding
+                decoded = base64.urlsafe_b64decode(payload)
+                claims = json.loads(decoded)
+                uid = claims.get("sub")
+                if uid:
+                    return f"user:{uid}"
+        except Exception:
+            pass
+    return get_remote_address(request)
+
+
+limiter = Limiter(key_func=_rate_limit_key, default_limits=["60/minute"])

bioai-platform/backend/app/routers/docking.py

@@ -3,13 +3,15 @@ from __future__ import annotations
 import json
 import math
 import re
-from fastapi import APIRouter, HTTPException, Depends
+from fastapi import APIRouter, HTTPException, Depends, Request
 from pydantic import BaseModel, Field
 from typing import Any, Optional
 
 from app.services.supabase import get_client
 from app.services.auth import require_user_id
 from app.services.ssrf import validate_url
+from app.services.rate_limit import check_daily_limit_docking
+from app.deps import limiter
 
 router = APIRouter(prefix="/api/docking", tags=["Docking"])
 _TABLE = "docking_jobs"
@@ -663,8 +665,9 @@ def _summarize_pose_interactions(protein_pdb: str, output_pdbqt: str) -> list[di
 # API endpoints
 # ---------------------------------------------------------------------------
 
-@router.post("/run", response_model=DockingJobResponse)
-async def create_docking_job(body: DockingJobCreate, user_id: str = Depends(require_user_id)):
+@router.post("/run", response_model=DockingJobResponse, dependencies=[Depends(check_daily_limit_docking)])
+@limiter.limit("5/minute")
+async def create_docking_job(request: Request, body: DockingJobCreate, user_id: str = Depends(require_user_id)):
     supabase = get_client()
     _prune_old(supabase)
 

bioai-platform/backend/app/routers/pipeline_v2.py

@@ -10,11 +10,13 @@ import uuid
 from datetime import datetime, timezone
 
 import httpx
-from fastapi import APIRouter, HTTPException
+from fastapi import APIRouter, HTTPException, Depends, Request
 from pydantic import BaseModel, Field
 from litellm import acompletion
 
 from app.config import settings
+from app.deps import limiter
+from app.services.rate_limit import check_daily_limit_pipelines
 from app.integrations.ncbi import blast as ncbi_blast
 from app.integrations.ncbi.parser import parse_blast_xml
 from app.services.validators import validate_fasta
@@ -59,8 +61,9 @@ class PipelineV2RunRequest(BaseModel):
     steps: list[str] = Field(default_factory=lambda: list(STEP_ORDER), description="Steps to run")
 
 
-@router.post("/run")
-async def run_pipeline_v2(req: PipelineV2RunRequest):
+@router.post("/run", dependencies=[Depends(check_daily_limit_pipelines)])
+@limiter.limit("3/minute")
+async def run_pipeline_v2(request: Request, req: PipelineV2RunRequest):
     validation = validate_fasta(req.sequence, "blast")
     if not validation.valid:
         raise HTTPException(status_code=400, detail=validation.error)

bioai-platform/backend/app/routers/pipelines.py

@@ -1,14 +1,16 @@
-from fastapi import APIRouter, HTTPException, Depends
+from fastapi import APIRouter, HTTPException, Depends, Request
 from pydantic import BaseModel
 from app.services.validators import validate_fasta
 from app.pipeline.definitions.protein_analysis import get_pipeline_definition
 from app.services.supabase import get_supabase
-from app.services.rate_limit import check_daily_limit
+from app.services.rate_limit import check_daily_limit, check_daily_limit_pipelines
 from app.services.auth import get_user_id
 from app.models.responses import PipelineRunResponse, PipelineDefinitionResponse
+from app.deps import limiter
 from datetime import datetime, timezone
 import uuid
 
+
 router = APIRouter()
 
 
@@ -20,8 +22,9 @@ class PipelineRunRequest(BaseModel):
     query_accession: str = ""
 
 
-@router.post("/run", response_model=PipelineRunResponse, dependencies=[Depends(check_daily_limit)])
-async def run_pipeline(req: PipelineRunRequest, user_id: str | None = Depends(get_user_id)):
+@router.post("/run", response_model=PipelineRunResponse, dependencies=[Depends(check_daily_limit_pipelines)])
+@limiter.limit("3/minute")
+async def run_pipeline(request: Request, req: PipelineRunRequest, user_id: str | None = Depends(get_user_id)):
     validation = validate_fasta(req.sequence, "blast")
     if not validation.valid:
         raise HTTPException(status_code=400, detail=validation.error)

bioai-platform/backend/app/routers/sequencing.py

@@ -6,13 +6,14 @@ import uuid
 from datetime import datetime, timezone, timedelta
 from typing import Optional
 
-from fastapi import APIRouter, HTTPException, Depends
+from fastapi import APIRouter, HTTPException, Depends, Request
 from pydantic import BaseModel
 
 from app.deps import limiter
 from app.services.supabase import get_supabase
 from app.services.auth import require_user_id
 from app.services.ssrf import validate_url
+from app.services.rate_limit import check_daily_limit_sequencing
 
 logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/api/sequencing", tags=["sequencing"])
@@ -131,8 +132,9 @@ async def _worker(job_id: str) -> None:
 VALID_DEMO = {"synthetic", "demo", "test"}
 
 
-@router.post("/run")
-async def run_sequencing(req: SequencingRequest, user_id: str = Depends(require_user_id)):
+@router.post("/run", dependencies=[Depends(check_daily_limit_sequencing)])
+@limiter.limit("3/minute")
+async def run_sequencing(request: Request, req: SequencingRequest, user_id: str = Depends(require_user_id)):
     if not req.fastq_url.strip():
         raise HTTPException(400, detail="fastq_url is required")
     if req.fastq_url.lower() not in VALID_DEMO:

bioai-platform/backend/app/services/rate_limit.py

@@ -1,3 +1,10 @@
+"""Rate limiting helpers.
+
+- `check_daily_limit` counts today's jobs across all three tables.
+- `check_daily_limit_pipelines`, `check_daily_limit_docking`, `check_daily_limit_sequencing`
+  count per-table for tighter per-feature caps.
+"""
+
 import base64
 import json
 import logging
@@ -31,16 +38,12 @@ def _extract_user_id_from_request(request: Request) -> str | None:
         return None
 
 
-async def check_daily_limit(request: Request) -> None:
-    user_id = _extract_user_id_from_request(request)
-    if not user_id:
-        logger.info("Rate limit check skipped: no authenticated user session")
-        return
-
+async def _count_today_jobs(user_id: str, table: str) -> int:
+    """Count today's jobs for a user in a specific table."""
     try:
         today = date.today().isoformat()
         url = (
-            f"{settings.SUPABASE_URL}/rest/v1/jobs"
+            f"{settings.SUPABASE_URL}/rest/v1/{table}"
             f"?user_id=eq.{user_id}"
             f"&created_at=gte.{today}T00:00:00"
             f"&select=id"
@@ -55,21 +58,62 @@ async def check_daily_limit(request: Request) -> None:
             resp.raise_for_status()
 
         content_range = resp.headers.get("content-range", "*/0")
-        total = int(content_range.split("/")[-1])
-
+        return int(content_range.split("/")[-1])
     except HTTPException:
         raise
     except Exception as e:
-        logger.warning(f"Rate limit check skipped (Supabase unreachable): {e}")
+        logger.warning("Daily count check skipped (table %s): %s", table, e)
+        return 0
+
+
+async def _enforce_limit(request: Request, table: str, limit: int, label: str) -> None:
+    user_id = _extract_user_id_from_request(request)
+    if not user_id:
         return
 
+    total = await _count_today_jobs(user_id, table)
+    if total >= limit:
+        raise HTTPException(
+            status_code=429,
+            detail={
+                "error": "daily_limit_exceeded",
+                "message": f"You've used all {limit} daily {label}. Resets at midnight UTC.",
+                "used": total,
+                "limit": limit,
+            },
+        )
+
+
+async def check_daily_limit(request: Request) -> None:
+    """Global daily limit across all tables (existing behavior)."""
+    user_id = _extract_user_id_from_request(request)
+    if not user_id:
+        return
+
+    tables = ["jobs", "docking_jobs", "sequencing_jobs"]
+    total = 0
+    for t in tables:
+        total += await _count_today_jobs(user_id, t)
+
     if total >= settings.DAILY_LIMIT:
         raise HTTPException(
             status_code=429,
             detail={
                 "error": "daily_limit_exceeded",
-                "message": f"You've used all {settings.DAILY_LIMIT} daily analyses. Resets at midnight IST.",
+                "message": f"You've used all {settings.DAILY_LIMIT} daily analyses. Resets at midnight UTC.",
                 "used": total,
                 "limit": settings.DAILY_LIMIT,
             },
         )
+
+
+async def check_daily_limit_pipelines(request: Request) -> None:
+    await _enforce_limit(request, "jobs", 10, "pipeline runs")
+
+
+async def check_daily_limit_docking(request: Request) -> None:
+    await _enforce_limit(request, "docking_jobs", 10, "docking jobs")
+
+
+async def check_daily_limit_sequencing(request: Request) -> None:
+    await _enforce_limit(request, "sequencing_jobs", 5, "sequencing jobs")