From fff207d3cb41c5a76db234196f16d28952dfd91c Mon Sep 17 00:00:00 2001
From: Fuwn <contact@fuwn.me>
Date: Thu, 6 Apr 2023 08:01:52 +0000
Subject: feat(response): async error handler

---
 src/handler/response/error.rs | 19 ++++++++++++++++---
 src/handler/response/route.rs |  4 +---
 src/router.rs                 | 40 ++++++++++++++++++++++++++--------------
 3 files changed, 43 insertions(+), 20 deletions(-)

(limited to 'src')

diff --git a/src/handler/response/error.rs b/src/handler/response/error.rs
index c433a46..a1a7885 100644
--- a/src/handler/response/error.rs
+++ b/src/handler/response/error.rs
@@ -16,10 +16,23 @@
 // Copyright (C) 2022-2022 Fuwn <contact@fuwn.me>
 // SPDX-License-Identifier: GPL-3.0-only
 
+use async_trait::async_trait;
+
 use crate::{context::ErrorContext, Response};
 
 #[allow(clippy::module_name_repetitions)]
-pub trait ErrorResponse: FnMut(ErrorContext) -> Response + Send + Sync {}
+#[async_trait]
+pub trait ErrorResponse: Send + Sync {
+  async fn call(&mut self, context: ErrorContext) -> Response;
+}
 
-impl<T> ErrorResponse for T where T: FnMut(ErrorContext) -> Response + Send + Sync
-{}
+#[async_trait]
+impl<T, F> ErrorResponse for T
+where
+  T: FnMut(ErrorContext) -> F + Send + Sync,
+  F: std::future::Future<Output = Response> + Send + 'static,
+{
+  async fn call(&mut self, context: ErrorContext) -> Response {
+    (*self)(context).await
+  }
+}
diff --git a/src/handler/response/route.rs b/src/handler/response/route.rs
index 0196363..a65b889 100644
--- a/src/handler/response/route.rs
+++ b/src/handler/response/route.rs
@@ -16,8 +16,6 @@
 // Copyright (C) 2022-2022 Fuwn <contact@fuwn.me>
 // SPDX-License-Identifier: GPL-3.0-only
 
-use std::future::Future;
-
 use async_trait::async_trait;
 
 use crate::{context::RouteContext, Response};
@@ -32,7 +30,7 @@ pub trait RouteResponse: Send + Sync {
 impl<T, F> RouteResponse for T
 where
   T: FnMut(RouteContext<'_>) -> F + Send + Sync,
-  F: Future<Output = Response> + Send + 'static,
+  F: std::future::Future<Output = Response> + Send + 'static,
 {
   async fn call(&mut self, context: RouteContext<'_>) -> Response {
     (*self)(context).await
diff --git a/src/router.rs b/src/router.rs
index c865fa9..030ffb8 100644
--- a/src/router.rs
+++ b/src/router.rs
@@ -75,7 +75,7 @@ macro_rules! or_error {
 #[derive(Clone)]
 pub struct Router {
   routes: matchit::Router<Arc<AsyncMutex<Box<dyn RouteResponse>>>>,
-  error_handler:         Arc<Mutex<Box<dyn ErrorResponse<Output = Response>>>>,
+  error_handler:         Arc<AsyncMutex<Box<dyn ErrorResponse>>>,
   private_key_file_name: String,
   ca_file_name:          String,
   headers:               Arc<Mutex<Vec<Box<dyn Partial<Output = String>>>>>,
@@ -190,11 +190,17 @@ impl Router {
   ///   windmark::success!("You have encountered an error!")
   /// });
   /// ```
-  pub fn set_error_handler(
+  pub fn set_error_handler<R>(
     &mut self,
-    handler: impl ErrorResponse + 'static,
-  ) -> &mut Self {
-    self.error_handler = Arc::new(Mutex::new(Box::new(handler)));
+    mut handler: impl FnMut(ErrorContext) -> R + Send + Sync + 'static,
+  ) -> &mut Self
+  where
+    R: IntoFuture<Output = Response> + Send + 'static,
+    <R as IntoFuture>::IntoFuture: Send,
+  {
+    self.error_handler = Arc::new(AsyncMutex::new(Box::new(move |context| {
+      handler(context).into_future()
+    })));
 
     self
   }
@@ -398,11 +404,15 @@ impl Router {
 
       handler.await
     } else {
-      (*self.error_handler).lock().unwrap()(ErrorContext::new(
-        stream.get_ref().peer_addr(),
-        url.clone(),
-        peer_certificate,
-      ))
+      (*self.error_handler)
+        .lock()
+        .await
+        .call(ErrorContext::new(
+          stream.get_ref().peer_addr(),
+          url.clone(),
+          peer_certificate,
+        ))
+        .await
     };
 
     for module in &mut *self.async_modules.lock().await {
@@ -855,10 +865,12 @@ impl Default for Router {
   fn default() -> Self {
     Self {
       routes: matchit::Router::new(),
-      error_handler: Arc::new(Mutex::new(Box::new(|_| {
-        Response::not_found(
-          "This capsule has not implemented an error handler...",
-        )
+      error_handler: Arc::new(AsyncMutex::new(Box::new(|_| {
+        async {
+          Response::not_found(
+            "This capsule has not implemented an error handler...",
+          )
+        }
       }))),
       private_key_file_name: String::new(),
       ca_file_name: String::new(),
-- 
cgit v1.2.3