Decouple core Yaak logic from Tauri (#354)

crates/yaak-ws/src/connect.rs

@@ -0,0 +1,50 @@
+use crate::error::Result;
+use http::HeaderMap;
+use log::info;
+use std::sync::Arc;
+use tokio::net::TcpStream;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_tungstenite::tungstenite::handshake::client::Response;
+use tokio_tungstenite::tungstenite::http::HeaderValue;
+use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
+use tokio_tungstenite::{
+    Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
+};
+use yaak_tls::{ClientCertificateConfig, get_tls_config};
+
+// Enabling ALPN breaks websocket requests
+const WITH_ALPN: bool = false;
+
+pub async fn ws_connect(
+    url: &str,
+    headers: HeaderMap<HeaderValue>,
+    validate_certificates: bool,
+    client_cert: Option<ClientCertificateConfig>,
+) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response)> {
+    info!("Connecting to WS {url}");
+    let tls_config = get_tls_config(validate_certificates, WITH_ALPN, client_cert.clone())?;
+
+    let mut req = url.into_client_request()?;
+    let req_headers = req.headers_mut();
+    for (name, value) in headers {
+        if let Some(name) = name {
+            req_headers.insert(name, value);
+        }
+    }
+
+    let (stream, response) = connect_async_tls_with_config(
+        req,
+        Some(WebSocketConfig::default()),
+        false,
+        Some(Connector::Rustls(Arc::new(tls_config))),
+    )
+    .await?;
+
+    info!(
+        "Connected to WS {url} validate_certificates={} client_cert={}",
+        validate_certificates,
+        client_cert.is_some()
+    );
+
+    Ok((stream, response))
+}

crates/yaak-ws/src/error.rs

@@ -0,0 +1,32 @@
+use serde::{Serialize, Serializer};
+use thiserror::Error;
+use tokio_tungstenite::tungstenite;
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("WebSocket error: {0}")]
+    WebSocketErr(#[from] tungstenite::Error),
+
+    #[error(transparent)]
+    ModelError(#[from] yaak_models::error::Error),
+
+    #[error(transparent)]
+    TemplateError(#[from] yaak_templates::error::Error),
+
+    #[error(transparent)]
+    TlsError(#[from] yaak_tls::error::Error),
+
+    #[error("WebSocket error: {0}")]
+    GenericError(String),
+}
+
+impl Serialize for Error {
+    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        serializer.serialize_str(self.to_string().as_ref())
+    }
+}
+
+pub type Result<T> = std::result::Result<T, Error>;

crates/yaak-ws/src/lib.rs

@@ -0,0 +1,12 @@
+mod connect;
+pub mod error;
+pub mod manager;
+pub mod render;
+
+pub use connect::ws_connect;
+pub use manager::WebsocketManager;
+pub use render::render_websocket_request;
+
+// Re-export http types needed by consumers
+pub use http::HeaderMap;
+pub use tokio_tungstenite::tungstenite::http::HeaderValue;

crates/yaak-ws/src/manager.rs

@@ -0,0 +1,101 @@
+use crate::connect::ws_connect;
+use crate::error::Result;
+use futures_util::stream::SplitSink;
+use futures_util::{SinkExt, StreamExt};
+use log::{debug, info, warn};
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::net::TcpStream;
+use tokio::sync::{Mutex, mpsc};
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::handshake::client::Response;
+use http::HeaderMap;
+use tokio_tungstenite::tungstenite::http::HeaderValue;
+use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
+use yaak_tls::ClientCertificateConfig;
+
+#[derive(Clone)]
+pub struct WebsocketManager {
+    connections:
+        Arc<Mutex<HashMap<String, SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>>,
+    read_tasks: Arc<Mutex<HashMap<String, tokio::task::JoinHandle<()>>>>,
+}
+
+impl WebsocketManager {
+    pub fn new() -> Self {
+        WebsocketManager { connections: Default::default(), read_tasks: Default::default() }
+    }
+
+    pub async fn connect(
+        &mut self,
+        id: &str,
+        url: &str,
+        headers: HeaderMap<HeaderValue>,
+        receive_tx: mpsc::Sender<Message>,
+        validate_certificates: bool,
+        client_cert: Option<ClientCertificateConfig>,
+    ) -> Result<Response> {
+        let tx = receive_tx.clone();
+
+        let (stream, response) =
+            ws_connect(url, headers, validate_certificates, client_cert).await?;
+        let (write, mut read) = stream.split();
+
+        self.connections.lock().await.insert(id.to_string(), write);
+
+        let handle = {
+            let connection_id = id.to_string();
+            let connections = self.connections.clone();
+            let read_tasks = self.read_tasks.clone();
+            tokio::task::spawn(async move {
+                while let Some(msg) = read.next().await {
+                    match msg {
+                        Err(e) => {
+                            warn!("Broken websocket connection: {}", e);
+                            break;
+                        }
+                        Ok(message) => tx.send(message).await.unwrap(),
+                    }
+                }
+                debug!("Connection {} closed", connection_id);
+                connections.lock().await.remove(&connection_id);
+                read_tasks.lock().await.remove(&connection_id);
+            })
+        };
+
+        self.read_tasks.lock().await.insert(id.to_string(), handle);
+
+        Ok(response)
+    }
+
+    pub async fn send(&mut self, id: &str, msg: Message) -> Result<()> {
+        debug!("Send websocket message {msg:?}");
+        let mut connections = self.connections.lock().await;
+        let connection = match connections.get_mut(id) {
+            None => return Ok(()),
+            Some(c) => c,
+        };
+        connection.send(msg).await?;
+        Ok(())
+    }
+
+    pub async fn close(&mut self, id: &str) -> Result<()> {
+        info!("Closing websocket");
+        if let Some(mut connection) = self.connections.lock().await.remove(id) {
+            // Wait a maximum of 1 second for the connection to close
+            if let Err(e) = connection.close().await {
+                warn!("Failed to close websocket connection {e:?}");
+            };
+        }
+
+        // Wait at short time for the server to close the connection, then stop
+        // reading.
+        tokio::time::sleep(Duration::from_millis(500)).await;
+        if let Some(handle) = self.read_tasks.lock().await.remove(id) {
+            handle.abort();
+        }
+
+        Ok(())
+    }
+}

crates/yaak-ws/src/render.rs

@@ -0,0 +1,74 @@
+use crate::error::Result;
+use log::info;
+use serde_json::Value;
+use std::collections::BTreeMap;
+use yaak_models::models::{Environment, HttpRequestHeader, HttpUrlParameter, WebsocketRequest};
+use yaak_models::render::make_vars_hashmap;
+use yaak_templates::{RenderOptions, TemplateCallback, parse_and_render, render_json_value_raw};
+
+pub async fn render_websocket_request<T: TemplateCallback>(
+    r: &WebsocketRequest,
+    environment_chain: Vec<Environment>,
+    cb: &T,
+    opt: &RenderOptions,
+) -> Result<WebsocketRequest> {
+    let vars = &make_vars_hashmap(environment_chain);
+
+    let mut url_parameters = Vec::new();
+    for p in r.url_parameters.clone() {
+        url_parameters.push(HttpUrlParameter {
+            enabled: p.enabled,
+            name: parse_and_render(&p.name, vars, cb, opt).await?,
+            value: parse_and_render(&p.value, vars, cb, opt).await?,
+            id: p.id,
+        })
+    }
+
+    let mut headers = Vec::new();
+    for p in r.headers.clone() {
+        headers.push(HttpRequestHeader {
+            enabled: p.enabled,
+            name: parse_and_render(&p.name, vars, cb, opt).await?,
+            value: parse_and_render(&p.value, vars, cb, opt).await?,
+            id: p.id,
+        })
+    }
+
+    let authentication = {
+        let mut disabled = false;
+        let mut auth = BTreeMap::new();
+        match r.authentication.get("disabled") {
+            Some(Value::Bool(true)) => {
+                disabled = true;
+            }
+            Some(Value::String(tmpl)) => {
+                disabled = parse_and_render(tmpl.as_str(), vars, cb, &opt)
+                    .await
+                    .unwrap_or_default()
+                    .is_empty();
+                info!(
+                    "Rendering authentication.disabled as a template: {disabled} from \"{tmpl}\""
+                );
+            }
+            _ => {}
+        }
+        if disabled {
+            auth.insert("disabled".to_string(), Value::Bool(true));
+        } else {
+            for (k, v) in r.authentication.clone() {
+                if k == "disabled" {
+                    auth.insert(k, Value::Bool(false));
+                } else {
+                    auth.insert(k, render_json_value_raw(v, vars, cb, &opt).await?);
+                }
+            }
+        }
+        auth
+    };
+
+    let url = parse_and_render(r.url.as_str(), vars, cb, opt).await?;
+
+    let message = parse_and_render(&r.message.clone(), vars, cb, opt).await?;
+
+    Ok(WebsocketRequest { url, url_parameters, headers, authentication, message, ..r.to_owned() })
+}