From 7dd36ae7b42ea20ddb0b218926031786154252ea Mon Sep 17 00:00:00 2001
From: Maddy <12754347+MaddyDev@users.noreply.github.com>
Date: Wed, 20 Jul 2022 13:58:49 -0700
Subject: [PATCH] fix: image disappear issue (#20039)

* check if base64 value is from image tag

* add test

* check image regex

* add comment

* update comment
---
 .../test/electron-browser/cell.test.ts        | 19 +++++++++++++++++++
 .../services/notebook/browser/models/cell.ts  |  7 ++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts b/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts
index 328dd5f50d..93cbd6172d 100644
--- a/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts
+++ b/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts
@@ -1080,6 +1080,25 @@ suite('Cell Model', function (): void {
 		assert.deepStrictEqual(serializedCell.attachments, undefined, 'JSON should not include attachments if attachments do not exist');
 	});
 
+	test('Should not include image in attachments if image is added in html image tag', async function () {
+		const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}');
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: '![ads.png](attachment:ads.png)',
+			attachments: cellAttachment
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		let imageElement = '<img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAggg=="';
+		model.source = '![ads.png](attachment:ads.png) \n Test image: ' + imageElement;
+
+		assert.deepStrictEqual(model.attachments, contents.attachments, 'Should not add the image represented in html tag to the attachments of cell source');
+	});
+
 	test('Should remove unused attachments name when updating cell source', async function () {
 		const cellAttachment = JSON.parse('{"ads.png":{"image/png":"iVBORw0KGgoAAAANSUhEUgAAAggg=="}}');
 		let notebookModel = new NotebookModelStub({
diff --git a/src/sql/workbench/services/notebook/browser/models/cell.ts b/src/sql/workbench/services/notebook/browser/models/cell.ts
index d18fc8455f..e7ef2be15f 100644
--- a/src/sql/workbench/services/notebook/browser/models/cell.ts
+++ b/src/sql/workbench/services/notebook/browser/models/cell.ts
@@ -350,7 +350,12 @@ export class CellModel extends Disposable implements ICellModel {
 	private attachImageFromSource(newSource: string | string[]): string | string[] {
 		if (!Array.isArray(newSource) && this.isValidBase64OctetStream(newSource)) {
 			let results;
-			while ((results = validBase64OctetStreamRegex.exec(newSource)) !== null) {
+			// only replace the base64 value if it's from markdown [](base64value) not html tags <img src="base64value">
+			let validImageTag = /<img\s+[^>]*src="([^"]*)"[^>]*>/;
+			let imageResults;
+			// Note: Currently this will not process any markdown image attachments that are below an HTML img element.
+			// This is acceptable for now given the low risk of this happening and an easy workaround being to just changing the img element to a markdown embedded image instead
+			while ((results = validBase64OctetStreamRegex.exec(newSource)) !== null && ((imageResults = validImageTag.exec(newSource)) !== null && this.isValidBase64OctetStream(imageResults[1]) && results[0] !== imageResults[1])) {
 				let imageName = this.addAttachment(results[1], results[0], 'image.png');
 				newSource = newSource.replace(validBase64OctetStreamRegex, `attachment:${imageName}`);
 			}