From f3e95fc1c6bd7f1e83f18f0cf4e0e3dbc315dabd Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 18:29:08 +0100 Subject: [PATCH 01/11] fix(core): fence scan worker leases (#303) Signed-off-by: Shaurya K Sharma --- .../e4f7a9b2c6d8_scan_leases_and_fencing.py | 68 ++++ api/models/finding.py | 316 ++++++++++++------ docs/async-scan-architecture.md | 11 +- scanner/worker.py | 144 +++++++- tests/test_async_scan_persistence.py | 81 ++++- tests/test_observability.py | 5 +- tests/test_scan_leases_postgres.py | 298 +++++++++++++++++ tests/test_severity_contract.py | 27 +- tests/test_worker.py | 25 +- 9 files changed, 841 insertions(+), 134 deletions(-) create mode 100644 alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py create mode 100644 tests/test_scan_leases_postgres.py diff --git a/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py b/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py new file mode 100644 index 00000000..a1ffbc93 --- /dev/null +++ b/alembic/versions/e4f7a9b2c6d8_scan_leases_and_fencing.py @@ -0,0 +1,68 @@ +"""Add renewable ownership leases and fencing tokens to scans. + +Revision ID: e4f7a9b2c6d8 +Revises: d8e4f6a1b2c3 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "e4f7a9b2c6d8" +down_revision: Union[str, Sequence[str], None] = "d8e4f6a1b2c3" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Add additive lease state and make legacy running work recoverable.""" + op.add_column("scans", sa.Column("lease_owner", sa.Text(), nullable=True)) + op.add_column("scans", sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True)) + op.add_column("scans", sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=True)) + op.add_column( + "scans", + sa.Column("fencing_token", sa.BigInteger(), server_default=sa.text("0"), nullable=False), + ) + + # A pre-lease running row belongs to an old worker that cannot satisfy the + # new fencing contract. Marking its lease expired preserves the row and + # lets the new worker recover it under a fresh owner/token. + op.execute( + """ + UPDATE scans + SET lease_expires_at = CURRENT_TIMESTAMP + WHERE status = 'running' AND lease_expires_at IS NULL + """ + ) + + # These indexes are additive and are created concurrently so a populated + # production scans table remains available while the migration runs. + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_scans_pending_started_at + ON scans (started_at ASC) + WHERE status = 'pending' + """ + ) + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_scans_running_lease_expires_at + ON scans (lease_expires_at ASC) + WHERE status = 'running' + """ + ) + + +def downgrade() -> None: + """Remove lease metadata; callers must be rolled back first.""" + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_scans_running_lease_expires_at") + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_scans_pending_started_at") + op.drop_column("scans", "fencing_token") + op.drop_column("scans", "last_heartbeat_at") + op.drop_column("scans", "lease_expires_at") + op.drop_column("scans", "lease_owner") diff --git a/api/models/finding.py b/api/models/finding.py index 0f366924..d064b3fc 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -11,6 +11,7 @@ import psycopg2 import psycopg2.extras import psycopg2.pool +from psycopg2 import extensions from openshield.severity import ( CONTRACT_VERSION, @@ -22,6 +23,11 @@ logger = logging.getLogger(__name__) + +class LostLease(RuntimeError): + """Raised when a worker no longer owns the scan it is trying to update.""" + + FRAMEWORKS_DIR = Path(__file__).parent.parent.parent / "compliance" / "frameworks" # One pool per DSN, shared across all DatabaseManager instances in this @@ -119,26 +125,59 @@ def __init__(self, dsn: Optional[str] = None) -> None: def connect(self) -> None: """Acquire a connection from this DSN's shared pool.""" + if self.conn is not None: + self._return_connection(close=bool(self.conn.closed)) self.conn = _get_pool(self.dsn).getconn() self.conn.autocommit = False logger.debug("Database connection acquired from pool") def _get_conn(self) -> Any: - if self.conn is None or self.conn.closed: + if self.conn is not None and self.conn.closed: + self._return_connection(close=True) + if self.conn is None: self.connect() return self.conn - def close(self) -> None: - """Return the connection to its pool (or discard it if broken).""" - if self.conn is None: + def _return_connection(self, close: bool = False, conn: Optional[Any] = None) -> None: + """Return the checked-out connection, discarding it when requested.""" + conn = conn or self.conn + if conn is None: return try: - _get_pool(self.dsn).putconn(self.conn, close=bool(self.conn.closed)) + _get_pool(self.dsn).putconn(conn, close=close or bool(conn.closed)) logger.debug("Database connection returned to pool") except Exception as exc: logger.error("Error returning connection to pool: %s", exc) + try: + conn.close() + except Exception: + pass + finally: + if self.conn is conn: + self.conn = None + + def rollback(self, conn: Optional[Any] = None) -> None: + """End a failed transaction, discarding a connection the server lost.""" + conn = conn or self.conn + if conn is None: + return + try: + if conn.closed or conn.info.transaction_status == extensions.TRANSACTION_STATUS_UNKNOWN: + self._return_connection(close=True, conn=conn) + return + conn.rollback() + except Exception as exc: + logger.warning("Database rollback failed; discarding connection: %s", exc) + self._return_connection(close=True, conn=conn) + + def close(self) -> None: + """Return the connection to its pool (or discard it if broken).""" + if self.conn is None: + return + try: + self.rollback() finally: - self.conn = None + self._return_connection() def ping(self) -> bool: """Execute a trivial query to confirm database connectivity. @@ -163,8 +202,13 @@ def init_db(self) -> None: # Write # # ------------------------------------------------------------------ # - def save_scan(self, scan_result: Dict[str, Any]) -> None: - """Persist a full scan result (scan header + all findings).""" + def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token: int) -> None: + """Persist a completed scan only while its worker still owns the lease. + + The ownership check and all authoritative result writes share one + transaction. A worker whose lease was reclaimed therefore cannot + delete or insert child findings after it has become stale. + """ from datetime import datetime, timezone # Validate and canonicalize the entire batch before issuing SQL. A bad @@ -181,32 +225,40 @@ def save_scan(self, scan_result: Dict[str, Any]) -> None: with conn.cursor() as cur: cur.execute( """ - INSERT INTO scans ( - scan_id, subscription_id, started_at, completed_at, - total_findings, score, cve_enrichment_status, status, - attempt_count, error_message, severity_contract_version - ) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) - ON CONFLICT (scan_id) DO UPDATE SET - completed_at = EXCLUDED.completed_at, - total_findings = EXCLUDED.total_findings, - score = EXCLUDED.score, - status = EXCLUDED.status, - error_message = EXCLUDED.error_message, - severity_contract_version = EXCLUDED.severity_contract_version + SELECT scan_id + FROM scans + WHERE scan_id = %s + AND status = 'running' + AND lease_owner = %s + AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + FOR UPDATE + """, + (scan_result["scan_id"], lease_owner, fencing_token), + ) + if cur.fetchone() is None: + raise LostLease(f"Scan {scan_result['scan_id']} is no longer owned by this worker") + cur.execute( + """ + UPDATE scans + SET completed_at = %s, + total_findings = %s, + score = %s, + cve_enrichment_status = %s, + status = 'completed', + error_message = NULL, + severity_contract_version = %s, + lease_owner = NULL, + lease_expires_at = NULL + WHERE scan_id = %s """, ( - scan_result["scan_id"], - scan_result["subscription_id"], - scan_result["started_at"], completed_at, len(findings), score_findings(findings), scan_result.get("cve_enrichment_status", "PENDING"), - scan_result.get("status", "completed"), - scan_result.get("attempt_count", 0), - scan_result.get("error_message"), CONTRACT_VERSION, + scan_result["scan_id"], ), ) # A worker retry replaces the previous result atomically. This @@ -251,7 +303,7 @@ def save_scan(self, scan_result: Dict[str, Any]) -> None: # psycopg2 connections remain in an aborted transaction after any # SQL error. Roll back here so the worker can record failure and # safely process subsequent scans on the same pooled connection. - conn.rollback() + self.rollback(conn) raise logger.info( "Saved scan %s with %d findings", @@ -373,94 +425,158 @@ def create_pending_scan(self, scan_id: str, subscription_id: str) -> None: conn.commit() logger.info("Created pending scan %s for %s", scan_id, subscription_id) - def update_scan_status(self, scan_id: str, status: str, error_message: Optional[str] = None) -> None: - """Update the status of a scan (running, completed, failed).""" + def update_scan_status( + self, + scan_id: str, + status: str, + error_message: Optional[str], + *, + lease_owner: str, + fencing_token: int, + ) -> None: + """Record a terminal failure while the caller still owns the lease.""" + if status != "failed": + raise ValueError("Fenced status updates only support terminal failures") conn = self._get_conn() - from datetime import datetime, timezone - - with conn.cursor() as cur: - if status == "completed": - completed_at = datetime.now(timezone.utc).isoformat() - cur.execute( - "UPDATE scans SET status = %s, completed_at = %s, error_message = NULL WHERE scan_id = %s", - (status, completed_at, scan_id), - ) - else: + try: + with conn.cursor() as cur: cur.execute( - "UPDATE scans SET status = %s, error_message = %s WHERE scan_id = %s", - (status, error_message, scan_id), + """ + UPDATE scans + SET status = 'failed', + error_message = %s, + lease_owner = NULL, + lease_expires_at = NULL + WHERE scan_id = %s + AND status = 'running' + AND lease_owner = %s + AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + """, + (error_message, scan_id, lease_owner, fencing_token), ) - conn.commit() + if cur.rowcount != 1: + raise LostLease(f"Scan {scan_id} is no longer owned by this worker") + conn.commit() + except Exception: + self.rollback(conn) + raise logger.info("Updated scan %s status to %s", scan_id, status) - def claim_next_pending_scan(self) -> Optional[Dict[str, Any]]: - """Atomically claim the next pending scan using SKIP LOCKED.""" + def claim_next_pending_scan(self, lease_owner: str, lease_seconds: int) -> Optional[Dict[str, Any]]: + """Atomically claim one pending scan and establish its renewable lease.""" + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") conn = self._get_conn() - from datetime import datetime, timezone + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + UPDATE scans + SET status = 'running', + claimed_at = CURRENT_TIMESTAMP, + lease_owner = %s, + lease_expires_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second'), + last_heartbeat_at = CURRENT_TIMESTAMP, + fencing_token = COALESCE(fencing_token, 0) + 1, + attempt_count = COALESCE(attempt_count, 0) + 1, + error_message = NULL + WHERE scan_id = ( + SELECT scan_id + FROM scans + WHERE status = 'pending' + ORDER BY started_at ASC + FOR UPDATE SKIP LOCKED + LIMIT 1 + ) + RETURNING * + """, + (lease_owner, lease_seconds), + ) + row = cur.fetchone() + conn.commit() + return dict(row) if row else None + except Exception: + self.rollback(conn) + raise - claimed_at = datetime.now(timezone.utc).isoformat() - with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: - cur.execute( - """ - UPDATE scans - SET status = 'running', - claimed_at = %s, - attempt_count = COALESCE(attempt_count, 0) + 1, - error_message = NULL - WHERE scan_id = ( - SELECT scan_id - FROM scans - WHERE status = 'pending' - ORDER BY started_at ASC - FOR UPDATE SKIP LOCKED - LIMIT 1 + def heartbeat_scan(self, scan_id: str, lease_owner: str, fencing_token: int, lease_seconds: int) -> Dict[str, Any]: + """Renew a still-valid lease, or raise :class:`LostLease`.""" + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") + conn = self._get_conn() + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + UPDATE scans + SET lease_expires_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second'), + last_heartbeat_at = CURRENT_TIMESTAMP + WHERE scan_id = %s + AND status = 'running' + AND lease_owner = %s + AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + RETURNING * + """, + (lease_seconds, scan_id, lease_owner, fencing_token), ) - RETURNING * - """, - (claimed_at,), - ) - row = cur.fetchone() - if row: - conn.commit() - return dict(row) - return None + row = cur.fetchone() + if row is None: + self.rollback(conn) + raise LostLease(f"Scan {scan_id} lease was lost before heartbeat") + conn.commit() + return dict(row) + except LostLease: + raise + except Exception: + self.rollback(conn) + raise - def recover_stale_scans(self, timeout_minutes: int = 60, max_attempts: int = 3) -> int: - """Recover scans left running after a worker crash or restart. + def recover_stale_scans(self, max_attempts: int = 3) -> int: + """Recover scans only after their renewable leases have expired. Stale scans are returned to pending while retry attempts remain. Once a scan has reached max_attempts, it is marked failed so it cannot loop forever on bad credentials or persistent Azure errors. """ conn = self._get_conn() - with conn.cursor() as cur: - cur.execute( - """ - UPDATE scans - SET status = 'failed', - error_message = 'Scan exceeded maximum retry attempts after worker interruption.' - WHERE status = 'running' - AND COALESCE(attempt_count, 1) >= %s - AND claimed_at < (CURRENT_TIMESTAMP - (%s * INTERVAL '1 minute')) - """, - (max_attempts, timeout_minutes), - ) - failed_count = cur.rowcount + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE scans + SET status = 'failed', + lease_owner = NULL, + lease_expires_at = NULL, + error_message = 'Scan exceeded maximum retry attempts after worker interruption.' + WHERE status = 'running' + AND COALESCE(attempt_count, 1) >= %s + AND lease_expires_at < CURRENT_TIMESTAMP + """, + (max_attempts,), + ) + failed_count = cur.rowcount - cur.execute( - """ - UPDATE scans - SET status = 'pending', - claimed_at = NULL, - error_message = 'Scan worker interrupted before completion. Queued for retry.' - WHERE status = 'running' - AND COALESCE(attempt_count, 0) < %s - AND claimed_at < (CURRENT_TIMESTAMP - (%s * INTERVAL '1 minute')) - """, - (max_attempts, timeout_minutes), - ) - retry_count = cur.rowcount - conn.commit() + cur.execute( + """ + UPDATE scans + SET status = 'pending', + claimed_at = NULL, + lease_owner = NULL, + lease_expires_at = NULL, + error_message = 'Scan worker interrupted before completion. Queued for retry.' + WHERE status = 'running' + AND COALESCE(attempt_count, 0) < %s + AND lease_expires_at < CURRENT_TIMESTAMP + """, + (max_attempts,), + ) + retry_count = cur.rowcount + conn.commit() + except Exception: + self.rollback(conn) + raise total_count = failed_count + retry_count if total_count > 0: logger.info( diff --git a/docs/async-scan-architecture.md b/docs/async-scan-architecture.md index 1d25d3a7..e9cf84f2 100644 --- a/docs/async-scan-architecture.md +++ b/docs/async-scan-architecture.md @@ -21,10 +21,17 @@ The scans table acts as a persistent task queue. This avoids the need for additi ### Render restart behavior Scan state is not stored in Flask memory. `POST /api/scans/trigger` inserts a `pending` row into PostgreSQL, and `GET /api/scans/` reads that same row back from PostgreSQL. If the Render web process restarts, queued scan state remains in the database and the dashboard can continue polling by `scan_id` after the app process comes back. -If the worker process restarts while a scan is marked `running`, `scanner/worker.py` calls `recover_stale_scans()` on each loop. Stale running scans are moved back to `pending` while retry attempts remain, so a Render restart can resume queued work instead of losing it. Once a scan reaches the maximum attempt count, it is marked `failed` so bad credentials or persistent Azure errors cannot retry forever. +Each claim is a renewable lease. The worker records a process-lifetime owner ID, +an expiry time, and a monotonically increasing fencing token, then renews the +lease while Azure work is running. `recover_stale_scans()` only requeues work +after its lease expires; a healthy worker is never reclaimed solely because its +original claim is old. A reclaimed scan receives a new token, so the previous +worker cannot persist completion, failure, or findings after it loses ownership. +Once a scan reaches the maximum attempt count, it is marked `failed` so bad +credentials or persistent Azure errors cannot retry forever. ### 3. The Worker (Python) -The scanner/worker.py process runs independently of the web server. Its lifecycle involves several steps. It queries the DB for scans where status is pending. It updates the status to running to prevent other workers from picking it up. It invokes ScanEngine.run_scan(scan_id). On success, it saves findings and sets status to completed. On failure, it captures the traceback and sets status to failed with the error_message. +The scanner/worker.py process runs independently of the web server. Its lifecycle involves several steps. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings and marks the scan complete atomically. On failure, it records a sanitized error only while it still owns the lease. ## Technical Rationale diff --git a/scanner/worker.py b/scanner/worker.py index 21a94c6f..f2f84b4a 100644 --- a/scanner/worker.py +++ b/scanner/worker.py @@ -7,11 +7,13 @@ import logging import os +import threading import time import traceback +import uuid from datetime import datetime, timezone -from api.models.finding import DatabaseManager +from api.models.finding import DatabaseManager, LostLease from api.observability import ( PENDING_SCANS, SCAN_DURATION_SECONDS, @@ -25,6 +27,102 @@ logger = logging.getLogger("scanner.worker") POLL_INTERVAL_SECONDS = 5 +DEFAULT_LEASE_SECONDS = 15 * 60 +DEFAULT_HEARTBEAT_SECONDS = 5 * 60 + + +def _positive_seconds(name: str, default: int) -> int: + """Read a positive interval without allowing malformed deploy config to stop work.""" + raw_value = os.environ.get(name) + if raw_value is None: + return default + try: + value = int(raw_value) + except ValueError: + logger.warning("Invalid %s=%r; using %d", name, raw_value, default) + return default + if value <= 0: + logger.warning("Invalid %s=%r; using %d", name, raw_value, default) + return default + return value + + +def lease_configuration() -> tuple[int, int]: + """Return a lease interval and a safely shorter heartbeat interval.""" + lease_seconds = _positive_seconds("SCAN_LEASE_SECONDS", DEFAULT_LEASE_SECONDS) + heartbeat_seconds = _positive_seconds("SCAN_HEARTBEAT_SECONDS", DEFAULT_HEARTBEAT_SECONDS) + if heartbeat_seconds >= lease_seconds: + heartbeat_seconds = max(1, lease_seconds // 3) + logger.warning( + "SCAN_HEARTBEAT_SECONDS must be less than SCAN_LEASE_SECONDS; using %d seconds", + heartbeat_seconds, + ) + return lease_seconds, heartbeat_seconds + + +class LeaseHeartbeat: + """Renew one claim through a dedicated database connection. + + Scan execution may block on Azure calls. The heartbeat intentionally owns + a separate DatabaseManager so it never shares a psycopg connection with + the worker's final persistence transaction. + """ + + def __init__( + self, + db_url: str, + scan_id: str, + lease_owner: str, + fencing_token: int, + lease_seconds: int, + heartbeat_seconds: int, + ) -> None: + self.db_url = db_url + self.scan_id = scan_id + self.lease_owner = lease_owner + self.fencing_token = fencing_token + self.lease_seconds = lease_seconds + self.heartbeat_seconds = heartbeat_seconds + self.lost = threading.Event() + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, name=f"scan-heartbeat-{scan_id}", daemon=True) + + def start(self) -> None: + self._thread.start() + + def stop(self) -> bool: + """Stop the heartbeat, reporting whether its thread actually exited.""" + self._stop.set() + self._thread.join(timeout=10) + if self._thread.is_alive(): + logger.critical("Heartbeat thread for scan %s did not stop promptly", self.scan_id) + return False + return True + + def _run(self) -> None: + db = DatabaseManager(self.db_url) + try: + while not self._stop.wait(self.heartbeat_seconds): + try: + db.heartbeat_scan( + self.scan_id, + self.lease_owner, + self.fencing_token, + self.lease_seconds, + ) + except LostLease: + self.lost.set() + logger.warning( + "Lease lost while scan %s was executing", self.scan_id, extra={"scan_id": self.scan_id} + ) + return + except Exception as exc: + # A transient DB failure must be visible and the next + # heartbeat must get a clean/reacquired connection. + logger.error("Heartbeat failed for scan %s: %s", self.scan_id, exc, exc_info=True) + db.rollback() + finally: + db.close() def run_worker(): @@ -38,24 +136,27 @@ def run_worker(): init_sentry() db = DatabaseManager(db_url) + worker_id = str(uuid.uuid4()) + lease_seconds, heartbeat_seconds = lease_configuration() logger.info("OpenShield Background Worker started. Polling every %ds", POLL_INTERVAL_SECONDS) while True: try: # 1. Cleanup stale scans from previous crashes - db.recover_stale_scans(timeout_minutes=60) + db.recover_stale_scans() # 2. Publish current queue depth PENDING_SCANS.set(len(db.get_pending_scans())) # 3. Atomic claim - scan = db.claim_next_pending_scan() + scan = db.claim_next_pending_scan(worker_id, lease_seconds) if not scan: time.sleep(POLL_INTERVAL_SECONDS) continue scan_id = str(scan["scan_id"]) subscription_id = scan["subscription_id"] + fencing_token = scan["fencing_token"] logger.info( "Starting scan %s for %s", @@ -65,6 +166,15 @@ def run_worker(): ) scan_start = time.perf_counter() + heartbeat = LeaseHeartbeat( + db_url, + scan_id, + worker_id, + fencing_token, + lease_seconds, + heartbeat_seconds, + ) + heartbeat.start() try: engine = ScanEngine(subscription_id) result = engine.run_scan(scan_id) @@ -73,14 +183,27 @@ def run_worker(): result["completed_at"] = datetime.now(timezone.utc).isoformat() result["status"] = "completed" - db.save_scan(result) + if not heartbeat.stop(): + # Do not start another scan while a database call in this + # heartbeat is still stuck; process supervision can restart + # us and the unconfirmed claim will safely expire. + return + if heartbeat.lost.is_set(): + raise LostLease(f"Scan {scan_id} lost its lease before completion") + db.save_scan(result, worker_id, fencing_token) SCANS_TOTAL.labels(status="completed").inc() logger.info( "Successfully completed scan %s", scan_id, extra={"scan_id": scan_id}, ) + except LostLease: + if not heartbeat.stop(): + return + logger.warning("Scan %s finished after its lease was lost; no result was persisted", scan_id) except Exception as exc: + if not heartbeat.stop(): + return error_msg = f"{str(exc)}\n{traceback.format_exc()}" SCANS_TOTAL.labels(status="failed").inc() logger.error( @@ -92,12 +215,23 @@ def run_worker(): # Sanitize public error message public_error = "An internal error occurred during the scan. Please check the logs." - db.update_scan_status(scan_id, "failed", error_message=public_error) + try: + db.update_scan_status( + scan_id, + "failed", + error_message=public_error, + lease_owner=worker_id, + fencing_token=fencing_token, + ) + except LostLease: + logger.warning("Scan %s failed after its lease was lost; failure was not persisted", scan_id) finally: + heartbeat.stop() SCAN_DURATION_SECONDS.observe(time.perf_counter() - scan_start) except Exception as exc: logger.error("Worker loop encountered an error: %s", exc) + db.rollback() time.sleep(POLL_INTERVAL_SECONDS) diff --git a/tests/test_async_scan_persistence.py b/tests/test_async_scan_persistence.py index 3bfbe4d6..dc022adc 100644 --- a/tests/test_async_scan_persistence.py +++ b/tests/test_async_scan_persistence.py @@ -7,7 +7,9 @@ from unittest.mock import MagicMock, patch -from api.models.finding import DatabaseManager +import pytest + +from api.models.finding import DatabaseManager, LostLease class _Cursor: @@ -105,20 +107,79 @@ def test_claim_next_pending_scan_increments_attempt_count(): """Claiming a pending scan should record a durable execution attempt.""" db = DatabaseManager.__new__(DatabaseManager) scan_id = "44444444-4444-4444-4444-444444444444" - cursor = _Cursor(rows=[{"scan_id": scan_id, "attempt_count": 1}]) + cursor = _Cursor(rows=[{"scan_id": scan_id, "attempt_count": 1, "fencing_token": 1}]) conn = MagicMock() conn.cursor.return_value = cursor with patch.object(db, "_get_conn", return_value=conn): - scan = db.claim_next_pending_scan() + scan = db.claim_next_pending_scan("worker-a", 900) executed_sql = cursor.calls[0][0] assert "attempt_count = COALESCE(attempt_count, 0) + 1" in executed_sql + assert "lease_owner = %s" in executed_sql + assert "lease_expires_at = CURRENT_TIMESTAMP" in executed_sql + assert "fencing_token = COALESCE(fencing_token, 0) + 1" in executed_sql assert "error_message = NULL" in executed_sql assert scan["scan_id"] == scan_id conn.commit.assert_called_once() +def test_empty_claim_commits_before_the_worker_can_sleep(): + """A no-work poll must not leave the long-lived worker transaction open.""" + db = DatabaseManager.__new__(DatabaseManager) + cursor = _Cursor(rows=[]) + conn = MagicMock() + conn.cursor.return_value = cursor + + with patch.object(db, "_get_conn", return_value=conn): + assert db.claim_next_pending_scan("worker-a", 900) is None + + conn.commit.assert_called_once() + + +def test_heartbeat_requires_current_owner_token_and_unexpired_lease(): + db = DatabaseManager.__new__(DatabaseManager) + cursor = _Cursor(rows=[]) + conn = MagicMock() + conn.closed = 0 + conn.info.transaction_status = 0 + conn.cursor.return_value = cursor + + with patch.object(db, "_get_conn", return_value=conn): + with pytest.raises(LostLease): + db.heartbeat_scan("scan-1", "worker-a", 3, 900) + + sql = cursor.calls[0][0] + assert "lease_owner = %s" in sql + assert "fencing_token = %s" in sql + assert "lease_expires_at > CURRENT_TIMESTAMP" in sql + conn.rollback.assert_called_once() + + +def test_fenced_failure_rejects_a_stale_owner(): + db = DatabaseManager.__new__(DatabaseManager) + cursor = _Cursor(rows=[]) + conn = MagicMock() + conn.closed = 0 + conn.info.transaction_status = 0 + conn.cursor.return_value = cursor + + with patch.object(db, "_get_conn", return_value=conn): + with pytest.raises(LostLease): + db.update_scan_status( + "scan-1", + "failed", + "failure", + lease_owner="worker-a", + fencing_token=3, + ) + + sql = cursor.calls[0][0] + assert "lease_owner = %s" in sql + assert "fencing_token = %s" in sql + assert "lease_expires_at > CURRENT_TIMESTAMP" in sql + + def test_recover_stale_scans_retries_before_max_attempts(): """Stale running scans should return to pending while attempts remain.""" db = DatabaseManager.__new__(DatabaseManager) @@ -127,15 +188,17 @@ def test_recover_stale_scans_retries_before_max_attempts(): conn.cursor.return_value = cursor with patch.object(db, "_get_conn", return_value=conn): - recovered = db.recover_stale_scans(timeout_minutes=15, max_attempts=3) + recovered = db.recover_stale_scans(max_attempts=3) failed_sql, failed_params = cursor.calls[0] retry_sql, retry_params = cursor.calls[1] assert "status = 'failed'" in failed_sql - assert failed_params == (3, 15) + assert failed_params == (3,) + assert "lease_expires_at < CURRENT_TIMESTAMP" in failed_sql assert "status = 'pending'" in retry_sql assert "claimed_at = NULL" in retry_sql - assert retry_params == (3, 15) + assert "lease_owner = NULL" in retry_sql + assert retry_params == (3,) assert recovered == 1 conn.commit.assert_called_once() @@ -148,13 +211,13 @@ def test_recover_stale_scans_fails_after_max_attempts(): conn.cursor.return_value = cursor with patch.object(db, "_get_conn", return_value=conn): - recovered = db.recover_stale_scans(timeout_minutes=60, max_attempts=3) + recovered = db.recover_stale_scans(max_attempts=3) failed_sql, failed_params = cursor.calls[0] retry_sql, retry_params = cursor.calls[1] assert "COALESCE(attempt_count, 1) >= %s" in failed_sql - assert failed_params == (3, 60) + assert failed_params == (3,) assert "COALESCE(attempt_count, 0) < %s" in retry_sql - assert retry_params == (3, 60) + assert retry_params == (3,) assert recovered == 1 conn.commit.assert_called_once() diff --git a/tests/test_observability.py b/tests/test_observability.py index 5fe30d90..ef9207ba 100644 --- a/tests/test_observability.py +++ b/tests/test_observability.py @@ -189,7 +189,7 @@ class _Stop(BaseException): mock_db = MagicMock() mock_db.recover_stale_scans.side_effect = [None, _Stop()] mock_db.claim_next_pending_scan.side_effect = [ - {"scan_id": scan_id, "subscription_id": sub_id}, + {"scan_id": scan_id, "subscription_id": sub_id, "fencing_token": 1}, None, ] mock_engine = MagicMock() @@ -203,6 +203,9 @@ class _Stop(BaseException): monkeypatch.setattr(worker, "DatabaseManager", MagicMock(return_value=mock_db)) monkeypatch.setattr(worker, "ScanEngine", MagicMock(return_value=mock_engine)) + heartbeat = MagicMock() + heartbeat.lost.is_set.return_value = False + monkeypatch.setattr(worker, "LeaseHeartbeat", MagicMock(return_value=heartbeat)) monkeypatch.setattr(worker.os.environ, "get", lambda *a, **k: "postgresql://x") monkeypatch.setattr(worker.time, "sleep", lambda *a, **k: None) diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py new file mode 100644 index 00000000..bb0c6efa --- /dev/null +++ b/tests/test_scan_leases_postgres.py @@ -0,0 +1,298 @@ +"""PostgreSQL-backed lease, fencing, and connection-recovery tests.""" + +import os +import threading +import uuid +from datetime import datetime, timezone + +import psycopg2 +from psycopg2 import extensions +import pytest + +from api.models.finding import DatabaseManager, LostLease + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + + +def _result(scan_id: str, subscription_id: str) -> dict: + return { + "scan_id": scan_id, + "subscription_id": subscription_id, + "started_at": datetime.now(timezone.utc).isoformat(), + "completed_at": datetime.now(timezone.utc).isoformat(), + "findings": [ + { + "rule_id": "AZ-LEASE-001", + "rule_name": "Lease test finding", + "severity": "HIGH", + "category": "Test", + "resource_id": f"/subscriptions/{subscription_id}/resourceGroups/test/providers/Test/resource", + "resource_name": "resource", + "resource_type": "Test/resource", + "description": "Lease test finding", + "remediation": "Fix the test resource", + "frameworks": {}, + "metadata": {}, + "detected_at": datetime.now(timezone.utc).isoformat(), + } + ], + } + + +class ScanRows: + def __init__(self, dsn: str): + self.dsn = dsn + self.scan_ids: list[str] = [] + + def create(self) -> tuple[str, str]: + scan_id = str(uuid.uuid4()) + subscription_id = str(uuid.uuid4()) + db = DatabaseManager(self.dsn) + try: + db.create_pending_scan(scan_id, subscription_id) + finally: + db.close() + self.scan_ids.append(scan_id) + return scan_id, subscription_id + + def expire(self, scan_id: str) -> None: + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "UPDATE scans SET lease_expires_at = CURRENT_TIMESTAMP - INTERVAL '1 second' WHERE scan_id = %s", + (scan_id,), + ) + + def scan(self, scan_id: str) -> dict: + db = DatabaseManager(self.dsn) + try: + scan = db.get_scan(scan_id) + assert scan is not None + return scan + finally: + db.close() + + def finding_count(self, scan_id: str) -> int: + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM findings WHERE scan_id = %s", (scan_id,)) + return cur.fetchone()[0] + + def cleanup(self) -> None: + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + for scan_id in self.scan_ids: + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) + + +@pytest.fixture +def scan_rows() -> ScanRows: + rows = ScanRows(os.environ["DATABASE_URL"]) + yield rows + rows.cleanup() + + +def _claim(dsn: str, owner: str) -> dict | None: + db = DatabaseManager(dsn) + try: + return db.claim_next_pending_scan(owner, 120) + finally: + db.close() + + +def _recover(dsn: str, max_attempts: int = 3) -> int: + db = DatabaseManager(dsn) + try: + return db.recover_stale_scans(max_attempts=max_attempts) + finally: + db.close() + + +def test_two_workers_race_to_claim_only_one_scan(scan_rows): + scan_rows.create() + barrier = threading.Barrier(2) + claims: list[dict | None] = [] + + def claim(owner: str) -> None: + barrier.wait() + claims.append(_claim(scan_rows.dsn, owner)) + + threads = [threading.Thread(target=claim, args=(owner,)) for owner in ("worker-a", "worker-b")] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + successful = [claim for claim in claims if claim is not None] + assert len(successful) == 1 + assert successful[0]["lease_owner"] in {"worker-a", "worker-b"} + assert successful[0]["fencing_token"] == 1 + + +def test_active_lease_cannot_be_reclaimed(scan_rows): + scan_rows.create() + assert _claim(scan_rows.dsn, "worker-a") is not None + assert _claim(scan_rows.dsn, "worker-b") is None + + +def test_heartbeat_extends_current_lease_without_changing_owner_or_token(scan_rows): + scan_id, _ = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a") + assert claim is not None + + db = DatabaseManager(scan_rows.dsn) + try: + renewed = db.heartbeat_scan(scan_id, "worker-a", claim["fencing_token"], 600) + finally: + db.close() + + assert renewed["lease_owner"] == "worker-a" + assert renewed["fencing_token"] == claim["fencing_token"] + assert renewed["lease_expires_at"] > claim["lease_expires_at"] + + +def test_expired_lease_is_reclaimed_with_new_fencing_token(scan_rows): + scan_id, _ = scan_rows.create() + first_claim = _claim(scan_rows.dsn, "worker-a") + assert first_claim is not None + scan_rows.expire(scan_id) + + assert _recover(scan_rows.dsn) == 1 + second_claim = _claim(scan_rows.dsn, "worker-b") + + assert second_claim is not None + assert second_claim["lease_owner"] == "worker-b" + assert second_claim["fencing_token"] > first_claim["fencing_token"] + + +def test_stale_worker_cannot_heartbeat_complete_fail_or_write_results(scan_rows): + scan_id, subscription_id = scan_rows.create() + first_claim = _claim(scan_rows.dsn, "worker-a") + assert first_claim is not None + scan_rows.expire(scan_id) + _recover(scan_rows.dsn) + second_claim = _claim(scan_rows.dsn, "worker-b") + assert second_claim is not None + + stale_db = DatabaseManager(scan_rows.dsn) + try: + with pytest.raises(LostLease): + stale_db.heartbeat_scan(scan_id, "worker-a", first_claim["fencing_token"], 120) + with pytest.raises(LostLease): + stale_db.update_scan_status( + scan_id, + "failed", + "stale worker failure", + lease_owner="worker-a", + fencing_token=first_claim["fencing_token"], + ) + with pytest.raises(LostLease): + stale_db.save_scan(_result(scan_id, subscription_id), "worker-a", first_claim["fencing_token"]) + finally: + stale_db.close() + + assert scan_rows.finding_count(scan_id) == 0 + scan = scan_rows.scan(scan_id) + assert scan["status"] == "running" + assert scan["lease_owner"] == "worker-b" + assert scan["fencing_token"] == second_claim["fencing_token"] + + +def test_current_owner_completion_persists_results_atomically(scan_rows): + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-b") + assert claim is not None + + db = DatabaseManager(scan_rows.dsn) + try: + db.save_scan(_result(scan_id, subscription_id), "worker-b", claim["fencing_token"]) + finally: + db.close() + + scan = scan_rows.scan(scan_id) + assert scan["status"] == "completed" + assert scan["lease_owner"] is None + assert scan_rows.finding_count(scan_id) == 1 + + +def test_sql_abort_rolls_back_and_the_connection_remains_usable(scan_rows): + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a") + assert claim is not None + broken_result = _result(scan_id, subscription_id) + broken_result["findings"][0]["detected_at"] = None + + db = DatabaseManager(scan_rows.dsn) + try: + with pytest.raises(psycopg2.Error): + db.save_scan(broken_result, "worker-a", claim["fencing_token"]) + db.update_scan_status( + scan_id, + "failed", + "controlled SQL failure", + lease_owner="worker-a", + fencing_token=claim["fencing_token"], + ) + finally: + db.close() + + assert scan_rows.scan(scan_id)["status"] == "failed" + assert scan_rows.finding_count(scan_id) == 0 + + +def test_terminated_backend_is_discarded_and_reacquired(scan_rows): + scan_id, _ = scan_rows.create() + db = DatabaseManager(scan_rows.dsn) + try: + conn = db._get_conn() + with conn.cursor() as cur: + cur.execute("SELECT pg_backend_pid()") + backend_pid = cur.fetchone()[0] + conn.commit() + + with psycopg2.connect(scan_rows.dsn) as terminator: + with terminator.cursor() as cur: + cur.execute("SELECT pg_terminate_backend(%s)", (backend_pid,)) + assert cur.fetchone()[0] is True + + with pytest.raises(psycopg2.OperationalError): + db.ping() + db.rollback() + + replacement = db._get_conn() + with replacement.cursor() as cur: + cur.execute("SELECT pg_backend_pid()") + assert cur.fetchone()[0] != backend_pid + db.rollback() + assert db.get_scan(scan_id) is not None + finally: + db.close() + + +def test_expired_restart_work_obeys_attempt_limit(scan_rows): + scan_id, _ = scan_rows.create() + assert _claim(scan_rows.dsn, "worker-a") is not None + scan_rows.expire(scan_id) + assert _recover(scan_rows.dsn, max_attempts=3) == 1 + + assert _claim(scan_rows.dsn, "worker-b") is not None + scan_rows.expire(scan_id) + assert _recover(scan_rows.dsn, max_attempts=3) == 1 + + assert _claim(scan_rows.dsn, "worker-c") is not None + scan_rows.expire(scan_id) + assert _recover(scan_rows.dsn, max_attempts=3) == 1 + assert scan_rows.scan(scan_id)["status"] == "failed" + + +def test_empty_claim_leaves_no_open_transaction(scan_rows): + db = DatabaseManager(scan_rows.dsn) + try: + assert db.claim_next_pending_scan("worker-a", 120) is None + assert db._get_conn().info.transaction_status == extensions.TRANSACTION_STATUS_IDLE + finally: + db.close() diff --git a/tests/test_severity_contract.py b/tests/test_severity_contract.py index 097ce33b..4a6822ec 100644 --- a/tests/test_severity_contract.py +++ b/tests/test_severity_contract.py @@ -138,7 +138,7 @@ def test_persistence_rejects_invalid_severity_before_opening_connection(): } with patch.object(db, "_get_conn") as get_conn: with pytest.raises(SeverityContractError): - db.save_scan(result) + db.save_scan(result, "worker-a", 1) get_conn.assert_not_called() @@ -172,14 +172,16 @@ def test_persistence_canonicalizes_alias_and_records_contract_version(): } with patch.object(db, "_get_conn", return_value=conn): - db.save_scan(result) - - scan_parameters = cursor.execute.call_args_list[0].args[1] - delete_parameters = cursor.execute.call_args_list[1].args[1] - finding_parameters = cursor.execute.call_args_list[2].args[1] - assert scan_parameters[4] == 1 - assert scan_parameters[5] == 100 - assert scan_parameters[10] == CONTRACT_VERSION + db.save_scan(result, "worker-a", 1) + + lock_parameters = cursor.execute.call_args_list[0].args[1] + scan_parameters = cursor.execute.call_args_list[1].args[1] + delete_parameters = cursor.execute.call_args_list[2].args[1] + finding_parameters = cursor.execute.call_args_list[3].args[1] + assert lock_parameters == (result["scan_id"], "worker-a", 1) + assert scan_parameters[1] == 1 + assert scan_parameters[2] == 100 + assert scan_parameters[4] == CONTRACT_VERSION assert delete_parameters == (result["scan_id"],) assert finding_parameters[0] == result["scan_id"] assert finding_parameters[3] == "INFO" @@ -191,8 +193,11 @@ def test_persistence_canonicalizes_alias_and_records_contract_version(): def test_persistence_rolls_back_a_failed_atomic_replacement(): db = _db() cursor = _cursor() - cursor.execute.side_effect = [None, None, RuntimeError("insert failed")] + cursor.fetchone.return_value = {"scan_id": "scan"} + cursor.execute.side_effect = [None, None, None, RuntimeError("insert failed")] conn = MagicMock() + conn.closed = 0 + conn.info.transaction_status = 0 conn.cursor.return_value = cursor result = { "scan_id": "00000000-0000-0000-0000-000000000000", @@ -203,7 +208,7 @@ def test_persistence_rolls_back_a_failed_atomic_replacement(): with patch.object(db, "_get_conn", return_value=conn): with pytest.raises(RuntimeError, match="insert failed"): - db.save_scan(result) + db.save_scan(result, "worker-a", 1) conn.rollback.assert_called_once_with() conn.commit.assert_not_called() diff --git a/tests/test_worker.py b/tests/test_worker.py index 29a00cf1..ddd0f5a0 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -8,7 +8,7 @@ """ import unittest -from unittest.mock import patch +from unittest.mock import ANY, patch from scanner.worker import run_worker, POLL_INTERVAL_SECONDS import uuid @@ -29,7 +29,10 @@ def setUp(self): @patch("scanner.worker.ScanEngine") @patch("scanner.worker.os.environ.get") @patch("scanner.worker.time.sleep") - def test_worker_processes_pending_scan_successfully(self, mock_sleep, mock_env, mock_engine_class, mock_db_class): + @patch("scanner.worker.LeaseHeartbeat") + def test_worker_processes_pending_scan_successfully( + self, mock_heartbeat_class, mock_sleep, mock_env, mock_engine_class, mock_db_class + ): """ Verify the happy path: 1. Worker claims a pending scan atomically. @@ -54,9 +57,10 @@ def test_worker_processes_pending_scan_successfully(self, mock_sleep, mock_env, # We need to stop the infinite loop. We'll raise StopWorker on the second call to recover_stale_scans. mock_db.recover_stale_scans.side_effect = [None, StopWorker()] mock_db.claim_next_pending_scan.side_effect = [ - {"scan_id": self.scan_id, "subscription_id": self.subscription_id}, + {"scan_id": self.scan_id, "subscription_id": self.subscription_id, "fencing_token": 1}, None, ] + mock_heartbeat_class.return_value.lost.is_set.return_value = False with self.assertRaises(StopWorker): run_worker() @@ -71,12 +75,16 @@ def test_worker_processes_pending_scan_successfully(self, mock_sleep, mock_env, saved_result = mock_db.save_scan.call_args[0][0] self.assertEqual(saved_result["status"], "completed") self.assertIn("completed_at", saved_result) + self.assertEqual(mock_db.save_scan.call_args[0][2], 1) @patch("scanner.worker.DatabaseManager") @patch("scanner.worker.ScanEngine") @patch("scanner.worker.os.environ.get") @patch("scanner.worker.time.sleep") - def test_worker_handles_scan_failure_gracefully(self, mock_sleep, mock_env, mock_engine_class, mock_db_class): + @patch("scanner.worker.LeaseHeartbeat") + def test_worker_handles_scan_failure_gracefully( + self, mock_heartbeat_class, mock_sleep, mock_env, mock_engine_class, mock_db_class + ): """ Verify the error path: 1. Worker claims a pending scan. @@ -88,9 +96,10 @@ def test_worker_handles_scan_failure_gracefully(self, mock_sleep, mock_env, mock mock_db.recover_stale_scans.side_effect = [None, StopWorker()] mock_db.claim_next_pending_scan.side_effect = [ - {"scan_id": self.scan_id, "subscription_id": self.subscription_id}, + {"scan_id": self.scan_id, "subscription_id": self.subscription_id, "fencing_token": 1}, None, ] + mock_heartbeat_class.return_value.lost.is_set.return_value = False # Mock Engine to fail mock_engine = mock_engine_class.return_value @@ -101,7 +110,11 @@ def test_worker_handles_scan_failure_gracefully(self, mock_sleep, mock_env, mock # Verify status was updated to failed with sanitized message mock_db.update_scan_status.assert_any_call( - self.scan_id, "failed", error_message="An internal error occurred during the scan. Please check the logs." + self.scan_id, + "failed", + error_message="An internal error occurred during the scan. Please check the logs.", + lease_owner=ANY, + fencing_token=1, ) # Ensure findings were NOT saved on failure mock_db.save_scan.assert_not_called() From d6883c194e34fe6522b971e831e1b87346ad5cdb Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 18:48:46 +0100 Subject: [PATCH 02/11] test(core): validate scan lease recovery Signed-off-by: Shaurya K Sharma --- api/models/finding.py | 8 ++- tests/test_database_manager_reliability.py | 21 +++++++ tests/test_worker.py | 73 +++++++++++++++++++++- 3 files changed, 100 insertions(+), 2 deletions(-) diff --git a/api/models/finding.py b/api/models/finding.py index d064b3fc..583f09b7 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -126,7 +126,13 @@ def __init__(self, dsn: Optional[str] = None) -> None: def connect(self) -> None: """Acquire a connection from this DSN's shared pool.""" if self.conn is not None: - self._return_connection(close=bool(self.conn.closed)) + previous_conn = self.conn + self.rollback(previous_conn) + # rollback() discards a dead connection and clears self.conn. A + # healthy connection still needs exactly one pool return before a + # replacement is borrowed. + if self.conn is previous_conn: + self._return_connection(close=bool(previous_conn.closed), conn=previous_conn) self.conn = _get_pool(self.dsn).getconn() self.conn.autocommit = False logger.debug("Database connection acquired from pool") diff --git a/tests/test_database_manager_reliability.py b/tests/test_database_manager_reliability.py index 06249280..d8a7573a 100644 --- a/tests/test_database_manager_reliability.py +++ b/tests/test_database_manager_reliability.py @@ -71,6 +71,27 @@ def test_connect_acquires_connection_from_shared_pool(): assert db.conn is fake_conn +def test_reconnect_rolls_back_and_returns_previous_connection_once(): + dsn = "postgresql://pool-test/db" + fake_pool = MagicMock() + old_conn = MagicMock() + old_conn.closed = 0 + old_conn.info.transaction_status = 0 + new_conn = MagicMock() + new_conn.closed = 0 + fake_pool.getconn.return_value = new_conn + db = _db(dsn) + db.conn = old_conn + + with patch.object(finding_module, "_get_pool", return_value=fake_pool): + db.connect() + + old_conn.rollback.assert_called_once() + fake_pool.putconn.assert_called_once_with(old_conn, close=False) + fake_pool.getconn.assert_called_once() + assert db.conn is new_conn + + def test_close_returns_healthy_connection_to_pool(): dsn = "postgresql://pool-test/db" fake_pool = MagicMock() diff --git a/tests/test_worker.py b/tests/test_worker.py index ddd0f5a0..73b051df 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -9,7 +9,8 @@ import unittest from unittest.mock import ANY, patch -from scanner.worker import run_worker, POLL_INTERVAL_SECONDS +from api.models.finding import LostLease +from scanner.worker import LeaseHeartbeat, POLL_INTERVAL_SECONDS, run_worker import uuid @@ -19,6 +20,20 @@ class StopWorker(BaseException): pass +class OneHeartbeatThenStop: + """Deterministically run one heartbeat loop iteration without sleeping.""" + + def __init__(self): + self.calls = 0 + + def wait(self, _seconds): + self.calls += 1 + return self.calls > 1 + + def set(self): + pass + + class TestWorker(unittest.TestCase): def setUp(self): self.mock_db_url = "postgresql://user:pass@localhost/db" @@ -76,6 +91,7 @@ def test_worker_processes_pending_scan_successfully( self.assertEqual(saved_result["status"], "completed") self.assertIn("completed_at", saved_result) self.assertEqual(mock_db.save_scan.call_args[0][2], 1) + self.assertGreaterEqual(mock_heartbeat_class.return_value.stop.call_count, 1) @patch("scanner.worker.DatabaseManager") @patch("scanner.worker.ScanEngine") @@ -118,6 +134,61 @@ def test_worker_handles_scan_failure_gracefully( ) # Ensure findings were NOT saved on failure mock_db.save_scan.assert_not_called() + self.assertGreaterEqual(mock_heartbeat_class.return_value.stop.call_count, 1) + + @patch("scanner.worker.DatabaseManager") + @patch("scanner.worker.ScanEngine") + @patch("scanner.worker.os.environ.get") + @patch("scanner.worker.time.sleep") + @patch("scanner.worker.LeaseHeartbeat") + def test_worker_does_not_persist_after_heartbeat_reports_lost_lease( + self, mock_heartbeat_class, mock_sleep, mock_env, mock_engine_class, mock_db_class + ): + mock_env.return_value = self.mock_db_url + mock_db = mock_db_class.return_value + mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_pending_scan.return_value = { + "scan_id": self.scan_id, + "subscription_id": self.subscription_id, + "fencing_token": 1, + } + mock_engine_class.return_value.run_scan.return_value = { + "scan_id": self.scan_id, + "subscription_id": self.subscription_id, + "findings": [], + } + mock_heartbeat_class.return_value.lost.is_set.return_value = True + + with self.assertRaises(StopWorker): + run_worker() + + mock_db.save_scan.assert_not_called() + mock_db.update_scan_status.assert_not_called() + + +@patch("scanner.worker.DatabaseManager") +def test_heartbeat_uses_and_closes_a_dedicated_database_manager(mock_db_class): + heartbeat = LeaseHeartbeat("postgresql://heartbeat-test/db", "scan-1", "worker-a", 7, 120, 1) + heartbeat._stop = OneHeartbeatThenStop() + + heartbeat._run() + + mock_db_class.assert_called_once_with("postgresql://heartbeat-test/db") + mock_db_class.return_value.heartbeat_scan.assert_called_once_with("scan-1", "worker-a", 7, 120) + mock_db_class.return_value.close.assert_called_once() + assert not heartbeat.lost.is_set() + + +@patch("scanner.worker.DatabaseManager") +def test_heartbeat_surfaces_lost_lease_and_closes_its_connection(mock_db_class): + mock_db_class.return_value.heartbeat_scan.side_effect = LostLease("stale") + heartbeat = LeaseHeartbeat("postgresql://heartbeat-test/db", "scan-1", "worker-a", 7, 120, 1) + heartbeat._stop = OneHeartbeatThenStop() + + heartbeat._run() + + assert heartbeat.lost.is_set() + mock_db_class.return_value.close.assert_called_once() @patch("scanner.worker.DatabaseManager") @patch("scanner.worker.os.environ.get") From bef9e3507e843a0231551112dadedf4c5d2de18f Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:25:03 +0100 Subject: [PATCH 03/11] fix(core): make scan result persistence idempotent Signed-off-by: Shaurya K Sharma --- ...4c9_idempotent_findings_and_evaluations.py | 66 ++++++++++ api/models/finding.py | 115 +++++++++++++++++- tests/test_scan_leases_postgres.py | 76 ++++++++++++ tests/test_severity_contract.py | 8 +- 4 files changed, 254 insertions(+), 11 deletions(-) create mode 100644 alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py diff --git a/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py b/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py new file mode 100644 index 00000000..7b4522a9 --- /dev/null +++ b/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py @@ -0,0 +1,66 @@ +"""Add database-enforced identities for scan results. + +Revision ID: f2b6d8e1a4c9 +Revises: e4f7a9b2c6d8 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + + +revision: str = "f2b6d8e1a4c9" +down_revision: Union[str, Sequence[str], None] = "e4f7a9b2c6d8" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Add stable finding keys and per-resource evaluation rows.""" + op.add_column("findings", sa.Column("finding_key", sa.Text(), nullable=True)) + # Existing records predate the identity contract. Preserve each record as + # distinct rather than attempting to infer equivalence from mutable text. + op.execute("UPDATE findings SET finding_key = 'legacy:' || id::text WHERE finding_key IS NULL") + op.alter_column("findings", "finding_key", nullable=False) + + with op.get_context().autocommit_block(): + op.execute("CREATE UNIQUE INDEX CONCURRENTLY uq_findings_scan_finding_key ON findings (scan_id, finding_key)") + + op.create_table( + "rule_evaluations", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("scan_id", postgresql.UUID(), nullable=False), + sa.Column("rule_id", sa.Text(), nullable=False), + sa.Column("resource_id", sa.Text(), nullable=False), + sa.Column("resource_type", sa.Text(), server_default=sa.text("''"), nullable=False), + sa.Column("status", sa.Text(), nullable=False), + sa.Column("reason_code", sa.Text(), nullable=True), + sa.Column("reason", sa.Text(), nullable=True), + sa.Column("evidence", postgresql.JSONB(), server_default=sa.text("'{}'::jsonb"), nullable=True), + sa.Column("finding_id", sa.Integer(), nullable=True), + sa.Column("evaluated_at", sa.DateTime(timezone=True), nullable=False), + sa.ForeignKeyConstraint(["scan_id"], ["scans.scan_id"], name="rule_evaluations_scan_id_fkey"), + sa.ForeignKeyConstraint( + ["finding_id"], ["findings.id"], name="rule_evaluations_finding_id_fkey", ondelete="SET NULL" + ), + sa.PrimaryKeyConstraint("id", name="rule_evaluations_pkey"), + sa.UniqueConstraint("scan_id", "rule_id", "resource_id", name="uq_rule_evaluations_scan_rule_resource"), + sa.CheckConstraint( + "status IN ('PASS', 'FAIL', 'UNKNOWN', 'ERROR', 'NOT_APPLICABLE')", + name="ck_rule_evaluations_status_v1", + ), + sa.CheckConstraint("resource_id <> ''", name="ck_rule_evaluations_resource_id_not_empty"), + ) + op.create_index("idx_rule_evaluations_scan_id", "rule_evaluations", ["scan_id"], unique=False) + + +def downgrade() -> None: + """Remove idempotent-result storage introduced by this revision.""" + op.drop_index("idx_rule_evaluations_scan_id", table_name="rule_evaluations") + op.drop_table("rule_evaluations") + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_findings_scan_finding_key") + op.drop_column("findings", "finding_key") diff --git a/api/models/finding.py b/api/models/finding.py index 583f09b7..8dc181f3 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -1,6 +1,7 @@ """Finding dataclass and PostgreSQL-backed DatabaseManager.""" import json +import hashlib import logging import os import threading @@ -41,6 +42,28 @@ class LostLease(RuntimeError): _POOL_MAX_CONN = int(os.environ.get("DB_POOL_MAX_CONN", "10")) +def stable_finding_key(scan_id: str, finding: Dict[str, Any]) -> str: + """Return an immutable identity for one logical finding in a scan. + + Rules that can report more than one violation for the same resource must + provide ``finding_discriminator``. Presentation fields such as severity, + description, and remediation are deliberately excluded so retries update + the existing authoritative finding rather than creating a duplicate. + """ + resource_scope = finding.get("resource_id") or { + "resource_type": finding.get("resource_type") or "", + "resource_name": finding.get("resource_name") or "", + } + identity = { + "scan_id": str(scan_id), + "rule_id": finding.get("rule_id") or "", + "resource_scope": resource_scope, + "discriminator": finding.get("finding_discriminator") or "default", + } + encoded = json.dumps(identity, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + def _get_pool(dsn: str) -> "psycopg2.pool.ThreadedConnectionPool": # Pool creation happens at most once per DSN per process lifetime, so a # plain lock (no unlocked fast path) is simpler and just as cheap here. @@ -223,7 +246,9 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token for raw_finding in scan_result.get("findings", []): finding = dict(raw_finding) finding["severity"] = normalize_severity(finding.get("severity")) + finding["finding_key"] = stable_finding_key(scan_result["scan_id"], finding) findings.append(finding) + evaluations = [dict(raw_evaluation) for raw_evaluation in scan_result.get("evaluations", [])] conn = self._get_conn() completed_at = scan_result.get("completed_at") or datetime.now(timezone.utc).isoformat() @@ -267,25 +292,39 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token scan_result["scan_id"], ), ) - # A worker retry replaces the previous result atomically. This - # keeps the scan header, child rows, and recomputed score in - # agreement instead of duplicating findings on every attempt. - cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_result["scan_id"],)) + finding_ids: Dict[tuple[str, str], int] = {} for f in findings: cur.execute( """ INSERT INTO findings - (scan_id, rule_id, rule_name, severity, category, + (scan_id, finding_key, rule_id, rule_name, severity, category, resource_id, resource_name, resource_type, description, remediation, playbook, frameworks, metadata, cve_references, cvss_score, exploit_available, detected_at) - VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) + VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) + ON CONFLICT (scan_id, finding_key) DO UPDATE SET + rule_name = EXCLUDED.rule_name, + severity = EXCLUDED.severity, + category = EXCLUDED.category, + resource_name = EXCLUDED.resource_name, + resource_type = EXCLUDED.resource_type, + description = EXCLUDED.description, + remediation = EXCLUDED.remediation, + playbook = EXCLUDED.playbook, + frameworks = EXCLUDED.frameworks, + metadata = EXCLUDED.metadata, + cve_references = EXCLUDED.cve_references, + cvss_score = EXCLUDED.cvss_score, + exploit_available = EXCLUDED.exploit_available, + detected_at = EXCLUDED.detected_at + RETURNING id """, ( # The parent scan owns every child in this batch. # Never trust a caller-supplied child scan_id. scan_result["scan_id"], + f["finding_key"], f.get("rule_id"), f.get("rule_name"), f.get("severity"), @@ -304,6 +343,70 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token f.get("detected_at"), ), ) + finding_row = cur.fetchone() + finding_id = finding_row["id"] if isinstance(finding_row, dict) else finding_row[0] + finding_ids[(f.get("rule_id") or "", f.get("resource_id") or "")] = finding_id + + finding_keys = [f["finding_key"] for f in findings] + if finding_keys: + cur.execute( + "DELETE FROM findings WHERE scan_id = %s AND NOT (finding_key = ANY(%s))", + (scan_result["scan_id"], finding_keys), + ) + else: + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_result["scan_id"],)) + + for evaluation in evaluations: + rule_id = evaluation.get("rule_id") + resource_id = evaluation.get("resource_id") + if not rule_id or not resource_id: + raise ValueError("evaluations require rule_id and resource_id") + cur.execute( + """ + INSERT INTO rule_evaluations + (scan_id, rule_id, resource_id, resource_type, status, + reason_code, reason, evidence, finding_id, evaluated_at) + VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) + ON CONFLICT (scan_id, rule_id, resource_id) DO UPDATE SET + resource_type = EXCLUDED.resource_type, + status = EXCLUDED.status, + reason_code = EXCLUDED.reason_code, + reason = EXCLUDED.reason, + evidence = EXCLUDED.evidence, + finding_id = EXCLUDED.finding_id, + evaluated_at = EXCLUDED.evaluated_at + """, + ( + scan_result["scan_id"], + rule_id, + resource_id, + evaluation.get("resource_type") or "", + evaluation.get("status"), + evaluation.get("reason_code"), + evaluation.get("reason"), + json.dumps(evaluation.get("evidence", {})), + finding_ids.get((rule_id, resource_id)), + completed_at, + ), + ) + + if evaluations: + cur.execute( + """ + DELETE FROM rule_evaluations + WHERE scan_id = %s + AND (rule_id, resource_id) NOT IN ( + SELECT * FROM unnest(%s::text[], %s::text[]) + ) + """, + ( + scan_result["scan_id"], + [evaluation["rule_id"] for evaluation in evaluations], + [evaluation["resource_id"] for evaluation in evaluations], + ), + ) + else: + cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_result["scan_id"],)) conn.commit() except Exception: # psycopg2 connections remain in an aborted transaction after any diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py index bb0c6efa..7b2fa7df 100644 --- a/tests/test_scan_leases_postgres.py +++ b/tests/test_scan_leases_postgres.py @@ -81,10 +81,31 @@ def finding_count(self, scan_id: str) -> int: cur.execute("SELECT COUNT(*) FROM findings WHERE scan_id = %s", (scan_id,)) return cur.fetchone()[0] + def rearm(self, scan_id: str, owner: str, fencing_token: int) -> None: + """Simulate duplicate delivery of the same claimed result for persistence tests.""" + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE scans + SET status = 'running', completed_at = NULL, lease_owner = %s, + fencing_token = %s, lease_expires_at = CURRENT_TIMESTAMP + INTERVAL '5 minutes' + WHERE scan_id = %s + """, + (owner, fencing_token, scan_id), + ) + + def evaluation_count(self, scan_id: str) -> int: + with psycopg2.connect(self.dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) + return cur.fetchone()[0] + def cleanup(self) -> None: with psycopg2.connect(self.dsn) as conn: with conn.cursor() as cur: for scan_id in self.scan_ids: + cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) @@ -296,3 +317,58 @@ def test_empty_claim_leaves_no_open_transaction(scan_rows): assert db._get_conn().info.transaction_status == extensions.TRANSACTION_STATUS_IDLE finally: db.close() + + +def test_duplicate_result_delivery_upserts_mutable_fields_and_evaluations(scan_rows): + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a") + assert claim is not None + result = _result(scan_id, subscription_id) + result["evaluations"] = [ + { + "rule_id": "AZ-LEASE-001", + "resource_id": result["findings"][0]["resource_id"], + "resource_type": "Test/resource", + "status": "FAIL", + "evidence": {"version": 1}, + } + ] + + db = DatabaseManager(scan_rows.dsn) + try: + db.save_scan(result, "worker-a", claim["fencing_token"]) + first_id = db.get_findings({"scan_id": scan_id})[0]["id"] + scan_rows.rearm(scan_id, "worker-a", claim["fencing_token"]) + result["findings"][0]["description"] = "Updated presentation text" + result["findings"][0]["severity"] = "CRITICAL" + result["evaluations"][0]["evidence"] = {"version": 2} + db.save_scan(result, "worker-a", claim["fencing_token"]) + persisted = db.get_findings({"scan_id": scan_id}) + finally: + db.close() + + assert len(persisted) == 1 + assert persisted[0]["id"] == first_id + assert persisted[0]["description"] == "Updated presentation text" + assert persisted[0]["severity"] == "CRITICAL" + assert scan_rows.evaluation_count(scan_id) == 1 + + +def test_distinct_finding_discriminators_preserve_multiple_violations(scan_rows): + scan_id, subscription_id = scan_rows.create() + claim = _claim(scan_rows.dsn, "worker-a") + assert claim is not None + result = _result(scan_id, subscription_id) + duplicate_scope = dict(result["findings"][0]) + result["findings"][0]["finding_discriminator"] = "network-rule-a" + duplicate_scope["finding_discriminator"] = "network-rule-b" + duplicate_scope["description"] = "A second violation on the same resource" + result["findings"].append(duplicate_scope) + + db = DatabaseManager(scan_rows.dsn) + try: + db.save_scan(result, "worker-a", claim["fencing_token"]) + finally: + db.close() + + assert scan_rows.finding_count(scan_id) == 2 diff --git a/tests/test_severity_contract.py b/tests/test_severity_contract.py index 4a6822ec..53388af6 100644 --- a/tests/test_severity_contract.py +++ b/tests/test_severity_contract.py @@ -176,15 +176,13 @@ def test_persistence_canonicalizes_alias_and_records_contract_version(): lock_parameters = cursor.execute.call_args_list[0].args[1] scan_parameters = cursor.execute.call_args_list[1].args[1] - delete_parameters = cursor.execute.call_args_list[2].args[1] - finding_parameters = cursor.execute.call_args_list[3].args[1] + finding_parameters = cursor.execute.call_args_list[2].args[1] assert lock_parameters == (result["scan_id"], "worker-a", 1) assert scan_parameters[1] == 1 assert scan_parameters[2] == 100 assert scan_parameters[4] == CONTRACT_VERSION - assert delete_parameters == (result["scan_id"],) assert finding_parameters[0] == result["scan_id"] - assert finding_parameters[3] == "INFO" + assert finding_parameters[4] == "INFO" assert raw_finding["severity"] == "INFORMATIONAL" conn.commit.assert_called_once_with() conn.rollback.assert_not_called() @@ -194,7 +192,7 @@ def test_persistence_rolls_back_a_failed_atomic_replacement(): db = _db() cursor = _cursor() cursor.fetchone.return_value = {"scan_id": "scan"} - cursor.execute.side_effect = [None, None, None, RuntimeError("insert failed")] + cursor.execute.side_effect = [None, None, RuntimeError("insert failed")] conn = MagicMock() conn.closed = 0 conn.info.transaction_status = 0 From 26d30a89668d8d89a3da1bcd74759b5e7b416f00 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:28:24 +0100 Subject: [PATCH 04/11] fix(core): make scan admission durable and idempotent Signed-off-by: Shaurya K Sharma --- ...a7c5e9d2f1b4_scan_admission_idempotency.py | 47 +++++++ api/models/finding.py | 115 +++++++++++++++--- api/routes/scans.py | 53 +++++++- tests/test_async_scan_persistence.py | 4 +- tests/test_error_exposure.py | 2 +- tests/test_input_validation.py | 3 +- tests/test_scan_admission_postgres.py | 100 +++++++++++++++ 7 files changed, 303 insertions(+), 21 deletions(-) create mode 100644 alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py create mode 100644 tests/test_scan_admission_postgres.py diff --git a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py new file mode 100644 index 00000000..a45eb609 --- /dev/null +++ b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py @@ -0,0 +1,47 @@ +"""Enforce durable scan admission and idempotency. + +Revision ID: a7c5e9d2f1b4 +Revises: f2b6d8e1a4c9 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "a7c5e9d2f1b4" +down_revision: Union[str, Sequence[str], None] = "f2b6d8e1a4c9" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Persist idempotency semantics and prevent more than one active scan.""" + op.add_column("scans", sa.Column("idempotency_key", sa.Text(), nullable=True)) + op.add_column("scans", sa.Column("request_fingerprint", sa.Text(), nullable=True)) + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE UNIQUE INDEX CONCURRENTLY uq_scans_subscription_idempotency_key + ON scans (subscription_id, idempotency_key) + WHERE idempotency_key IS NOT NULL + """ + ) + op.execute( + """ + CREATE UNIQUE INDEX CONCURRENTLY uq_scans_one_active_per_subscription + ON scans (subscription_id) + WHERE status IN ('pending', 'running') + """ + ) + + +def downgrade() -> None: + """Remove scan admission metadata and constraints.""" + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_scans_one_active_per_subscription") + op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_scans_subscription_idempotency_key") + op.drop_column("scans", "request_fingerprint") + op.drop_column("scans", "idempotency_key") diff --git a/api/models/finding.py b/api/models/finding.py index 8dc181f3..adccf1f3 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -29,6 +29,14 @@ class LostLease(RuntimeError): """Raised when a worker no longer owns the scan it is trying to update.""" +class ScanAdmissionConflict(RuntimeError): + """Raised when an idempotency key is reused for different scan semantics.""" + + +class ScanQuotaExceeded(RuntimeError): + """Raised when an explicitly configured subscription scan quota is exhausted.""" + + FRAMEWORKS_DIR = Path(__file__).parent.parent.parent / "compliance" / "frameworks" # One pool per DSN, shared across all DatabaseManager instances in this @@ -517,22 +525,101 @@ def update_scan_enrichment_status(self, scan_id: str, status: str) -> None: conn.commit() logger.info("Updated scan %s enrichment status to %s", scan_id, status) - def create_pending_scan(self, scan_id: str, subscription_id: str) -> None: - """Create a scan record in the 'pending' state.""" + def admit_scan( + self, + scan_id: str, + subscription_id: str, + *, + idempotency_key: Optional[str] = None, + request_fingerprint: Optional[str] = None, + max_scans_per_hour: int = 0, + ) -> tuple[Dict[str, Any], bool]: + """Atomically admit one scan or return its durable logical predecessor. + + The PostgreSQL advisory transaction lock serializes admission decisions + for one subscription. The partial unique index added by the migration + remains the final database enforcement of the one-active-scan rule. + """ + if max_scans_per_hour < 0: + raise ValueError("max_scans_per_hour must not be negative") conn = self._get_conn() - from datetime import datetime, timezone + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (subscription_id,)) + if idempotency_key: + cur.execute( + """ + SELECT * FROM scans + WHERE subscription_id = %s AND idempotency_key = %s + """, + (subscription_id, idempotency_key), + ) + existing = cur.fetchone() + if existing: + existing = dict(existing) + if existing.get("request_fingerprint") != request_fingerprint: + raise ScanAdmissionConflict("Idempotency-Key was reused with different request semantics") + conn.commit() + return existing, False - started_at = datetime.now(timezone.utc).isoformat() - with conn.cursor() as cur: - cur.execute( - """ - INSERT INTO scans (scan_id, subscription_id, started_at, status, attempt_count) - VALUES (%s, %s, %s, 'pending', 0) - """, - (scan_id, subscription_id, started_at), - ) - conn.commit() - logger.info("Created pending scan %s for %s", scan_id, subscription_id) + cur.execute( + """ + SELECT * FROM scans + WHERE subscription_id = %s AND status IN ('pending', 'running') + ORDER BY started_at ASC + LIMIT 1 + """, + (subscription_id,), + ) + active_scan = cur.fetchone() + if active_scan: + conn.commit() + return dict(active_scan), False + + if max_scans_per_hour: + cur.execute( + """ + SELECT COUNT(*) FROM scans + WHERE subscription_id = %s + AND started_at >= CURRENT_TIMESTAMP - INTERVAL '1 hour' + """, + (subscription_id,), + ) + quota_row = cur.fetchone() + scan_count = quota_row["count"] if isinstance(quota_row, dict) else quota_row[0] + if scan_count >= max_scans_per_hour: + raise ScanQuotaExceeded("Configured hourly scan quota has been reached") + + from datetime import datetime, timezone + + cur.execute( + """ + INSERT INTO scans ( + scan_id, subscription_id, started_at, status, attempt_count, + idempotency_key, request_fingerprint + ) + VALUES (%s, %s, %s, 'pending', 0, %s, %s) + RETURNING * + """, + ( + scan_id, + subscription_id, + datetime.now(timezone.utc).isoformat(), + idempotency_key, + request_fingerprint, + ), + ) + admitted = dict(cur.fetchone()) + conn.commit() + logger.info("Admitted pending scan %s for %s", scan_id, subscription_id) + return admitted, True + except Exception: + self.rollback(conn) + raise + + def create_pending_scan(self, scan_id: str, subscription_id: str) -> None: + """Create a pending scan for older internal callers without a key.""" + self.admit_scan(scan_id, subscription_id) def update_scan_status( self, diff --git a/api/routes/scans.py b/api/routes/scans.py index 9ec2a289..109635f7 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -2,11 +2,13 @@ import logging import os +import hashlib +import json import threading import uuid from flask import Blueprint, g, jsonify, request -from api.models.finding import DatabaseManager +from api.models.finding import DatabaseManager, ScanAdmissionConflict, ScanQuotaExceeded from api.validation import ( VALIDATION_ERROR_MESSAGE, ValidationError, @@ -20,6 +22,7 @@ logger = logging.getLogger(__name__) _AUTHORIZED_SUBSCRIPTIONS_ENV = "OPENSHIELD_AUTHORIZED_SUBSCRIPTIONS" +_MAX_SCANS_PER_HOUR_ENV = "OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR" def _subscription_is_authorized(subscription_id: str) -> bool: @@ -53,6 +56,17 @@ def _get_db() -> DatabaseManager: return g.db +def _configured_hourly_quota() -> int: + """Return an optional policy quota; zero preserves existing no-limit policy.""" + raw_value = os.environ.get(_MAX_SCANS_PER_HOUR_ENV, "0") + try: + quota = int(raw_value) + except ValueError: + logger.warning("Invalid %s=%r; disabling hourly quota", _MAX_SCANS_PER_HOUR_ENV, raw_value) + return 0 + return max(0, quota) + + @scans_bp.get("/api/scans") def list_scans(): """Return all historical scan results ordered by most recent first.""" @@ -105,18 +119,49 @@ def trigger_scan(): logger.warning("Scan trigger rejected: subscription %s is not on the authorized allowlist", subscription_id) return jsonify({"error": "Subscription is not authorized for this deployment"}), 403 + idempotency_key = request.headers.get("Idempotency-Key") + if idempotency_key is not None: + idempotency_key = idempotency_key.strip() + if not idempotency_key or len(idempotency_key) > 200: + return jsonify({"error": "Idempotency-Key must be between 1 and 200 characters"}), 400 + request_fingerprint = hashlib.sha256( + json.dumps({"subscription_id": subscription_id}, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() scan_id = str(uuid.uuid4()) - logger.info("Async scan triggered for subscription %s (id: %s)", subscription_id, scan_id) try: db = _get_db() - db.create_pending_scan(scan_id, subscription_id) + admitted, created = db.admit_scan( + scan_id, + subscription_id, + idempotency_key=idempotency_key, + request_fingerprint=request_fingerprint, + max_scans_per_hour=_configured_hourly_quota(), + ) + except ScanAdmissionConflict as exc: + return jsonify({"error": str(exc)}), 409 + except ScanQuotaExceeded as exc: + return jsonify({"error": str(exc)}), 429 except Exception as exc: logger.error("Failed to create pending scan: %s", exc, exc_info=True) return jsonify({"error": "Database error"}), 500 + response_scan_id = str(admitted["scan_id"]) + if not created: + return jsonify( + { + "scan_id": response_scan_id, + "status": admitted["status"], + "message": "Existing logical scan returned.", + } + ), 200 + logger.info("Async scan admitted for subscription %s (id: %s)", subscription_id, response_scan_id) return jsonify( - {"scan_id": scan_id, "status": "pending", "message": "Scan has been queued and will start shortly."} + { + "scan_id": response_scan_id, + "status": "pending", + "message": "Scan has been queued and will start shortly.", + } ), 202 except ValidationError: diff --git a/tests/test_async_scan_persistence.py b/tests/test_async_scan_persistence.py index dc022adc..2883377d 100644 --- a/tests/test_async_scan_persistence.py +++ b/tests/test_async_scan_persistence.py @@ -42,6 +42,7 @@ def test_trigger_scan_persists_pending_scan_to_database(client, auth_headers, mo scan_id = "11111111-1111-1111-1111-111111111111" subscription_id = "00000000-0000-0000-0000-000000000000" mock_db = MagicMock() + mock_db.admit_scan.return_value = ({"scan_id": scan_id, "status": "pending"}, True) with patch("api.routes.scans.DatabaseManager", return_value=mock_db) as db_class: with patch("api.routes.scans.uuid.uuid4", return_value=scan_id): @@ -59,7 +60,8 @@ def test_trigger_scan_persists_pending_scan_to_database(client, auth_headers, mo } db_class.assert_called_once_with("postgresql://ci:ci@localhost/ci_db") mock_db.connect.assert_called_once() - mock_db.create_pending_scan.assert_called_once_with(scan_id, subscription_id) + mock_db.admit_scan.assert_called_once() + assert mock_db.admit_scan.call_args.args[:2] == (scan_id, subscription_id) def test_get_scan_status_reads_from_database(client, auth_headers, monkeypatch): diff --git a/tests/test_error_exposure.py b/tests/test_error_exposure.py index f098ccdb..cea508a5 100644 --- a/tests/test_error_exposure.py +++ b/tests/test_error_exposure.py @@ -46,7 +46,7 @@ def test_get_scan_status_error_does_not_leak_exception(client, auth_headers): def test_trigger_scan_db_error_does_not_leak_exception(client, auth_headers): - with patch.object(scans_route, "_get_db", return_value=_raising_db("create_pending_scan")): + with patch.object(scans_route, "_get_db", return_value=_raising_db("admit_scan")): resp = client.post( "/api/scans/trigger", json={"subscription_id": "00000000-0000-0000-0000-000000000001"}, diff --git a/tests/test_input_validation.py b/tests/test_input_validation.py index 86931256..8bca8054 100644 --- a/tests/test_input_validation.py +++ b/tests/test_input_validation.py @@ -60,10 +60,11 @@ def test_trigger_rejects_malformed_subscription_id(client, auth_headers): def test_trigger_accepts_canonical_subscription_uuid(client, auth_headers): db = MagicMock() + db.admit_scan.return_value = ({"scan_id": _SCAN_ID, "status": "pending"}, True) with patch.object(scans_route, "_get_db", return_value=db): response = client.post("/api/scans/trigger", json={"subscription_id": _SUBSCRIPTION_ID}, headers=auth_headers) assert response.status_code == 202 - assert db.create_pending_scan.call_args.args[1] == _SUBSCRIPTION_ID + assert db.admit_scan.call_args.args[1] == _SUBSCRIPTION_ID @pytest.mark.parametrize( diff --git a/tests/test_scan_admission_postgres.py b/tests/test_scan_admission_postgres.py new file mode 100644 index 00000000..9c0fa829 --- /dev/null +++ b/tests/test_scan_admission_postgres.py @@ -0,0 +1,100 @@ +"""Real PostgreSQL tests for durable scan admission invariants.""" + +import os +import threading +import uuid + +import psycopg2 +import pytest + +from api.models.finding import DatabaseManager, ScanAdmissionConflict, ScanQuotaExceeded + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + + +@pytest.fixture +def admitted_scans(): + dsn = os.environ["DATABASE_URL"] + scan_ids: list[str] = [] + yield dsn, scan_ids + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + for scan_id in scan_ids: + cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) + + +def _admit(dsn: str, subscription_id: str, key: str | None = None, fingerprint: str = "same"): + db = DatabaseManager(dsn) + try: + return db.admit_scan( + str(uuid.uuid4()), + subscription_id, + idempotency_key=key, + request_fingerprint=fingerprint if key else None, + ) + finally: + db.close() + + +def test_concurrent_admission_returns_one_active_scan(admitted_scans): + dsn, scan_ids = admitted_scans + subscription_id = str(uuid.uuid4()) + barrier = threading.Barrier(2) + outcomes = [] + + def admit() -> None: + barrier.wait() + outcomes.append(_admit(dsn, subscription_id)) + + threads = [threading.Thread(target=admit) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + scan_ids.append(str(outcomes[0][0]["scan_id"])) + assert {str(scan["scan_id"]) for scan, _created in outcomes} == {scan_ids[0]} + assert sum(created for _scan, created in outcomes) == 1 + + +def test_idempotency_key_replays_or_rejects_changed_semantics(admitted_scans): + dsn, scan_ids = admitted_scans + subscription_id = str(uuid.uuid4()) + first, created = _admit(dsn, subscription_id, "request-1", "fingerprint-a") + scan_ids.append(str(first["scan_id"])) + replay, replay_created = _admit(dsn, subscription_id, "request-1", "fingerprint-a") + + assert created is True + assert replay_created is False + assert replay["scan_id"] == first["scan_id"] + with pytest.raises(ScanAdmissionConflict): + _admit(dsn, subscription_id, "request-1", "fingerprint-b") + + +def test_completed_scan_allows_a_later_admission_and_configured_quota(admitted_scans): + dsn, scan_ids = admitted_scans + subscription_id = str(uuid.uuid4()) + first, _ = _admit(dsn, subscription_id) + scan_ids.append(str(first["scan_id"])) + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("UPDATE scans SET status = 'completed' WHERE scan_id = %s", (first["scan_id"],)) + + second, created = _admit(dsn, subscription_id) + scan_ids.append(str(second["scan_id"])) + assert created is True + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("UPDATE scans SET status = 'completed' WHERE scan_id = %s", (second["scan_id"],)) + + db = DatabaseManager(dsn) + try: + with pytest.raises(ScanQuotaExceeded): + db.admit_scan(str(uuid.uuid4()), subscription_id, max_scans_per_hour=2) + finally: + db.close() From b2ae11fc1032983d46cad64016fa9aba31d3f55f Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:37:45 +0100 Subject: [PATCH 05/11] fix(core): make CVE enrichment durable Signed-off-by: Shaurya K Sharma --- .../c9e1a5b7d3f2_durable_enrichment_jobs.py | 69 +++++ api/models/finding.py | 268 ++++++++++++++++++ api/routes/scans.py | 99 +------ scanner/cve_correlator.py | 18 +- scanner/enrichment_worker.py | 50 ++++ scanner/nvd_client.py | 123 ++++---- scanner/worker.py | 10 +- tests/test_enrichment_jobs_postgres.py | 186 ++++++++++++ tests/test_nvd_client.py | 12 + tests/test_scans_enrich.py | 124 ++------ tests/test_worker.py | 4 + 11 files changed, 714 insertions(+), 249 deletions(-) create mode 100644 alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py create mode 100644 scanner/enrichment_worker.py create mode 100644 tests/test_enrichment_jobs_postgres.py diff --git a/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py b/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py new file mode 100644 index 00000000..ab9c3a0a --- /dev/null +++ b/alembic/versions/c9e1a5b7d3f2_durable_enrichment_jobs.py @@ -0,0 +1,69 @@ +"""Add durable, fenced CVE enrichment jobs. + +Revision ID: c9e1a5b7d3f2 +Revises: a7c5e9d2f1b4 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + + +revision: str = "c9e1a5b7d3f2" +down_revision: Union[str, Sequence[str], None] = "a7c5e9d2f1b4" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Create one resumable enrichment job per scan.""" + op.create_table( + "enrichment_jobs", + sa.Column("job_id", postgresql.UUID(), nullable=False), + sa.Column("scan_id", postgresql.UUID(), nullable=False), + sa.Column("status", sa.Text(), nullable=False, server_default=sa.text("'pending'")), + sa.Column("lease_owner", sa.Text(), nullable=True), + sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("fencing_token", sa.BigInteger(), nullable=False, server_default=sa.text("0")), + sa.Column("attempt_count", sa.Integer(), nullable=False, server_default=sa.text("0")), + sa.Column( + "next_retry_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("CURRENT_TIMESTAMP") + ), + sa.Column("checkpoint", sa.Integer(), nullable=False, server_default=sa.text("0")), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column( + "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("CURRENT_TIMESTAMP") + ), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.ForeignKeyConstraint(["scan_id"], ["scans.scan_id"], name="enrichment_jobs_scan_id_fkey"), + sa.PrimaryKeyConstraint("job_id", name="enrichment_jobs_pkey"), + sa.UniqueConstraint("scan_id", name="uq_enrichment_jobs_scan_id"), + sa.CheckConstraint("status IN ('pending', 'running', 'completed', 'failed')", name="ck_enrichment_jobs_status"), + ) + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_enrichment_jobs_pending_retry + ON enrichment_jobs (next_retry_at ASC) + WHERE status = 'pending' + """ + ) + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_enrichment_jobs_running_lease + ON enrichment_jobs (lease_expires_at ASC) + WHERE status = 'running' + """ + ) + + +def downgrade() -> None: + """Remove durable enrichment work state.""" + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_enrichment_jobs_running_lease") + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_enrichment_jobs_pending_retry") + op.drop_table("enrichment_jobs") diff --git a/api/models/finding.py b/api/models/finding.py index adccf1f3..fe3eaf94 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -5,6 +5,7 @@ import logging import os import threading +import uuid from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Optional @@ -525,6 +526,273 @@ def update_scan_enrichment_status(self, scan_id: str, status: str) -> None: conn.commit() logger.info("Updated scan %s enrichment status to %s", scan_id, status) + def enqueue_enrichment_job(self, scan_id: str) -> tuple[Dict[str, Any], bool]: + """Durably enqueue exactly one CVE enrichment job for a scan.""" + conn = self._get_conn() + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + INSERT INTO enrichment_jobs (job_id, scan_id, status, attempt_count, checkpoint) + VALUES (%s, %s, 'pending', 0, 0) + ON CONFLICT (scan_id) DO NOTHING + RETURNING * + """, + (str(uuid.uuid4()), scan_id), + ) + job = cur.fetchone() + if job is None: + cur.execute("SELECT * FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + existing = cur.fetchone() + if existing is None: + raise RuntimeError("enrichment job conflict did not return an existing job") + conn.commit() + return dict(existing), False + cur.execute( + "UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", + (scan_id,), + ) + conn.commit() + return dict(job), True + except Exception: + self.rollback(conn) + raise + + def claim_next_enrichment_job(self, lease_owner: str, lease_seconds: int) -> Optional[Dict[str, Any]]: + """Atomically claim the next retry-ready enrichment job.""" + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") + conn = self._get_conn() + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + UPDATE enrichment_jobs + SET status = 'running', lease_owner = %s, + lease_expires_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second'), + last_heartbeat_at = CURRENT_TIMESTAMP, + fencing_token = fencing_token + 1, + attempt_count = attempt_count + 1, + error_message = NULL + WHERE job_id = ( + SELECT job_id FROM enrichment_jobs + WHERE status = 'pending' + AND next_retry_at <= CURRENT_TIMESTAMP + ORDER BY created_at ASC + FOR UPDATE SKIP LOCKED + LIMIT 1 + ) + RETURNING * + """, + (lease_owner, lease_seconds), + ) + job = cur.fetchone() + if job: + cur.execute( + "UPDATE scans SET cve_enrichment_status = 'ENRICHING' WHERE scan_id = %s", (job["scan_id"],) + ) + conn.commit() + return dict(job) if job else None + except Exception: + self.rollback(conn) + raise + + def heartbeat_enrichment_job(self, job_id: str, lease_owner: str, fencing_token: int, lease_seconds: int) -> None: + """Renew an enrichment claim or raise LostLease.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE enrichment_jobs + SET lease_expires_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second'), + last_heartbeat_at = CURRENT_TIMESTAMP + WHERE job_id = %s AND status = 'running' + AND lease_owner = %s AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + """, + (lease_seconds, job_id, lease_owner, fencing_token), + ) + if cur.rowcount != 1: + raise LostLease(f"Enrichment job {job_id} is no longer owned by this worker") + conn.commit() + except Exception: + self.rollback(conn) + raise + + def get_enrichment_findings(self, scan_id: str) -> List[Dict[str, Any]]: + """Return a deterministic snapshot ordered for checkpointed enrichment.""" + conn = self._get_conn() + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute("SELECT * FROM findings WHERE scan_id = %s ORDER BY id ASC", (scan_id,)) + return [dict(row) for row in cur.fetchall()] + + def persist_enrichment_progress( + self, + job_id: str, + lease_owner: str, + fencing_token: int, + finding: Dict[str, Any], + checkpoint: int, + ) -> None: + """Persist one finding and checkpoint it under the current job fence.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + SELECT scan_id FROM enrichment_jobs + WHERE job_id = %s AND status = 'running' + AND lease_owner = %s AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + FOR UPDATE + """, + (job_id, lease_owner, fencing_token), + ) + if cur.fetchone() is None: + raise LostLease(f"Enrichment job {job_id} lost its lease before checkpointing") + cur.execute( + """ + UPDATE findings SET cve_references = %s, cvss_score = %s, exploit_available = %s + WHERE id = %s + """, + ( + json.dumps(finding.get("cve_references", [])), + finding.get("cvss_score"), + finding.get("exploit_available", False), + finding["id"], + ), + ) + cur.execute("UPDATE enrichment_jobs SET checkpoint = %s WHERE job_id = %s", (checkpoint, job_id)) + conn.commit() + except Exception: + self.rollback(conn) + raise + + def complete_enrichment_job(self, job_id: str, lease_owner: str, fencing_token: int) -> None: + """Atomically complete a fenced job and its scan's enrichment state.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE enrichment_jobs + SET status = 'completed', completed_at = CURRENT_TIMESTAMP, + lease_owner = NULL, lease_expires_at = NULL + WHERE job_id = %s AND status = 'running' + AND lease_owner = %s AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + RETURNING scan_id + """, + (job_id, lease_owner, fencing_token), + ) + row = cur.fetchone() + if row is None: + raise LostLease(f"Enrichment job {job_id} lost its lease before completion") + scan_id = row[0] + cur.execute("UPDATE scans SET cve_enrichment_status = 'COMPLETED' WHERE scan_id = %s", (scan_id,)) + conn.commit() + except Exception: + self.rollback(conn) + raise + + def fail_enrichment_job( + self, + job_id: str, + lease_owner: str, + fencing_token: int, + error_message: str, + *, + max_attempts: int = 3, + retry_seconds: int = 30, + ) -> str: + """Record bounded retry state under the job's current fencing token.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + SELECT scan_id, attempt_count FROM enrichment_jobs + WHERE job_id = %s AND status = 'running' + AND lease_owner = %s AND fencing_token = %s + AND lease_expires_at > CURRENT_TIMESTAMP + FOR UPDATE + """, + (job_id, lease_owner, fencing_token), + ) + row = cur.fetchone() + if row is None: + raise LostLease(f"Enrichment job {job_id} is no longer owned by this worker") + scan_id, attempt_count = row + terminal = attempt_count >= max_attempts + if terminal: + cur.execute( + """ + UPDATE enrichment_jobs SET status = 'failed', completed_at = CURRENT_TIMESTAMP, + error_message = %s, lease_owner = NULL, lease_expires_at = NULL + WHERE job_id = %s + """, + (error_message, job_id), + ) + scan_status = "FAILED" + outcome = "failed" + else: + cur.execute( + """ + UPDATE enrichment_jobs SET status = 'pending', error_message = %s, + lease_owner = NULL, lease_expires_at = NULL, + next_retry_at = CURRENT_TIMESTAMP + (%s * INTERVAL '1 second') + WHERE job_id = %s + """, + (error_message, retry_seconds, job_id), + ) + scan_status = "PENDING" + outcome = "retry" + cur.execute("UPDATE scans SET cve_enrichment_status = %s WHERE scan_id = %s", (scan_status, scan_id)) + conn.commit() + return outcome + except Exception: + self.rollback(conn) + raise + + def recover_stale_enrichment_jobs(self, max_attempts: int = 3) -> int: + """Return expired enrichment claims to pending or terminally fail them.""" + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE enrichment_jobs SET status = 'failed', completed_at = CURRENT_TIMESTAMP, + lease_owner = NULL, lease_expires_at = NULL, + error_message = 'Enrichment exceeded maximum retry attempts after worker interruption.' + WHERE status = 'running' AND attempt_count >= %s + AND lease_expires_at < CURRENT_TIMESTAMP + RETURNING scan_id + """, + (max_attempts,), + ) + failed_scans = [row[0] for row in cur.fetchall()] + cur.execute( + """ + UPDATE enrichment_jobs SET status = 'pending', lease_owner = NULL, lease_expires_at = NULL, + error_message = 'Enrichment worker interrupted; queued for retry.' + WHERE status = 'running' AND attempt_count < %s + AND lease_expires_at < CURRENT_TIMESTAMP + RETURNING scan_id + """, + (max_attempts,), + ) + retried_scans = [row[0] for row in cur.fetchall()] + for scan_id in failed_scans: + cur.execute("UPDATE scans SET cve_enrichment_status = 'FAILED' WHERE scan_id = %s", (scan_id,)) + for scan_id in retried_scans: + cur.execute("UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", (scan_id,)) + conn.commit() + return len(failed_scans) + len(retried_scans) + except Exception: + self.rollback(conn) + raise + def admit_scan( self, scan_id: str, diff --git a/api/routes/scans.py b/api/routes/scans.py index 109635f7..ec90033a 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -4,7 +4,6 @@ import os import hashlib import json -import threading import uuid from flask import Blueprint, g, jsonify, request @@ -16,7 +15,6 @@ require_json_object, uuid_string, ) -from scanner.cve_correlator import enrich_findings scans_bp = Blueprint("scans", __name__) logger = logging.getLogger(__name__) @@ -171,81 +169,14 @@ def trigger_scan(): return jsonify({"error": "Critical route failure"}), 500 -def _run_enrichment_in_background(scan_id: str, findings: list, db_url: str) -> None: - """Run CVE enrichment off the request thread and persist the result. - - Runs outside the Flask request/app context (it's started via - threading.Thread), so it opens its own DatabaseManager rather than - reusing flask.g. - """ - db = DatabaseManager(db_url) - try: - enriched = enrich_findings(findings) - db.update_cve_fields(enriched) - db.update_scan_enrichment_status(scan_id, "COMPLETED") - logger.info("Background CVE enrichment complete for scan %s (%d findings)", scan_id, len(enriched)) - except Exception as exc: - logger.error("Background enrichment failed for scan %s: %s", scan_id, exc) - try: - # A failed write (e.g. in update_cve_fields) can leave db.conn in - # an aborted-transaction state. Roll back first, or this status - # update itself raises InFailedSqlTransaction and gets swallowed - # below, leaving the scan stuck at ENRICHING forever. - if db.conn is not None: - db.conn.rollback() - db.update_scan_enrichment_status(scan_id, "FAILED") - except Exception as status_exc: - logger.error("Failed to record FAILED status for scan %s: %s", scan_id, status_exc) - finally: - db.close() - - -def _run_enrichment_in_background(scan_id: str, findings: list, db_url: str) -> None: - """Run CVE enrichment off the request thread and persist the result. - - Runs outside the Flask request/app context (it's started via - threading.Thread), so it opens its own DatabaseManager rather than - reusing flask.g. - """ - db = DatabaseManager(db_url) - try: - enriched = enrich_findings(findings) - db.update_cve_fields(enriched) - db.update_scan_enrichment_status(scan_id, "COMPLETED") - logger.info("Background CVE enrichment complete for scan %s (%d findings)", scan_id, len(enriched)) - except Exception as exc: - logger.error("Background enrichment failed for scan %s: %s", scan_id, exc) - try: - # A failed write (e.g. in update_cve_fields) can leave db.conn in - # an aborted-transaction state. Roll back first, or this status - # update itself raises InFailedSqlTransaction and gets swallowed - # below, leaving the scan stuck at ENRICHING forever. - if db.conn is not None: - db.conn.rollback() - db.update_scan_enrichment_status(scan_id, "FAILED") - except Exception as status_exc: - logger.error("Failed to record FAILED status for scan %s: %s", scan_id, status_exc) - finally: - db.close() - - @scans_bp.post("/api/scans//enrich") def enrich_scan(scan_id): - """Kick off CVE enrichment for an existing scan in the background. - - Returns immediately with 202 and status ENRICHING; poll - GET /api/scans/ (cve_enrichment_status) for completion. - Running enrichment synchronously here previously caused the request to - time out on scans spanning many rule categories, since NVD lookups are - rate-limited to one every ~7 seconds. - """ + """Enqueue durable CVE enrichment; no request-owned thread is created.""" try: scan_id = uuid_string(scan_id, "scan_id") db = _get_db() - # Check current status to avoid redundant NVD calls - scans = db.get_scans() - current_scan = next((s for s in scans if str(s["scan_id"]) == scan_id), None) + current_scan = db.get_scan(scan_id) if not current_scan: return jsonify({"error": "Scan not found"}), 404 @@ -253,27 +184,27 @@ def enrich_scan(scan_id): status = current_scan.get("cve_enrichment_status") if status == "COMPLETED": return jsonify({"message": "Scan already enriched", "scan_id": scan_id}), 200 - if status == "ENRICHING": - return jsonify({"message": "Enrichment already in progress", "scan_id": scan_id}), 202 - findings = db.get_findings({"scan_id": scan_id}) if not findings: return jsonify({"error": "No findings found for this scan"}), 404 - logger.info("Starting background CVE enrichment for %d findings in scan %s", len(findings), scan_id) - db.update_scan_enrichment_status(scan_id, "ENRICHING") - - threading.Thread( - target=_run_enrichment_in_background, - args=(scan_id, findings, db.dsn), - daemon=True, - ).start() + job, created = db.enqueue_enrichment_job(scan_id) + if not created: + return jsonify( + { + "job_id": str(job["job_id"]), + "scan_id": scan_id, + "status": job["status"], + "message": "Existing enrichment job returned.", + } + ), 202 return jsonify( { "scan_id": scan_id, - "status": "ENRICHING", - "message": "CVE enrichment started; poll GET /api/scans/ for completion.", + "job_id": str(job["job_id"]), + "status": "PENDING", + "message": "CVE enrichment queued; poll GET /api/scans/ for completion.", } ), 202 diff --git a/scanner/cve_correlator.py b/scanner/cve_correlator.py index 0ac3de4a..94d4af0b 100644 --- a/scanner/cve_correlator.py +++ b/scanner/cve_correlator.py @@ -10,7 +10,7 @@ import logging from typing import Optional -from scanner.nvd_client import query_nvd +from scanner.nvd_client import query_nvd, query_nvd_strict logger = logging.getLogger(__name__) @@ -132,3 +132,19 @@ def enrich_findings(findings: list[dict]) -> list[dict]: enriched = [_enrich_single_finding(f) for f in findings] logger.info("CVE enrichment complete.") return enriched + + +def enrich_finding_durable(finding: dict) -> dict: + """Enrich one persisted finding and propagate retriable NVD failures.""" + keyword = _get_nvd_keyword(finding.get("rule_id", "")) + if not keyword: + finding["cve_references"] = [] + finding["cvss_score"] = None + finding["exploit_available"] = False + return finding + cves = query_nvd_strict(keyword) + finding["cve_references"] = cves + scores = [c["cvss_score"] for c in cves if c.get("cvss_score") is not None] + finding["cvss_score"] = max(scores) if scores else None + finding["exploit_available"] = any(c.get("exploit_available") for c in cves) + return finding diff --git a/scanner/enrichment_worker.py b/scanner/enrichment_worker.py new file mode 100644 index 00000000..4ca92c44 --- /dev/null +++ b/scanner/enrichment_worker.py @@ -0,0 +1,50 @@ +"""Durable CVE enrichment job execution for the scan worker.""" + +import logging + +from api.models.finding import DatabaseManager, LostLease +from scanner.cve_correlator import enrich_finding_durable + + +logger = logging.getLogger("scanner.enrichment_worker") + + +def process_enrichment_job( + db: DatabaseManager, + job: dict, + lease_owner: str, + lease_seconds: int, + *, + max_attempts: int = 3, +) -> str: + """Run a claimed job, checkpointing each finding under its fence. + + A restart resumes at the stored finding offset. Replaying the last item is + safe because persistence updates the existing finding row by primary key. + """ + job_id = str(job["job_id"]) + fencing_token = job["fencing_token"] + try: + findings = db.get_enrichment_findings(str(job["scan_id"])) + for index, finding in enumerate(findings[job["checkpoint"] :], start=job["checkpoint"]): + db.heartbeat_enrichment_job(job_id, lease_owner, fencing_token, lease_seconds) + enriched = enrich_finding_durable(finding) + db.persist_enrichment_progress(job_id, lease_owner, fencing_token, enriched, index + 1) + db.complete_enrichment_job(job_id, lease_owner, fencing_token) + logger.info("Completed enrichment job %s", job_id) + return "completed" + except LostLease: + logger.warning("Enrichment job %s lost its lease; no further writes attempted", job_id) + return "lost_lease" + except Exception as exc: + retry_seconds = min(300, 30 * (2 ** max(0, job["attempt_count"] - 1))) + outcome = db.fail_enrichment_job( + job_id, + lease_owner, + fencing_token, + "CVE enrichment failed; see worker logs for details.", + max_attempts=max_attempts, + retry_seconds=retry_seconds, + ) + logger.error("Enrichment job %s failed (%s): %s", job_id, outcome, exc, exc_info=True) + return outcome diff --git a/scanner/nvd_client.py b/scanner/nvd_client.py index 989ea0de..8ca65666 100644 --- a/scanner/nvd_client.py +++ b/scanner/nvd_client.py @@ -32,13 +32,17 @@ _NVD_BASE_URL = "https://services.nvd.nist.gov/rest/json/cves/2.0" _REQUEST_DELAY_SECONDS = 7.0 # Stay under 5 req/30 sec limit _MAX_RETRIES = 3 -_RESULTS_PER_PAGE = 5 # Top 5 CVEs per finding is enough for display +_RESULTS_PER_PAGE = 2000 # NVD's documented maximum; fetch every matching page. # In-memory cache. Keyed by "keyword:results_per_page". # Resets each process - intentional, NVD data changes slowly. _cache: dict[str, list[dict]] = {} +class NvdRequestError(RuntimeError): + """Raised by the durable worker when an NVD page cannot be retrieved.""" + + class _RateLimiter: """Tracks the last NVD request time, guarded by a lock so concurrent callers (e.g. background enrichment threads) can't both read a stale @@ -128,25 +132,12 @@ def _parse_cve_item(item: dict) -> Optional[dict]: return None -def query_nvd(keyword: str, results_per_page: int = _RESULTS_PER_PAGE) -> list[dict]: - """ - Query NVD for CVEs matching a keyword. - - Returns a list of parsed CVE dicts (may be empty). - Never raises - all failures return []. - - Args: - keyword: Search term, e.g. "Azure Storage Account" - results_per_page: Max CVEs to fetch (default 5) - """ - cache_key = f"{keyword}:{results_per_page}" - if cache_key in _cache: - logger.debug("NVD cache hit for: %s", keyword) - return _cache[cache_key] - +def _fetch_nvd_page(keyword: str, start_index: int, results_per_page: int) -> dict: + """Fetch one NVD page, raising only after bounded retries are exhausted.""" params = urllib.parse.urlencode( { "keywordSearch": keyword, + "startIndex": start_index, "resultsPerPage": results_per_page, } ) @@ -157,58 +148,74 @@ def query_nvd(keyword: str, results_per_page: int = _RESULTS_PER_PAGE) -> list[d or parsed_url.hostname != "services.nvd.nist.gov" or parsed_url.port not in (None, 443) ): - logger.error("Refusing request to an untrusted NVD endpoint") - return [] + raise NvdRequestError("Refusing request to an untrusted NVD endpoint") + last_error: Optional[Exception] = None for attempt in range(1, _MAX_RETRIES + 1): try: _wait_for_rate_limit() - logger.debug("NVD query (attempt %d): %s", attempt, keyword) - req = urllib.request.Request( url, headers={"User-Agent": "OpenShield/0.1 (github.com/openshield-org/openshield)"}, ) - # URL host is the hardcoded NVD API base, not user-controlled with NVD_REQUEST_LATENCY_SECONDS.time(): - # The URL is built from the fixed HTTPS NVD endpoint; only its query string varies. - with urllib.request.urlopen( # nosec B310 # nosemgrep: python.lang.security.audit.dynamic-urllib-use-detected.dynamic-urllib-use-detected # noqa: E501 + with urllib.request.urlopen( # nosec B310 # noqa: E501 req, timeout=10 ) as resp: - data = json.loads(resp.read()) - - vulnerabilities = data.get("vulnerabilities", []) - results = [parsed for item in vulnerabilities if (parsed := _parse_cve_item(item)) is not None] - - _cache[cache_key] = results - logger.info("NVD returned %d CVEs for: %s", len(results), keyword) - return results - - except urllib.error.HTTPError as e: - if e.code == 429: - wait = 30 * attempt # Back off harder each retry - logger.warning( - "NVD rate limited (429). Waiting %ds before retry %d/%d", - wait, - attempt, - _MAX_RETRIES, - ) - time.sleep(wait) - else: - logger.warning("NVD HTTP %d for keyword '%s': %s", e.code, keyword, e) - break # Non-rate-limit HTTP errors won't improve on retry - - except Exception as e: - logger.warning( - "NVD query failed (attempt %d/%d) for '%s': %s", - attempt, - _MAX_RETRIES, - keyword, - e, - ) + return json.loads(resp.read()) + except urllib.error.HTTPError as exc: + last_error = exc + if exc.code != 429: + break + time.sleep(30 * attempt) + except Exception as exc: + last_error = exc if attempt < _MAX_RETRIES: time.sleep(2**attempt) + raise NvdRequestError(f"NVD request failed for {keyword!r}: {last_error}") + - logger.warning("NVD lookup failed for '%s' - returning empty list", keyword) - _cache[cache_key] = [] # Cache the failure to avoid hammering NVD - return [] +def query_nvd_strict(keyword: str, results_per_page: int = _RESULTS_PER_PAGE) -> list[dict]: + """Return every NVD result, propagating terminal retrieval failure to the job queue.""" + cache_key = f"strict:{keyword}:{results_per_page}" + if cache_key in _cache: + return _cache[cache_key] + + start_index = 0 + total_results: Optional[int] = None + results: list[dict] = [] + while total_results is None or start_index < total_results: + page = _fetch_nvd_page(keyword, start_index, results_per_page) + vulnerabilities = page.get("vulnerabilities", []) + results.extend(parsed for item in vulnerabilities if (parsed := _parse_cve_item(item)) is not None) + total_results = int(page.get("totalResults", len(vulnerabilities))) + if not vulnerabilities: + break + start_index += len(vulnerabilities) + _cache[cache_key] = results + return results + + +def query_nvd(keyword: str, results_per_page: int = _RESULTS_PER_PAGE) -> list[dict]: + """ + Query NVD for CVEs matching a keyword. + + Returns a list of parsed CVE dicts (may be empty). + Never raises - all failures return []. + + Args: + keyword: Search term, e.g. "Azure Storage Account" + results_per_page: Max CVEs to fetch (default 5) + """ + cache_key = f"{keyword}:{results_per_page}" + if cache_key in _cache: + logger.debug("NVD cache hit for: %s", keyword) + return _cache[cache_key] + + try: + results = query_nvd_strict(keyword, results_per_page) + except NvdRequestError as exc: + logger.warning("NVD lookup failed for '%s': %s", keyword, exc) + results = [] + _cache[cache_key] = results + return results diff --git a/scanner/worker.py b/scanner/worker.py index f2f84b4a..0a3404cf 100644 --- a/scanner/worker.py +++ b/scanner/worker.py @@ -22,6 +22,7 @@ init_sentry, ) from scanner.engine import ScanEngine +from scanner.enrichment_worker import process_enrichment_job configure_logging() logger = logging.getLogger("scanner.worker") @@ -144,11 +145,18 @@ def run_worker(): try: # 1. Cleanup stale scans from previous crashes db.recover_stale_scans() + db.recover_stale_enrichment_jobs() # 2. Publish current queue depth PENDING_SCANS.set(len(db.get_pending_scans())) - # 3. Atomic claim + # 3. Run one durable enrichment job before taking another scan. + enrichment_job = db.claim_next_enrichment_job(worker_id, lease_seconds) + if enrichment_job: + process_enrichment_job(db, enrichment_job, worker_id, lease_seconds) + continue + + # 4. Atomic scan claim scan = db.claim_next_pending_scan(worker_id, lease_seconds) if not scan: time.sleep(POLL_INTERVAL_SECONDS) diff --git a/tests/test_enrichment_jobs_postgres.py b/tests/test_enrichment_jobs_postgres.py new file mode 100644 index 00000000..ff659efe --- /dev/null +++ b/tests/test_enrichment_jobs_postgres.py @@ -0,0 +1,186 @@ +"""PostgreSQL lease, retry, and resume tests for durable enrichment jobs.""" + +import os +import threading +import uuid +from unittest.mock import patch + +import psycopg2 +import pytest + +from api.models.finding import DatabaseManager, LostLease +from scanner.enrichment_worker import process_enrichment_job + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + + +@pytest.fixture +def enrichment_scan(): + dsn = os.environ["DATABASE_URL"] + scan_id, subscription_id = str(uuid.uuid4()), str(uuid.uuid4()) + db = DatabaseManager(dsn) + try: + db.create_pending_scan(scan_id, subscription_id) + claim = db.claim_next_pending_scan("seed", 120) + result = { + "scan_id": scan_id, + "subscription_id": subscription_id, + "findings": [ + { + "rule_id": "AZ-STOR-001", + "rule_name": "test", + "severity": "HIGH", + "resource_id": f"/subscriptions/{subscription_id}/resources/one", + "resource_name": "one", + "resource_type": "Test/resource", + "detected_at": "2026-08-29T00:00:00+00:00", + }, + { + "rule_id": "AZ-STOR-001", + "rule_name": "test", + "severity": "HIGH", + "resource_id": f"/subscriptions/{subscription_id}/resources/two", + "resource_name": "two", + "resource_type": "Test/resource", + "detected_at": "2026-08-29T00:00:00+00:00", + }, + ], + } + db.save_scan(result, "seed", claim["fencing_token"]) + job, _ = db.enqueue_enrichment_job(scan_id) + yield dsn, scan_id, job + finally: + db.close() + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) + + +def _claim(dsn): + db = DatabaseManager(dsn) + try: + return db.claim_next_enrichment_job("worker-a", 120) + finally: + db.close() + + +def test_duplicate_enqueue_and_claim_race(enrichment_scan): + dsn, scan_id, first_job = enrichment_scan + db = DatabaseManager(dsn) + try: + replay, created = db.enqueue_enrichment_job(scan_id) + finally: + db.close() + assert created is False + assert replay["job_id"] == first_job["job_id"] + + barrier = threading.Barrier(2) + claims = [] + + def claim(): + barrier.wait() + claims.append(_claim(dsn)) + + threads = [threading.Thread(target=claim) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert len([job for job in claims if job]) == 1 + + +def test_checkpoint_resume_and_completion(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + job = _claim(dsn) + assert job is not None + db = DatabaseManager(dsn) + try: + with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: + calls = 0 + + def enrich_once_then_fail(finding): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("transient NVD failure") + return {**finding, "cve_references": [{"cve_id": "CVE-1"}]} + + enrich.side_effect = enrich_once_then_fail + assert process_enrichment_job(db, job, "worker-a", 120) == "retry" + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT status, checkpoint FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone() == ("pending", 1) + cur.execute( + "UPDATE enrichment_jobs SET next_retry_at = CURRENT_TIMESTAMP WHERE scan_id = %s", (scan_id,) + ) + resumed = db.claim_next_enrichment_job("worker-b", 120) + with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: + enrich.side_effect = lambda finding: {**finding, "cve_references": [{"cve_id": "CVE-1"}]} + assert process_enrichment_job(db, resumed, "worker-b", 120) == "completed" + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT status, checkpoint FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone() == ("completed", 2) + cur.execute( + "SELECT COUNT(*) FROM findings WHERE scan_id = %s AND cve_references <> '[]'::jsonb", (scan_id,) + ) + assert cur.fetchone()[0] == 2 + finally: + db.close() + + +def test_enrichment_retry_limit_becomes_terminal(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + for attempt in range(1, 4): + job = db.claim_next_enrichment_job("worker-a", 120) + assert job is not None + with patch("scanner.enrichment_worker.enrich_finding_durable", side_effect=RuntimeError("NVD unavailable")): + expected = "failed" if attempt == 3 else "retry" + assert process_enrichment_job(db, job, "worker-a", 120) == expected + if attempt < 3: + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "UPDATE enrichment_jobs SET next_retry_at = CURRENT_TIMESTAMP WHERE scan_id = %s", + (scan_id,), + ) + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT status FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone()[0] == "failed" + finally: + db.close() + + +def test_expired_job_is_recovered_with_new_token_and_stale_owner_is_rejected(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + first = _claim(dsn) + assert first is not None + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + UPDATE enrichment_jobs + SET lease_expires_at = CURRENT_TIMESTAMP - INTERVAL '1 second' + WHERE scan_id = %s + """, + (scan_id,), + ) + db = DatabaseManager(dsn) + try: + assert db.recover_stale_enrichment_jobs() == 1 + second = db.claim_next_enrichment_job("worker-b", 120) + assert second["fencing_token"] > first["fencing_token"] + with pytest.raises(LostLease): + db.heartbeat_enrichment_job(str(first["job_id"]), "worker-a", first["fencing_token"], 120) + finally: + db.close() diff --git a/tests/test_nvd_client.py b/tests/test_nvd_client.py index 03ce424a..60100668 100644 --- a/tests/test_nvd_client.py +++ b/tests/test_nvd_client.py @@ -220,6 +220,18 @@ def test_second_call_uses_cache(self, mock_wait, mock_urlopen): query_nvd("Azure Storage Account") # Should be served from cache self.assertEqual(mock_urlopen.call_count, 1) + @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client._wait_for_rate_limit") + def test_fetches_every_nvd_page(self, mock_wait, mock_urlopen): + first_page = {"totalResults": 2, "vulnerabilities": [_SAMPLE_NVD_RESPONSE["vulnerabilities"][0]]} + second_page = {"totalResults": 2, "vulnerabilities": [_SAMPLE_NVD_RESPONSE["vulnerabilities"][1]]} + mock_urlopen.side_effect = [_make_mock_urlopen_response(first_page), _make_mock_urlopen_response(second_page)] + + results = query_nvd("Azure Storage Account", results_per_page=1) + + self.assertEqual([item["cve_id"] for item in results], ["CVE-2023-12345", "CVE-2022-99999"]) + self.assertEqual(mock_urlopen.call_count, 2) + @patch("scanner.nvd_client.urllib.request.urlopen") @patch("scanner.nvd_client._wait_for_rate_limit") def test_returns_empty_list_on_network_error(self, mock_wait, mock_urlopen): diff --git a/tests/test_scans_enrich.py b/tests/test_scans_enrich.py index 44c42c82..80addb6d 100644 --- a/tests/test_scans_enrich.py +++ b/tests/test_scans_enrich.py @@ -1,59 +1,41 @@ -"""Tests for POST /api/scans//enrich backgrounding (REL-004).""" +"""Tests for durable POST /api/scans//enrich job admission.""" from unittest.mock import MagicMock, patch import api.routes.scans as scans_route + _SCAN_ID = "00000000-0000-0000-0000-000000000001" def _mock_db(current_scan=None, findings=None): db = MagicMock() - db.get_scans.return_value = [current_scan] if current_scan else [] + db.get_scan.return_value = current_scan db.get_findings.return_value = findings if findings is not None else [] return db -class _FakeThread: - """Records what would have been threaded, and runs synchronously on start().""" - - last_instance = None - - def __init__(self, target, args, daemon): - self.target = target - self.args = args - self.daemon = daemon - self.started = False - _FakeThread.last_instance = self - - def start(self): - self.started = True - - -def test_enrich_returns_202_and_schedules_background_thread(client, auth_headers, monkeypatch): - monkeypatch.setenv("DATABASE_URL", "postgresql://mock/mock") +def test_enrich_returns_202_and_enqueues_durable_job(client, auth_headers): scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} - findings = [{"id": 1, "rule_id": "AZ-STOR-001"}] - db = _mock_db(current_scan=scan, findings=findings) + db = _mock_db(current_scan=scan, findings=[{"id": 1, "rule_id": "AZ-STOR-001"}]) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "pending"}, True) - with ( - patch.object(scans_route, "_get_db", return_value=db), - patch.object(scans_route.threading, "Thread", _FakeThread), - ): + with patch.object(scans_route, "_get_db", return_value=db): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 202 - body = resp.get_json() - assert body["status"] == "ENRICHING" + assert resp.get_json()["status"] == "PENDING" + db.enqueue_enrichment_job.assert_called_once_with(_SCAN_ID) - db.update_scan_enrichment_status.assert_called_once_with(_SCAN_ID, "ENRICHING") - thread = _FakeThread.last_instance - assert thread.started is True - assert thread.daemon is True - assert thread.target is scans_route._run_enrichment_in_background - assert thread.args[0] == _SCAN_ID - assert thread.args[1] == findings +def test_enrich_reuses_existing_durable_job(client, auth_headers): + scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} + db = _mock_db(current_scan=scan, findings=[{"id": 1}]) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "running"}, False) + with patch.object(scans_route, "_get_db", return_value=db): + resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) + assert resp.status_code == 202 + assert resp.get_json()["status"] == "running" def test_enrich_already_completed_returns_200(client, auth_headers): @@ -65,82 +47,14 @@ def test_enrich_already_completed_returns_200(client, auth_headers): assert "already enriched" in resp.get_json()["message"] -def test_enrich_already_in_progress_returns_202(client, auth_headers): - scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "ENRICHING"} - db = _mock_db(current_scan=scan) - with patch.object(scans_route, "_get_db", return_value=db): - resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) - assert resp.status_code == 202 - assert "in progress" in resp.get_json()["message"] - - def test_enrich_missing_scan_returns_404(client, auth_headers): - db = _mock_db(current_scan=None) - with patch.object(scans_route, "_get_db", return_value=db): + with patch.object(scans_route, "_get_db", return_value=_mock_db(current_scan=None)): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 404 def test_enrich_no_findings_returns_404(client, auth_headers): scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} - db = _mock_db(current_scan=scan, findings=[]) - with patch.object(scans_route, "_get_db", return_value=db): + with patch.object(scans_route, "_get_db", return_value=_mock_db(current_scan=scan, findings=[])): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 404 - - -def test_background_enrichment_marks_completed_on_success(): - fake_db = MagicMock() - with ( - patch.object(scans_route, "DatabaseManager", return_value=fake_db), - patch.object(scans_route, "enrich_findings", return_value=[{"id": 1}]), - ): - scans_route._run_enrichment_in_background("scan-1", [{"id": 1}], "postgresql://mock/mock") - - fake_db.update_cve_fields.assert_called_once_with([{"id": 1}]) - fake_db.update_scan_enrichment_status.assert_called_once_with("scan-1", "COMPLETED") - fake_db.close.assert_called_once() - - -def test_background_enrichment_marks_failed_on_error(): - fake_db = MagicMock() - with ( - patch.object(scans_route, "DatabaseManager", return_value=fake_db), - patch.object(scans_route, "enrich_findings", side_effect=RuntimeError("boom")), - ): - scans_route._run_enrichment_in_background("scan-1", [{"id": 1}], "postgresql://mock/mock") - - fake_db.update_scan_enrichment_status.assert_called_once_with("scan-1", "FAILED") - fake_db.close.assert_called_once() - - -def test_background_enrichment_rolls_back_before_marking_failed(): - """A write failure (e.g. in update_cve_fields) can leave db.conn in an - aborted-transaction state. Without a rollback first, the follow-up - update_scan_enrichment_status(FAILED) call would itself raise - InFailedSqlTransaction and get silently swallowed, leaving the scan - stuck at ENRICHING forever.""" - fake_db = MagicMock() - fake_db.conn = MagicMock() # an open connection, left mid-transaction - with ( - patch.object(scans_route, "DatabaseManager", return_value=fake_db), - patch.object(scans_route, "enrich_findings", side_effect=RuntimeError("boom")), - ): - scans_route._run_enrichment_in_background("scan-1", [{"id": 1}], "postgresql://mock/mock") - - fake_db.conn.rollback.assert_called_once() - fake_db.update_scan_enrichment_status.assert_called_once_with("scan-1", "FAILED") - - -def test_background_enrichment_skips_rollback_when_never_connected(): - """If enrich_findings() fails before any DB write, db.conn is still None - — rollback() must not be called on it.""" - fake_db = MagicMock() - fake_db.conn = None - with ( - patch.object(scans_route, "DatabaseManager", return_value=fake_db), - patch.object(scans_route, "enrich_findings", side_effect=RuntimeError("boom")), - ): - scans_route._run_enrichment_in_background("scan-1", [{"id": 1}], "postgresql://mock/mock") - - fake_db.update_scan_enrichment_status.assert_called_once_with("scan-1", "FAILED") diff --git a/tests/test_worker.py b/tests/test_worker.py index 73b051df..f973b15f 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -71,6 +71,7 @@ def test_worker_processes_pending_scan_successfully( # We need to stop the infinite loop. We'll raise StopWorker on the second call to recover_stale_scans. mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.side_effect = [ {"scan_id": self.scan_id, "subscription_id": self.subscription_id, "fencing_token": 1}, None, @@ -111,6 +112,7 @@ def test_worker_handles_scan_failure_gracefully( mock_db = mock_db_class.return_value mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.side_effect = [ {"scan_id": self.scan_id, "subscription_id": self.subscription_id, "fencing_token": 1}, None, @@ -147,6 +149,7 @@ def test_worker_does_not_persist_after_heartbeat_reports_lost_lease( mock_env.return_value = self.mock_db_url mock_db = mock_db_class.return_value mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.return_value = { "scan_id": self.scan_id, "subscription_id": self.subscription_id, @@ -199,6 +202,7 @@ def test_worker_sleeps_when_no_scans_pending(self, mock_sleep, mock_env, mock_db mock_db = mock_db_class.return_value mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.return_value = None with self.assertRaises(StopWorker): From 648c920bb069b156b4dd733c8ec50ccaa5563c2d Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:41:13 +0100 Subject: [PATCH 06/11] feat(core): expose durable worker metrics Signed-off-by: Shaurya K Sharma --- ...d4a8c1e6b2f9_operational_worker_metrics.py | 38 +++++++++ api/models/finding.py | 84 +++++++++++++++++++ api/observability.py | 47 +++++++++++ scanner/worker.py | 4 + tests/test_observability.py | 1 + tests/test_operational_metrics_postgres.py | 39 +++++++++ 6 files changed, 213 insertions(+) create mode 100644 alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py create mode 100644 tests/test_operational_metrics_postgres.py diff --git a/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py new file mode 100644 index 00000000..66c2262b --- /dev/null +++ b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py @@ -0,0 +1,38 @@ +"""Persist worker liveness used by bounded operational metrics. + +Revision ID: d4a8c1e6b2f9 +Revises: c9e1a5b7d3f2 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "d4a8c1e6b2f9" +down_revision: Union[str, Sequence[str], None] = "c9e1a5b7d3f2" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Store one liveness timestamp per worker process.""" + op.create_table( + "worker_heartbeats", + sa.Column("worker_id", sa.Text(), nullable=False), + sa.Column("worker_type", sa.Text(), nullable=False), + sa.Column("last_seen_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("worker_id", "worker_type", name="worker_heartbeats_pkey"), + sa.CheckConstraint("worker_type IN ('scan', 'enrichment')", name="ck_worker_heartbeats_type"), + ) + op.create_index( + "idx_worker_heartbeats_type_seen", "worker_heartbeats", ["worker_type", "last_seen_at"], unique=False + ) + + +def downgrade() -> None: + """Remove durable worker heartbeat state.""" + op.drop_index("idx_worker_heartbeats_type_seen", table_name="worker_heartbeats") + op.drop_table("worker_heartbeats") diff --git a/api/models/finding.py b/api/models/finding.py index fe3eaf94..afb73f52 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -1073,6 +1073,90 @@ def get_scans(self) -> List[Dict[str, Any]]: cur.execute("SELECT * FROM scans ORDER BY started_at DESC LIMIT 100") return [dict(row) for row in cur.fetchall()] + def record_worker_heartbeat(self, worker_id: str, worker_type: str) -> None: + """Persist liveness without using worker IDs as metric labels.""" + if worker_type not in {"scan", "enrichment"}: + raise ValueError("unsupported worker type") + conn = self._get_conn() + try: + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO worker_heartbeats (worker_id, worker_type, last_seen_at) + VALUES (%s, %s, CURRENT_TIMESTAMP) + ON CONFLICT (worker_id, worker_type) DO UPDATE SET + last_seen_at = EXCLUDED.last_seen_at + """, + (worker_id, worker_type), + ) + conn.commit() + except Exception: + self.rollback(conn) + raise + + def get_operational_metrics(self) -> Dict[str, Any]: + """Return bounded aggregates used by the public Prometheus endpoint.""" + conn = self._get_conn() + try: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(started_at)), 0) AS value + FROM scans WHERE status = 'pending' + """ + ) + scan_queue_age = cur.fetchone()["value"] + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at)), 0) AS value + FROM enrichment_jobs WHERE status = 'pending' + """ + ) + enrichment_queue_age = cur.fetchone()["value"] + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(claimed_at)), 0) AS value + FROM scans WHERE status = 'running' + """ + ) + scan_lease_age = cur.fetchone()["value"] + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(last_heartbeat_at)), 0) AS value + FROM enrichment_jobs WHERE status = 'running' + """ + ) + enrichment_lease_age = cur.fetchone()["value"] + cur.execute("SELECT COALESCE(SUM(GREATEST(attempt_count - 1, 0)), 0) AS value FROM scans") + scan_retries = cur.fetchone()["value"] + cur.execute("SELECT COALESCE(SUM(GREATEST(attempt_count - 1, 0)), 0) AS value FROM enrichment_jobs") + enrichment_retries = cur.fetchone()["value"] + cur.execute( + """ + SELECT worker_type, EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MAX(last_seen_at)) AS age + FROM worker_heartbeats GROUP BY worker_type + """ + ) + heartbeat_age = {row["worker_type"]: float(row["age"]) for row in cur.fetchall()} + cur.execute( + """ + SELECT COALESCE(EXTRACT(EPOCH FROM MAX(completed_at)), 0) AS value + FROM scans WHERE status = 'completed' + """ + ) + last_success = cur.fetchone()["value"] + conn.commit() + return { + "oldest_queue_age": {"scan": float(scan_queue_age), "enrichment": float(enrichment_queue_age)}, + "oldest_lease_age": {"scan": float(scan_lease_age), "enrichment": float(enrichment_lease_age)}, + "retry_attempts": {"scan": float(scan_retries), "enrichment": float(enrichment_retries)}, + "worker_heartbeat_age": heartbeat_age, + "last_successful_scan_timestamp": float(last_success), + } + except Exception: + self.rollback(conn) + raise + # ------------------------------------------------------------------ # # Scoring # # ------------------------------------------------------------------ # diff --git a/api/observability.py b/api/observability.py index 4b7bcfa6..a5aae461 100644 --- a/api/observability.py +++ b/api/observability.py @@ -68,6 +68,30 @@ "openshield_pending_scans", "Number of scans currently waiting in the pending queue.", ) +WORKER_HEARTBEAT_AGE_SECONDS = Gauge( + "openshield_worker_heartbeat_age_seconds", + "Seconds since the most recent durable worker heartbeat.", + ["worker_type"], +) +OLDEST_QUEUE_AGE_SECONDS = Gauge( + "openshield_oldest_queue_age_seconds", + "Age in seconds of the oldest pending durable work item.", + ["queue"], +) +OLDEST_ACTIVE_LEASE_AGE_SECONDS = Gauge( + "openshield_oldest_active_lease_age_seconds", + "Age in seconds of the oldest active durable lease.", + ["queue"], +) +RETRY_ATTEMPTS = Gauge( + "openshield_retry_attempts", + "Aggregate retry attempts currently recorded for durable work.", + ["queue"], +) +LAST_SUCCESSFUL_SCAN_TIMESTAMP = Gauge( + "openshield_last_successful_scan_timestamp_seconds", + "Unix timestamp of the most recently completed scan, or zero when none exist.", +) RULE_ERRORS_TOTAL = Counter( "openshield_rule_errors_total", "Total number of times a scanner rule raised an exception.", @@ -180,4 +204,27 @@ def _record_observability(response: Response) -> Response: @app.get("/metrics") def metrics() -> Response: + # Metrics are derived from durable state so the API can expose worker + # liveness even when scan workers run in separate processes. + db_url = os.environ.get("DATABASE_URL") + if db_url: + try: + from api.models.finding import DatabaseManager + + db = DatabaseManager(db_url) + try: + snapshot = db.get_operational_metrics() + finally: + db.close() + for queue in ("scan", "enrichment"): + OLDEST_QUEUE_AGE_SECONDS.labels(queue=queue).set(snapshot["oldest_queue_age"].get(queue, 0)) + OLDEST_ACTIVE_LEASE_AGE_SECONDS.labels(queue=queue).set(snapshot["oldest_lease_age"].get(queue, 0)) + RETRY_ATTEMPTS.labels(queue=queue).set(snapshot["retry_attempts"].get(queue, 0)) + for worker_type in ("scan", "enrichment"): + WORKER_HEARTBEAT_AGE_SECONDS.labels(worker_type=worker_type).set( + snapshot["worker_heartbeat_age"].get(worker_type, 0) + ) + LAST_SUCCESSFUL_SCAN_TIMESTAMP.set(snapshot["last_successful_scan_timestamp"]) + except Exception as exc: + logger.warning("Unable to refresh durable operational metrics: %s", exc) return Response(generate_latest(), content_type=CONTENT_TYPE_LATEST) diff --git a/scanner/worker.py b/scanner/worker.py index 0a3404cf..b5021e11 100644 --- a/scanner/worker.py +++ b/scanner/worker.py @@ -143,6 +143,10 @@ def run_worker(): while True: try: + db.record_worker_heartbeat(worker_id, "scan") + # This process executes both durable queue types; record both + # liveness signals without exporting the worker UUID as a label. + db.record_worker_heartbeat(worker_id, "enrichment") # 1. Cleanup stale scans from previous crashes db.recover_stale_scans() db.recover_stale_enrichment_jobs() diff --git a/tests/test_observability.py b/tests/test_observability.py index ef9207ba..2745e940 100644 --- a/tests/test_observability.py +++ b/tests/test_observability.py @@ -188,6 +188,7 @@ class _Stop(BaseException): mock_db = MagicMock() mock_db.recover_stale_scans.side_effect = [None, _Stop()] + mock_db.claim_next_enrichment_job.return_value = None mock_db.claim_next_pending_scan.side_effect = [ {"scan_id": scan_id, "subscription_id": sub_id, "fencing_token": 1}, None, diff --git a/tests/test_operational_metrics_postgres.py b/tests/test_operational_metrics_postgres.py new file mode 100644 index 00000000..f4f63c7c --- /dev/null +++ b/tests/test_operational_metrics_postgres.py @@ -0,0 +1,39 @@ +"""PostgreSQL coverage for durable operational metric aggregates.""" + +import os +import uuid + +import psycopg2 +import pytest + +from api.models.finding import DatabaseManager + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + + +def test_durable_operational_metrics_cover_queue_lease_retries_and_heartbeat(): + dsn = os.environ["DATABASE_URL"] + scan_id, subscription_id = str(uuid.uuid4()), str(uuid.uuid4()) + db = DatabaseManager(dsn) + try: + db.create_pending_scan(scan_id, subscription_id) + db.record_worker_heartbeat("worker-test", "scan") + db.record_worker_heartbeat("worker-test", "enrichment") + claim = db.claim_next_pending_scan("worker-a", 120) + assert claim is not None + snapshot = db.get_operational_metrics() + assert snapshot["oldest_lease_age"]["scan"] >= 0 + assert snapshot["retry_attempts"]["scan"] == 0 + assert snapshot["worker_heartbeat_age"]["scan"] >= 0 + assert snapshot["worker_heartbeat_age"]["enrichment"] >= 0 + assert snapshot["last_successful_scan_timestamp"] >= 0 + finally: + db.close() + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("DELETE FROM worker_heartbeats WHERE worker_id = 'worker-test'") + cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) + cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) From e1a1e4bbda95229d669a9b59548ba6e906bbf258 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 20:01:53 +0100 Subject: [PATCH 07/11] fix(core): complete scan durability integration Signed-off-by: Shaurya K Sharma --- api/models/finding.py | 12 ++++++++++++ docs/async-scan-architecture.md | 25 ++++++++++++++++++++++-- docs/cve_correlation_feature.md | 7 ++++--- tests/test_enrichment_jobs_postgres.py | 3 ++- tests/test_scan_leases_postgres.py | 1 + tests/test_subscription_authorization.py | 7 ++++++- 6 files changed, 48 insertions(+), 7 deletions(-) diff --git a/api/models/finding.py b/api/models/finding.py index afb73f52..02cb6cbc 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -416,6 +416,18 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token ) else: cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_result["scan_id"],)) + + # Completion creates one durable job in this same fenced + # transaction. The scan_id uniqueness constraint makes result + # replay harmless and prevents duplicate enrichment delivery. + cur.execute( + """ + INSERT INTO enrichment_jobs (job_id, scan_id, status, attempt_count, checkpoint) + VALUES (%s, %s, 'pending', 0, 0) + ON CONFLICT (scan_id) DO NOTHING + """, + (str(uuid.uuid4()), scan_result["scan_id"]), + ) conn.commit() except Exception: # psycopg2 connections remain in an aborted transaction after any diff --git a/docs/async-scan-architecture.md b/docs/async-scan-architecture.md index e9cf84f2..757aec61 100644 --- a/docs/async-scan-architecture.md +++ b/docs/async-scan-architecture.md @@ -13,7 +13,7 @@ In the legacy synchronous model, POST /api/scans/trigger would block the HTTP re OpenShield now employs a decoupled, database backed worker architecture. This is the industry standard for long running security tasks where reliability and state persistence are critical. ### 1. The API (Flask) -When a scan is triggered, the API performs minimal work. It validates the subscription_id, creates a record in the scans table with status set to pending, and returns 202 Accepted and the scan_id immediately. +When a scan is triggered, the API validates the subscription and creates a durable pending record. PostgreSQL permits at most one `pending` or `running` scan per subscription. `Idempotency-Key` replays return the same logical scan when the request fingerprint matches; reuse with different semantics returns a conflict. The optional `OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR` policy enables an explicit time-window quota. A zero/unset value preserves the current no-business-limit policy while the one-active-scan concurrency quota remains enforced. ### 2. The Queue (PostgreSQL) The scans table acts as a persistent task queue. This avoids the need for additional infrastructure like Redis or RabbitMQ while providing ACID compliance, visibility, and auditability. Scan states are never lost during crashes, status polling is a simple SQL query, and every scan has a persistent record of its error state. @@ -30,8 +30,22 @@ worker cannot persist completion, failure, or findings after it loses ownership. Once a scan reaches the maximum attempt count, it is marked `failed` so bad credentials or persistent Azure errors cannot retry forever. +Findings use a stable database-enforced identity (`scan`, rule, canonical +resource scope, and an optional rule-specific discriminator). Result retries +use PostgreSQL upserts, so mutable text or severity is updated rather than +creating a second authoritative finding. Rule evaluations use the same +database-first uniqueness model. + ### 3. The Worker (Python) -The scanner/worker.py process runs independently of the web server. Its lifecycle involves several steps. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings and marks the scan complete atomically. On failure, it records a sanitized error only while it still owns the lease. +The scanner/worker.py process runs independently of the web server. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings, evaluations, and one durable enrichment job atomically. On failure, it records a sanitized error only while it still owns the lease. + +### Durable CVE enrichment + +`POST /api/scans//enrich` enqueues (or returns) the one durable PostgreSQL enrichment job for the scan; it never starts a request-owned daemon thread. The scan worker claims those jobs with the same owner/expiry/fencing model, checkpoints after each finding, and retries transient failures with bounded exponential backoff. An expired job is recovered or terminally failed after its attempt limit. NVD retrieval follows every `totalResults` page; replaying a checkpoint updates the existing finding instead of duplicating CVE data. + +### Operational signals + +`/metrics` derives bounded-cardinality operational gauges from PostgreSQL: worker heartbeat age, oldest queue age, oldest active lease age, aggregate retry attempts, and the last successful scan timestamp. Labels are limited to `queue` (`scan` or `enrichment`) and `worker_type`; scan, job, subscription, and worker identifiers are never metric labels. ## Technical Rationale @@ -41,6 +55,13 @@ While Celery is powerful, it introduces external dependencies and operational co ### Why not Threading Python background threads are ephemeral. If the web server process restarts, all in flight scans are killed instantly and marked as running forever in the DB. A separate worker process ensures that the scan lifecycle is independent of the web server lifecycle. +## Deployment order + +This release is not safe for mixed old and new scan workers. Stop or drain old +workers, apply Alembic migrations, then start the fenced worker version. Legacy +workers do not carry ownership/fencing state; legacy running scans are retained +and made recoverable by the lease migration rather than deleted. + ## Testing Suite The asynchronous transition is verified through a multi layered testing strategy. diff --git a/docs/cve_correlation_feature.md b/docs/cve_correlation_feature.md index a0820815..eac8ddec 100644 --- a/docs/cve_correlation_feature.md +++ b/docs/cve_correlation_feature.md @@ -19,7 +19,8 @@ The CVE Correlation feature integrates the MITRE National Vulnerability Database | File | Change | Why | |---|---|---| | scanner/engine.py | Decoupled Scan. Removed synchronous enrichment from the scan lifecycle. | Performance: Azure scans now return immediately without waiting for NVD rate limits (7s per resource type). | -| api/routes/scans.py | New Endpoint. Added `POST /api/scans//enrich`. | Flexibility: CVE enrichment can now be triggered on-demand or by a background job after the scan completes. | +| api/routes/scans.py | Durable endpoint. `POST /api/scans//enrich` enqueues or returns a PostgreSQL job. | No Gunicorn daemon thread owns authoritative work. | +| scanner/enrichment_worker.py | Checkpointed durable enrichment execution. | Reclaims expired work safely and resumes at a finding checkpoint. | | api/models/finding.py | Updated Scan model and added enrichment status tracking. | Persistence: Adds `cve_enrichment_status` to track `PENDING`, `COMPLETED`, or `FAILED` states. | | alembic/versions/ | Defines CVE columns in the versioned database schema. | Deployment: Alembic owns schema changes independently of Flask application startup. | | api/routes/score.py | Added GET /api/score/cve-summary endpoint. | Dashboard UI: Provides the frontend with high-level data like Total Known Exploits and enrichment status. | @@ -30,10 +31,10 @@ The CVE Correlation feature integrates the MITRE National Vulnerability Database To ensure the frontend dashboard works perfectly, the architecture uses a Decoupled Enrichment model: 1. Fast Dashboard Loads: The scan engine completes rapidly. The dashboard can check the enrichment status of the latest scan. -2. Manual/Job Enrichment: A "Trigger Enrichment" button or a background task calls `POST /api/scans//enrich` to populate CVE data. +2. Durable Enrichment: Completion creates one job; `POST /api/scans//enrich` returns that job idempotently. The worker claims, retries, and checkpoints it in PostgreSQL. 3. Dashboard-Ready Summary Endpoint: The /api/score/cve-summary endpoint includes the `status` field, allowing the UI to show a "Scan Enriched" badge or a "Pending" spinner. 4. Actionable Risk (CISA KEV): The exploit_available flag uses the CISA Known Exploited Vulnerabilities catalogue, allowing the dashboard to highlight high-priority risks that are being exploited in the wild. -5. Persistent Historical State: Enrichment happens at the time of the enrichment call, and the result is persisted. +5. Persistent Historical State: Enrichment checkpoint state and results survive API/worker restarts. NVD pagination continues through every `totalResults` page. ## Security and Compliance Audit diff --git a/tests/test_enrichment_jobs_postgres.py b/tests/test_enrichment_jobs_postgres.py index ff659efe..b98b73fe 100644 --- a/tests/test_enrichment_jobs_postgres.py +++ b/tests/test_enrichment_jobs_postgres.py @@ -50,7 +50,8 @@ def enrichment_scan(): ], } db.save_scan(result, "seed", claim["fencing_token"]) - job, _ = db.enqueue_enrichment_job(scan_id) + job, created = db.enqueue_enrichment_job(scan_id) + assert created is False yield dsn, scan_id, job finally: db.close() diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py index 7b2fa7df..fd5df2f2 100644 --- a/tests/test_scan_leases_postgres.py +++ b/tests/test_scan_leases_postgres.py @@ -105,6 +105,7 @@ def cleanup(self) -> None: with psycopg2.connect(self.dsn) as conn: with conn.cursor() as cur: for scan_id in self.scan_ids: + cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) diff --git a/tests/test_subscription_authorization.py b/tests/test_subscription_authorization.py index d79f21a2..640d9686 100644 --- a/tests/test_subscription_authorization.py +++ b/tests/test_subscription_authorization.py @@ -14,7 +14,12 @@ def _trigger(client, auth_headers, subscription_id, monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://ci:ci@localhost/ci_db") - with patch("api.routes.scans.DatabaseManager", return_value=MagicMock()): + database = MagicMock() + database.admit_scan.return_value = ( + {"scan_id": "test-scan", "status": "pending"}, + True, + ) + with patch("api.routes.scans.DatabaseManager", return_value=database): return client.post( "/api/scans/trigger", json={"subscription_id": subscription_id}, From f1b1724da789b89541ad80a402996c9c5742b631 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 20:09:16 +0100 Subject: [PATCH 08/11] fix(api): avoid exposing scan admission errors Signed-off-by: Shaurya K Sharma --- api/routes/scans.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/api/routes/scans.py b/api/routes/scans.py index ec90033a..a84d233a 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -136,10 +136,10 @@ def trigger_scan(): request_fingerprint=request_fingerprint, max_scans_per_hour=_configured_hourly_quota(), ) - except ScanAdmissionConflict as exc: - return jsonify({"error": str(exc)}), 409 - except ScanQuotaExceeded as exc: - return jsonify({"error": str(exc)}), 429 + except ScanAdmissionConflict: + return jsonify({"error": "Idempotency-Key is already associated with a different request."}), 409 + except ScanQuotaExceeded: + return jsonify({"error": "Scan quota exceeded for this subscription."}), 429 except Exception as exc: logger.error("Failed to create pending scan: %s", exc, exc_info=True) return jsonify({"error": "Database error"}), 500 From bc76dfb8c3be6264cbe5231f65869e9a8767b41c Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 20:32:32 +0100 Subject: [PATCH 09/11] fix(scanner): use safe NVD request transport Signed-off-by: Shaurya K Sharma --- scanner/nvd_client.py | 50 ++++++++++------------- tests/test_nvd_client.py | 87 ++++++++++++++++++++-------------------- 2 files changed, 65 insertions(+), 72 deletions(-) diff --git a/scanner/nvd_client.py b/scanner/nvd_client.py index 8ca65666..8cefbcc1 100644 --- a/scanner/nvd_client.py +++ b/scanner/nvd_client.py @@ -19,12 +19,10 @@ import threading import time import logging -import urllib.request -import urllib.error -import urllib.parse -import json from typing import Optional +import requests + from api.observability import NVD_REQUEST_LATENCY_SECONDS logger = logging.getLogger(__name__) @@ -134,40 +132,34 @@ def _parse_cve_item(item: dict) -> Optional[dict]: def _fetch_nvd_page(keyword: str, start_index: int, results_per_page: int) -> dict: """Fetch one NVD page, raising only after bounded retries are exhausted.""" - params = urllib.parse.urlencode( - { - "keywordSearch": keyword, - "startIndex": start_index, - "resultsPerPage": results_per_page, - } - ) - url = f"{_NVD_BASE_URL}?{params}" - parsed_url = urllib.parse.urlsplit(url) - if ( - parsed_url.scheme != "https" - or parsed_url.hostname != "services.nvd.nist.gov" - or parsed_url.port not in (None, 443) - ): - raise NvdRequestError("Refusing request to an untrusted NVD endpoint") + params = { + "keywordSearch": keyword, + "startIndex": start_index, + "resultsPerPage": results_per_page, + } last_error: Optional[Exception] = None for attempt in range(1, _MAX_RETRIES + 1): try: _wait_for_rate_limit() - req = urllib.request.Request( - url, - headers={"User-Agent": "OpenShield/0.1 (github.com/openshield-org/openshield)"}, - ) with NVD_REQUEST_LATENCY_SECONDS.time(): - with urllib.request.urlopen( # nosec B310 # noqa: E501 - req, timeout=10 - ) as resp: - return json.loads(resp.read()) - except urllib.error.HTTPError as exc: + response = requests.get( + _NVD_BASE_URL, + params=params, + headers={"User-Agent": "OpenShield/0.1 (github.com/openshield-org/openshield)"}, + timeout=10, + ) + response.raise_for_status() + return response.json() + except requests.HTTPError as exc: last_error = exc - if exc.code != 429: + if exc.response is None or exc.response.status_code != 429: break time.sleep(30 * attempt) + except requests.RequestException as exc: + last_error = exc + if attempt < _MAX_RETRIES: + time.sleep(2**attempt) except Exception as exc: last_error = exc if attempt < _MAX_RETRIES: diff --git a/tests/test_nvd_client.py b/tests/test_nvd_client.py index 60100668..06818b63 100644 --- a/tests/test_nvd_client.py +++ b/tests/test_nvd_client.py @@ -12,13 +12,13 @@ TestQueryNvd - query_nvd() HTTP behaviour (mocked urlopen) """ -import json import threading import time import unittest -import urllib.error from unittest.mock import patch, MagicMock +import requests + # Clear the module cache before import so previous test runs don't bleed in from scanner.nvd_client import query_nvd, _parse_cve_item, _cache, _wait_for_rate_limit @@ -68,21 +68,17 @@ _EMPTY_NVD_RESPONSE = {"vulnerabilities": []} -def _make_mock_urlopen_response(data: dict) -> MagicMock: +def _make_mock_requests_response(data: dict, status: int = 200) -> MagicMock: """ - Return a MagicMock that behaves like urllib.request.urlopen()'s - context manager return value. - - urlopen() is used as: - with urllib.request.urlopen(req, timeout=10) as resp: - data = json.loads(resp.read()) + Return a MagicMock that behaves like a requests NVD response. - So the mock needs __enter__/__exit__ and a .read() method. + `raise_for_status()` raises the same requests exception the client handles. """ mock_resp = MagicMock() - mock_resp.read.return_value = json.dumps(data).encode("utf-8") - mock_resp.__enter__ = lambda s: s - mock_resp.__exit__ = MagicMock(return_value=False) + mock_resp.json.return_value = data + mock_resp.status_code = status + if status >= 400: + mock_resp.raise_for_status.side_effect = requests.HTTPError(response=mock_resp) return mock_resp @@ -175,14 +171,14 @@ def test_falls_back_to_cvss_v2_when_v31_absent(self): # --------------------------------------------------------------------------- # TestQueryNvd -# Tests for query_nvd() - mocks urllib.request.urlopen to prevent live calls. +# Tests for query_nvd() - mocks requests.get to prevent live calls. # Also mocks _wait_for_rate_limit to keep tests fast. # --------------------------------------------------------------------------- class TestQueryNvd(unittest.TestCase): """ - query_nvd() builds a URL, calls urlopen, parses the response, caches it, + query_nvd() uses the fixed NVD HTTPS URL, parses the response, caches it, and handles errors gracefully. All HTTP is mocked. """ @@ -190,77 +186,82 @@ def setUp(self): """Clear the module-level cache before each test.""" _cache.clear() - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_returns_parsed_cves_on_success(self, mock_wait, mock_urlopen): + def test_returns_parsed_cves_on_success(self, mock_wait, mock_get): """Successful response is parsed into a list of CVE dicts.""" - mock_urlopen.return_value = _make_mock_urlopen_response(_SAMPLE_NVD_RESPONSE) + mock_get.return_value = _make_mock_requests_response(_SAMPLE_NVD_RESPONSE) results = query_nvd("Azure Storage Account") self.assertEqual(len(results), 2) self.assertEqual(results[0]["cve_id"], "CVE-2023-12345") self.assertEqual(results[1]["cve_id"], "CVE-2022-99999") + mock_get.assert_called_once_with( + "https://services.nvd.nist.gov/rest/json/cves/2.0", + params={"keywordSearch": "Azure Storage Account", "startIndex": 0, "resultsPerPage": 2000}, + headers={"User-Agent": "OpenShield/0.1 (github.com/openshield-org/openshield)"}, + timeout=10, + ) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_returns_empty_list_on_empty_nvd_response(self, mock_wait, mock_urlopen): + def test_returns_empty_list_on_empty_nvd_response(self, mock_wait, mock_get): """An empty vulnerabilities list returns [] without error.""" - mock_urlopen.return_value = _make_mock_urlopen_response(_EMPTY_NVD_RESPONSE) + mock_get.return_value = _make_mock_requests_response(_EMPTY_NVD_RESPONSE) results = query_nvd("nonexistent-resource-xyz") self.assertEqual(results, []) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_second_call_uses_cache(self, mock_wait, mock_urlopen): + def test_second_call_uses_cache(self, mock_wait, mock_get): """ - Calling query_nvd twice with the same keyword only hits urlopen once. + Calling query_nvd twice with the same keyword only makes one request. The second call must return from cache without a network request. """ - mock_urlopen.return_value = _make_mock_urlopen_response(_SAMPLE_NVD_RESPONSE) + mock_get.return_value = _make_mock_requests_response(_SAMPLE_NVD_RESPONSE) query_nvd("Azure Storage Account") query_nvd("Azure Storage Account") # Should be served from cache - self.assertEqual(mock_urlopen.call_count, 1) + self.assertEqual(mock_get.call_count, 1) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_fetches_every_nvd_page(self, mock_wait, mock_urlopen): + def test_fetches_every_nvd_page(self, mock_wait, mock_get): first_page = {"totalResults": 2, "vulnerabilities": [_SAMPLE_NVD_RESPONSE["vulnerabilities"][0]]} second_page = {"totalResults": 2, "vulnerabilities": [_SAMPLE_NVD_RESPONSE["vulnerabilities"][1]]} - mock_urlopen.side_effect = [_make_mock_urlopen_response(first_page), _make_mock_urlopen_response(second_page)] + mock_get.side_effect = [ + _make_mock_requests_response(first_page), + _make_mock_requests_response(second_page), + ] results = query_nvd("Azure Storage Account", results_per_page=1) self.assertEqual([item["cve_id"] for item in results], ["CVE-2023-12345", "CVE-2022-99999"]) - self.assertEqual(mock_urlopen.call_count, 2) + self.assertEqual(mock_get.call_count, 2) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_returns_empty_list_on_network_error(self, mock_wait, mock_urlopen): + def test_returns_empty_list_on_network_error(self, mock_wait, mock_get): """A network exception returns [] and does not propagate the error.""" - mock_urlopen.side_effect = Exception("Connection refused") + mock_get.side_effect = requests.ConnectionError("Connection refused") results = query_nvd("Azure Storage Account") self.assertEqual(results, []) - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_returns_empty_list_on_http_503(self, mock_wait, mock_urlopen): + def test_returns_empty_list_on_http_503(self, mock_wait, mock_get): """An HTTP 503 returns [] and does not propagate the error.""" - mock_urlopen.side_effect = urllib.error.HTTPError( - url=None, code=503, msg="Service Unavailable", hdrs=None, fp=None - ) + mock_get.return_value = _make_mock_requests_response({}, status=503) results = query_nvd("Azure Storage Account") self.assertEqual(results, []) @patch("scanner.nvd_client.time.sleep") - @patch("scanner.nvd_client.urllib.request.urlopen") + @patch("scanner.nvd_client.requests.get") @patch("scanner.nvd_client._wait_for_rate_limit") - def test_backs_off_and_retries_on_429(self, mock_wait, mock_urlopen, mock_sleep): + def test_backs_off_and_retries_on_429(self, mock_wait, mock_get, mock_sleep): """ A 429 response triggers a sleep and retry. After MAX_RETRIES 429s, returns [] gracefully. """ - mock_urlopen.side_effect = urllib.error.HTTPError( - url=None, code=429, msg="Too Many Requests", hdrs=None, fp=None - ) + mock_get.return_value = _make_mock_requests_response({}, status=429) results = query_nvd("Azure Storage Account") self.assertEqual(results, []) # time.sleep should have been called (back-off logic) From 8f27e8f11b6b02d1257c4ce4823f4885cbaf51d2 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Wed, 2 Sep 2026 01:08:47 +0100 Subject: [PATCH 10/11] fix(core): address review of scan durability hardening (#303) Resolves the two must-fix items and the follow-ups raised in review of #325. Integration: leave the #263 evaluation contract to #321 ------------------------------------------------------ This branch created a second `rule_evaluations` table, near-identical to the one PR #321 adds, that no production code ever populated: `run_scan()` emits no `evaluations` key, so the storage, the upsert and its tests described a contract that could not be observed in production. Two `CREATE TABLE rule_evaluations` statements would also have broken whichever of #321/#325 merged second, independently of the Alembic head ordering. Issue #263 owns that contract - the table, the PASS/FAIL/UNKNOWN/ERROR/ NOT_APPLICABLE semantics, engine emission and the compliance-score fix - so this branch drops it entirely and keeps only what #303 asks for: the stable `finding_key` identity and its unique index. `save_scan()` marks where #321's evaluation writes belong, inside the fenced completion transaction, so they inherit the lease/ownership check without re-implementing it. Terminally failed enrichment jobs are recoverable again ------------------------------------------------------- After three failed attempts a job became `failed` and nothing could move it: enqueue used ON CONFLICT DO NOTHING, claim only selected `pending`, and stale recovery only handled expired `running` leases. That regressed the operator retry the old thread-based path gave for free. `enqueue_enrichment_job()` now returns an explicit outcome - `created`, `requeued`, `active` or `completed` - and atomically resets a `failed` job to `pending` with a fresh retry budget. It keeps the same job row, its last error message (audit) and its checkpoint (so the retry resumes), never revives a `completed` job, and never disturbs a live `running` lease. The conditional UPDATE is the whole guard, so concurrent re-POSTs converge on one logical job. Admission migration cannot leave an INVALID index ------------------------------------------------- `CREATE UNIQUE INDEX CONCURRENTLY uq_scans_one_active_per_subscription` fails, and leaves an unusable index behind, on a deployment that already holds several active scans for one subscription. The migration now preflights, changes nothing, and names the offending subscriptions in an actionable error; deciding which production scan is authoritative stays an operator call, and no scan history is deleted. A retry is safe: an INVALID index from an interrupted build is dropped before rebuilding. PostgreSQL tests no longer race the shared queue ------------------------------------------------ Fixtures called `claim_next_pending_scan()`, which takes the globally oldest pending scan, then persisted against the scan they had just created - so any unrelated pending row made `save_scan()` raise LostLease. Reproduced on a real database: with one older pending scan present, the old fixture claims someone else's row and fails; the new one does not. `claim_next_pending_scan()` and `claim_next_enrichment_job()` take an optional `scan_id` so a caller can claim a known row under identical lease and fencing semantics, and the fixtures use it. Queue-wide `recover_stale_*()` counts are asserted as progression of the test's own row rather than as global totals. Follow-ups ---------- - worker_heartbeats is bounded: rows past WORKER_HEARTBEAT_RETENTION_SECONDS are pruned on the beat that registers a new worker identity - once per process, not on every beat, and never for a live worker. - The scan and enrichment queues alternate one item per loop iteration instead of draining enrichment first, so neither can starve the other. - Added the partial index that keeps /metrics' last-successful-scan lookup from degrading into a sequential scan as history grows. - Documented SCAN_LEASE_SECONDS, SCAN_HEARTBEAT_SECONDS, OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR and WORKER_HEARTBEAT_RETENTION_SECONDS in .env.example. - docs/api-reference.md now documents the real admission and enrichment contracts, including which status codes are actually returned. Tests ----- New PostgreSQL coverage for terminal-failure requeue (explicit requeue, no restart of a completed job, a live lease is not stolen, concurrent requeues converge, a stale token cannot write after reclaim, a requeued job completes), the duplicate-active-scan migration preflight, stale heartbeat pruning, and worker fairness with both queues backlogged. Full suite on postgres:16-alpine: 897 passed, 2 skipped. The one failure, test_vector_store_purity, is a local checkout artifact - it needs a BM25 index this checkout's ai/vectorstore lacks, and passes in a clean worktree. Signed-off-by: Shaurya K Sharma --- .env.example | 20 ++ ...a7c5e9d2f1b4_scan_admission_idempotency.py | 65 ++++- ...d4a8c1e6b2f9_operational_worker_metrics.py | 14 + ...6d8e1a4c9_idempotent_finding_identities.py | 41 +++ ...4c9_idempotent_findings_and_evaluations.py | 66 ----- api/models/finding.py | 200 ++++++++------ api/routes/scans.py | 38 +-- docs/api-reference.md | 85 +++++- docs/async-scan-architecture.md | 43 ++- scanner/worker.py | 25 +- tests/test_enrichment_jobs_postgres.py | 215 +++++++++++++-- tests/test_operational_metrics_postgres.py | 47 +++- .../test_scan_admission_migration_postgres.py | 244 ++++++++++++++++++ tests/test_scan_admission_postgres.py | 1 - tests/test_scan_leases_postgres.py | 85 +++--- tests/test_scans_enrich.py | 38 ++- tests/test_worker.py | 62 +++++ 17 files changed, 1039 insertions(+), 250 deletions(-) create mode 100644 alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py delete mode 100644 alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py create mode 100644 tests/test_scan_admission_migration_postgres.py diff --git a/.env.example b/.env.example index 05912d7f..3ffca381 100644 --- a/.env.example +++ b/.env.example @@ -9,6 +9,26 @@ JWT_SECRET=change-me-in-production # logs a startup warning when this is unset. See docs/api-reference.md. OPENSHIELD_AUTHORIZED_SUBSCRIPTIONS= +# Optional - durable scan worker tuning. A claim is held for SCAN_LEASE_SECONDS +# and renewed every SCAN_HEARTBEAT_SECONDS; the heartbeat MUST be shorter than +# the lease or the worker loses its own claim mid-scan. A non-numeric or +# non-positive value logs a warning and uses the default; a heartbeat >= the +# lease logs a warning and uses one third of the lease. Defaults: 900 and 300. +SCAN_LEASE_SECONDS=900 +SCAN_HEARTBEAT_SECONDS=300 + +# Optional - explicit per-subscription hourly admission quota. Unset or 0 keeps +# the historical no-time-window policy (a non-numeric value logs a warning and +# disables the quota); one active (pending/running) scan per subscription is +# always enforced regardless of this value. +OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR=0 + +# Optional - how long a retired worker's heartbeat row is kept before it is +# pruned. Worker identities are per-process, so this bounds worker_heartbeats +# across restarts. Must stay far above SCAN_HEARTBEAT_SECONDS so that a live +# worker is never pruned. Default: 604800 (7 days). +WORKER_HEARTBEAT_RETENTION_SECONDS=604800 + # AI providers - add at least one ANTHROPIC_API_KEY= GROQ_API_KEY= diff --git a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py index a45eb609..e4daf0a2 100644 --- a/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py +++ b/alembic/versions/a7c5e9d2f1b4_scan_admission_idempotency.py @@ -16,22 +16,75 @@ branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None +_ACTIVE_INDEX = "uq_scans_one_active_per_subscription" +_KEY_INDEX = "uq_scans_subscription_idempotency_key" + + +def _assert_one_active_scan_per_subscription() -> None: + """Fail with an actionable error instead of an unusable index. + + ``CREATE UNIQUE INDEX CONCURRENTLY`` on a table that already violates the + constraint fails *and* leaves an INVALID index behind. A deployment that + predates the one-active-scan rule can legitimately hold several + ``pending``/``running`` rows for one subscription, so this checks first and + reports exactly which subscriptions block the upgrade. Choosing which of + those scans is authoritative is an operator decision -- deleting or + completing production scan history automatically is never this migration's + call. + """ + rows = ( + op.get_bind() + .execute( + sa.text( + """ + SELECT subscription_id, COUNT(*) AS active + FROM scans + WHERE status IN ('pending', 'running') + GROUP BY subscription_id + HAVING COUNT(*) > 1 + ORDER BY active DESC, subscription_id + """ + ) + ) + .fetchall() + ) + if not rows: + return + detail = ", ".join(f"{subscription_id} ({active} active)" for subscription_id, active in rows) + raise RuntimeError( + "Cannot enforce one active scan per subscription: " + f"{len(rows)} subscription(s) already have more than one pending/running scan: {detail}. " + "Resolve them first (let the scans finish, or mark the superseded rows " + "'failed'), then re-run this migration. " + "See docs/async-scan-architecture.md for the documented cleanup order." + ) + def upgrade() -> None: """Persist idempotency semantics and prevent more than one active scan.""" op.add_column("scans", sa.Column("idempotency_key", sa.Text(), nullable=True)) op.add_column("scans", sa.Column("request_fingerprint", sa.Text(), nullable=True)) + + # Checked before either index is built so a blocked upgrade leaves the + # schema exactly as it was, with the added columns unused and harmless. + _assert_one_active_scan_per_subscription() + with op.get_context().autocommit_block(): + # An earlier interrupted or failed CONCURRENTLY build leaves an INVALID + # index that cannot serve queries but does occupy the name. Drop both + # names first so retrying this migration is deterministic. + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_KEY_INDEX}") + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_ACTIVE_INDEX}") op.execute( - """ - CREATE UNIQUE INDEX CONCURRENTLY uq_scans_subscription_idempotency_key + f""" + CREATE UNIQUE INDEX CONCURRENTLY {_KEY_INDEX} ON scans (subscription_id, idempotency_key) WHERE idempotency_key IS NOT NULL """ ) op.execute( - """ - CREATE UNIQUE INDEX CONCURRENTLY uq_scans_one_active_per_subscription + f""" + CREATE UNIQUE INDEX CONCURRENTLY {_ACTIVE_INDEX} ON scans (subscription_id) WHERE status IN ('pending', 'running') """ @@ -41,7 +94,7 @@ def upgrade() -> None: def downgrade() -> None: """Remove scan admission metadata and constraints.""" with op.get_context().autocommit_block(): - op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_scans_one_active_per_subscription") - op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_scans_subscription_idempotency_key") + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_ACTIVE_INDEX}") + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_KEY_INDEX}") op.drop_column("scans", "request_fingerprint") op.drop_column("scans", "idempotency_key") diff --git a/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py index 66c2262b..598877dd 100644 --- a/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py +++ b/alembic/versions/d4a8c1e6b2f9_operational_worker_metrics.py @@ -31,8 +31,22 @@ def upgrade() -> None: "idx_worker_heartbeats_type_seen", "worker_heartbeats", ["worker_type", "last_seen_at"], unique=False ) + # /metrics reports the last successful scan on every scrape. Without this + # the aggregate degrades into a sequential scan of the whole scans table as + # scan history grows; the partial index keeps it an index-only lookup. + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE INDEX CONCURRENTLY idx_scans_completed_completed_at + ON scans (completed_at DESC) + WHERE status = 'completed' + """ + ) + def downgrade() -> None: """Remove durable worker heartbeat state.""" + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_scans_completed_completed_at") op.drop_index("idx_worker_heartbeats_type_seen", table_name="worker_heartbeats") op.drop_table("worker_heartbeats") diff --git a/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py b/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py new file mode 100644 index 00000000..7c067286 --- /dev/null +++ b/alembic/versions/f2b6d8e1a4c9_idempotent_finding_identities.py @@ -0,0 +1,41 @@ +"""Add database-enforced identities for scan findings. + +Revision ID: f2b6d8e1a4c9 +Revises: e4f7a9b2c6d8 +Create Date: 2026-08-29 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "f2b6d8e1a4c9" +down_revision: Union[str, Sequence[str], None] = "e4f7a9b2c6d8" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +_UNIQUE_INDEX = "uq_findings_scan_finding_key" + + +def upgrade() -> None: + """Give every finding a stable identity so replayed results upsert.""" + op.add_column("findings", sa.Column("finding_key", sa.Text(), nullable=True)) + # Existing records predate the identity contract. Preserve each record as + # distinct rather than attempting to infer equivalence from mutable text. + op.execute("UPDATE findings SET finding_key = 'legacy:' || id::text WHERE finding_key IS NULL") + op.alter_column("findings", "finding_key", nullable=False) + + with op.get_context().autocommit_block(): + # A previous interrupted CONCURRENTLY build leaves an unusable index + # behind that would make this statement fail with "already exists". + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_UNIQUE_INDEX}") + op.execute(f"CREATE UNIQUE INDEX CONCURRENTLY {_UNIQUE_INDEX} ON findings (scan_id, finding_key)") + + +def downgrade() -> None: + """Remove the finding identity introduced by this revision.""" + with op.get_context().autocommit_block(): + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_UNIQUE_INDEX}") + op.drop_column("findings", "finding_key") diff --git a/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py b/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py deleted file mode 100644 index 7b4522a9..00000000 --- a/alembic/versions/f2b6d8e1a4c9_idempotent_findings_and_evaluations.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Add database-enforced identities for scan results. - -Revision ID: f2b6d8e1a4c9 -Revises: e4f7a9b2c6d8 -Create Date: 2026-08-29 00:00:00.000000 -""" - -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -from sqlalchemy.dialects import postgresql - - -revision: str = "f2b6d8e1a4c9" -down_revision: Union[str, Sequence[str], None] = "e4f7a9b2c6d8" -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - """Add stable finding keys and per-resource evaluation rows.""" - op.add_column("findings", sa.Column("finding_key", sa.Text(), nullable=True)) - # Existing records predate the identity contract. Preserve each record as - # distinct rather than attempting to infer equivalence from mutable text. - op.execute("UPDATE findings SET finding_key = 'legacy:' || id::text WHERE finding_key IS NULL") - op.alter_column("findings", "finding_key", nullable=False) - - with op.get_context().autocommit_block(): - op.execute("CREATE UNIQUE INDEX CONCURRENTLY uq_findings_scan_finding_key ON findings (scan_id, finding_key)") - - op.create_table( - "rule_evaluations", - sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), - sa.Column("scan_id", postgresql.UUID(), nullable=False), - sa.Column("rule_id", sa.Text(), nullable=False), - sa.Column("resource_id", sa.Text(), nullable=False), - sa.Column("resource_type", sa.Text(), server_default=sa.text("''"), nullable=False), - sa.Column("status", sa.Text(), nullable=False), - sa.Column("reason_code", sa.Text(), nullable=True), - sa.Column("reason", sa.Text(), nullable=True), - sa.Column("evidence", postgresql.JSONB(), server_default=sa.text("'{}'::jsonb"), nullable=True), - sa.Column("finding_id", sa.Integer(), nullable=True), - sa.Column("evaluated_at", sa.DateTime(timezone=True), nullable=False), - sa.ForeignKeyConstraint(["scan_id"], ["scans.scan_id"], name="rule_evaluations_scan_id_fkey"), - sa.ForeignKeyConstraint( - ["finding_id"], ["findings.id"], name="rule_evaluations_finding_id_fkey", ondelete="SET NULL" - ), - sa.PrimaryKeyConstraint("id", name="rule_evaluations_pkey"), - sa.UniqueConstraint("scan_id", "rule_id", "resource_id", name="uq_rule_evaluations_scan_rule_resource"), - sa.CheckConstraint( - "status IN ('PASS', 'FAIL', 'UNKNOWN', 'ERROR', 'NOT_APPLICABLE')", - name="ck_rule_evaluations_status_v1", - ), - sa.CheckConstraint("resource_id <> ''", name="ck_rule_evaluations_resource_id_not_empty"), - ) - op.create_index("idx_rule_evaluations_scan_id", "rule_evaluations", ["scan_id"], unique=False) - - -def downgrade() -> None: - """Remove idempotent-result storage introduced by this revision.""" - op.drop_index("idx_rule_evaluations_scan_id", table_name="rule_evaluations") - op.drop_table("rule_evaluations") - with op.get_context().autocommit_block(): - op.execute("DROP INDEX CONCURRENTLY IF EXISTS uq_findings_scan_finding_key") - op.drop_column("findings", "finding_key") diff --git a/api/models/finding.py b/api/models/finding.py index 02cb6cbc..eda23b0b 100644 --- a/api/models/finding.py +++ b/api/models/finding.py @@ -25,6 +25,11 @@ logger = logging.getLogger(__name__) +# Worker identities are per-process UUIDs. Heartbeat rows older than this are +# long-dead workers that no metric reads, so they are pruned to keep +# worker_heartbeats bounded. Must stay far above any heartbeat interval. +DEFAULT_WORKER_HEARTBEAT_RETENTION_SECONDS = 7 * 24 * 60 * 60 + class LostLease(RuntimeError): """Raised when a worker no longer owns the scan it is trying to update.""" @@ -257,7 +262,6 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token finding["severity"] = normalize_severity(finding.get("severity")) finding["finding_key"] = stable_finding_key(scan_result["scan_id"], finding) findings.append(finding) - evaluations = [dict(raw_evaluation) for raw_evaluation in scan_result.get("evaluations", [])] conn = self._get_conn() completed_at = scan_result.get("completed_at") or datetime.now(timezone.utc).isoformat() @@ -301,7 +305,6 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token scan_result["scan_id"], ), ) - finding_ids: Dict[tuple[str, str], int] = {} for f in findings: cur.execute( """ @@ -327,7 +330,6 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token cvss_score = EXCLUDED.cvss_score, exploit_available = EXCLUDED.exploit_available, detected_at = EXCLUDED.detected_at - RETURNING id """, ( # The parent scan owns every child in this batch. @@ -352,10 +354,6 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token f.get("detected_at"), ), ) - finding_row = cur.fetchone() - finding_id = finding_row["id"] if isinstance(finding_row, dict) else finding_row[0] - finding_ids[(f.get("rule_id") or "", f.get("resource_id") or "")] = finding_id - finding_keys = [f["finding_key"] for f in findings] if finding_keys: cur.execute( @@ -365,57 +363,11 @@ def save_scan(self, scan_result: Dict[str, Any], lease_owner: str, fencing_token else: cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_result["scan_id"],)) - for evaluation in evaluations: - rule_id = evaluation.get("rule_id") - resource_id = evaluation.get("resource_id") - if not rule_id or not resource_id: - raise ValueError("evaluations require rule_id and resource_id") - cur.execute( - """ - INSERT INTO rule_evaluations - (scan_id, rule_id, resource_id, resource_type, status, - reason_code, reason, evidence, finding_id, evaluated_at) - VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) - ON CONFLICT (scan_id, rule_id, resource_id) DO UPDATE SET - resource_type = EXCLUDED.resource_type, - status = EXCLUDED.status, - reason_code = EXCLUDED.reason_code, - reason = EXCLUDED.reason, - evidence = EXCLUDED.evidence, - finding_id = EXCLUDED.finding_id, - evaluated_at = EXCLUDED.evaluated_at - """, - ( - scan_result["scan_id"], - rule_id, - resource_id, - evaluation.get("resource_type") or "", - evaluation.get("status"), - evaluation.get("reason_code"), - evaluation.get("reason"), - json.dumps(evaluation.get("evidence", {})), - finding_ids.get((rule_id, resource_id)), - completed_at, - ), - ) - - if evaluations: - cur.execute( - """ - DELETE FROM rule_evaluations - WHERE scan_id = %s - AND (rule_id, resource_id) NOT IN ( - SELECT * FROM unnest(%s::text[], %s::text[]) - ) - """, - ( - scan_result["scan_id"], - [evaluation["rule_id"] for evaluation in evaluations], - [evaluation["resource_id"] for evaluation in evaluations], - ), - ) - else: - cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_result["scan_id"],)) + # Per-resource rule evaluations (PASS/FAIL/UNKNOWN/ERROR/ + # NOT_APPLICABLE) are the #263 coverage contract and are owned + # by PR #321, not by this change. When that lands, its + # evaluation writes belong here, inside this transaction, so + # they inherit the lease/fencing check above for free. # Completion creates one durable job in this same fenced # transaction. The scan_id uniqueness constraint makes result @@ -538,8 +490,23 @@ def update_scan_enrichment_status(self, scan_id: str, status: str) -> None: conn.commit() logger.info("Updated scan %s enrichment status to %s", scan_id, status) - def enqueue_enrichment_job(self, scan_id: str) -> tuple[Dict[str, Any], bool]: - """Durably enqueue exactly one CVE enrichment job for a scan.""" + def enqueue_enrichment_job(self, scan_id: str) -> tuple[Dict[str, Any], str]: + """Durably enqueue, or explicitly requeue, one CVE enrichment job. + + Returns the job row and one of four outcomes: + + ``created`` a new job row was inserted for this scan. + ``requeued`` a terminally ``failed`` job was reset to ``pending``. + ``active`` a ``pending``/``running`` job already exists. + ``completed`` enrichment already finished; nothing was changed. + + There is never more than one job per scan: the ``scan_id`` unique + constraint makes the insert idempotent, and the requeue is a single + conditional ``UPDATE`` so concurrent callers converge on one row. A + ``running`` job is left strictly alone -- only its lease expiring + (:meth:`recover_stale_enrichment_jobs`) may take it away from its + current owner, so an operator retry can never steal a live claim. + """ conn = self._get_conn() try: with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: @@ -553,25 +520,69 @@ def enqueue_enrichment_job(self, scan_id: str) -> tuple[Dict[str, Any], bool]: (str(uuid.uuid4()), scan_id), ) job = cur.fetchone() - if job is None: - cur.execute("SELECT * FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) - existing = cur.fetchone() - if existing is None: - raise RuntimeError("enrichment job conflict did not return an existing job") + if job is not None: + cur.execute( + "UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", + (scan_id,), + ) conn.commit() - return dict(existing), False + return dict(job), "created" + + # A job already exists. Only a terminally failed one is + # revived, and the WHERE clause is the whole guard: a + # concurrent requeue that lost the race sees status + # 'pending' and matches nothing, so both callers end up + # with the same single pending job. + # + # attempt_count restarts because the operator is explicitly + # granting a fresh retry budget -- leaving it at the limit + # would make the job fail again on its first attempt. The + # previous error_message is kept as the audit trail, and + # checkpoint is kept so the retry resumes rather than + # re-enriching findings that already succeeded. cur.execute( - "UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", + """ + UPDATE enrichment_jobs + SET status = 'pending', + attempt_count = 0, + next_retry_at = CURRENT_TIMESTAMP, + lease_owner = NULL, + lease_expires_at = NULL, + last_heartbeat_at = NULL, + completed_at = NULL + WHERE scan_id = %s AND status = 'failed' + RETURNING * + """, (scan_id,), ) + requeued = cur.fetchone() + if requeued is not None: + cur.execute( + "UPDATE scans SET cve_enrichment_status = 'PENDING' WHERE scan_id = %s", + (scan_id,), + ) + conn.commit() + return dict(requeued), "requeued" + + cur.execute("SELECT * FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + existing = cur.fetchone() + if existing is None: + raise RuntimeError("enrichment job conflict did not return an existing job") conn.commit() - return dict(job), True + existing = dict(existing) + return existing, "completed" if existing["status"] == "completed" else "active" except Exception: self.rollback(conn) raise - def claim_next_enrichment_job(self, lease_owner: str, lease_seconds: int) -> Optional[Dict[str, Any]]: - """Atomically claim the next retry-ready enrichment job.""" + def claim_next_enrichment_job( + self, lease_owner: str, lease_seconds: int, scan_id: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """Atomically claim the next retry-ready enrichment job. + + ``scan_id`` restricts the claim to one scan's job, with the same lease + and fencing semantics as an unrestricted claim. + """ if lease_seconds <= 0: raise ValueError("lease_seconds must be positive") conn = self._get_conn() @@ -590,13 +601,14 @@ def claim_next_enrichment_job(self, lease_owner: str, lease_seconds: int) -> Opt SELECT job_id FROM enrichment_jobs WHERE status = 'pending' AND next_retry_at <= CURRENT_TIMESTAMP + AND (%s IS NULL OR scan_id = %s::uuid) ORDER BY created_at ASC FOR UPDATE SKIP LOCKED LIMIT 1 ) RETURNING * """, - (lease_owner, lease_seconds), + (lease_owner, lease_seconds, scan_id, scan_id), ) job = cur.fetchone() if job: @@ -939,8 +951,16 @@ def update_scan_status( raise logger.info("Updated scan %s status to %s", scan_id, status) - def claim_next_pending_scan(self, lease_owner: str, lease_seconds: int) -> Optional[Dict[str, Any]]: - """Atomically claim one pending scan and establish its renewable lease.""" + def claim_next_pending_scan( + self, lease_owner: str, lease_seconds: int, scan_id: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """Atomically claim one pending scan and establish its renewable lease. + + ``scan_id`` restricts the claim to one specific scan instead of the + oldest pending one. The lease, fencing and skip-locked semantics are + identical either way; callers that must act on a known scan use it so + they neither depend on nor disturb the shared queue order. + """ if lease_seconds <= 0: raise ValueError("lease_seconds must be positive") conn = self._get_conn() @@ -961,13 +981,14 @@ def claim_next_pending_scan(self, lease_owner: str, lease_seconds: int) -> Optio SELECT scan_id FROM scans WHERE status = 'pending' + AND (%s IS NULL OR scan_id = %s::uuid) ORDER BY started_at ASC FOR UPDATE SKIP LOCKED LIMIT 1 ) RETURNING * """, - (lease_owner, lease_seconds), + (lease_owner, lease_seconds, scan_id, scan_id), ) row = cur.fetchone() conn.commit() @@ -1085,10 +1106,24 @@ def get_scans(self) -> List[Dict[str, Any]]: cur.execute("SELECT * FROM scans ORDER BY started_at DESC LIMIT 100") return [dict(row) for row in cur.fetchall()] - def record_worker_heartbeat(self, worker_id: str, worker_type: str) -> None: - """Persist liveness without using worker IDs as metric labels.""" + def record_worker_heartbeat( + self, + worker_id: str, + worker_type: str, + retention_seconds: int = DEFAULT_WORKER_HEARTBEAT_RETENTION_SECONDS, + ) -> None: + """Persist liveness without using worker IDs as metric labels. + + Worker identities are per-process, so a long-lived deployment would + otherwise accumulate one dead row per restart forever. Retired rows + are pruned only on the heartbeat that actually inserts a new worker + identity -- i.e. once per worker process -- so the table stays bounded + without a scheduler and without a full-table sweep on every beat. + """ if worker_type not in {"scan", "enrichment"}: raise ValueError("unsupported worker type") + if retention_seconds <= 0: + raise ValueError("retention_seconds must be positive") conn = self._get_conn() try: with conn.cursor() as cur: @@ -1098,9 +1133,22 @@ def record_worker_heartbeat(self, worker_id: str, worker_type: str) -> None: VALUES (%s, %s, CURRENT_TIMESTAMP) ON CONFLICT (worker_id, worker_type) DO UPDATE SET last_seen_at = EXCLUDED.last_seen_at + RETURNING (xmax = 0) AS inserted """, (worker_id, worker_type), ) + row = cur.fetchone() + inserted = row["inserted"] if isinstance(row, dict) else row[0] + if inserted: + # Retention is far longer than any heartbeat interval, so a + # live worker can never prune itself or a peer. + cur.execute( + """ + DELETE FROM worker_heartbeats + WHERE last_seen_at < CURRENT_TIMESTAMP - (%s * INTERVAL '1 second') + """, + (retention_seconds,), + ) conn.commit() except Exception: self.rollback(conn) diff --git a/api/routes/scans.py b/api/routes/scans.py index a84d233a..89907e33 100644 --- a/api/routes/scans.py +++ b/api/routes/scans.py @@ -169,6 +169,14 @@ def trigger_scan(): return jsonify({"error": "Critical route failure"}), 500 +_ENRICH_MESSAGES = { + "created": "CVE enrichment queued; poll GET /api/scans/ for completion.", + "requeued": "Previously failed enrichment job requeued; poll GET /api/scans/ for completion.", + "active": "Existing enrichment job returned.", + "completed": "Scan already enriched", +} + + @scans_bp.post("/api/scans//enrich") def enrich_scan(scan_id): """Enqueue durable CVE enrichment; no request-owned thread is created.""" @@ -188,25 +196,17 @@ def enrich_scan(scan_id): if not findings: return jsonify({"error": "No findings found for this scan"}), 404 - job, created = db.enqueue_enrichment_job(scan_id) - if not created: - return jsonify( - { - "job_id": str(job["job_id"]), - "scan_id": scan_id, - "status": job["status"], - "message": "Existing enrichment job returned.", - } - ), 202 - - return jsonify( - { - "scan_id": scan_id, - "job_id": str(job["job_id"]), - "status": "PENDING", - "message": "CVE enrichment queued; poll GET /api/scans/ for completion.", - } - ), 202 + job, outcome = db.enqueue_enrichment_job(scan_id) + body = { + "scan_id": scan_id, + "job_id": str(job["job_id"]), + "status": job["status"], + "outcome": outcome, + "message": _ENRICH_MESSAGES[outcome], + } + # A job that already finished is reported as-is rather than restarted; + # every other outcome leaves exactly one queued or running job. + return jsonify(body), 200 if outcome == "completed" else 202 except ValidationError: return jsonify({"error": VALIDATION_ERROR_MESSAGE}), 400 diff --git a/docs/api-reference.md b/docs/api-reference.md index 0ddbf8fd..ceb8930f 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -200,9 +200,9 @@ Example response: ## POST /api/scans/trigger -Triggers an asynchronous scan against the configured subscription. Returns `202 Accepted` with the `scan_id` immediately. The actual scan execution happens in a background worker process. +Admits an asynchronous scan against the configured subscription. Execution happens in a background worker process; the response returns as soon as the scan is durably recorded. -Request body: +Request body (optional — falls back to `AZURE_SUBSCRIPTION_ID`): ```json { @@ -210,7 +210,26 @@ Request body: } ``` -Example response: +### Admission semantics + +Admission is serialized per subscription and enforced by the database, so concurrent and replayed triggers converge on one logical scan rather than creating duplicates: + +- **At most one active scan per subscription.** While a `pending` or `running` scan exists, a further trigger returns that existing scan instead of queueing another. +- **`Idempotency-Key` (optional request header, 1–200 characters).** A repeat of the same key for the same subscription returns the original scan. The key is scoped to the subscription; the same key under a different subscription is a different request. +- **`OPENSHIELD_MAX_SCANS_PER_SUBSCRIPTION_PER_HOUR`** adds an optional hourly admission quota. Unset or `0` (the default) applies no time-window limit; the one-active-scan rule still applies. + +### Responses + +| Status | When | Body | +| --- | --- | --- | +| `202 Accepted` | A new scan was admitted and queued. | `scan_id`, `status: "pending"`, `message` | +| `200 OK` | The request resolved to an existing logical scan — an `Idempotency-Key` replay of the same request, or a trigger while a scan is already active for the subscription. | `scan_id`, `status` (the existing scan's `pending`/`running`), `message: "Existing logical scan returned."` | +| `400 Bad Request` | Malformed body, invalid `subscription_id`, missing subscription, or an `Idempotency-Key` outside 1–200 characters. | `error` | +| `403 Forbidden` | `subscription_id` is not on the `OPENSHIELD_AUTHORIZED_SUBSCRIPTIONS` allowlist. | `error` | +| `409 Conflict` | The `Idempotency-Key` was reused with a different request payload. | `error: "Idempotency-Key is already associated with a different request."` | +| `429 Too Many Requests` | The configured hourly quota for this subscription is exhausted. | `error: "Scan quota exceeded for this subscription."` | + +New scan (`202`): ```json { @@ -220,7 +239,17 @@ Example response: } ``` -Missing subscription response: +Replay or already-active scan (`200`): + +```json +{ + "scan_id": "6f4a08ac-7d3a-4d9a-a4b4-2a26e5f63c8a", + "status": "running", + "message": "Existing logical scan returned." +} +``` + +Missing subscription response (`400`): ```json { @@ -230,6 +259,54 @@ Missing subscription response: --- +## POST /api/scans/<scan_id>/enrich + +Queues durable CVE enrichment for a completed scan's findings. Enrichment runs as a database-backed job claimed by the background worker — the request never owns a thread, so the work survives an API restart. + +There is **never more than one enrichment job per scan**. Repeat calls are safe: they report the state of the single job rather than creating another. + +### Responses + +Every response carries `scan_id`, `job_id`, `status` (the job row's state) and an `outcome` naming what this call did: + +| Status | `outcome` | When | +| --- | --- | --- | +| `202 Accepted` | `created` | No job existed; one was queued. | +| `202 Accepted` | `requeued` | A previously **failed** job was reset to `pending` and will be retried. | +| `202 Accepted` | `active` | A `pending` or `running` job already exists and was returned unchanged. A live claim is never interrupted. | +| `200 OK` | `completed` | Enrichment already finished; nothing was restarted. | +| `404 Not Found` | — | Unknown `scan_id`, or the scan has no findings to enrich. | + +A job that exhausts its retry budget becomes `failed`. Re-POSTing this endpoint is the supported operator recovery: it atomically returns the job to `pending` with a fresh retry budget, clears the lease, and keeps the last `error_message` and the `checkpoint` so the retry resumes rather than re-enriching findings that already succeeded. Concurrent re-POSTs converge — exactly one reports `requeued` and the rest report `active`. + +Newly queued (`202`): + +```json +{ + "scan_id": "6f4a08ac-7d3a-4d9a-a4b4-2a26e5f63c8a", + "job_id": "1f2e3d4c-5b6a-4790-8123-456789abcdef", + "status": "pending", + "outcome": "created", + "message": "CVE enrichment queued; poll GET /api/scans/ for completion." +} +``` + +Requeued after terminal failure (`202`): + +```json +{ + "scan_id": "6f4a08ac-7d3a-4d9a-a4b4-2a26e5f63c8a", + "job_id": "1f2e3d4c-5b6a-4790-8123-456789abcdef", + "status": "pending", + "outcome": "requeued", + "message": "Previously failed enrichment job requeued; poll GET /api/scans/ for completion." +} +``` + +Poll `GET /api/scans/` for `cve_enrichment_status` (`PENDING`, `ENRICHING`, `COMPLETED`, `FAILED`). + +--- + ## GET /api/score Returns the overall security posture score from 0 to 100. Under [severity contract v1](severity-contract.md), the score starts at 100 and deducts 20 per CRITICAL finding, 10 per HIGH finding, 5 per MEDIUM finding, and 2 per LOW finding. INFO findings deduct zero. diff --git a/docs/async-scan-architecture.md b/docs/async-scan-architecture.md index 757aec61..e8c6eeb6 100644 --- a/docs/async-scan-architecture.md +++ b/docs/async-scan-architecture.md @@ -33,20 +33,32 @@ credentials or persistent Azure errors cannot retry forever. Findings use a stable database-enforced identity (`scan`, rule, canonical resource scope, and an optional rule-specific discriminator). Result retries use PostgreSQL upserts, so mutable text or severity is updated rather than -creating a second authoritative finding. Rule evaluations use the same -database-first uniqueness model. +creating a second authoritative finding. Per-resource rule evaluations are a +separate contract (issue #263) and are not persisted by this architecture yet; +when they land, they are written inside the same fenced completion transaction +and inherit its ownership check. ### 3. The Worker (Python) -The scanner/worker.py process runs independently of the web server. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings, evaluations, and one durable enrichment job atomically. On failure, it records a sanitized error only while it still owns the lease. +The scanner/worker.py process runs independently of the web server. It atomically claims a pending scan, starts a dedicated lease-heartbeat connection, and invokes `ScanEngine.run_scan(scan_id)`. Completion and failure are fenced database transactions: they only succeed when the current worker still owns the same unexpired token. On success, it persists findings and one durable enrichment job atomically. On failure, it records a sanitized error only while it still owns the lease. ### Durable CVE enrichment `POST /api/scans//enrich` enqueues (or returns) the one durable PostgreSQL enrichment job for the scan; it never starts a request-owned daemon thread. The scan worker claims those jobs with the same owner/expiry/fencing model, checkpoints after each finding, and retries transient failures with bounded exponential backoff. An expired job is recovered or terminally failed after its attempt limit. NVD retrieval follows every `totalResults` page; replaying a checkpoint updates the existing finding instead of duplicating CVE data. +A job that exhausts its retries becomes `failed`. Re-POSTing the endpoint is the supported recovery path: it atomically returns that job to `pending` with a fresh retry budget while keeping the same job row, its last `error_message`, and its `checkpoint`. A `running` job with a live lease is never disturbed by a re-POST — only lease expiry can move it — and a `completed` job is never restarted. See [api-reference.md](api-reference.md) for the exact response contract. + +### Worker scheduling + +One worker process serves both durable queues. Each loop iteration takes **at most one** enrichment job and **at most one** scan, so neither queue can starve the other no matter how deep either gets: a large enrichment backlog delays scans by one job per iteration rather than blocking them until it drains. An iteration that did any work polls again immediately; only a completely idle iteration sleeps for the poll interval. + ### Operational signals `/metrics` derives bounded-cardinality operational gauges from PostgreSQL: worker heartbeat age, oldest queue age, oldest active lease age, aggregate retry attempts, and the last successful scan timestamp. Labels are limited to `queue` (`scan` or `enrichment`) and `worker_type`; scan, job, subscription, and worker identifiers are never metric labels. +These gauges are recomputed on every scrape rather than cached. The queue, lease and heartbeat aggregates are served by partial indexes and touch only active rows; the last-successful-scan lookup is served by a partial index on completed scans. The two `retry_attempts` sums scan their whole table and therefore grow with scan history — acceptable at current volumes, and the thing to revisit first (a short-TTL in-process cache) if `/metrics` scrape latency ever becomes visible. + +Worker identities are per-process, so `worker_heartbeats` would otherwise gain a permanent row on every restart. Rows older than `WORKER_HEARTBEAT_RETENTION_SECONDS` (default 7 days) are pruned on the heartbeat that registers a new worker identity — once per worker process, never on every beat. Retention is far longer than any heartbeat interval, so a live worker is never pruned. + ## Technical Rationale ### Why not Celery or Redis @@ -62,6 +74,31 @@ workers, apply Alembic migrations, then start the fenced worker version. Legacy workers do not carry ownership/fencing state; legacy running scans are retained and made recoverable by the lease migration rather than deleted. +### Migration prerequisite: one active scan per subscription + +`a7c5e9d2f1b4` adds the unique index that enforces one `pending`/`running` scan +per subscription. A deployment that predates that rule may already hold several, +in which case the migration **stops before creating any index** and names the +offending subscriptions: + +``` +Cannot enforce one active scan per subscription: 1 subscription(s) already have +more than one pending/running scan: (2 active). ... +``` + +Nothing is changed when this happens — the migration will not decide which of +your production scans is authoritative. Resolve it, then re-run: + +1. Drain the old workers first (step 1 above), so no new active scans appear. +2. Let the in-flight scans finish, or mark the superseded rows `failed` + (`UPDATE scans SET status = 'failed', error_message = '...' WHERE scan_id = ...`). + Never delete the rows: scan history is retained, and only `pending`/`running` + rows are constrained — any number of `completed`/`failed` scans per + subscription remains valid. +3. Re-run `alembic upgrade head`. A retry is safe: the migration drops the + INVALID index that an interrupted `CREATE INDEX CONCURRENTLY` leaves behind + before rebuilding it. + ## Testing Suite The asynchronous transition is verified through a multi layered testing strategy. diff --git a/scanner/worker.py b/scanner/worker.py index b5021e11..2ae431ae 100644 --- a/scanner/worker.py +++ b/scanner/worker.py @@ -13,7 +13,11 @@ import uuid from datetime import datetime, timezone -from api.models.finding import DatabaseManager, LostLease +from api.models.finding import ( + DEFAULT_WORKER_HEARTBEAT_RETENTION_SECONDS, + DatabaseManager, + LostLease, +) from api.observability import ( PENDING_SCANS, SCAN_DURATION_SECONDS, @@ -139,14 +143,17 @@ def run_worker(): db = DatabaseManager(db_url) worker_id = str(uuid.uuid4()) lease_seconds, heartbeat_seconds = lease_configuration() + heartbeat_retention_seconds = _positive_seconds( + "WORKER_HEARTBEAT_RETENTION_SECONDS", DEFAULT_WORKER_HEARTBEAT_RETENTION_SECONDS + ) logger.info("OpenShield Background Worker started. Polling every %ds", POLL_INTERVAL_SECONDS) while True: try: - db.record_worker_heartbeat(worker_id, "scan") + db.record_worker_heartbeat(worker_id, "scan", heartbeat_retention_seconds) # This process executes both durable queue types; record both # liveness signals without exporting the worker UUID as a label. - db.record_worker_heartbeat(worker_id, "enrichment") + db.record_worker_heartbeat(worker_id, "enrichment", heartbeat_retention_seconds) # 1. Cleanup stale scans from previous crashes db.recover_stale_scans() db.recover_stale_enrichment_jobs() @@ -154,16 +161,22 @@ def run_worker(): # 2. Publish current queue depth PENDING_SCANS.set(len(db.get_pending_scans())) - # 3. Run one durable enrichment job before taking another scan. + # 3. Take at most one job from each durable queue per iteration. + # Draining enrichment first and restarting the loop would let a + # sustained enrichment backlog hold off every pending scan, so + # the two queues alternate instead: neither can starve the + # other regardless of how deep either one gets. enrichment_job = db.claim_next_enrichment_job(worker_id, lease_seconds) if enrichment_job: process_enrichment_job(db, enrichment_job, worker_id, lease_seconds) - continue # 4. Atomic scan claim scan = db.claim_next_pending_scan(worker_id, lease_seconds) if not scan: - time.sleep(POLL_INTERVAL_SECONDS) + # Only idle when there was no work at all; an iteration that + # ran an enrichment job polls again immediately. + if not enrichment_job: + time.sleep(POLL_INTERVAL_SECONDS) continue scan_id = str(scan["scan_id"]) diff --git a/tests/test_enrichment_jobs_postgres.py b/tests/test_enrichment_jobs_postgres.py index b98b73fe..0b4780b0 100644 --- a/tests/test_enrichment_jobs_postgres.py +++ b/tests/test_enrichment_jobs_postgres.py @@ -6,6 +6,7 @@ from unittest.mock import patch import psycopg2 +import psycopg2.extras import pytest from api.models.finding import DatabaseManager, LostLease @@ -24,7 +25,11 @@ def enrichment_scan(): db = DatabaseManager(dsn) try: db.create_pending_scan(scan_id, subscription_id) - claim = db.claim_next_pending_scan("seed", 120) + # Claim this scan explicitly. An unrestricted claim takes the globally + # oldest pending scan, which in a shared test database is very often a + # row another test admitted first. + claim = db.claim_next_pending_scan("seed", 120, scan_id=scan_id) + assert claim is not None and str(claim["scan_id"]) == scan_id result = { "scan_id": scan_id, "subscription_id": subscription_id, @@ -50,23 +55,23 @@ def enrichment_scan(): ], } db.save_scan(result, "seed", claim["fencing_token"]) - job, created = db.enqueue_enrichment_job(scan_id) - assert created is False + job, outcome = db.enqueue_enrichment_job(scan_id) + assert outcome == "active" yield dsn, scan_id, job finally: db.close() with psycopg2.connect(dsn) as conn: with conn.cursor() as cur: cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) - cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) -def _claim(dsn): +def _claim(dsn, scan_id, owner="worker-a"): + """Claim this scan's enrichment job, never another test's.""" db = DatabaseManager(dsn) try: - return db.claim_next_enrichment_job("worker-a", 120) + return db.claim_next_enrichment_job(owner, 120, scan_id=scan_id) finally: db.close() @@ -75,10 +80,10 @@ def test_duplicate_enqueue_and_claim_race(enrichment_scan): dsn, scan_id, first_job = enrichment_scan db = DatabaseManager(dsn) try: - replay, created = db.enqueue_enrichment_job(scan_id) + replay, outcome = db.enqueue_enrichment_job(scan_id) finally: db.close() - assert created is False + assert outcome == "active" assert replay["job_id"] == first_job["job_id"] barrier = threading.Barrier(2) @@ -86,7 +91,7 @@ def test_duplicate_enqueue_and_claim_race(enrichment_scan): def claim(): barrier.wait() - claims.append(_claim(dsn)) + claims.append(_claim(dsn, scan_id)) threads = [threading.Thread(target=claim) for _ in range(2)] for thread in threads: @@ -98,7 +103,7 @@ def claim(): def test_checkpoint_resume_and_completion(enrichment_scan): dsn, scan_id, _ = enrichment_scan - job = _claim(dsn) + job = _claim(dsn, scan_id) assert job is not None db = DatabaseManager(dsn) try: @@ -121,7 +126,7 @@ def enrich_once_then_fail(finding): cur.execute( "UPDATE enrichment_jobs SET next_retry_at = CURRENT_TIMESTAMP WHERE scan_id = %s", (scan_id,) ) - resumed = db.claim_next_enrichment_job("worker-b", 120) + resumed = db.claim_next_enrichment_job("worker-b", 120, scan_id=scan_id) with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: enrich.side_effect = lambda finding: {**finding, "cve_references": [{"cve_id": "CVE-1"}]} assert process_enrichment_job(db, resumed, "worker-b", 120) == "completed" @@ -142,7 +147,7 @@ def test_enrichment_retry_limit_becomes_terminal(enrichment_scan): db = DatabaseManager(dsn) try: for attempt in range(1, 4): - job = db.claim_next_enrichment_job("worker-a", 120) + job = db.claim_next_enrichment_job("worker-a", 120, scan_id=scan_id) assert job is not None with patch("scanner.enrichment_worker.enrich_finding_durable", side_effect=RuntimeError("NVD unavailable")): expected = "failed" if attempt == 3 else "retry" @@ -164,7 +169,7 @@ def test_enrichment_retry_limit_becomes_terminal(enrichment_scan): def test_expired_job_is_recovered_with_new_token_and_stale_owner_is_rejected(enrichment_scan): dsn, scan_id, _ = enrichment_scan - first = _claim(dsn) + first = _claim(dsn, scan_id) assert first is not None with psycopg2.connect(dsn) as conn: with conn.cursor() as cur: @@ -178,10 +183,190 @@ def test_expired_job_is_recovered_with_new_token_and_stale_owner_is_rejected(enr ) db = DatabaseManager(dsn) try: - assert db.recover_stale_enrichment_jobs() == 1 - second = db.claim_next_enrichment_job("worker-b", 120) + # Other tests may share this database; assert this job specifically + # was recovered rather than that it was the only one. + assert db.recover_stale_enrichment_jobs() >= 1 + second = db.claim_next_enrichment_job("worker-b", 120, scan_id=scan_id) assert second["fencing_token"] > first["fencing_token"] with pytest.raises(LostLease): db.heartbeat_enrichment_job(str(first["job_id"]), "worker-a", first["fencing_token"], 120) finally: db.close() + + +def _fail_terminally(dsn, db, scan_id): + """Drive a job through its whole retry budget until it is 'failed'.""" + for attempt in range(1, 4): + job = db.claim_next_enrichment_job("worker-a", 120, scan_id=scan_id) + assert job is not None + with patch("scanner.enrichment_worker.enrich_finding_durable", side_effect=RuntimeError("NVD unavailable")): + process_enrichment_job(db, job, "worker-a", 120) + if attempt < 3: + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "UPDATE enrichment_jobs SET next_retry_at = CURRENT_TIMESTAMP WHERE scan_id = %s", + (scan_id,), + ) + return _job_row(dsn, scan_id) + + +def _job_row(dsn, scan_id): + with psycopg2.connect(dsn) as conn: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute("SELECT * FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + return dict(cur.fetchone()) + + +def test_terminally_failed_job_can_be_explicitly_requeued(enrichment_scan): + dsn, scan_id, first_job = enrichment_scan + db = DatabaseManager(dsn) + try: + failed = _fail_terminally(dsn, db, scan_id) + assert failed["status"] == "failed" + assert failed["attempt_count"] >= 3 + + job, outcome = db.enqueue_enrichment_job(scan_id) + finally: + db.close() + + assert outcome == "requeued" + # Same logical job, now retryable again. + assert str(job["job_id"]) == str(first_job["job_id"]) + assert job["status"] == "pending" + assert job["attempt_count"] == 0 + assert job["lease_owner"] is None + assert job["lease_expires_at"] is None + assert job["completed_at"] is None + # The failure reason is kept as the audit trail, and the checkpoint is kept + # so the retry resumes instead of re-enriching what already succeeded. + assert job["error_message"] + assert job["checkpoint"] == failed["checkpoint"] + assert _job_row(dsn, scan_id)["next_retry_at"] is not None + + +def test_completed_job_is_never_restarted(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + job = db.claim_next_enrichment_job("worker-a", 120, scan_id=scan_id) + with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: + enrich.side_effect = lambda finding: {**finding, "cve_references": [{"cve_id": "CVE-1"}]} + assert process_enrichment_job(db, job, "worker-a", 120) == "completed" + + requeued, outcome = db.enqueue_enrichment_job(scan_id) + finally: + db.close() + + assert outcome == "completed" + assert requeued["status"] == "completed" + assert _job_row(dsn, scan_id)["status"] == "completed" + + +def test_requeue_does_not_steal_a_valid_running_lease(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + running = db.claim_next_enrichment_job("worker-a", 120, scan_id=scan_id) + assert running is not None + + job, outcome = db.enqueue_enrichment_job(scan_id) + assert outcome == "active" + assert job["status"] == "running" + + # The live owner keeps its lease and can still heartbeat and complete. + db.heartbeat_enrichment_job(str(running["job_id"]), "worker-a", running["fencing_token"], 120) + after = _job_row(dsn, scan_id) + assert after["lease_owner"] == "worker-a" + assert after["fencing_token"] == running["fencing_token"] + finally: + db.close() + + +def test_concurrent_requeues_converge_on_one_logical_job(enrichment_scan): + dsn, scan_id, first_job = enrichment_scan + db = DatabaseManager(dsn) + try: + _fail_terminally(dsn, db, scan_id) + finally: + db.close() + + barrier = threading.Barrier(4) + outcomes: list[str] = [] + + def requeue() -> None: + worker_db = DatabaseManager(dsn) + try: + barrier.wait() + outcomes.append(worker_db.enqueue_enrichment_job(scan_id)[1]) + finally: + worker_db.close() + + threads = [threading.Thread(target=requeue) for _ in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + # Exactly one caller performs the failed -> pending transition; the rest + # observe the job that is already queued. Nobody creates a second job. + assert sorted(outcomes) == ["active", "active", "active", "requeued"] + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone()[0] == 1 + assert str(_job_row(dsn, scan_id)["job_id"]) == str(first_job["job_id"]) + + +def test_stale_token_cannot_write_after_requeue_and_reclaim(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + # worker-a held the job through its last failed attempt; that is the + # token a stale process would still be carrying. + failed = _fail_terminally(dsn, db, scan_id) + stale_job_id, stale_token = str(failed["job_id"]), failed["fencing_token"] + + assert db.enqueue_enrichment_job(scan_id)[1] == "requeued" + reclaimed = db.claim_next_enrichment_job("worker-b", 120, scan_id=scan_id) + assert reclaimed is not None + assert reclaimed["fencing_token"] > stale_token + + findings = db.get_enrichment_findings(scan_id) + with pytest.raises(LostLease): + db.heartbeat_enrichment_job(stale_job_id, "worker-a", stale_token, 120) + with pytest.raises(LostLease): + db.persist_enrichment_progress(stale_job_id, "worker-a", stale_token, findings[0], 99) + with pytest.raises(LostLease): + db.complete_enrichment_job(stale_job_id, "worker-a", stale_token) + + # None of the rejected writes landed. + current = _job_row(dsn, scan_id) + assert current["checkpoint"] != 99 + assert current["status"] == "running" + assert current["lease_owner"] == "worker-b" + finally: + db.close() + + +def test_requeued_job_can_eventually_complete(enrichment_scan): + dsn, scan_id, _ = enrichment_scan + db = DatabaseManager(dsn) + try: + _fail_terminally(dsn, db, scan_id) + _, outcome = db.enqueue_enrichment_job(scan_id) + assert outcome == "requeued" + + job = db.claim_next_enrichment_job("worker-b", 120, scan_id=scan_id) + assert job is not None + with patch("scanner.enrichment_worker.enrich_finding_durable") as enrich: + enrich.side_effect = lambda finding: {**finding, "cve_references": [{"cve_id": "CVE-2"}]} + assert process_enrichment_job(db, job, "worker-b", 120) == "completed" + finally: + db.close() + + assert _job_row(dsn, scan_id)["status"] == "completed" + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT cve_enrichment_status FROM scans WHERE scan_id = %s", (scan_id,)) + assert cur.fetchone()[0] == "COMPLETED" diff --git a/tests/test_operational_metrics_postgres.py b/tests/test_operational_metrics_postgres.py index f4f63c7c..c8df812b 100644 --- a/tests/test_operational_metrics_postgres.py +++ b/tests/test_operational_metrics_postgres.py @@ -22,11 +22,13 @@ def test_durable_operational_metrics_cover_queue_lease_retries_and_heartbeat(): db.create_pending_scan(scan_id, subscription_id) db.record_worker_heartbeat("worker-test", "scan") db.record_worker_heartbeat("worker-test", "enrichment") - claim = db.claim_next_pending_scan("worker-a", 120) + claim = db.claim_next_pending_scan("worker-a", 120, scan_id=scan_id) assert claim is not None snapshot = db.get_operational_metrics() assert snapshot["oldest_lease_age"]["scan"] >= 0 - assert snapshot["retry_attempts"]["scan"] == 0 + # Aggregated over the whole table, so other rows in a shared test + # database may contribute; only the shape is this test's contract. + assert snapshot["retry_attempts"]["scan"] >= 0 assert snapshot["worker_heartbeat_age"]["scan"] >= 0 assert snapshot["worker_heartbeat_age"]["enrichment"] >= 0 assert snapshot["last_successful_scan_timestamp"] >= 0 @@ -37,3 +39,44 @@ def test_durable_operational_metrics_cover_queue_lease_retries_and_heartbeat(): cur.execute("DELETE FROM worker_heartbeats WHERE worker_id = 'worker-test'") cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) + + +def test_stale_worker_heartbeats_are_pruned_without_dropping_live_workers(): + """Retired worker rows must not accumulate, and live ones must survive.""" + dsn = os.environ["DATABASE_URL"] + live_worker, dead_worker = f"live-{uuid.uuid4()}", f"dead-{uuid.uuid4()}" + db = DatabaseManager(dsn) + try: + db.record_worker_heartbeat(dead_worker, "scan") + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "UPDATE worker_heartbeats SET last_seen_at = CURRENT_TIMESTAMP - INTERVAL '30 days' " + "WHERE worker_id = %s", + (dead_worker,), + ) + + # Pruning happens on the beat that inserts a new worker identity. + db.record_worker_heartbeat(live_worker, "scan", retention_seconds=7 * 24 * 60 * 60) + + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM worker_heartbeats WHERE worker_id = %s", (dead_worker,)) + assert cur.fetchone()[0] == 0 + cur.execute("SELECT COUNT(*) FROM worker_heartbeats WHERE worker_id = %s", (live_worker,)) + assert cur.fetchone()[0] == 1 + + # A repeat beat from an existing worker only refreshes its timestamp. + db.record_worker_heartbeat(live_worker, "scan") + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM worker_heartbeats WHERE worker_id = %s", (live_worker,)) + assert cur.fetchone()[0] == 1 + finally: + db.close() + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + "DELETE FROM worker_heartbeats WHERE worker_id IN (%s, %s)", + (live_worker, dead_worker), + ) diff --git a/tests/test_scan_admission_migration_postgres.py b/tests/test_scan_admission_migration_postgres.py new file mode 100644 index 00000000..bb10ad10 --- /dev/null +++ b/tests/test_scan_admission_migration_postgres.py @@ -0,0 +1,244 @@ +"""Migration-time safety for the one-active-scan unique index. + +``CREATE UNIQUE INDEX CONCURRENTLY`` fails -- and leaves an INVALID index +behind -- when a deployment already holds several active scans for one +subscription. These tests run the real migration against a throwaway +PostgreSQL database seeded with exactly that legacy shape. +""" + +import os +import uuid +from contextlib import contextmanager + +import psycopg2 +import pytest +from alembic import command +from alembic.config import Config + + +pytestmark = pytest.mark.skipif( + not os.environ.get("DATABASE_URL"), reason="DATABASE_URL is required for PostgreSQL tests" +) + +# The revision immediately before scan admission is introduced. +_BEFORE_ADMISSION = "f2b6d8e1a4c9" +_ADMISSION = "a7c5e9d2f1b4" +_ACTIVE_INDEX = "uq_scans_one_active_per_subscription" +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +@contextmanager +def _scratch_database(): + """Create and drop a database dedicated to one migration test.""" + base = os.environ["DATABASE_URL"].rsplit("/", 1)[0] + name = f"openshield_mig_{uuid.uuid4().hex[:12]}" + admin = psycopg2.connect(f"{base}/postgres") + admin.autocommit = True + try: + with admin.cursor() as cur: + cur.execute(f'CREATE DATABASE "{name}"') + yield f"{base}/{name}" + finally: + with admin.cursor() as cur: + cur.execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = %s AND pid <> pg_backend_pid()", + (name,), + ) + cur.execute(f'DROP DATABASE IF EXISTS "{name}"') + admin.close() + + +def _alembic_config() -> Config: + """Alembic config that does not touch this process's logging setup. + + Passing alembic.ini would make ``env.py`` call ``fileConfig()``, which + disables every logger that already exists -- silencing the rest of the + test session. Only ``script_location`` is actually needed here; ``env.py`` + reads the database URL from the environment. + """ + config = Config() + config.set_main_option("script_location", os.path.join(_REPO_ROOT, "alembic")) + return config + + +def _upgrade(dsn: str, revision: str) -> None: + """Run alembic against ``dsn`` without disturbing the ambient env.""" + previous = os.environ.get("DATABASE_URL") + os.environ["DATABASE_URL"] = dsn + try: + command.upgrade(_alembic_config(), revision) + finally: + if previous is None: + os.environ.pop("DATABASE_URL", None) + else: + os.environ["DATABASE_URL"] = previous + + +def _seed_active_scans(dsn: str, subscription_id: str, count: int) -> list[str]: + scan_ids = [str(uuid.uuid4()) for _ in range(count)] + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + for offset, scan_id in enumerate(scan_ids): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status) + VALUES (%s, %s, CURRENT_TIMESTAMP - (%s * INTERVAL '1 minute'), 'pending') + """, + (scan_id, subscription_id, offset), + ) + return scan_ids + + +def _index_is_valid(dsn: str, name: str): + """Return True/False for a present index, or None when it does not exist.""" + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute( + """ + SELECT i.indisvalid + FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid + WHERE c.relname = %s + """, + (name,), + ) + row = cur.fetchone() + return None if row is None else row[0] + + +def test_legacy_duplicate_active_scans_fail_the_migration_with_an_actionable_error(): + subscription_id = str(uuid.uuid4()) + with _scratch_database() as dsn: + _upgrade(dsn, _BEFORE_ADMISSION) + _seed_active_scans(dsn, subscription_id, 2) + + with pytest.raises(RuntimeError) as excinfo: + _upgrade(dsn, _ADMISSION) + + message = str(excinfo.value) + assert subscription_id in message + assert "2 active" in message + # The operator is told what to do, and no unusable index was left behind. + assert "Resolve them first" in message + assert _index_is_valid(dsn, _ACTIVE_INDEX) is None + + # Historical rows are untouched: the migration decides nothing for us. + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("SELECT COUNT(*) FROM scans WHERE subscription_id = %s", (subscription_id,)) + assert cur.fetchone()[0] == 2 + + +def test_migration_succeeds_and_enforces_the_contract_after_cleanup(): + subscription_id = str(uuid.uuid4()) + with _scratch_database() as dsn: + _upgrade(dsn, _BEFORE_ADMISSION) + scan_ids = _seed_active_scans(dsn, subscription_id, 2) + + with pytest.raises(RuntimeError): + _upgrade(dsn, _ADMISSION) + + # The documented cleanup: retire the superseded scan, keep the record. + with psycopg2.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("UPDATE scans SET status = 'failed' WHERE scan_id = %s", (scan_ids[1],)) + + _upgrade(dsn, _ADMISSION) + assert _index_is_valid(dsn, _ACTIVE_INDEX) is True + + conn = psycopg2.connect(dsn) + try: + with conn.cursor() as cur: + # Both historical rows survived the migration. + cur.execute("SELECT COUNT(*) FROM scans WHERE subscription_id = %s", (subscription_id,)) + assert cur.fetchone()[0] == 2 + + # A second active scan for the same subscription is now refused. + with pytest.raises(psycopg2.errors.UniqueViolation): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'running') + """, + (str(uuid.uuid4()), subscription_id), + ) + conn.rollback() + + with conn.cursor() as cur: + # Terminal scans stay unconstrained: history keeps accumulating. + for _ in range(3): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed') + """, + (str(uuid.uuid4()), subscription_id), + ) + # Idempotency-key uniqueness is per subscription and ignores NULLs. + for _ in range(2): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status, idempotency_key) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed', NULL) + """, + (str(uuid.uuid4()), subscription_id), + ) + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status, idempotency_key) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed', 'key-1') + """, + (str(uuid.uuid4()), subscription_id), + ) + conn.commit() + + with conn.cursor() as cur: + with pytest.raises(psycopg2.errors.UniqueViolation): + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status, idempotency_key) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed', 'key-1') + """, + (str(uuid.uuid4()), subscription_id), + ) + conn.rollback() + + # The same key under a different subscription is still admissible. + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO scans (scan_id, subscription_id, started_at, status, idempotency_key) + VALUES (%s, %s, CURRENT_TIMESTAMP, 'completed', 'key-1') + """, + (str(uuid.uuid4()), str(uuid.uuid4())), + ) + conn.commit() + finally: + conn.close() + + +def test_an_invalid_index_left_by_a_failed_build_is_replaced_not_inherited(): + """A retry after a failed CONCURRENTLY build must not trip over its debris.""" + subscription_id = str(uuid.uuid4()) + with _scratch_database() as dsn: + _upgrade(dsn, _BEFORE_ADMISSION) + _seed_active_scans(dsn, subscription_id, 1) + + # Reproduce the debris a failed concurrent build leaves behind. + conn = psycopg2.connect(dsn) + conn.autocommit = True + try: + with conn.cursor() as cur: + cur.execute( + f"CREATE UNIQUE INDEX CONCURRENTLY {_ACTIVE_INDEX} " + "ON scans (subscription_id) WHERE status IN ('pending', 'running')" + ) + cur.execute( + "UPDATE pg_index SET indisvalid = false WHERE indexrelid = %s::regclass", + (_ACTIVE_INDEX,), + ) + finally: + conn.close() + assert _index_is_valid(dsn, _ACTIVE_INDEX) is False + + _upgrade(dsn, _ADMISSION) + assert _index_is_valid(dsn, _ACTIVE_INDEX) is True diff --git a/tests/test_scan_admission_postgres.py b/tests/test_scan_admission_postgres.py index 9c0fa829..05ed909f 100644 --- a/tests/test_scan_admission_postgres.py +++ b/tests/test_scan_admission_postgres.py @@ -23,7 +23,6 @@ def admitted_scans(): with psycopg2.connect(dsn) as conn: with conn.cursor() as cur: for scan_id in scan_ids: - cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) diff --git a/tests/test_scan_leases_postgres.py b/tests/test_scan_leases_postgres.py index fd5df2f2..2d19e6cd 100644 --- a/tests/test_scan_leases_postgres.py +++ b/tests/test_scan_leases_postgres.py @@ -95,18 +95,11 @@ def rearm(self, scan_id: str, owner: str, fencing_token: int) -> None: (owner, fencing_token, scan_id), ) - def evaluation_count(self, scan_id: str) -> int: - with psycopg2.connect(self.dsn) as conn: - with conn.cursor() as cur: - cur.execute("SELECT COUNT(*) FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) - return cur.fetchone()[0] - def cleanup(self) -> None: with psycopg2.connect(self.dsn) as conn: with conn.cursor() as cur: for scan_id in self.scan_ids: cur.execute("DELETE FROM enrichment_jobs WHERE scan_id = %s", (scan_id,)) - cur.execute("DELETE FROM rule_evaluations WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM findings WHERE scan_id = %s", (scan_id,)) cur.execute("DELETE FROM scans WHERE scan_id = %s", (scan_id,)) @@ -118,10 +111,16 @@ def scan_rows() -> ScanRows: rows.cleanup() -def _claim(dsn: str, owner: str) -> dict | None: +def _claim(dsn: str, owner: str, scan_id: str | None = None) -> dict | None: + """Claim a pending scan, restricted to ``scan_id`` when given. + + Tests always pass their own scan: an unrestricted claim takes the globally + oldest pending row, so in a shared database it would otherwise claim + whatever unrelated scan another test admitted first. + """ db = DatabaseManager(dsn) try: - return db.claim_next_pending_scan(owner, 120) + return db.claim_next_pending_scan(owner, 120, scan_id=scan_id) finally: db.close() @@ -135,13 +134,13 @@ def _recover(dsn: str, max_attempts: int = 3) -> int: def test_two_workers_race_to_claim_only_one_scan(scan_rows): - scan_rows.create() + scan_id, _ = scan_rows.create() barrier = threading.Barrier(2) claims: list[dict | None] = [] def claim(owner: str) -> None: barrier.wait() - claims.append(_claim(scan_rows.dsn, owner)) + claims.append(_claim(scan_rows.dsn, owner, scan_id)) threads = [threading.Thread(target=claim, args=(owner,)) for owner in ("worker-a", "worker-b")] for thread in threads: @@ -156,14 +155,14 @@ def claim(owner: str) -> None: def test_active_lease_cannot_be_reclaimed(scan_rows): - scan_rows.create() - assert _claim(scan_rows.dsn, "worker-a") is not None - assert _claim(scan_rows.dsn, "worker-b") is None + scan_id, _ = scan_rows.create() + assert _claim(scan_rows.dsn, "worker-a", scan_id) is not None + assert _claim(scan_rows.dsn, "worker-b", scan_id) is None def test_heartbeat_extends_current_lease_without_changing_owner_or_token(scan_rows): scan_id, _ = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-a") + claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert claim is not None db = DatabaseManager(scan_rows.dsn) @@ -179,12 +178,15 @@ def test_heartbeat_extends_current_lease_without_changing_owner_or_token(scan_ro def test_expired_lease_is_reclaimed_with_new_fencing_token(scan_rows): scan_id, _ = scan_rows.create() - first_claim = _claim(scan_rows.dsn, "worker-a") + first_claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert first_claim is not None scan_rows.expire(scan_id) - assert _recover(scan_rows.dsn) == 1 - second_claim = _claim(scan_rows.dsn, "worker-b") + # Recovery is queue-wide, so assert this scan was recovered rather than + # that it was the only scan recovered. + assert _recover(scan_rows.dsn) >= 1 + assert scan_rows.scan(scan_id)["status"] == "pending" + second_claim = _claim(scan_rows.dsn, "worker-b", scan_id) assert second_claim is not None assert second_claim["lease_owner"] == "worker-b" @@ -193,11 +195,11 @@ def test_expired_lease_is_reclaimed_with_new_fencing_token(scan_rows): def test_stale_worker_cannot_heartbeat_complete_fail_or_write_results(scan_rows): scan_id, subscription_id = scan_rows.create() - first_claim = _claim(scan_rows.dsn, "worker-a") + first_claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert first_claim is not None scan_rows.expire(scan_id) _recover(scan_rows.dsn) - second_claim = _claim(scan_rows.dsn, "worker-b") + second_claim = _claim(scan_rows.dsn, "worker-b", scan_id) assert second_claim is not None stale_db = DatabaseManager(scan_rows.dsn) @@ -226,7 +228,7 @@ def test_stale_worker_cannot_heartbeat_complete_fail_or_write_results(scan_rows) def test_current_owner_completion_persists_results_atomically(scan_rows): scan_id, subscription_id = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-b") + claim = _claim(scan_rows.dsn, "worker-b", scan_id) assert claim is not None db = DatabaseManager(scan_rows.dsn) @@ -243,7 +245,7 @@ def test_current_owner_completion_persists_results_atomically(scan_rows): def test_sql_abort_rolls_back_and_the_connection_remains_usable(scan_rows): scan_id, subscription_id = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-a") + claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert claim is not None broken_result = _result(scan_id, subscription_id) broken_result["findings"][0]["detected_at"] = None @@ -297,43 +299,31 @@ def test_terminated_backend_is_discarded_and_reacquired(scan_rows): def test_expired_restart_work_obeys_attempt_limit(scan_rows): scan_id, _ = scan_rows.create() - assert _claim(scan_rows.dsn, "worker-a") is not None - scan_rows.expire(scan_id) - assert _recover(scan_rows.dsn, max_attempts=3) == 1 - - assert _claim(scan_rows.dsn, "worker-b") is not None - scan_rows.expire(scan_id) - assert _recover(scan_rows.dsn, max_attempts=3) == 1 - - assert _claim(scan_rows.dsn, "worker-c") is not None - scan_rows.expire(scan_id) - assert _recover(scan_rows.dsn, max_attempts=3) == 1 + # Recovery is queue-wide, so assert this scan's own progression rather + # than a global count another test could contribute to. + for owner in ("worker-a", "worker-b", "worker-c"): + assert _claim(scan_rows.dsn, owner, scan_id) is not None + scan_rows.expire(scan_id) + assert _recover(scan_rows.dsn, max_attempts=3) >= 1 assert scan_rows.scan(scan_id)["status"] == "failed" def test_empty_claim_leaves_no_open_transaction(scan_rows): + # A scan_id with no pending row exercises the same "claimed nothing" path + # without depending on the shared queue being empty. db = DatabaseManager(scan_rows.dsn) try: - assert db.claim_next_pending_scan("worker-a", 120) is None + assert db.claim_next_pending_scan("worker-a", 120, scan_id=str(uuid.uuid4())) is None assert db._get_conn().info.transaction_status == extensions.TRANSACTION_STATUS_IDLE finally: db.close() -def test_duplicate_result_delivery_upserts_mutable_fields_and_evaluations(scan_rows): +def test_duplicate_result_delivery_upserts_mutable_fields(scan_rows): scan_id, subscription_id = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-a") + claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert claim is not None result = _result(scan_id, subscription_id) - result["evaluations"] = [ - { - "rule_id": "AZ-LEASE-001", - "resource_id": result["findings"][0]["resource_id"], - "resource_type": "Test/resource", - "status": "FAIL", - "evidence": {"version": 1}, - } - ] db = DatabaseManager(scan_rows.dsn) try: @@ -342,22 +332,21 @@ def test_duplicate_result_delivery_upserts_mutable_fields_and_evaluations(scan_r scan_rows.rearm(scan_id, "worker-a", claim["fencing_token"]) result["findings"][0]["description"] = "Updated presentation text" result["findings"][0]["severity"] = "CRITICAL" - result["evaluations"][0]["evidence"] = {"version": 2} db.save_scan(result, "worker-a", claim["fencing_token"]) persisted = db.get_findings({"scan_id": scan_id}) finally: db.close() + # Replay updates the same row in place instead of duplicating it. assert len(persisted) == 1 assert persisted[0]["id"] == first_id assert persisted[0]["description"] == "Updated presentation text" assert persisted[0]["severity"] == "CRITICAL" - assert scan_rows.evaluation_count(scan_id) == 1 def test_distinct_finding_discriminators_preserve_multiple_violations(scan_rows): scan_id, subscription_id = scan_rows.create() - claim = _claim(scan_rows.dsn, "worker-a") + claim = _claim(scan_rows.dsn, "worker-a", scan_id) assert claim is not None result = _result(scan_id, subscription_id) duplicate_scope = dict(result["findings"][0]) diff --git a/tests/test_scans_enrich.py b/tests/test_scans_enrich.py index 80addb6d..ae1dcd3a 100644 --- a/tests/test_scans_enrich.py +++ b/tests/test_scans_enrich.py @@ -18,24 +18,54 @@ def _mock_db(current_scan=None, findings=None): def test_enrich_returns_202_and_enqueues_durable_job(client, auth_headers): scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} db = _mock_db(current_scan=scan, findings=[{"id": 1, "rule_id": "AZ-STOR-001"}]) - db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "pending"}, True) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "pending"}, "created") with patch.object(scans_route, "_get_db", return_value=db): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 202 - assert resp.get_json()["status"] == "PENDING" + body = resp.get_json() + assert body["outcome"] == "created" + assert body["status"] == "pending" + assert body["job_id"] == _SCAN_ID db.enqueue_enrichment_job.assert_called_once_with(_SCAN_ID) def test_enrich_reuses_existing_durable_job(client, auth_headers): scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} db = _mock_db(current_scan=scan, findings=[{"id": 1}]) - db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "running"}, False) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "running"}, "active") with patch.object(scans_route, "_get_db", return_value=db): resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) assert resp.status_code == 202 - assert resp.get_json()["status"] == "running" + body = resp.get_json() + assert body["outcome"] == "active" + assert body["status"] == "running" + + +def test_enrich_requeues_a_terminally_failed_job(client, auth_headers): + """A failed job is the case a re-POST exists to recover from.""" + scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "FAILED"} + db = _mock_db(current_scan=scan, findings=[{"id": 1}]) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "pending"}, "requeued") + with patch.object(scans_route, "_get_db", return_value=db): + resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) + assert resp.status_code == 202 + body = resp.get_json() + assert body["outcome"] == "requeued" + assert body["status"] == "pending" + assert "requeued" in body["message"] + + +def test_enrich_reports_an_already_completed_job_without_restarting_it(client, auth_headers): + # The scan header still reads PENDING, so the job row is the authority. + scan = {"scan_id": _SCAN_ID, "cve_enrichment_status": "PENDING"} + db = _mock_db(current_scan=scan, findings=[{"id": 1}]) + db.enqueue_enrichment_job.return_value = ({"job_id": _SCAN_ID, "status": "completed"}, "completed") + with patch.object(scans_route, "_get_db", return_value=db): + resp = client.post(f"/api/scans/{_SCAN_ID}/enrich", headers=auth_headers) + assert resp.status_code == 200 + assert resp.get_json()["outcome"] == "completed" def test_enrich_already_completed_returns_200(client, auth_headers): diff --git a/tests/test_worker.py b/tests/test_worker.py index f973b15f..8bbcae94 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -210,6 +210,68 @@ def test_worker_sleeps_when_no_scans_pending(self, mock_sleep, mock_env, mock_db mock_sleep.assert_called_with(POLL_INTERVAL_SECONDS) + @patch("scanner.worker.process_enrichment_job") + @patch("scanner.worker.DatabaseManager") + @patch("scanner.worker.ScanEngine") + @patch("scanner.worker.os.environ.get") + @patch("scanner.worker.time.sleep") + @patch("scanner.worker.LeaseHeartbeat") + def test_enrichment_backlog_does_not_starve_pending_scans( + self, mock_heartbeat_class, mock_sleep, mock_env, mock_engine_class, mock_db_class, mock_process + ): + """With both queues permanently backlogged, each still gets served. + + Draining enrichment and restarting the loop would mean a scan is never + claimed while enrichment work keeps arriving. + """ + mock_env.return_value = self.mock_db_url + mock_db = mock_db_class.return_value + mock_engine_class.return_value.run_scan.return_value = { + "scan_id": self.scan_id, + "subscription_id": self.subscription_id, + "findings": [], + "started_at": "2026-06-05T12:00:00Z", + } + mock_heartbeat_class.return_value.lost.is_set.return_value = False + + # Both queues always have work; stop after two full iterations. + mock_db.recover_stale_scans.side_effect = [None, None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = {"job_id": "job-1", "scan_id": self.scan_id} + mock_db.claim_next_pending_scan.return_value = { + "scan_id": self.scan_id, + "subscription_id": self.subscription_id, + "fencing_token": 1, + } + + with self.assertRaises(StopWorker): + run_worker() + + # Each iteration served one job from each queue -- neither starved. + self.assertEqual(mock_process.call_count, 2) + self.assertEqual(mock_db.save_scan.call_count, 2) + # A busy worker never idles. + mock_sleep.assert_not_called() + + @patch("scanner.worker.process_enrichment_job") + @patch("scanner.worker.DatabaseManager") + @patch("scanner.worker.os.environ.get") + @patch("scanner.worker.time.sleep") + def test_enrichment_only_backlog_polls_again_without_idling( + self, mock_sleep, mock_env, mock_db_class, mock_process + ): + """Enrichment work must not be paced by the empty-queue poll interval.""" + mock_env.return_value = self.mock_db_url + mock_db = mock_db_class.return_value + mock_db.recover_stale_scans.side_effect = [None, StopWorker()] + mock_db.claim_next_enrichment_job.return_value = {"job_id": "job-1", "scan_id": self.scan_id} + mock_db.claim_next_pending_scan.return_value = None + + with self.assertRaises(StopWorker): + run_worker() + + mock_process.assert_called_once() + mock_sleep.assert_not_called() + if __name__ == "__main__": unittest.main() From 9da3d3e13c7b768e9873daca7f08965704d4d53b Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Wed, 2 Sep 2026 01:35:55 +0100 Subject: [PATCH 11/11] test(worker): cover the lease/heartbeat interval clamp `lease_configuration()` guarantees the heartbeat interval stays strictly shorter than the lease - a worker that heartbeats no more often than its lease expires would lose its own claim mid-scan and have its results fenced out. .env.example documents that constraint, but nothing tested it. Covers the defaults, valid overrides, heartbeat == lease, heartbeat > lease, a lease small enough that `lease // 3` would floor to a zero-second heartbeat, and malformed/non-positive values falling back to the defaults. Verified the tests fail when the clamp is removed (3 failures) and pass when it is restored. Signed-off-by: Shaurya K Sharma --- tests/test_worker.py | 52 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/tests/test_worker.py b/tests/test_worker.py index 8bbcae94..2d75c948 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -10,7 +10,14 @@ import unittest from unittest.mock import ANY, patch from api.models.finding import LostLease -from scanner.worker import LeaseHeartbeat, POLL_INTERVAL_SECONDS, run_worker +from scanner.worker import ( + DEFAULT_HEARTBEAT_SECONDS, + DEFAULT_LEASE_SECONDS, + LeaseHeartbeat, + POLL_INTERVAL_SECONDS, + lease_configuration, + run_worker, +) import uuid @@ -273,5 +280,48 @@ def test_enrichment_only_backlog_polls_again_without_idling( mock_sleep.assert_not_called() +class TestLeaseConfiguration(unittest.TestCase): + """A worker that heartbeats no more often than its lease expires would + lose its own claim mid-scan, so the configuration is clamped rather than + trusted. .env.example documents this; these tests hold it to it.""" + + def _configure(self, **env): + with patch.dict("scanner.worker.os.environ", env, clear=True): + return lease_configuration() + + def test_defaults_keep_the_heartbeat_shorter_than_the_lease(self): + lease, heartbeat = self._configure() + self.assertEqual((lease, heartbeat), (DEFAULT_LEASE_SECONDS, DEFAULT_HEARTBEAT_SECONDS)) + self.assertLess(heartbeat, lease) + + def test_valid_overrides_are_used_as_given(self): + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS="600", SCAN_HEARTBEAT_SECONDS="60") + self.assertEqual((lease, heartbeat), (600, 60)) + + def test_heartbeat_equal_to_the_lease_is_clamped_below_it(self): + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS="300", SCAN_HEARTBEAT_SECONDS="300") + self.assertEqual(lease, 300) + self.assertEqual(heartbeat, 100) + self.assertLess(heartbeat, lease) + + def test_heartbeat_longer_than_the_lease_is_clamped_below_it(self): + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS="300", SCAN_HEARTBEAT_SECONDS="9000") + self.assertEqual((lease, heartbeat), (300, 100)) + self.assertLess(heartbeat, lease) + + def test_a_tiny_lease_still_yields_a_positive_heartbeat(self): + # lease // 3 would floor to 0 and make the heartbeat thread spin. + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS="2", SCAN_HEARTBEAT_SECONDS="2") + self.assertEqual((lease, heartbeat), (2, 1)) + self.assertGreater(heartbeat, 0) + + def test_malformed_or_non_positive_values_fall_back_to_defaults(self): + for bad in ("not-a-number", "0", "-30", ""): + with self.subTest(value=bad): + lease, heartbeat = self._configure(SCAN_LEASE_SECONDS=bad, SCAN_HEARTBEAT_SECONDS=bad) + self.assertEqual((lease, heartbeat), (DEFAULT_LEASE_SECONDS, DEFAULT_HEARTBEAT_SECONDS)) + self.assertLess(heartbeat, lease) + + if __name__ == "__main__": unittest.main()