Merge pull request #37 from NVIDIA-NeMo/security/viewer-ingest-resource-limits

Viewer Ingest Resource Limits

Author: rdasilveiracabral

src/nooa/viewer/main.py

@@ -10,6 +10,7 @@
 import asyncio
 import hmac
 import ipaddress
+import json
 import logging
 import os
 import time
@@ -34,6 +35,30 @@
 
 log = logging.getLogger(__name__)
 
+_INGEST_MAX_BODY_BYTES = 16 * 1024 * 1024
+_INGEST_QUEUE_MAX_ITEMS = 64
+
+
+async def _read_limited_body(request: Request) -> bytes:
+    """Read at most ``_INGEST_MAX_BODY_BYTES`` without buffering more."""
+    raw_length = request.headers.get("content-length")
+    if raw_length and raw_length.isdecimal() and int(raw_length) > _INGEST_MAX_BODY_BYTES:
+        raise HTTPException(status_code=413, detail="ingest body too large")
+
+    chunks: list[bytes] = []
+    body_size = 0
+    async for chunk in request.stream():
+        body_size += len(chunk)
+        if body_size > _INGEST_MAX_BODY_BYTES:
+            raise HTTPException(status_code=413, detail="ingest body too large")
+        chunks.append(chunk)
+    return b"".join(chunks)
+
+
+async def _read_limited_json(request: Request) -> object:
+    return json.loads(await _read_limited_body(request))
+
+
 # ---------------------------------------------------------------------------
 # Write queue — decouple HTTP ingest latency from SQLite write latency.
 #
@@ -46,8 +71,8 @@
 # at a time, the event loop is never blocked, and HTTP latency is near-zero.
 # ---------------------------------------------------------------------------
 
-_ingest_queue: asyncio.Queue[bytes] = asyncio.Queue()
-_QUEUE_WARN_THRESHOLD = 500  # log a warning if the backlog grows this large
+_ingest_queue: asyncio.Queue[bytes] = asyncio.Queue(maxsize=_INGEST_QUEUE_MAX_ITEMS)
+_QUEUE_WARN_THRESHOLD = 48
 
 # Single-writer thread pool: exactly one thread owns the write connection.
 # Using max_workers=1 ensures serial SQLite writes with no concurrent access.
@@ -216,21 +241,17 @@ async def otlp_ingest(request: Request):
     (3-5 MB payloads) don't block the event loop while parsing.
     """
     try:
-        body_bytes = await request.body()
+        body_bytes = await _read_limited_body(request)
     except ClientDisconnect:
-        # BSP exporter timed out and closed the connection before we read the
-        # body — log at WARNING (not ERROR) since this is a transient backpressure
-        # signal, not a bug.  The BSP will retry on the next export cycle.
         log.warning("[otlp_ingest] Client disconnected before body was read — BSP may retry")
         return JSONResponse(status_code=499, content={"error": "client disconnected"})
     qsize = _ingest_queue.qsize()
     if qsize >= _QUEUE_WARN_THRESHOLD:
-        log.warning(
-            "[otlp_ingest] Write queue backlog: %d pending — "
-            "SQLite may not be keeping up with ingest rate.",
-            qsize,
-        )
-    await _ingest_queue.put(body_bytes)
+        log.warning("[otlp_ingest] Write queue backlog: %d pending", qsize)
+    try:
+        _ingest_queue.put_nowait(body_bytes)
+    except asyncio.QueueFull:
+        return JSONResponse(status_code=503, content={"error": "ingest queue is full"})
     return JSONResponse(content={"queued": True})
 
 
@@ -278,7 +299,8 @@ async def journal_messages_ingest(request: Request):
     """
     import sqlite3
 
-    body = await request.json()
+    body = await _read_limited_json(request)
+
     if not isinstance(body, list):
         return JSONResponse(
             status_code=400,
@@ -315,7 +337,8 @@ async def journal_call_ingest(request: Request):
     """
     import sqlite3
 
-    body = await request.json()
+    body = await _read_limited_json(request)
+
     if not isinstance(body, dict) or not body.get("call_id") or not body.get("session_id"):
         return JSONResponse(
             status_code=400,
@@ -361,7 +384,8 @@ async def journal_blocks_ingest(request: Request):
     """
     import sqlite3
 
-    body = await request.json()
+    body = await _read_limited_json(request)
+
     if not isinstance(body, list):
         return JSONResponse(
             status_code=400,

tests/viewer/test_main.py

@@ -139,9 +139,7 @@ def test_api_routes_require_configured_bearer_token(client, monkeypatch):
     assert write_response.status_code == 401
     assert read_response.headers["www-authenticate"] == "Bearer"
 
-    authorized = client.get(
-        "/api/version", headers={"Authorization": "Bearer test-viewer-token"}
-    )
+    authorized = client.get("/api/version", headers={"Authorization": "Bearer test-viewer-token"})
     assert authorized.status_code == 200
 
 
@@ -258,6 +256,17 @@ def test_dedup_idempotent(self, client):
         assert r1.status_code == 200
         assert r2.status_code == 200
 
+    def test_oversized_body_is_rejected_before_json_decode(self, client, monkeypatch):
+        monkeypatch.setattr(main_module, "_INGEST_MAX_BODY_BYTES", 32)
+
+        response = client.post(
+            "/v1/journal/messages",
+            content=b"x" * 33,
+            headers={"Content-Type": "application/json"},
+        )
+
+        assert response.status_code == 413
+
 
 # ---------------------------------------------------------------------------
 # POST /v1/journal/calls

tests/viewer/test_otlp_ingest.py

@@ -46,9 +46,13 @@ async def http_client(mock_store, tmp_path):
     The lifespan is NOT triggered (httpx limitation), so no worker task runs.
     Suitable for testing the endpoint handler in isolation.
     """
-    with patch("nooa.viewer.main.otlp_store", mock_store):
-        from nooa.viewer.main import app
+    from nooa.viewer.main import app
 
+    fresh_queue: asyncio.Queue = asyncio.Queue(maxsize=256)
+    with (
+        patch("nooa.viewer.main.otlp_store", mock_store),
+        patch("nooa.viewer.main._ingest_queue", fresh_queue),
+    ):
         async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c:
             yield c
 
@@ -70,6 +74,38 @@ async def test_does_not_call_ingest_synchronously(self, http_client, mock_store)
         # Without a running worker, ingest_batch_write_bytes should NOT have been called yet.
         mock_store.ingest_batch_write_bytes.assert_not_called()
 
+    async def test_rejects_declared_body_over_limit_before_queueing(self, http_client):
+        from nooa.viewer import main
+
+        with patch.object(main, "_INGEST_MAX_BODY_BYTES", 32):
+            response = await http_client.post("/v1/traces", content=b"x" * 33)
+
+        assert response.status_code == 413
+        assert main._ingest_queue.empty()
+
+    async def test_rejects_chunked_body_over_limit(self, http_client):
+        from nooa.viewer import main
+
+        async def chunks():
+            yield b"x" * 20
+            yield b"x" * 20
+
+        with patch.object(main, "_INGEST_MAX_BODY_BYTES", 32):
+            response = await http_client.post("/v1/traces", content=chunks())
+
+        assert response.status_code == 413
+        assert main._ingest_queue.empty()
+
+    async def test_full_queue_returns_503(self, http_client):
+        from nooa.viewer import main
+
+        full_queue: asyncio.Queue = asyncio.Queue(maxsize=1)
+        full_queue.put_nowait(b"existing")
+        with patch.object(main, "_ingest_queue", full_queue):
+            response = await http_client.post("/v1/traces", content=b"{}")
+
+        assert response.status_code == 503
+
 
 # ---------------------------------------------------------------------------
 # Worker tests — verify the background task drains the queue correctly