Update code to work properly with borrowed errors

the10thWiz · the10thWiz · commit 6d06ac7a63e1 · 2024-09-02T20:25:59.000-05:00
diff --git a/core/lib/src/catcher/types.rs b/core/lib/src/catcher/types.rs
@@ -93,7 +93,6 @@ impl<'r, L, R> TypedError<'r> for Either<L, R>
     fn name(&self) -> &'static str { std::any::type_name::<Self>() }
 
     fn source(&'r self) -> Option<&'r (dyn TypedError<'r> + 'r)> { 
-        println!("Downcasting either");
         match self {
             Self::Left(v) => Some(v),
             Self::Right(v) => Some(v),
diff --git a/core/lib/src/erased.rs b/core/lib/src/erased.rs
@@ -8,6 +8,7 @@ use futures::future::BoxFuture;
 use http::request::Parts;
 use tokio::io::{AsyncRead, ReadBuf};
 
+use crate::catcher::TypedError;
 use crate::data::{Data, IoHandler, RawStream};
 use crate::{Request, Response, Rocket, Orbit};
 
@@ -34,10 +35,12 @@ impl Drop for ErasedRequest {
     fn drop(&mut self) { }
 }
 
-#[derive(Debug)]
+// TODO: #[derive(Debug)]
 pub struct ErasedResponse {
     // XXX: SAFETY: This (dependent) field must come first due to drop order!
     response: Response<'static>,
+    // XXX: SAFETY: This (dependent) field must come first due to drop order!
+    error: Option<Box<dyn TypedError<'static> + 'static>>,
     _request: Arc<ErasedRequest>,
 }
 
@@ -94,7 +97,8 @@ impl ErasedRequest {
             T,
             &'r Rocket<Orbit>,
             &'r Request<'r>,
-            Data<'r>
+            Data<'r>,
+            &'r mut Option<Box<dyn TypedError<'r> + 'r>>,
         ) -> BoxFuture<'r, Response<'r>>,
     ) -> ErasedResponse
         where T: Send + Sync + 'static,
@@ -111,15 +115,19 @@ impl ErasedRequest {
         };
 
         let parent = parent;
+        let mut error_ptr: Option<Box<dyn TypedError<'static> + 'static>> = None;
         let response: Response<'_> = {
             let parent: &ErasedRequest = &parent;
             let parent: &'static ErasedRequest = unsafe { transmute(parent) };
             let rocket: &Rocket<Orbit> = &parent._rocket;
             let request: &Request<'_> = &parent.request;
-            dispatch(token, rocket, request, data).await
+            // SAFETY: error_ptr is transmuted into the same type, with the same lifetime as the request.
+            // It is kept alive by the erased response, so that the response type can borrow from it
+            dispatch(token, rocket, request, data, unsafe { transmute(&mut error_ptr)}).await
         };
 
         ErasedResponse {
+            error: error_ptr,
             _request: parent,
             response,
         }
diff --git a/core/lib/src/lifecycle.rs b/core/lib/src/lifecycle.rs
@@ -94,6 +94,7 @@ impl Rocket<Orbit> {
         _token: RequestToken,
         request: &'r Request<'s>,
         data: Data<'r>,
+        error_ptr: &'r mut Option<Box<dyn TypedError<'r> + 'r>>,
         // io_stream: impl Future<Output = io::Result<IoStream>> + Send,
     ) -> Response<'r> {
         // Remember if the request is `HEAD` for later body stripping.
@@ -109,16 +110,24 @@ impl Rocket<Orbit> {
                 request._set_method(Method::Get);
                 match self.route(request, data).await {
                     Outcome::Success(response) => response,
-                    Outcome::Error((status, error))
-                        => self.dispatch_error(status, request, error).await,
-                    Outcome::Forward((_, status, error))
-                        => self.dispatch_error(status, request, error).await,
+                    Outcome::Error((status, error)) => {
+                        *error_ptr = error;
+                        self.dispatch_error(status, request, error_ptr.as_ref().map(|b| b.as_ref())).await
+                    },
+                    Outcome::Forward((_, status, error)) => {
+                        *error_ptr = error;
+                        self.dispatch_error(status, request, error_ptr.as_ref().map(|b| b.as_ref())).await
+                    },
                 }
             }
-            Outcome::Forward((_, status, error))
-                => self.dispatch_error(status, request, error).await,
-            Outcome::Error((status, error))
-                => self.dispatch_error(status, request, error).await,
+            Outcome::Forward((_, status, error)) => {
+                *error_ptr = error;
+                self.dispatch_error(status, request, error_ptr.as_ref().map(|b| b.as_ref())).await
+            },
+            Outcome::Error((status, error)) => {
+                *error_ptr = error;
+                self.dispatch_error(status, request, error_ptr.as_ref().map(|b| b.as_ref())).await
+            },
         };
 
         // Set the cookies. Note that error responses will only include cookies
@@ -236,7 +245,7 @@ impl Rocket<Orbit> {
         &'s self,
         mut status: Status,
         req: &'r Request<'s>,
-        mut error: Option<Box<dyn TypedError<'r> + 'r>>,
+        mut error: Option<&'r dyn TypedError<'r>>,
     ) -> Response<'r> {
         // We may wish to relax this in the future.
         req.cookies().reset_delta();
@@ -273,48 +282,64 @@ impl Rocket<Orbit> {
     async fn invoke_catcher<'s, 'r: 's>(
         &'s self,
         status: Status,
-        error: Option<Box<dyn TypedError<'r> + 'r>>,
+        error: Option<&'r dyn TypedError<'r>>,
         req: &'r Request<'s>
     ) -> Result<Response<'r>, Option<Status>> {
-        let error_ty = error.as_ref().map(|e| e.as_any().type_id());
-        println!("Catching {:?}", error.as_ref().map(|e| e.name()));
-        if let Some(catcher) = self.router.catch(status, req, error_ty) {
-            self.invoke_specific_catcher(catcher, status, error.as_ref().map(|e| e.as_ref()), req).await
-        } else if let Some(source) = error.as_ref().and_then(|e| e.source()) {
-            println!("Catching {:?}", source.name());
-            let error_ty = source.as_any().type_id();
-            if let Some(catcher) = self.router.catch(status, req, Some(error_ty)) {
-                self.invoke_specific_catcher(catcher, status, error.as_ref().and_then(|e| e.source()), req).await
-            } else {
-                info!(name: "catcher", name = "rocket::default", "uri.base" = "/", code = status.code,
-                    "no registered catcher: using Rocket default");
-                Ok(catcher::default_handler(status, req))
+        let mut error_copy  = error;
+        let mut counter = 0;
+        // Matches error [.source ...] type
+        while error_copy.is_some() && counter < 5 {
+            if let Some(catcher) = self.router.catch(status, req, error_copy.map(|e| e.trait_obj_typeid())) {
+                return self.invoke_specific_catcher(catcher, status, error_copy, req).await;
             }
-        } else {
-            info!(name: "catcher", name = "rocket::default", "uri.base" = "/", code = status.code,
-                "no registered catcher: using Rocket default");
-            Ok(catcher::default_handler(status, req))
+            error_copy = error_copy.and_then(|e| e.source());
+            counter += 1;
         }
-        // if let Some(catcher) = self.router.catch(status, req, error.as_ref().map(|t| t.as_any().type_id())) {
-        //     catcher.trace_info();
-        //     catch_handle(
-        //         catcher.name.as_deref(),
-        //         || catcher.handler.handle(status, req, error)
-        //     ).await
-        //         .map(|result| result.map_err(|(s, e)| (Some(s), e)))
-        //         .unwrap_or_else(|| Err((None, None)))
-        // } else {
-        //     info!(name: "catcher", name = "rocket::default", "uri.base" = "/", code = status.code,
-        //         "no registered catcher: using Rocket default");
-        //     Ok(catcher::default_handler(status, req))
-        // }
+        // Matches None type
+        if let Some(catcher) = self.router.catch(status, req, None) {
+            return self.invoke_specific_catcher(catcher, status, None, req).await;
+        }
+        let mut error_copy  = error;
+        let mut counter = 0;
+        // Matches error [.source ...] type, and any status
+        while error_copy.is_some() && counter < 5 {
+            if let Some(catcher) = self.router.catch_any(status, req, error_copy.map(|e| e.trait_obj_typeid())) {
+                return self.invoke_specific_catcher(catcher, status, error_copy, req).await;
+            }
+            error_copy = error_copy.and_then(|e| e.source());
+            counter += 1;
+        }
+        // Matches None type, and any status
+        if let Some(catcher) = self.router.catch_any(status, req, None) {
+            return self.invoke_specific_catcher(catcher, status, None, req).await;
+        }
+        if let Some(error) = error {
+            if let Ok(res) = error.respond_to(req) {
+                return Ok(res);
+                // TODO: this ignores the returned status.
+            }
+        }
+        // Rocket default catcher
+        info!(name: "catcher", name = "rocket::default", "uri.base" = "/", code = status.code,
+            "no registered catcher: using Rocket default");
+        Ok(catcher::default_handler(status, req))
+        // TODO: document:
+        // Set of matching catchers, tried in order:
+        // - Matches error type
+        // - Matches error.source type
+        // - Matches error.source.source type
+        // - ... etc
+        // - Matches None type
+        // - Registered default handler
+        // - Rocket default handler
+        // At each step, the catcher with the longest path is selected
     }
 
     async fn invoke_specific_catcher<'s, 'r: 's>(
         &'s self,
         catcher: &Catcher,
         status: Status,
-        error: Option<&'r (dyn TypedError<'r> + 'r)>,
+        error: Option<&'r dyn TypedError<'r>>,
         req: &'r Request<'s>
     ) -> Result<Response<'r>, Option<Status>> {
         catcher.trace_info();
diff --git a/core/lib/src/local/asynchronous/request.rs b/core/lib/src/local/asynchronous/request.rs
@@ -86,18 +86,19 @@ impl<'c> LocalRequest<'c> {
             // _shouldn't_ error. Check that now and error only if not.
             if self.inner().uri() == invalid {
                 error!("invalid request URI: {:?}", invalid.path());
-                return LocalResponse::new(self.request, move |req| {
+                return LocalResponse::new(self.request, move |req, error_ptr| {
                     // TODO: Ideally the RequestErrors should contain actual information.
-                    rocket.dispatch_error(Status::BadRequest, req, Some(Box::new(RequestErrors::new(&[]))))
+                    *error_ptr = Some(Box::new(RequestErrors::new(&[])));
+                    rocket.dispatch_error(Status::BadRequest, req, error_ptr.as_ref().map(|b| b.as_ref()))
                 }).await
             }
         }
 
         // Actually dispatch the request.
         let mut data = Data::local(self.data);
         let token = rocket.preprocess(&mut self.request, &mut data).await;
-        let response = LocalResponse::new(self.request, move |req| {
-            rocket.dispatch(token, req, data)
+        let response = LocalResponse::new(self.request, move |req, error_ptr| {
+            rocket.dispatch(token, req, data, error_ptr)
         }).await;
 
         // If the client is tracking cookies, updates the internal cookie jar
diff --git a/core/lib/src/local/asynchronous/response.rs b/core/lib/src/local/asynchronous/response.rs
@@ -4,6 +4,7 @@ use std::{pin::Pin, task::{Context, Poll}};
 
 use tokio::io::{AsyncRead, ReadBuf};
 
+use crate::catcher::TypedError;
 use crate::http::CookieJar;
 use crate::{Request, Response};
 
@@ -55,6 +56,7 @@ use crate::{Request, Response};
 pub struct LocalResponse<'c> {
     // XXX: SAFETY: This (dependent) field must come first due to drop order!
     response: Response<'c>,
+    _error: Option<Box<dyn TypedError<'c> + 'c>>,
     cookies: CookieJar<'c>,
     _request: Box<Request<'c>>,
 }
@@ -65,8 +67,8 @@ impl Drop for LocalResponse<'_> {
 
 impl<'c> LocalResponse<'c> {
     pub(crate) fn new<F, O>(req: Request<'c>, f: F) -> impl Future<Output = LocalResponse<'c>>
-        where F: FnOnce(&'c Request<'c>) -> O + Send,
-              O: Future<Output = Response<'c>> + Send
+        where F: FnOnce(&'c Request<'c>, &'c mut Option<Box<dyn TypedError<'c> + 'c>>) -> O + Send,
+              O: Future<Output = Response<'c>> + Send + 'c
     {
         // `LocalResponse` is a self-referential structure. In particular,
         // `response` and `cookies` can refer to `_request` and its contents. As
@@ -93,17 +95,21 @@ impl<'c> LocalResponse<'c> {
         let request: &'c Request<'c> = unsafe { &*(&*boxed_req as *const _) };
 
         async move {
+            use std::mem::transmute;
+            let mut error: Option<Box<dyn TypedError<'c> + 'c>> = None;
             // NOTE: The cookie jar `secure` state will not reflect the last
             // known value in `request.cookies()`. This is okay: new cookies
             // should never be added to the resulting jar which is the only time
             // the value is used to set cookie defaults.
-            let response: Response<'c> = f(request).await;
+            // SAFETY: Much like request above, error can borrow from request, and
+            // response can borrow from request. TODO
+            let response: Response<'c> = f(request, unsafe { transmute(&mut error) }).await;
             let mut cookies = CookieJar::new(None, request.rocket());
             for cookie in response.cookies() {
                 cookies.add_original(cookie.into_owned());
             }
 
-            LocalResponse { _request: boxed_req, cookies, response, }
+            LocalResponse { _request: boxed_req, _error: error, cookies, response, }
         }
     }
 }
diff --git a/core/lib/src/router/matcher.rs b/core/lib/src/router/matcher.rs
@@ -137,19 +137,12 @@ impl Catcher {
     /// let b_count = b.base().segments().filter(|s| !s.is_empty()).count();
     /// assert!(b_count > a_count);
     /// ```
+    // TODO: document error matching
     pub fn matches(&self, status: Status, request: &Request<'_>, error: Option<TypeId>) -> bool {
         self.code.map_or(true, |code| code == status.code)
-            && self.error_matches(error)
+            && error == self.error_type.map(|(ty, _)| ty)
             && self.base().segments().prefix_of(request.uri().path().segments())
     }
-
-    fn error_matches(&self, error: Option<TypeId>) -> bool {
-        if let Some((ty, _)) = self.error_type {
-            error.map_or(false, |t| t == ty)
-        } else {
-            true
-        }
-    }
 }
 
 
diff --git a/core/lib/src/router/router.rs b/core/lib/src/router/router.rs
@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 use transient::TypeId;
 
+use crate::catcher::TypedError;
 use crate::request::Request;
 use crate::http::{Method, Status};
 
@@ -54,16 +55,18 @@ impl Router {
     }
 
     // For many catchers, using aho-corasick or similar should be much faster.
-    pub fn catch<'r>(&self, status: Status, req: &'r Request<'r>, error_ty: Option<TypeId>) -> Option<&Catcher> {
+    // TODO: document difference between catch, and catch_any
+    pub fn catch<'r>(&self, status: Status, req: &'r Request<'r>, error: Option<TypeId>) -> Option<&Catcher> {
         // Note that catchers are presorted by descending base length.
         self.catchers.get(&Some(status.code))
-            .and_then(|c| c.iter().find(|c| c.matches(status, req, error_ty)))
+            .and_then(|c| c.iter().find(|c| c.matches(status, req, error)))
     }
 
-    pub fn catch_any<'r>(&self, status: Status, req: &'r Request<'r>, error_ty: Option<TypeId>) -> Option<&Catcher> {
+    // For many catchers, using aho-corasick or similar should be much faster.
+    pub fn catch_any<'r>(&self, status: Status, req: &'r Request<'r>, error: Option<TypeId>) -> Option<&Catcher> {
         // Note that catchers are presorted by descending base length.
         self.catchers.get(&None)
-            .and_then(|c| c.iter().find(|c| c.matches(status, req, error_ty)))
+            .and_then(|c| c.iter().find(|c| c.matches(status, req, error)))
     }
 
     fn collisions<'a, I, T>(&self, items: I) -> impl Iterator<Item = (T, T)> + 'a
diff --git a/core/lib/src/server.rs b/core/lib/src/server.rs
@@ -43,16 +43,17 @@ impl Rocket<Orbit> {
         let mut response = request.into_response(
             stream,
             |rocket, request, data| Box::pin(rocket.preprocess(request, data)),
-            |token, rocket, request, data| Box::pin(async move {
+            |token, rocket, request, data, error_ptr| Box::pin(async move {
                 if !request.errors.is_empty() {
+                    *error_ptr =  Some(Box::new(RequestErrors::new(&request.errors)));
                     return rocket.dispatch_error(
                         Status::BadRequest,
                         request,
-                        Some(Box::new(RequestErrors::new(&request.errors)))
+                        error_ptr.as_ref().map(|b| b.as_ref())
                     ).await;
                 }
 
-                rocket.dispatch(token, request, data).await
+                rocket.dispatch(token, request, data, error_ptr).await
             })
         ).await;
 
