fix(core): share webcontext between webviews (#11043)

* fix(core): share webcontext between webviews

closes #10981

* update wry version

* Update crates/tauri-runtime-wry/src/lib.rs [skip ci]

* on linux, only register protocol once per context

---------

Co-authored-by: Lucas Fernandes Nogueira <lucas@tauri.app>
Co-authored-by: Lucas Nogueira <lucas@tauri.studio>

Author: amrbashir

.changes/share-webcontext.md

@@ -0,0 +1,7 @@
+---
+"tauri": "patch:bug"
+"tauri-runtime-wry": "patch:bug"
+---
+
+Fix `localStorage` not shared between webviews that use the same data directory.
+

crates/tauri-runtime-wry/src/lib.rs

@@ -66,8 +66,8 @@ use tauri_utils::TitleBarStyle;
 use tauri_utils::{config::WindowConfig, Theme};
 use url::Url;
 use wry::{
-  DragDropEvent as WryDragDropEvent, ProxyConfig, ProxyEndpoint, WebContext, WebView,
-  WebViewBuilder,
+  DragDropEvent as WryDragDropEvent, ProxyConfig, ProxyEndpoint, WebContext as WryWebContext,
+  WebView, WebViewBuilder,
 };
 
 pub use tao;
@@ -101,7 +101,7 @@ use std::{
   cell::RefCell,
   collections::{
     hash_map::Entry::{Occupied, Vacant},
-    BTreeMap, HashMap,
+    BTreeMap, HashMap, HashSet,
   },
   fmt,
   ops::Deref,
@@ -131,6 +131,15 @@ mod undecorated_resizing;
 mod webview;
 pub use webview::Webview;
 
+#[derive(Debug)]
+pub struct WebContext {
+  pub inner: WryWebContext,
+  pub referenced_by_webviews: HashSet<String>,
+  // on Linux the custom protocols are associated with the context
+  // and you cannot register a URI scheme more than once
+  pub registered_custom_protocols: HashSet<String>,
+}
+
 pub type WebContextStore = Arc<Mutex<HashMap<Option<PathBuf>, WebContext>>>;
 // window
 pub type WindowEventHandler = Box<dyn Fn(&WindowEvent) + Send>;
@@ -216,7 +225,6 @@ pub struct Context<T: UserEvent> {
   next_webview_id: Arc<AtomicU32>,
   next_window_event_id: Arc<AtomicU32>,
   next_webview_event_id: Arc<AtomicU32>,
-  next_webcontext_id: Arc<AtomicU32>,
 }
 
 impl<T: UserEvent> Context<T> {
@@ -246,10 +254,6 @@ impl<T: UserEvent> Context<T> {
   fn next_webview_event_id(&self) -> u32 {
     self.next_webview_event_id.fetch_add(1, Ordering::Relaxed)
   }
-
-  fn next_webcontext_id(&self) -> u32 {
-    self.next_webcontext_id.fetch_add(1, Ordering::Relaxed)
-  }
 }
 
 impl<T: UserEvent> Context<T> {
@@ -2048,7 +2052,15 @@ impl Deref for WebviewWrapper {
 impl Drop for WebviewWrapper {
   fn drop(&mut self) {
     if Rc::get_mut(&mut self.inner).is_some() {
-      self.context_store.lock().unwrap().remove(&self.context_key);
+      let mut context_store = self.context_store.lock().unwrap();
+
+      if let Some(web_context) = context_store.get_mut(&self.context_key) {
+        web_context.referenced_by_webviews.remove(&self.label);
+
+        if web_context.referenced_by_webviews.is_empty() {
+          context_store.remove(&self.context_key);
+        }
+      }
     }
   }
 }
@@ -2357,7 +2369,6 @@ impl<T: UserEvent> Wry<T> {
       next_webview_id: Default::default(),
       next_window_event_id: Default::default(),
       next_webview_event_id: Default::default(),
-      next_webcontext_id: Default::default(),
     };
 
     Ok(Self {
@@ -4104,16 +4115,6 @@ fn create_webview<T: UserEvent>(
     ipc_handler,
   ));
 
-  for (scheme, protocol) in uri_scheme_protocols {
-    webview_builder =
-      webview_builder.with_asynchronous_custom_protocol(scheme, move |request, responder| {
-        protocol(
-          request,
-          Box::new(move |response| responder.respond(response)),
-        )
-      });
-  }
-
   for script in webview_attributes.initialization_scripts {
     webview_builder = webview_builder.with_initialization_script(&script);
   }
@@ -4124,30 +4125,59 @@ fn create_webview<T: UserEvent>(
     .lock()
     .expect("poisoned WebContext store");
   let is_first_context = web_context.is_empty();
+  // the context must be stored on the HashMap because it must outlive the WebView on macOS
   let automation_enabled = std::env::var("TAURI_WEBVIEW_AUTOMATION").as_deref() == Ok("true");
-  let web_context_key = // force a unique WebContext when automation is false;
-    // the context must be stored on the HashMap because it must outlive the WebView on macOS
-    if automation_enabled {
-      webview_attributes.data_directory.clone()
-    } else {
-      // unique key
-      let key = context.next_webcontext_id().to_string().into();
-      Some(key)
-    };
+  let web_context_key = webview_attributes.data_directory;
   let entry = web_context.entry(web_context_key.clone());
   let web_context = match entry {
-    Occupied(occupied) => occupied.into_mut(),
+    Occupied(occupied) => {
+      let occupied = occupied.into_mut();
+      occupied.referenced_by_webviews.insert(label.clone());
+      occupied
+    }
     Vacant(vacant) => {
-      let mut web_context = WebContext::new(webview_attributes.data_directory);
+      let mut web_context = WryWebContext::new(web_context_key.clone());
       web_context.set_allows_automation(if automation_enabled {
         is_first_context
       } else {
         false
       });
-      vacant.insert(web_context)
+      vacant.insert(WebContext {
+        inner: web_context,
+        referenced_by_webviews: [label.clone()].into(),
+        registered_custom_protocols: HashSet::new(),
+      })
     }
   };
 
+  for (scheme, protocol) in uri_scheme_protocols {
+    // on Linux the custom protocols are associated with the web context
+    // and you cannot register a scheme more than once
+    if cfg!(any(
+      target_os = "linux",
+      target_os = "dragonfly",
+      target_os = "freebsd",
+      target_os = "netbsd",
+      target_os = "openbsd"
+    )) {
+      if web_context.registered_custom_protocols.contains(&scheme) {
+        continue;
+      }
+
+      web_context
+        .registered_custom_protocols
+        .insert(scheme.clone());
+    }
+
+    webview_builder =
+      webview_builder.with_asynchronous_custom_protocol(scheme, move |request, responder| {
+        protocol(
+          request,
+          Box::new(move |response| responder.respond(response)),
+        )
+      });
+  }
+
   if webview_attributes.clipboard {
     webview_builder.attrs.clipboard = true;
   }
@@ -4175,7 +4205,7 @@ fn create_webview<T: UserEvent>(
   }
 
   let webview = webview_builder
-    .with_web_context(web_context)
+    .with_web_context(&mut web_context.inner)
     .build()
     .map_err(|e| Error::CreateWebview(Box::new(e)))?;