From dfdb9cc41197a750850b60c36a9f6855c4481168 Mon Sep 17 00:00:00 2001
From: Yunfei He <i.heyunfei@gmail.com>
Date: Fri, 27 Aug 2021 22:26:04 +0800
Subject: [PATCH] feat: allow custom vite env prefix (#4676)

---
 docs/config/index.md                          | 12 ++++++
 docs/guide/env-and-mode.md                    |  2 +
 packages/playground/env/.env                  |  1 +
 packages/playground/env/__tests__/env.spec.ts |  5 +++
 packages/playground/env/index.html            |  2 +
 packages/playground/env/vite.config.js        |  5 +++
 .../vite/src/node/__tests__/config.spec.ts    | 27 +++++++++++++-
 packages/vite/src/node/config.ts              | 37 ++++++++++++++++---
 packages/vite/src/node/utils.ts               |  4 ++
 9 files changed, 88 insertions(+), 7 deletions(-)
 create mode 100644 packages/playground/env/vite.config.js

diff --git a/docs/config/index.md b/docs/config/index.md
index f1cada641797fe..428dc74cdf6bc4 100644
--- a/docs/config/index.md
+++ b/docs/config/index.md
@@ -333,6 +333,18 @@ export default defineConfig(async ({ command, mode }) => {
 
   See [here](/guide/env-and-mode#env-files) for more about environment files.
 
+### envPrefix
+
+- **Type:** `string | string[]`
+- **Default:** `VITE_`
+
+  Env variables starts with `envPrefix` will be exposed to your client source code via import.meta.env.
+
+:::warning SECURITY NOTES
+
+- `envPrefix` should not be set as `''`, which will expose all your env variables and cause unexpected leaking of of sensitive information. Vite will throw error when detecting `''`.
+  :::
+
 ## Server Options
 
 ### server.host
diff --git a/docs/guide/env-and-mode.md b/docs/guide/env-and-mode.md
index 85ff820b1d05db..92c326e413247e 100644
--- a/docs/guide/env-and-mode.md
+++ b/docs/guide/env-and-mode.md
@@ -44,6 +44,8 @@ VITE_SOME_KEY=123
 
 Only `VITE_SOME_KEY` will be exposed as `import.meta.env.VITE_SOME_KEY` to your client source code, but `DB_PASSWORD` will not.
 
+If you want to customize env variables prefix, see [envPrefix](/config/index#envPrefix) option.
+
 :::warning SECURITY NOTES
 
 - `.env.*.local` files are local-only and can contain sensitive variables. You should add `.local` to your `.gitignore` to avoid them being checked into git.
diff --git a/packages/playground/env/.env b/packages/playground/env/.env
index 7dc2f2192b70ec..ad0e78511b8662 100644
--- a/packages/playground/env/.env
+++ b/packages/playground/env/.env
@@ -1,2 +1,3 @@
 VITE_CUSTOM_ENV_VARIABLE=1
+CUSTOM_PREFIX_ENV_VARIABLE=1
 VITE_EFFECTIVE_MODE_FILE_NAME=.env
\ No newline at end of file
diff --git a/packages/playground/env/__tests__/env.spec.ts b/packages/playground/env/__tests__/env.spec.ts
index f2c265d825ed80..907cebc8037ce9 100644
--- a/packages/playground/env/__tests__/env.spec.ts
+++ b/packages/playground/env/__tests__/env.spec.ts
@@ -22,6 +22,10 @@ test('custom', async () => {
   expect(await page.textContent('.custom')).toBe('1')
 })
 
+test('custom-prefix', async () => {
+  expect(await page.textContent('.custom-prefix')).toBe('1')
+})
+
 test('mode file override', async () => {
   expect(await page.textContent('.mode-file')).toBe(`.env.${mode}`)
 })
@@ -40,6 +44,7 @@ test('env object', async () => {
   const envText = await page.textContent('.env-object')
   expect(JSON.parse(envText)).toMatchObject({
     VITE_EFFECTIVE_MODE_FILE_NAME: `.env.${mode}`,
+    CUSTOM_PREFIX_ENV_VARIABLE: '1',
     VITE_CUSTOM_ENV_VARIABLE: '1',
     BASE_URL: '/',
     MODE: mode,
diff --git a/packages/playground/env/index.html b/packages/playground/env/index.html
index 202a0fae9a3624..77b3f7da5f5105 100644
--- a/packages/playground/env/index.html
+++ b/packages/playground/env/index.html
@@ -4,6 +4,7 @@ <h1>Environment Variables</h1>
 <p>import.meta.env.DEV: <code class="dev"></code></p>
 <p>import.meta.env.PROD: <code class="prod"></code></p>
 <p>import.meta.env.VITE_CUSTOM_ENV_VARIABLE: <code class="custom"></code></p>
+<p>import.meta.env.CUSTOM_PREFIX_ENV_VARIABLE: <code class="custom-prefix"></code></p>
 <p>
   import.meta.env.VITE_EFFECTIVE_MODE_FILE_NAME: <code class="mode-file"></code>
 </p>
@@ -17,6 +18,7 @@ <h1>Environment Variables</h1>
   text('.dev', import.meta.env.DEV)
   text('.prod', import.meta.env.PROD)
   text('.custom', import.meta.env.VITE_CUSTOM_ENV_VARIABLE)
+  text('.custom-prefix', import.meta.env.CUSTOM_PREFIX_ENV_VARIABLE)
   text('.mode-file', import.meta.env.VITE_EFFECTIVE_MODE_FILE_NAME)
   text('.inline', import.meta.env.VITE_INLINE)
   text('.node-env', process.env.NODE_ENV)
diff --git a/packages/playground/env/vite.config.js b/packages/playground/env/vite.config.js
new file mode 100644
index 00000000000000..9d6809cba8e39b
--- /dev/null
+++ b/packages/playground/env/vite.config.js
@@ -0,0 +1,5 @@
+const { defineConfig } = require('vite')
+
+module.exports = defineConfig({
+  envPrefix: ['VITE_', 'CUSTOM_PREFIX_']
+})
diff --git a/packages/vite/src/node/__tests__/config.spec.ts b/packages/vite/src/node/__tests__/config.spec.ts
index 450e74aeee4bb5..9fa1a61b0b4710 100644
--- a/packages/vite/src/node/__tests__/config.spec.ts
+++ b/packages/vite/src/node/__tests__/config.spec.ts
@@ -1,5 +1,11 @@
 import { InlineConfig } from '..'
-import { mergeConfig, resolveConfig, UserConfigExport } from '../config'
+import {
+  mergeConfig,
+  resolveConfig,
+  UserConfigExport,
+  resolveEnvPrefix,
+  UserConfig
+} from '../config'
 
 describe('mergeConfig', () => {
   test('handles configs with different alias schemas', () => {
@@ -139,3 +145,22 @@ describe('resolveConfig', () => {
     })
   })
 })
+
+describe('resolveEnvPrefix', () => {
+  test(`use 'VITE_' as default value`, () => {
+    const config: UserConfig = {}
+    expect(resolveEnvPrefix(config)).toMatchObject(['VITE_'])
+  })
+
+  test(`throw error if envPrefix contains ''`, () => {
+    let config: UserConfig = { envPrefix: '' }
+    expect(() => resolveEnvPrefix(config)).toThrow()
+    config = { envPrefix: ['', 'CUSTOM_'] }
+    expect(() => resolveEnvPrefix(config)).toThrow()
+  })
+
+  test('should work correctly for valid envPrefix value', () => {
+    const config: UserConfig = { envPrefix: [' ', 'CUSTOM_'] }
+    expect(resolveEnvPrefix(config)).toMatchObject([' ', 'CUSTOM_'])
+  })
+})
diff --git a/packages/vite/src/node/config.ts b/packages/vite/src/node/config.ts
index 5d04918ed48f42..539b26205d0bc6 100644
--- a/packages/vite/src/node/config.ts
+++ b/packages/vite/src/node/config.ts
@@ -9,6 +9,7 @@ import {
 } from './server'
 import { CSSOptions } from './plugins/css'
 import {
+  arraify,
   createDebugger,
   isExternalUrl,
   isObject,
@@ -165,6 +166,11 @@ export interface UserConfig {
    * @default root
    */
   envDir?: string
+  /**
+   * Env variables starts with `envPrefix` will be exposed to your client source code via import.meta.env.
+   * @default 'VITE_'
+   */
+  envPrefix?: string | string[]
   /**
    * Import aliases
    * @deprecated use `resolve.alias` instead
@@ -323,7 +329,9 @@ export async function resolveConfig(
   const envDir = config.envDir
     ? normalizePath(path.resolve(resolvedRoot, config.envDir))
     : resolvedRoot
-  const userEnv = inlineConfig.envFile !== false && loadEnv(mode, envDir)
+  const userEnv =
+    inlineConfig.envFile !== false &&
+    loadEnv(mode, envDir, resolveEnvPrefix(config))
 
   // Note it is possible for user to have a custom mode, e.g. `staging` where
   // production-like behavior is expected. This is indicated by NODE_ENV=production
@@ -942,7 +950,7 @@ async function loadConfigFromBundledFile(
 export function loadEnv(
   mode: string,
   envDir: string,
-  prefix = 'VITE_'
+  prefixes: string | string[] = 'VITE_'
 ): Record<string, string> {
   if (mode === 'local') {
     throw new Error(
@@ -950,7 +958,7 @@ export function loadEnv(
         `the .local postfix for .env files.`
     )
   }
-
+  prefixes = arraify(prefixes)
   const env: Record<string, string> = {}
   const envFiles = [
     /** mode local file */ `.env.${mode}.local`,
@@ -962,7 +970,10 @@ export function loadEnv(
   // check if there are actual env variables starting with VITE_*
   // these are typically provided inline and should be prioritized
   for (const key in process.env) {
-    if (key.startsWith(prefix) && env[key] === undefined) {
+    if (
+      prefixes.some((prefix) => key.startsWith(prefix)) &&
+      env[key] === undefined
+    ) {
       env[key] = process.env[key] as string
     }
   }
@@ -983,7 +994,10 @@ export function loadEnv(
 
       // only keys that start with prefix are exposed to client
       for (const [key, value] of Object.entries(parsed)) {
-        if (key.startsWith(prefix) && env[key] === undefined) {
+        if (
+          prefixes.some((prefix) => key.startsWith(prefix)) &&
+          env[key] === undefined
+        ) {
           env[key] = value
         } else if (key === 'NODE_ENV') {
           // NODE_ENV override in .env file
@@ -992,6 +1006,17 @@ export function loadEnv(
       }
     }
   }
-
   return env
 }
+
+export function resolveEnvPrefix({
+  envPrefix = 'VITE_'
+}: UserConfig): string[] {
+  envPrefix = arraify(envPrefix)
+  if (envPrefix.some((prefix) => prefix === '')) {
+    throw new Error(
+      `envPrefix option contains value '', which could lead unexpected exposure of sensitive information.`
+    )
+  }
+  return envPrefix
+}
diff --git a/packages/vite/src/node/utils.ts b/packages/vite/src/node/utils.ts
index 305e08ce37796a..56baeedb884f8a 100644
--- a/packages/vite/src/node/utils.ts
+++ b/packages/vite/src/node/utils.ts
@@ -512,3 +512,7 @@ export function resolveHostname(
 
   return { host, name }
 }
+
+export function arraify<T>(target: T | T[]): T[] {
+  return Array.isArray(target) ? target : [target]
+}