From 671727ba83fbd509877ac8507c1269f9c5fc267a Mon Sep 17 00:00:00 2001
From: Nate Moore <natemoo-re@users.noreply.github.com>
Date: Thu, 24 Mar 2022 15:10:08 -0500
Subject: [PATCH] fix(#2846): handle destructured private env usage during SSR
 (#2861)

* fix(#2846): handle destructured private env usage during SSR

* test: add destructured env vars test

* fix: support destructured env vars

* fix: only inline referenced vars
---
 packages/astro/src/vite-plugin-env/index.ts   | 29 +++++++++++++++----
 packages/astro/test/astro-envs.test.js        |  7 +++++
 .../astro-envs/src/pages/destructured.astro   |  5 ++++
 3 files changed, 36 insertions(+), 5 deletions(-)
 create mode 100644 packages/astro/test/fixtures/astro-envs/src/pages/destructured.astro

diff --git a/packages/astro/src/vite-plugin-env/index.ts b/packages/astro/src/vite-plugin-env/index.ts
index ff6cacc23..df11534ca 100644
--- a/packages/astro/src/vite-plugin-env/index.ts
+++ b/packages/astro/src/vite-plugin-env/index.ts
@@ -34,11 +34,14 @@ function getPrivateEnv(viteConfig: vite.ResolvedConfig, astroConfig: AstroConfig
 	return Object.fromEntries(privateKeys.map((key) => [key, JSON.stringify(fullEnv[key])]));
 }
 
-function referencesPrivateKey(source: string, privateEnv: Record<string, any>) {
+function getReferencedPrivateKeys(source: string, privateEnv: Record<string, any>): Set<string> {
+	const references = new Set<string>();
 	for (const key of Object.keys(privateEnv)) {
-		if (source.includes(key)) return true;
+		if (source.includes(key)) {
+			references.add(key);
+		}
 	}
-	return false;
+	return references;
 }
 
 export default function envVitePlugin({ config: astroConfig }: EnvPluginOptions): vite.PluginOption {
@@ -68,6 +71,13 @@ export default function envVitePlugin({ config: astroConfig }: EnvPluginOptions)
 				if (privateEnv) {
 					const entries = Object.entries(privateEnv).map(([key, value]) => [`import.meta.env.${key}`, value]);
 					replacements = Object.fromEntries(entries);
+					// These additional replacements are needed to match Vite
+					replacements = Object.assign(replacements, {
+						// This catches destructed `import.meta.env` calls,
+						// BUT we only want to inject private keys referenced in the file.
+						// We overwrite this value on a per-file basis.
+						'import.meta.env': `({})`,
+					})
 					pattern = new RegExp(
 						// Do not allow preceding '.', but do allow preceding '...' for spread operations
 						'(?<!(?<!\\.\\.)\\.)\\b(' +
@@ -84,7 +94,8 @@ export default function envVitePlugin({ config: astroConfig }: EnvPluginOptions)
 			}
 
 			if (!privateEnv || !pattern) return source;
-			if (!referencesPrivateKey(source, privateEnv)) return source;
+			const references = getReferencedPrivateKeys(source, privateEnv);
+			if (references.size === 0) return source;
 
 			// Find matches for *private* env and do our own replacement.
 			const s = new MagicString(source);
@@ -93,7 +104,15 @@ export default function envVitePlugin({ config: astroConfig }: EnvPluginOptions)
 			while ((match = pattern.exec(source))) {
 				const start = match.index;
 				const end = start + match[0].length;
-				const replacement = '' + replacements[match[1]];
+				let replacement = '' + replacements[match[1]];
+				// If we match exactly `import.meta.env`, define _only_ referenced private variables
+				if (match[0] === 'import.meta.env') {
+					replacement = `(Object.assign(import.meta.env,{`
+					for (const key of references.values()) {
+						replacement += `${key}:${privateEnv[key]},`
+					}
+					replacement += '}))'
+				}
 				s.overwrite(start, end, replacement);
 			}
 
diff --git a/packages/astro/test/astro-envs.test.js b/packages/astro/test/astro-envs.test.js
index 3f4720c33..5cf3dc84a 100644
--- a/packages/astro/test/astro-envs.test.js
+++ b/packages/astro/test/astro-envs.test.js
@@ -23,6 +23,13 @@ describe('Environment Variables', () => {
 		expect(indexHtml).to.include('BLUE_BAYOU');
 	});
 
+	it('does render destructured public env and private env', async () => {
+		let indexHtml = await fixture.readFile('/destructured/index.html');
+
+		expect(indexHtml).to.include('CLUB_33');
+		expect(indexHtml).to.include('BLUE_BAYOU');
+	});
+
 	it('includes public env in client-side JS', async () => {
 		let dirs = await fixture.readdir('/');
 		let found = false;
diff --git a/packages/astro/test/fixtures/astro-envs/src/pages/destructured.astro b/packages/astro/test/fixtures/astro-envs/src/pages/destructured.astro
new file mode 100644
index 000000000..25fb95954
--- /dev/null
+++ b/packages/astro/test/fixtures/astro-envs/src/pages/destructured.astro
@@ -0,0 +1,5 @@
+---
+const { PUBLIC_PLACE, SECRET_PLACE } = import.meta.env;
+---
+<environment-variable>{PUBLIC_PLACE}</environment-variable>
+<environment-variable>{SECRET_PLACE}</environment-variable>