Skip to content

Commit

Permalink
feat(websocket): Add TLS configuration (#925)
Browse files Browse the repository at this point in the history
* Allow TLS configuration

* refactor: Added builder pattern

---------

Co-authored-by: FabianLars <[email protected]>
  • Loading branch information
jnugh and FabianLars authored Apr 15, 2024
1 parent 4210cf3 commit 5c97db9
Showing 1 changed file with 41 additions and 10 deletions.
51 changes: 41 additions & 10 deletions plugins/websocket/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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();
Expand Down Expand Up @@ -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()
}
}

0 comments on commit 5c97db9

Please sign in to comment.