Cookie Support (#19)

src-tauri/src/http.rs

@@ -2,16 +2,18 @@ use std::fs;
 use std::fs::{create_dir_all, File};
 use std::io::Write;
 use std::path::PathBuf;
+use std::str::FromStr;
+use std::sync::Arc;
 use std::time::Duration;
 
 use base64::Engine;
-use http::{HeaderMap, HeaderName, HeaderValue, Method};
 use http::header::{ACCEPT, USER_AGENT};
+use http::{HeaderMap, HeaderName, HeaderValue, Method};
 use log::{error, info, warn};
-use reqwest::multipart;
 use reqwest::redirect::Policy;
+use reqwest::{multipart, Url};
+use sqlx::types::{Json, JsonValue};
 use sqlx::{Pool, Sqlite};
-use sqlx::types::Json;
 use tauri::{AppHandle, Wry};
 
 use crate::{emit_side_effect, models, render, response_err};
@@ -19,13 +21,13 @@ use crate::{emit_side_effect, models, render, response_err};
 pub async fn send_http_request(
     request: models::HttpRequest,
     response: &models::HttpResponse,
-    environment_id: &str,
+    environment: Option<models::Environment>,
+    cookie_jar: Option<models::CookieJar>,
     app_handle: &AppHandle<Wry>,
     pool: &Pool<Sqlite>,
     download_path: Option<PathBuf>,
 ) -> Result<models::HttpResponse, String> {
     let start = std::time::Instant::now();
-    let environment = models::get_environment(environment_id, pool).await.ok();
     let environment_ref = environment.as_ref();
     let workspace = models::get_workspace(&request.workspace_id, pool)
         .await
@@ -49,6 +51,32 @@ pub async fn send_http_request(
         .danger_accept_invalid_certs(!workspace.setting_validate_certificates)
         .tls_info(true);
 
+    // Add cookie store if specified
+    let maybe_cookie_manager = match cookie_jar.clone() {
+        Some(cj) => {
+            // HACK: Can't construct Cookie without serde, so we have to do this
+            let cookies = cj
+                .cookies
+                .0
+                .iter()
+                .map(|json_cookie| {
+                    serde_json::from_value(json_cookie.clone())
+                        .expect("Failed to deserialize cookie")
+                })
+                .map(|c| Ok(c))
+                .collect::<Vec<Result<_, ()>>>();
+
+            let store = reqwest_cookie_store::CookieStore::from_cookies(cookies, true)
+                .expect("Failed to create cookie store");
+            let cookie_store = reqwest_cookie_store::CookieStoreMutex::new(store);
+            let cookie_store = Arc::new(cookie_store);
+            client_builder = client_builder.cookie_provider(Arc::clone(&cookie_store));
+
+            Some((cookie_store, cj))
+        }
+        None => None,
+    };
+
     if workspace.setting_request_timeout > 0 {
         client_builder = client_builder.timeout(Duration::from_millis(
             workspace.setting_request_timeout.unsigned_abs(),
@@ -58,14 +86,37 @@ pub async fn send_http_request(
     // .use_rustls_tls() // TODO: Make this configurable (maybe)
     let client = client_builder.build().expect("Failed to build client");
 
+    let url = match Url::from_str(url_string.as_str()) {
+        Ok(u) => u,
+        Err(e) => {
+            return response_err(response, e.to_string(), app_handle, pool).await;
+        }
+    };
+
     let m = Method::from_bytes(request.method.to_uppercase().as_bytes())
         .expect("Failed to create method");
-    let mut request_builder = client.request(m, url_string.to_string());
+    let mut request_builder = client.request(m, url.clone());
 
     let mut headers = HeaderMap::new();
     headers.insert(USER_AGENT, HeaderValue::from_static("yaak"));
     headers.insert(ACCEPT, HeaderValue::from_static("*/*"));
 
+    // TODO: Set cookie header ourselves once we also handle redirects. We need to do this
+    //  because reqwest doesn't give us a way to inspect the headers it sent (we have to do
+    //  everything manually to know that).
+    // if let Some(cookie_store) = maybe_cookie_store.clone() {
+    //     let values1 = cookie_store.get_request_values(&url);
+    //     println!("COOKIE VLUAES: {:?}", values1.collect::<Vec<_>>());
+    //     let raw_value = cookie_store.get_request_values(&url)
+    //         .map(|(name, value)| format!("{}={}", name, value))
+    //         .collect::<Vec<_>>()
+    //         .join("; ");
+    //     headers.insert(
+    //         COOKIE,
+    //         HeaderValue::from_str(&raw_value).expect("Failed to create cookie header"),
+    //     );
+    // }
+
     for h in request.headers.0 {
         if h.name.is_empty() && h.value.is_empty() {
             continue;
@@ -252,10 +303,11 @@ pub async fn send_http_request(
     match raw_response {
         Ok(v) => {
             let mut response = response.clone();
+            let response_headers = v.headers().clone();
             response.status = v.status().as_u16() as i64;
             response.status_reason = v.status().canonical_reason().map(|s| s.to_string());
             response.headers = Json(
-                v.headers()
+                response_headers
                     .iter()
                     .map(|(k, v)| models::HttpResponseHeader {
                         name: k.as_str().to_string(),
@@ -304,15 +356,42 @@ pub async fn send_http_request(
             match (download_path, response.body_path.clone()) {
                 (Some(dl_path), Some(body_path)) => {
                     info!("Downloading response body to {}", dl_path.display());
-                    fs::copy(body_path, dl_path).expect("Failed to copy file for response download");
+                    fs::copy(body_path, dl_path)
+                        .expect("Failed to copy file for response download");
                 }
                 _ => {}
             };
 
+            // Add cookie store if specified
+            if let Some((cookie_store, mut cookie_jar)) = maybe_cookie_manager {
+                // let cookies = response_headers.get_all(SET_COOKIE).iter().map(|h| {
+                //     println!("RESPONSE COOKIE: {}", h.to_str().unwrap());
+                //     cookie_store::RawCookie::from_str(h.to_str().unwrap())
+                //         .expect("Failed to parse cookie")
+                // });
+                // store.store_response_cookies(cookies, &url);
+
+                let json_cookies: Json<Vec<JsonValue>> = Json(
+                    cookie_store
+                        .lock()
+                        .unwrap()
+                        .iter_any()
+                        .map(|c| serde_json::to_value(&c).expect("Failed to serialize cookie"))
+                        .collect::<Vec<_>>(),
+                );
+                cookie_jar.cookies = json_cookies;
+                match models::upsert_cookie_jar(pool, &cookie_jar).await {
+                    Ok(updated_jar) => {
+                        emit_side_effect(app_handle, "updated_model", &updated_jar);
+                    }
+                    Err(e) => {
+                        error!("Failed to update cookie jar: {}", e);
+                    }
+                };
+            }
+
             Ok(response)
         }
-        Err(e) => {
-            response_err(response, e.to_string(), app_handle, pool).await
-        }
+        Err(e) => response_err(response, e.to_string(), app_handle, pool).await,
     }
 }