feat(core): inject invoke key on <script type="module"> (#2120)

lucasfernog · web-flow · commit f03eea9c9b96 · 2021-06-29T20:40:41.000-03:00
diff --git a/.changes/inject-invoke-key-module-script.md b/.changes/inject-invoke-key-module-script.md
@@ -0,0 +1,7 @@
+---
+"tauri": patch
+"tauri-codegen": patch
+"tauri-utils": patch
+---
+
+Inject invoke key on `script` tags with `type="module"`.
diff --git a/core/tauri-codegen/Cargo.toml b/core/tauri-codegen/Cargo.toml
@@ -21,3 +21,4 @@ tauri-utils = { version = "1.0.0-beta.1", path = "../tauri-utils", features = [
 thiserror = "1"
 walkdir = "2"
 zstd = "0.9"
+kuchiki = "0.8"
diff --git a/core/tauri-codegen/src/embedded_assets.rs b/core/tauri-codegen/src/embedded_assets.rs
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-License-Identifier: MIT
 
+use kuchiki::traits::*;
 use proc_macro2::TokenStream;
 use quote::{quote, ToTokens, TokenStreamExt};
 use std::{
@@ -10,7 +11,10 @@ use std::{
   fs::File,
   path::{Path, PathBuf},
 };
-use tauri_utils::{assets::AssetKey, html::inject_csp};
+use tauri_utils::{
+  assets::AssetKey,
+  html::{inject_csp, inject_invoke_key_token},
+};
 use thiserror::Error;
 use walkdir::WalkDir;
 
@@ -170,45 +174,47 @@ impl EmbeddedAssets {
       path: path.to_owned(),
       error,
     })?;
-    if let Some(csp) = &options.csp {
-      if path.extension() == Some(OsStr::new("html")) {
-        input = inject_csp(String::from_utf8_lossy(&input).into_owned(), csp)
+    if path.extension() == Some(OsStr::new("html")) {
+      let mut document = kuchiki::parse_html().one(String::from_utf8_lossy(&input).into_owned());
+      if let Some(csp) = &options.csp {
+        inject_csp(&mut document, csp);
+      }
+      inject_invoke_key_token(&mut document);
+      input = document.to_string().as_bytes().to_vec();
+    } else {
+      let is_javascript = ["js", "cjs", "mjs"]
+        .iter()
+        .any(|e| path.extension() == Some(OsStr::new(e)));
+      if is_javascript {
+        let js = String::from_utf8_lossy(&input).into_owned();
+        input = if [
+          "import{", "import*", "import ", "export{", "export*", "export ",
+        ]
+        .iter()
+        .any(|t| js.contains(t))
+        {
+          format!(
+            r#"
+              const __TAURI_INVOKE_KEY__ = __TAURI__INVOKE_KEY_TOKEN__;
+              {}
+            "#,
+            js
+          )
+          .as_bytes()
+          .to_vec()
+        } else {
+          format!(
+            r#"(function () {{
+              const __TAURI_INVOKE_KEY__ = __TAURI__INVOKE_KEY_TOKEN__;
+              {}
+            }})()"#,
+            js
+          )
           .as_bytes()
-          .to_vec();
+          .to_vec()
+        };
       }
     }
-    let is_javascript = ["js", "cjs", "mjs"]
-      .iter()
-      .any(|e| path.extension() == Some(OsStr::new(e)));
-    if is_javascript {
-      let js = String::from_utf8_lossy(&input).into_owned();
-      input = if [
-        "import{", "import*", "import ", "export{", "export*", "export ",
-      ]
-      .iter()
-      .any(|t| js.contains(t))
-      {
-        format!(
-          r#"
-            const __TAURI_INVOKE_KEY__ = __TAURI__INVOKE_KEY_TOKEN__;
-            {}
-          "#,
-          js
-        )
-        .as_bytes()
-        .to_vec()
-      } else {
-        format!(
-          r#"(function () {{
-            const __TAURI_INVOKE_KEY__ = __TAURI__INVOKE_KEY_TOKEN__;
-            {}
-          }})()"#,
-          js
-        )
-        .as_bytes()
-        .to_vec()
-      };
-    }
 
     // we must canonicalize the base of our paths to allow long paths on windows
     let out_dir = std::env::var("OUT_DIR")
diff --git a/core/tauri-utils/src/html.rs b/core/tauri-utils/src/html.rs
@@ -2,17 +2,55 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-License-Identifier: MIT
 
-use html5ever::{
-  interface::QualName,
-  namespace_url, ns,
-  tendril::{fmt::UTF8, NonAtomic, Tendril},
-  LocalName,
-};
-use kuchiki::{traits::*, Attribute, ExpandedName, NodeRef};
+use html5ever::{interface::QualName, namespace_url, ns, LocalName};
+use kuchiki::{Attribute, ExpandedName, NodeRef};
+
+/// Injects the invoke key token to each script on the document.
+///
+/// The invoke key token is replaced at runtime with the actual invoke key value.
+pub fn inject_invoke_key_token(document: &mut NodeRef) {
+  let mut targets = vec![];
+  if let Ok(scripts) = document.select("script") {
+    for target in scripts {
+      targets.push(target);
+    }
+    for target in targets {
+      let node = target.as_node();
+      let element = node.as_element().unwrap();
+
+      let attrs = element.attributes.borrow();
+      // if the script is external (has `src`) or its type is not "module", we won't inject the token
+      if attrs.get("src").is_some() || attrs.get("type") != Some("module") {
+        continue;
+      }
+
+      let replacement_node = NodeRef::new_element(
+        QualName::new(None, ns!(html), "script".into()),
+        element
+          .attributes
+          .borrow()
+          .clone()
+          .map
+          .into_iter()
+          .collect::<Vec<_>>(),
+      );
+      let script = node.text_contents();
+      replacement_node.append(NodeRef::new_text(format!(
+        r#"
+          const __TAURI_INVOKE_KEY__ = __TAURI__INVOKE_KEY_TOKEN__;
+          {}
+        "#,
+        script
+      )));
+
+      node.insert_after(replacement_node);
+      node.detach();
+    }
+  }
+}
 
 /// Injects a content security policy to the HTML.
-pub fn inject_csp<H: Into<Tendril<UTF8, NonAtomic>>>(html: H, csp: &str) -> String {
-  let document = kuchiki::parse_html().one(html);
+pub fn inject_csp(document: &mut NodeRef, csp: &str) {
   if let Ok(ref head) = document.select_first("head") {
     head.as_node().append(create_csp_meta_tag(csp));
   } else {
@@ -23,7 +61,6 @@ pub fn inject_csp<H: Into<Tendril<UTF8, NonAtomic>>>(html: H, csp: &str) -> Stri
     head.append(create_csp_meta_tag(csp));
     document.prepend(head);
   }
-  document.to_string()
 }
 
 fn create_csp_meta_tag(csp: &str) -> NodeRef {
@@ -50,17 +87,19 @@ fn create_csp_meta_tag(csp: &str) -> NodeRef {
 
 #[cfg(test)]
 mod tests {
+  use kuchiki::traits::*;
   #[test]
   fn csp() {
     let htmls = vec![
       "<html><head></head></html>".to_string(),
       "<html></html>".to_string(),
     ];
     for html in htmls {
+      let mut document = kuchiki::parse_html().one(html);
       let csp = "default-src 'self'; img-src https://*; child-src 'none';";
-      let new = super::inject_csp(html, csp);
+      super::inject_csp(&mut document, csp);
       assert_eq!(
-        new,
+        document.to_string(),
         format!(
           r#"<html><head><meta content="{}" http-equiv="Content-Security-Policy"></head><body></body></html>"#,
           csp
diff --git a/core/tauri/src/manager.rs b/core/tauri/src/manager.rs
@@ -459,6 +459,7 @@ impl<P: Params> WindowManager<P> {
         };
         let is_javascript =
           path.ends_with(".js") || path.ends_with(".cjs") || path.ends_with(".mjs");
+        let is_html = path.ends_with(".html");
 
         let asset_response = assets
           .get(&path)
@@ -471,16 +472,17 @@ impl<P: Params> WindowManager<P> {
           .map(Cow::into_owned);
         match asset_response {
           Ok(asset) => {
-            if is_javascript {
-              let js = String::from_utf8_lossy(&asset).into_owned();
+            if is_javascript || is_html {
+              let contents = String::from_utf8_lossy(&asset).into_owned();
               Ok(
-                js.replacen(
-                  "__TAURI__INVOKE_KEY_TOKEN__",
-                  &manager.generate_invoke_key().to_string(),
-                  1,
-                )
-                .as_bytes()
-                .to_vec(),
+                contents
+                  .replacen(
+                    "__TAURI__INVOKE_KEY_TOKEN__",
+                    &manager.generate_invoke_key().to_string(),
+                    1,
+                  )
+                  .as_bytes()
+                  .to_vec(),
               )
             } else {
               Ok(asset)
