skillhub/scanner/skillhub_scanner_app.py

"""Runtime safeguards around the upstream Cisco Skill Scanner ASGI application."""

import asyncio
import logging
import os
import shutil
import tempfile
from pathlib import Path
from typing import NoReturn

from fastapi import Request
from fastapi.responses import JSONResponse
from skill_scanner.api.api import app


_MAX_CONCURRENT_SCANS = max(1, int(os.getenv("SKILLHUB_SCANNER_MAX_CONCURRENT_SCANS", "1")))
_HARD_TIMEOUT_SECONDS = max(1, int(os.getenv("SKILLHUB_SCANNER_HARD_TIMEOUT_SECONDS", "930")))
_active_scans = 0
_active_scans_guard = asyncio.Lock()
_SCAN_PATHS = {"/scan", "/scan-upload"}
_log = logging.getLogger(__name__)


def _cleanup_stale_scan_directories(temp_root: Path | None = None) -> None:
    """Remove incomplete upstream extraction directories left by a process restart."""
    root = temp_root or Path(tempfile.gettempdir())
    for candidate in root.glob("skill_scanner_*"):
        if not candidate.is_dir():
            continue
        try:
            shutil.rmtree(candidate)
        except OSError as error:
            _log.warning("Could not remove stale scanner directory %s: %s", candidate, error)


def _restart_after_hard_timeout(request_path: str) -> NoReturn:
    """Terminate the single-scan worker so the container runtime can recover it."""
    _log.critical(
        "Security scan exceeded the %s second hard timeout: path=%s; restarting scanner",
        _HARD_TIMEOUT_SECONDS,
        request_path,
    )
    logging.shutdown()
    os._exit(124)


async def _await_scan_until(scan_task: asyncio.Task, deadline: float, request_path: str):
    remaining = deadline - asyncio.get_running_loop().time()
    if remaining <= 0:
        _restart_after_hard_timeout(request_path)
    try:
        return await asyncio.wait_for(asyncio.shield(scan_task), timeout=remaining)
    except asyncio.TimeoutError:
        _restart_after_hard_timeout(request_path)


app.router.add_event_handler("startup", _cleanup_stale_scan_directories)


@app.middleware("http")
async def limit_concurrent_scans(request: Request, call_next):
    """Reject excess scan work so timed-out client retries cannot multiply memory use."""
    global _active_scans
    if request.method != "POST" or request.url.path not in _SCAN_PATHS:
        return await call_next(request)

    async with _active_scans_guard:
        if _active_scans >= _MAX_CONCURRENT_SCANS:
            return JSONResponse(
                status_code=503,
                content={"detail": "Scanner is busy; retry later"},
                headers={"Retry-After": "30"},
            )
        _active_scans += 1

    scan_task = asyncio.create_task(call_next(request))
    deadline = asyncio.get_running_loop().time() + _HARD_TIMEOUT_SECONDS
    try:
        # Keep the capacity slot until upstream work really ends, even if the HTTP client
        # disconnects while the scanner's worker thread is still running.
        return await _await_scan_until(scan_task, deadline, request.url.path)
    except asyncio.CancelledError:
        await _await_scan_until(scan_task, deadline, request.url.path)
        raise
    finally:
        async with _active_scans_guard:
            _active_scans -= 1