feat(router): modify content from post-route callback

3 files changed, 50 insertions, 41 deletions

diff --git a/examples/windmark.rs b/examples/windmark.rs
index fca1c97..b33f686 100644
--- a/examples/windmark.rs
+++ b/examples/windmark.rs
@@ -82,7 +82,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
       context.url.to_string()
     )
   }));
-  router.set_post_route_callback(Box::new(|context| {
+  router.set_post_route_callback(Box::new(|context, content| {
+    *content = content.replace("Welcome!", "Welcome to Windmark!");
+
     info!(
       "closed connection from {}",
       context.tcp.peer_addr().unwrap().ip()
diff --git a/src/handler.rs b/src/handler.rs
index 15dfd6d..c5e23e8 100644
--- a/src/handler.rs
+++ b/src/handler.rs
@@ -16,12 +16,17 @@
 // Copyright (C) 2022-2022 Fuwn <[email protected]>
 // SPDX-License-Identifier: GPL-3.0-only
 
-use crate::{returnable, returnable::RouteContext, Response};
+use crate::{
+  returnable,
+  returnable::{CallbackContext, RouteContext},
+  Response,
+};
 
 pub type RouteResponse =
   Box<dyn FnMut(RouteContext<'_>) -> Response<'_> + Send + Sync>;
 pub type ErrorResponse =
   Box<dyn FnMut(returnable::ErrorContext<'_>) -> Response<'_> + Send + Sync>;
-pub type Callback =
-  Box<dyn FnMut(returnable::CallbackContext<'_>) + Send + Sync>;
+pub type Callback = Box<dyn FnMut(CallbackContext<'_>) + Send + Sync>;
+pub type CleanupCallback =
+  Box<dyn FnMut(CallbackContext<'_>, &mut String) + Send + Sync>;
 pub type Partial = Box<dyn FnMut(RouteContext<'_>) -> String + Send + Sync>;
diff --git a/src/router.rs b/src/router.rs
index dc81d40..a28eb67 100644
--- a/src/router.rs
+++ b/src/router.rs
@@ -27,7 +27,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use url::Url;
 
 use crate::{
-  handler::{Callback, ErrorResponse, Partial, RouteResponse},
+  handler::{Callback, CleanupCallback, ErrorResponse, Partial, RouteResponse},
   module::Module,
   response::{to_value_set_status, Response},
   returnable::{CallbackContext, ErrorContext, RouteContext},
@@ -64,7 +64,7 @@ pub struct Router {
   #[cfg(feature = "logger")]
   default_logger:        bool,
   pre_route_callback:    Arc<Mutex<Callback>>,
-  post_route_callback:   Arc<Mutex<Callback>>,
+  post_route_callback:   Arc<Mutex<CleanupCallback>>,
   charset:               String,
   language:              String,
   port:                  i32,
@@ -323,7 +323,7 @@ impl Router {
       &stream.ssl().peer_certificate(),
     ));
 
-    let content = if let Ok(ref route) = route {
+    let mut content = if let Ok(ref route) = route {
       let footers_length = (*self.footers.lock().unwrap()).len();
 
       for partial_header in &mut *self.headers.lock().unwrap() {
@@ -378,6 +378,25 @@ impl Router {
       )
     };
 
+    for module in &mut *self.modules.lock().unwrap() {
+      module.on_post_route(CallbackContext::new(
+        stream.get_ref(),
+        &url,
+        route.as_ref().map_or(None, |route| Some(&route.params)),
+        &stream.ssl().peer_certificate(),
+      ));
+    }
+
+    (*self.post_route_callback).lock().unwrap()(
+      CallbackContext::new(
+        stream.get_ref(),
+        &url,
+        route.as_ref().map_or(None, |route| Some(&route.params)),
+        &stream.ssl().peer_certificate(),
+      ),
+      &mut content,
+    );
+
     stream
       .write_all(
         format!(
@@ -412,22 +431,6 @@ impl Router {
       )
       .await?;
 
-    for module in &mut *self.modules.lock().unwrap() {
-      module.on_post_route(CallbackContext::new(
-        stream.get_ref(),
-        &url,
-        route.as_ref().map_or(None, |route| Some(&route.params)),
-        &stream.ssl().peer_certificate(),
-      ));
-    }
-
-    (*self.post_route_callback).lock().unwrap()(CallbackContext::new(
-      stream.get_ref(),
-      &url,
-      route.as_ref().map_or(None, |route| Some(&route.params)),
-      &stream.ssl().peer_certificate(),
-    ));
-
     stream.shutdown().await?;
 
     Ok(())
@@ -539,14 +542,12 @@ impl Router {
   /// ```rust
   /// use log::info;
   ///
-  /// windmark::Router::new().set_pre_route_callback(Box::new(
-  ///   |stream, _url, _| {
-  ///     info!(
-  ///       "accepted connection from {}",
-  ///       stream.peer_addr().unwrap().ip(),
-  ///     )
-  ///   },
-  /// ));
+  /// windmark::Router::new().set_pre_route_callback(Box::new(|context| {
+  ///   info!(
+  ///     "accepted connection from {}",
+  ///     context.stream.peer_addr().unwrap().ip(),
+  ///   )
+  /// }));
   /// ```
   pub fn set_pre_route_callback(&mut self, callback: Callback) -> &mut Self {
     self.pre_route_callback = Arc::new(Mutex::new(callback));
@@ -561,16 +562,17 @@ impl Router {
   /// ```rust
   /// use log::info;
   ///
-  /// windmark::Router::new().set_post_route_callback(Box::new(
-  ///   |stream, _url, _| {
-  ///     info!(
-  ///       "closed connection from {}",
-  ///       stream.peer_addr().unwrap().ip(),
-  ///     )
-  ///   },
-  /// ));
+  /// windmark::Router::new().set_post_route_callback(Box::new(|context, _| {
+  ///   info!(
+  ///     "closed connection from {}",
+  ///     context.stream.peer_addr().unwrap().ip(),
+  ///   )
+  /// }));
   /// ```
-  pub fn set_post_route_callback(&mut self, callback: Callback) -> &mut Self {
+  pub fn set_post_route_callback(
+    &mut self,
+    callback: CleanupCallback,
+  ) -> &mut Self {
     self.post_route_callback = Arc::new(Mutex::new(callback));
 
     self
@@ -763,7 +765,7 @@ impl Default for Router {
       #[cfg(feature = "logger")]
       default_logger: false,
       pre_route_callback: Arc::new(Mutex::new(Box::new(|_| {}))),
-      post_route_callback: Arc::new(Mutex::new(Box::new(|_| {}))),
+      post_route_callback: Arc::new(Mutex::new(Box::new(|_, _| {}))),
       charset: "utf-8".to_string(),
       language: "en".to_string(),
       port: 1965,