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::storage::{
8    ClearCookiesParams as StorageClearCookiesParams,
9    DeleteCookiesParams as StorageDeleteCookiesParams, GetCookiesParams as StorageGetCookiesParams,
10    GetCookiesResult as StorageGetCookiesResult, SetCookiesParams as StorageSetCookiesParams,
11};
12use viewpoint_cdp::protocol::{CookieParam, CookieSameSite};
13
14use super::BrowserContext;
15use super::types::{Cookie, SameSite};
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(
112                    StorageGetCookiesParams::new()
113                        .browser_context_id(self.context_id().to_string()),
114                ),
115                None,
116            )
117            .await?;
118
119        let cookies = result
120            .cookies
121            .into_iter()
122            .map(|c| Cookie {
123                name: c.name,
124                value: c.value,
125                domain: Some(c.domain),
126                path: Some(c.path),
127                url: None,
128                expires: if c.expires > 0.0 {
129                    Some(c.expires)
130                } else {
131                    None
132                },
133                http_only: Some(c.http_only),
134                secure: Some(c.secure),
135                same_site: c.same_site.map(|s| match s {
136                    CookieSameSite::Strict => SameSite::Strict,
137                    CookieSameSite::Lax => SameSite::Lax,
138                    CookieSameSite::None => SameSite::None,
139                }),
140            })
141            .collect();
142
143        Ok(cookies)
144    }
145
146    /// Get cookies for specific URLs.
147    ///
148    /// Note: This method gets all cookies and filters client-side by URL domain matching.
149    ///
150    /// # Errors
151    ///
152    /// Returns an error if getting cookies fails.
153    #[instrument(level = "debug", skip(self, urls))]
154    pub async fn cookies_for_urls(&self, urls: Vec<String>) -> Result<Vec<Cookie>, ContextError> {
155        if self.is_closed() {
156            return Err(ContextError::Closed);
157        }
158
159        // Get all cookies and filter by URL domains
160        let all_cookies = self.cookies().await?;
161
162        // Extract domains from URLs for filtering
163        let domains: Vec<String> = urls
164            .iter()
165            .filter_map(|url| {
166                url::Url::parse(url)
167                    .ok()
168                    .and_then(|u| u.host_str().map(std::string::ToString::to_string))
169            })
170            .collect();
171
172        // Filter cookies that match any of the URL domains
173        let cookies = all_cookies
174            .into_iter()
175            .filter(|c| {
176                if let Some(ref cookie_domain) = c.domain {
177                    domains.iter().any(|d| {
178                        // Cookie domain can start with '.' for subdomain matching
179                        let cookie_domain = cookie_domain.trim_start_matches('.');
180                        d == cookie_domain || d.ends_with(&format!(".{cookie_domain}"))
181                    })
182                } else {
183                    false
184                }
185            })
186            .collect();
187
188        Ok(cookies)
189    }
190
191    /// Get cookies for a specific URL.
192    ///
193    /// # Errors
194    ///
195    /// Returns an error if getting cookies fails.
196    pub async fn cookies_for_url(
197        &self,
198        url: impl Into<String>,
199    ) -> Result<Vec<Cookie>, ContextError> {
200        self.cookies_for_urls(vec![url.into()]).await
201    }
202
203    /// Clear all cookies in the browser context.
204    ///
205    /// # Errors
206    ///
207    /// Returns an error if clearing cookies fails.
208    #[instrument(level = "debug", skip(self))]
209    pub async fn clear_cookies(&self) -> Result<(), ContextError> {
210        if self.is_closed() {
211            return Err(ContextError::Closed);
212        }
213
214        self.connection()
215            .send_command::<_, serde_json::Value>(
216                "Storage.clearCookies",
217                Some(
218                    StorageClearCookiesParams::new()
219                        .browser_context_id(self.context_id().to_string()),
220                ),
221                None,
222            )
223            .await?;
224
225        Ok(())
226    }
227
228    /// Create a builder for clearing cookies with filters.
229    pub fn clear_cookies_builder(&self) -> ClearCookiesBuilder<'_> {
230        ClearCookiesBuilder::new(self)
231    }
232}
233
234/// Builder for clearing cookies with filters.
235#[derive(Debug)]
236pub struct ClearCookiesBuilder<'a> {
237    context: &'a BrowserContext,
238    name: Option<String>,
239    domain: Option<String>,
240    path: Option<String>,
241}
242
243impl<'a> ClearCookiesBuilder<'a> {
244    pub(crate) fn new(context: &'a BrowserContext) -> Self {
245        Self {
246            context,
247            name: None,
248            domain: None,
249            path: None,
250        }
251    }
252
253    /// Filter by cookie name.
254    #[must_use]
255    pub fn name(mut self, name: impl Into<String>) -> Self {
256        self.name = Some(name.into());
257        self
258    }
259
260    /// Filter by domain.
261    #[must_use]
262    pub fn domain(mut self, domain: impl Into<String>) -> Self {
263        self.domain = Some(domain.into());
264        self
265    }
266
267    /// Filter by path.
268    #[must_use]
269    pub fn path(mut self, path: impl Into<String>) -> Self {
270        self.path = Some(path.into());
271        self
272    }
273
274    /// Execute the clear operation.
275    ///
276    /// # Errors
277    ///
278    /// Returns an error if clearing cookies fails.
279    pub async fn execute(self) -> Result<(), ContextError> {
280        if self.context.is_closed() {
281            return Err(ContextError::Closed);
282        }
283
284        // If no filters, clear all cookies
285        if self.name.is_none() && self.domain.is_none() && self.path.is_none() {
286            return self.context.clear_cookies().await;
287        }
288
289        // Get all cookies and filter
290        let cookies = self.context.cookies().await?;
291
292        for cookie in cookies {
293            let matches_name = self.name.as_ref().is_none_or(|n| cookie.name == *n);
294            let matches_domain = self
295                .domain
296                .as_ref()
297                .is_none_or(|d| cookie.domain.as_deref() == Some(d.as_str()));
298            let matches_path = self
299                .path
300                .as_ref()
301                .is_none_or(|p| cookie.path.as_deref() == Some(p.as_str()));
302
303            if matches_name && matches_domain && matches_path {
304                let mut params = StorageDeleteCookiesParams::new(&cookie.name)
305                    .browser_context_id(self.context.context_id().to_string());
306                if let Some(domain) = &cookie.domain {
307                    params = params.domain(domain);
308                }
309                if let Some(path) = &cookie.path {
310                    params = params.path(path);
311                }
312
313                self.context
314                    .connection()
315                    .send_command::<_, serde_json::Value>(
316                        "Storage.deleteCookies",
317                        Some(params),
318                        None,
319                    )
320                    .await?;
321            }
322        }
323
324        Ok(())
325    }
326}
viewpoint_core/context/cookies/mod.rs

viewpoint_core/context/cookies/
mod.rs
