Typed catchers

core/codegen/src/attribute/catch/mod.rs

@@ -1,19 +1,73 @@
 mod parse;
 
-use devise::ext::SpanDiagnosticExt;
-use devise::{Spanned, Result};
+use devise::{Result, Spanned};
 use proc_macro2::{TokenStream, Span};
 
 use crate::http_codegen::Optional;
-use crate::syn_ext::ReturnTypeExt;
+use crate::syn_ext::{IdentExt, ReturnTypeExt};
 use crate::exports::*;
 
+use self::parse::ErrorGuard;
+
+use super::param::Guard;
+
+fn error_type(guard: &ErrorGuard) -> TokenStream {
+    let ty = &guard.ty;
+    quote! {
+        (#_catcher::TypeId::of::<#ty>(), ::std::any::type_name::<#ty>())
+    }
+}
+
+fn error_guard_decl(guard: &ErrorGuard) -> TokenStream {
+    let (ident, ty) = (guard.ident.rocketized(), &guard.ty);
+    quote_spanned! { ty.span() =>
+        let #ident: &#ty = match #_catcher::downcast(__error_init.as_ref()) {
+            Some(v) => v,
+            None => return #_Result::Err((#__status, __error_init)),
+        };
+    }
+}
+
+fn request_guard_decl(guard: &Guard) -> TokenStream {
+    let (ident, ty) = (guard.fn_ident.rocketized(), &guard.ty);
+    quote_spanned! { ty.span() =>
+        let #ident: #ty = match <#ty as #FromRequest>::from_request(#__req).await {
+            #Outcome::Success(__v) => __v,
+            #Outcome::Forward(__e) => {
+                ::rocket::trace::info!(
+                    name: "forward",
+                    target: concat!("rocket::codegen::catch::", module_path!()),
+                    parameter = stringify!(#ident),
+                    type_name = stringify!(#ty),
+                    status = __e.code,
+                    "request guard forwarding; trying next catcher"
+                );
+
+                return #_Err((#__status, __error_init));
+            },
+            #[allow(unreachable_code)]
+            #Outcome::Error((__c, __e)) => {
+                ::rocket::trace::info!(
+                    name: "failure",
+                    target: concat!("rocket::codegen::catch::", module_path!()),
+                    parameter = stringify!(#ident),
+                    type_name = stringify!(#ty),
+                    reason = %#display_hack!(&__e),
+                    "request guard failed; forwarding to 500 handler"
+                );
+
+                return #_Err((#Status::InternalServerError, __error_init));
+            }
+        };
+    }
+}
+
 pub fn _catch(
     args: proc_macro::TokenStream,
     input: proc_macro::TokenStream
 ) -> Result<TokenStream> {
     // Parse and validate all of the user's input.
-    let catch = parse::Attribute::parse(args.into(), input)?;
+    let catch = parse::Attribute::parse(args.into(), input.into())?;
 
     // Gather everything we'll need to generate the catcher.
     let user_catcher_fn = &catch.function;
@@ -22,35 +76,28 @@ pub fn _catch(
     let status_code = Optional(catch.status.map(|s| s.code));
     let deprecated = catch.function.attrs.iter().find(|a| a.path().is_ident("deprecated"));
 
-    // Determine the number of parameters that will be passed in.
-    if catch.function.sig.inputs.len() > 2 {
-        return Err(catch.function.sig.paren_token.span.join()
-            .error("invalid number of arguments: must be zero, one, or two")
-            .help("catchers optionally take `&Request` or `Status, &Request`"));
-    }
-
     // This ensures that "Responder not implemented" points to the return type.
     let return_type_span = catch.function.sig.output.ty()
         .map(|ty| ty.span())
         .unwrap_or_else(Span::call_site);
 
-    // Set the `req` and `status` spans to that of their respective function
-    // arguments for a more correct `wrong type` error span. `rev` to be cute.
-    let codegen_args = &[__req, __status];
-    let inputs = catch.function.sig.inputs.iter().rev()
-        .zip(codegen_args.iter())
-        .map(|(fn_arg, codegen_arg)| match fn_arg {
-            syn::FnArg::Receiver(_) => codegen_arg.respanned(fn_arg.span()),
-            syn::FnArg::Typed(a) => codegen_arg.respanned(a.ty.span())
-        }).rev();
+    let status_guard = catch.status_guard.as_ref().map(|(_, s)| {
+        let ident = s.rocketized();
+        quote! { let #ident = #__status; }
+    });
+    let error_guard = catch.error_guard.as_ref().map(error_guard_decl);
+    let error_type = Optional(catch.error_guard.as_ref().map(error_type));
+    let request_guards = catch.request_guards.iter().map(request_guard_decl);
+    let parameter_names = catch.arguments.map.values()
+        .map(|(ident, _)| ident.rocketized());
 
     // We append `.await` to the function call if this is `async`.
     let dot_await = catch.function.sig.asyncness
         .map(|a| quote_spanned!(a.span() => .await));
 
     let catcher_response = quote_spanned!(return_type_span => {
-        let ___responder = #user_catcher_fn_name(#(#inputs),*) #dot_await;
-        #_response::Responder::respond_to(___responder, #__req)?
+        let ___responder = #user_catcher_fn_name(#(#parameter_names),*) #dot_await;
+        #_response::Responder::respond_to(___responder, #__req).map_err(|s| (s, __error_init))?
     });
 
     // Generate the catcher, keeping the user's input around.
@@ -68,20 +115,27 @@ pub fn _catch(
             fn into_info(self) -> #_catcher::StaticInfo {
                 fn monomorphized_function<'__r>(
                     #__status: #Status,
-                    #__req: &'__r #Request<'_>
+                    #__req: &'__r #Request<'_>,
+                    __error_init: #ErasedError<'__r>,
                 ) -> #_catcher::BoxFuture<'__r> {
                     #_Box::pin(async move {
+                        #error_guard
+                        #status_guard
+                        #(#request_guards)*
                         let __response = #catcher_response;
-                        #Response::build()
-                            .status(#__status)
-                            .merge(__response)
-                            .ok()
+                        #_Result::Ok(
+                            #Response::build()
+                                .status(#__status)
+                                .merge(__response)
+                                .finalize()
+                        )
                     })
                 }
 
                 #_catcher::StaticInfo {
                     name: ::core::stringify!(#user_catcher_fn_name),
                     code: #status_code,
+                    error_type: #error_type,
                     handler: monomorphized_function,
                     location: (::core::file!(), ::core::line!(), ::core::column!()),
                 }

core/codegen/src/attribute/catch/parse.rs

@@ -1,7 +1,12 @@
-use devise::ext::SpanDiagnosticExt;
-use devise::{MetaItem, Spanned, Result, FromMeta, Diagnostic};
-use proc_macro2::TokenStream;
+use devise::ext::{SpanDiagnosticExt, TypeExt};
+use devise::{Diagnostic, FromMeta, MetaItem, Result, SpanWrapped, Spanned};
+use proc_macro2::{Span, TokenStream, Ident};
+use quote::ToTokens;
 
+use crate::attribute::param::{Dynamic, Guard};
+use crate::name::{ArgumentMap, Arguments, Name};
+use crate::proc_macro_ext::Diagnostics;
+use crate::syn_ext::FnArgExt;
 use crate::{http, http_codegen};
 
 /// This structure represents the parsed `catch` attribute and associated items.
@@ -10,13 +15,66 @@ pub struct Attribute {
     pub status: Option<http::Status>,
     /// The function that was decorated with the `catch` attribute.
     pub function: syn::ItemFn,
+    pub arguments: Arguments,
+    pub error_guard: Option<ErrorGuard>,
+    pub status_guard: Option<(Name, syn::Ident)>,
+    pub request_guards: Vec<Guard>,
+}
+
+pub struct ErrorGuard {
+    pub span: Span,
+    pub name: Name,
+    pub ident: syn::Ident,
+    pub ty: syn::Type,
+}
+
+impl ErrorGuard {
+    fn new(param: SpanWrapped<Dynamic>, args: &Arguments) -> Result<Self> {
+        if let Some((ident, ty)) = args.map.get(&param.name) {
+            match ty {
+                syn::Type::Reference(syn::TypeReference { elem, .. }) => Ok(Self {
+                    span: param.span(),
+                    name: param.name.clone(),
+                    ident: ident.clone(),
+                    ty: elem.as_ref().clone(),
+                }),
+                ty => {
+                    let msg = format!(
+                        "Error argument must be a reference, found `{}`",
+                        ty.to_token_stream()
+                    );
+                    let diag = param.span()
+                        .error("invalid type")
+                        .span_note(ty.span(), msg)
+                        .help(format!("Perhaps use `&{}` instead", ty.to_token_stream()));
+                    Err(diag)
+                }
+            }
+        } else {
+            let msg = format!("expected argument named `{}` here", param.name);
+            let diag = param.span().error("unused parameter").span_note(args.span, msg);
+            Err(diag)
+        }
+    }
+}
+
+fn status_guard(param: SpanWrapped<Dynamic>, args: &Arguments) -> Result<(Name, Ident)> {
+    if let Some((ident, _)) = args.map.get(&param.name) {
+        Ok((param.name.clone(), ident.clone()))
+    } else {
+        let msg = format!("expected argument named `{}` here", param.name);
+        let diag = param.span().error("unused parameter").span_note(args.span, msg);
+        Err(diag)
+    }
 }
 
 /// We generate a full parser for the meta-item for great error messages.
 #[derive(FromMeta)]
 struct Meta {
     #[meta(naked)]
     code: Code,
+    error: Option<SpanWrapped<Dynamic>>,
+    status: Option<SpanWrapped<Dynamic>>,
 }
 
 /// `Some` if there's a code, `None` if it's `default`.
@@ -43,16 +101,66 @@ impl FromMeta for Code {
 
 impl Attribute {
     pub fn parse(args: TokenStream, input: proc_macro::TokenStream) -> Result<Self> {
+        let mut diags = Diagnostics::new();
+
         let function: syn::ItemFn = syn::parse(input)
             .map_err(Diagnostic::from)
             .map_err(|diag| diag.help("`#[catch]` can only be used on functions"))?;
 
         let attr: MetaItem = syn::parse2(quote!(catch(#args)))?;
-        let status = Meta::from_meta(&attr)
-            .map(|meta| meta.code.0)
+        let attr = Meta::from_meta(&attr)
+            .map(|meta| meta)
             .map_err(|diag| diag.help("`#[catch]` expects a status code int or `default`: \
                         `#[catch(404)]` or `#[catch(default)]`"))?;
 
-        Ok(Attribute { status, function })
+        let span = function.sig.paren_token.span.join();
+        let mut arguments = Arguments { map: ArgumentMap::new(), span };
+        for arg in function.sig.inputs.iter() {
+            if let Some((ident, ty)) = arg.typed() {
+                let value = (ident.clone(), ty.with_stripped_lifetimes());
+                arguments.map.insert(Name::from(ident), value);
+            } else {
+                let span = arg.span();
+                let diag = if arg.wild().is_some() {
+                    span.error("handler arguments must be named")
+                        .help("to name an ignored handler argument, use `_name`")
+                } else {
+                    span.error("handler arguments must be of the form `ident: Type`")
+                };
+
+                diags.push(diag);
+            }
+        }
+        // let mut error_guard = None;
+        let error_guard = attr.error.clone()
+            .map(|p| ErrorGuard::new(p, &arguments))
+            .and_then(|p| p.map_err(|e| diags.push(e)).ok());
+        let status_guard = attr.status.clone()
+            .map(|n| status_guard(n, &arguments))
+            .and_then(|p| p.map_err(|e| diags.push(e)).ok());
+        let request_guards = arguments.map.iter()
+            .filter(|(name, _)| {
+                let mut all_other_guards = error_guard.iter()
+                    .map(|g| &g.name)
+                    .chain(status_guard.iter().map(|(n, _)| n));
+
+                all_other_guards.all(|n| n != *name)
+            })
+            .enumerate()
+            .map(|(index, (name, (ident, ty)))| Guard {
+                source: Dynamic { index, name: name.clone(), trailing: false },
+                fn_ident: ident.clone(),
+                ty: ty.clone(),
+            })
+            .collect();
+
+        diags.head_err_or(Attribute {
+            status: attr.code.0,
+            function,
+            arguments,
+            error_guard,
+            status_guard,
+            request_guards,
+        })
     }
 }