Middleware isn't running when called through `use_server_future`

[patrik-cihal]

Problem

When invoking middleware through use_server_future call, the middlware implementation cannot access the server context. This includes potential database headers etc...

• Dioxus version: 0.6.3
• App platform: web

[ealmloff]

Can you share a reproduction for this issue? I cannot reproduce the issue with this code: https://github.com/ealmloff/fail-repro-4117

[patrik-cihal]

That link doesn't lead anywhere.

[patrik-cihal]

This is the middleware that doesn't work. It doesn't find the cookie. And when I had DB in server context it didn't find it. Running the middleware manually in the server function logic, resolves the issue.

use crate::{DB, prelude::*};
use std::convert::Infallible;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use tower::{Layer, Service};

#[derive(Clone)]
pub struct AuthLayer;

impl<S> Layer<S> for AuthLayer {
    type Service = AuthMiddleware<S>;

    fn layer(&self, service: S) -> Self::Service {
        AuthMiddleware { inner: service }
    }
}

#[derive(Clone)]
pub struct AuthMiddleware<S> {
    inner: S,
}

impl<S, Request> Service<Request> for AuthMiddleware<S>
where
    S: Service<Request, Error = Infallible>,
    S::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = Infallible;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        let future = self.inner.call(request);
        Box::pin(async move {
            let server_context = server_context();
            let headers: axum::http::HeaderMap = server_context.extract().await.unwrap();
            let conn = &mut DB.get().await.unwrap();

            info!("TEST {headers:?}");

            // Extract session cookie
            let session = if let Some(cookie) = headers.get("Cookie") {
                let cookie_str = cookie.to_str().unwrap_or("");

                if let Some(session_token) = cookie_str
                    .split(';')
                    .find(|s| s.trim().starts_with("session="))
                    .and_then(|s| s.trim().strip_prefix("session="))
                {
                    // Hash token for comparison
                    let token_hash = sha256::digest(session_token.as_bytes());

                    // Find and validate session
                    let session = Session::by_token(conn, token_hash).await.unwrap();

                    session
                } else {
                    None
                }
            } else {
                warn!("No cookie found :/");
                None
            };

            server_context.insert(session);
            future.await
        })
    }
}

[ealmloff]

Link should be fixed. My git token expired

[patrik-cihal]

Btw I don't expect you to necessarily attempt to fix it when I don't provide a reproducible example... It's mostly so that if someone has the same issue then they have a place to report it.

I think it's better to report it in a non-complete form than not do that at all. Maybe that could be explicitly stated in the issue boilerplate. Something like "reproducible example is not mandatory but we might not focus on it or fix it too late".

Nevertheless I still think there is a bug somewhere and I'll provide a reproducible example tomorrow.

[patrik-cihal]

So I guess true is equal to false... This doesn't fail either... Seems like middleware doens't throw panics.

#![allow(non_snake_case)]
use dioxus::prelude::*;

fn main() {
    dioxus::LaunchBuilder::new()
        .with_context(1234u32)
        .launch(app);
}

fn app() -> Element {
    let result = use_server_future(get_server_data)?;

    rsx! {
        "{result:?}"
    }
}

#[cfg(feature = "server")]
async fn assert_server_context_provided() {
    assert_eq!(true, false);
    let FromContext(i): FromContext<u32> = extract().await.unwrap();
    assert_eq!(i, 1234u32);
}

#[cfg(feature = "server")]
struct ServerContextProvidedService<S> {
    inner: S,
}

#[cfg(feature = "server")]
impl<S: tower::Service<R>, R> tower::Service<R> for ServerContextProvidedService<S>
where
    S::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
    >;

    fn call(&mut self, req: R) -> Self::Future {
        let fut = self.inner.call(req);
        Box::pin(async move {
            assert_server_context_provided().await;
            fut.await
        })
    }

    fn poll_ready(
        &mut self,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }
}

#[cfg(feature = "server")]
fn assert_server_context_provided_layer_fn<
    S: tower::Service<axum::extract::Request<axum::body::Body>>,
>(
    service: S,
) -> ServerContextProvidedService<S> {
    ServerContextProvidedService { inner: service }
}

#[server(GetServerData)]
#[middleware(tower::layer::layer_fn(assert_server_context_provided_layer_fn))]
async fn get_server_data() -> Result<String, ServerFnError> {
    // assert_server_context_provided().await;
    Ok("Hello from the server!".to_string())
}

[patrik-cihal]

Or that middleware isn't even running.

[patrik-cihal]

Yep your middleware isn't even running.

[patrik-cihal]

I see the issue is that the middleware isn't running when called through use_server_future...

use dioxus::logger::tracing::*;
use dioxus::prelude::*;

static mut DID_RUN: bool = false;

fn main() {
    dioxus::LaunchBuilder::new()
        .with_context(1234u32)
        .launch(app);
}

fn app() -> Element {
    let result = use_server_future(get_server_data)?; // replacing this with use_resource fixes it

    rsx! {
        "{result:?}"
    }
}

#[cfg(feature = "server")]
mod middleware {
    use super::*;
    use dioxus::prelude::FromContext;
    use std::future::Future;
    use std::pin::Pin;
    use std::task::{Context, Poll};
    use tower::{Layer, Service};

    #[derive(Clone)]
    pub struct AuthLayer;

    impl<S> Layer<S> for AuthLayer {
        type Service = AuthMiddleware<S>;

        fn layer(&self, service: S) -> Self::Service {
            AuthMiddleware { inner: service }
        }
    }

    #[derive(Clone)]
    pub struct AuthMiddleware<S> {
        inner: S,
    }

    impl<S, Request> Service<Request> for AuthMiddleware<S>
    where
        S: Service<Request, Error = ServerFnError>,
        S::Future: Send + 'static,
    {
        type Response = S::Response;
        type Error = ServerFnError;
        type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

        fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            self.inner.poll_ready(cx)
        }

        fn call(&mut self, request: Request) -> Self::Future {
            let future = self.inner.call(request);
            Box::pin(async move {
                let FromContext::<u32>(test) = extract().await.unwrap();

                unsafe { DID_RUN = true };

                assert_eq!(test, 1234u32);

                future.await
            })
        }
    }
}

#[server]
#[middleware(middleware::AuthLayer {})]
async fn get_server_data() -> Result<String, ServerFnError> {
    // assert_server_context_provided().await;
    let did_run = unsafe { DID_RUN };
    info!("{did_run}");
    Ok("Hello from the server!".to_string())
}

[ealmloff]

Server functions run in two different places:

1. If server functions are called during the initial server-side render or during streaming, they are called like a normal async function with the server context. Middleware does not run, and the HTTP response for the HTML stream has already been started, so response headers cannot be modified
2. If server functions are called from a client, they get their own request. Middleware run,s and the HTTP response can be modified

The two different modes of execution are pretty standard across metaframeworks. I'm unsure how they interact with the HTTP request/response.

[patrik-cihal]

I see. But when added as an axum middleware, not associated with individual server functions it should work. But it doesn't.

[lanlin]

Does this mean that custom middleware like the following will only work for manually registered methods, but not for methods automatically registered through #[server]?

    let router = axum::Router::new()
        .serve_dioxus_application(
            ServeConfig::new().unwrap(),
            || rsx!{}
        )
    .route("/the/custom/api", axum::routing::get(custom_api_handler))
    .layer(middleware::from_fn(my_middleware));

[ealmloff]

Does this mean that custom middleware like the following will only work for manually registered methods, but not for methods automatically registered through #[server]?

let router = axum::Router::new()
    .serve_dioxus_application(
        ServeConfig::new().unwrap(),
        || rsx!{}
    )
.route("/the/custom/api", axum::routing::get(custom_api_handler))
.layer(middleware::from_fn(my_middleware));

When called from the server, server functions act like a normal function and don't have a http request. Because of that, they will never go through your axum router at all or run the layer. If the function is called from the client, it will run the layer

[acheronte]

Just confirming what @ealmloff said.

I have split my project into frontend and backend crates, so can't speak of server functions executed from the client, but server functions, called from the server, were working, yet the middleware layer wasn't being applied to them.

I didn't like having middleware-level code on individual server functions so I got rid of server functions altogether and now using pure Axum handlers, and the middleware layer works fine on these.