From 29b8e77ae7e4a6c50a9a2a7d34956c31295530e9 Mon Sep 17 00:00:00 2001
From: Nate Moore <nate@astro.build>
Date: Fri, 5 Jan 2024 12:19:18 -0600
Subject: [PATCH] fix(html): properly handle escape sequences

---
 .../src/vite-plugin-html/transform/escape.ts  |  8 ++--
 .../src/vite-plugin-html/transform/slots.ts   |  4 +-
 .../src/vite-plugin-html/transform/utils.ts   | 36 +++++++++++++--
 .../html-escape-bug/src/pages/index.html      | 28 ------------
 .../package.json                              |  0
 .../html-escape-complex/src/pages/index.html  | 32 +++++++++++++
 packages/astro/test/html-escape-bug.test.js   | 38 ----------------
 .../astro/test/html-escape-complex.test.js    | 45 +++++++++++++++++++
 8 files changed, 116 insertions(+), 75 deletions(-)
 delete mode 100644 packages/astro/test/fixtures/html-escape-bug/src/pages/index.html
 rename packages/astro/test/fixtures/{html-escape-bug => html-escape-complex}/package.json (100%)
 create mode 100644 packages/astro/test/fixtures/html-escape-complex/src/pages/index.html
 delete mode 100644 packages/astro/test/html-escape-bug.test.js
 create mode 100644 packages/astro/test/html-escape-complex.test.js

diff --git a/packages/astro/src/vite-plugin-html/transform/escape.ts b/packages/astro/src/vite-plugin-html/transform/escape.ts
index 1c250d43dc..6bf5260554 100644
--- a/packages/astro/src/vite-plugin-html/transform/escape.ts
+++ b/packages/astro/src/vite-plugin-html/transform/escape.ts
@@ -3,19 +3,19 @@ import type MagicString from 'magic-string';
 import type { Plugin } from 'unified';
 import { visit } from 'unist-util-visit';
 
-import { escape, needsEscape, replaceAttribute } from './utils.js';
+import { escapeTemplateLiteralCharacters, needsEscape, replaceAttribute } from './utils.js';
 
 const rehypeEscape: Plugin<[{ s: MagicString }], Root> = ({ s }) => {
 	return (tree) => {
 		visit(tree, (node: Root | RootContent) => {
 			if (node.type === 'text' || node.type === 'comment') {
 				if (needsEscape(node.value)) {
-					s.overwrite(node.position!.start.offset!, node.position!.end.offset!, escape(node.value));
+					s.overwrite(node.position!.start.offset!, node.position!.end.offset!, escapeTemplateLiteralCharacters(node.value));
 				}
 			} else if (node.type === 'element') {
 				for (const [key, value] of Object.entries(node.properties ?? {})) {
-					const newKey = needsEscape(key) ? escape(key) : key;
-					const newValue = needsEscape(value) ? escape(value) : value;
+					const newKey = needsEscape(key) ? escapeTemplateLiteralCharacters(key) : key;
+					const newValue = needsEscape(value) ? escapeTemplateLiteralCharacters(value) : value;
 					if (newKey === key && newValue === value) continue;
 					replaceAttribute(s, node, key, value === '' ? newKey : `${newKey}="${newValue}"`);
 				}
diff --git a/packages/astro/src/vite-plugin-html/transform/slots.ts b/packages/astro/src/vite-plugin-html/transform/slots.ts
index a549a48c91..e324edda42 100644
--- a/packages/astro/src/vite-plugin-html/transform/slots.ts
+++ b/packages/astro/src/vite-plugin-html/transform/slots.ts
@@ -3,7 +3,7 @@ import type { Plugin } from 'unified';
 
 import type MagicString from 'magic-string';
 import { visit } from 'unist-util-visit';
-import { escape } from './utils.js';
+import { escapeTemplateLiteralCharacters } from './utils.js';
 
 const rehypeSlots: Plugin<[{ s: MagicString }], Root> = ({ s }) => {
 	return (tree, file) => {
@@ -18,7 +18,7 @@ const rehypeSlots: Plugin<[{ s: MagicString }], Root> = ({ s }) => {
 				const text = file.value
 					.slice(first.position?.start.offset ?? 0, last.position?.end.offset ?? 0)
 					.toString();
-				s.overwrite(start, end, `\${${SLOT_PREFIX}["${name}"] ?? \`${escape(text).trim()}\`}`);
+				s.overwrite(start, end, `\${${SLOT_PREFIX}["${name}"] ?? \`${escapeTemplateLiteralCharacters(text).trim()}\`}`);
 			}
 		});
 	};
diff --git a/packages/astro/src/vite-plugin-html/transform/utils.ts b/packages/astro/src/vite-plugin-html/transform/utils.ts
index 822cc595e0..10771d364e 100644
--- a/packages/astro/src/vite-plugin-html/transform/utils.ts
+++ b/packages/astro/src/vite-plugin-html/transform/utils.ts
@@ -21,9 +21,39 @@ export function replaceAttribute(s: MagicString, node: Element, key: string, new
 		s.overwrite(start, end, newValue);
 	}
 }
+
+// Embedding in our own template literal expression requires escaping
+// any meaningful template literal characters in the user's code!
+const NEEDS_ESCAPE_RE = /[`\\]|\$\{/g
+
 export function needsEscape(value: any): value is string {
-	return typeof value === 'string' && (value.includes('`') || value.includes('${'));
+	// Reset the RegExp's global state
+	NEEDS_ESCAPE_RE.lastIndex = 0;
+	return typeof value === 'string' && NEEDS_ESCAPE_RE.test(value);
 }
-export function escape(value: string) {
-	return value.replace(/(\\*)\`/g, '\\`').replace(/\$\{/g, '\\${');
+
+export function escapeTemplateLiteralCharacters(value: string) {
+	// Reset the RegExp's global state
+	NEEDS_ESCAPE_RE.lastIndex = 0;
+
+	let char: string | undefined;
+	let startIndex = 0;
+	let segment = '';
+	let text = '';
+	
+	// Rather than a naive `String.replace()`, we have to iterate through
+	// the raw contents to properly handle existing backslashes
+	while ([char] = NEEDS_ESCAPE_RE.exec(value) ?? []) {
+		// Final loop when char === undefined, append trailing content
+		if (!char) {
+			text += value.slice(startIndex);
+			break;
+		}
+		const endIndex = NEEDS_ESCAPE_RE.lastIndex - char.length;
+		const prefix = segment === '\\' ? '' : '\\';
+		segment = prefix + char;
+		text += value.slice(startIndex, endIndex) + segment;
+		startIndex = NEEDS_ESCAPE_RE.lastIndex;
+	}
+	return text;
 }
diff --git a/packages/astro/test/fixtures/html-escape-bug/src/pages/index.html b/packages/astro/test/fixtures/html-escape-bug/src/pages/index.html
deleted file mode 100644
index 6ef2759ba4..0000000000
--- a/packages/astro/test/fixtures/html-escape-bug/src/pages/index.html
+++ /dev/null
@@ -1,28 +0,0 @@
-<html lang="en">
-
-<head>
-	<meta charset="utf-8" />
-	<link rel="icon" type="image/svg+xml" href="/favicon.svg" />
-	<meta name="viewport" content="width=device-width" />
-	<meta name="generator" content="{Astro.generator}" />
-	<title>Astro</title>
-</head>
-
-<body>
-	<h1>Astro</h1>
-	<p>
-		Text should be red, but isn’t because the script breaks before it gets to
-		the point of making it red.
-	</p>
-	<p id="normal"></p>
-	<p id="content"></p>
-	<script>
-		const count = 6;
-		const normal = `There are ${count} things!`;
-		const content = `There are \`${count}\` things!`;
-		document.getElementById('normal').innerText = normal;
-		document.getElementById('content').innerText = content;
-	</script>
-</body>
-
-</html>
diff --git a/packages/astro/test/fixtures/html-escape-bug/package.json b/packages/astro/test/fixtures/html-escape-complex/package.json
similarity index 100%
rename from packages/astro/test/fixtures/html-escape-bug/package.json
rename to packages/astro/test/fixtures/html-escape-complex/package.json
diff --git a/packages/astro/test/fixtures/html-escape-complex/src/pages/index.html b/packages/astro/test/fixtures/html-escape-complex/src/pages/index.html
new file mode 100644
index 0000000000..00977998ee
--- /dev/null
+++ b/packages/astro/test/fixtures/html-escape-complex/src/pages/index.html
@@ -0,0 +1,32 @@
+<html lang="en">
+
+	<head>
+		<meta charset="utf-8" />
+		<link rel="icon" type="image/svg+xml" href="/favicon.svg" />
+		<meta name="viewport" content="width=device-width" />
+		<meta name="generator" content="{Astro.generator}" />
+		<title>Astro</title>
+	</head>
+
+	<body>
+		<h1>Astro</h1>
+		<p>
+			Text should be red, but isn’t because the script breaks before it gets to
+			the point of making it red.
+		</p>
+		<p id="normal"></p>
+		<p id="content"></p>
+
+		<script>
+			const normal = `There are ${count} things!`;
+			const content = `There are \`${count}\` things!`;
+			const a = "\`${a}\`";
+			const b = "\\`${b}\\`";
+			const c = "\\\`${c}\\\`";
+			const d = "\\\\`${d}\\\\`";
+			const e = "\\\\\`${e}\\\\\`";
+			const f = "\\\\\\`${f}\\\\\\`";
+		</script>
+	</body>
+
+</html>
diff --git a/packages/astro/test/html-escape-bug.test.js b/packages/astro/test/html-escape-bug.test.js
deleted file mode 100644
index fea8625b3d..0000000000
--- a/packages/astro/test/html-escape-bug.test.js
+++ /dev/null
@@ -1,38 +0,0 @@
-import { expect } from 'chai';
-import * as cheerio from 'cheerio';
-import { loadFixture } from './test-utils.js';
-import { describe } from 'node:test';
-
-describe('Html Escape Bug', () => {
-	let fixture;
-
-	before(async () => {
-		fixture = await loadFixture({
-			root: './fixtures/html-escape-bug/',
-		});
-	});
-
-	describe('build', () => {
-		before(async () => {
-			await fixture.build();
-		});
-
-		it('work', async () => {
-			const html = await fixture.readFile('/index.html');
-			const $ = cheerio.load(html);
-			const h1 = $('h1');
-			const script = $('script');
-
-			expect(h1.text()).to.equal('Astro');
-			expect(script.text()).to.equal(
-				[
-					'\n\t\tconst count = 6;',
-					'const normal = `There are ${count} things!`;',
-					'const content = `There are `${count}` things!`;',
-					`document.getElementById('normal').innerText = normal;`,
-					`document.getElementById('content').innerText = content;`,
-				].join('\n\t\t') + '\n\t'
-			);
-		});
-	});
-});
diff --git a/packages/astro/test/html-escape-complex.test.js b/packages/astro/test/html-escape-complex.test.js
new file mode 100644
index 0000000000..475636564e
--- /dev/null
+++ b/packages/astro/test/html-escape-complex.test.js
@@ -0,0 +1,45 @@
+import { expect } from 'chai';
+import * as cheerio from 'cheerio';
+import { loadFixture } from './test-utils.js';
+
+describe('HTML Escape (Complex)', () => {
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({
+			root: './fixtures/html-escape-complex/',
+		});
+	});
+
+	describe('build', () => {
+		before(async () => {
+			await fixture.build();
+		});
+
+		it('properly escapes user code', async () => {
+			const html = await fixture.readFile('/index.html');
+			const $ = cheerio.load(html);
+			const h1 = $('h1');
+			const script = $('script');
+
+			expect(h1.text()).to.equal('Astro');
+			// Ignore whitespace in text
+			const text = script.text().trim().split('\n').map(ln => ln.trim());
+
+			// HANDY FOR DEBUGGING BACKSLASHES:
+			// The logged output should exactly match the way <script> is authored in `index.html`
+			// console.log(text.join('\n'));
+
+			expect(text).to.deep.equal([
+				"const normal = `There are ${count} things!`;",
+				"const content = `There are \\`${count}\\` things!`;",
+				'const a = "\\`${a}\\`";',
+				'const b = "\\\\`${b}\\\\`";',
+				'const c = "\\\\\\`${c}\\\\\\`";',
+				'const d = "\\\\\\\\`${d}\\\\\\\\`";',
+				'const e = "\\\\\\\\\\`${e}\\\\\\\\\\`";',
+				'const f = "\\\\\\\\\\\\`${f}\\\\\\\\\\\\`";',
+			])
+		});
+	});
+});