refactor(cookie): de-duplicate cookie support as CookieService middleware (#2787)

Signed-off-by: Yihai Lin <yihai-lin@foxmai.com>
Co-authored-by: Yihai Lin <yihai-lin@foxmai.com>

Author: linyihai

src/async_impl/client.rs

@@ -26,6 +26,8 @@ use crate::connect::{
 };
 #[cfg(feature = "cookies")]
 use crate::cookie;
+#[cfg(feature = "cookies")]
+use crate::cookie::service::CookieService;
 #[cfg(feature = "hickory-dns")]
 use crate::dns::hickory::HickoryDnsResolver;
 use crate::dns::{gai::GaiResolver, DnsResolverWithOverrides, DynResolver, Resolve};
@@ -96,8 +98,6 @@ enum HttpVersionPref {
 
 #[derive(Clone)]
 struct HyperService {
-    #[cfg(feature = "cookies")]
-    cookie_store: Option<Arc<dyn cookie::CookieStore>>,
     hyper: HyperClient,
 }
 
@@ -110,43 +110,11 @@ impl Service<hyper::Request<crate::async_impl::body::Body>> for HyperService {
         self.hyper.poll_ready(cx).map_err(crate::error::request)
     }
 
-    #[cfg(not(feature = "cookies"))]
     fn call(&mut self, req: hyper::Request<crate::async_impl::body::Body>) -> Self::Future {
         let clone = self.hyper.clone();
         let mut inner = std::mem::replace(&mut self.hyper, clone);
         Box::pin(async move { inner.call(req).await.map_err(crate::error::request) })
     }
-
-    #[cfg(feature = "cookies")]
-    fn call(&mut self, mut req: hyper::Request<crate::async_impl::body::Body>) -> Self::Future {
-        let clone = self.hyper.clone();
-        let mut inner = std::mem::replace(&mut self.hyper, clone);
-        let url = Url::parse(req.uri().to_string().as_str()).expect("invalid URL");
-
-        if let Some(cookie_store) = self.cookie_store.as_ref() {
-            if req.headers().get(crate::header::COOKIE).is_none() {
-                let headers = req.headers_mut();
-                crate::util::add_cookie_header(headers, &**cookie_store, &url);
-            }
-        }
-
-        let cookie_store = self.cookie_store.clone();
-        Box::pin(async move {
-            let res = inner.call(req).await.map_err(crate::error::request);
-
-            if let Some(ref cookie_store) = cookie_store {
-                if let Ok(res) = &res {
-                    let mut cookies =
-                        cookie::extract_response_cookie_headers(res.headers()).peekable();
-                    if cookies.peek().is_some() {
-                        cookie_store.set_cookies(&mut cookies, &url);
-                    }
-                }
-            }
-
-            res
-        })
-    }
 }
 
 struct Config {
@@ -993,8 +961,6 @@ impl ClientBuilder {
 
         let hyper_client = builder.build(connector_builder.build(config.connector_layers));
         let hyper_service = HyperService {
-            #[cfg(feature = "cookies")]
-            cookie_store: config.cookie_store.clone(),
             hyper: hyper_client,
         };
 
@@ -1007,9 +973,11 @@ impl ClientBuilder {
 
         let retry_policy = config.retry_policy.into_policy();
 
-        let retries = tower::retry::Retry::new(retry_policy.clone(), hyper_service);
+        let svc = tower::retry::Retry::new(retry_policy.clone(), hyper_service);
 
-        let hyper = FollowRedirect::with_policy(retries, redirect_policy.clone());
+        #[cfg(feature = "cookies")]
+        let svc = CookieService::new(svc, config.cookie_store.clone());
+        let hyper = FollowRedirect::with_policy(svc, redirect_policy.clone());
 
         Ok(Client {
             inner: Arc::new(ClientRef {
@@ -1021,16 +989,11 @@ impl ClientBuilder {
                 #[cfg(feature = "http3")]
                 h3_client: match h3_connector {
                     Some(h3_connector) => {
-                        #[cfg(not(feature = "cookies"))]
                         let h3_service = H3Client::new(h3_connector, config.pool_idle_timeout);
+                        let svc = tower::retry::Retry::new(retry_policy, h3_service);
                         #[cfg(feature = "cookies")]
-                        let h3_service = H3Client::new(
-                            h3_connector,
-                            config.pool_idle_timeout,
-                            config.cookie_store,
-                        );
-                        let retries = tower::retry::Retry::new(retry_policy, h3_service);
-                        Some(FollowRedirect::with_policy(retries, redirect_policy))
+                        let svc = CookieService::new(svc, config.cookie_store);
+                        Some(FollowRedirect::with_policy(svc, redirect_policy))
                     }
                     None => None,
                 },
@@ -2812,8 +2775,14 @@ impl Config {
     }
 }
 
+#[cfg(not(feature = "cookies"))]
 type LayeredService<T> =
     FollowRedirect<tower::retry::Retry<crate::retry::Policy, T>, TowerRedirectPolicy>;
+#[cfg(feature = "cookies")]
+type LayeredService<T> = FollowRedirect<
+    CookieService<tower::retry::Retry<crate::retry::Policy, T>>,
+    TowerRedirectPolicy,
+>;
 type LayeredFuture<T> = <LayeredService<T> as Service<http::Request<Body>>>::Future;
 
 struct ClientRef {

src/async_impl/h3_client/mod.rs

@@ -6,17 +6,13 @@ mod pool;
 
 use crate::async_impl::body::ResponseBody;
 use crate::async_impl::h3_client::pool::{Key, Pool, PoolClient};
-#[cfg(feature = "cookies")]
-use crate::cookie;
 use crate::error::{BoxError, Error, Kind};
 use crate::{error, Body};
 use connect::H3Connector;
 use http::{Request, Response};
 use log::trace;
 use std::future::{self, Future};
 use std::pin::Pin;
-#[cfg(feature = "cookies")]
-use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::time::Duration;
 use sync_wrapper::SyncWrapper;
@@ -26,32 +22,16 @@ use tower::Service;
 pub(crate) struct H3Client {
     pool: Pool,
     connector: H3Connector,
-    #[cfg(feature = "cookies")]
-    cookie_store: Option<Arc<dyn cookie::CookieStore>>,
 }
 
 impl H3Client {
-    #[cfg(not(feature = "cookies"))]
     pub fn new(connector: H3Connector, pool_timeout: Option<Duration>) -> Self {
         H3Client {
             pool: Pool::new(pool_timeout),
             connector,
         }
     }
 
-    #[cfg(feature = "cookies")]
-    pub fn new(
-        connector: H3Connector,
-        pool_timeout: Option<Duration>,
-        cookie_store: Option<Arc<dyn cookie::CookieStore>>,
-    ) -> Self {
-        H3Client {
-            pool: Pool::new(pool_timeout),
-            connector,
-            cookie_store,
-        }
-    }
-
     async fn get_pooled_client(&mut self, key: Key) -> Result<PoolClient, BoxError> {
         if let Some(client) = self.pool.try_pool(&key) {
             trace!("getting client from pool with key {key:?}");
@@ -79,7 +59,6 @@ impl H3Client {
         Ok(self.pool.new_connection(lock, driver, tx))
     }
 
-    #[cfg(not(feature = "cookies"))]
     async fn send_request(
         mut self,
         key: Key,
@@ -95,42 +74,6 @@ impl H3Client {
             .map_err(|e| Error::new(Kind::Request, Some(e)))
     }
 
-    #[cfg(feature = "cookies")]
-    async fn send_request(
-        mut self,
-        key: Key,
-        mut req: Request<Body>,
-    ) -> Result<Response<ResponseBody>, Error> {
-        let mut pooled = match self.get_pooled_client(key).await {
-            Ok(client) => client,
-            Err(e) => return Err(error::request(e)),
-        };
-
-        let url = url::Url::parse(req.uri().to_string().as_str()).unwrap();
-        if let Some(cookie_store) = self.cookie_store.as_ref() {
-            if req.headers().get(crate::header::COOKIE).is_none() {
-                let headers = req.headers_mut();
-                crate::util::add_cookie_header(headers, &**cookie_store, &url);
-            }
-        }
-
-        let res = pooled
-            .send_request(req)
-            .await
-            .map_err(|e| Error::new(Kind::Request, Some(e)));
-
-        if let Some(ref cookie_store) = self.cookie_store {
-            if let Ok(res) = &res {
-                let mut cookies = cookie::extract_response_cookie_headers(res.headers()).peekable();
-                if cookies.peek().is_some() {
-                    cookie_store.set_cookies(&mut cookies, &url);
-                }
-            }
-        }
-
-        res
-    }
-
     pub fn request(&self, mut req: Request<Body>) -> H3ResponseFuture {
         let pool_key = match pool::extract_domain(req.uri_mut()) {
             Ok(s) => s,

src/cookie.rs

@@ -1,13 +1,12 @@
 //! HTTP Cookies
 
+use crate::header::{HeaderValue, SET_COOKIE};
+use bytes::Bytes;
 use std::convert::TryInto;
 use std::fmt;
 use std::sync::RwLock;
 use std::time::SystemTime;
 
-use crate::header::{HeaderValue, SET_COOKIE};
-use bytes::Bytes;
-
 /// Actions for a persistent cookie store providing session support.
 pub trait CookieStore: Send + Sync {
     /// Store a set of Set-Cookie header values received from `url`
@@ -189,3 +188,106 @@ impl CookieStore for Jar {
         HeaderValue::from_maybe_shared(Bytes::from(s)).ok()
     }
 }
+
+pub(crate) mod service {
+    use crate::cookie;
+    use http::{Request, Response};
+    use http_body::Body;
+    use pin_project_lite::pin_project;
+    use std::future::Future;
+    use std::pin::Pin;
+    use std::sync::Arc;
+    use std::task::ready;
+    use std::task::Context;
+    use std::task::Poll;
+    use tower::Service;
+    use url::Url;
+
+    /// A [`Service`] that adds cookie support to a lower-level [`Service`].
+    #[derive(Clone)]
+    pub struct CookieService<S> {
+        inner: S,
+        cookie_store: Option<Arc<dyn cookie::CookieStore>>,
+    }
+
+    impl<S> CookieService<S> {
+        /// Create a new [`CookieService`].
+        pub fn new(inner: S, cookie_store: Option<Arc<dyn cookie::CookieStore>>) -> Self {
+            Self {
+                inner,
+                cookie_store,
+            }
+        }
+    }
+
+    impl<ReqBody, ResBody, S> Service<Request<ReqBody>> for CookieService<S>
+    where
+        S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone,
+        ReqBody: Body + Default,
+    {
+        type Response = Response<ResBody>;
+        type Error = S::Error;
+        type Future = ResponseFuture<S, ReqBody>;
+
+        #[inline]
+        fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            self.inner.poll_ready(cx)
+        }
+
+        fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
+            let clone = self.inner.clone();
+            let mut inner = std::mem::replace(&mut self.inner, clone);
+            let url = Url::parse(req.uri().to_string().as_str()).expect("invalid URL");
+            if let Some(cookie_store) = self.cookie_store.as_ref() {
+                if req.headers().get(crate::header::COOKIE).is_none() {
+                    let headers = req.headers_mut();
+                    crate::util::add_cookie_header(headers, &**cookie_store, &url);
+                }
+            }
+
+            let cookie_store = self.cookie_store.clone();
+            ResponseFuture {
+                future: inner.call(req),
+                cookie_store,
+                url,
+            }
+        }
+    }
+
+    pin_project! {
+        #[allow(missing_debug_implementations)]
+        #[derive(Clone)]
+        /// A [`Future`] that adds cookie support to a lower-level [`Future`].
+        pub struct ResponseFuture<S, B>
+        where
+            S: Service<Request<B>>,
+        {
+            #[pin]
+            future: S::Future,
+            cookie_store: Option<Arc<dyn cookie::CookieStore>>,
+            url: Url,
+        }
+    }
+
+    impl<S, ReqBody, ResBody> Future for ResponseFuture<S, ReqBody>
+    where
+        S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone,
+        ReqBody: Body + Default,
+    {
+        type Output = Result<Response<ResBody>, S::Error>;
+
+        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            let cookie_store = self.cookie_store.clone();
+            let url = self.url.clone();
+            let res = ready!(self.project().future.as_mut().poll(cx)?);
+
+            if let Some(cookie_store) = cookie_store.as_ref() {
+                let mut cookies = cookie::extract_response_cookie_headers(res.headers()).peekable();
+                if cookies.peek().is_some() {
+                    cookie_store.set_cookies(&mut cookies, &url);
+                }
+            }
+            Poll::Ready(Ok(res))
+        }
+    }
+}