feat(package-vulnerability-scanner): query vulnerabilities in parallel chunks

extensions/package-vulnerability-scanner/main.py

@@ -120,6 +120,14 @@ def _upstream_error(message: str, exc: Exception) -> HTTPException:
 # deployment doesn't produce an unwieldy payload.
 PPM_QUERY_CHUNK = 100
 
+# How many times to try a Package Manager request before giving up, so a
+# transient network or server blip doesn't fail the whole scan.
+PPM_MAX_ATTEMPTS = 3
+
+# Bound how many Package Manager requests run at once so a large scan does not
+# overload the public service.
+PPM_QUERY_CONCURRENCY = 5
+
 # How many times to try a Connect API call before giving up, so a transient 5xx
 # (e.g. a 504 while Connect is under load) doesn't fail the whole scan.
 CONTENT_MAX_ATTEMPTS = 3
@@ -177,6 +185,21 @@ async def _fetch_with_retry(fn, timeout=None):
             raise
 
 
+# Like asyncio.gather, but if one coroutine raises, cancel the rest and wait for
+# them to actually stop before propagating, instead of leaving them running
+# in the background with their results discarded (and, for callers that close a
+# resource like an httpx.AsyncClient right after gathering, out from under them).
+async def _gather_and_cancel_on_error(*aws):
+    tasks = [asyncio.ensure_future(aw) for aw in aws]
+    try:
+        return await asyncio.gather(*tasks)
+    except Exception:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        raise
+
+
 @app.get("/api/content")
 async def search_content(request: Request, show_all: bool = False):
     try:
@@ -262,7 +285,7 @@ async def fetch_one(guid: str) -> tuple[str, dict]:
                     "error": "Couldn't read this content's packages.",
                 }
 
-    pairs = await asyncio.gather(*(fetch_one(guid) for guid in guids))
+    pairs = await _gather_and_cancel_on_error(*(fetch_one(guid) for guid in guids))
     return dict(pairs)
 
 
@@ -304,6 +327,7 @@ async def get_vulnerabilities(installed: InstalledPackages):
     # which only flags packages whose latest version is vulnerable and so misses
     # older deployed versions that were patched later.
     results = {}
+    ppm_semaphore = asyncio.Semaphore(PPM_QUERY_CONCURRENCY)
 
     async def fetch_repo_vulns(repo, specifiers):
         # name -> {vuln id -> vuln}; the same package may be requested at
@@ -312,29 +336,67 @@ async def fetch_repo_vulns(repo, specifiers):
         specs = sorted(set(specifiers))
         if not specs:
             return repo, {}
-        async with httpx.AsyncClient() as http:
-            for start in range(0, len(specs), PPM_QUERY_CHUNK):
-                payload = {
-                    "repo": repo,
-                    "names": specs[start : start + PPM_QUERY_CHUNK],
-                    "omit_downloads": True,
-                    "omit_dependencies": True,
-                }
-                response = await http.post(PPM_URL, json=payload)
-                response.raise_for_status()
-                for line in response.text.strip().split("\n"):
-                    if not line:
-                        continue
-                    found = json.loads(line)
-                    for vuln in found.get("vulns") or []:
-                        merged.setdefault(found["name"], {})[vuln["id"]] = vuln
+        chunks = [
+            specs[start : start + PPM_QUERY_CHUNK]
+            for start in range(0, len(specs), PPM_QUERY_CHUNK)
+        ]
+
+        async def fetch_chunk(http, names):
+            payload = {
+                "repo": repo,
+                "names": names,
+                "omit_downloads": True,
+                "omit_dependencies": True,
+            }
+            # Retry a transient failure (network error or 5xx) with a short
+            # backoff so one blip against the public service doesn't fail the
+            # whole scan. A 4xx or a persistent failure still raises, so the scan
+            # fails loudly rather than silently under-reporting vulnerabilities.
+            async with ppm_semaphore:
+                for attempt in range(PPM_MAX_ATTEMPTS):
+                    try:
+                        response = await http.post(PPM_URL, json=payload)
+                        response.raise_for_status()
+                        return response.text
+                    except httpx.HTTPStatusError as e:
+                        if (
+                            e.response.status_code < 500
+                            or attempt == PPM_MAX_ATTEMPTS - 1
+                        ):
+                            raise
+                    except httpx.TransportError:
+                        if attempt == PPM_MAX_ATTEMPTS - 1:
+                            raise
+                    await asyncio.sleep(0.5 * (attempt + 1))
+
+        # Query the chunks in parallel so a large repo doesn't pay one round-trip
+        # after another. The timeout is generous enough for a full chunk yet still
+        # bounds a stalled request so the scan can't hang.
+        async with httpx.AsyncClient(timeout=30.0) as http:
+            texts = await _gather_and_cancel_on_error(
+                *(fetch_chunk(http, chunk) for chunk in chunks)
+            )
+        for text in texts:
+            for line in text.strip().split("\n"):
+                if not line:
+                    continue
+                found = json.loads(line)
+                for vuln in found.get("vulns") or []:
+                    merged.setdefault(found["name"], {})[vuln["id"]] = vuln
         return repo, {name: list(v.values()) for name, v in merged.items()}
 
-    for repo, data in await asyncio.gather(
-        fetch_repo_vulns("pypi", installed.pypi),
-        fetch_repo_vulns("cran", installed.cran),
-    ):
-        results[repo] = data
+    # A failure (after retries) is reported so the scan fails loudly rather than
+    # silently under-reporting vulnerabilities, matching the packages endpoint.
+    try:
+        for repo, data in await _gather_and_cancel_on_error(
+            fetch_repo_vulns("pypi", installed.pypi),
+            fetch_repo_vulns("cran", installed.cran),
+        ):
+            results[repo] = data
+    except Exception as e:
+        raise _upstream_error(
+            "Couldn't fetch vulnerabilities from Package Manager.", e
+        )
 
     return results
 

extensions/package-vulnerability-scanner/manifest.json

@@ -20,17 +20,17 @@
   },
   "packages": {},
   "files": {
-    "dist/assets/index-Bd7WoAqv.css": {
-      "checksum": "980b7708188eee5f84707dcc6cf3b01d"
+    "dist/assets/index-CKUnK_gI.css": {
+      "checksum": "910900dfbf81f7ac4caf125b844a495c"
     },
-    "dist/assets/index-BEWOoybn.js": {
-      "checksum": "da5ddbd1d9affbb66e1449a6f1e9392b"
+    "dist/assets/index-CgkWwSXn.js": {
+      "checksum": "31a12fd9c6e9d80d692a587a0bc20ec5"
     },
     "dist/index.html": {
-      "checksum": "1a6d2fe30d21376a0668ad781de538e6"
+      "checksum": "d8f8de453602d81829635911a04d612d"
     },
     "main.py": {
-      "checksum": "8697dad866d91d30ec56cf62707bc817"
+      "checksum": "ea913b00aaac84531ee46021969f4eaa"
     },
     "requirements.txt": {
       "checksum": "9fc5cc5fb559eded1f323793b73e82c5"

extensions/package-vulnerability-scanner/src/components/ContentList.vue

@@ -3,11 +3,12 @@ import { computed, ref, watch } from "vue";
 import { storeToRefs } from "pinia";
 
 import { usePackagesStore } from "../stores/packages";
-import { useVulnsStore, type InstalledPackages } from "../stores/vulns";
+import { useVulnsStore } from "../stores/vulns";
 import { useContentStore } from "../stores/content";
 import { useScannerStore } from "../stores/scanner";
 import { useUserStore } from "../stores/user";
 import type { User } from "../stores/user";
+import { collectInstalledPackages } from "../lib/collectInstalledPackages";
 import StatusMessage from "./ui/StatusMessage.vue";
 import SkeletonText from "./ui/SkeletonText.vue";
 import BadgeTabs, { type Tab } from "./ui/BadgeTabs.vue";
@@ -76,27 +77,15 @@ async function fetchPackagesInBatches(batchSize = 3) {
   // Packages are loaded; look up vulnerabilities for the exact installed
   // versions. Skip on a plain remount where nothing new was fetched.
   if (contentToFetch.length > 0 || !vulnStore.isFetched) {
-    await vulnStore.fetchVulns(collectInstalledPackages());
+    await vulnStore.fetchVulns(
+      collectInstalledPackages(
+        contentStore.contentList,
+        packagesStore.contentItems,
+      ),
+    );
   }
 }
 
-// Gather unique installed packages as "name==version", grouped by repo.
-function collectInstalledPackages(): InstalledPackages {
-  const pypi = new Set<string>();
-  const cran = new Set<string>();
-  for (const item of Object.values(packagesStore.contentItems)) {
-    for (const pkg of item.packages) {
-      const language = pkg.language.toLowerCase();
-      if (language === "python") {
-        pypi.add(`${pkg.name}==${pkg.version}`);
-      } else if (language === "r") {
-        cran.add(`${pkg.name}==${pkg.version}`);
-      }
-    }
-  }
-  return { pypi: [...pypi], cran: [...cran] };
-}
-
 fetchPackagesInBatches();
 
 const tabs = computed<Tab[]>(() => {

extensions/package-vulnerability-scanner/src/lib/collectInstalledPackages.test.ts

@@ -0,0 +1,77 @@
+import { describe, it, expect } from "vitest";
+
+import { collectInstalledPackages } from "./collectInstalledPackages";
+import type { ContentListItem } from "../stores/content";
+import type { ContentPackages, Package } from "../stores/packages";
+
+function content(guid: string): ContentListItem {
+  return { guid, title: guid };
+}
+
+function pkg(name: string, version: string, language: string): Package {
+  return { name, version, language, hash: null };
+}
+
+function item(guid: string, packages: Package[]): ContentPackages {
+  return {
+    guid,
+    packages,
+    isLoading: false,
+    error: null,
+    isFetched: true,
+    lastFetchTime: null,
+  };
+}
+
+describe("collectInstalledPackages", () => {
+  it("groups python into pypi and r into cran as name==version", () => {
+    const result = collectInstalledPackages([content("g1")], {
+      g1: item("g1", [pkg("flask", "2.0", "python"), pkg("dplyr", "1.1", "r")]),
+    });
+
+    expect(result.pypi).toEqual(["flask==2.0"]);
+    expect(result.cran).toEqual(["dplyr==1.1"]);
+  });
+
+  it("dedupes the same specifier across content items", () => {
+    const result = collectInstalledPackages([content("g1"), content("g2")], {
+      g1: item("g1", [pkg("flask", "2.0", "python")]),
+      g2: item("g2", [pkg("flask", "2.0", "python")]),
+    });
+
+    expect(result.pypi).toEqual(["flask==2.0"]);
+  });
+
+  it("treats language case-insensitively", () => {
+    const result = collectInstalledPackages([content("g1")], {
+      g1: item("g1", [pkg("flask", "2.0", "Python"), pkg("dplyr", "1.1", "R")]),
+    });
+
+    expect(result.pypi).toEqual(["flask==2.0"]);
+    expect(result.cran).toEqual(["dplyr==1.1"]);
+  });
+
+  it("ignores languages that are neither python nor r", () => {
+    const result = collectInstalledPackages([content("g1")], {
+      g1: item("g1", [pkg("some-lib", "1.0", "javascript")]),
+    });
+
+    expect(result).toEqual({ pypi: [], cran: [] });
+  });
+
+  it("ignores cached packages for content not in the visible list", () => {
+    // g2 has packages cached but is not in the list, so its packages must not
+    // leak into the query.
+    const result = collectInstalledPackages([content("g1")], {
+      g1: item("g1", [pkg("flask", "2.0", "python")]),
+      g2: item("g2", [pkg("django", "3.0", "python")]),
+    });
+
+    expect(result.pypi).toEqual(["flask==2.0"]);
+  });
+
+  it("skips content in the list that has no packages entry yet", () => {
+    const result = collectInstalledPackages([content("g1")], {});
+    expect(result).toEqual({ pypi: [], cran: [] });
+  });
+});

extensions/package-vulnerability-scanner/src/lib/collectInstalledPackages.ts

@@ -0,0 +1,27 @@
+import type { ContentListItem } from "../stores/content";
+import type { ContentPackages } from "../stores/packages";
+import type { InstalledPackages } from "../stores/vulns";
+
+// Gather unique installed packages as "name==version", grouped by repo, across
+// the content currently in view. Content not in the list (cached packages for
+// hidden content) is ignored, so the vuln lookup matches what's on screen.
+export function collectInstalledPackages(
+  contentList: ContentListItem[],
+  contentItems: Record<string, ContentPackages>,
+): InstalledPackages {
+  const pypi = new Set<string>();
+  const cran = new Set<string>();
+  for (const content of contentList) {
+    const item = contentItems[content.guid];
+    if (!item) continue;
+    for (const pkg of item.packages) {
+      const language = pkg.language.toLowerCase();
+      if (language === "python") {
+        pypi.add(`${pkg.name}==${pkg.version}`);
+      } else if (language === "r") {
+        cran.add(`${pkg.name}==${pkg.version}`);
+      }
+    }
+  }
+  return { pypi: [...pypi], cran: [...cran] };
+}