introduce the ideas of Before and After middlewares

jbr · jbr · commit 1cb9af2fcefc · 2020-05-28T16:12:44.000-07:00
diff --git a/examples/middleware.rs b/examples/middleware.rs
@@ -2,14 +2,14 @@ use std::future::Future;
 use std::pin::Pin;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
-use tide::{Middleware, Next, Request, Response, Result, StatusCode};
+use tide::{After, Before, Middleware, Next, Request, Response, Result, StatusCode};
 
 #[derive(Debug)]
 struct User {
     name: String,
 }
 
-#[derive(Default)]
+#[derive(Default, Debug)]
 struct UserDatabase;
 impl UserDatabase {
     async fn find_user(&self) -> Option<User> {
@@ -77,14 +77,37 @@ impl<State: Send + Sync + 'static> Middleware<State> for RequestCounterMiddlewar
         })
     }
 }
+const NOT_FOUND_HTML_PAGE: &str = "<html><body>
+  <h1>uh oh, we couldn't find that document</h1>
+  <p>
+    probably, this would be served from the file system or
+    included with `include_bytes!`
+  </p>
+</body></html>";
 
 #[async_std::main]
 async fn main() -> Result<()> {
     tide::log::start();
     let mut app = tide::with_state(UserDatabase::default());
 
+    app.middleware(After::new(|result: Result| async move {
+        let response = result.unwrap_or_else(|e| Response::new(e.status()));
+        match response.status() {
+            StatusCode::NotFound => Ok(Response::new(response.status())
+                .set_mime(tide::http::mime::HTML)
+                .body_string(NOT_FOUND_HTML_PAGE.into())),
+            _ => Ok(response),
+        }
+    }));
+
     app.middleware(user_loader);
     app.middleware(RequestCounterMiddleware::new(0));
+    app.middleware(Before::new(
+        |mut request: Request<UserDatabase>| async move {
+            request.set_ext(std::time::Instant::now());
+            request
+        },
+    ));
 
     app.at("/").get(|req: Request<_>| async move {
         let count: &RequestCount = req.ext().unwrap();
diff --git a/src/lib.rs b/src/lib.rs
@@ -211,7 +211,7 @@ pub mod security;
 pub mod sse;
 
 pub use endpoint::Endpoint;
-pub use middleware::{Middleware, Next};
+pub use middleware::{After, Before, Middleware, Next};
 pub use redirect::Redirect;
 pub use request::Request;
 pub use response::Response;
diff --git a/src/middleware.rs b/src/middleware.rs
@@ -27,6 +27,58 @@ pub trait Middleware<State>: 'static + Send + Sync {
     }
 }
 
+#[derive(Debug)]
+pub struct Before<F>(F);
+impl<F> Before<F> {
+    pub fn new(f: F) -> Self {
+        Self(f)
+    }
+}
+
+impl<State, F, Fut> Middleware<State> for Before<F>
+where
+    State: Send + Sync + 'static,
+    F: Fn(Request<State>) -> Fut + Send + Sync + 'static,
+    Fut: std::future::Future<Output = Request<State>> + Send + Sync,
+{
+    fn handle<'a>(
+        &'a self,
+        request: Request<State>,
+        next: Next<'a, State>,
+    ) -> BoxFuture<'a, crate::Result> {
+        Box::pin(async move {
+            let request = (self.0)(request).await;
+            next.run(request).await
+        })
+    }
+}
+
+#[derive(Debug)]
+pub struct After<F>(F);
+impl<F> After<F> {
+    pub fn new(f: F) -> Self {
+        Self(f)
+    }
+}
+
+impl<State, F, Fut> Middleware<State> for After<F>
+where
+    State: Send + Sync + 'static,
+    F: Fn(crate::Result) -> Fut + Send + Sync + 'static,
+    Fut: std::future::Future<Output = crate::Result> + Send + Sync,
+{
+    fn handle<'a>(
+        &'a self,
+        request: Request<State>,
+        next: Next<'a, State>,
+    ) -> BoxFuture<'a, crate::Result> {
+        Box::pin(async move {
+            let result = next.run(request).await;
+            (self.0)(result).await
+        })
+    }
+}
+
 impl<State, F> Middleware<State> for F
 where
     F: Send
