From 939d2e745c9ae3e5ba1f240f0613efa8db66390c Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Mon, 25 Mar 2019 12:47:58 -0700
Subject: [PATCH] rename Resource::middleware to Resource::wrap and add wrap_fn
 for fn middlewares

---
 examples/basic.rs |   2 +-
 src/lib.rs        |  26 ++++++++++
 src/resource.rs   | 125 ++++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 149 insertions(+), 4 deletions(-)

diff --git a/examples/basic.rs b/examples/basic.rs
index 756f1b796..ee7e4c967 100644
--- a/examples/basic.rs
+++ b/examples/basic.rs
@@ -33,7 +33,7 @@ fn main() -> std::io::Result<()> {
             .service(no_params)
             .service(
                 web::resource("/resource2/index.html")
-                    .middleware(
+                    .wrap(
                         middleware::DefaultHeaders::new().header("X-Version-R2", "0.3"),
                     )
                     .default_resource(|r| {
diff --git a/src/lib.rs b/src/lib.rs
index 79e1ba34a..8ae7156c3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -340,4 +340,30 @@ pub mod web {
     {
         blocking::run(f).from_err()
     }
+
+    use actix_service::{fn_transform, Service, Transform};
+
+    use crate::service::{ServiceRequest, ServiceResponse};
+
+    /// Create middleare
+    pub fn md<F, R, S, P, B>(
+        f: F,
+    ) -> impl Transform<
+        S,
+        Request = ServiceRequest<P>,
+        Response = ServiceResponse<B>,
+        Error = Error,
+        InitError = (),
+    >
+    where
+        S: Service<
+            Request = ServiceRequest<P>,
+            Response = ServiceResponse<B>,
+            Error = Error,
+        >,
+        F: FnMut(ServiceRequest<P>, &mut S) -> R + Clone,
+        R: IntoFuture<Item = ServiceResponse<B>, Error = Error>,
+    {
+        fn_transform(f)
+    }
 }
diff --git a/src/resource.rs b/src/resource.rs
index 29ff07857..5d5671310 100644
--- a/src/resource.rs
+++ b/src/resource.rs
@@ -230,7 +230,9 @@ where
     /// This is similar to `App's` middlewares, but middleware get invoked on resource level.
     /// Resource level middlewares are not allowed to change response
     /// type (i.e modify response's body).
-    pub fn middleware<M, F>(
+    ///
+    /// **Note**: middlewares get called in opposite order of middlewares registration.
+    pub fn wrap<M, F>(
         self,
         mw: F,
     ) -> Resource<
@@ -264,6 +266,57 @@ where
         }
     }
 
+    /// Register a resource middleware function.
+    ///
+    /// This function accepts instance of `ServiceRequest` type and
+    /// mutable reference to the next middleware in chain.
+    ///
+    /// This is similar to `App's` middlewares, but middleware get invoked on resource level.
+    /// Resource level middlewares are not allowed to change response
+    /// type (i.e modify response's body).
+    ///
+    /// ```rust
+    /// use actix_service::Service;
+    /// # use futures::Future;
+    /// use actix_web::{web, App};
+    /// use actix_web::http::{header::CONTENT_TYPE, HeaderValue};
+    ///
+    /// fn index() -> &'static str {
+    ///     "Welcome!"
+    /// }
+    ///
+    /// fn main() {
+    ///     let app = App::new().service(
+    ///         web::resource("/index.html")
+    ///             .wrap_fn(|req, srv|
+    ///                 srv.call(req).map(|mut res| {
+    ///                     res.headers_mut().insert(
+    ///                        CONTENT_TYPE, HeaderValue::from_static("text/plain"),
+    ///                     );
+    ///                     res
+    ///                 }))
+    ///             .route(web::get().to(index)));
+    /// }
+    /// ```
+    pub fn wrap_fn<F, R>(
+        self,
+        mw: F,
+    ) -> Resource<
+        P,
+        impl NewService<
+            Request = ServiceRequest<P>,
+            Response = ServiceResponse,
+            Error = Error,
+            InitError = (),
+        >,
+    >
+    where
+        F: FnMut(ServiceRequest<P>, &mut T::Service) -> R + Clone,
+        R: IntoFuture<Item = ServiceResponse, Error = Error>,
+    {
+        self.wrap(mw)
+    }
+
     /// Default resource to be used if no matching route could be found.
     /// By default *405* response get returned. Resource does not use
     /// default handler from `App` or `Scope`.
@@ -489,9 +542,75 @@ impl<P: 'static> NewService for ResourceEndpoint<P> {
 
 #[cfg(test)]
 mod tests {
-    use crate::http::{Method, StatusCode};
+    use actix_service::Service;
+    use futures::{Future, IntoFuture};
+
+    use crate::http::{header, HeaderValue, Method, StatusCode};
+    use crate::service::{ServiceRequest, ServiceResponse};
     use crate::test::{call_success, init_service, TestRequest};
-    use crate::{web, App, HttpResponse};
+    use crate::{web, App, Error, HttpResponse};
+
+    fn md1<S, P, B>(
+        req: ServiceRequest<P>,
+        srv: &mut S,
+    ) -> impl IntoFuture<Item = ServiceResponse<B>, Error = Error>
+    where
+        S: Service<
+            Request = ServiceRequest<P>,
+            Response = ServiceResponse<B>,
+            Error = Error,
+        >,
+    {
+        srv.call(req).map(|mut res| {
+            res.headers_mut()
+                .insert(header::CONTENT_TYPE, HeaderValue::from_static("0001"));
+            res
+        })
+    }
+
+    #[test]
+    fn test_middleware() {
+        let mut srv = init_service(
+            App::new().service(
+                web::resource("/test")
+                    .wrap(md1)
+                    .route(web::get().to(|| HttpResponse::Ok())),
+            ),
+        );
+        let req = TestRequest::with_uri("/test").to_request();
+        let resp = call_success(&mut srv, req);
+        assert_eq!(resp.status(), StatusCode::OK);
+        assert_eq!(
+            resp.headers().get(header::CONTENT_TYPE).unwrap(),
+            HeaderValue::from_static("0001")
+        );
+    }
+
+    #[test]
+    fn test_middleware_fn() {
+        let mut srv = init_service(
+            App::new().service(
+                web::resource("/test")
+                    .wrap_fn(|req, srv| {
+                        srv.call(req).map(|mut res| {
+                            res.headers_mut().insert(
+                                header::CONTENT_TYPE,
+                                HeaderValue::from_static("0001"),
+                            );
+                            res
+                        })
+                    })
+                    .route(web::get().to(|| HttpResponse::Ok())),
+            ),
+        );
+        let req = TestRequest::with_uri("/test").to_request();
+        let resp = call_success(&mut srv, req);
+        assert_eq!(resp.status(), StatusCode::OK);
+        assert_eq!(
+            resp.headers().get(header::CONTENT_TYPE).unwrap(),
+            HeaderValue::from_static("0001")
+        );
+    }
 
     #[test]
     fn test_default_resource() {