Use sdk for batchScan

Author: jdalton

.dep-stats.json

@@ -9,7 +9,7 @@
     "@socketregistry/packageurl-js": "1.0.4",
     "@socketsecurity/config": "2.1.3",
     "@socketsecurity/registry": "1.0.133",
-    "@socketsecurity/sdk": "1.4.11",
+    "@socketsecurity/sdk": "1.4.12",
     "blessed": "0.1.81",
     "blessed-contrib": "4.11.0",
     "browserslist": "4.24.4",

src/commands/package/fetch-purls-shallow-score.ts

@@ -12,16 +12,17 @@ import type {
 export async function fetchPurlsShallowScore(
   purls: string[]
 ): Promise<SocketSdkReturnType<'batchPackageFetch'>> {
-  const socketSdk = await setupSdk(getPublicToken())
+  logger.error(
+    `Requesting shallow score data for ${purls.length} package urls (purl): ${purls.join(', ')}`
+  )
 
   // Lazily access constants.spinner.
   const { spinner } = constants
 
-  logger.error(
-    `Requesting shallow score data for ${purls.length} package urls (purl): ${purls.join(', ')}`
-  )
   spinner.start(`Requesting data ...`)
 
+  const socketSdk = await setupSdk(getPublicToken())
+
   const result: Awaited<SocketSdkResultType<'batchPackageFetch'>> =
     await handleApiCall(
       socketSdk.batchPackageFetch(

src/utils/alert/artifact.ts

@@ -1,78 +1,7 @@
-import events from 'node:events'
-import https from 'node:https'
-import readline from 'node:readline'
-
 import constants from '../../constants'
-import { getPublicToken } from '../sdk'
 
 import type { Remap } from '@socketsecurity/registry/lib/objects'
-import type { IncomingMessage } from 'node:http'
-
-export type CveAlertType = 'cve' | 'mediumCVE' | 'mildCVE' | 'criticalCVE'
-
-export type SocketArtifactAlert = {
-  key: string
-  type: string
-  severity: string
-  category: string
-  action?: string | undefined
-  actionPolicyIndex?: number | undefined
-  file?: string | undefined
-  props?: any | undefined
-  start?: number | undefined
-  end?: number | undefined
-}
-
-export type SocketArtifact = {
-  type: string
-  name: string
-  namespace?: string | undefined
-  version?: string | undefined
-  subpath?: string | undefined
-  release?: string | undefined
-  id?: string | undefined
-  author?: string[]
-  license?: string | undefined
-  licenseDetails?: Array<{
-    spdxDisj: string
-    provenance: string
-    filepath: string
-    match_strength: number
-  }>
-  licenseAttrib?: Array<{
-    attribText: string
-    attribData: Array<{
-      purl: string
-      foundInFilepath: string
-      spdxExpr: string
-      foundAuthors: string[]
-    }>
-  }>
-  score?: {
-    supplyChain: number
-    quality: number
-    maintenance: number
-    vulnerability: number
-    license: number
-    overall: number
-  }
-  alerts?: SocketArtifactAlert[]
-  size?: number | undefined
-  batchIndex?: number | undefined
-}
-
-export type CompactSocketArtifactAlert = Remap<
-  Omit<
-    SocketArtifactAlert,
-    'action' | 'actionPolicyIndex' | 'category' | 'end' | 'file' | 'start'
-  >
->
-
-export type CompactSocketArtifact = Remap<
-  Omit<SocketArtifact, 'alerts' | 'batchIndex' | 'size'> & {
-    alerts: CompactSocketArtifactAlert[]
-  }
->
+import type { components } from '@socketsecurity/sdk/types/api'
 
 export type ArtifactAlertCve = Remap<
   Omit<CompactSocketArtifactAlert, 'type'> & {
@@ -97,136 +26,39 @@ export type ArtifactAlertUpgrade = Remap<
   }
 >
 
+export type CveAlertType = 'cve' | 'mediumCVE' | 'mildCVE' | 'criticalCVE'
+
+export type CompactSocketArtifactAlert = Remap<
+  Omit<
+    SocketArtifactAlert,
+    'action' | 'actionPolicyIndex' | 'category' | 'end' | 'file' | 'start'
+  >
+>
+
+export type CompactSocketArtifact = Remap<
+  Omit<SocketArtifact, 'alerts' | 'batchIndex' | 'size'> & {
+    alerts: CompactSocketArtifactAlert[]
+  }
+>
+
+export type SocketArtifact = components['schemas']['SocketArtifact']
+
+export type SocketArtifactAlert = Remap<
+  Omit<components['schemas']['SocketAlert'], 'props'> & {
+    props?: any | undefined
+  }
+>
+
 const {
   ALERT_TYPE_CRITICAL_CVE,
   ALERT_TYPE_CVE,
   ALERT_TYPE_MEDIUM_CVE,
   ALERT_TYPE_MILD_CVE,
   ALERT_TYPE_SOCKET_UPGRADE_AVAILABLE,
-  API_V0_URL,
   CVE_ALERT_PROPS_FIRST_PATCHED_VERSION_IDENTIFIER,
-  CVE_ALERT_PROPS_VULNERABLE_VERSION_RANGE,
-  abortSignal
+  CVE_ALERT_PROPS_VULNERABLE_VERSION_RANGE
 } = constants
 
-async function* createBatchGenerator(
-  chunk: string[]
-): AsyncGenerator<CompactSocketArtifact> {
-  // Adds the first 'abort' listener to abortSignal.
-  const req = https
-    .request(
-      `${API_V0_URL}purl?${new URLSearchParams([
-        ['alerts', 'true'],
-        ['compact', 'true']
-      ])}`,
-      {
-        method: 'POST',
-        headers: {
-          Authorization: `Basic ${btoa(`${getPublicToken()}:`)}`
-        }
-        // TODO: Fix to not abort process on network abort.
-        // signal: abortSignal
-      }
-    )
-    .end(
-      JSON.stringify({
-        components: chunk.map(id => ({ purl: `pkg:npm/${id}` }))
-      })
-    )
-  // Adds the second 'abort' listener to abortSignal.
-  const { 0: res } = (await events.once(req, 'response', {
-    signal: abortSignal
-  })) as [IncomingMessage]
-  const ok = res.statusCode! >= 200 && res.statusCode! <= 299
-  if (!ok) {
-    throw new Error(`Socket API Error: ${res.statusCode}`)
-  }
-  const rli = readline.createInterface({
-    input: res,
-    crlfDelay: Infinity,
-    signal: abortSignal
-  })
-  for await (const line of rli) {
-    yield JSON.parse(line) as CompactSocketArtifact
-  }
-}
-
-export async function* batchScan(
-  pkgIds: string[],
-  concurrencyLimit = 50
-): AsyncGenerator<CompactSocketArtifact> {
-  type GeneratorStep = {
-    generator: AsyncGenerator<CompactSocketArtifact>
-    iteratorResult: IteratorResult<CompactSocketArtifact>
-  }
-  type GeneratorEntry = {
-    generator: AsyncGenerator<CompactSocketArtifact>
-    promise: Promise<GeneratorStep>
-  }
-  type ResolveFn = (value: GeneratorStep) => void
-
-  // The createBatchGenerator method will add 2 'abort' event listeners to
-  // abortSignal so we multiply the concurrencyLimit by 2.
-  const neededMaxListeners = concurrencyLimit * 2
-  // Increase abortSignal max listeners count to avoid Node's MaxListenersExceededWarning.
-  const oldAbortSignalMaxListeners = events.getMaxListeners(abortSignal)
-  let abortSignalMaxListeners = oldAbortSignalMaxListeners
-  if (oldAbortSignalMaxListeners < neededMaxListeners) {
-    abortSignalMaxListeners = oldAbortSignalMaxListeners + neededMaxListeners
-    events.setMaxListeners(abortSignalMaxListeners, abortSignal)
-  }
-  const { length: pkgIdsCount } = pkgIds
-  const running: GeneratorEntry[] = []
-  let index = 0
-  const enqueueGen = () => {
-    if (index >= pkgIdsCount) {
-      // No more work to do.
-      return
-    }
-    const chunk = pkgIds.slice(index, index + 25)
-    index += 25
-    const generator = createBatchGenerator(chunk)
-    continueGen(generator)
-  }
-  const continueGen = (generator: AsyncGenerator<CompactSocketArtifact>) => {
-    let resolveFn: ResolveFn
-    running.push({
-      generator,
-      promise: new Promise<GeneratorStep>(resolve => (resolveFn = resolve))
-    })
-    void generator
-      .next()
-      .then(res => resolveFn!({ generator, iteratorResult: res }))
-  }
-  // Start initial batch of generators.
-  while (running.length < concurrencyLimit && index < pkgIdsCount) {
-    enqueueGen()
-  }
-  while (running.length > 0) {
-    // eslint-disable-next-line no-await-in-loop
-    const { generator, iteratorResult } = await Promise.race(
-      running.map(entry => entry.promise)
-    )
-    // Remove generator.
-    running.splice(
-      running.findIndex(entry => entry.generator === generator),
-      1
-    )
-    if (iteratorResult.done) {
-      // Start a new generator if available.
-      enqueueGen()
-    } else {
-      yield iteratorResult.value
-      // Keep fetching values from this generator.
-      continueGen(generator)
-    }
-  }
-  // Reset abortSignal max listeners count.
-  if (abortSignalMaxListeners > oldAbortSignalMaxListeners) {
-    events.setMaxListeners(oldAbortSignalMaxListeners, abortSignal)
-  }
-}
-
 export function isArtifactAlertCve(
   alert: CompactSocketArtifactAlert
 ): alert is ArtifactAlertCve {

src/utils/api.ts

@@ -6,8 +6,8 @@ import { logger } from '@socketsecurity/registry/lib/logger'
 import { isNonEmptyString } from '@socketsecurity/registry/lib/strings'
 
 import { AuthError } from './errors'
-import constants from '../constants'
 import { getSetting } from './settings'
+import constants from '../constants'
 
 import type {
   SocketSdkErrorType,

src/utils/lockfile/package-lock-json.ts

@@ -11,7 +11,8 @@ import constants from '../../constants'
 import { SafeArborist } from '../../shadow/npm/arborist/lib/arborist'
 import { DiffAction } from '../../shadow/npm/arborist/lib/arborist/types'
 import { Edge } from '../../shadow/npm/arborist/lib/edge'
-import { batchScan } from '../alert/artifact'
+import { getPublicToken, setupSdk } from '../../utils/sdk'
+import { CompactSocketArtifact } from '../alert/artifact'
 import {
   type AlertsByPkgId,
   addArtifactToAlertsMap
@@ -244,12 +245,30 @@ export async function getAlertsMapFromArborist(
       })
     )
   }
+
+  const socketSdk = await setupSdk(getPublicToken())
+
   const toAlertsMapOptions = {
     overrides,
     ...options
   }
-  for await (const artifact of batchScan(pkgIds)) {
-    await addArtifactToAlertsMap(artifact, alertsByPkgId, toAlertsMapOptions)
+
+  for await (const batchPackageFetchResult of socketSdk.batchPackageStream(
+    {
+      alerts: 'true',
+      compact: 'true'
+    },
+    {
+      components: pkgIds.map(id => ({ purl: `pkg:npm/${id}` }))
+    }
+  )) {
+    if (batchPackageFetchResult.success) {
+      await addArtifactToAlertsMap(
+        batchPackageFetchResult.data as CompactSocketArtifact,
+        alertsByPkgId,
+        toAlertsMapOptions
+      )
+    }
     remaining -= 1
     if (spinner && remaining > 0) {
       spinner.start()

src/utils/lockfile/pnpm-lock-yaml.ts

@@ -1,8 +1,10 @@
 import { detectDepTypes } from '@pnpm/lockfile.detect-dep-types'
 
-import { batchScan } from '../alert/artifact'
-import { AlertsByPkgId, addArtifactToAlertsMap } from '../socket-package-alert'
+import { getPublicToken, setupSdk } from '../sdk'
+import { addArtifactToAlertsMap } from '../socket-package-alert'
 
+import type { CompactSocketArtifact } from '../alert/artifact'
+import type { AlertsByPkgId } from '../socket-package-alert'
 import type { Lockfile } from '@pnpm/lockfile-file'
 import type { Spinner } from '@socketsecurity/registry/lib/spinner'
 
@@ -31,6 +33,7 @@ export async function getAlertsMapFromPnpmLockfile(
 
   const depTypes = detectDepTypes(lockfile)
   const pkgIds = Object.keys(depTypes)
+
   let { length: remaining } = pkgIds
   const alertsByPkgId: AlertsByPkgId = new Map()
   if (!remaining) {
@@ -40,12 +43,29 @@ export async function getAlertsMapFromPnpmLockfile(
 
   spinner?.start(getText())
 
+  const socketSdk = await setupSdk(getPublicToken())
+
   const toAlertsMapOptions = {
     overrides: lockfile.overrides,
     ...options
   }
-  for await (const artifact of batchScan(pkgIds)) {
-    await addArtifactToAlertsMap(artifact, alertsByPkgId, toAlertsMapOptions)
+
+  for await (const batchPackageFetchResult of socketSdk.batchPackageStream(
+    {
+      alerts: 'true',
+      compact: 'true'
+    },
+    {
+      components: pkgIds.map(id => ({ purl: `pkg:npm/${id}` }))
+    }
+  )) {
+    if (batchPackageFetchResult.success) {
+      await addArtifactToAlertsMap(
+        batchPackageFetchResult.data as CompactSocketArtifact,
+        alertsByPkgId,
+        toAlertsMapOptions
+      )
+    }
     remaining -= 1
     if (spinner && remaining > 0) {
       spinner.start()