Add domain filter to cookie template function (#452)

Author: gschier

crates-cli/yaak-cli/src/plugin_events.rs

@@ -14,6 +14,7 @@ use yaak::plugin_events::{
 use yaak::render::{render_grpc_request, render_http_request};
 use yaak::send::{SendHttpRequestWithPluginsParams, send_http_request_with_plugins};
 use yaak_crypto::manager::EncryptionManager;
+use yaak_http::cookies::get_cookie_value_from_jar;
 use yaak_models::blob_manager::BlobManager;
 use yaak_models::models::Environment;
 use yaak_models::queries::any_request::AnyRequest;
@@ -496,10 +497,8 @@ async fn build_plugin_reply(
                         }
                     };
 
-                let value = cookie_jar.cookies.into_iter().find_map(|c| {
-                    let (name, value) = parse_cookie_name_value(&c.raw_cookie)?;
-                    if name == req.name { Some(value) } else { None }
-                });
+                let value =
+                    get_cookie_value_from_jar(cookie_jar.cookies, &req.name, req.domain.as_deref());
                 Some(InternalEventPayload::GetCookieValueResponse(GetCookieValueResponse { value }))
             }
             HostRequest::WindowInfo(req) => {
@@ -532,7 +531,6 @@ async fn render_json_value_for_cli<T: TemplateCallback>(
     render_json_value_raw(value, vars, cb, opt).await
 }
 
-
 fn parse_cookie_name_value(raw_cookie: &str) -> Option<(String, String)> {
     let first_part = raw_cookie.split(';').next()?.trim();
     let (name, value) = first_part.split_once('=')?;

crates-tauri/yaak-app/src/plugin_events.rs

@@ -19,6 +19,7 @@ use yaak::plugin_events::{
     GroupedPluginEvent, HostRequest, SharedPluginEventContext, handle_shared_plugin_event,
 };
 use yaak_crypto::manager::EncryptionManager;
+use yaak_http::cookies::get_cookie_value_from_jar;
 use yaak_models::models::{HttpResponse, Plugin};
 use yaak_models::queries::any_request::AnyRequest;
 use yaak_models::util::UpdateSource;
@@ -420,12 +421,7 @@ async fn handle_host_plugin_request<R: Runtime>(
             let window = get_window_from_plugin_context(app_handle, plugin_context)?;
             let value = match cookie_jar_from_window(&window) {
                 None => None,
-                Some(j) => j.cookies.into_iter().find_map(|c| match Cookie::parse(c.raw_cookie) {
-                    Ok(c) if c.name().to_string().eq(&req.name) => {
-                        Some(c.value_trimmed().to_string())
-                    }
-                    _ => None,
-                }),
+                Some(j) => get_cookie_value_from_jar(j.cookies, &req.name, req.domain.as_deref()),
             };
             Ok(Some(InternalEventPayload::GetCookieValueResponse(GetCookieValueResponse { value })))
         }

crates/yaak-http/src/cookies.rs

@@ -124,6 +124,30 @@ impl CookieStore {
     }
 }
 
+/// Get a stored cookie value by name, optionally scoped to an exact stored domain.
+pub fn get_cookie_value_from_jar(
+    cookies: impl IntoIterator<Item = Cookie>,
+    name: &str,
+    domain: Option<&str>,
+) -> Option<String> {
+    let domain = domain.and_then(normalize_cookie_domain_filter);
+
+    cookies.into_iter().find_map(|cookie| {
+        let (cookie_name, value) = parse_cookie_name_value(&cookie.raw_cookie)?;
+        if cookie_name != name {
+            return None;
+        }
+
+        if let Some(domain) = domain.as_deref() {
+            if !cookie_domain_matches_filter(&cookie.domain, domain) {
+                return None;
+            }
+        }
+
+        Some(value)
+    })
+}
+
 /// Parse name=value from a cookie string (raw_cookie format)
 fn parse_cookie_name_value(raw_cookie: &str) -> Option<(String, String)> {
     // The raw_cookie typically looks like "name=value" or "name=value; attr1; attr2=..."
@@ -135,6 +159,20 @@ fn parse_cookie_name_value(raw_cookie: &str) -> Option<(String, String)> {
     if name.is_empty() { None } else { Some((name, value)) }
 }
 
+fn normalize_cookie_domain_filter(domain: &str) -> Option<String> {
+    let domain = domain.trim().trim_start_matches('.').to_lowercase();
+    if domain.is_empty() { None } else { Some(domain) }
+}
+
+fn cookie_domain_matches_filter(cookie_domain: &CookieDomain, domain: &str) -> bool {
+    match cookie_domain {
+        CookieDomain::HostOnly(cookie_domain) | CookieDomain::Suffix(cookie_domain) => {
+            normalize_cookie_domain_filter(cookie_domain).is_some_and(|d| d == domain)
+        }
+        CookieDomain::NotPresent | CookieDomain::Empty => false,
+    }
+}
+
 /// Parse a Set-Cookie header into a Cookie
 fn parse_set_cookie(header_value: &str, request_url: &Url) -> Option<Cookie> {
     let parsed = cookie::Cookie::parse(header_value).ok()?;
@@ -278,6 +316,15 @@ fn is_localhost(domain: &str) -> bool {
 mod tests {
     use super::*;
 
+    fn cookie(raw_cookie: &str, domain: CookieDomain) -> Cookie {
+        Cookie {
+            raw_cookie: raw_cookie.to_string(),
+            domain,
+            expires: CookieExpires::SessionEnd,
+            path: ("/".to_string(), false),
+        }
+    }
+
     #[test]
     fn test_parse_cookie_name_value() {
         assert_eq!(
@@ -387,6 +434,52 @@ mod tests {
         assert_eq!(store.get_all_cookies().len(), 1);
     }
 
+    #[test]
+    fn test_get_cookie_value_preserves_name_only_first_match() {
+        let cookies = vec![
+            cookie("co-auth=", CookieDomain::HostOnly("foo.example.com".to_string())),
+            cookie("co-auth=token", CookieDomain::Suffix("example.com".to_string())),
+        ];
+
+        assert_eq!(get_cookie_value_from_jar(cookies, "co-auth", None), Some("".to_string()));
+    }
+
+    #[test]
+    fn test_get_cookie_value_matches_domain() {
+        let cookies = vec![
+            cookie("co-auth=", CookieDomain::HostOnly("foo.example.com".to_string())),
+            cookie("co-auth=token", CookieDomain::Suffix("example.com".to_string())),
+        ];
+
+        assert_eq!(
+            get_cookie_value_from_jar(cookies, "co-auth", Some("example.com")),
+            Some("token".to_string())
+        );
+    }
+
+    #[test]
+    fn test_get_cookie_value_normalizes_domain_filter() {
+        let cookies = vec![cookie(
+            "co-auth=token",
+            CookieDomain::Suffix("Example.COM".to_string()),
+        )];
+
+        assert_eq!(
+            get_cookie_value_from_jar(cookies, "co-auth", Some(" .example.com ")),
+            Some("token".to_string())
+        );
+    }
+
+    #[test]
+    fn test_get_cookie_value_requires_exact_stored_domain_match() {
+        let cookies = vec![cookie(
+            "co-auth=token",
+            CookieDomain::HostOnly("foo.example.com".to_string()),
+        )];
+
+        assert_eq!(get_cookie_value_from_jar(cookies, "co-auth", Some("example.com")), None);
+    }
+
     #[test]
     fn test_is_single_component_domain() {
         // Single-component domains (TLDs)

crates/yaak-plugins/bindings/gen_events.ts

@@ -307,6 +307,9 @@ pub struct ListCookieNamesResponse {
 #[ts(export, export_to = "gen_events.ts")]
 pub struct GetCookieValueRequest {
     pub name: String,
+
+    #[ts(optional = nullable)]
+    pub domain: Option<String>,
 }
 
 #[derive(Debug, Clone, Default, Serialize, Deserialize, TS)]

crates/yaak-plugins/src/events.rs

@@ -11,12 +11,27 @@ export const plugin: PluginDefinition = {
           type: "text",
           name: "name",
           label: "Cookie Name",
+          placeholder: "cookie_name",
+        },
+        {
+          type: "text",
+          name: "domain",
+          label: "Domain",
+          placeholder: "example.com",
+          description:
+            "Optionally filter by domain, useful if multiple cookies with the same name.",
+          optional: true,
         },
       ],
       async onRender(ctx: Context, args: CallTemplateFunctionArgs): Promise<string | null> {
         // The legacy name was cookie_name, but we changed it
         const name = args.values.cookie_name ?? args.values.name;
-        return ctx.cookies.getValue({ name: String(name) });
+        const domain = String(args.values.domain ?? "").trim();
+
+        return ctx.cookies.getValue({
+          name: String(name),
+          ...(domain.length > 0 ? { domain } : {}),
+        });
       },
     },
   ],