Support diagnostics in a untrusted workspace

Author: Kartik Raj

src/client/activation/activationManager.ts

@@ -97,9 +97,7 @@ export class ExtensionActivationManager implements IExtensionActivationManager {
         }
         await sendActivationTelemetry(this.fileSystem, this.workspaceService, resource);
         await Promise.all(this.activationServices.map((item) => item.activate(resource)));
-        if (this.workspaceService.isTrusted) {
-            await this.appDiagnostics.performPreStartupHealthCheck(resource);
-        }
+        await this.appDiagnostics.performPreStartupHealthCheck(resource);
     }
 
     public async initialize(): Promise<void> {

src/client/application/diagnostics/applicationDiagnostics.ts

@@ -5,6 +5,7 @@
 
 import { inject, injectable, named } from 'inversify';
 import { DiagnosticSeverity } from 'vscode';
+import { IWorkspaceService } from '../../common/application/types';
 import { isTestExecution, STANDARD_OUTPUT_CHANNEL } from '../../common/constants';
 import { IOutputChannel, Resource } from '../../common/types';
 import { IServiceContainer } from '../../ioc/types';
@@ -26,7 +27,11 @@ export class ApplicationDiagnostics implements IApplicationDiagnostics {
         if (isTestExecution()) {
             return;
         }
-        const services = this.serviceContainer.getAll<IDiagnosticsService>(IDiagnosticsService);
+        let services = this.serviceContainer.getAll<IDiagnosticsService>(IDiagnosticsService);
+        const workspaceService = this.serviceContainer.get<IWorkspaceService>(IWorkspaceService);
+        if (!workspaceService.isTrusted) {
+            services = services.filter((item) => item.runInUntrustedWorkspace);
+        }
         // Perform these validation checks in the foreground.
         await this.runDiagnostics(
             services.filter((item) => !item.runInBackground),

src/client/application/diagnostics/base.ts

@@ -35,6 +35,7 @@ export abstract class BaseDiagnosticsService implements IDiagnosticsService, IDi
         @unmanaged() protected serviceContainer: IServiceContainer,
         @unmanaged() disposableRegistry: IDisposableRegistry,
         @unmanaged() public readonly runInBackground: boolean = false,
+        @unmanaged() public readonly runInUntrustedWorkspace: boolean = false,
     ) {
         this.filterService = serviceContainer.get<IDiagnosticFilterService>(IDiagnosticFilterService);
         disposableRegistry.push(this);

src/client/application/diagnostics/checks/envPathVariable.ts

@@ -43,7 +43,13 @@ export class EnvironmentPathVariableDiagnosticsService extends BaseDiagnosticsSe
         @inject(IServiceContainer) serviceContainer: IServiceContainer,
         @inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry,
     ) {
-        super([DiagnosticCodes.InvalidEnvironmentPathVariableDiagnostic], serviceContainer, disposableRegistry, true);
+        super(
+            [DiagnosticCodes.InvalidEnvironmentPathVariableDiagnostic],
+            serviceContainer,
+            disposableRegistry,
+            true,
+            true,
+        );
         this.platform = this.serviceContainer.get<IPlatformService>(IPlatformService);
         this.messageService = serviceContainer.get<IDiagnosticHandlerService<MessageCommandPrompt>>(
             IDiagnosticHandlerService,

src/client/application/diagnostics/checks/pylanceDefault.ts

@@ -40,7 +40,7 @@ export class PylanceDefaultDiagnosticService extends BaseDiagnosticsService {
         protected readonly messageService: IDiagnosticHandlerService<MessageCommandPrompt>,
         @inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry,
     ) {
-        super([DiagnosticCodes.PylanceDefaultDiagnostic], serviceContainer, disposableRegistry, true);
+        super([DiagnosticCodes.PylanceDefaultDiagnostic], serviceContainer, disposableRegistry, true, true);
 
         this.initialMementoValue = this.context.globalState.get(EXTENSION_VERSION_MEMENTO);
     }

src/client/application/diagnostics/checks/switchToDefaultLS.ts

@@ -38,7 +38,7 @@ export class SwitchToDefaultLanguageServerDiagnosticService extends BaseDiagnost
         protected readonly messageService: IDiagnosticHandlerService<MessageCommandPrompt>,
         @inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry,
     ) {
-        super([DiagnosticCodes.JediPython27NotSupportedDiagnostic], serviceContainer, disposableRegistry, true);
+        super([DiagnosticCodes.JediPython27NotSupportedDiagnostic], serviceContainer, disposableRegistry, true, true);
     }
 
     public diagnose(resource: Resource): Promise<IDiagnostic[]> {

src/client/application/diagnostics/checks/upgradeCodeRunner.ts

@@ -43,7 +43,7 @@ export class UpgradeCodeRunnerDiagnosticService extends BaseDiagnosticsService {
         @inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry,
         @inject(IExtensions) private readonly extensions: IExtensions,
     ) {
-        super([DiagnosticCodes.UpgradeCodeRunnerDiagnostic], serviceContainer, disposableRegistry, true);
+        super([DiagnosticCodes.UpgradeCodeRunnerDiagnostic], serviceContainer, disposableRegistry, true, true);
         this.workspaceService = this.serviceContainer.get<IWorkspaceService>(IWorkspaceService);
     }
 

src/client/application/diagnostics/types.ts

@@ -27,6 +27,7 @@ export const IDiagnosticsService = Symbol('IDiagnosticsService');
 
 export interface IDiagnosticsService {
     readonly runInBackground: boolean;
+    readonly runInUntrustedWorkspace: boolean;
     diagnose(resource: Resource): Promise<IDiagnostic[]>;
     canHandle(diagnostic: IDiagnostic): Promise<boolean>;
     handle(diagnostics: IDiagnostic[]): Promise<void>;

src/test/activation/activationManager.unit.test.ts

@@ -159,7 +159,7 @@ suite('Activation Manager', () => {
             appDiagnostics
                 .setup((a) => a.performPreStartupHealthCheck(resource))
                 .returns(() => Promise.resolve())
-                .verifiable(typemoq.Times.never());
+                .verifiable(typemoq.Times.once());
 
             managerTest = new ExtensionActivationManagerTest(
                 [instance(activationService1), instance(activationService2)],

src/test/application/diagnostics/applicationDiagnostics.unit.test.ts

@@ -17,6 +17,7 @@ import {
     ISourceMapSupportService,
 } from '../../../client/application/diagnostics/types';
 import { IApplicationDiagnostics } from '../../../client/application/types';
+import { IWorkspaceService } from '../../../client/common/application/types';
 import { STANDARD_OUTPUT_CHANNEL } from '../../../client/common/constants';
 import { IOutputChannel } from '../../../client/common/types';
 import { createDeferred, createDeferredFromPromise } from '../../../client/common/utils/async';
@@ -30,6 +31,7 @@ suite('Application Diagnostics - ApplicationDiagnostics', () => {
     let lsNotSupportedCheck: typemoq.IMock<IDiagnosticsService>;
     let pythonInterpreterCheck: typemoq.IMock<IDiagnosticsService>;
     let outputChannel: typemoq.IMock<IOutputChannel>;
+    let workspaceService: typemoq.IMock<IWorkspaceService>;
     let appDiagnostics: IApplicationDiagnostics;
     const oldValueOfVSC_PYTHON_UNIT_TEST = process.env.VSC_PYTHON_UNIT_TEST;
     const oldValueOfVSC_PYTHON_CI_TEST = process.env.VSC_PYTHON_CI_TEST;
@@ -44,14 +46,20 @@ suite('Application Diagnostics - ApplicationDiagnostics', () => {
         lsNotSupportedCheck.setup((service) => service.runInBackground).returns(() => false);
         pythonInterpreterCheck = typemoq.Mock.ofType<IDiagnosticsService>();
         pythonInterpreterCheck.setup((service) => service.runInBackground).returns(() => false);
+        pythonInterpreterCheck.setup((service) => service.runInUntrustedWorkspace).returns(() => false);
         outputChannel = typemoq.Mock.ofType<IOutputChannel>();
+        workspaceService = typemoq.Mock.ofType<IWorkspaceService>();
+        workspaceService.setup((w) => w.isTrusted).returns(() => true);
 
         serviceContainer
             .setup((d) => d.getAll(typemoq.It.isValue(IDiagnosticsService)))
             .returns(() => [envHealthCheck.object, lsNotSupportedCheck.object, pythonInterpreterCheck.object]);
         serviceContainer
             .setup((d) => d.get(typemoq.It.isValue(IOutputChannel), typemoq.It.isValue(STANDARD_OUTPUT_CHANNEL)))
             .returns(() => outputChannel.object);
+        serviceContainer
+            .setup((d) => d.get(typemoq.It.isValue(IWorkspaceService)))
+            .returns(() => workspaceService.object);
 
         appDiagnostics = new ApplicationDiagnostics(serviceContainer.object, outputChannel.object);
     });
@@ -95,6 +103,29 @@ suite('Application Diagnostics - ApplicationDiagnostics', () => {
         pythonInterpreterCheck.verifyAll();
     });
 
+    test('When running in a untrusted workspace skip diagnosing validation checks which do not support it', async () => {
+        workspaceService.reset();
+        workspaceService.setup((w) => w.isTrusted).returns(() => false);
+        envHealthCheck
+            .setup((e) => e.diagnose(typemoq.It.isAny()))
+            .returns(() => Promise.resolve([]))
+            .verifiable(typemoq.Times.once());
+        lsNotSupportedCheck
+            .setup((p) => p.diagnose(typemoq.It.isAny()))
+            .returns(() => Promise.resolve([]))
+            .verifiable(typemoq.Times.once());
+        pythonInterpreterCheck
+            .setup((p) => p.diagnose(typemoq.It.isAny()))
+            .returns(() => Promise.resolve([]))
+            .verifiable(typemoq.Times.never());
+
+        await appDiagnostics.performPreStartupHealthCheck(undefined);
+
+        envHealthCheck.verifyAll();
+        lsNotSupportedCheck.verifyAll();
+        pythonInterpreterCheck.verifyAll();
+    });
+
     test('Performing Pre Startup Health Check must handles all validation checks only once either in background or foreground', async () => {
         const diagnostic: IDiagnostic = {
             code: 'Error' as any,
@@ -215,9 +246,12 @@ suite('Application Diagnostics - ApplicationDiagnostics', () => {
         const foreGroundService = mock(InvalidPythonInterpreterService);
         const backGroundService = mock(EnvironmentPathVariableDiagnosticsService);
         const svcContainer = mock(ServiceContainer);
+        const workspaceService = mock<IWorkspaceService>();
         const foreGroundDeferred = createDeferred<IDiagnostic[]>();
         const backgroundGroundDeferred = createDeferred<IDiagnostic[]>();
 
+        when(svcContainer.get<IWorkspaceService>(IWorkspaceService)).thenReturn(workspaceService);
+        when(workspaceService.isTrusted).thenReturn(true);
         when(svcContainer.getAll<IDiagnosticsService>(IDiagnosticsService)).thenReturn([
             instance(foreGroundService),
             instance(backGroundService),