From 5c97db96bdbf7bcc5c40076986c54a0ae8964d1b Mon Sep 17 00:00:00 2001
From: Jonas Osburg <accounts@jnugh.de>
Date: Mon, 15 Apr 2024 16:41:49 +0200
Subject: [PATCH] feat(websocket): Add TLS configuration (#925)

* Allow TLS configuration

* refactor: Added builder pattern

---------

Co-authored-by: FabianLars <fabianlars@fabianlars.de>
---
 plugins/websocket/src/lib.rs | 51 +++++++++++++++++++++++++++++-------
 1 file changed, 41 insertions(+), 10 deletions(-)

diff --git a/plugins/websocket/src/lib.rs b/plugins/websocket/src/lib.rs
index f7922c5b..e5d1007e 100644
--- a/plugins/websocket/src/lib.rs
+++ b/plugins/websocket/src/lib.rs
@@ -8,13 +8,13 @@ use tauri::{
 };
 use tokio::{net::TcpStream, sync::Mutex};
 use tokio_tungstenite::{
-    connect_async_with_config,
+    connect_async_tls_with_config,
     tungstenite::{
         client::IntoClientRequest,
         protocol::{CloseFrame as ProtocolCloseFrame, WebSocketConfig},
         Message,
     },
-    MaybeTlsStream, WebSocketStream,
+    Connector, MaybeTlsStream, WebSocketStream,
 };
 
 use std::collections::HashMap;
@@ -49,6 +49,8 @@ impl Serialize for Error {
 #[derive(Default)]
 struct ConnectionManager(Mutex<HashMap<Id, WebSocketWriter>>);
 
+struct TlsConnector(Mutex<Option<Connector>>);
+
 #[derive(Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ConnectionConfig {
@@ -103,6 +105,10 @@ async fn connect<R: Runtime>(
 ) -> Result<Id> {
     let id = rand::random();
     let mut request = url.into_client_request()?;
+    let tls_connector = match window.try_state::<TlsConnector>() {
+        Some(tls_connector) => tls_connector.0.lock().await.clone(),
+        None => None,
+    };
 
     if let Some(headers) = config.as_ref().and_then(|c| c.headers.as_ref()) {
         for (k, v) in headers {
@@ -112,7 +118,9 @@ async fn connect<R: Runtime>(
         }
     }
 
-    let (ws_stream, _) = connect_async_with_config(request, config.map(Into::into), false).await?;
+    let (ws_stream, _) =
+        connect_async_tls_with_config(request, config.map(Into::into), false, tls_connector)
+            .await?;
 
     tauri::async_runtime::spawn(async move {
         let (write, read) = ws_stream.split();
@@ -186,11 +194,34 @@ async fn send(
 }
 
 pub fn init<R: Runtime>() -> TauriPlugin<R> {
-    PluginBuilder::new("websocket")
-        .invoke_handler(tauri::generate_handler![connect, send])
-        .setup(|app| {
-            app.manage(ConnectionManager::default());
-            Ok(())
-        })
-        .build()
+    Builder::default().build()
+}
+
+#[derive(Default)]
+pub struct Builder {
+    tls_connector: Option<Connector>,
+}
+
+impl Builder {
+    pub fn new() -> Self {
+        Self {
+            tls_connector: None,
+        }
+    }
+
+    pub fn tls_connector(mut self, connector: Connector) -> Self {
+        self.tls_connector.replace(connector);
+        self
+    }
+
+    pub fn build<R: Runtime>(self) -> TauriPlugin<R> {
+        PluginBuilder::new("websocket")
+            .invoke_handler(tauri::generate_handler![connect, send])
+            .setup(|app| {
+                app.manage(ConnectionManager::default());
+                app.manage(TlsConnector(Mutex::new(self.tls_connector)));
+                Ok(())
+            })
+            .build()
+    }
 }