volo_http/client/
cookie.rs

1//! Cookie implementation
2//!
3//! This module provides [`CookieLayer`] for extracting and setting cookies.
4//!
5//! See [`CookieLayer`] for more details.
6
7use motore::{Service, layer::Layer};
8use tokio::sync::RwLock;
9
10use crate::{
11    context::ClientContext,
12    error::ClientError,
13    request::{Request, RequestPartsExt},
14    response::Response,
15    utils::cookie::CookieStore,
16};
17
18/// [`CookieLayer`] generated [`Service`]
19///
20/// See [`CookieLayer`] for more details.
21pub struct CookieService<S> {
22    inner: S,
23    cookie_store: RwLock<CookieStore>,
24}
25
26impl<S> CookieService<S> {
27    fn new(inner: S, cookie_store: RwLock<CookieStore>) -> Self {
28        Self {
29            inner,
30            cookie_store,
31        }
32    }
33}
34
35impl<S, ReqBody, RespBody> Service<ClientContext, Request<ReqBody>> for CookieService<S>
36where
37    S: Service<ClientContext, Request<ReqBody>, Response = Response<RespBody>, Error = ClientError>
38        + Send
39        + Sync
40        + 'static,
41    ReqBody: Send,
42    RespBody: Send,
43{
44    type Response = S::Response;
45    type Error = S::Error;
46
47    async fn call(
48        &self,
49        cx: &mut ClientContext,
50        mut req: Request<ReqBody>,
51    ) -> Result<Self::Response, Self::Error> {
52        let url = req.url();
53
54        if let Some(url) = &url {
55            let (mut parts, body) = req.into_parts();
56            if parts.headers.get(http::header::COOKIE).is_none() {
57                self.cookie_store
58                    .read()
59                    .await
60                    .add_cookie_header(&mut parts.headers, url);
61            }
62            req = Request::from_parts(parts, body);
63        }
64
65        let resp = self.inner.call(cx, req).await?;
66
67        if let Some(url) = &url {
68            self.cookie_store
69                .write()
70                .await
71                .store_response_headers(resp.headers(), url);
72        }
73
74        Ok(resp)
75    }
76}
77
78/// [`Layer`] for extracting and setting cookies.
79///
80/// See [`CookieLayer::new`] for more details.
81pub struct CookieLayer {
82    cookie_store: RwLock<CookieStore>,
83}
84
85impl CookieLayer {
86    /// Create a new [`CookieLayer`] with the given [` CookieStore`](cookie_store::CookieStore).
87    ///
88    /// It will set cookies from the [`CookieStore`](cookie_store::CookieStore) into the request
89    /// header before sending the request,
90    ///
91    /// and store cookies after receiving the response.
92    ///
93    /// It is recommended to use [`CookieLayer`] as the innermost layer in the client stack
94    /// since it will extract cookies from the request header and store them before and after call
95    /// the transport layer.
96    ///
97    /// # Example
98    ///
99    /// ```rust
100    /// use volo_http::{Client, client::cookie::CookieLayer};
101    ///
102    /// let client: Client = Client::builder()
103    ///     .layer_inner(CookieLayer::new(Default::default()))
104    ///     .build()
105    ///     .unwrap();
106    /// ```
107    pub fn new(cookie_store: cookie_store::CookieStore) -> Self {
108        Self {
109            cookie_store: RwLock::new(CookieStore::new(cookie_store)),
110        }
111    }
112}
113
114impl<S> Layer<S> for CookieLayer {
115    type Service = CookieService<S>;
116
117    fn layer(self, inner: S) -> Self::Service {
118        CookieService::new(inner, self.cookie_store)
119    }
120}
volo_http/client/cookie.rs

volo_http/client/
cookie.rs
