middleware

rpc-toolkit-macro-internals/src/rpc_server/build.rs

@@ -1,4 +1,4 @@
-use proc_macro2::TokenStream;
+use proc_macro2::{Span, TokenStream};
 use quote::quote;
 use syn::spanned::Spanned;
 
@@ -18,7 +18,20 @@ pub fn build(args: RpcServerArgs) -> TokenStream {
     let status_fn = args.status_fn.unwrap_or_else(|| {
         syn::parse2(quote! { |_| ::rpc_toolkit::hyper::StatusCode::OK }).unwrap()
     });
-    quote! {
+    let middleware_name_pre = (0..)
+        .map(|i| Ident::new(&format!("middleware_pre_{}", i), Span::call_site()))
+        .take(args.middleware.len());
+    let middleware_name_pre2 = middleware_name_pre.clone();
+    let middleware_name = (0..)
+        .map(|i| Ident::new(&format!("middleware_{}", i), Span::call_site()))
+        .take(args.middleware.len());
+    let middleware_name_inv = middleware_name
+        .clone()
+        .collect::<Vec<_>>()
+        .into_iter()
+        .rev();
+    let middleware = args.middleware.iter();
+    let res = quote! {
         {
             let ctx = #ctx;
             let status_fn = #status_fn;
@@ -29,28 +42,52 @@ pub fn build(args: RpcServerArgs) -> TokenStream {
                     Ok::<_, ::rpc_toolkit::hyper::Error>(::rpc_toolkit::hyper::service::service_fn(move |mut req| {
                         let ctx = ctx.clone();
                         async move {
+                            #(
+                                let #middleware_name_pre = match ::rpc_toolkit::rpc_server_helpers::constrain_middleware(#middleware)(&mut req).await? {
+                                    Ok(a) => a,
+                                    Err(res) => return Ok(res),
+                                };
+                            )*
                             let rpc_req = ::rpc_toolkit::rpc_server_helpers::make_request(&mut req).await;
-                            ::rpc_toolkit::rpc_server_helpers::to_response(
-                                &req,
-                                match rpc_req {
-                                    Ok(rpc_req) => Ok((
-                                        rpc_req.id,
-                                        #command(
-                                            ctx,
-                                            ::rpc_toolkit::yajrc::RpcMethod::as_str(&rpc_req.method),
-                                            rpc_req.params,
-                                        )
-                                        .await,
-                                    )),
-                                    Err(e) => Err(e),
-                                },
-                                status_fn,
-                            )
+                            match rpc_req {
+                                Ok(mut rpc_req) => {
+                                    #(
+                                        let #middleware_name = match #middleware_name_pre2(&mut rpc_req).await? {
+                                            Ok(a) => a,
+                                            Err(res) => return Ok(res),
+                                        };
+                                    )*
+                                    let mut rpc_res = ::rpc_toolkit::rpc_server_helpers::to_response(
+                                        &req,
+                                        Ok((
+                                            rpc_req.id,
+                                            #command(
+                                                ctx,
+                                                ::rpc_toolkit::yajrc::RpcMethod::as_str(&rpc_req.method),
+                                                rpc_req.params,
+                                            )
+                                            .await,
+                                        )),
+                                        status_fn,
+                                    )?;
+                                    #(
+                                        #middleware_name_inv(&mut rpc_res).await?;
+                                    )*
+                                    Ok::<_, ::rpc_toolkit::hyper::http::Error>(rpc_res)
+                                }
+                                Err(e) => ::rpc_toolkit::rpc_server_helpers::to_response(
+                                    &req,
+                                    Err(e),
+                                    status_fn,
+                                ),
+                            }
                         }
                     }))
                 }
             });
             builder.serve(make_svc)
         }
-    }
+    };
+    // panic!("{}", res);
+    res
 }

rpc-toolkit-macro-internals/src/rpc_server/mod.rs

@@ -4,6 +4,7 @@ pub struct RpcServerArgs {
     command: Path,
     ctx: Expr,
     status_fn: Option<Expr>,
+    middleware: punctuated::Punctuated<Expr, token::Comma>,
 }
 
 pub mod build;

rpc-toolkit-macro-internals/src/rpc_server/parse.rs

@@ -1,24 +1,45 @@
 use syn::parse::{Parse, ParseStream};
+use syn::punctuated::Punctuated;
 
 use super::*;
 
 impl Parse for RpcServerArgs {
     fn parse(input: ParseStream) -> Result<Self> {
-        let command = input.parse()?;
-        let _: token::Comma = input.parse()?;
-        let ctx = input.parse()?;
-        if !input.is_empty() {
-            let _: token::Comma = input.parse()?;
+        let args;
+        braced!(args in input);
+        let mut command = None;
+        let mut ctx = None;
+        let mut status_fn = None;
+        let mut middleware = Punctuated::new();
+        while !args.is_empty() {
+            let arg_name: syn::Ident = args.parse()?;
+            let _: token::Colon = args.parse()?;
+            match arg_name.to_string().as_str() {
+                "command" => {
+                    command = Some(args.parse()?);
+                }
+                "context" => {
+                    ctx = Some(args.parse()?);
+                }
+                "status" => {
+                    status_fn = Some(args.parse()?);
+                }
+                "middleware" => {
+                    let middlewares;
+                    bracketed!(middlewares in args);
+                    middleware = middlewares.parse_terminated(Expr::parse)?;
+                }
+                _ => return Err(Error::new(arg_name.span(), "unknown argument")),
+            }
+            if !args.is_empty() {
+                let _: token::Comma = args.parse()?;
+            }
         }
-        let status_fn = if !input.is_empty() {
-            Some(input.parse()?)
-        } else {
-            None
-        };
         Ok(RpcServerArgs {
-            command,
-            ctx,
+            command: command.expect("`command` is required"),
+            ctx: ctx.expect("`context` is required"),
             status_fn,
+            middleware,
         })
     }
 }