Per-route middleware

Author: tirr-c

src/endpoint.rs

@@ -1,7 +1,10 @@
+use std::sync::Arc;
+
 use async_std::future::Future;
 
+use crate::middleware::Next;
 use crate::utils::BoxFuture;
-use crate::{response::IntoResponse, Request, Response};
+use crate::{response::IntoResponse, Middleware, Request, Response};
 
 /// An HTTP request handler.
 ///
@@ -63,3 +66,52 @@ where
         Box::pin(async move { fut.await.into_response() })
     }
 }
+
+pub struct MiddlewareEndpoint<E, State> {
+    endpoint: E,
+    middleware: Vec<Arc<dyn Middleware<State>>>,
+}
+
+impl<E: Clone, State> Clone for MiddlewareEndpoint<E, State> {
+    fn clone(&self) -> Self {
+        Self {
+            endpoint: self.endpoint.clone(),
+            middleware: self.middleware.clone(),
+        }
+    }
+}
+
+impl<E, State> std::fmt::Debug for MiddlewareEndpoint<E, State> {
+    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            fmt,
+            "MiddlewareEndpoint (length: {})",
+            self.middleware.len(),
+        )
+    }
+}
+
+impl<E, State> MiddlewareEndpoint<E, State>
+where
+    E: Endpoint<State>,
+{
+    pub fn wrap_with_middleware(ep: E, middleware: &[Arc<dyn Middleware<State>>]) -> Self {
+        Self {
+            endpoint: ep,
+            middleware: middleware.to_vec(),
+        }
+    }
+}
+
+impl<E, State: 'static> Endpoint<State> for MiddlewareEndpoint<E, State>
+where
+    E: Endpoint<State>,
+{
+    fn call<'a>(&'a self, req: Request<State>) -> BoxFuture<'a, Response> {
+        let next = Next {
+            endpoint: &self.endpoint,
+            next_middleware: &self.middleware,
+        };
+        next.run(req)
+    }
+}

src/router.rs

@@ -1,7 +1,7 @@
 use route_recognizer::{Match, Params, Router as MethodRouter};
 use std::collections::HashMap;
 
-use crate::endpoint::{DynEndpoint, Endpoint};
+use crate::endpoint::DynEndpoint;
 use crate::utils::BoxFuture;
 use crate::{Request, Response};
 
@@ -29,15 +29,15 @@ impl<State: 'static> Router<State> {
         }
     }
 
-    pub(crate) fn add(&mut self, path: &str, method: http::Method, ep: impl Endpoint<State>) {
+    pub(crate) fn add(&mut self, path: &str, method: http::Method, ep: Box<DynEndpoint<State>>) {
         self.method_map
             .entry(method)
             .or_insert_with(MethodRouter::new)
-            .add(path, Box::new(ep))
+            .add(path, ep)
     }
 
-    pub(crate) fn add_all(&mut self, path: &str, ep: impl Endpoint<State>) {
-        self.all_method_router.add(path, Box::new(ep))
+    pub(crate) fn add_all(&mut self, path: &str, ep: Box<DynEndpoint<State>>) {
+        self.all_method_router.add(path, ep)
     }
 
     pub(crate) fn route(&self, path: &str, method: http::Method) -> Selection<'_, State> {

src/server/route.rs

@@ -1,5 +1,8 @@
+use std::sync::Arc;
+
+use crate::endpoint::MiddlewareEndpoint;
 use crate::utils::BoxFuture;
-use crate::{router::Router, Endpoint, Response};
+use crate::{router::Router, Endpoint, Middleware, Response};
 
 /// A handle to a route.
 ///
@@ -13,6 +16,7 @@ use crate::{router::Router, Endpoint, Response};
 pub struct Route<'a, State> {
     router: &'a mut Router<State>,
     path: String,
+    middleware: Vec<Arc<dyn Middleware<State>>>,
     /// Indicates whether the path of current route is treated as a prefix. Set by
     /// [`strip_prefix`].
     ///
@@ -25,11 +29,14 @@ impl<'a, State: 'static> Route<'a, State> {
         Route {
             router,
             path,
+            middleware: Vec::new(),
             prefix: false,
         }
     }
 
     /// Extend the route with the given `path`.
+    ///
+    /// The returned route won't have any middleware applied.
     pub fn at<'b>(&'b mut self, path: &str) -> Route<'b, State> {
         let mut p = self.path.clone();
 
@@ -44,6 +51,7 @@ impl<'a, State: 'static> Route<'a, State> {
         Route {
             router: &mut self.router,
             path: p,
+            middleware: Vec::new(),
             prefix: false,
         }
     }
@@ -60,6 +68,18 @@ impl<'a, State: 'static> Route<'a, State> {
         self
     }
 
+    /// Apply the given middleware to the current route.
+    pub fn middleware(&mut self, middleware: impl Middleware<State>) -> &mut Self {
+        self.middleware.push(Arc::new(middleware));
+        self
+    }
+
+    /// Reset the middleware chain for the current route, if any.
+    pub fn reset_middleware(&mut self) -> &mut Self {
+        self.middleware.clear();
+        self
+    }
+
     /// Nest a [`Server`] at the current path.
     ///
     /// [`Server`]: struct.Server.html
@@ -78,10 +98,29 @@ impl<'a, State: 'static> Route<'a, State> {
     pub fn method(&mut self, method: http::Method, ep: impl Endpoint<State>) -> &mut Self {
         if self.prefix {
             let ep = StripPrefixEndpoint::new(ep);
-            self.router.add(&self.path, method.clone(), ep.clone());
+            let (ep1, ep2): (Box<dyn Endpoint<_>>, Box<dyn Endpoint<_>>) =
+                if self.middleware.is_empty() {
+                    let ep = Box::new(ep);
+                    (ep.clone(), ep)
+                } else {
+                    let ep = Box::new(MiddlewareEndpoint::wrap_with_middleware(
+                        ep,
+                        &self.middleware,
+                    ));
+                    (ep.clone(), ep)
+                };
+            self.router.add(&self.path, method.clone(), ep1);
             let wildcard = self.at("*--tide-path-rest");
-            wildcard.router.add(&wildcard.path, method, ep);
+            wildcard.router.add(&wildcard.path, method, ep2);
         } else {
+            let ep: Box<dyn Endpoint<_>> = if self.middleware.is_empty() {
+                Box::new(ep)
+            } else {
+                Box::new(MiddlewareEndpoint::wrap_with_middleware(
+                    ep,
+                    &self.middleware,
+                ))
+            };
             self.router.add(&self.path, method, ep);
         }
         self
@@ -93,10 +132,29 @@ impl<'a, State: 'static> Route<'a, State> {
     pub fn all(&mut self, ep: impl Endpoint<State>) -> &mut Self {
         if self.prefix {
             let ep = StripPrefixEndpoint::new(ep);
-            self.router.add_all(&self.path, ep.clone());
+            let (ep1, ep2): (Box<dyn Endpoint<_>>, Box<dyn Endpoint<_>>) =
+                if self.middleware.is_empty() {
+                    let ep = Box::new(ep);
+                    (ep.clone(), ep)
+                } else {
+                    let ep = Box::new(MiddlewareEndpoint::wrap_with_middleware(
+                        ep,
+                        &self.middleware,
+                    ));
+                    (ep.clone(), ep)
+                };
+            self.router.add_all(&self.path, ep1);
             let wildcard = self.at("*--tide-path-rest");
-            wildcard.router.add_all(&wildcard.path, ep);
+            wildcard.router.add_all(&wildcard.path, ep2);
         } else {
+            let ep: Box<dyn Endpoint<_>> = if self.middleware.is_empty() {
+                Box::new(ep)
+            } else {
+                Box::new(MiddlewareEndpoint::wrap_with_middleware(
+                    ep,
+                    &self.middleware,
+                ))
+            };
             self.router.add_all(&self.path, ep);
         }
         self

tests/route_middleware.rs

@@ -0,0 +1,81 @@
+use async_std::io::prelude::*;
+use futures::executor::block_on;
+use futures::future::BoxFuture;
+use http_service::Body;
+use http_service_mock::make_server;
+use tide::Middleware;
+
+struct TestMiddleware(&'static str);
+
+impl<State: Send + Sync + 'static> Middleware<State> for TestMiddleware {
+    fn handle<'a>(
+        &'a self,
+        req: tide::Request<State>,
+        next: tide::Next<'a, State>,
+    ) -> BoxFuture<'a, tide::Response> {
+        Box::pin(async move {
+            let res = next.run(req).await;
+            res.set_header("X-Tide-Test", self.0)
+        })
+    }
+}
+
+async fn echo_path<State>(req: tide::Request<State>) -> String {
+    req.uri().path().to_string()
+}
+
+#[test]
+fn route_middleware() {
+    let mut app = tide::new();
+    let mut foo_route = app.at("/foo");
+    foo_route.middleware(TestMiddleware("foo"))
+        .get(echo_path);
+    foo_route.at("/bar")
+        .middleware(TestMiddleware("bar"))
+        .get(echo_path);
+    foo_route.post(echo_path)
+        .reset_middleware()
+        .put(echo_path);
+    let mut server = make_server(app.into_http_service()).unwrap();
+
+    let mut buf = Vec::new();
+    let req = http::Request::get("/foo").body(Body::empty()).unwrap();
+    let res = server.simulate(req).unwrap();
+    assert_eq!(
+        res.headers().get("X-Tide-Test"),
+        Some(&"foo".parse().unwrap())
+    );
+    assert_eq!(res.status(), 200);
+    block_on(res.into_body().read_to_end(&mut buf)).unwrap();
+    assert_eq!(&*buf, &*b"/foo");
+
+    buf.clear();
+    let req = http::Request::post("/foo").body(Body::empty()).unwrap();
+    let res = server.simulate(req).unwrap();
+    assert_eq!(
+        res.headers().get("X-Tide-Test"),
+        Some(&"foo".parse().unwrap())
+    );
+    assert_eq!(res.status(), 200);
+    block_on(res.into_body().read_to_end(&mut buf)).unwrap();
+    assert_eq!(&*buf, &*b"/foo");
+
+    buf.clear();
+    let req = http::Request::put("/foo").body(Body::empty()).unwrap();
+    let res = server.simulate(req).unwrap();
+    assert_eq!(res.headers().get("X-Tide-Test"), None);
+    assert_eq!(res.status(), 200);
+    block_on(res.into_body().read_to_end(&mut buf)).unwrap();
+    assert_eq!(&*buf, &*b"/foo");
+
+    buf.clear();
+    let req = http::Request::get("/foo/bar").body(Body::empty()).unwrap();
+    let res = server.simulate(req).unwrap();
+    assert_eq!(
+        res.headers().get("X-Tide-Test"),
+        Some(&"bar".parse().unwrap())
+    );
+    assert_eq!(res.status(), 200);
+    block_on(res.into_body().read_to_end(&mut buf)).unwrap();
+    assert_eq!(&*buf, &*b"/foo/bar");
+}