viewpoint_core/context/cookies/
mod.rs

1//! Cookie management for BrowserContext.
2//!
3//! This module provides methods for managing cookies in an isolated browser context.
4
5use tracing::{debug, instrument};
6
7use viewpoint_cdp::protocol::{CookieParam, CookieSameSite};
8use viewpoint_cdp::protocol::storage::{
9    ClearCookiesParams as StorageClearCookiesParams,
10    DeleteCookiesParams as StorageDeleteCookiesParams, GetCookiesParams as StorageGetCookiesParams,
11    GetCookiesResult as StorageGetCookiesResult, SetCookiesParams as StorageSetCookiesParams,
12};
13
14use super::types::{Cookie, SameSite};
15use super::BrowserContext;
16use crate::error::ContextError;
17
18impl BrowserContext {
19    /// Add cookies to the browser context.
20    ///
21    /// # Example
22    ///
23    /// ```no_run
24    /// use viewpoint_core::{BrowserContext, context::{Cookie, SameSite}};
25    ///
26    /// # async fn example(context: &BrowserContext) -> Result<(), viewpoint_core::CoreError> {
27    /// context.add_cookies(vec![
28    ///     Cookie::new("session", "abc123")
29    ///         .domain("example.com")
30    ///         .path("/")
31    ///         .secure(true)
32    ///         .http_only(true),
33    /// ]).await?;
34    /// # Ok(())
35    /// # }
36    /// ```
37    ///
38    /// # Errors
39    ///
40    /// Returns an error if setting cookies fails.
41    #[instrument(level = "debug", skip(self, cookies))]
42    pub async fn add_cookies(&self, cookies: Vec<Cookie>) -> Result<(), ContextError> {
43        if self.is_closed() {
44            return Err(ContextError::Closed);
45        }
46
47        debug!(count = cookies.len(), "Adding cookies");
48
49        let cdp_cookies: Vec<CookieParam> = cookies
50            .into_iter()
51            .map(|c| {
52                let mut param = CookieParam::new(&c.name, &c.value);
53                if let Some(url) = c.url {
54                    param = param.url(url);
55                }
56                if let Some(domain) = c.domain {
57                    param = param.domain(domain);
58                }
59                if let Some(path) = c.path {
60                    param = param.path(path);
61                }
62                if let Some(secure) = c.secure {
63                    param = param.secure(secure);
64                }
65                if let Some(http_only) = c.http_only {
66                    param = param.http_only(http_only);
67                }
68                if let Some(expires) = c.expires {
69                    param = param.expires(expires);
70                }
71                if let Some(same_site) = c.same_site {
72                    param.same_site = Some(match same_site {
73                        SameSite::Strict => CookieSameSite::Strict,
74                        SameSite::Lax => CookieSameSite::Lax,
75                        SameSite::None => CookieSameSite::None,
76                    });
77                }
78                param
79            })
80            .collect();
81
82        self.connection()
83            .send_command::<_, serde_json::Value>(
84                "Storage.setCookies",
85                Some(
86                    StorageSetCookiesParams::new(cdp_cookies)
87                        .browser_context_id(self.context_id().to_string()),
88                ),
89                None,
90            )
91            .await?;
92
93        Ok(())
94    }
95
96    /// Get all cookies in the browser context.
97    ///
98    /// # Errors
99    ///
100    /// Returns an error if getting cookies fails.
101    #[instrument(level = "debug", skip(self))]
102    pub async fn cookies(&self) -> Result<Vec<Cookie>, ContextError> {
103        if self.is_closed() {
104            return Err(ContextError::Closed);
105        }
106
107        let result: StorageGetCookiesResult = self
108            .connection()
109            .send_command(
110                "Storage.getCookies",
111                Some(StorageGetCookiesParams::new().browser_context_id(self.context_id().to_string())),
112                None,
113            )
114            .await?;
115
116        let cookies = result
117            .cookies
118            .into_iter()
119            .map(|c| Cookie {
120                name: c.name,
121                value: c.value,
122                domain: Some(c.domain),
123                path: Some(c.path),
124                url: None,
125                expires: if c.expires > 0.0 {
126                    Some(c.expires)
127                } else {
128                    None
129                },
130                http_only: Some(c.http_only),
131                secure: Some(c.secure),
132                same_site: c.same_site.map(|s| match s {
133                    CookieSameSite::Strict => SameSite::Strict,
134                    CookieSameSite::Lax => SameSite::Lax,
135                    CookieSameSite::None => SameSite::None,
136                }),
137            })
138            .collect();
139
140        Ok(cookies)
141    }
142
143    /// Get cookies for specific URLs.
144    ///
145    /// Note: This method gets all cookies and filters client-side by URL domain matching.
146    ///
147    /// # Errors
148    ///
149    /// Returns an error if getting cookies fails.
150    #[instrument(level = "debug", skip(self, urls))]
151    pub async fn cookies_for_urls(&self, urls: Vec<String>) -> Result<Vec<Cookie>, ContextError> {
152        if self.is_closed() {
153            return Err(ContextError::Closed);
154        }
155
156        // Get all cookies and filter by URL domains
157        let all_cookies = self.cookies().await?;
158
159        // Extract domains from URLs for filtering
160        let domains: Vec<String> = urls
161            .iter()
162            .filter_map(|url| {
163                url::Url::parse(url)
164                    .ok()
165                    .and_then(|u| u.host_str().map(std::string::ToString::to_string))
166            })
167            .collect();
168
169        // Filter cookies that match any of the URL domains
170        let cookies = all_cookies
171            .into_iter()
172            .filter(|c| {
173                if let Some(ref cookie_domain) = c.domain {
174                    domains.iter().any(|d| {
175                        // Cookie domain can start with '.' for subdomain matching
176                        let cookie_domain = cookie_domain.trim_start_matches('.');
177                        d == cookie_domain || d.ends_with(&format!(".{cookie_domain}"))
178                    })
179                } else {
180                    false
181                }
182            })
183            .collect();
184
185        Ok(cookies)
186    }
187
188    /// Get cookies for a specific URL.
189    ///
190    /// # Errors
191    ///
192    /// Returns an error if getting cookies fails.
193    pub async fn cookies_for_url(&self, url: impl Into<String>) -> Result<Vec<Cookie>, ContextError> {
194        self.cookies_for_urls(vec![url.into()]).await
195    }
196
197    /// Clear all cookies in the browser context.
198    ///
199    /// # Errors
200    ///
201    /// Returns an error if clearing cookies fails.
202    #[instrument(level = "debug", skip(self))]
203    pub async fn clear_cookies(&self) -> Result<(), ContextError> {
204        if self.is_closed() {
205            return Err(ContextError::Closed);
206        }
207
208        self.connection()
209            .send_command::<_, serde_json::Value>(
210                "Storage.clearCookies",
211                Some(StorageClearCookiesParams::new().browser_context_id(self.context_id().to_string())),
212                None,
213            )
214            .await?;
215
216        Ok(())
217    }
218
219    /// Create a builder for clearing cookies with filters.
220    pub fn clear_cookies_builder(&self) -> ClearCookiesBuilder<'_> {
221        ClearCookiesBuilder::new(self)
222    }
223}
224
225/// Builder for clearing cookies with filters.
226#[derive(Debug)]
227pub struct ClearCookiesBuilder<'a> {
228    context: &'a BrowserContext,
229    name: Option<String>,
230    domain: Option<String>,
231    path: Option<String>,
232}
233
234impl<'a> ClearCookiesBuilder<'a> {
235    pub(crate) fn new(context: &'a BrowserContext) -> Self {
236        Self {
237            context,
238            name: None,
239            domain: None,
240            path: None,
241        }
242    }
243
244    /// Filter by cookie name.
245    #[must_use]
246    pub fn name(mut self, name: impl Into<String>) -> Self {
247        self.name = Some(name.into());
248        self
249    }
250
251    /// Filter by domain.
252    #[must_use]
253    pub fn domain(mut self, domain: impl Into<String>) -> Self {
254        self.domain = Some(domain.into());
255        self
256    }
257
258    /// Filter by path.
259    #[must_use]
260    pub fn path(mut self, path: impl Into<String>) -> Self {
261        self.path = Some(path.into());
262        self
263    }
264
265    /// Execute the clear operation.
266    ///
267    /// # Errors
268    ///
269    /// Returns an error if clearing cookies fails.
270    pub async fn execute(self) -> Result<(), ContextError> {
271        if self.context.is_closed() {
272            return Err(ContextError::Closed);
273        }
274
275        // If no filters, clear all cookies
276        if self.name.is_none() && self.domain.is_none() && self.path.is_none() {
277            return self.context.clear_cookies().await;
278        }
279
280        // Get all cookies and filter
281        let cookies = self.context.cookies().await?;
282
283        for cookie in cookies {
284            let matches_name = self
285                .name
286                .as_ref()
287                .is_none_or(|n| cookie.name == *n);
288            let matches_domain = self
289                .domain
290                .as_ref()
291                .is_none_or(|d| cookie.domain.as_deref() == Some(d.as_str()));
292            let matches_path = self
293                .path
294                .as_ref()
295                .is_none_or(|p| cookie.path.as_deref() == Some(p.as_str()));
296
297            if matches_name && matches_domain && matches_path {
298                let mut params = StorageDeleteCookiesParams::new(&cookie.name)
299                    .browser_context_id(self.context.context_id().to_string());
300                if let Some(domain) = &cookie.domain {
301                    params = params.domain(domain);
302                }
303                if let Some(path) = &cookie.path {
304                    params = params.path(path);
305                }
306
307                self.context
308                    .connection()
309                    .send_command::<_, serde_json::Value>("Storage.deleteCookies", Some(params), None)
310                    .await?;
311            }
312        }
313
314        Ok(())
315    }
316}
viewpoint_core/context/cookies/mod.rs

viewpoint_core/context/cookies/
mod.rs
