feat: add vulnerability quick-fix + hint (#39)

* feat: add vulnerability quick-fix + hint

* chore: remove unused variable

* refactor: improve vulnerability diagnostic messaging

* cleanup

* refactor: simplify

* Update src/providers/diagnostics/rules/vulnerability.ts

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

* refactor: move test to code-actions

* test: update

* refactor: simplify `getBestFixedInVersion`

---------

Co-authored-by: Vida Xie <vida_2020@163.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

Author: 3 people

src/index.ts

@@ -6,6 +6,7 @@ import { openFileInNpmx } from './commands/open-file-in-npmx'
 import { openInBrowser } from './commands/open-in-browser'
 import { commands, displayName, version } from './generated-meta'
 import { UpgradeProvider } from './providers/code-actions/upgrade'
+import { VulnerabilityCodeActionProvider } from './providers/code-actions/vulnerability'
 import { VersionCompletionItemProvider } from './providers/completion-item/version'
 import { useDiagnostics } from './providers/diagnostics'
 import { NpmxHoverProvider } from './providers/hover/npmx'
@@ -53,6 +54,19 @@ export const { activate, deactivate } = defineExtension(() => {
     onCleanup(() => Disposable.from(...disposables).dispose())
   })
 
+  watchEffect((onCleanup) => {
+    if (!config.diagnostics.vulnerability)
+      return
+
+    const provider = new VulnerabilityCodeActionProvider()
+    const options = { providedCodeActionKinds: [CodeActionKind.QuickFix] }
+    const disposables = extractorEntries.map(({ pattern }) =>
+      languages.registerCodeActionsProvider({ pattern }, provider, options),
+    )
+
+    onCleanup(() => Disposable.from(...disposables).dispose())
+  })
+
   useDiagnostics()
 
   useCommands({

src/providers/code-actions/vulnerability.ts

@@ -0,0 +1,49 @@
+import type { CodeActionContext, CodeActionProvider, Diagnostic, Range, TextDocument } from 'vscode'
+import { CodeAction, CodeActionKind, WorkspaceEdit } from 'vscode'
+
+const FIXED_VERSION_MESSAGE_PATTERN = / Upgrade to (?<fixedInVersion>\S+) to fix\.$/
+
+function getDiagnosticCodeValue(diagnostic: Diagnostic): string | null {
+  if (typeof diagnostic.code === 'string')
+    return diagnostic.code
+
+  if (typeof diagnostic.code === 'object' && typeof diagnostic.code.value === 'string')
+    return diagnostic.code.value
+
+  return null
+}
+
+function isVulnerabilityDiagnostic(diagnostic: Diagnostic): boolean {
+  return getDiagnosticCodeValue(diagnostic) === 'vulnerability'
+}
+
+function getFixedInVersion(diagnostic: Diagnostic): string | null {
+  const fixedInVersionMatch = FIXED_VERSION_MESSAGE_PATTERN.exec(diagnostic.message)
+  const fixedInVersion = fixedInVersionMatch?.groups?.fixedInVersion
+  return fixedInVersion && fixedInVersion.length > 0 ? fixedInVersion : null
+}
+
+function createUpdateVersionAction(document: TextDocument, range: Range, fixedInVersion: string): CodeAction {
+  const codeAction = new CodeAction(`Update to ${fixedInVersion} to fix vulnerabilities`, CodeActionKind.QuickFix)
+  codeAction.isPreferred = true
+  const workspaceEdit = new WorkspaceEdit()
+  workspaceEdit.replace(document.uri, range, fixedInVersion)
+  codeAction.edit = workspaceEdit
+
+  return codeAction
+}
+
+export class VulnerabilityCodeActionProvider implements CodeActionProvider {
+  provideCodeActions(document: TextDocument, _range: Range, context: CodeActionContext): CodeAction[] {
+    return context.diagnostics.flatMap((diagnostic) => {
+      if (!isVulnerabilityDiagnostic(diagnostic))
+        return []
+
+      const fixedInVersion = getFixedInVersion(diagnostic)
+      if (!fixedInVersion)
+        return []
+
+      return [createUpdateVersionAction(document, diagnostic.range, fixedInVersion)]
+    })
+  }
+}

src/providers/diagnostics/rules/vulnerability.ts

@@ -1,8 +1,8 @@
-import type { OsvSeverityLevel } from '#utils/api/vulnerability'
+import type { OsvSeverityLevel, PackageVulnerabilityInfo } from '#utils/api/vulnerability'
 import type { DiagnosticRule } from '..'
 import { getVulnerability, SEVERITY_LEVELS } from '#utils/api/vulnerability'
 import { npmxPackageUrl } from '#utils/links'
-import { isSupportedProtocol, parseVersion } from '#utils/version'
+import { formatVersion, isSupportedProtocol, lt, parseVersion } from '#utils/version'
 import { DiagnosticSeverity, Uri } from 'vscode'
 
 const DIAGNOSTIC_MAPPING: Record<Exclude<OsvSeverityLevel, 'unknown'>, DiagnosticSeverity> = {
@@ -12,6 +12,19 @@ const DIAGNOSTIC_MAPPING: Record<Exclude<OsvSeverityLevel, 'unknown'>, Diagnosti
   low: DiagnosticSeverity.Hint,
 }
 
+function getBigestFixedInVersion(vulnerablePackages: PackageVulnerabilityInfo[]): string | undefined {
+  let bigest: string | undefined
+  for (const { depth, vulnerabilities } of vulnerablePackages) {
+    if (depth !== 'root')
+      continue
+    for (const { fixedIn } of vulnerabilities) {
+      if (fixedIn && (!bigest || lt(bigest, fixedIn)))
+        bigest = fixedIn
+    }
+  }
+  return bigest
+}
+
 export const checkVulnerability: DiagnosticRule = async (dep, pkg) => {
   const parsed = parseVersion(dep.version)
   if (!parsed || !isSupportedProtocol(parsed.protocol))
@@ -26,7 +39,7 @@ export const checkVulnerability: DiagnosticRule = async (dep, pkg) => {
   if (!result)
     return
 
-  const { totalCounts } = result
+  const { totalCounts, vulnerablePackages } = result
   const message: string[] = []
   let severity: DiagnosticSeverity | null = null
 
@@ -45,10 +58,15 @@ export const checkVulnerability: DiagnosticRule = async (dep, pkg) => {
   if (!message.length)
     return
 
+  const fixedInVersion = getBigestFixedInVersion(vulnerablePackages)
+  const messageSuffix = fixedInVersion
+    ? ` Upgrade to ${formatVersion({ ...parsed, semver: fixedInVersion })} to fix.`
+    : ''
+
   return {
     node: dep.versionNode,
-    message: `This version has ${message.join(', ')} ${message.length === 1 ? 'vulnerability' : 'vulnerabilities'}`,
-    severity: DiagnosticSeverity.Error,
+    message: `This version has ${message.join(', ')} ${message.length === 1 ? 'vulnerability' : 'vulnerabilities'}.${messageSuffix}`,
+    severity: severity ?? DiagnosticSeverity.Error,
     code: {
       value: 'vulnerability',
       target: Uri.parse(npmxPackageUrl(dep.name, semver)),

src/utils/api/vulnerability.ts

@@ -23,6 +23,7 @@ export interface VulnerabilitySummary {
   severity: OsvSeverityLevel
   aliases: string[]
   url: string
+  fixedIn?: string
 }
 
 /** Depth in dependency tree */

tests/__mocks__/vscode.ts

@@ -9,6 +9,9 @@ export const Range = vscode.Range
 export const Position = vscode.Position
 export const Location = vscode.Location
 export const Selection = vscode.Selection
+export const CodeAction = vscode.CodeAction
+export const CodeActionKind = vscode.CodeActionKind
+export const WorkspaceEdit = vscode.WorkspaceEdit
 export const ThemeColor = vscode.ThemeColor
 export const ThemeIcon = vscode.ThemeIcon
 export const TreeItem = vscode.TreeItem

tests/code-actions/vulnerability.test.ts

@@ -0,0 +1,70 @@
+import type { CodeActionContext, Diagnostic, TextDocument } from 'vscode'
+import { describe, expect, it, vi } from 'vitest'
+import { Range, Uri } from 'vscode'
+import { VulnerabilityCodeActionProvider } from '../../src/providers/code-actions/vulnerability'
+
+function createDiagnostic(options: { code: string | { value: string }, message: string }): Diagnostic {
+  return {
+    code: options.code,
+    message: options.message,
+    range: new Range(0, 0, 0, 6),
+  } as Diagnostic
+}
+
+function createTextDocument(versionText: string): TextDocument {
+  return {
+    uri: Uri.parse('file:///package.json'),
+    getText: vi.fn(() => versionText),
+  } as unknown as TextDocument
+}
+
+function createCodeActionContext(diagnostics: Diagnostic[]): CodeActionContext {
+  return {
+    diagnostics,
+    triggerKind: 1 as CodeActionContext['triggerKind'],
+    only: undefined,
+  }
+}
+
+describe('vulnerability code action provider', () => {
+  it('provides a quick fix when vulnerability message includes upgrade version', () => {
+    const provider = new VulnerabilityCodeActionProvider()
+    const textDocument = createTextDocument('^1.0.0')
+
+    const diagnostic = createDiagnostic({
+      code: { value: 'vulnerability' },
+      message: 'This version has 1 high vulnerability. Upgrade to ^1.2.3 to fix.',
+    })
+
+    const codeActions = provider.provideCodeActions(
+      textDocument,
+      diagnostic.range,
+      createCodeActionContext([diagnostic]),
+    )
+
+    expect(codeActions).toEqual([
+      expect.objectContaining({
+        title: 'Update to ^1.2.3 to fix vulnerabilities',
+        isPreferred: true,
+      }),
+    ])
+  })
+
+  it('does not provide a quick fix when vulnerability message has no upgrade target', () => {
+    const provider = new VulnerabilityCodeActionProvider()
+    const textDocument = createTextDocument('^1.0.0')
+
+    const diagnostic = createDiagnostic({
+      code: { value: 'vulnerability' },
+      message: 'This version has 1 high vulnerability.',
+    })
+
+    const codeActions = provider.provideCodeActions(
+      textDocument,
+      diagnostic.range,
+      createCodeActionContext([diagnostic]),
+    )
+
+    expect(codeActions).toHaveLength(0)
+  })
+})

vitest.config.ts

@@ -9,6 +9,7 @@ export default defineConfig({
     alias: {
       '#constants': join(rootDir, '/src/constants.ts'),
       '#state': join(rootDir, '/src/state.ts'),
+      '#utils': join(rootDir, '/src/utils'),
       '#types/*': join(rootDir, '/src/types/*'),
       '#utils/*': join(rootDir, '/src/utils/*'),
       'vscode': join(rootDir, '/tests/__mocks__/vscode.ts'),