Fix/enroll

app/core/constant.py

@@ -50,5 +50,10 @@ class AuditEventType(str, Enum):
 GOOGLE_DRIVE_FILES_URL = "https://www.googleapis.com/drive/v3/files/{file_id}"
 
 MAX_IMAGE_SIZE = 5 * 1024 * 1024
+MIN_IMAGE_DIM = 64
+MAX_IMAGE_DIM = 4096
 MIN_ENROLL_IMAGES = 3
 MAX_ENROLL_IMAGES = 5
+
+ENROLL_RATE_LIMIT_MAX = 5
+ENROLL_RATE_LIMIT_WINDOW = 3600

app/router/mobile/enrollement.py

@@ -1,75 +1,163 @@
+import re
+import uuid
+from io import BytesIO
 from typing import Annotated, List
 
-from fastapi import APIRouter, File, UploadFile,  Depends
+import filetype  # type: ignore[import-untyped]
+import pillow_heif # type: ignore[import-untyped]
+from fastapi import APIRouter, Depends, File, UploadFile
+from fastapi.concurrency import run_in_threadpool
+from PIL import Image
+from pydantic import BaseModel
 
 from app.container import Container, get_container
-from app.deps.token_auth import MobileUserSchema, get_current_mobile_user
-from app.core.exceptions import AppException
 from app.core.constant import (
-    DEFAULT_CONTENT_TYPE,
+    ENROLL_RATE_LIMIT_MAX,
+    ENROLL_RATE_LIMIT_WINDOW,
     IMAGE_ALLOWED_TYPES,
     MAX_ENROLL_IMAGES,
     MAX_IMAGE_SIZE,
+    MAX_IMAGE_DIM,
     MIN_ENROLL_IMAGES,
+    MIN_IMAGE_DIM,
 )
+from app.core.exceptions import AppException
+from app.deps.token_auth import MobileUserSchema, get_current_mobile_user
 from app.service.face_embedding import FaceImagePayload
-from db.generated.models import User
+
+
+pillow_heif.register_heif_opener()
+
+
+Image.MAX_IMAGE_PIXELS = MAX_IMAGE_DIM * MAX_IMAGE_DIM
+
+
+class EnrollmentResponse(BaseModel):
+    id: uuid.UUID
+
+    model_config = {"from_attributes": True}
+
 
 router = APIRouter()
 
 
-@router.post("/enroll")
+def _sanitise_filename(raw: str | None, extension: str) -> str:
+    prefix = str(uuid.uuid4())
+    if not raw:
+        return f"{prefix}.{extension}"
+    name = re.sub(r'[\\/:*?"<>|\x00-\x1f]', "_", raw)
+    name = name.lstrip(".")[:128]
+    return f"{prefix}_{name}"
+
+
+def _validate_dimensions(contents: bytes) -> None:
+
+    try:
+        img = Image.open(BytesIO(contents))
+        w, h = img.size
+    except Exception as e:
+        raise AppException.image_format_error(
+            "File could not be decoded as a valid image"
+        ) from e
+
+    max_pixels = Image.MAX_IMAGE_PIXELS
+
+    if max_pixels is not None and w * h > max_pixels:
+        raise AppException.bad_request(
+            f"Image exceeds maximum allowed resolution of {max_pixels} total pixels."
+        )
+
+    try:
+        img.load()
+    except Exception as e:
+        raise AppException.image_format_error(
+            "File contains corrupted or incomplete pixel data"
+        ) from e
+
+    if w < MIN_IMAGE_DIM or h < MIN_IMAGE_DIM:
+        raise AppException.bad_request(
+            f"Image too small — minimum {MIN_IMAGE_DIM}x{MIN_IMAGE_DIM} px"
+        )
+    if w > MAX_IMAGE_DIM or h > MAX_IMAGE_DIM:
+        raise AppException.bad_request(
+            f"Image too large — maximum {MAX_IMAGE_DIM}x{MAX_IMAGE_DIM} px"
+        )
+
+
+async def read_limited(file: UploadFile, limit: int) -> bytes:
+    chunks: list[bytes] = []
+    total = 0
+    while True:
+        chunk = await file.read(65536)
+        if not chunk:
+            break
+        total += len(chunk)
+        if total > limit:
+            raise AppException.bad_request(
+                f"File exceeds maximum allowed size of {limit} bytes"
+            )
+        chunks.append(chunk)
+
+    await file.seek(0)
+    return b"".join(chunks)
+
+
+@router.post("/enroll", response_model=EnrollmentResponse)
 async def enroll_face(
-   files: Annotated[
+    files: Annotated[
         List[UploadFile],
         File(
-            description="Upload one or more face images",
-            openapi_examples={
-                "single_file": {
-                    "summary": "One file example",
-                    "description": "Example of uploading one file",
-                    "value": "example.jpg"
-                },
-                "multiple_files": {
-                    "summary": "Multiple files example",
-                    "description": "Example of uploading multiple files",
-                    "value": ["face1.png", "face2.png"]
-                },
-            },
+            description=(
+                f"Between {MIN_ENROLL_IMAGES} and {MAX_ENROLL_IMAGES} face images "
+                f"(JPEG, PNG, HEIC, or HEIF). "
+                f"Each file must be under {MAX_IMAGE_SIZE // (1024 * 1024)} MB "
+                f"and at least {MIN_IMAGE_DIM}x{MIN_IMAGE_DIM} px."
+            ),
         ),
     ],
     container: Container = Depends(get_container),
     user: MobileUserSchema = Depends(get_current_mobile_user),
-) -> User:
+) -> EnrollmentResponse:
+    await container.auth_service.check_rate_limit(
+        redis=container.redis,
+        key=f"rate:enroll:{user.user_id}",
+        max_requests=ENROLL_RATE_LIMIT_MAX,
+        window_seconds=ENROLL_RATE_LIMIT_WINDOW,
+    )
 
     if not (MIN_ENROLL_IMAGES <= len(files) <= MAX_ENROLL_IMAGES):
         raise AppException.bad_request(
             f"You must upload between {MIN_ENROLL_IMAGES} and {MAX_ENROLL_IMAGES} images for enrollment."
         )
 
-
     image_payloads: list[FaceImagePayload] = []
+
     for file in files:
-        if file.content_type not in IMAGE_ALLOWED_TYPES:
+        contents = await read_limited(file, MAX_IMAGE_SIZE)
+
+        kind = filetype.guess(contents)
+        if kind is None or kind.mime not in IMAGE_ALLOWED_TYPES:
             raise AppException.image_format_error(
-                f"File {file.filename} has unsupported format {file.content_type}"
+                f"Unsupported format. Allowed types: {', '.join(IMAGE_ALLOWED_TYPES)}"
             )
 
-        contents = await file.read()
-        if len(contents) > MAX_IMAGE_SIZE:
-            raise AppException.bad_request(
-                f"File {file.filename} exceeds maximum size of {MAX_IMAGE_SIZE} bytes"
-            )
+        await run_in_threadpool(_validate_dimensions, contents)
 
-        payload: FaceImagePayload = FaceImagePayload(
-            filename=file.filename or "unknown",
-            content_type=file.content_type or DEFAULT_CONTENT_TYPE,
-            bytes=contents,
+        image_payloads.append(
+            FaceImagePayload(
+                filename=_sanitise_filename(file.filename, kind.extension),
+                content_type=kind.mime,
+                bytes=contents,
+            )
         )
 
-        image_payloads.append(payload)
-
-    return await container.auth_service.add_embbed_user(
-        user.user_id,
-        image_payloads,
-    )
+    try:
+        updated_user = await container.auth_service.add_embbed_user(
+            user.user_id,
+            image_payloads,
+        )
+        return EnrollmentResponse.model_validate(updated_user)
+    except Exception as e:
+        raise AppException.internal_error(
+            "Enrollment failed due to an internal error"
+        ) from e

app/service/users.py

@@ -216,7 +216,6 @@ async def _create_mobile_session(
         expiry = Get_expiry_time()
         logger.info("session_created session_id=%s user_id=%s", session.id, user_id)
 
-        # Populate Redis auth cache for fast-path validation
         await SessionService.cache_session_for_auth(
             redis=redis,
             session_id=session.id,
@@ -290,6 +289,15 @@ async def add_embbed_user(
     ) -> User:
         logger.info("Generating face embeddings for user %s", user_id)
 
+        existing = await self.user_querier.get_user_by_id(id=user_id)
+        if not existing:
+            raise AppException.not_found("User not found")
+        if existing.face_embedding is not None:
+            raise AppException.conflict(
+                "User already has an active face enrollment. "
+                "Delete the existing enrollment before re-enrolling."
+            )
+
         averaging = await self.face_embedding_service.compute_average_embedding(
             image_payloads
         )
@@ -410,7 +418,6 @@ async def delete_user(self, *, redis: RedisClient, user_id: uuid.UUID) -> User:
             session_key = constant.RedisKey.UserSessionByUser.value.format(
                 user_id=user_id
             )
-            # Best-effort: also invalidate the per-session MobileSessionCache.
             raw_session_id = await redis.get(session_key)
             if raw_session_id:
                 try:
@@ -431,15 +438,13 @@ async def block_user(self, *, redis: RedisClient, user_id: uuid.UUID) -> User:
                 raise AppException.not_found("User not found")
 
             session_key = constant.RedisKey.UserSessionByUser.value.format(user_id=user_id)
-            # Best-effort: retrieve the session_id from UserSessionByUser cache to also
-            # invalidate the per-session MobileSessionCache entry.
             raw_session_id = await redis.get(session_key)
             if raw_session_id:
                 try:
                     session_id = uuid.UUID(raw_session_id)
                     await SessionService.delete_session_cache(redis=redis, session_id=session_id)
                 except (ValueError, Exception):
-                    pass  # non-blocking: session cache will expire naturally
+                    pass
             await redis.delete(session_key)
 
             return user
@@ -474,9 +479,9 @@ async def check_rate_limit(
     ) -> None:
         """Enforce rate limiting using Redis INCR + EXPIRE.
 
-        Increments a counter for ``key``.  On the first increment the key
+        Increments a counter for ``key``. On the first increment the key
         is given a TTL of ``window_seconds`` so the window resets
-        automatically.  If the counter exceeds ``max_requests`` a 429
+        automatically. If the counter exceeds ``max_requests`` a 429
         response is raised with a ``Retry-After`` header.
         """
         current_count = await redis.incr(key)

pyproject.toml

@@ -29,6 +29,8 @@ dependencies = [
     "firebase-admin>=6.8.0",
     "pywebpush>=2.3.0",
     "opencv-python>=4.13.0.92",
+    "filetype>=1.2.0",
+    "pillow-heif>=1.3.0",
 ]
 
 [tool.ruff]