add server function middleware

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())?
+                }
+            })
+        })
+    }
+}