🔒 [security fix] Fix Unbounded In-Memory Job Storage Leading to Resource Exhaustion

src/api/main.py

@@ -14,6 +14,7 @@
 import logging
 
 import time
+import threading
 from src.core.engine import NDLOCREngine
 from src.schemas.ocr import OCRResponse, OCRPage, OCRLine, OCRRequest, OCRJobResponse, OCRJobResult
 from src.schemas.openai import ChatCompletionRequest, ChatCompletionResponse, ChatCompletionChoice, ChatCompletionChoiceMessage
@@ -27,21 +28,63 @@ class InMemoryJobStore:
     In-memory job store for OCR results.
     Used to track the status and final output of asynchronous background jobs.
     In a production environment, this should be replaced with a persistent store like Redis.
+    Includes time-to-live (TTL) eviction and LRU eviction to prevent unbounded memory growth.
     """
-    def __init__(self):
+    def __init__(self, max_size: int = 1000, ttl_seconds: float = 3600.0):
         self._jobs: Dict[str, OCRJobResult] = {}
+        self._timestamps: Dict[str, float] = {}  # job_id -> last access timestamp
+        self._lock = threading.Lock()
+        self._max_size = max_size
+        self._ttl_seconds = ttl_seconds
+
+    def _evict_expired(self, now: float):
+        """Removes expired jobs from the store. Must be called under lock."""
+        expired_ids = [
+            job_id for job_id, ts in list(self._timestamps.items())
+            if now - ts > self._ttl_seconds
+        ]
+        for job_id in expired_ids:
+            self._jobs.pop(job_id, None)
+            self._timestamps.pop(job_id, None)
+
+    def _evict_lru(self):
+        """Removes the oldest (least recently accessed) job. Must be called under lock."""
+        if not self._timestamps:
+            return
+        # Find key with the minimum timestamp
+        oldest_job_id = min(self._timestamps, key=self._timestamps.get)
+        self._jobs.pop(oldest_job_id, None)
+        self._timestamps.pop(oldest_job_id, None)
 
     def get(self, job_id: str) -> Optional[OCRJobResult]:
-        """Retrieves a job result by its ID."""
-        return self._jobs.get(job_id)
+        """Retrieves a job result by its ID and updates its access time."""
+        with self._lock:
+            now = time.time()
+            self._evict_expired(now)
+            if job_id in self._jobs:
+                self._timestamps[job_id] = now
+                return self._jobs[job_id]
+            return None
 
     def set(self, job_id: str, result: OCRJobResult):
-        """Stores or updates a job result."""
-        self._jobs[job_id] = result
+        """Stores or updates a job result, evicting expired or LRU items if full."""
+        with self._lock:
+            now = time.time()
+            self._evict_expired(now)
+
+            # If still over limit and adding a new key would exceed max_size
+            if len(self._jobs) >= self._max_size and job_id not in self._jobs:
+                self._evict_lru()
+
+            self._jobs[job_id] = result
+            self._timestamps[job_id] = now
 
     def exists(self, job_id: str) -> bool:
         """Checks if a job with the given ID exists."""
-        return job_id in self._jobs
+        with self._lock:
+            now = time.time()
+            self._evict_expired(now)
+            return job_id in self._jobs
 
 def _engine_result_to_ocr_page(result: Dict[str, Any], index: int = 0) -> OCRPage:
     """
@@ -74,6 +117,8 @@ def _engine_result_to_ocr_page(result: Dict[str, Any], index: int = 0) -> OCRPag
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", 10 * 1024 * 1024)) # Default 10MB
 MAX_BODY_SIZE = int(os.getenv("MAX_BODY_SIZE", 15 * 1024 * 1024))   # Default 15MB
 MAX_PIXELS = int(os.getenv("MAX_PIXELS", 100_000_000))            # Default 100MP
+JOB_STORE_MAX_SIZE = int(os.getenv("JOB_STORE_MAX_SIZE", "1000"))
+JOB_STORE_TTL_SECONDS = float(os.getenv("JOB_STORE_TTL_SECONDS", "3600.0"))
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -87,7 +132,7 @@ async def lifespan(app: FastAPI):
     # Models are loaded once and stored in app.state for sharing across requests.
     enable_tcy = os.getenv("ENABLE_TCY", "false").lower() == "true"
     app.state.engine = NDLOCREngine(device="cpu", enable_tcy=enable_tcy)
-    app.state.job_store = InMemoryJobStore()
+    app.state.job_store = InMemoryJobStore(max_size=JOB_STORE_MAX_SIZE, ttl_seconds=JOB_STORE_TTL_SECONDS)
     yield
     logger.info("Shutting down...")
     # Clean up engine resources (e.g., ThreadPoolExecutor)

tests/test_job_store_limits.py

@@ -0,0 +1,81 @@
+import time
+import threading
+from src.api.main import InMemoryJobStore
+from src.schemas.ocr import OCRJobResult
+
+def test_job_store_max_size_eviction():
+    # Store with max_size=3
+    store = InMemoryJobStore(max_size=3, ttl_seconds=3600.0)
+
+    job1 = OCRJobResult(job_id="job1", status="pending")
+    job2 = OCRJobResult(job_id="job2", status="pending")
+    job3 = OCRJobResult(job_id="job3", status="pending")
+    job4 = OCRJobResult(job_id="job4", status="pending")
+
+    store.set("job1", job1)
+    store.set("job2", job2)
+    store.set("job3", job3)
+
+    # All should exist
+    assert store.exists("job1")
+    assert store.exists("job2")
+    assert store.exists("job3")
+
+    # Access job1 and job2 to update their timestamps (making job3 the least recently used)
+    # Note: Because they are created very close in time, we want to make sure order is clear.
+    # To be extremely safe, we will simulate or just set them with slight artificial delays or let the system access them.
+    store.get("job1")
+    time.sleep(0.01)
+    store.get("job2")
+    time.sleep(0.01)
+
+    # Now we insert job4, which should trigger LRU eviction on job3
+    store.set("job4", job4)
+
+    assert store.exists("job1")
+    assert store.exists("job2")
+    assert store.exists("job4")
+    assert not store.exists("job3")
+
+def test_job_store_ttl_eviction():
+    # Store with max_size=10, ttl_seconds=0.05
+    store = InMemoryJobStore(max_size=10, ttl_seconds=0.05)
+
+    job1 = OCRJobResult(job_id="job1", status="pending")
+    store.set("job1", job1)
+
+    assert store.exists("job1")
+
+    # Sleep past TTL
+    time.sleep(0.06)
+
+    # exists/get/set should trigger eviction and return None / False
+    assert not store.exists("job1")
+    assert store.get("job1") is None
+
+def test_job_store_thread_safety():
+    # Ensure no race conditions under high concurrent set/get operations
+    store = InMemoryJobStore(max_size=10, ttl_seconds=10.0)
+
+    num_threads = 20
+    ops_per_thread = 100
+    errors = []
+
+    def worker(thread_idx):
+        try:
+            for i in range(ops_per_thread):
+                job_id = f"thread-{thread_idx}-job-{i}"
+                job = OCRJobResult(job_id=job_id, status="pending")
+                store.set(job_id, job)
+                store.get(job_id)
+                store.exists(job_id)
+        except Exception as e:
+            errors.append(e)
+
+    threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_threads)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert not errors, f"Encountered thread safety errors: {errors}"