From d7576eea81d84470de9de2e50db80d1597ee2a97 Mon Sep 17 00:00:00 2001
From: Rob Lourens <roblourens@gmail.com>
Date: Sat, 17 Dec 2022 15:42:29 -0800
Subject: [PATCH] Disallow cell URIs with NotebookEdit Fix #146690

---
 src/vs/workbench/api/common/extHostTypes.ts   |  6 +++-
 .../api/test/browser/extHostTypes.test.ts     | 28 +++++++++++++++----
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/vs/workbench/api/common/extHostTypes.ts b/src/vs/workbench/api/common/extHostTypes.ts
index afb0d03fe11f..a2786f18059d 100644
--- a/src/vs/workbench/api/common/extHostTypes.ts
+++ b/src/vs/workbench/api/common/extHostTypes.ts
@@ -17,7 +17,7 @@ import { IExtensionDescription } from 'vs/platform/extensions/common/extensions'
 import { FileSystemProviderErrorCode, markAsFileSystemProviderError } from 'vs/platform/files/common/files';
 import { RemoteAuthorityResolverErrorCode } from 'vs/platform/remote/common/remoteAuthorityResolver';
 import { IRelativePatternDto } from 'vs/workbench/api/common/extHost.protocol';
-import { CellEditType, ICellPartialMetadataEdit, IDocumentMetadataEdit, isTextStreamMime } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellEditType, CellUri, ICellPartialMetadataEdit, IDocumentMetadataEdit, isTextStreamMime } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { checkProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions';
 import type * as vscode from 'vscode';
 
@@ -862,6 +862,10 @@ export class WorkspaceEdit implements vscode.WorkspaceEdit {
 					edit = editOrTuple;
 				}
 				if (NotebookEdit.isNotebookCellEdit(edit)) {
+					if (uri.scheme === CellUri.scheme) {
+						throw new Error('set must be called with a notebook document URI, not a cell URI.');
+					}
+
 					if (edit.newCellMetadata) {
 						this.replaceNotebookCellMetadata(uri, edit.range.start, edit.newCellMetadata, metadata);
 					} else if (edit.newNotebookMetadata) {
diff --git a/src/vs/workbench/api/test/browser/extHostTypes.test.ts b/src/vs/workbench/api/test/browser/extHostTypes.test.ts
index 5bf6668b6499..4bca6a1c4c43 100644
--- a/src/vs/workbench/api/test/browser/extHostTypes.test.ts
+++ b/src/vs/workbench/api/test/browser/extHostTypes.test.ts
@@ -4,13 +4,14 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as assert from 'assert';
-import { URI } from 'vs/base/common/uri';
-import * as types from 'vs/workbench/api/common/extHostTypes';
+import { CancellationError } from 'vs/base/common/errors';
+import { MarshalledId } from 'vs/base/common/marshallingIds';
+import { Mimes } from 'vs/base/common/mime';
 import { isWindows } from 'vs/base/common/platform';
 import { assertType } from 'vs/base/common/types';
-import { Mimes } from 'vs/base/common/mime';
-import { MarshalledId } from 'vs/base/common/marshallingIds';
-import { CancellationError } from 'vs/base/common/errors';
+import { URI } from 'vs/base/common/uri';
+import * as types from 'vs/workbench/api/common/extHostTypes';
+import { CellUri } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 
 function assertToJSON(a: any, expected: any) {
 	const raw = JSON.stringify(a);
@@ -432,6 +433,23 @@ suite('ExtHostTypes', function () {
 		assert.strictEqual(second.edit.newText, 'Foo');
 	});
 
+	test('WorkspaceEdit - NotebookEdits', () => {
+		const edit = new types.WorkspaceEdit();
+		const notebookEdit = types.NotebookEdit.insertCells(0, [new types.NotebookCellData(types.NotebookCellKind.Code, '// hello', 'javascript')]) as types.NotebookEdit;
+		const notebookUri = URI.parse('/foo/notebook.ipynb');
+		edit.set(notebookUri, [notebookEdit]);
+
+		const cellUri = CellUri.generate(notebookUri, 123);
+		try {
+			edit.set(cellUri, [notebookEdit]);
+		} catch (err) {
+			assert.ok(err.message.includes('set must be called with a notebook document URI'), err.toString());
+			return;
+		}
+
+		throw new Error('Expected set to throw with cell URI');
+	});
+
 	test('DocumentLink', () => {
 		assert.throws(() => new types.DocumentLink(null!, null!));
 		assert.throws(() => new types.DocumentLink(new types.Range(1, 1, 1, 1), null!));