Make npm fix on par with pnpm fix

src/commands/fix/npm-fix.ts

@@ -1,4 +1,5 @@
 import { getManifestData } from '@socketsecurity/registry'
+import { arrayUnique } from '@socketsecurity/registry/lib/arrays'
 import { logger } from '@socketsecurity/registry/lib/logger'
 import { runScript } from '@socketsecurity/registry/lib/npm'
 import {
@@ -14,6 +15,7 @@ import {
   SafeArborist
 } from '../../shadow/npm/arborist/lib/arborist'
 import {
+  findPackageNode,
   findPackageNodes,
   getAlertsMapFromArborist,
   updateNode,
@@ -23,22 +25,23 @@ import { getCveInfoByAlertsMap } from '../../utils/socket-package-alert'
 
 import type { SafeNode } from '../../shadow/npm/arborist/lib/node'
 import type { EnvDetails } from '../../utils/package-environment'
+import type { PackageJson } from '@socketsecurity/registry/lib/packages'
 import type { Spinner } from '@socketsecurity/registry/lib/spinner'
 
 const { CI, NPM } = constants
 
-type SaveAndInstallOptions = {
+type InstallOptions = {
   cwd?: string | undefined
 }
 
-async function saveAndInstall(
+async function install(
   idealTree: SafeNode,
-  options: SaveAndInstallOptions
+  options: InstallOptions
 ): Promise<void> {
   const { cwd = process.cwd() } = {
     __proto__: null,
     ...options
-  } as SaveAndInstallOptions
+  } as InstallOptions
   const arb2 = new Arborist({ path: cwd })
   arb2.idealTree = idealTree
   await arb2.reify()
@@ -68,7 +71,7 @@ export async function npmFix(
     path: cwd,
     ...SAFE_ARBORIST_REIFY_OPTIONS_OVERRIDES
   })
-
+  // Calling arb.reify() creates the arb.diff object and nulls-out arb.idealTree.
   await arb.reify()
 
   const alertsMap = await getAlertsMapFromArborist(arb, {
@@ -87,42 +90,48 @@ export async function npmFix(
     return
   }
 
-  await arb.buildIdealTree()
-
   const editablePkgJson = await readPackageJson(cwd, { editable: true })
+  const { content: pkgJson } = editablePkgJson
 
-  for (const { 0: name, 1: infos } of infoByPkg) {
-    const revertToIdealTree = arb.idealTree!
-    arb.idealTree = null
+  await arb.buildIdealTree()
 
-    // eslint-disable-next-line no-await-in-loop
-    await arb.buildIdealTree()
-    const tree = arb.idealTree!
+  spinner?.stop()
 
+  for (const { 0: name, 1: infos } of infoByPkg) {
     const hasUpgrade = !!getManifestData(NPM, name)
     if (hasUpgrade) {
       spinner?.info(`Skipping ${name}. Socket Optimize package exists.`)
       continue
     }
-
-    const nodes = findPackageNodes(tree, name)
-
+    const specs = arrayUnique(
+      findPackageNodes(arb.idealTree!, name).map(n => `${n.name}@${n.version}`)
+    )
     const packument =
-      nodes.length && infos.length
+      specs.length && infos.length
         ? // eslint-disable-next-line no-await-in-loop
           await fetchPackagePackument(name)
         : null
     if (!packument) {
       continue
     }
 
-    for (const node of nodes) {
+    for (const spec of specs) {
+      const lastAtSignIndex = spec.lastIndexOf('@')
+      const name = spec.slice(0, lastAtSignIndex)
+      const oldVersion = spec.slice(lastAtSignIndex + 1)
       for (const {
         firstPatchedVersionIdentifier,
         vulnerableVersionRange
       } of infos) {
+        const revertTree = arb.idealTree!
+        arb.idealTree = null
+        // eslint-disable-next-line no-await-in-loop
+        await arb.buildIdealTree()
+        const node = findPackageNode(arb.idealTree!, name, oldVersion)
+        if (!node) {
+          continue
+        }
         spinner?.stop()
-        const { version: oldVersion } = node
         const oldSpec = `${name}@${oldVersion}`
         if (
           updateNode(
@@ -134,38 +143,54 @@ export async function npmFix(
         ) {
           const targetVersion = node.package.version!
           const fixSpec = `${name}@^${targetVersion}`
+          const revertData = {
+            ...(pkgJson.dependencies
+              ? { dependencies: pkgJson.dependencies }
+              : undefined),
+            ...(pkgJson.optionalDependencies
+              ? { optionalDependencies: pkgJson.optionalDependencies }
+              : undefined),
+            ...(pkgJson.peerDependencies
+              ? { peerDependencies: pkgJson.peerDependencies }
+              : undefined)
+          } as PackageJson
 
           spinner?.start()
           spinner?.info(`Installing ${fixSpec}`)
 
           let saved = false
           let installed = false
           try {
+            updatePackageJsonFromNode(editablePkgJson, arb.idealTree!, node)
             // eslint-disable-next-line no-await-in-loop
-            await saveAndInstall(arb.idealTree!, { cwd })
+            await editablePkgJson.save()
             saved = true
+
+            // eslint-disable-next-line no-await-in-loop
+            await install(arb.idealTree!, { cwd })
             installed = true
 
             if (test) {
               spinner?.info(`Testing ${fixSpec}`)
               // eslint-disable-next-line no-await-in-loop
               await runScript(testScript, [], { spinner, stdio: 'ignore' })
             }
+            spinner?.info(`Fixed ${name}`)
             // Lazily access constants.ENV[CI].
             if (constants.ENV[CI]) {
               // eslint-disable-next-line no-await-in-loop
               await openGitHubPullRequest(name, targetVersion, cwd)
             }
-            updatePackageJsonFromNode(editablePkgJson, tree, node)
-            // eslint-disable-next-line no-await-in-loop
-            await editablePkgJson.save()
-            spinner?.info(`Fixed ${name}`)
           } catch {
             spinner?.error(`Reverting ${fixSpec}`)
-            if (saved || installed) {
-              arb.idealTree = revertToIdealTree
+            if (saved) {
+              editablePkgJson.update(revertData)
+              // eslint-disable-next-line no-await-in-loop
+              await editablePkgJson.save()
+            }
+            if (installed) {
               // eslint-disable-next-line no-await-in-loop
-              await saveAndInstall(arb.idealTree!, { cwd })
+              await install(revertTree, { cwd })
             }
             spinner?.stop()
             logger.error(`Failed to fix ${oldSpec}`)

src/commands/fix/pnpm-fix.ts

@@ -1,6 +1,7 @@
 import { readWantedLockfile } from '@pnpm/lockfile.fs'
 
 import { getManifestData } from '@socketsecurity/registry'
+import { arrayUnique } from '@socketsecurity/registry/lib/arrays'
 import { logger } from '@socketsecurity/registry/lib/logger'
 import { runScript } from '@socketsecurity/registry/lib/npm'
 import {
@@ -16,6 +17,7 @@ import {
 } from '../../shadow/npm/arborist/lib/arborist'
 import {
   findBestPatchVersion,
+  findPackageNode,
   findPackageNodes,
   updatePackageJsonFromNode
 } from '../../utils/arborist-helpers'
@@ -81,42 +83,47 @@ export async function pnpmFix(
     return spinner?.stop()
   }
 
+  const editablePkgJson = await readPackageJson(cwd, { editable: true })
+  const { content: pkgJson } = editablePkgJson
+
   const arb = new SafeArborist({
     path: cwd,
     ...SAFE_ARBORIST_REIFY_OPTIONS_OVERRIDES
   })
   await arb.loadActual()
 
-  const editablePkgJson = await readPackageJson(cwd, { editable: true })
-  const { content: pkgJson } = editablePkgJson
-
   spinner?.stop()
 
   for (const { 0: name, 1: infos } of infoByPkg) {
-    const tree = arb.actualTree!
-
     if (getManifestData(NPM, name)) {
       logger.info(`Skipping ${name}. Socket Optimize package exists.`)
       continue
     }
-
-    const nodes = findPackageNodes(tree, name)
+    const specs = arrayUnique(
+      findPackageNodes(arb.actualTree!, name).map(n => `${n.name}@${n.version}`)
+    )
     const packument =
-      nodes.length && infos.length
+      specs.length && infos.length
         ? // eslint-disable-next-line no-await-in-loop
           await fetchPackagePackument(name)
         : null
     if (!packument) {
       continue
     }
 
-    for (const node of nodes) {
+    for (const spec of specs) {
+      const lastAtSignIndex = spec.lastIndexOf('@')
+      const name = spec.slice(0, lastAtSignIndex)
+      const oldVersion = spec.slice(lastAtSignIndex + 1)
       for (const {
         firstPatchedVersionIdentifier,
         vulnerableVersionRange
       } of infos) {
+        const node = findPackageNode(arb.idealTree!, name, oldVersion)
+        if (!node) {
+          continue
+        }
         spinner?.stop()
-        const { version: oldVersion } = node
         const oldSpec = `${name}@${oldVersion}`
         const availableVersions = Object.keys(packument.versions)
         const targetVersion = findBestPatchVersion(
@@ -180,7 +187,7 @@ export async function pnpmFix(
           let installed = false
           try {
             editablePkgJson.update(updateData)
-            updatePackageJsonFromNode(editablePkgJson, tree, node)
+            updatePackageJsonFromNode(editablePkgJson, arb.actualTree!, node)
             // eslint-disable-next-line no-await-in-loop
             await editablePkgJson.save()
             saved = true
@@ -194,14 +201,14 @@ export async function pnpmFix(
               // eslint-disable-next-line no-await-in-loop
               await runScript(testScript, [], { spinner, stdio: 'ignore' })
             }
+
+            spinner?.info(`Fixed ${name}`)
+
             // Lazily access constants.ENV[CI].
             if (constants.ENV[CI]) {
               // eslint-disable-next-line no-await-in-loop
               await openGitHubPullRequest(name, targetVersion, cwd)
             }
-            // eslint-disable-next-line no-await-in-loop
-            await editablePkgJson.save()
-            spinner?.info(`Fixed ${name}`)
           } catch {
             spinner?.error(`Reverting ${fixSpec}`)
             if (saved) {
@@ -212,6 +219,9 @@ export async function pnpmFix(
             if (installed) {
               // eslint-disable-next-line no-await-in-loop
               await install(pkgEnvDetails, { spinner })
+              arb.actualTree = null
+              // eslint-disable-next-line no-await-in-loop
+              await arb.loadActual()
             }
             spinner?.stop()
             logger.error(`Failed to fix ${oldSpec}`)

src/utils/arborist-helpers.ts

@@ -163,9 +163,33 @@ export function findBestPatchVersion(
   return semver.maxSatisfying(eligibleVersions, '*')
 }
 
+export function findPackageNode(
+  tree: SafeNode,
+  name: string,
+  version?: string | undefined
+): SafeNode | undefined {
+  const queue: Array<{ node: typeof tree }> = [{ node: tree }]
+  let sentinel = 0
+  while (queue.length) {
+    if (sentinel++ === LOOP_SENTINEL) {
+      throw new Error('Detected infinite loop in findPackageNodes')
+    }
+    const { node: currentNode } = queue.pop()!
+    const node = currentNode.children.get(name)
+    if (node && (typeof version !== 'string' || node.version === version)) {
+      return node as unknown as SafeNode
+    }
+    const children = [...currentNode.children.values()]
+    for (let i = children.length - 1; i >= 0; i -= 1) {
+      queue.push({ node: children[i] as unknown as SafeNode })
+    }
+  }
+}
+
 export function findPackageNodes(
   tree: SafeNode,
-  packageName: string
+  name: string,
+  version?: string | undefined
 ): SafeNode[] {
   const queue: Array<{ node: typeof tree }> = [{ node: tree }]
   const matches: SafeNode[] = []
@@ -175,8 +199,8 @@ export function findPackageNodes(
       throw new Error('Detected infinite loop in findPackageNodes')
     }
     const { node: currentNode } = queue.pop()!
-    const node = currentNode.children.get(packageName)
-    if (node) {
+    const node = currentNode.children.get(name)
+    if (node && (typeof version !== 'string' || node.version === version)) {
       matches.push(node as unknown as SafeNode)
     }
     const children = [...currentNode.children.values()]
@@ -321,39 +345,49 @@ export function updateNode(
     // No suitable patch version found.
     return false
   }
-  // Use Object.defineProperty to override the version.
+  // Object.defineProperty is needed to set the version property and replace
+  // the old value with targetVersion.
   Object.defineProperty(node, 'version', {
     configurable: true,
     enumerable: true,
     get: () => targetVersion
   })
+  // Update package.version associated with the node.
   node.package.version = targetVersion
-  // Update resolved and clear integrity for the new version.
+  // Update node.resolved.
   const purlObj = PackageURL.fromString(`pkg:npm/${node.name}`)
   node.resolved = `${NPM_REGISTRY_URL}/${node.name}/-/${purlObj.name}-${targetVersion}.tgz`
+  // Update node.integrity with the targetPackument.dist.integrity value if available
+  // else delete node.integrity so a new value is resolved for the target version.
   const { integrity } = targetPackument.dist
   if (integrity) {
     node.integrity = integrity
   } else {
     delete node.integrity
   }
-  if ('deprecated' in targetPackument) {
+  // Update node.package.deprecated based on targetPackument.deprecated.
+  if (hasOwn(targetPackument, 'deprecated')) {
     node.package['deprecated'] = targetPackument.deprecated as string
   } else {
     delete node.package['deprecated']
   }
+  // Update node.package.dependencies.
   const newDeps = { ...targetPackument.dependencies }
   const { dependencies: oldDeps } = node.package
   node.package.dependencies = newDeps
   if (oldDeps) {
     for (const oldDepName of Object.keys(oldDeps)) {
       if (!hasOwn(newDeps, oldDepName)) {
+        // Detach old edges for dependencies that don't exist on the updated
+        // node.package.dependencies.
         node.edgesOut.get(oldDepName)?.detach()
       }
     }
   }
   for (const newDepName of Object.keys(newDeps)) {
     if (!hasOwn(oldDeps, newDepName)) {
+      // Add new edges for dependencies that don't exist on the old
+      // node.package.dependencies.
       node.addEdgeOut(
         new Edge({
           from: node,