use crate::error::Error::GenericError; use crate::error::Result; use crate::manager::WebsocketManager; use crate::render::render_request; use chrono::Utc; use log::{info, warn}; use std::str::FromStr; use tauri::http::{HeaderMap, HeaderName}; use tauri::{AppHandle, Manager, Runtime, State, Url, WebviewWindow}; use tokio::sync::{mpsc, Mutex}; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::Message; use yaak_http::apply_path_placeholders; use yaak_models::models::{ HttpResponseHeader, WebsocketConnection, WebsocketConnectionState, WebsocketEvent, WebsocketEventType, WebsocketRequest, }; use yaak_models::queries; use yaak_models::queries::{ get_base_environment, get_cookie_jar, get_environment, get_websocket_connection, get_websocket_request, upsert_websocket_connection, upsert_websocket_event, UpdateSource, }; use yaak_plugins::events::{ CallHttpAuthenticationRequest, HttpHeader, RenderPurpose, WindowContext, }; use yaak_plugins::manager::PluginManager; use yaak_plugins::template_callback::PluginTemplateCallback; #[tauri::command] pub(crate) async fn upsert_request( request: WebsocketRequest, w: WebviewWindow, ) -> Result { Ok(queries::upsert_websocket_request(&w, request, &UpdateSource::Window).await?) } #[tauri::command] pub(crate) async fn duplicate_request( request_id: &str, w: WebviewWindow, ) -> Result { Ok(queries::duplicate_websocket_request(&w, request_id, &UpdateSource::Window).await?) } #[tauri::command] pub(crate) async fn delete_request( request_id: &str, w: WebviewWindow, ) -> Result { Ok(queries::delete_websocket_request(&w, request_id, &UpdateSource::Window).await?) } #[tauri::command] pub(crate) async fn delete_connection( connection_id: &str, w: WebviewWindow, ) -> Result { Ok(queries::delete_websocket_connection(&w, connection_id, &UpdateSource::Window).await?) } #[tauri::command] pub(crate) async fn delete_connections( request_id: &str, w: WebviewWindow, ) -> Result<()> { Ok(queries::delete_all_websocket_connections(&w, request_id, &UpdateSource::Window).await?) } #[tauri::command] pub(crate) async fn list_events( connection_id: &str, app_handle: AppHandle, ) -> Result> { Ok(queries::list_websocket_events(&app_handle, connection_id).await?) } #[tauri::command] pub(crate) async fn list_requests( workspace_id: &str, app_handle: AppHandle, ) -> Result> { Ok(queries::list_websocket_requests(&app_handle, workspace_id).await?) } #[tauri::command] pub(crate) async fn list_connections( workspace_id: &str, app_handle: AppHandle, ) -> Result> { Ok(queries::list_websocket_connections_for_workspace(&app_handle, workspace_id).await?) } #[tauri::command] pub(crate) async fn send( connection_id: &str, environment_id: Option<&str>, window: WebviewWindow, ws_manager: State<'_, Mutex>, ) -> Result { let connection = get_websocket_connection(&window, connection_id).await?; let unrendered_request = get_websocket_request(&window, &connection.request_id) .await? .ok_or(GenericError("WebSocket Request not found".to_string()))?; let environment = match environment_id { Some(id) => Some(get_environment(&window, id).await?), None => None, }; let base_environment = get_base_environment(&window, &unrendered_request.workspace_id).await?; let request = render_request( &unrendered_request, &base_environment, environment.as_ref(), &PluginTemplateCallback::new( window.app_handle(), &WindowContext::from_window(&window), RenderPurpose::Send, ), ) .await; let mut ws_manager = ws_manager.lock().await; ws_manager.send(&connection.id, Message::Text(request.message.clone().into())).await?; upsert_websocket_event( &window, WebsocketEvent { connection_id: connection.id.clone(), request_id: request.id.clone(), workspace_id: connection.workspace_id.clone(), is_server: false, message_type: WebsocketEventType::Text, message: request.message.into(), ..Default::default() }, &UpdateSource::Window, ) .await .unwrap(); Ok(connection) } #[tauri::command] pub(crate) async fn close( connection_id: &str, window: WebviewWindow, ws_manager: State<'_, Mutex>, ) -> Result { let connection = get_websocket_connection(&window, connection_id).await?; let request = get_websocket_request(&window, &connection.request_id) .await? .ok_or(GenericError("WebSocket Request not found".to_string()))?; let mut ws_manager = ws_manager.lock().await; if let Err(e) = ws_manager.send(&connection.id, Message::Close(None)).await { warn!("Failed to close WebSocket connection: {e:?}"); }; upsert_websocket_event( &window, WebsocketEvent { connection_id: connection.id.clone(), request_id: request.id.clone(), workspace_id: request.workspace_id.clone(), is_server: false, message_type: WebsocketEventType::Close, ..Default::default() }, &UpdateSource::Window, ) .await .unwrap(); let connection = upsert_websocket_connection( &window, &WebsocketConnection { state: WebsocketConnectionState::Closed, elapsed: Utc::now() .naive_utc() .signed_duration_since(connection.created_at) .num_milliseconds() as i32, ..connection.clone() }, &UpdateSource::Window, ) .await?; Ok(connection) } #[tauri::command] pub(crate) async fn connect( request_id: &str, environment_id: Option<&str>, cookie_jar_id: Option<&str>, window: WebviewWindow, plugin_manager: State<'_, PluginManager>, ws_manager: State<'_, Mutex>, ) -> Result { let unrendered_request = get_websocket_request(&window, request_id) .await? .ok_or(GenericError("Failed to find GRPC request".to_string()))?; let environment = match environment_id { Some(id) => Some(get_environment(&window, id).await?), None => None, }; let base_environment = get_base_environment(&window, &unrendered_request.workspace_id).await?; let request = render_request( &unrendered_request, &base_environment, environment.as_ref(), &PluginTemplateCallback::new( window.app_handle(), &WindowContext::from_window(&window), RenderPurpose::Send, ), ) .await; let mut headers = HeaderMap::new(); if let Some(auth_name) = request.authentication_type.clone() { let auth = request.authentication.clone(); let plugin_req = CallHttpAuthenticationRequest { context_id: format!("{:x}", md5::compute(request_id.to_string())), values: serde_json::from_value(serde_json::to_value(&auth).unwrap()).unwrap(), method: "POST".to_string(), url: request.url.clone(), headers: request .headers .clone() .into_iter() .map(|h| HttpHeader { name: h.name, value: h.value, }) .collect(), }; let plugin_result = plugin_manager.call_http_authentication(&window, &auth_name, plugin_req).await?; for header in plugin_result.set_headers { headers.insert( HeaderName::from_str(&header.name).unwrap(), HeaderValue::from_str(&header.value).unwrap(), ); } } // TODO: Handle cookies let _cookie_jar = match cookie_jar_id { Some(id) => Some(get_cookie_jar(&window, id).await?), None => None, }; let connection = upsert_websocket_connection( &window, &WebsocketConnection { workspace_id: request.workspace_id.clone(), request_id: request_id.to_string(), ..Default::default() }, &UpdateSource::Window, ) .await?; let (receive_tx, mut receive_rx) = mpsc::channel::(128); let mut ws_manager = ws_manager.lock().await; { let connection_id = connection.id.clone(); let request_id = request.id.to_string(); let workspace_id = request.workspace_id.clone(); let window = window.clone(); tokio::spawn(async move { while let Some(message) = receive_rx.recv().await { upsert_websocket_event( &window, WebsocketEvent { connection_id: connection_id.clone(), request_id: request_id.clone(), workspace_id: workspace_id.clone(), is_server: true, message_type: match message { Message::Text(_) => WebsocketEventType::Text, Message::Binary(_) => WebsocketEventType::Binary, Message::Ping(_) => WebsocketEventType::Ping, Message::Pong(_) => WebsocketEventType::Pong, Message::Close(_) => WebsocketEventType::Close, // Raw frame will never happen during a read Message::Frame(_) => WebsocketEventType::Frame, }, message: message.into_data().into(), ..Default::default() }, &UpdateSource::Window, ) .await .unwrap(); } info!("Websocket connection closed"); }); } let (url, url_parameters) = apply_path_placeholders(&request.url, request.url_parameters); // Add URL parameters to URL let mut url = Url::parse(&url).unwrap(); { let mut query_pairs = url.query_pairs_mut(); for p in url_parameters.clone() { if !p.enabled || p.name.is_empty() { continue; } query_pairs.append_pair(p.name.as_str(), p.value.as_str()); } } let response = match ws_manager.connect(&connection.id, url.as_str(), headers, receive_tx).await { Ok(r) => r, Err(e) => { return Ok(upsert_websocket_connection( &window, &WebsocketConnection { error: Some(format!("{e:?}")), state: WebsocketConnectionState::Closed, ..connection }, &UpdateSource::Window, ) .await?); } }; let response_headers = response .headers() .into_iter() .map(|(name, value)| HttpResponseHeader { name: name.to_string(), value: value.to_str().unwrap().to_string(), }) .collect::>(); let connection = upsert_websocket_connection( &window, &WebsocketConnection { state: WebsocketConnectionState::Connected, headers: response_headers, status: response.status().as_u16() as i32, url: request.url.clone(), ..connection }, &UpdateSource::Window, ) .await?; Ok(connection) }