fix: fix layer highlights for Dockerfiles (#8)

Author: Jujuyeh

package-lock.json

@@ -3,7 +3,7 @@
   "publisher": "sysdig",
   "displayName": "Sysdig Scanner",
   "description": "Sysdig Scanner for Visual Studio Code",
-  "version": "0.2.6",
+  "version": "0.2.7",
   "icon": "img/logo.png",
   "repository": "https://github.com/sysdiglabs/vscode-extension",
   "engines": {

package.json

@@ -31,7 +31,7 @@ export async function scanDockerfile(document: vscode.TextDocument, buildAndScan
         if (!report) {
             vscode.window.showErrorMessage('Failed to scan base image ' + baseImage);
         } else {
-            highlightImage(report, baseImage, document, baseImageRange);
+            highlightImage(report, document, baseImageRange);
         }
     }
 
@@ -51,7 +51,7 @@ export async function scanDockerfile(document: vscode.TextDocument, buildAndScan
             const isBuilt = await imageExists(imageName);
             if (buildResult && isBuilt) {
                 // 3. Scan the built image
-                report = await vscode.commands.executeCommand('sysdig-vscode-ext.scanImage', imageName, /* updateTrees: */ true, /* source: */ document, baseImageRange);
+                report = await vscode.commands.executeCommand('sysdig-vscode-ext.scanImage', imageName, /* updateTrees: */ true, /* source: */ document);
                 if (!report) {
                     vscode.window.showErrorMessage('Failed to scan built image ' + imageName);
                 } else if (!report.result.layers) {

src/fileScanners/Dockerfile/dockerfileScanner.ts

@@ -18,7 +18,7 @@ export async function scanComposeFile(document: vscode.TextDocument) {
       continue;
     }
 
-    highlightImage(report, image, document, range);
+    highlightImage(report, document, range);
   }
 }
 

src/fileScanners/docker-compose/dockercomposeScanner.ts

@@ -2,6 +2,7 @@ import * as vscode from 'vscode';
 import { createMarkdownSummary, Report } from '../types';
 import { Instruction } from 'dockerfile-ast';
 import { createMarkdownVulnsForLayer } from '../types/report';
+import { vulnTreeDataProvider } from '../extension';
 
 interface DecorationsMap {
     [key: string]: [
@@ -55,7 +56,7 @@ export function clearDecorations(document: vscode.TextDocument) {
     }
 }
 
-export function highlightImage(report: Report, image: string, document: vscode.TextDocument, baseImageRange?: vscode.Range) {
+export function highlightImage(report: Report, document: vscode.TextDocument, baseImageRange?: vscode.Range) {
     if (!baseImageRange) {
         return;
     }
@@ -95,14 +96,17 @@ export function highlightLayer(report: Report, instructions : Instruction[], doc
     let instructionIndex = instructions.length - 1;
     let layerIndex = layers.length - 1;
     let decorations : vscode.DecorationOptions[] = [];
-
+    let imageRange : vscode.Range | undefined;
+    let layerRanges : Map<string, vscode.Range> = new Map<string, vscode.Range>();
+    
     while (instructionIndex >= 0 && layerIndex >= 0) {
         const instruction = instructions[instructionIndex];
         const layer = layers[layerIndex];
 
         // Skip FROM instructions as base image is already scanned
         if (instruction.getInstruction() === 'FROM') {
             instructionIndex = -1;
+            imageRange = instruction.getRange() as vscode.Range;
             break;
         }
 
@@ -150,6 +154,10 @@ export function highlightLayer(report: Report, instructions : Instruction[], doc
                 }
             };
 
+            if (layer.digest) {
+                layerRanges = layerRanges.set(layer.digest, instruction.getRange() as vscode.Range);
+            }
+
             decorations.push(decorationOptions);
         }
     }
@@ -160,6 +168,8 @@ export function highlightLayer(report: Report, instructions : Instruction[], doc
     });
 
     addDecorations(document, decorations, decorationType);
+
+    vulnTreeDataProvider.updateVulnTree(report.result.packages, report.info.resultUrl, document, imageRange, layerRanges);
 }
 
 export function grepString(document: vscode.TextDocument, searchString: string): vscode.Range[] | undefined {

src/fileScanners/highlighters.ts

@@ -34,7 +34,7 @@ export async function scanKubernetesFile(document: vscode.TextDocument) {
         }
 
         scannedImages.push(image);
-        highlightImage(report, image, document);
+        highlightImage(report, document);
     }
 }
 

src/fileScanners/kubernetes-yaml/kubernetesScanner.ts

@@ -90,7 +90,7 @@ suite('Highlighters Tests', () => {
             editBuilder.insert(new vscode.Position(0, 0), image);
         });
 
-        highlightImage(report, image, document, range);
+        highlightImage(report, document, range);
 
         assert.deepEqual(decorationsMap[document.uri.toString()][0].decorations as vscode.DecorationOptions[], expectedDecorations);
     });
@@ -248,7 +248,7 @@ suite('Highlighters Tests', () => {
             }
         };
         const image = 'not-found-image';
-        highlightImage(report, image, document);
+        highlightImage(report, document);
 
         assert.equal(decorationsMap[document.uri.toString()], undefined);
     });

src/fileScanners/test/highlighters.test.ts

@@ -65,7 +65,7 @@ export async function runScan(context: vscode.ExtensionContext, binaryPath: stri
     let outputScanFile : string = `${tempDir}/${VM_SCAN_FILE}`;
     let dbPath : string = `${tempDir}/main.db`;
     let cachePath : string = `${tempDir}/cache`;
-    let skipUpload : boolean = !(configuration.get('vulnerabilityManagement.uploadResults') || false) || imageOverride !== undefined;
+    let skipUpload : boolean = !(configuration.get('vulnerabilityManagement.uploadResults') || false);
     let policies : Array<string> = configuration.get('vulnerabilityManagement.addPolicies') || [];
     let imageToScan : string = imageOverride || configuration.get('vulnerabilityManagement.imageToScan') || "";
     let standalone : string = configuration.get('vulnerabilityManagement.standaloneMode') || "Never";
@@ -164,10 +164,14 @@ function updateVulnerabilities(statusBar: vscode.StatusBarItem, report: Report,
     statusBar.text = `$(shield) C ${summary.critical}  H ${summary.high}  M ${summary.medium}  L ${summary.low}  N ${summary.negligible}`;
     statusBar.show();
 
-    vulnTreeDataProvider.updateVulnTree(report.result.packages, report.info.resultUrl, report.result.layers, source, range);
+
+    vulnTreeDataProvider.updateVulnTree(report.result.packages, report.info.resultUrl, source, range);
     policyTreeDataProvider.addPolicies(report.result.policyEvaluations || []);
     if (source) {
-        clearDecorations(source);
-        highlightImage(report, "", source, range);
+        if (range) { /* If new highlight range is provided, clear the previous decorations */
+            clearDecorations(source);
+        } /* If not, just add upon the previous decorations */
+
+        highlightImage(report, source, range);
     }
 }

src/runners/vmScanRunner.ts

@@ -3,7 +3,6 @@ import * as vscode from 'vscode';
 import {
     Vulnerability, vulnToColor, vulnToMarkdownString, doesVulnPassFilter,
     Package, sortPackages,
-    Layer
 } from '../types';
 
 export class VulnTreeItem extends vscode.TreeItem {
@@ -78,8 +77,8 @@ export class VulnTreeDataProvider implements vscode.TreeDataProvider<VulnTreeIte
 
     private backlink : string = "";
     private source : vscode.TextDocument | undefined = undefined;
-    private range: vscode.Range | undefined = undefined;
-    private layers : Layer[] | undefined;
+    private imageRange: vscode.Range | undefined = undefined;
+    private ranges: Map<string, vscode.Range> | undefined = undefined;
 
     constructor(source? : vscode.TextDocument) {
         this.source = source;
@@ -150,8 +149,9 @@ export class VulnTreeDataProvider implements vscode.TreeDataProvider<VulnTreeIte
 
             // Filter out packages without vulnerabilities 
             return Promise.resolve(this.filteredPackages.map(pkg => {
-                if (this.source && pkg.layerDigest && this.layers) {
-                    return new TreePackage(pkg, undefined, undefined, this.source, this.range);
+                if (this.source && pkg.layerDigest) {
+                    const range = this.ranges?.get(pkg.layerDigest) || this.imageRange;
+                    return new TreePackage(pkg, undefined, undefined, this.source, range);
                 } else {
                     return new TreePackage(pkg);
                 }
@@ -164,10 +164,10 @@ export class VulnTreeDataProvider implements vscode.TreeDataProvider<VulnTreeIte
         }
     }
 
-    updateVulnTree(packages: Package[], backlink: string, layers?: Layer[], source?: vscode.TextDocument, range?: vscode.Range) {
-        this.layers = layers;
+    updateVulnTree(packages: Package[], backlink: string, source?: vscode.TextDocument, imageRange?: vscode.Range, layerRanges?:  Map<string, vscode.Range>) {
         this.source = source;
-        this.range = range;
+        this.imageRange = imageRange;
+        this.ranges = layerRanges;
         this.addPackages(packages);
         vscode.commands.executeCommand('setContext', 'sysdig-vscode-ext.showBacklink', false);