Merge pull request #6 from adamlaw669/feat/security-and-isolation-tests

Feat/security and isolation tests

Author: adamlawal

somba/api/middleware/idempotency.py

@@ -1,33 +1,222 @@
-"""Idempotency validation middleware for mutating API requests."""
+"""Idempotency middleware.
+
+Requires an Idempotency-Key on mutating requests, and — for authenticated
+requests — stores the response the first time a key is seen and *replays* it
+when the same key is reused. This makes POST/PATCH/DELETE safe to retry:
+
+  - same key + same request body   -> the stored response is replayed
+  - same key + different body       -> 409 (the key was reused for a new request)
+  - key still in flight             -> 409 (a concurrent request holds it)
+
+The store lives in idempotency_records (one row per merchant+key+method+path).
+"""
 
 from __future__ import annotations
 
-from collections.abc import Callable
+import hashlib
+import json
+import logging
 
 from fastapi import Request
+from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import Response
 
 from somba.api.errors import APIError, error_response
+from somba.db.models import IdempotencyRecord, IdempotencyRecordStatus, Merchant
+from somba.db.session import get_db
+from somba.security import parse_api_key, verify_api_key_secret
+
+log = logging.getLogger(__name__)
 
 MUTATING_METHODS = {"POST", "PUT", "PATCH", "DELETE"}
 IDEMPOTENCY_EXEMPT = {"/v1/webhooks/nomba"}
 
 
 class IdempotencyMiddleware(BaseHTTPMiddleware):
-    """Require an idempotency key for mutating requests."""
+    """Require an idempotency key and replay stored responses on repeat keys."""
+
+    async def dispatch(self, request: Request, call_next):
+        if request.method not in MUTATING_METHODS or request.url.path in IDEMPOTENCY_EXEMPT:
+            return await call_next(request)
 
-    async def dispatch(self, request: Request, call_next: Callable[[Request], Response]) -> Response:
-        if request.method in MUTATING_METHODS and request.url.path not in IDEMPOTENCY_EXEMPT:
-            key = request.headers.get("Idempotency-Key", "").strip()
-            if not key:
-                error = APIError(
+        key = request.headers.get("Idempotency-Key", "").strip()
+        if not key:
+            return error_response(
+                APIError(
                     code="missing_idempotency_key",
                     message="Mutating requests require an Idempotency-Key header",
                     status_code=400,
                     param="Idempotency-Key",
                 )
-                return error_response(error)
-            request.state.idempotency_key = key
+            )
+        request.state.idempotency_key = key
+
+        merchant_id = self._resolve_merchant_id(request)
+        if merchant_id is None:
+            # Unauthenticated / invalid key: let the route's auth return 401.
+            return await call_next(request)
+
+        # Reading the body here is safe — Starlette's BaseHTTPMiddleware caches
+        # it and replays it to the downstream handler.
+        body = await request.body()
+        request_hash = hashlib.sha256(
+            request.method.encode() + b"|" + request.url.path.encode() + b"|" + body
+        ).hexdigest()
+        lookup = (merchant_id, key, request.method, request.url.path)
+
+        # --- claim the key, or replay an existing result ---
+        db, gen = self._session(request)
+        try:
+            existing = self._fetch(db, lookup)
+            if existing is not None:
+                replay = self._maybe_replay(existing, request_hash)
+                if replay is not None:
+                    return replay
+                # A 'failed' record: drop it and let this attempt retry.
+                db.delete(existing)
+                db.commit()
+
+            db.add(
+                IdempotencyRecord(
+                    merchant_id=merchant_id,
+                    idempotency_key=key,
+                    method=request.method,
+                    path=request.url.path,
+                    request_hash=request_hash,
+                    status=IdempotencyRecordStatus.in_progress,
+                )
+            )
+            try:
+                db.commit()
+            except IntegrityError:
+                # A concurrent request claimed the same key first.
+                db.rollback()
+                return self._conflict("A request with this Idempotency-Key is already in progress")
+        finally:
+            self._close(gen)
+
+        # --- run the real handler ---
+        try:
+            response = await call_next(request)
+        except Exception:
+            # Handler blew up: release the key so the client can retry.
+            self._release(request, lookup)
+            raise
+
+        raw = b"".join([chunk async for chunk in response.body_iterator])
+
+        # --- persist the response (2xx) or release the key (anything else) ---
+        db, gen = self._session(request)
+        try:
+            rec = self._fetch(db, lookup)
+            if rec is not None:
+                if 200 <= response.status_code < 300:
+                    rec.status = IdempotencyRecordStatus.completed
+                    rec.response_status = response.status_code
+                    rec.response_body = _safe_json(raw)
+                else:
+                    db.delete(rec)
+                db.commit()
+        finally:
+            self._close(gen)
+
+        headers = dict(response.headers)
+        headers.pop("content-length", None)
+        return Response(
+            content=raw,
+            status_code=response.status_code,
+            headers=headers,
+            media_type=response.media_type,
+        )
+
+    # ------------------------------------------------------------------ helpers
+
+    def _session(self, request: Request):
+        """Resolve a DB session the same way FastAPI does — honouring overrides."""
+        dep = request.app.dependency_overrides.get(get_db, get_db)
+        gen = dep()
+        return next(gen), gen
+
+    @staticmethod
+    def _close(gen) -> None:
+        try:
+            next(gen)
+        except StopIteration:
+            pass
+
+    @staticmethod
+    def _fetch(db, lookup):
+        merchant_id, key, method, path = lookup
+        return db.scalar(
+            select(IdempotencyRecord).where(
+                IdempotencyRecord.merchant_id == merchant_id,
+                IdempotencyRecord.idempotency_key == key,
+                IdempotencyRecord.method == method,
+                IdempotencyRecord.path == path,
+            )
+        )
+
+    def _release(self, request: Request, lookup) -> None:
+        db, gen = self._session(request)
+        try:
+            rec = self._fetch(db, lookup)
+            if rec is not None:
+                db.delete(rec)
+                db.commit()
+        finally:
+            self._close(gen)
+
+    def _resolve_merchant_id(self, request: Request) -> int | None:
+        header = request.headers.get("Authorization", "")
+        if not header.startswith("Bearer "):
+            return None
+        token = header.removeprefix("Bearer ").strip()
+        try:
+            public_id, secret = parse_api_key(token)
+        except ValueError:
+            return None
+        db, gen = self._session(request)
+        try:
+            merchant = db.scalar(select(Merchant).where(Merchant.api_key_id == public_id))
+            if merchant is None or not verify_api_key_secret(secret, merchant.api_key_hash):
+                return None
+            return merchant.id
+        finally:
+            self._close(gen)
+
+    def _maybe_replay(self, existing: IdempotencyRecord, request_hash: str):
+        if existing.status == IdempotencyRecordStatus.completed:
+            if existing.request_hash == request_hash:
+                body = existing.response_body
+                content = json.dumps(body).encode() if body is not None else b""
+                resp = Response(
+                    content=content,
+                    status_code=existing.response_status or 200,
+                    media_type="application/json",
+                )
+                resp.headers["Idempotency-Replayed"] = "true"
+                return resp
+            return self._conflict(
+                "Idempotency-Key was reused with a different request body",
+                code="idempotency_key_reuse",
+            )
+        if existing.status == IdempotencyRecordStatus.in_progress:
+            return self._conflict("A request with this Idempotency-Key is already in progress")
+        return None  # failed -> caller drops it and retries
+
+    @staticmethod
+    def _conflict(message: str, code: str = "idempotency_conflict"):
+        return error_response(
+            APIError(code=code, message=message, status_code=409, param="Idempotency-Key")
+        )
+
 
-        return await call_next(request)
+def _safe_json(raw: bytes):
+    if not raw:
+        return None
+    try:
+        return json.loads(raw)
+    except ValueError:
+        return None

tests/integration/test_idempotency.py

@@ -0,0 +1,70 @@
+"""Idempotency end-to-end: a repeated Idempotency-Key replays the stored
+response instead of performing the action twice."""
+
+from __future__ import annotations
+
+from sqlalchemy import func, select
+
+from somba.db.models import Customer
+
+
+def _headers(token: str, key: str) -> dict:
+    return {"Authorization": f"Bearer {token}", "Idempotency-Key": key}
+
+
+def _customer_count(db) -> int:
+    return db.scalar(select(func.count()).select_from(Customer))
+
+
+def test_same_key_same_body_replays_response(api_client, merchant_and_token, db):
+    _, token = merchant_and_token
+    key = "idem-fixed-123"
+    body = {"email": "alice@gym.com", "name": "Alice"}
+
+    first = api_client.post("/v1/customers", json=body, headers=_headers(token, key))
+    assert first.status_code == 201
+    assert "Idempotency-Replayed" not in first.headers
+
+    second = api_client.post("/v1/customers", json=body, headers=_headers(token, key))
+    assert second.status_code == 201
+    assert second.headers.get("Idempotency-Replayed") == "true"
+
+    # Identical response, and only ONE customer was actually created.
+    assert second.json() == first.json()
+    assert _customer_count(db) == 1
+
+
+def test_same_key_different_body_conflicts(api_client, merchant_and_token, db):
+    _, token = merchant_and_token
+    key = "idem-fixed-456"
+
+    first = api_client.post(
+        "/v1/customers", json={"email": "a@gym.com", "name": "A"}, headers=_headers(token, key)
+    )
+    assert first.status_code == 201
+
+    second = api_client.post(
+        "/v1/customers", json={"email": "different@gym.com", "name": "B"}, headers=_headers(token, key)
+    )
+    assert second.status_code == 409
+    assert second.json()["error"]["code"] == "idempotency_key_reuse"
+    assert _customer_count(db) == 1  # the conflicting request created nothing
+
+
+def test_different_key_creates_distinct_resources(api_client, merchant_and_token, db):
+    _, token = merchant_and_token
+    body = {"email": "a@gym.com", "name": "A"}
+
+    api_client.post("/v1/customers", json=body, headers=_headers(token, "k1"))
+    api_client.post("/v1/customers", json={"email": "b@gym.com", "name": "B"}, headers=_headers(token, "k2"))
+
+    assert _customer_count(db) == 2  # distinct keys -> two creations
+
+
+def test_missing_idempotency_key_rejected(api_client, merchant_and_token):
+    _, token = merchant_and_token
+    resp = api_client.post(
+        "/v1/customers", json={"email": "x@gym.com"}, headers={"Authorization": f"Bearer {token}"}
+    )
+    assert resp.status_code == 400
+    assert resp.json()["error"]["code"] == "missing_idempotency_key"