From 3e52440d28b639eb60d9d1ea815081e9b989d0e8 Mon Sep 17 00:00:00 2001
From: Tom Gao <tom@tomgao.cc>
Date: Mon, 3 Jun 2024 22:18:19 +0800
Subject: [PATCH] feat: add custom csp option to webview

---
 README.md                   |  22 ++++++-
 README.zh_CN.md             |  22 ++++++-
 examples/vue/vite.config.ts |   6 +-
 src/index.ts                | 117 ++++++++++++++++++++----------------
 src/types.ts                |  16 ++++-
 5 files changed, 127 insertions(+), 56 deletions(-)

diff --git a/README.md b/README.md
index c3dc4db..afc55c4 100644
--- a/README.md
+++ b/README.md
@@ -213,7 +213,7 @@ function __getWebviewHtml__(
 | --- | --- | --- | --- |
 | recommended | `boolean` | `true` | This option is intended to provide recommended default parameters and behavior. |
 | extension | [ExtensionOptions](#ExtensionOptions) |  | Configuration options for the vscode extension. |
-| webview | `boolean` | `false` | Inject [@tomjs/vscode-extension-webview](https://github.com/tomjs/vscode-extension-webview) into vscode extension code and web client code, so that webview can support HMR during the development stage. |
+| webview | `boolean` \| `string` \| [WebviewOption](#WebviewOption) | `__getWebviewHtml__` | Inject html code |
 
 **Notice**
 
@@ -222,6 +222,16 @@ The `recommended` option is used to set the default configuration and behavior,
 - The output directory is based on the `build.outDir` parameter of `vite`, and outputs `extension` and `src` to `dist/extension` and `dist/webview` respectively.
 - Other behaviors to be implemented
 
+**Webview**
+
+Inject [@tomjs/vscode-extension-webview](https://github.com/tomjs/vscode-extension-webview) into vscode extension code and web client code, so that `webview` can support `HMR` during the development stage.
+
+- vite serve
+  - extension: Inject `import __getWebviewHtml__ from '@tomjs/vscode-extension-webview';` above the file that calls the `__getWebviewHtml__` method
+  - web: Add `<script>` tag to index.html and inject `@tomjs/vscode-extension-webview/client` code
+- vite build
+  - extension: Inject `import __getWebviewHtml__ from '@tomjs/vite-plugin-vscode-inject';` above the file that calls the `__getWebviewHtml__` method If is string, will set inject method name. Default is `__getWebviewHtml__`.
+
 ### ExtensionOptions
 
 Based on [Options](https://paka.dev/npm/tsup) of [tsup](https://tsup.egoist.dev/), some default values are added for ease of use.
@@ -232,6 +242,16 @@ Based on [Options](https://paka.dev/npm/tsup) of [tsup](https://tsup.egoist.dev/
 | outDir | `string` | `dist-extension/main` | The output directory for the vscode extension file |
 | onSuccess | `() => Promise<void \| undefined \| (() => void \| Promise<void>)>` | `undefined` | A function that will be executed after the build succeeds. |
 
+### WebviewOption
+
+| Property | Type | Default | Description |
+| --- | --- | --- | --- |
+| name | `string` | `__getWebviewHtml__` | The inject method name |
+| csp | `string` | `<meta http-equiv="Content-Security-Policy" content="default-src 'none'; style-src {{cspSource}} 'unsafe-inline'; script-src 'nonce-{{nonce}}' 'unsafe-eval';">` | The `CSP` meta for the webview |
+
+- `{{cspSource}}`: [webview.cspSource](https://code.visualstudio.com/api/references/vscode-api#Webview)
+- `{{nonce}}`: uuid
+
 ### Additional Information
 
 - Default values for `extension` when the relevant parameters are not configured
diff --git a/README.zh_CN.md b/README.zh_CN.md
index 0ddda33..22c751e 100644
--- a/README.zh_CN.md
+++ b/README.zh_CN.md
@@ -217,7 +217,7 @@ function __getWebviewHtml__(
 | --- | --- | --- | --- |
 | recommended | `boolean` | `true` | 这个选项是为了提供推荐的默认参数和行为 |
 | extension | [ExtensionOptions](#ExtensionOptions) |  | vscode extension 可选配置 |
-| webview | `boolean` | `false` | 在vscode扩展代码和Web客户端代码中注入[@tomjs/vscode-extension-webview](https://github.com/tomjs/vscode-extension-webview)，以便webview在开发阶段可以支持HMR。 |
+| webview | `boolean` \| `string` \| [WebviewOption](#WebviewOption) | `__getWebviewHtml__` | 注入 html 代码 |
 
 **Notice**
 
@@ -227,6 +227,16 @@ function __getWebviewHtml__(
 
 - 其他待实现的行为
 
+**Webview**
+
+在 vscode 扩展代码和 web 客户端代码中注入 [@tomjs/vscode-extension-webview](https://github.com/tomjs/vscode-extension-webview)，使 `webview` 在开发阶段能够支持 `HMR`。
+
+- vite serve
+  - extension: 在调用 `__getWebviewHtml__` 方法的文件上方注入 `import __getWebviewHtml__ from '@tomjs/vscode-extension-webview';`
+  - web: 在 index.html 中添加 `<script>` 标签，注入 `@tomjs/vscode-extension-webview/client` 代码
+- vite build
+  - extension: 在调用 `__getWebviewHtml__` 方法的文件上方注入 `import __getWebviewHtml__ from '@tomjs/vite-plugin-vscode-inject';` 如果为字符串，则设置注入方法名，默认为 `__getWebviewHtml__`。
+
 ### ExtensionOptions
 
 继承自 [tsup](https://tsup.egoist.dev/) 的 [Options](https://paka.dev/npm/tsup)，添加了一些默认值，方便使用。
@@ -237,6 +247,16 @@ function __getWebviewHtml__(
 | outDir | `string` | `dist-extension/main` | 输出文件夹 |
 | onSuccess | `() => Promise<void \| undefined \| (() => void \| Promise<void>)>` | `undefined` | 构建成功后运行的回调函数 |
 
+### WebviewOption
+
+| 参数名 | 类型 | 默认值 | 说明 |
+| --- | --- | --- | --- |
+| name | `string` | `__getWebviewHtml__` | 注入的方法名 |
+| csp | `string` | `<meta http-equiv="Content-Security-Policy" content="default-src 'none'; style-src {{cspSource}} 'unsafe-inline'; script-src 'nonce-{{nonce}}' 'unsafe-eval';">` | webview 的 `CSP` |
+
+- `{{cspSource}}`: [webview.cspSource](https://code.visualstudio.com/api/references/vscode-api#Webview)
+- `{{nonce}}`: uuid
+
 ### 补充说明
 
 - `extension` 未配置相关参数时的默认值
diff --git a/examples/vue/vite.config.ts b/examples/vue/vite.config.ts
index 596b69d..a7608aa 100644
--- a/examples/vue/vite.config.ts
+++ b/examples/vue/vite.config.ts
@@ -12,6 +12,10 @@ export default defineConfig({
         },
       },
     }),
-    vscode(),
+    vscode({
+      webview: {
+        // csp: '<meta http-equiv="Content-Security-Policy" />',
+      },
+    }),
   ],
 });
diff --git a/src/index.ts b/src/index.ts
index 35b8df9..68688b4 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,5 +1,5 @@
 import type { Plugin, ResolvedConfig, UserConfig } from 'vite';
-import type { ExtensionOptions, PluginOptions } from './types';
+import type { ExtensionOptions, PluginOptions, WebviewOption } from './types';
 import fs from 'node:fs';
 import path from 'node:path';
 import { cwd } from 'node:process';
@@ -80,48 +80,55 @@ function preMergeOptions(options?: PluginOptions): PluginOptions {
 
   opts.extension = opt;
 
-  if (opts.webview === true) {
-    opts.webview = WEBVIEW_METHOD_NAME;
+  if (opts.webview !== false) {
+    let name = WEBVIEW_METHOD_NAME;
+    if (typeof opts.webview === 'string') {
+      name = opts.webview ?? WEBVIEW_METHOD_NAME;
+    }
+    opts.webview = Object.assign({ name }, opts.webview);
   }
-  opts.webview = opts.webview ?? WEBVIEW_METHOD_NAME;
 
   return opts;
 }
 
 const prodCachePkgName = `${PACKAGE_NAME}-inject`;
-function genProdWebviewCode(cache: Record<string, string>) {
+function genProdWebviewCode(cache: Record<string, string>, webview?: WebviewOption) {
+  webview = Object.assign({}, webview);
+
   const prodCacheFolder = path.join(cwd(), 'node_modules', prodCachePkgName);
   emptyPath(prodCacheFolder);
   const destFile = path.join(prodCacheFolder, 'index.ts');
 
-  function handleHtmlCode(html) {
+  function handleHtmlCode(html: string) {
     const root = htmlParser(html);
     const head = root.querySelector('head')!;
     if (!head) {
       root?.insertAdjacentHTML('beforeend', '<head></head>');
     }
 
-    head.insertAdjacentHTML(
-      'afterbegin',
-      `\n<meta http-equiv="Content-Security-Policy" content="default-src 'none'; style-src {{cspSource}} 'unsafe-inline'; script-src 'nonce-{{nonce}}' 'unsafe-eval';">`,
-    );
-
-    const tags = {
-      script: 'src',
-      link: 'href',
-    };
-
-    Object.keys(tags).forEach(tag => {
-      const elements = root.querySelectorAll(tag);
-      elements.forEach(element => {
-        const attr = element.getAttribute(tags[tag]);
-        if (attr) {
-          element.setAttribute(tags[tag], `{{baseUri}}${attr}`);
-        }
-
-        element.setAttribute('nonce', '{{nonce}}');
+    const csp =
+      webview?.csp ||
+      `<meta http-equiv="Content-Security-Policy" content="default-src 'none'; style-src {{cspSource}} 'unsafe-inline'; script-src 'nonce-{{nonce}}' 'unsafe-eval';">`;
+    head.insertAdjacentHTML('afterbegin', csp);
+
+    if (csp && csp.includes('{{nonce}}')) {
+      const tags = {
+        script: 'src',
+        link: 'href',
+      };
+
+      Object.keys(tags).forEach(tag => {
+        const elements = root.querySelectorAll(tag);
+        elements.forEach(element => {
+          const attr = element.getAttribute(tags[tag]);
+          if (attr) {
+            element.setAttribute(tags[tag], `{{baseUri}}${attr}`);
+          }
+
+          element.setAttribute('nonce', '{{nonce}}');
+        });
       });
-    });
+    }
 
     return root.toString();
   }
@@ -149,7 +156,7 @@ export default function getWebviewHtml(webview: Webview, context: ExtensionConte
   const nonce = uuid();
   const baseUri = webview.asWebviewUri(Uri.joinPath(context.extensionUri, (process.env.VITE_WEBVIEW_DIST || 'dist')));
   const html = htmlCode[inputName || 'index'] || '';
-  return html.replaceAll('{{cspSource}}',webview.cspSource).replaceAll('{{nonce}}', nonce).replaceAll('{{baseUri}}', baseUri);
+  return html.replaceAll('{{cspSource}}', webview.cspSource).replaceAll('{{nonce}}', nonce).replaceAll('{{baseUri}}', baseUri);
 }
   `;
   fs.writeFileSync(destFile, code, { encoding: 'utf8' });
@@ -236,32 +243,36 @@ export function useVSCodePlugin(options?: PluginOptions): Plugin[] {
 
           let buildCount = 0;
 
+          const webview = opts?.webview as WebviewOption;
+
           const { onSuccess: _onSuccess, ...tsupOptions } = opts.extension || {};
           await tsupBuild(
             merge(tsupOptions, {
               watch: true,
               env,
               silent: true,
-              esbuildPlugins: [
-                {
-                  name: '@tomjs:vscode:inject',
-                  setup(build) {
-                    build.onLoad({ filter: /\.ts$/ }, async args => {
-                      const file = fs.readFileSync(args.path, 'utf-8');
-                      if (file.includes(`${opts.webview}(`)) {
-                        return {
-                          contents:
-                            `import ${opts.webview} from '@tomjs/vscode-extension-webview';\n` +
-                            file,
-                          loader: 'ts',
-                        };
-                      }
-
-                      return {};
-                    });
-                  },
-                },
-              ],
+              esbuildPlugins: !webview
+                ? []
+                : [
+                    {
+                      name: '@tomjs:vscode:inject',
+                      setup(build) {
+                        build.onLoad({ filter: /\.ts$/ }, async args => {
+                          const file = fs.readFileSync(args.path, 'utf-8');
+                          if (file.includes(`${webview.name}(`)) {
+                            return {
+                              contents:
+                                `import ${webview.name} from '@tomjs/vscode-extension-webview';\n` +
+                                file,
+                              loader: 'ts',
+                            };
+                          }
+
+                          return {};
+                        });
+                      },
+                    },
+                  ],
               async onSuccess() {
                 if (typeof _onSuccess === 'function') {
                   await _onSuccess();
@@ -305,8 +316,10 @@ export function useVSCodePlugin(options?: PluginOptions): Plugin[] {
       },
       closeBundle() {
         let webviewPath: string;
-        if (opts.webview) {
-          webviewPath = genProdWebviewCode(prodHtmlCache);
+
+        const webview = opts?.webview as WebviewOption;
+        if (webview) {
+          webviewPath = genProdWebviewCode(prodHtmlCache, webview);
         }
 
         let outDir = buildConfig.build.outDir.replace(cwd(), '').replaceAll('\\', '/');
@@ -325,7 +338,7 @@ export function useVSCodePlugin(options?: PluginOptions): Plugin[] {
           merge(tsupOptions, {
             env,
             silent: true,
-            esbuildPlugins: !opts.webview
+            esbuildPlugins: !webview
               ? []
               : [
                   {
@@ -333,9 +346,9 @@ export function useVSCodePlugin(options?: PluginOptions): Plugin[] {
                     setup(build) {
                       build.onLoad({ filter: /\.ts$/ }, async args => {
                         const file = fs.readFileSync(args.path, 'utf-8');
-                        if (file.includes(`${opts.webview}(`)) {
+                        if (file.includes(`${webview.name}(`)) {
                           return {
-                            contents: `import ${opts.webview} from \`${webviewPath}\`;\n` + file,
+                            contents: `import ${webview.name} from \`${webviewPath}\`;\n` + file,
                             loader: 'ts',
                           };
                         }
diff --git a/src/types.ts b/src/types.ts
index d584f47..bde9ed4 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -32,6 +32,20 @@ export interface ExtensionOptions
   onSuccess?: () => Promise<void | undefined | (() => void | Promise<void>)>;
 }
 
+/**
+ * vscode webview options.
+ */
+export interface WebviewOption {
+  /**
+   * The method name to inject. Default is '__getWebviewHtml__'
+   */
+  name?: string;
+  /**
+   * The CSP meta for the webview. Default is `<meta http-equiv="Content-Security-Policy" content="default-src 'none'; style-src {{cspSource}} 'unsafe-inline'; script-src 'nonce-{{nonce}}' 'unsafe-eval';">`
+   */
+  csp?: string;
+}
+
 /**
  * vite plugin options.
  */
@@ -74,7 +88,7 @@ export interface PluginOptions {
    * </html>
    * ```
    */
-  webview?: boolean | string;
+  webview?: boolean | string | WebviewOption;
   /**
    * extension vite config.
    */