From 8a9fcddb3cad417bba5119111e044bb6ff11771d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=CE=BAeen?= <3han5chou7@gmail.com>
Date: Mon, 9 Sep 2019 15:26:38 +0900
Subject: [PATCH] Condition middleware (#1075)

* add condition middleware

* write tests

* update changes

* Update src/middleware/condition.rs

Co-Authored-By: Yuki Okushi <huyuumi.dev@gmail.com>

* Update src/middleware/condition.rs

Co-Authored-By: Yuki Okushi <huyuumi.dev@gmail.com>

* Update src/middleware/condition.rs

Co-Authored-By: Yuki Okushi <huyuumi.dev@gmail.com>

* Update src/middleware/condition.rs

Co-Authored-By: Yuki Okushi <huyuumi.dev@gmail.com>
---
 CHANGES.md                  |   6 ++
 src/middleware/condition.rs | 143 ++++++++++++++++++++++++++++++++++++
 src/middleware/mod.rs       |   2 +
 3 files changed, 151 insertions(+)
 create mode 100644 src/middleware/condition.rs

diff --git a/CHANGES.md b/CHANGES.md
index 5f8f489fc..57304d083 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,4 +1,10 @@
 # Changes
+## not released yet
+
+### Added
+
+*  Add `middleware::Conditon` that conditionally enables another middleware
+
 
 ## [1.0.7] - 2019-08-29
 
diff --git a/src/middleware/condition.rs b/src/middleware/condition.rs
new file mode 100644
index 000000000..ddc5fdd42
--- /dev/null
+++ b/src/middleware/condition.rs
@@ -0,0 +1,143 @@
+//! `Middleware` for conditionally enables another middleware.
+use actix_service::{Service, Transform};
+use futures::future::{ok, Either, FutureResult, Map};
+use futures::{Future, Poll};
+
+/// `Middleware` for conditionally enables another middleware.
+/// The controled middleware must not change the `Service` interfaces.
+/// This means you cannot control such middlewares like `Logger` or `Compress`.
+///
+/// ## Usage
+///
+/// ```rust
+/// use actix_web::middleware::{Condition, NormalizePath};
+/// use actix_web::App;
+///
+/// fn main() {
+///     let enable_normalize = std::env::var("NORMALIZE_PATH") == Ok("true".into());
+///     let app = App::new()
+///         .wrap(Condition::new(enable_normalize, NormalizePath));
+/// }
+/// ```
+pub struct Condition<T> {
+    trans: T,
+    enable: bool,
+}
+
+impl<T> Condition<T> {
+    pub fn new(enable: bool, trans: T) -> Self {
+        Self { trans, enable }
+    }
+}
+
+impl<S, T> Transform<S> for Condition<T>
+where
+    S: Service,
+    T: Transform<S, Request = S::Request, Response = S::Response, Error = S::Error>,
+{
+    type Request = S::Request;
+    type Response = S::Response;
+    type Error = S::Error;
+    type InitError = T::InitError;
+    type Transform = ConditionMiddleware<T::Transform, S>;
+    type Future = Either<
+        Map<T::Future, fn(T::Transform) -> Self::Transform>,
+        FutureResult<Self::Transform, Self::InitError>,
+    >;
+
+    fn new_transform(&self, service: S) -> Self::Future {
+        if self.enable {
+            let f = self
+                .trans
+                .new_transform(service)
+                .map(ConditionMiddleware::Enable as fn(T::Transform) -> Self::Transform);
+            Either::A(f)
+        } else {
+            Either::B(ok(ConditionMiddleware::Disable(service)))
+        }
+    }
+}
+
+pub enum ConditionMiddleware<E, D> {
+    Enable(E),
+    Disable(D),
+}
+
+impl<E, D> Service for ConditionMiddleware<E, D>
+where
+    E: Service,
+    D: Service<Request = E::Request, Response = E::Response, Error = E::Error>,
+{
+    type Request = E::Request;
+    type Response = E::Response;
+    type Error = E::Error;
+    type Future = Either<E::Future, D::Future>;
+
+    fn poll_ready(&mut self) -> Poll<(), Self::Error> {
+        use ConditionMiddleware::*;
+        match self {
+            Enable(service) => service.poll_ready(),
+            Disable(service) => service.poll_ready(),
+        }
+    }
+
+    fn call(&mut self, req: E::Request) -> Self::Future {
+        use ConditionMiddleware::*;
+        match self {
+            Enable(service) => Either::A(service.call(req)),
+            Disable(service) => Either::B(service.call(req)),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use actix_service::IntoService;
+
+    use super::*;
+    use crate::dev::{ServiceRequest, ServiceResponse};
+    use crate::error::Result;
+    use crate::http::{header::CONTENT_TYPE, HeaderValue, StatusCode};
+    use crate::middleware::errhandlers::*;
+    use crate::test::{self, TestRequest};
+    use crate::HttpResponse;
+
+    fn render_500<B>(mut res: ServiceResponse<B>) -> Result<ErrorHandlerResponse<B>> {
+        res.response_mut()
+            .headers_mut()
+            .insert(CONTENT_TYPE, HeaderValue::from_static("0001"));
+        Ok(ErrorHandlerResponse::Response(res))
+    }
+
+    #[test]
+    fn test_handler_enabled() {
+        let srv = |req: ServiceRequest| {
+            req.into_response(HttpResponse::InternalServerError().finish())
+        };
+
+        let mw =
+            ErrorHandlers::new().handler(StatusCode::INTERNAL_SERVER_ERROR, render_500);
+
+        let mut mw =
+            test::block_on(Condition::new(true, mw).new_transform(srv.into_service()))
+                .unwrap();
+        let resp = test::call_service(&mut mw, TestRequest::default().to_srv_request());
+        assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "0001");
+    }
+    #[test]
+    fn test_handler_disabled() {
+        let srv = |req: ServiceRequest| {
+            req.into_response(HttpResponse::InternalServerError().finish())
+        };
+
+        let mw =
+            ErrorHandlers::new().handler(StatusCode::INTERNAL_SERVER_ERROR, render_500);
+
+        let mut mw =
+            test::block_on(Condition::new(false, mw).new_transform(srv.into_service()))
+                .unwrap();
+
+        let resp = test::call_service(&mut mw, TestRequest::default().to_srv_request());
+        assert_eq!(resp.headers().get(CONTENT_TYPE), None);
+    }
+}
diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs
index 814993f0c..311d0ee99 100644
--- a/src/middleware/mod.rs
+++ b/src/middleware/mod.rs
@@ -6,7 +6,9 @@ mod defaultheaders;
 pub mod errhandlers;
 mod logger;
 mod normalize;
+mod condition;
 
 pub use self::defaultheaders::DefaultHeaders;
 pub use self::logger::Logger;
 pub use self::normalize::NormalizePath;
+pub use self::condition::Condition;