feat(http): add cookie feature for client

Cargo.toml

@@ -50,6 +50,7 @@ chrono = { version = "0.4", default-features = false, features = [
 clap = "4"
 colored = "2"
 cookie = "0.18"
+cookie_store = "0.21"
 dashmap = "6"
 dirs = "5"
 faststr = { version = "0.2.21", features = ["serde"] }

volo-http/Cargo.toml

@@ -58,6 +58,7 @@ tokio = { workspace = true, features = [
 ] }
 tokio-util = { workspace = true, features = ["io"] }
 tracing.workspace = true
+url.workspace = true
 
 # =====optional=====
 multer = { workspace = true, optional = true }
@@ -75,6 +76,7 @@ tokio-native-tls = { workspace = true, optional = true }
 
 # cookie support
 cookie = { workspace = true, optional = true, features = ["percent-encode"] }
+cookie_store = { workspace = true, optional = true }
 
 # serde and form, query, json
 serde = { workspace = true, optional = true }
@@ -87,7 +89,6 @@ libc.workspace = true
 serde = { workspace = true, features = ["derive"] }
 reqwest = { workspace = true, features = ["multipart"] }
 tokio-test.workspace = true
-url.workspace = true
 
 [features]
 default = []
@@ -109,7 +110,7 @@ rustls = ["__tls", "dep:tokio-rustls", "volo/rustls"]
 native-tls = ["__tls", "dep:tokio-native-tls", "volo/native-tls"]
 native-tls-vendored = ["native-tls", "volo/native-tls-vendored"]
 
-cookie = ["dep:cookie"]
+cookie = ["dep:cookie", "dep:cookie_store"]
 
 __serde = ["dep:serde"] # a private feature for enabling `serde` by `serde_xxx`
 query = ["__serde", "dep:serde_urlencoded"]

volo-http/src/client/cookie.rs

@@ -0,0 +1,119 @@
+//! Cookie implementation
+//!
+//! This module provides [`CookieLayer`] for extracting and setting cookies.
+//!
+//! See [`CookieLayer`] for more details.
+
+use motore::{layer::Layer, Service};
+use tokio::sync::RwLock;
+
+use crate::{
+    context::ClientContext,
+    error::ClientError,
+    request::{ClientRequest, RequestPartsExt},
+    response::ClientResponse,
+    utils::cookie::CookieStore,
+};
+
+/// [`CookieLayer`] generated [`Service`]
+///
+/// See [`CookieLayer`] for more details.
+pub struct CookieService<S> {
+    inner: S,
+    cookie_store: RwLock<CookieStore>,
+}
+
+impl<S> CookieService<S> {
+    fn new(inner: S, cookie_store: RwLock<CookieStore>) -> Self {
+        Self {
+            inner,
+            cookie_store,
+        }
+    }
+}
+
+impl<S, B> Service<ClientContext, ClientRequest<B>> for CookieService<S>
+where
+    S: Service<ClientContext, ClientRequest<B>, Response = ClientResponse, Error = ClientError>
+        + Send
+        + Sync
+        + 'static,
+    B: Send + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+
+    async fn call(
+        &self,
+        cx: &mut ClientContext,
+        mut req: ClientRequest<B>,
+    ) -> Result<Self::Response, Self::Error> {
+        let url = req.url();
+
+        if let Some(url) = &url {
+            let (mut parts, body) = req.into_parts();
+            if parts.headers.get(http::header::COOKIE).is_none() {
+                self.cookie_store
+                    .read()
+                    .await
+                    .add_cookie_header(&mut parts.headers, url);
+            }
+            req = ClientRequest::from_parts(parts, body);
+        }
+
+        let resp = self.inner.call(cx, req).await?;
+
+        if let Some(url) = &url {
+            self.cookie_store
+                .write()
+                .await
+                .store_response_headers(resp.headers(), url);
+        }
+
+        Ok(resp)
+    }
+}
+
+/// [`Layer`] for extracting and setting cookies.
+///
+/// See [`CookieLayer::new`] for more details.
+pub struct CookieLayer {
+    cookie_store: RwLock<CookieStore>,
+}
+
+impl CookieLayer {
+    /// Create a new [`CookieLayer`] with the given [` CookieStore`](cookie_store::CookieStore).
+    ///
+    /// It will set cookies from the [`CookieStore`](cookie_store::CookieStore) into the request
+    /// header before sending the request,
+    ///
+    /// and store cookies after receiving the response.
+    ///
+    /// It is recommended to use [`CookieLayer`] as the innermost layer in the client stack
+    /// since it will extract cookies from the request header and store them before and after call
+    /// the transport layer.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use volo_http::{client::cookie::CookieLayer, Client};
+    ///
+    /// let builder = Client::builder();
+    /// let client = builder
+    ///     .layer_inner(CookieLayer::new(Default::default()))
+    ///     .build();
+    /// ```
+    pub fn new(cookie_store: cookie_store::CookieStore) -> Self {
+        Self {
+            cookie_store: RwLock::new(CookieStore::new(cookie_store)),
+        }
+    }
+}
+
+impl<S> Layer<S> for CookieLayer {
+    type Service = CookieService<S>;
+
+    fn layer(self, inner: S) -> Self::Service {
+        CookieService::new(inner, self.cookie_store)
+    }
+}

volo-http/src/client/mod.rs

@@ -7,7 +7,7 @@ use std::{cell::RefCell, error::Error, sync::Arc, time::Duration};
 use faststr::FastStr;
 use http::{
     header::{self, HeaderMap, HeaderName, HeaderValue},
-    uri::Uri,
+    uri::{Scheme, Uri},
     Method,
 };
 use metainfo::{MetaInfo, METAINFO};
@@ -48,6 +48,8 @@ use crate::{
 };
 
 pub mod callopt;
+#[cfg(feature = "cookie")]
+pub mod cookie;
 pub mod dns;
 pub mod loadbalance;
 mod meta;
@@ -880,6 +882,12 @@ impl<S> Client<S> {
             request.headers_mut().insert(header::HOST, host);
         }
 
+        let scheme = match target.is_https() {
+            true => Scheme::HTTPS,
+            false => Scheme::HTTP,
+        };
+        request.extensions_mut().insert(scheme);
+
         let mut cx = ClientContext::new();
         cx.rpc_info_mut().caller_mut().set_service_name(caller_name);
         cx.rpc_info_mut().callee_mut().set_service_name(callee_name);
@@ -947,9 +955,9 @@ where
 #[cfg(feature = "json")]
 #[cfg(test)]
 mod client_tests {
-
     use std::{collections::HashMap, future::Future};
 
+    use cookie::Cookie;
     use http::{header, StatusCode};
     use motore::{
         layer::{Layer, Stack},
@@ -963,6 +971,8 @@ mod client_tests {
         dns::{parse_target, DnsResolver},
         get, Client, DefaultClient, Target,
     };
+    #[cfg(feature = "cookie")]
+    use crate::client::cookie::CookieLayer;
     use crate::{
         body::BodyConversion, error::client::status_error, utils::consts::HTTP_DEFAULT_PORT,
         ClientBuilder,
@@ -1284,4 +1294,60 @@ mod client_tests {
         let resp = client.get(HTTPBIN_GET).send().await;
         assert!(resp.is_ok());
     }
+
+    #[cfg(feature = "cookie")]
+    #[tokio::test]
+    async fn cookie_store() {
+        let mut builder = Client::builder().layer_inner(CookieLayer::new(Default::default()));
+
+        builder.host("httpbin.org");
+
+        let client = builder.build();
+
+        // test server add cookie
+        let resp = client
+            .get("http://httpbin.org/cookies/set?key=value")
+            .send()
+            .await
+            .unwrap();
+        let cookies = resp
+            .headers()
+            .get_all(http::header::SET_COOKIE)
+            .iter()
+            .filter_map(|value| {
+                std::str::from_utf8(value.as_bytes())
+                    .ok()
+                    .and_then(|val| Cookie::parse(val).map(|c| c.into_owned()).ok())
+            })
+            .collect::<Vec<_>>();
+        assert_eq!(cookies[0].name(), "key");
+        assert_eq!(cookies[0].value(), "value");
+
+        #[derive(serde::Deserialize)]
+        struct CookieResponse {
+            #[serde(default)]
+            cookies: HashMap<String, String>,
+        }
+        let resp = client
+            .get("http://httpbin.org/cookies")
+            .send()
+            .await
+            .unwrap();
+        let json = resp.into_json::<CookieResponse>().await.unwrap();
+        assert_eq!(json.cookies["key"], "value");
+
+        // test server delete cookie
+        _ = client
+            .get("http://httpbin.org/cookies/delete?key")
+            .send()
+            .await
+            .unwrap();
+        let resp = client
+            .get("http://httpbin.org/cookies")
+            .send()
+            .await
+            .unwrap();
+        let json = resp.into_json::<CookieResponse>().await.unwrap();
+        assert_eq!(json.cookies.len(), 0);
+    }
 }