From cb25dbb4973f8c6d720d99861e5823cd6e2260c7 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 18:29:08 +0100 Subject: [PATCH 1/9] 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 dc71df514b14a2388ba6f784dc6e0aecc0cd2b35 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 18:48:46 +0100 Subject: [PATCH 2/9] 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 4ce0b004cd99ecf2d6154cd43225759e1f47fa3c Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:25:03 +0100 Subject: [PATCH 3/9] 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 75d61e99c5be3e44a6e3eed781f8fd46c2aa145e Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:28:24 +0100 Subject: [PATCH 4/9] 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 9443914ff0833a4c13799b64a90da64a6a44e0cf Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:37:45 +0100 Subject: [PATCH 5/9] 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 bfa2d5150e5c2b6c6fa3a290b814830d75331040 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 19:41:13 +0100 Subject: [PATCH 6/9] 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 ca6a9acfeaf8e84bf784f62b9a590aacd6aba0fb Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 20:01:53 +0100 Subject: [PATCH 7/9] 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 259aadde179bf9d46019434f51a8cffe16282847 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 20:09:16 +0100 Subject: [PATCH 8/9] 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 3b35e369ef17cb0e775bfd45a3f72543e542ac19 Mon Sep 17 00:00:00 2001 From: Shaurya K Sharma Date: Sat, 29 Aug 2026 20:32:32 +0100 Subject: [PATCH 9/9] 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)