feat: Allow injecting JavaScript code into subframes (#1365)

Co-authored-by: Geometrically <[email protected]>
Co-authored-by: amrbashir <[email protected]>

.changes/allow-injecting-into-subframes.md

@@ -0,0 +1,5 @@
+---
+"wry": minor
+---
+
+Add `WebViewBuilder::with_initialization_script_for_main_only` to enable injecting JavaScript code into main frame only or all subframes.

src/android/main_pipe.rs

@@ -73,7 +73,7 @@ impl<'a> MainPipe<'a> {
             self.env.set_object_array_element(
               &initialization_scripts_array,
               i as i32,
-              self.env.new_string(script)?,
+              self.env.new_string(script.0)?,
             )?;
           }
 
@@ -427,7 +427,7 @@ pub(crate) struct CreateWebViewAttributes {
   pub autoplay: bool,
   pub on_webview_created: Option<Box<dyn Fn(super::Context) -> JniResult<()> + Send>>,
   pub user_agent: Option<String>,
-  pub initialization_scripts: Vec<String>,
+  pub initialization_scripts: Vec<(String, bool)>,
 }
 
 // SAFETY: only use this when you are sure the span will be dropped on the same thread it was entered

src/android/mod.rs

@@ -260,10 +260,10 @@ impl InnerWebView {
                           QualName::new(None, ns!(html), "script".into()),
                           None,
                         );
-                        script_el.append(NodeRef::new_text(script));
+                        script_el.append(NodeRef::new_text(script.0.as_str()));
                         head.prepend(script_el);
                         if csp.is_some() {
-                          hashes.push(hash_script(script));
+                          hashes.push(hash_script(script.0.as_str()));
                         }
                       }
                     });

src/lib.rs

@@ -357,15 +357,19 @@ pub struct WebViewAttributes<'a> {
   /// - **Windows:** the string can not be larger than 2 MB (2 * 1024 * 1024 bytes) in total size
   pub html: Option<String>,
 
-  /// Initialize javascript code when loading new pages. When webview load a new page, this
-  /// initialization code will be executed. It is guaranteed that code is executed before
-  /// `window.onload`.
+  /// A list of initialization javascript scripts to run when loading new pages.
+  /// When webview load a new page, this initialization code will be executed.
+  /// It is guaranteed that code is executed before `window.onload`.
+  ///
+  /// Second parameter represents if script should be added to main frame only or sub frames also.
+  /// `true` for main frame only, `false` for sub frames.
   ///
   /// ## Platform-specific
   ///
+  /// - **Windows**: scripts are injected into sub frames.
   /// - **Android:** The Android WebView does not provide an API for initialization scripts,
   /// so we prepend them to each HTML head. They are only implemented on custom protocol URLs.
-  pub initialization_scripts: Vec<String>,
+  pub initialization_scripts: Vec<(String, bool)>,
 
   /// A list of custom loading protocols with pairs of scheme uri string and a handling
   /// closure.
@@ -693,16 +697,28 @@ impl<'a> WebViewBuilder<'a> {
   ///
   /// ## Platform-specific
   ///
+  /// - **Windows:** scripts are added to subframes as well.
   /// - **Android:** When [addDocumentStartJavaScript] is not supported,
   /// we prepend them to each HTML head (implementation only supported on custom protocol URLs).
   /// For remote URLs, we use [onPageStarted] which is not guaranteed to run before other scripts.
   ///
   /// [addDocumentStartJavaScript]: https://developer.android.com/reference/androidx/webkit/WebViewCompat#addDocumentStartJavaScript(android.webkit.WebView,java.lang.String,java.util.Set%3Cjava.lang.String%3E)
   /// [onPageStarted]: https://developer.android.com/reference/android/webkit/WebViewClient#onPageStarted(android.webkit.WebView,%20java.lang.String,%20android.graphics.Bitmap)
   pub fn with_initialization_script(self, js: &str) -> Self {
+    self.with_initialization_script_for_main_only(js, true)
+  }
+
+  /// Same as [`with_initialization_script`](Self::with_initialization_script) but with option to inject into main frame only or sub frames.
+  ///
+  /// ## Platform-specific:
+  ///
+  /// - **Windows:** scripts are always added to subframes regardless of the option.
+  pub fn with_initialization_script_for_main_only(self, js: &str, main_only: bool) -> Self {
     self.and_then(|mut b| {
       if !js.is_empty() {
-        b.attrs.initialization_scripts.push(js.to_string());
+        b.attrs
+          .initialization_scripts
+          .push((js.to_string(), main_only));
       }
       Ok(b)
     })

src/webkitgtk/mod.rs

@@ -304,11 +304,11 @@ impl InnerWebView {
     };
 
     // Initialize message handler
-    w.init("Object.defineProperty(window, 'ipc', { value: Object.freeze({ postMessage: function(x) { window.webkit.messageHandlers['ipc'].postMessage(x) } }) })")?;
+    w.init("Object.defineProperty(window, 'ipc', { value: Object.freeze({ postMessage: function(x) { window.webkit.messageHandlers['ipc'].postMessage(x) } }) })", true)?;
 
     // Initialize scripts
-    for js in attributes.initialization_scripts {
-      w.init(&js)?;
+    for (js, for_main_only) in attributes.initialization_scripts {
+      w.init(&js, for_main_only)?;
     }
 
     // Run pending webview.eval() scripts once webview loads.
@@ -614,12 +614,15 @@ impl InnerWebView {
     Ok(())
   }
 
-  fn init(&self, js: &str) -> Result<()> {
+  fn init(&self, js: &str, for_main_only: bool) -> Result<()> {
     if let Some(manager) = self.webview.user_content_manager() {
       let script = UserScript::new(
         js,
-        // TODO: feature to allow injecting into subframes
-        UserContentInjectedFrames::TopFrame,
+        if for_main_only {
+          UserContentInjectedFrames::TopFrame
+        } else {
+          UserContentInjectedFrames::AllFrames
+        },
         UserScriptInjectionTime::Start,
         &[],
         &[],

src/webview2/mod.rs

@@ -440,9 +440,9 @@ impl InnerWebView {
       };
     }
 
-    // Initialize scripts
-    for js in attributes.initialization_scripts {
-      Self::add_script_to_execute_on_document_created(&webview, js)?;
+    // Initialize main frame scripts
+    for (js, for_main_only) in attributes.initialization_scripts {
+      Self::add_script_to_execute_on_document_created(&webview, js, for_main_only)?;
     }
 
     // Enable clipboard
@@ -589,6 +589,8 @@ impl InnerWebView {
     if let Some(on_page_load_handler) = attributes.on_page_load_handler.take() {
       let on_page_load_handler = Rc::new(on_page_load_handler);
       let on_page_load_handler_ = on_page_load_handler.clone();
+      let scripts = attributes.initialization_scripts.clone();
+
       webview.add_ContentLoading(
         &ContentLoadingEventHandler::create(Box::new(move |webview, _| {
           let Some(webview) = webview else {
@@ -597,6 +599,12 @@ impl InnerWebView {
 
           on_page_load_handler_(PageLoadEvent::Started, Self::url_from_webview(&webview)?);
 
+          for (script, inject_into_sub_frames) in &scripts {
+            if *inject_into_sub_frames {
+              Self::execute_script(&webview, script.clone(), |_| ())?;
+            }
+          }
+
           Ok(())
         })),
         token,
@@ -757,6 +765,7 @@ impl InnerWebView {
       String::from(
         r#"Object.defineProperty(window, 'ipc', { value: Object.freeze({ postMessage: s=> window.chrome.webview.postMessage(s) }) });"#,
       ),
+      true,
     )?;
 
     let ipc_handler = attributes.ipc_handler.take();
@@ -1140,9 +1149,12 @@ impl InnerWebView {
     );
   }
 
-  // TODO: feature to allow injecting into (specific) subframes
   #[inline]
-  fn add_script_to_execute_on_document_created(webview: &ICoreWebView2, js: String) -> Result<()> {
+  fn add_script_to_execute_on_document_created(
+    webview: &ICoreWebView2,
+    js: String,
+    _for_main_only: bool,
+  ) -> Result<()> {
     let webview = webview.clone();
     AddScriptToExecuteOnDocumentCreatedCompletedHandler::wait_for_async_operation(
       Box::new(move |handler| unsafe {

src/wkwebview/mod.rs

@@ -480,9 +480,10 @@ impl InnerWebView {
 r#"Object.defineProperty(window, 'ipc', {
   value: Object.freeze({postMessage: function(s) {window.webkit.messageHandlers.ipc.postMessage(s);}})
 });"#,
+      true
       );
-      for js in attributes.initialization_scripts {
-        w.init(&js);
+      for (js, for_main_only) in attributes.initialization_scripts {
+        w.init(&js, for_main_only);
       }
 
       // Set user agent
@@ -602,16 +603,15 @@ r#"Object.defineProperty(window, 'ipc', {
     Ok(())
   }
 
-  fn init(&self, js: &str) {
+  fn init(&self, js: &str, for_main_only: bool) {
     // Safety: objc runtime calls are unsafe
     unsafe {
       let userscript = WKUserScript::alloc();
-      // TODO: feature to allow injecting into subframes
       let script = WKUserScript::initWithSource_injectionTime_forMainFrameOnly(
         userscript,
         &NSString::from_str(js),
         WKUserScriptInjectionTime::AtDocumentStart,
-        true,
+        for_main_only,
       );
       self.manager.addUserScript(&script);
     }