integrate extractors with the macro

packages/fullstack/server-macro/src/lib.rs

@@ -1,6 +1,7 @@
 use proc_macro::TokenStream;
-use quote::ToTokens;
+use quote::{ToTokens, __private::TokenStream as TokenStream2};
 use server_fn_macro::*;
+use syn::{parse::Parse, spanned::Spanned, ItemFn};
 
 /// Declares that a function is a [server function](dioxus_fullstack). This means that
 /// its body will only run on the server, i.e., when the `ssr` feature is enabled.
@@ -54,9 +55,45 @@ use server_fn_macro::*;
 ///   or response or other server-only dependencies, but it does *not* have access to reactive state that exists in the client.
 #[proc_macro_attribute]
 pub fn server(args: proc_macro::TokenStream, s: TokenStream) -> TokenStream {
+    // before we pass this off to the server function macro, we apply extractors and middleware
+    let mut function: syn::ItemFn = match syn::parse(s).map_err(|e| e.to_compile_error()) {
+        Ok(f) => f,
+        Err(e) => return e.into(),
+    };
+
+    // find all arguments with the #[extract] attribute
+    let mut extractors: Vec<Extractor> = vec![];
+    function.sig.inputs = function
+        .sig
+        .inputs
+        .into_iter()
+        .filter(|arg| {
+            if let Ok(extractor) = syn::parse2(arg.clone().into_token_stream()) {
+                extractors.push(extractor);
+                false
+            } else {
+                true
+            }
+        })
+        .collect();
+
+    let ItemFn {
+        attrs,
+        vis,
+        sig,
+        block,
+    } = function;
+    let mapped_body = quote::quote! {
+        #(#attrs)*
+        #vis #sig {
+            #(#extractors)*
+            #block
+        }
+    };
+
     match server_macro_impl(
         args.into(),
-        s.into(),
+        mapped_body,
         syn::parse_quote!(::dioxus_fullstack::prelude::ServerFnTraitObj),
         None,
         Some(syn::parse_quote!(::dioxus_fullstack::prelude::server_fn)),
@@ -65,3 +102,42 @@ pub fn server(args: proc_macro::TokenStream, s: TokenStream) -> TokenStream {
         Ok(s) => s.to_token_stream().into(),
     }
 }
+
+struct Extractor {
+    pat: syn::PatType,
+}
+
+impl ToTokens for Extractor {
+    fn to_tokens(&self, tokens: &mut TokenStream2) {
+        let pat = &self.pat;
+        tokens.extend(quote::quote! {
+            let #pat = ::dioxus_fullstack::prelude::extract_server_context().await?;
+        });
+    }
+}
+
+impl Parse for Extractor {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        let arg: syn::FnArg = input.parse()?;
+        match arg {
+            syn::FnArg::Typed(mut pat_type) => {
+                let mut contains_extract = false;
+                pat_type.attrs.retain(|attr| {
+                    let is_extract = attr.path().is_ident("extract");
+                    if is_extract {
+                        contains_extract = true;
+                    }
+                    !is_extract
+                });
+                if !contains_extract {
+                    return Err(syn::Error::new(
+                        pat_type.span(),
+                        "expected an argument with the #[extract] attribute",
+                    ));
+                }
+                Ok(Extractor { pat: pat_type })
+            }
+            _ => Err(syn::Error::new(arg.span(), "expected a typed argument")),
+        }
+    }
+}