integrate extractors with the macro

packages/fullstack/examples/axum-hello-world/Cargo.toml

@@ -19,3 +19,10 @@ execute = "0.2.12"
 default = []
 ssr = ["axum", "tokio", "dioxus-fullstack/axum"]
 web = ["dioxus-web"]
+
+[profile.release]
+lto = true
+panic = "abort"
+opt-level = 'z'
+strip = true
+codegen-units = 1

packages/fullstack/examples/axum-hello-world/src/main.rs

@@ -41,13 +41,12 @@ fn app(cx: Scope<AppProps>) -> Element {
 }
 
 #[server(PostServerData)]
-async fn post_server_data(data: String) -> Result<(), ServerFnError> {
+async fn post_server_data(
-    // The server context contains information about the current request and allows you to modify the response.
+    #[extract] Axum(axum::extract::Host(host), _): Axum<_, _>,
-    let cx = server_context();
+    data: String,
-    cx.response_headers_mut()
+) -> Result<(), ServerFnError> {
-        .insert("Set-Cookie", "foo=bar".parse().unwrap());
     println!("Server received: {}", data);
-    println!("Request parts are {:?}", cx.request_parts());
+    println!("{:?}", host);
 
     Ok(())
 }

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")),
+        }
+    }
+}

packages/fullstack/src/lib.rs

@@ -38,11 +38,12 @@ pub mod prelude {
     pub use crate::render::SSRState;
     #[cfg(feature = "ssr")]
     pub use crate::serve_config::{ServeConfig, ServeConfigBuilder};
-    #[cfg(feature = "ssr")]
+    #[cfg(all(feature = "ssr", feature = "axum"))]
     pub use crate::server_context::Axum;
     #[cfg(feature = "ssr")]
     pub use crate::server_context::{
-        server_context, DioxusServerContext, FromServerContext, ProvideServerContext,
+        extract_server_context, server_context, DioxusServerContext, FromServerContext,
+        ProvideServerContext,
     };
     pub use crate::server_fn::DioxusServerFn;
     #[cfg(feature = "ssr")]

packages/fullstack/src/server_context.rs

@@ -126,7 +126,12 @@ std::thread_local! {
 ///
 /// This function will only provide the current server context if it is called from a server function.
 pub fn server_context() -> DioxusServerContext {
-    SERVER_CONTEXT.with(|ctx| *ctx.borrow_mut().clone())
+    SERVER_CONTEXT.with(|ctx| *ctx.borrow().clone())
+}
+
+/// Extract some part from the current server request.
+pub async fn extract_server_context<E: FromServerContext>() -> Result<E, E::Rejection> {
+    E::from_request(&server_context()).await
 }
 
 pub(crate) fn with_server_context<O>(
@@ -221,8 +226,9 @@ impl<T: Send + Sync + Clone + 'static> FromServerContext for FromContext<T> {
 pub struct Axum<
     I: axum::extract::FromRequestParts<(), Rejection = R>,
     R: axum::response::IntoResponse + std::error::Error,
->(pub(crate) I, std::marker::PhantomData<R>);
+>(pub I, pub std::marker::PhantomData<R>);
 
+#[cfg(feature = "axum")]
 impl<
         I: axum::extract::FromRequestParts<(), Rejection = R>,
         R: axum::response::IntoResponse + std::error::Error,
@@ -235,6 +241,7 @@ impl<
     }
 }
 
+#[cfg(feature = "axum")]
 impl<
         I: axum::extract::FromRequestParts<(), Rejection = R>,
         R: axum::response::IntoResponse + std::error::Error,