feat(layer): 分离请求和响应头移除功能 (#27)

Author: sodekim

Cargo.lock

@@ -19,6 +19,7 @@ serde = { workspace = true, features = ["derive"] }
 tower = { workspace = true, features = ["limit"] }
 tower-http = { workspace = true, features = ["trace", "set-header", "timeout", "cors"] }
 tracing = { workspace = true }
+pin-project-lite = { workspace = true }
 
 [dev-dependencies]
 serde_yaml = { workspace = true }

crates/layer/Cargo.toml

@@ -0,0 +1,33 @@
+mod remove_request_header;
+mod remove_response_header;
+
+use http::{HeaderName, Request, Response};
+pub use remove_request_header::{
+    MakeRemoveRequestHeaderRouteLayer, RemoveRequestHeader, RemoveRequestHeaderLayer,
+};
+pub use remove_response_header::{
+    MakeRemoveResponseHeaderRouteLayer, RemoveResponseHeader, RemoveResponseHeaderLayer,
+    ResponseFuture,
+};
+
+pub trait Removable {
+    fn remove(&mut self, name: &HeaderName);
+}
+
+impl<ResBody> Removable for Response<ResBody> {
+    fn remove(&mut self, name: &HeaderName) {
+        let headers = self.headers_mut();
+        if headers.contains_key(name) {
+            headers.remove(name);
+        }
+    }
+}
+
+impl<ReqBody> Removable for Request<ReqBody> {
+    fn remove(&mut self, name: &HeaderName) {
+        let headers = self.headers_mut();
+        if headers.contains_key(name) {
+            headers.remove(name);
+        }
+    }
+}

crates/layer/src/remove_header/mod.rs

@@ -1,64 +1,64 @@
 use std::{str::FromStr, sync::Arc};
 
 use crate::make::MakeRouteLayer;
-use http::{HeaderName, Request};
+use crate::remove_header::Removable;
+use http::HeaderName;
 use satex_core::component::Args;
 use satex_core::{component::Configurable, Error};
 use satex_macro::make;
 use tower::{Layer, Service};
 
-#[make(kind = RemoveHeader)]
-pub struct MakeRemoveHeaderRouteLayer {
+#[make(kind = RemoveRequestHeader)]
+pub struct MakeRemoveRequestHeaderRouteLayer {
     name: String,
 }
 
-impl MakeRouteLayer for MakeRemoveHeaderRouteLayer {
-    type Layer = RemoveHeaderRouteLayer;
+impl MakeRouteLayer for MakeRemoveRequestHeaderRouteLayer {
+    type Layer = RemoveRequestHeaderLayer;
 
     fn make(&self, args: Args) -> Result<Self::Layer, Error> {
         Config::with_args(args).and_then(|config| {
             HeaderName::from_str(&config.name)
                 .map_err(Error::new)
-                .map(|name| RemoveHeaderRouteLayer {
-                    name: Arc::new(name),
-                })
+                .map(RemoveRequestHeaderLayer::new)
         })
     }
 }
 
 #[derive(Debug, Clone)]
-pub struct RemoveHeaderRouteLayer {
+pub struct RemoveRequestHeaderLayer {
     name: Arc<HeaderName>,
 }
 
-impl RemoveHeaderRouteLayer {
-    pub fn new(name: &str) -> Self {
+impl RemoveRequestHeaderLayer {
+    pub fn new(name: HeaderName) -> Self {
         Self {
-            name: Arc::new(HeaderName::from_str(name).unwrap()),
+            name: Arc::new(name),
         }
     }
 }
 
-impl<S> Layer<S> for RemoveHeaderRouteLayer {
-    type Service = RemoveHeader<S>;
+impl<S> Layer<S> for RemoveRequestHeaderLayer {
+    type Service = RemoveRequestHeader<S>;
 
     fn layer(&self, inner: S) -> Self::Service {
-        RemoveHeader {
+        RemoveRequestHeader {
             name: self.name.clone(),
             inner,
         }
     }
 }
 
 #[derive(Debug, Clone)]
-pub struct RemoveHeader<S> {
+pub struct RemoveRequestHeader<S> {
     inner: S,
     name: Arc<HeaderName>,
 }
 
-impl<S, ReqBody> Service<Request<ReqBody>> for RemoveHeader<S>
+impl<S, Req> Service<Req> for RemoveRequestHeader<S>
 where
-    S: Service<Request<ReqBody>>,
+    S: Service<Req>,
+    Req: Removable,
 {
     type Response = S::Response;
 
@@ -73,9 +73,8 @@ where
         self.inner.poll_ready(cx)
     }
 
-    fn call(&mut self, mut request: Request<ReqBody>) -> Self::Future {
-        let headers = request.headers_mut();
-        headers.remove(self.name.as_ref());
-        self.inner.call(request)
+    fn call(&mut self, mut req: Req) -> Self::Future {
+        req.remove(self.name.as_ref());
+        self.inner.call(req)
     }
 }

crates/layer/src/remove_header.rs …c/remove_header/remove_request_header.rscrates/layer/src/remove_header.rs renamed to crates/layer/src/remove_header/remove_request_header.rs crates/layer/src/remove_header.rs renamed to crates/layer/src/remove_header/remove_request_header.rs

@@ -0,0 +1,124 @@
+use crate::make::MakeRouteLayer;
+use crate::remove_header::Removable;
+use http::HeaderName;
+use pin_project_lite::pin_project;
+use satex_core::component::Args;
+use satex_core::{component::Configurable, Error};
+use satex_macro::make;
+use std::pin::Pin;
+use std::task::{ready, Context, Poll};
+use std::{str::FromStr, sync::Arc};
+use tower::{Layer, Service};
+
+#[make(kind = RemoveResponseHeader)]
+pub struct MakeRemoveResponseHeaderRouteLayer {
+    name: String,
+}
+
+impl MakeRouteLayer for MakeRemoveResponseHeaderRouteLayer {
+    type Layer = RemoveResponseHeaderLayer;
+
+    fn make(&self, args: Args) -> Result<Self::Layer, Error> {
+        Config::with_args(args).and_then(|config| {
+            HeaderName::from_str(&config.name)
+                .map_err(Error::new)
+                .map(RemoveResponseHeaderLayer::new)
+        })
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct RemoveResponseHeaderLayer {
+    name: Arc<HeaderName>,
+}
+
+impl RemoveResponseHeaderLayer {
+    pub fn new(name: HeaderName) -> Self {
+        Self {
+            name: Arc::new(name),
+        }
+    }
+}
+
+impl<S> Layer<S> for RemoveResponseHeaderLayer {
+    type Service = RemoveResponseHeader<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        RemoveResponseHeader {
+            name: self.name.clone(),
+            inner,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct RemoveResponseHeader<S> {
+    inner: S,
+    name: Arc<HeaderName>,
+}
+
+impl<S, Req, Res> Service<Req> for RemoveResponseHeader<S>
+where
+    S: Service<Req, Response=Res>,
+    Res: Removable,
+{
+    type Response = S::Response;
+
+    type Error = S::Error;
+
+    type Future = ResponseFuture<S::Future>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, request: Req) -> Self::Future {
+        let name = self.name.clone();
+        ResponseFuture::new(name, self.inner.call(request))
+    }
+}
+
+pin_project! {
+    #[project = ResponseFutureProj]
+    #[project_replace = MapProjReplace]
+    pub enum ResponseFuture<F> {
+        Incomplete {
+            name: Arc<HeaderName>,
+            #[pin]
+            future: F,
+        },
+        Complete,
+    }
+}
+
+impl<F> ResponseFuture<F> {
+    pub fn new(name: Arc<HeaderName>, future: F) -> Self {
+        Self::Incomplete { name, future }
+    }
+}
+
+impl<F, Res, E> Future for ResponseFuture<F>
+where
+    F: Future<Output=Result<Res, E>>,
+    Res: Removable,
+{
+    type Output = F::Output;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.as_mut().project() {
+            ResponseFutureProj::Incomplete { future, .. } => match ready!(future.poll(cx)) {
+                Ok(mut res) => match self.project_replace(ResponseFuture::Complete) {
+                    MapProjReplace::Incomplete { name, .. } => {
+                        res.remove(name.as_ref());
+                        Poll::Ready(Ok(res))
+                    }
+                    MapProjReplace::Complete => unreachable!(),
+                },
+                Err(e) => Poll::Ready(Err(e)),
+            },
+            ResponseFutureProj::Complete { .. } => {
+                panic!("polled after completion")
+            }
+        }
+    }
+}

crates/layer/src/remove_header/remove_response_header.rs

@@ -0,0 +1,39 @@
+mod identify;
+
+use crate::identify::Identify;
+use http::header::HOST;
+use http::{HeaderValue, Request};
+use satex_core::body::Body;
+use satex_core::component::Args;
+use satex_layer::make::MakeRouteLayer;
+use satex_layer::remove_header::{MakeRemoveRequestHeaderRouteLayer, RemoveRequestHeaderLayer};
+use tower::{Layer, Service};
+
+async fn _layer(layer: RemoveRequestHeaderLayer) {
+    let mut service = layer.layer(Identify::new(|request: Request<Body>| {
+        !request.headers().contains_key(HOST)
+    }));
+    let mut request = Request::builder().uri("/").body(Body::empty()).unwrap();
+    request
+        .headers_mut()
+        .insert(HOST, HeaderValue::from_static("127.0.0.1:3000"));
+
+    let result = service.call(request).await;
+    assert!(result.is_ok());
+}
+
+#[tokio::test]
+async fn make_with_shortcut() {
+    let args = Args::shortcut("host");
+    let layer = MakeRemoveRequestHeaderRouteLayer.make(args).unwrap();
+    _layer(layer).await;
+}
+
+#[tokio::test]
+async fn make_with_full() {
+    let yaml = "name: host";
+    let value = serde_yaml::from_str(yaml).unwrap();
+    let args = Args::full(&value);
+    let layer = MakeRemoveRequestHeaderRouteLayer.make(args).unwrap();
+    _layer(layer).await;
+}

crates/layer/tests/remove_request_header.rs

@@ -0,0 +1,40 @@
+use http::header::CONTENT_TYPE;
+use http::{HeaderValue, Request, Response};
+use satex_core::body::Body;
+use satex_core::component::Args;
+use satex_layer::make::MakeRouteLayer;
+use satex_layer::remove_header::{
+    MakeRemoveResponseHeaderRouteLayer, RemoveResponseHeaderLayer,
+};
+use std::convert::Infallible;
+use tower::{service_fn, Layer, Service};
+
+async fn _layer(layer: RemoveResponseHeaderLayer) {
+    let mut service = layer.layer(service_fn(|_| async move {
+        let mut response = Response::new(Body::empty());
+        response
+            .headers_mut()
+            .insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
+        Ok::<_, Infallible>(response)
+    }));
+    let request = Request::new(Body::empty());
+
+    let response = service.call(request).await.unwrap();
+    assert!(!response.headers().contains_key(CONTENT_TYPE))
+}
+
+#[tokio::test]
+async fn make_with_shortcut() {
+    let args = Args::shortcut("content-type");
+    let layer = MakeRemoveResponseHeaderRouteLayer.make(args).unwrap();
+    _layer(layer).await;
+}
+
+#[tokio::test]
+async fn make_with_full() {
+    let yaml = "name: content-type";
+    let value = serde_yaml::from_str(yaml).unwrap();
+    let args = Args::full(&value);
+    let layer = MakeRemoveResponseHeaderRouteLayer.make(args).unwrap();
+    _layer(layer).await;
+}

crates/layer/tests/remove_response_header.rs

@@ -6,7 +6,9 @@ use satex_core::BoxError;
 use satex_layer::concurrency_limit::MakeConcurrencyLimitRouteLayer;
 use satex_layer::cors::MakeCorsRouteLayer;
 use satex_layer::make::{ArcMakeRouteLayer, MakeRouteLayer};
-use satex_layer::remove_header::MakeRemoveHeaderRouteLayer;
+use satex_layer::remove_header::{
+    MakeRemoveRequestHeaderRouteLayer, MakeRemoveResponseHeaderRouteLayer,
+};
 use satex_layer::set_header::{MakeSetRequestHeaderRouteLayer, MakeSetResponseHeaderRouteLayer};
 use satex_layer::set_method::MakeSetMethodRouteLayer;
 use satex_layer::set_prefix::MakeSetPrefixRouteLayer;
@@ -34,7 +36,8 @@ impl MakeRouteLayerRegistry {
             MakeSetMethodRouteLayer,
             MakeSetRequestHeaderRouteLayer,
             MakeSetResponseHeaderRouteLayer,
-            MakeRemoveHeaderRouteLayer,
+            MakeRemoveRequestHeaderRouteLayer,
+            MakeRemoveResponseHeaderRouteLayer,
             MakeStripPrefixRouteLayer,
             MakeTraceRouteLayer,
             MakeTimeoutRouteLayer,