Websockets for plugin runtime communication (#156)

src-tauri/yaak-plugins/src/server_ws.rs

@@ -0,0 +1,132 @@
+use crate::events::{ErrorResponse, InternalEvent, InternalEventPayload, InternalEventRawPayload};
+use futures_util::{SinkExt, StreamExt};
+use log::{error, info, warn};
+use std::sync::Arc;
+use tokio::net::{TcpListener, TcpStream};
+use tokio::sync::{mpsc, Mutex};
+use tokio_tungstenite::accept_async_with_config;
+use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
+use tokio_tungstenite::tungstenite::Message;
+
+#[derive(Clone)]
+pub(crate) struct PluginRuntimeServerWebsocket {
+    pub(crate) app_to_plugin_events_tx: Arc<Mutex<Option<mpsc::Sender<InternalEvent>>>>,
+    client_disconnect_tx: mpsc::Sender<bool>,
+    client_connect_tx: tokio::sync::watch::Sender<bool>,
+    plugin_to_app_events_tx: mpsc::Sender<InternalEvent>,
+}
+
+impl PluginRuntimeServerWebsocket {
+    pub fn new(
+        events_tx: mpsc::Sender<InternalEvent>,
+        disconnect_tx: mpsc::Sender<bool>,
+        connect_tx: tokio::sync::watch::Sender<bool>,
+    ) -> Self {
+        PluginRuntimeServerWebsocket {
+            app_to_plugin_events_tx: Arc::new(Mutex::new(None)),
+            client_disconnect_tx: disconnect_tx,
+            client_connect_tx: connect_tx,
+            plugin_to_app_events_tx: events_tx,
+        }
+    }
+
+    pub async fn listen(&self, listener: TcpListener) {
+        while let Ok((stream, _)) = listener.accept().await {
+            self.accept_connection(stream).await;
+        }
+    }
+
+    async fn accept_connection(&self, stream: TcpStream) {
+        let (to_plugin_tx, mut to_plugin_rx) = mpsc::channel::<InternalEvent>(128);
+        let mut app_to_plugin_events_tx = self.app_to_plugin_events_tx.lock().await;
+        *app_to_plugin_events_tx = Some(to_plugin_tx);
+
+        let plugin_to_app_events_tx = self.plugin_to_app_events_tx.clone();
+        let client_disconnect_tx = self.client_disconnect_tx.clone();
+        let client_connect_tx = self.client_connect_tx.clone();
+
+        let addr = stream.peer_addr().expect("connected streams should have a peer address");
+
+        let conf = WebSocketConfig::default();
+        let ws_stream = accept_async_with_config(stream, Some(conf))
+            .await
+            .expect("Error during the websocket handshake occurred");
+
+        let (mut ws_sender, mut ws_receiver) = ws_stream.split();
+
+        tauri::async_runtime::spawn(async move {
+            client_connect_tx.send(true).expect("Failed to send client ready event");
+
+            info!("New plugin runtime websocket connection: {}", addr);
+
+            loop {
+                tokio::select! {
+                    msg = ws_receiver.next() => {
+                        let msg = match msg {
+                            Some(Ok(msg)) => msg,
+                            Some(Err(e)) => {
+                                warn!("Websocket error {e:?}");
+                                continue;
+                            }
+                            None => break,
+                        };
+
+                        // Skip non-text messages
+                        if !msg.is_text() {
+                            return;
+                        }
+
+                        let event = match serde_json::from_str::<InternalEventRawPayload>(&msg.into_text().unwrap()) {
+                            Ok(e) => e,
+                            Err(e) => {
+                                error!("Failed to decode plugin event {e:?}");
+                                continue;
+                            }
+                        };
+
+                        // Parse everything but the payload so we can catch errors on that, specifically
+                        let payload = serde_json::from_value::<InternalEventPayload>(event.payload)
+                            .unwrap_or_else(|e| {
+                                InternalEventPayload::ErrorResponse(ErrorResponse {
+                                    error: format!("Plugin error from {}: {e:?}", event.plugin_name),
+                                })
+                            });
+
+                        let event = InternalEvent{
+                            id: event.id,
+                            payload,
+                            plugin_ref_id: event.plugin_ref_id,
+                            plugin_name: event.plugin_name,
+                            window_context: event.window_context,
+                            reply_id: event.reply_id,
+                        };
+
+                        // Send event to subscribers
+                        // Emit event to the channel for server to handle
+                        if let Err(e) = plugin_to_app_events_tx.try_send(event) {
+                            warn!("Failed to send to channel. Receiver probably isn't listening: {:?}", e);
+                        }
+                    }
+
+                    event_for_plugin = to_plugin_rx.recv() => {
+                        match event_for_plugin {
+                            None => {
+                                error!("Plugin runtime client WS channel closed");
+                                return;
+                            },
+                            Some(event) => {
+                                let event_bytes = serde_json::to_string(&event).unwrap();
+                                let msg = Message::text(event_bytes);
+                                ws_sender.send(msg).await.unwrap();
+                            }
+                        }
+                    }
+                }
+            }
+
+            if let Err(e) = client_disconnect_tx.send(true).await {
+                warn!("Failed to send killed event {:?}", e);
+            }
+        });
+    }
+}