microsoft · joyceerhl · Jul 6, 2020 · Jul 2, 2020 · Jul 2, 2020 · Jul 2, 2020
diff --git a/src/client/common/experiments/groups.ts b/src/client/common/experiments/groups.ts
@@ -94,3 +94,8 @@ export enum EnableStartPage {
 export enum RemoveKernelToolbarInInteractiveWindow {
     experiment = 'RemoveKernelToolbarInInteractiveWindow'
 }
+
+// Experiment to turn on trusted notebooks checks
+export enum EnableTrustedNotebooks {
+    experiment = 'EnableTrustedNotebooks'
+}
diff --git a/src/client/datascience/interactive-common/interactiveWindowTypes.ts b/src/client/datascience/interactive-common/interactiveWindowTypes.ts
@@ -333,6 +333,7 @@ export interface IRefreshVariablesRequest {
 export interface ILoadAllCells {
     cells: ICell[];
     isNotebookTrusted?: boolean;
+    shouldShowTrustMessage?: boolean;
 }
 
 export interface IScrollToCell {

diff --git a/src/client/datascience/interactive-ipynb/nativeEditor.ts b/src/client/datascience/interactive-ipynb/nativeEditor.ts
@@ -89,6 +89,7 @@ import cloneDeep = require('lodash/cloneDeep');
 import { concatMultilineStringInput, splitMultilineString } from '../../../datascience-ui/common';
 import { ServerStatus } from '../../../datascience-ui/interactive-common/mainState';
 import { isTestExecution, PYTHON_LANGUAGE, UseCustomEditorApi } from '../../common/constants';
+import { EnableTrustedNotebooks } from '../../common/experiments/groups';
 import { translateKernelLanguageToMonaco } from '../common';
 import { IDataViewerFactory } from '../data-viewing/types';
 import { getCellHashProvider } from '../editor-integration/cellhashprovider';
@@ -184,7 +185,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         @inject(INotebookProvider) notebookProvider: INotebookProvider,
         @inject(UseCustomEditorApi) useCustomEditorApi: boolean,
         @inject(ITrustService) private trustService: ITrustService,
-        @inject(IExperimentService) expService: IExperimentService
+        @inject(IExperimentService) private expService: IExperimentService
     ) {
         super(
             listeners,
@@ -716,7 +717,13 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
 
     private async sendInitialCellsToWebView(cells: ICell[], isNotebookTrusted: boolean): Promise<void> {
         sendTelemetryEvent(Telemetry.CellCount, undefined, { count: cells.length });
-        return this.postMessage(InteractiveWindowMessages.LoadAllCells, { cells, isNotebookTrusted });
+
+        const shouldShowTrustMessage = await this.expService.inExperiment(EnableTrustedNotebooks.experiment);
+        return this.postMessage(InteractiveWindowMessages.LoadAllCells, {
+            cells,
+            isNotebookTrusted,
+            shouldShowTrustMessage
+        });
     }
 
     private async exportAs(): Promise<void> {

diff --git a/src/client/datascience/interactive-ipynb/trustService.ts b/src/client/datascience/interactive-ipynb/trustService.ts
@@ -1,7 +1,8 @@
 import { createHmac } from 'crypto';
 import { inject, injectable } from 'inversify';
 import { EventEmitter, Uri } from 'vscode';
-import { IConfigurationService } from '../../common/types';
+import { EnableTrustedNotebooks } from '../../common/experiments/groups';
+import { IConfigurationService, IExperimentService } from '../../common/types';
 import { IDigestStorage, ITrustService } from '../types';
 
 @injectable()
@@ -13,22 +14,25 @@ export class TrustService implements ITrustService {
         return this.configService.getSettings().datascience.alwaysTrustNotebooks;
     }
     protected readonly _onDidSetNotebookTrust = new EventEmitter<void>();
+    private enabled: Promise<boolean>;
     constructor(
-        // @inject(IExperimentsManager) private readonly experiment: IExperimentsManager,
+        @inject(IExperimentService) private readonly experimentService: IExperimentService,
         @inject(IDigestStorage) private readonly digestStorage: IDigestStorage,
         @inject(IConfigurationService) private configService: IConfigurationService
-    ) {}
+    ) {
+        this.enabled = this.isInExperiment();
+    }
 
     /**
      * When a notebook is opened, we check the database to see if a trusted checkpoint
      * for this notebook exists by computing and looking up its digest.
-     * If the digest does not exist, we mark all the cells untrusted.
+     * If the digest does not exist, the notebook is marked untrusted.
      * Once a notebook is loaded in an untrusted state, no code will be executed and no
      * markdown will be rendered until notebook as a whole is marked trusted
      */
     public async isNotebookTrusted(uri: Uri, notebookContents: string) {
-        if (this.alwaysTrustNotebooks) {
-            return true; // User manually overrode our trust checking
+        if (this.alwaysTrustNotebooks || !(await this.enabled)) {
+            return true; // Skip check if user manually overrode our trust checking, or if user is not in experiment
         }
         // Compute digest and see if notebook is trusted
         const digest = await this.computeDigest(notebookContents);
@@ -41,7 +45,7 @@ export class TrustService implements ITrustService {
      * I.e. if the notebook has already been trusted by the user
      */
     public async trustNotebook(uri: Uri, notebookContents: string) {
-        if (!this.alwaysTrustNotebooks) {
+        if (!this.alwaysTrustNotebooks && (await this.enabled)) {
             // Only update digest store if the user wants us to check trust
             const digest = await this.computeDigest(notebookContents);
             await this.digestStorage.saveDigest(uri, digest);
@@ -54,4 +58,8 @@ export class TrustService implements ITrustService {
         hmac.update(notebookContents);
         return hmac.digest('hex');
     }
+
+    private async isInExperiment() {
+        return this.experimentService.inExperiment(EnableTrustedNotebooks.experiment);
+    }
 }
diff --git a/src/datascience-ui/history-react/interactivePanel.tsx b/src/datascience-ui/history-react/interactivePanel.tsx
@@ -223,6 +223,7 @@ ${buildSettingsCss(this.props.settings)}`}</style>
                         kernel={this.props.kernel}
                         selectServer={this.props.selectServer}
                         selectKernel={this.props.selectKernel}
+                        shouldShowTrustMessage={false}
                     />
                 );
             } else if (this.props.kernel.localizedUri === getLocString('DataScience.localJupyterServer', 'local')) {
@@ -237,6 +238,7 @@ ${buildSettingsCss(this.props.settings)}`}</style>
                 kernel={this.props.kernel}
                 selectServer={this.props.selectServer}
                 selectKernel={this.props.selectKernel}
+                shouldShowTrustMessage={false}
             />
         );
     }

diff --git a/src/datascience-ui/interactive-common/jupyterInfo.tsx b/src/datascience-ui/interactive-common/jupyterInfo.tsx
@@ -10,7 +10,8 @@ export interface IJupyterInfoProps {
     baseTheme: string;
     font: IFont;
     kernel: IServerState;
-    isNotebookTrusted?: boolean; // Native editor-specific
+    isNotebookTrusted?: boolean;
+    shouldShowTrustMessage: boolean;
     selectServer(): void;
     launchNotebookTrustPrompt?(): void; // Native editor-specific
     selectKernel(): void;
@@ -90,7 +91,7 @@ export class JupyterInfo extends React.Component<IJupyterInfoProps> {
     }
 
     private renderTrustMessage() {
-        if (this.props.isNotebookTrusted !== undefined) {
+        if (this.props.shouldShowTrustMessage) {
             const text = this.props.isNotebookTrusted
                 ? getLocString('DataScience.notebookIsTrusted', 'Trusted')
                 : getLocString('DataScience.notebookIsNotTrusted', 'Not Trusted');

diff --git a/src/datascience-ui/interactive-common/mainState.ts b/src/datascience-ui/interactive-common/mainState.ts
@@ -89,6 +89,7 @@ export type IMainState = {
     loaded: boolean;
     kernel: IServerState;
     isNotebookTrusted: boolean;
+    shouldShowTrustMessage: boolean;
 };
 
 export type SelectionAndFocusedInfo = {
@@ -196,7 +197,8 @@ export function generateTestState(filePath: string = '', editable: boolean = fal
             jupyterServerStatus: ServerStatus.NotStarted,
             language: PYTHON_LANGUAGE
         },
-        isNotebookTrusted: true
+        isNotebookTrusted: true,
+        shouldShowTrustMessage: true
     };
 }
 

diff --git a/src/datascience-ui/interactive-common/redux/store.ts b/src/datascience-ui/interactive-common/redux/store.ts
@@ -74,7 +74,8 @@ function generateDefaultState(
             },
             settings: testMode ? getDefaultSettings() : undefined, // When testing, we don't send (or wait) for the real settings.
             editorOptions: testMode ? computeEditorOptions(getDefaultSettings()) : undefined,
-            isNotebookTrusted: true
+            isNotebookTrusted: true,
+            shouldShowTrustMessage: false
         };
     }
 }

diff --git a/src/datascience-ui/native-editor/nativeEditor.tsx b/src/datascience-ui/native-editor/nativeEditor.tsx
@@ -126,7 +126,12 @@ ${buildSettingsCss(this.props.settings)}`}</style>
         setTimeout(() => this.props.insertAboveFirst(), 1);
     }
     private renderToolbarPanel() {
-        return <ToolbarComponent isNotebookTrusted={this.props.isNotebookTrusted}></ToolbarComponent>;
+        return (
+            <ToolbarComponent
+                shouldShowTrustMessage={this.props.shouldShowTrustMessage}
+                isNotebookTrusted={this.props.isNotebookTrusted}
+            ></ToolbarComponent>
+        );
     }
 
     private renderVariablePanel(baseTheme: string) {

diff --git a/src/datascience-ui/native-editor/redux/reducers/creation.ts b/src/datascience-ui/native-editor/redux/reducers/creation.ts
@@ -332,7 +332,8 @@ export namespace Creation {
             undoStack: [],
             cellVMs: vms,
             loaded: true,
-            isNotebookTrusted: arg.payload.data.isNotebookTrusted!
+            isNotebookTrusted: arg.payload.data.isNotebookTrusted!,
+            shouldShowTrustMessage: arg.payload.data.shouldShowTrustMessage!
         };
     }
 

diff --git a/src/datascience-ui/native-editor/toolbar.tsx b/src/datascience-ui/native-editor/toolbar.tsx
@@ -47,6 +47,7 @@ export type INativeEditorToolbarProps = INativeEditorDataProps & {
     selectServer: typeof actionCreators.selectServer;
     launchNotebookTrustPrompt: typeof actionCreators.launchNotebookTrustPrompt;
     isNotebookTrusted: boolean;
+    shouldShowTrustMessage: boolean;
 };
 
 function mapStateToProps(state: IStore): INativeEditorDataProps {
@@ -263,6 +264,7 @@ export class Toolbar extends React.PureComponent<INativeEditorToolbarProps> {
                         kernel={this.props.kernel}
                         selectServer={selectServer}
                         selectKernel={selectKernel}
+                        shouldShowTrustMessage={this.props.shouldShowTrustMessage}
                         isNotebookTrusted={this.props.isNotebookTrusted}
                         launchNotebookTrustPrompt={launchNotebookTrustPrompt}
                     />

diff --git a/src/test/datascience/interactivePanel.functional.test.tsx b/src/test/datascience/interactivePanel.functional.test.tsx
@@ -114,7 +114,8 @@ suite('DataScience Interactive Panel', () => {
             setVariableExplorerHeight: noopAny,
             editorOptions: {},
             settings: { showCellInputCode: true, allowInput: true, extraSettings: { editor: {} } } as any,
-            isNotebookTrusted: true
+            isNotebookTrusted: true,
+            shouldShowTrustMessage: false
         };
     });
     test('Input Cell is displayed', () => {

diff --git a/src/test/datascience/nativeEditor.toolbar.functional.test.tsx b/src/test/datascience/nativeEditor.toolbar.functional.test.tsx
@@ -62,7 +62,8 @@ suite('DataScience Native Toolbar', () => {
             setVariableExplorerHeight: sinon.stub(),
             launchNotebookTrustPrompt: sinon.stub(),
             variablesVisible: false,
-            isNotebookTrusted: true
+            isNotebookTrusted: true,
+            shouldShowTrustMessage: true
         };
     });
     function mountToolbar() {

diff --git a/src/test/datascience/trustService.unit.test.ts b/src/test/datascience/trustService.unit.test.ts
@@ -9,6 +9,7 @@ import { anything, instance, mock, when } from 'ts-mockito';
 import * as typemoq from 'typemoq';
 import { Uri } from 'vscode';
 import { ConfigurationService } from '../../client/common/configuration/service';
+import { ExperimentService } from '../../client/common/experiments/service';
 import { FileSystem } from '../../client/common/platform/fileSystem';
 import { IExtensionContext } from '../../client/common/types';
 import { DigestStorage } from '../../client/datascience/interactive-ipynb/digestStorage';
@@ -20,6 +21,7 @@ suite('DataScience - TrustService', () => {
     setup(() => {
         alwaysTrustNotebooks = false;
         const configService = mock(ConfigurationService);
+        const experimentService = mock(ExperimentService);
         const fileSystem = mock(FileSystem);
         const context = typemoq.Mock.ofType<IExtensionContext>();
         context.setup((c) => c.globalStoragePath).returns(() => os.tmpdir());
@@ -31,9 +33,10 @@ suite('DataScience - TrustService', () => {
         when(fileSystem.readFile(anything())).thenCall((f) => fs.readFile(f));
         when(fileSystem.createDirectory(anything())).thenCall((d) => fs.mkdir(d));
         when(fileSystem.directoryExists(anything())).thenCall((d) => fs.pathExists(d));
+        when(experimentService.inExperiment(anything())).thenResolve(true); // Pretend we're in an experiment so that trust checks proceed
 
         const digestStorage = new DigestStorage(instance(fileSystem), context.object);
-        trustService = new TrustService(digestStorage, instance(configService));
+        trustService = new TrustService(instance(experimentService), digestStorage, instance(configService));
     });
 
     test('Trusting a notebook', async () => {