add server function middleware

packages/fullstack/Cargo.toml

@@ -19,11 +19,12 @@ dioxus_server_macro = { path = "server-macro" }
 # warp
 warp = { version = "0.3.5", features = ["compression-gzip"], optional = true }
 http-body = { version = "0.4.5", optional = true }
+http-body-util = "0.1.0-rc.2"
 
 # axum
 axum = { version = "0.6.1", features = ["ws", "macros"], optional = true }
 tower-http = { version = "0.4.0", optional = true, features = ["fs", "compression-gzip"] }
-tower = { version = "0.4.13", optional = true }
+tower = { version = "0.4.13", features = ["util"], optional = true }
 axum-macros = "0.3.7"
 
 # salvo
@@ -54,6 +55,8 @@ base64 = "0.21.0"
 
 pin-project = "1.1.2"
 async-trait = "0.1.71"
+bytes = "1.4.0"
+tower-layer = "0.3.2"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 dioxus-hot-reload = { workspace = true }
@@ -65,10 +68,10 @@ web-sys = { version = "0.3.61", features = ["Window", "Document", "Element", "Ht
 default = ["hot-reload", "default-tls"]
 router = ["dioxus-router"]
 hot-reload = ["serde_json", "futures-util"]
-warp = ["dep:warp", "http-body", "ssr"]
-axum = ["dep:axum", "tower-http", "tower", "ssr"]
+warp = ["dep:warp", "ssr"]
+axum = ["dep:axum", "tower-http", "ssr"]
 salvo = ["dep:salvo", "ssr"]
-ssr = ["server_fn/ssr", "tokio", "dioxus-ssr", "hyper", "http", "dioxus-router/ssr", "tokio-stream"]
+ssr = ["server_fn/ssr", "tokio", "dioxus-ssr", "tower", "hyper", "http", "http-body", "dioxus-router/ssr", "tokio-stream"]
 default-tls = ["server_fn/default-tls"]
 rustls = ["server_fn/rustls"]
 

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

@@ -14,6 +14,7 @@ axum = { version = "0.6.12", optional = true }
 tokio = { workspace = true, features = ["full"], optional = true }
 serde = "1.0.159"
 execute = "0.2.12"
+tower-http = { version = "0.4.1", features = ["auth"] }
 
 [features]
 default = []

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

@@ -41,9 +41,9 @@ fn app(cx: Scope<AppProps>) -> Element {
 }
 
 #[server(PostServerData)]
-#[middleware(Auth(AuthLevel::Admin))]
+#[middleware(tower_http::validate_request::ValidateRequestHeaderLayer::bearer("passwordlol"))]
 async fn post_server_data(
-    #[extract] Axum(axum::extract::Host(host), _): Axum<_, _>,
+    #[extract] axum::extract::Host(host): axum::extract::Host,
     data: String,
 ) -> Result<(), ServerFnError> {
     println!("Server received: {}", data);

packages/fullstack/server-macro/Cargo.toml

@@ -6,9 +6,11 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+proc-macro2 = "1.0.63"
 quote = "1.0.26"
 server_fn_macro = { git = "https://github.com/leptos-rs/leptos", rev = "6b90e1babd425c9a324181c86e3fd1b942c9b10f" }
 syn = { version = "2", features = ["full"] }
+convert_case = "0.6.0"
 
 [lib]
 proc-macro = true

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

@@ -1,7 +1,13 @@
+use convert_case::{Case, Converter};
 use proc_macro::TokenStream;
+use proc_macro2::Literal;
 use quote::{ToTokens, __private::TokenStream as TokenStream2};
 use server_fn_macro::*;
-use syn::{parse::Parse, spanned::Spanned, ItemFn};
+use syn::{
+    parse::{Parse, ParseStream},
+    spanned::Spanned,
+    Ident, ItemFn, Token,
+};
 
 /// 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.
@@ -77,6 +83,25 @@ pub fn server(args: proc_macro::TokenStream, s: TokenStream) -> TokenStream {
         })
         .collect();
 
+    // extract all #[middleware] attributes
+    let mut middlewares: Vec<Middleware> = vec![];
+    function.attrs = function
+        .attrs
+        .into_iter()
+        .filter(|attr| {
+            if attr.meta.path().is_ident("middleware") {
+                if let Ok(middleware) = attr.parse_args() {
+                    middlewares.push(middleware);
+                    false
+                } else {
+                    true
+                }
+            } else {
+                true
+            }
+        })
+        .collect();
+
     let ItemFn {
         attrs,
         vis,
@@ -91,15 +116,49 @@ pub fn server(args: proc_macro::TokenStream, s: TokenStream) -> TokenStream {
         }
     };
 
+    let server_fn_path: syn::Path = syn::parse_quote!(::dioxus_fullstack::prelude::server_fn);
+    let trait_obj_wrapper: syn::Type =
+        syn::parse_quote!(::dioxus_fullstack::prelude::ServerFnTraitObj);
+    let mut args: ServerFnArgs = match syn::parse(args) {
+        Ok(args) => args,
+        Err(e) => return e.to_compile_error().into(),
+    };
+    if args.struct_name.is_none() {
+        let upper_cammel_case_name = Converter::new()
+            .from_case(Case::Snake)
+            .to_case(Case::UpperCamel)
+            .convert(&sig.ident.to_string());
+        args.struct_name = Some(Ident::new(
+            &format!("{}", upper_cammel_case_name),
+            sig.ident.span(),
+        ));
+    }
+    let struct_name = args.struct_name.as_ref().unwrap();
     match server_macro_impl(
-        args.into(),
+        quote::quote!(#args),
         mapped_body,
-        syn::parse_quote!(::dioxus_fullstack::prelude::ServerFnTraitObj),
+        trait_obj_wrapper.clone(),
         None,
-        Some(syn::parse_quote!(::dioxus_fullstack::prelude::server_fn)),
+        Some(server_fn_path.clone()),
     ) {
         Err(e) => e.to_compile_error().into(),
-        Ok(s) => s.to_token_stream().into(),
+        Ok(tokens) => quote::quote! {
+            #tokens
+            #[cfg(feature = "ssr")]
+            #server_fn_path::inventory::submit! {
+                ::dioxus_fullstack::prelude::ServerFnMiddleware {
+                    prefix: #struct_name::PREFIX,
+                    url: #struct_name::URL,
+                    middleware: || vec![
+                        #(
+                            std::sync::Arc::new(#middlewares),
+                        ),*
+                    ]
+                }
+            }
+        }
+        .to_token_stream()
+        .into(),
     }
 }
 
@@ -141,3 +200,71 @@ impl Parse for Extractor {
         }
     }
 }
+
+#[derive(Debug)]
+struct Middleware {
+    expr: syn::Expr,
+}
+
+impl ToTokens for Middleware {
+    fn to_tokens(&self, tokens: &mut TokenStream2) {
+        let expr = &self.expr;
+        tokens.extend(quote::quote! {
+            #expr
+        });
+    }
+}
+
+impl Parse for Middleware {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        let arg: syn::Expr = input.parse()?;
+        Ok(Middleware { expr: arg })
+    }
+}
+
+struct ServerFnArgs {
+    struct_name: Option<Ident>,
+    _comma: Option<Token![,]>,
+    prefix: Option<Literal>,
+    _comma2: Option<Token![,]>,
+    encoding: Option<Literal>,
+    _comma3: Option<Token![,]>,
+    fn_path: Option<Literal>,
+}
+
+impl ToTokens for ServerFnArgs {
+    fn to_tokens(&self, tokens: &mut TokenStream2) {
+        let struct_name = self.struct_name.as_ref().map(|s| quote::quote! { #s, });
+        let prefix = self.prefix.as_ref().map(|p| quote::quote! { #p, });
+        let encoding = self.encoding.as_ref().map(|e| quote::quote! { #e, });
+        let fn_path = self.fn_path.as_ref().map(|f| quote::quote! { #f, });
+        tokens.extend(quote::quote! {
+            #struct_name
+            #prefix
+            #encoding
+            #fn_path
+        })
+    }
+}
+
+impl Parse for ServerFnArgs {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let struct_name = input.parse()?;
+        let _comma = input.parse()?;
+        let prefix = input.parse()?;
+        let _comma2 = input.parse()?;
+        let encoding = input.parse()?;
+        let _comma3 = input.parse()?;
+        let fn_path = input.parse()?;
+
+        Ok(Self {
+            struct_name,
+            _comma,
+            prefix,
+            _comma2,
+            encoding,
+            _comma3,
+            fn_path,
+        })
+    }
+}

packages/fullstack/src/adapters/axum_adapter.rs

@@ -255,11 +255,22 @@ where
 
     fn register_server_fns(self, server_fn_route: &'static str) -> Self {
         self.register_server_fns_with_handler(server_fn_route, |func| {
-            move |req: Request<Body>| async move {
-                let (parts, body) = req.into_parts();
-                let parts: Arc<RwLock<http::request::Parts>> = Arc::new(parts.into());
-                let server_context = DioxusServerContext::new(parts.clone());
-                server_fn_handler(server_context, func.clone(), parts, body).await
+            use crate::layer::Service;
+            move |req: Request<Body>| {
+                let mut service = crate::server_fn_service(Default::default(), func);
+                async move {
+                    let (req, body) = req.into_parts();
+                    let req = Request::from_parts(req, body);
+                    let res = service.run(req);
+                    match res.await {
+                        Ok(res) => Ok::<_, std::convert::Infallible>(res.map(|b| b.into())),
+                        Err(e) => {
+                            let mut res = Response::new(Body::empty());
+                            *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
+                            Ok(res)
+                        }
+                    }
+                }
             }
         })
     }
@@ -373,7 +384,7 @@ async fn render_handler<P: Clone + serde::Serialize + Send + Sync + 'static>(
             let crate::render::RenderResponse { html, freshness } = rendered;
             let mut response = axum::response::Html::from(html).into_response();
             freshness.write(response.headers_mut());
-            let headers = server_context.take_response_headers();
+            let headers = server_context.response_parts().unwrap().headers.clone();
             apply_request_parts_to_response(headers, &mut response);
             response
         }
@@ -384,90 +395,6 @@ async fn render_handler<P: Clone + serde::Serialize + Send + Sync + 'static>(
     }
 }
 
-/// A default handler for server functions. It will deserialize the request, call the server function, and serialize the response.
-pub async fn server_fn_handler(
-    server_context: DioxusServerContext,
-    function: server_fn::ServerFnTraitObj<()>,
-    parts: Arc<RwLock<http::request::Parts>>,
-    body: Body,
-) -> impl IntoResponse {
-    let body = hyper::body::to_bytes(body).await;
-    let Ok(body) = body else {
-        return report_err(body.err().unwrap());
-    };
-
-    // Because the future returned by `server_fn_handler` is `Send`, and the future returned by this function must be send, we need to spawn a new runtime
-    let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
-    let query_string = parts
-        .read()
-        .unwrap()
-        .uri
-        .query()
-        .unwrap_or_default()
-        .to_string();
-    spawn_blocking({
-        move || {
-            tokio::runtime::Runtime::new()
-                .expect("couldn't spawn runtime")
-                .block_on(async {
-                    let query = &query_string.into();
-                    let data = match &function.encoding() {
-                        Encoding::Url | Encoding::Cbor => &body,
-                        Encoding::GetJSON | Encoding::GetCBOR => query,
-                    };
-                    let server_function_future = function.call((), data);
-                    let server_function_future =
-                        ProvideServerContext::new(server_function_future, server_context.clone());
-                    let resp = match server_function_future.await {
-                        Ok(serialized) => {
-                            // if this is Accept: application/json then send a serialized JSON response
-                            let parts = parts.read().unwrap();
-                            let accept_header = parts
-                                .headers
-                                .get("Accept")
-                                .and_then(|value| value.to_str().ok());
-                            let mut res = Response::builder();
-                            *res.headers_mut().expect("empty response should be valid") =
-                                server_context.take_response_headers();
-                            if accept_header == Some("application/json")
-                                || accept_header
-                                    == Some(
-                                        "application/\
-                                                 x-www-form-urlencoded",
-                                    )
-                                || accept_header == Some("application/cbor")
-                            {
-                                res = res.status(StatusCode::OK);
-                            }
-
-                            let resp = match serialized {
-                                Payload::Binary(data) => res
-                                    .header("Content-Type", "application/cbor")
-                                    .body(body::boxed(Full::from(data))),
-                                Payload::Url(data) => res
-                                    .header(
-                                        "Content-Type",
-                                        "application/\
-                                        x-www-form-urlencoded",
-                                    )
-                                    .body(body::boxed(data)),
-                                Payload::Json(data) => res
-                                    .header("Content-Type", "application/json")
-                                    .body(body::boxed(data)),
-                            };
-
-                            resp.unwrap()
-                        }
-                        Err(e) => report_err(e),
-                    };
-
-                    resp_tx.send(resp).unwrap();
-                })
-        }
-    });
-    resp_rx.await.unwrap()
-}
-
 fn report_err<E: std::fmt::Display>(e: E) -> Response<BoxBody> {
     Response::builder()
         .status(StatusCode::INTERNAL_SERVER_ERROR)

packages/fullstack/src/adapters/mod.rs

@@ -16,3 +16,146 @@ pub mod axum_adapter;
 pub mod salvo_adapter;
 #[cfg(feature = "warp")]
 pub mod warp_adapter;
+
+use std::sync::{Arc, RwLock};
+
+use http::StatusCode;
+use http_body_util::BodyExt;
+use hyper::body::Body;
+use server_fn::{Encoding, Payload};
+use tokio::task::spawn_blocking;
+
+use crate::{
+    layer::{BoxedService, Service},
+    prelude::{DioxusServerContext, ProvideServerContext},
+};
+
+/// Create a server function handler with the given server context and server function.
+pub fn server_fn_service(
+    context: DioxusServerContext,
+    function: server_fn::ServerFnTraitObj<()>,
+) -> crate::layer::BoxedService {
+    let prefix = function.prefix().to_string();
+    let url = function.url().to_string();
+    if let Some(middleware) = crate::server_fn::MIDDLEWARE.get(&(&prefix, &url)) {
+        let mut service = BoxedService(Box::new(ServerFnHandler::new(context, function)));
+        for middleware in middleware {
+            service = middleware.layer(service);
+        }
+        service
+    } else {
+        BoxedService(Box::new(ServerFnHandler::new(context, function)))
+    }
+}
+
+#[derive(Clone)]
+/// A default handler for server functions. It will deserialize the request body, call the server function, and serialize the response.
+pub struct ServerFnHandler {
+    server_context: DioxusServerContext,
+    function: server_fn::ServerFnTraitObj<()>,
+}
+
+impl ServerFnHandler {
+    /// Create a new server function handler with the given server context and server function.
+    pub fn new(
+        server_context: impl Into<DioxusServerContext>,
+        function: server_fn::ServerFnTraitObj<()>,
+    ) -> Self {
+        let server_context = server_context.into();
+        Self {
+            server_context,
+            function,
+        }
+    }
+}
+
+impl Service for ServerFnHandler {
+    fn run(
+        &mut self,
+        req: http::Request<hyper::body::Body>,
+    ) -> std::pin::Pin<
+        Box<
+            dyn std::future::Future<
+                    Output = Result<http::Response<hyper::body::Body>, server_fn::ServerFnError>,
+                > + Send,
+        >,
+    > {
+        let Self {
+            server_context,
+            function,
+        } = self.clone();
+        Box::pin(async move {
+            let query = req.uri().query().unwrap_or_default().as_bytes().to_vec();
+            let (parts, body) = req.into_parts();
+            let body = hyper::body::to_bytes(body).await?.to_vec();
+            let headers = &parts.headers;
+            let accept_header = headers.get("Accept").cloned();
+            let parts = Arc::new(RwLock::new(parts));
+
+            // Because the future returned by `server_fn_handler` is `Send`, and the future returned by this function must be send, we need to spawn a new runtime
+            let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
+            spawn_blocking({
+                let function = function.clone();
+                let mut server_context = server_context.clone();
+                server_context.parts = parts;
+                move || {
+                    tokio::runtime::Runtime::new()
+                        .expect("couldn't spawn runtime")
+                        .block_on(async move {
+                            let data = match function.encoding() {
+                                Encoding::Url | Encoding::Cbor => &body,
+                                Encoding::GetJSON | Encoding::GetCBOR => &query,
+                            };
+                            let server_function_future = function.call((), data);
+                            let server_function_future = ProvideServerContext::new(
+                                server_function_future,
+                                server_context.clone(),
+                            );
+                            let resp = server_function_future.await;
+
+                            resp_tx.send(resp).unwrap();
+                        })
+                }
+            });
+            let result = resp_rx.await.unwrap();
+            let mut res = http::Response::builder();
+
+            // Set the headers from the server context
+            let parts = server_context.parts.read().unwrap();
+            *res.headers_mut().expect("empty headers should be valid") = parts.headers.clone();
+
+            let serialized = result?;
+            // if this is Accept: application/json then send a serialized JSON response
+            let accept_header = accept_header.as_ref().and_then(|value| value.to_str().ok());
+            if accept_header == Some("application/json")
+                || accept_header
+                    == Some(
+                        "application/\
+                                x-www-form-urlencoded",
+                    )
+                || accept_header == Some("application/cbor")
+            {
+                res = res.status(StatusCode::OK);
+            }
+
+            Ok(match serialized {
+                Payload::Binary(data) => {
+                    res = res.header("Content-Type", "application/cbor");
+                    res.body(data.into())?
+                }
+                Payload::Url(data) => {
+                    res = res.header(
+                        "Content-Type",
+                        "application/\
+                                    x-www-form-urlencoded",
+                    );
+                    res.body(data.into())?
+                }
+                Payload::Json(data) => {
+                    res = res.header("Content-Type", "application/json");
+                    res.body(data.into())?
+                }
+            })
+        })
+    }
+}

packages/fullstack/src/layer.rs

@@ -0,0 +1,77 @@
+use std::pin::Pin;
+
+use http::{Request, Response};
+use hyper::Body;
+
+pub trait Layer: Send + Sync + 'static {
+    fn layer(&self, inner: BoxedService) -> BoxedService;
+}
+
+impl<L> Layer for L
+where
+    L: tower_layer::Layer<BoxedService> + Sync + Send + 'static,
+    L::Service: Service + Send + 'static,
+{
+    fn layer(&self, inner: BoxedService) -> BoxedService {
+        BoxedService(Box::new(self.layer(inner)))
+    }
+}
+
+pub trait Service {
+    fn run(
+        &mut self,
+        req: http::Request<hyper::body::Body>,
+    ) -> Pin<
+        Box<
+            dyn std::future::Future<
+                    Output = Result<Response<hyper::body::Body>, server_fn::ServerFnError>,
+                > + Send,
+        >,
+    >;
+}
+
+impl<S> Service for S
+where
+    S: tower::Service<http::Request<hyper::body::Body>, Response = Response<hyper::body::Body>>,
+    S::Future: Send + 'static,
+    S::Error: Into<server_fn::ServerFnError>,
+{
+    fn run(
+        &mut self,
+        req: http::Request<hyper::body::Body>,
+    ) -> Pin<
+        Box<
+            dyn std::future::Future<
+                    Output = Result<Response<hyper::body::Body>, server_fn::ServerFnError>,
+                > + Send,
+        >,
+    > {
+        let fut = self.call(req);
+        Box::pin(async move { fut.await.map_err(|err| err.into()) })
+    }
+}
+
+pub struct BoxedService(pub Box<dyn Service + Send>);
+
+impl tower::Service<http::Request<hyper::body::Body>> for BoxedService {
+    type Response = http::Response<hyper::body::Body>;
+    type Error = server_fn::ServerFnError;
+    type Future = Pin<
+        Box<
+            dyn std::future::Future<
+                    Output = Result<http::Response<hyper::body::Body>, server_fn::ServerFnError>,
+                > + Send,
+        >,
+    >;
+
+    fn poll_ready(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), Self::Error>> {
+        Ok(()).into()
+    }
+
+    fn call(&mut self, req: Request<hyper::body::Body>) -> Self::Future {
+        self.0.run(req)
+    }
+}

packages/fullstack/src/lib.rs

@@ -3,18 +3,23 @@
 #![doc(html_favicon_url = "https://avatars.githubusercontent.com/u/79236386")]
 #![deny(missing_docs)]
 
-pub use adapters::*;
+pub use once_cell;
 
 mod props_html;
 
 #[cfg(feature = "router")]
 pub mod router;
 
+#[cfg(feature = "ssr")]
 mod adapters;
+#[cfg(feature = "ssr")]
+pub use adapters::*;
 #[cfg(all(debug_assertions, feature = "hot-reload", feature = "ssr"))]
 mod hot_reload;
 pub mod launch;
 #[cfg(feature = "ssr")]
+mod layer;
+#[cfg(feature = "ssr")]
 mod render;
 #[cfg(feature = "ssr")]
 mod serve_config;
@@ -47,7 +52,7 @@ pub mod prelude {
     };
     pub use crate::server_fn::DioxusServerFn;
     #[cfg(feature = "ssr")]
-    pub use crate::server_fn::{ServerFnTraitObj, ServerFunction};
+    pub use crate::server_fn::{ServerFnMiddleware, ServerFnTraitObj, ServerFunction};
     pub use crate::{launch, launch_router};
     pub use dioxus_server_macro::*;
     #[cfg(feature = "ssr")]

packages/fullstack/src/router.rs

@@ -26,7 +26,7 @@ where
                         #[cfg(feature = "ssr")]
                         let history = dioxus_router::prelude::MemoryHistory::with_initial_path(
                             context
-                                .request_parts()
+                                .request_parts().unwrap()
                                 .uri
                                 .to_string()
                                 .parse()

packages/fullstack/src/server_context.rs

@@ -2,7 +2,7 @@ pub use server_fn_impl::*;
 use std::sync::Arc;
 use std::sync::RwLock;
 
-/// A shared context for server functions that contains infomation about the request and middleware state.
+/// A shared context for server functions that contains information about the request and middleware state.
 /// This allows you to pass data between your server framework and the server functions. This can be used to pass request information or information about the state of the server. For example, you could pass authentication data though this context to your server functions.
 ///
 /// You should not construct this directly inside components. Instead use the `HasServerContext` trait to get the server context from the scope.
@@ -11,7 +11,7 @@ pub struct DioxusServerContext {
     shared_context: std::sync::Arc<
         std::sync::RwLock<anymap::Map<dyn anymap::any::Any + Send + Sync + 'static>>,
     >,
-    headers: std::sync::Arc<std::sync::RwLock<hyper::header::HeaderMap>>,
+    response_parts: std::sync::Arc<std::sync::RwLock<http::response::Parts>>,
     pub(crate) parts: Arc<RwLock<http::request::Parts>>,
 }
 
@@ -20,7 +20,9 @@ impl Default for DioxusServerContext {
     fn default() -> Self {
         Self {
             shared_context: std::sync::Arc::new(std::sync::RwLock::new(anymap::Map::new())),
-            headers: Default::default(),
+            response_parts: std::sync::Arc::new(RwLock::new(
+                http::response::Response::new(()).into_parts().0,
+            )),
             parts: std::sync::Arc::new(RwLock::new(http::request::Request::new(()).into_parts().0)),
         }
     }
@@ -40,7 +42,9 @@ mod server_fn_impl {
             Self {
                 parts: parts.into(),
                 shared_context: Arc::new(RwLock::new(SendSyncAnyMap::new())),
-                headers: Default::default(),
+                response_parts: std::sync::Arc::new(RwLock::new(
+                    http::response::Response::new(()).into_parts().0,
+                )),
             }
         }
 
@@ -60,35 +64,16 @@ mod server_fn_impl {
                 .map(|_| ())
         }
 
-        /// Get the headers from the server context
-        pub fn response_headers(&self) -> RwLockReadGuard<'_, hyper::header::HeaderMap> {
-            self.try_response_headers()
-                .expect("Failed to get headers from server context")
+        /// Get the response parts from the server context
+        pub fn response_parts(&self) -> LockResult<RwLockReadGuard<'_, http::response::Parts>> {
+            self.response_parts.read()
         }
 
-        /// Try to get the headers from the server context
-        pub fn try_response_headers(
+        /// Get the response parts from the server context
+        pub fn response_parts_mut(
             &self,
-        ) -> LockResult<RwLockReadGuard<'_, hyper::header::HeaderMap>> {
-            self.headers.read()
-        }
-
-        /// Get the headers mutably from the server context
-        pub fn response_headers_mut(&self) -> RwLockWriteGuard<'_, hyper::header::HeaderMap> {
-            self.try_response_headers_mut()
-                .expect("Failed to get headers mutably from server context")
-        }
-
-        /// Try to get the headers mut from the server context
-        pub fn try_response_headers_mut(
-            &self,
-        ) -> LockResult<RwLockWriteGuard<'_, hyper::header::HeaderMap>> {
-            self.headers.write()
-        }
-
-        pub(crate) fn take_response_headers(&self) -> hyper::header::HeaderMap {
-            let mut headers = self.headers.write().unwrap();
-            std::mem::take(&mut *headers)
+        ) -> LockResult<RwLockWriteGuard<'_, http::response::Parts>> {
+            self.response_parts.write()
         }
 
         /// Get the request that triggered:
@@ -124,13 +109,15 @@ std::thread_local! {
 
 /// Get information about the current server request.
 ///
-/// This function will only provide the current server context if it is called from a server function.
+/// This function will only provide the current server context if it is called from a server function or on the server rendering a request.
 pub fn server_context() -> DioxusServerContext {
     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> {
+///
+/// This function will only provide the current server context if it is called from a server function or on the server rendering a request.
+pub async fn extract_server_context<E: FromServerContext<I>, I>() -> Result<E, E::Rejection> {
     E::from_request(&server_context()).await
 }
 
@@ -181,8 +168,8 @@ impl<F: std::future::Future> std::future::Future for ProvideServerContext<F> {
 
 /// A trait for extracting types from the server context
 #[async_trait::async_trait(?Send)]
-pub trait FromServerContext: Sized {
-    /// The error type returned when extraction fails. This type must implement `IntoResponse`.
+pub trait FromServerContext<I = ()>: Sized {
+    /// The error type returned when extraction fails. This type must implement `std::error::Error`.
     type Rejection: std::error::Error;
 
     /// Extract this type from the server context.
@@ -223,48 +210,18 @@ impl<T: Send + Sync + Clone + 'static> FromServerContext for FromContext<T> {
 
 #[cfg(feature = "axum")]
 /// An adapter for axum extractors for the server context
-pub struct Axum<
-    I: axum::extract::FromRequestParts<(), Rejection = R>,
-    R: axum::response::IntoResponse + std::error::Error,
->(pub I, pub std::marker::PhantomData<R>);
-
-#[cfg(feature = "axum")]
-impl<
-        I: axum::extract::FromRequestParts<(), Rejection = R>,
-        R: axum::response::IntoResponse + std::error::Error,
-    > std::ops::Deref for Axum<I, R>
-{
-    type Target = I;
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-#[cfg(feature = "axum")]
-impl<
-        I: axum::extract::FromRequestParts<(), Rejection = R>,
-        R: axum::response::IntoResponse + std::error::Error,
-    > std::ops::DerefMut for Axum<I, R>
-{
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.0
-    }
-}
+pub struct Axum;
 
 #[cfg(feature = "axum")]
 #[async_trait::async_trait(?Send)]
 impl<
         I: axum::extract::FromRequestParts<(), Rejection = R>,
         R: axum::response::IntoResponse + std::error::Error,
-    > FromServerContext for Axum<I, R>
+    > FromServerContext<Axum> for I
 {
     type Rejection = R;
 
     async fn from_request(req: &DioxusServerContext) -> Result<Self, Self::Rejection> {
-        Ok(Self(
-            I::from_request_parts(&mut *req.request_parts_mut().unwrap(), &()).await?,
-            std::marker::PhantomData,
-        ))
+        Ok(I::from_request_parts(&mut *req.request_parts_mut().unwrap(), &()).await?)
     }
 }

packages/fullstack/src/server_fn.rs

@@ -39,6 +39,39 @@ impl ServerFnTraitObj {
 #[cfg(any(feature = "ssr", doc))]
 server_fn::inventory::collect!(ServerFnTraitObj);
 
+#[cfg(any(feature = "ssr", doc))]
+/// Middleware for a server function
+pub struct ServerFnMiddleware {
+    /// The prefix of the server function.
+    pub prefix: &'static str,
+    /// The url of the server function.
+    pub url: &'static str,
+    /// The middleware layers.
+    pub middleware: fn() -> Vec<std::sync::Arc<dyn crate::layer::Layer>>,
+}
+
+#[cfg(any(feature = "ssr", doc))]
+pub(crate) static MIDDLEWARE: once_cell::sync::Lazy<
+    std::collections::HashMap<
+        (&'static str, &'static str),
+        Vec<std::sync::Arc<dyn crate::layer::Layer>>,
+    >,
+> = once_cell::sync::Lazy::new(|| {
+    let mut map: std::collections::HashMap<
+        (&'static str, &'static str),
+        Vec<std::sync::Arc<dyn crate::layer::Layer>>,
+    > = std::collections::HashMap::new();
+    for middleware in server_fn::inventory::iter::<ServerFnMiddleware> {
+        map.entry((middleware.prefix, middleware.url))
+            .or_default()
+            .extend((middleware.middleware)().iter().cloned());
+    }
+    map
+});
+
+#[cfg(any(feature = "ssr", doc))]
+server_fn::inventory::collect!(ServerFnMiddleware);
+
 #[cfg(any(feature = "ssr", doc))]
 /// A server function that can be called on serializable arguments and returns a serializable result.
 pub type ServerFunction = server_fn::SerializedFnTraitObj<()>;