From f7bd6eeedcd0705b31c2bca71fe92092179b9a4f Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Thu, 7 Jun 2018 19:46:38 -0700
Subject: [PATCH] add application filters

---
 src/application.rs | 66 +++++++++++++++++++++++++++++++++++++++++++---
 src/httprequest.rs |  6 +++++
 2 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/src/application.rs b/src/application.rs
index 90c70bd1..c9171008 100644
--- a/src/application.rs
+++ b/src/application.rs
@@ -22,6 +22,7 @@ pub struct HttpApplication<S = ()> {
     prefix_len: usize,
     router: Router,
     inner: Rc<UnsafeCell<Inner<S>>>,
+    filters: Option<Vec<Box<Predicate<S>>>>,
     middlewares: Rc<RefCell<Vec<Box<Middleware<S>>>>>,
 }
 
@@ -143,11 +144,21 @@ impl<S: 'static> HttpHandler for HttpApplication<S> {
                     || path.split_at(self.prefix_len).1.starts_with('/'))
         };
         if m {
-            let mut req = req.with_state(Rc::clone(&self.state), self.router.clone());
-            let tp = self.get_handler(&mut req);
+            let mut req2 =
+                req.clone_with_state(Rc::clone(&self.state), self.router.clone());
+
+            if let Some(ref filters) = self.filters {
+                for filter in filters {
+                    if !filter.check(&mut req2) {
+                        return Err(req);
+                    }
+                }
+            }
+
+            let tp = self.get_handler(&mut req2);
             let inner = Rc::clone(&self.inner);
             Ok(Box::new(Pipeline::new(
-                req,
+                req2,
                 Rc::clone(&self.middlewares),
                 inner,
                 tp,
@@ -168,6 +179,7 @@ struct ApplicationParts<S> {
     external: HashMap<String, Resource>,
     encoding: ContentEncoding,
     middlewares: Vec<Box<Middleware<S>>>,
+    filters: Vec<Box<Predicate<S>>>,
 }
 
 /// Structure that follows the builder pattern for building application
@@ -190,6 +202,7 @@ impl App<()> {
                 handlers: Vec::new(),
                 external: HashMap::new(),
                 encoding: ContentEncoding::Auto,
+                filters: Vec::new(),
                 middlewares: Vec::new(),
             }),
         }
@@ -229,6 +242,7 @@ where
                 handlers: Vec::new(),
                 external: HashMap::new(),
                 middlewares: Vec::new(),
+                filters: Vec::new(),
                 encoding: ContentEncoding::Auto,
             }),
         }
@@ -285,6 +299,26 @@ where
         self
     }
 
+    /// Add match predicate to application.
+    ///
+    /// ```rust
+    /// # extern crate actix_web;
+    /// # use actix_web::*;
+    /// # fn main() {
+    /// App::new()
+    ///     .filter(pred::Get())
+    ///     .resource("/path", |r| r.f(|_| HttpResponse::Ok()))
+    /// #      .finish();
+    /// # }
+    /// ```
+    pub fn filter<T: Predicate<S> + 'static>(mut self, p: T) -> App<S> {
+        {
+            let parts = self.parts.as_mut().expect("Use after finish");
+            parts.filters.push(Box::new(p));
+        }
+        self
+    }
+
     /// Configure route for a specific path.
     ///
     /// This is a simplified version of the `App::resource()` method.
@@ -608,6 +642,11 @@ where
             handlers: parts.handlers,
             resources,
         }));
+        let filters = if parts.filters.is_empty() {
+            None
+        } else {
+            Some(parts.filters)
+        };
 
         HttpApplication {
             state: Rc::new(parts.state),
@@ -616,6 +655,7 @@ where
             prefix,
             prefix_len,
             inner,
+            filters,
         }
     }
 
@@ -700,7 +740,8 @@ mod tests {
     use http::StatusCode;
     use httprequest::HttpRequest;
     use httpresponse::HttpResponse;
-    use test::TestRequest;
+    use pred;
+    use test::{TestRequest, TestServer};
 
     #[test]
     fn test_default_resource() {
@@ -899,4 +940,21 @@ mod tests {
         let resp = app.run(req);
         assert_eq!(resp.as_msg().status(), StatusCode::NOT_FOUND);
     }
+
+    #[test]
+    fn test_filter() {
+        let mut srv = TestServer::with_factory(|| {
+            App::new()
+                .filter(pred::Get())
+                .handler("/test", |_| HttpResponse::Ok())
+        });
+
+        let request = srv.get().uri(srv.url("/test")).finish().unwrap();
+        let response = srv.execute(request.send()).unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+
+        let request = srv.post().uri(srv.url("/test")).finish().unwrap();
+        let response = srv.execute(request.send()).unwrap();
+        assert_eq!(response.status(), StatusCode::NOT_FOUND);
+    }
 }
diff --git a/src/httprequest.rs b/src/httprequest.rs
index d852bc74..a54a9958 100644
--- a/src/httprequest.rs
+++ b/src/httprequest.rs
@@ -141,6 +141,12 @@ impl HttpRequest<()> {
     pub fn with_state<S>(self, state: Rc<S>, router: Router) -> HttpRequest<S> {
         HttpRequest(self.0, Some(state), Some(router))
     }
+
+    pub(crate) fn clone_with_state<S>(
+        &self, state: Rc<S>, router: Router,
+    ) -> HttpRequest<S> {
+        HttpRequest(self.0.clone(), Some(state), Some(router))
+    }
 }
 
 impl<S> HttpMessage for HttpRequest<S> {