Add --limit flag to fix command (#498)

Author: jdalton

src/commands/fix/cmd-fix.mts

@@ -19,11 +19,6 @@ const config: CliCommandConfig = {
   hidden: false,
   flags: {
     ...commonFlags,
-    autopilot: {
-      type: 'boolean',
-      default: false,
-      description: `Shorthand for --autoMerge --test`
-    },
     autoMerge: {
       type: 'boolean',
       default: false,
@@ -32,6 +27,16 @@ const config: CliCommandConfig = {
         'https://docs.github.com/en/repositories/configuring-branches-and-merges-in-your-repository/configuring-pull-request-merges/managing-auto-merge-for-pull-requests-in-your-repository'
       )} for managing auto-merge for pull requests in your repository.`
     },
+    autopilot: {
+      type: 'boolean',
+      default: false,
+      description: `Shorthand for --autoMerge --test`
+    },
+    limit: {
+      type: 'number',
+      default: Infinity,
+      description: 'The number of fixes to attempt at a time'
+    },
     purl: {
       type: 'string',
       default: [],
@@ -112,6 +117,7 @@ async function run(
   await runFix({
     autoMerge: Boolean(cli.flags['autoMerge']),
     autopilot: Boolean(cli.flags['autopilot']),
+    limit: Number(cli.flags['limit']),
     dryRun: Boolean(cli.flags['dryRun']),
     purls: Array.isArray(cli.flags['purl']) ? cli.flags['purl'] : [],
     rangeStyle: (cli.flags['rangeStyle'] ?? undefined) as

src/commands/fix/npm-fix.mts

@@ -24,7 +24,7 @@ import {
   getGitHubEnvRepoInfo,
   openGitHubPullRequest
 } from './open-pr.mts'
-import { alertMapOptions } from './shared.mts'
+import { getAlertMapOptions } from './shared.mts'
 import constants from '../../constants.mts'
 import {
   Arborist,
@@ -78,6 +78,7 @@ export async function npmFix(
     autoMerge,
     cwd,
     dryRun,
+    limit,
     purls,
     rangeStyle,
     test,
@@ -102,10 +103,10 @@ export async function npmFix(
   await arb.reify()
 
   const alertsMap = purls.length
-    ? await getAlertsMapFromPurls(purls, alertMapOptions)
-    : await getAlertsMapFromArborist(arb, alertMapOptions)
+    ? await getAlertsMapFromPurls(purls, getAlertMapOptions({ limit }))
+    : await getAlertsMapFromArborist(arb, getAlertMapOptions({ limit }))
 
-  const infoByPkg = getCveInfoByAlertsMap(alertsMap)
+  const infoByPkg = getCveInfoByAlertsMap(alertsMap, { limit })
   if (!infoByPkg) {
     spinner?.stop()
     logger.info('No fixable vulnerabilities found.')
@@ -124,7 +125,8 @@ export async function npmFix(
     pkgEnvDetails.editablePkgJson.filename!
   ]
 
-  for (const { 0: name, 1: infos } of infoByPkg) {
+  let count = 0
+  infoByPkgLoop: for (const { 0: name, 1: infos } of infoByPkg) {
     debugLog(`Processing vulnerable package: ${name}`)
 
     if (getManifestData(NPM, name)) {
@@ -334,6 +336,9 @@ export async function npmFix(
               error
             )
           }
+          if (++count >= limit) {
+            break infoByPkgLoop
+          }
         }
       }
     }

src/commands/fix/pnpm-fix.mts

@@ -26,7 +26,7 @@ import {
   getGitHubEnvRepoInfo,
   openGitHubPullRequest
 } from './open-pr.mts'
-import { alertMapOptions } from './shared.mts'
+import { getAlertMapOptions } from './shared.mts'
 import constants from '../../constants.mts'
 import {
   SAFE_ARBORIST_REIFY_OPTIONS_OVERRIDES,
@@ -97,6 +97,7 @@ export async function pnpmFix(
     autoMerge,
     cwd,
     dryRun,
+    limit,
     purls,
     rangeStyle,
     test,
@@ -126,10 +127,14 @@ export async function pnpmFix(
   }
 
   const alertsMap = purls.length
-    ? await getAlertsMapFromPurls(purls, alertMapOptions)
-    : await getAlertsMapFromPnpmLockfile(lockfile, alertMapOptions)
+    ? await getAlertsMapFromPurls(purls, getAlertMapOptions({ limit }))
+    : await getAlertsMapFromPnpmLockfile(
+        lockfile,
+        getAlertMapOptions({ limit })
+      )
 
-  const infoByPkg = getCveInfoByAlertsMap(alertsMap)
+  const infoByPkg = getCveInfoByAlertsMap(alertsMap, { limit })
+  console.log(infoByPkg)
   if (!infoByPkg) {
     spinner?.stop()
     logger.info('No fixable vulnerabilities found.')
@@ -148,7 +153,8 @@ export async function pnpmFix(
     pkgEnvDetails.editablePkgJson.filename!
   ]
 
-  for (const { 0: name, 1: infos } of infoByPkg) {
+  let count = 0
+  infoByPkgLoop: for (const { 0: name, 1: infos } of infoByPkg) {
     debugLog(`Processing vulnerable package: ${name}`)
 
     if (getManifestData(NPM, name)) {
@@ -397,6 +403,9 @@ export async function pnpmFix(
               error
             )
           }
+          if (++count >= limit) {
+            break infoByPkgLoop
+          }
         }
       }
     }

src/commands/fix/shared.mts

@@ -1,16 +1,28 @@
 import type { FixOptions, NormalizedFixOptions } from './types.mts'
+import type { GetAlertsMapFromPurlsOptions } from '../../utils/alerts-map.mts'
+import type { Remap } from '@socketsecurity/registry/lib/objects'
 
 export const CMD_NAME = 'socket fix'
 
-export const alertMapOptions = Object.freeze({
-  consolidate: true,
-  include: {
-    existing: true,
-    unfixable: false,
-    upgradable: false
-  },
-  nothrow: true
-})
+export function getAlertMapOptions(options: GetAlertsMapFromPurlsOptions = {}) {
+  return {
+    __proto__: null,
+    consolidate: true,
+    nothrow: true,
+    ...options,
+    include: {
+      __proto__: null,
+      existing: true,
+      unfixable: false,
+      upgradable: false,
+      ...options?.include
+    }
+  } as Remap<
+    Omit<GetAlertsMapFromPurlsOptions, 'include' | 'overrides' | 'spinner'> & {
+      include: Exclude<GetAlertsMapFromPurlsOptions['include'], undefined>
+    }
+  >
+}
 
 export function normalizeFixOptions(
   options_: FixOptions
@@ -28,6 +40,13 @@ export function normalizeFixOptions(
   if (typeof options.cwd !== 'string') {
     options.cwd = process.cwd()
   }
+  const limit =
+    typeof options.limit === 'number'
+      ? options.limit
+      : parseInt(`${options.limit || ''}`, 10)
+
+  options.limit = Number.isNaN(limit) ? Infinity : limit
+
   options.purls = Array.isArray(options.purls)
     ? options.purls.flatMap(p => p.split(/, */))
     : []

src/commands/fix/types.mts

@@ -9,6 +9,7 @@ export type FixOptions = {
   autopilot?: boolean | undefined
   cwd?: string | undefined
   dryRun?: boolean | undefined
+  limit?: number | undefined
   purls?: string[] | readonly string[] | undefined
   rangeStyle?: RangeStyle | undefined
   test?: boolean | undefined

src/utils/alerts-map.mts

@@ -29,6 +29,7 @@ export async function getAlertsMapFromArborist(
   const options = {
     __proto__: null,
     consolidate: false,
+    limit: Infinity,
     nothrow: false,
     ...options_
   } as GetAlertsMapFromArboristOptions
@@ -87,6 +88,7 @@ export async function getAlertsMapFromPnpmLockfile(
   const options = {
     __proto__: null,
     consolidate: false,
+    limit: Infinity,
     nothrow: false,
     ...options_
   } as GetAlertsMapFromPnpmLockfileOptions
@@ -100,6 +102,7 @@ export async function getAlertsMapFromPnpmLockfile(
 export type GetAlertsMapFromPurlsOptions = {
   consolidate?: boolean | undefined
   include?: AlertIncludeFilter | undefined
+  limit?: number | undefined
   overrides?: { [key: string]: string } | undefined
   nothrow?: boolean | undefined
   spinner?: Spinner | undefined
@@ -112,6 +115,7 @@ export async function getAlertsMapFromPurls(
   const options = {
     __proto__: null,
     consolidate: false,
+    limit: Infinity,
     nothrow: false,
     ...options_
   } as GetAlertsMapFromPurlsOptions

src/utils/socket-package-alert.mts

@@ -321,18 +321,27 @@ export type CveInfoByPkgId = Map<
 
 export type GetCveInfoByPackageOptions = {
   exclude?: CveExcludeFilter | undefined
+  limit?: number | undefined
 }
 
 export function getCveInfoByAlertsMap(
   alertsMap: AlertsByPkgId,
   options?: GetCveInfoByPackageOptions | undefined
 ): CveInfoByPkgId | null {
+  const { exclude: _exclude, limit = Infinity } = {
+    __proto__: null,
+    ...options
+  } as GetCveInfoByPackageOptions
+  console.log('limit', limit)
   const exclude = {
+    __proto__: null,
     upgradable: true,
-    ...({ __proto__: null, ...options } as GetCveInfoByPackageOptions).exclude
-  }
+    ..._exclude
+  } as CveExcludeFilter
+
+  let count = 0
   let infoByPkg: CveInfoByPkgId | null = null
-  for (const [pkgId, sockPkgAlerts] of alertsMap) {
+  alertsMapLoop: for (const [pkgId, sockPkgAlerts] of alertsMap) {
     const purlObj = PackageURL.fromString(idToPurl(pkgId))
     const name = resolvePackageName(purlObj)
     for (const sockPkgAlert of sockPkgAlerts) {
@@ -362,6 +371,9 @@ export function getCveInfoByAlertsMap(
             vulnerableVersionRange.replace(/, +/g, ' ')
           ).format()
         })
+        if (++count >= limit) {
+          break alertsMapLoop
+        }
       } catch (e) {
         debugLog('getCveInfoByAlertsMap', {
           firstPatchedVersionIdentifier,