From f1acb3c92538ca3967d7ac0d3ef4776cfe374107 Mon Sep 17 00:00:00 2001
From: Carter Costic <cartercostic@gmail.com>
Date: Wed, 23 Jul 2025 16:14:15 -0400
Subject: [PATCH] Merge pull request #245

* Attach cookies to WS Upgrade

* Merge branch 'main' into main

* Move reqwest_cookie_store to workspace dep
---
 src-tauri/Cargo.lock              |  1 +
 src-tauri/Cargo.toml              |  3 +-
 src-tauri/yaak-ws/Cargo.toml      |  1 +
 src-tauri/yaak-ws/src/commands.rs | 66 ++++++++++++++++++++++++++++---
 4 files changed, 64 insertions(+), 7 deletions(-)

diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock
index 78d91c32..d57aa295 100644
--- a/src-tauri/Cargo.lock
+++ b/src-tauri/Cargo.lock
@@ -8103,6 +8103,7 @@ dependencies = [
  "futures-util",
  "log",
  "md5 0.7.0",
+ "reqwest_cookie_store",
  "serde",
  "serde_json",
  "tauri",
diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml
index c9cc0579..bd2d199a 100644
--- a/src-tauri/Cargo.toml
+++ b/src-tauri/Cargo.toml
@@ -50,7 +50,7 @@ md5 = "0.8.0"
 mime_guess = "2.0.5"
 rand = "0.9.0"
 reqwest = { workspace = true, features = ["multipart", "cookies", "gzip", "brotli", "deflate", "json", "rustls-tls-manual-roots-no-provider", "socks"] }
-reqwest_cookie_store = "0.8.0"
+reqwest_cookie_store = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true, features = ["raw_value"] }
 tauri = { workspace = true, features = ["devtools", "protocol-asset"] }
@@ -97,6 +97,7 @@ tauri-plugin-shell = "2.3.0"
 tokio = "1.45.1"
 thiserror = "2.0.12"
 ts-rs = "11.0.1"
+reqwest_cookie_store = "0.8.0"
 rustls = { version = "0.23.27", default-features = false }
 rustls-platform-verifier = "0.6.0"
 sha2 = "0.10.9"
diff --git a/src-tauri/yaak-ws/Cargo.toml b/src-tauri/yaak-ws/Cargo.toml
index cbc06abb..ddd8d269 100644
--- a/src-tauri/yaak-ws/Cargo.toml
+++ b/src-tauri/yaak-ws/Cargo.toml
@@ -9,6 +9,7 @@ publish = false
 futures-util = "0.3.31"
 log = "0.4.20"
 md5 = "0.7.0"
+reqwest_cookie_store = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
 tauri = { workspace = true }
diff --git a/src-tauri/yaak-ws/src/commands.rs b/src-tauri/yaak-ws/src/commands.rs
index c1a3a0a2..294549be 100644
--- a/src-tauri/yaak-ws/src/commands.rs
+++ b/src-tauri/yaak-ws/src/commands.rs
@@ -2,6 +2,7 @@ use crate::error::Result;
 use crate::manager::WebsocketManager;
 use crate::render::render_websocket_request;
 use crate::resolve::resolve_websocket_request;
+use log::debug;
 use log::{info, warn};
 use std::str::FromStr;
 use tauri::http::{HeaderMap, HeaderName};
@@ -300,11 +301,44 @@ pub(crate) async fn connect<R: Runtime>(
         }
     }
 
-    // TODO: Handle cookies
-    let _cookie_jar = match cookie_jar_id {
-        Some(id) => Some(app_handle.db().get_cookie_jar(id)?),
-        None => None,
-    };
+    // Add cookies to WS HTTP Upgrade
+    if let Some(id) = cookie_jar_id {
+        let cookie_jar = app_handle.db().get_cookie_jar(id)?;
+
+        let cookies = cookie_jar
+            .cookies
+            .iter()
+            .filter_map(|cookie| {
+                // HACK: same as in src-tauri/src/http_request.rs
+                let json_cookie = serde_json::to_value(cookie).ok()?;
+                match serde_json::from_value(json_cookie) {
+                    Ok(cookie) => Some(Ok(cookie)),
+                    Err(_e) => None,
+                }
+            })
+            .collect::<Vec<Result<_>>>();
+
+        let store = reqwest_cookie_store::CookieStore::from_cookies(cookies, true)?;
+
+        // Convert WS URL -> HTTP URL bc reqwest_cookie_store's `get_request_values`
+        // strictly matches based on Path/HttpOnly/Secure attributes even though WS upgrades are HTTP requests
+        let http_url = convert_ws_url_to_http(&url);
+        let pairs: Vec<_> = store.get_request_values(&http_url).collect();
+        debug!("Inserting {} cookies into WS upgrade to {}", pairs.len(), url);
+
+        let cookie_header_value = pairs
+            .into_iter()
+            .map(|(name, value)| format!("{}={}", name, value))
+            .collect::<Vec<_>>()
+            .join("; ");
+
+        if !cookie_header_value.is_empty() {
+            headers.insert(
+                HeaderName::from_static("cookie"),
+                HeaderValue::from_str(&cookie_header_value).unwrap(),
+            );
+        }
+    }
 
     let (receive_tx, mut receive_rx) = mpsc::channel::<Message>(128);
     let mut ws_manager = ws_manager.lock().await;
@@ -337,7 +371,7 @@ pub(crate) async fn connect<R: Runtime>(
         Err(e) => {
             return Ok(app_handle.db().upsert_websocket_connection(
                 &WebsocketConnection {
-                    error: Some(format!("{e:?}")),
+                    error: Some(e.to_string()),
                     state: WebsocketConnectionState::Closed,
                     ..connection
                 },
@@ -448,3 +482,23 @@ pub(crate) async fn connect<R: Runtime>(
 
     Ok(connection)
 }
+
+/// Convert WS URL to HTTP URL for cookie filtering
+/// WebSocket upgrade requests are HTTP requests initially, so HttpOnly cookies should apply
+fn convert_ws_url_to_http(ws_url: &Url) -> Url {
+    let mut http_url = ws_url.clone();
+
+    match ws_url.scheme() {
+        "ws" => {
+            http_url.set_scheme("http").expect("Failed to set http scheme");
+        }
+        "wss" => {
+            http_url.set_scheme("https").expect("Failed to set https scheme");
+        }
+        _ => {
+            // Already HTTP/HTTPS, no conversion needed
+        }
+    }
+
+    http_url
+}