From e7f05fc3b45e3e068e8a8373782e7f519f40690d Mon Sep 17 00:00:00 2001
From: Fuwn <contact@fuwn.me>
Date: Wed, 29 Mar 2023 08:28:56 +0000
Subject: feat(router): modify content from post-route callback

---
 src/handler.rs | 11 ++++++---
 src/router.rs  | 76 ++++++++++++++++++++++++++++++----------------------------
 2 files changed, 47 insertions(+), 40 deletions(-)

(limited to 'src')

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 <contact@fuwn.me>
 // 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,
-- 
cgit v1.2.3