Add context to axum Middleware

[zhuravkov]

Help me solve the issue:How does it reset the context in Middleware? I'm trying both through Extension and through use_context, but it's still not there.

use axum::{
    extract::{Request, Extension},

    http::StatusCode,
    middleware::Next,
    response::{Response, IntoResponse, Redirect},
};
use crate::extractor::CurrentUser;
use crate::service::AuthService;
use leptos::prelude::use_context;


pub async fn auth_required(
    auth_service: Option<Extension<AuthService>>,
    req: Request,
    next: Next,
) -> impl IntoResponse {

    let uri = req.uri();
    let path = uri.path();

    let protected_prefixes = vec![
        "/users",
        "/api/protected",
    ];
    
    // Check if path starts with protected prefix
    let is_protected = protected_prefixes.iter().any(|prefix| path.starts_with(prefix));
    let mut auth_error:Option<String> = None;
    
    if !is_protected {
        println!("Skipping public route: {}", path);
        return next.run(req).await;
    }

    println!("Protected route: {}", path);
    
    // Extract token from cookies
    let access_token = extract_token_from_cookies(&req);
    if access_token.is_err() {
        println!("No token in cookies, redirecting to /auth");
        auth_error = Some("No token in cookies".to_string());
    }
    // Get AuthService from Leptos context
    if auth_service.is_none() { 
        println!("AuthService not found in Extension, redirecting to /auth");
        auth_error = Some("AuthService not found in Extension".to_string());
    }
    else {
        let auth_service = auth_service.unwrap();
    }
    let auth_service_context = use_context::<AuthService>();
    if auth_service_context.is_none() {
        println!("AuthService not found in Context, redirecting to /auth");
        auth_error = Some("AuthService not found in Context".to_string());
    }
    else {
        let auth_service_context = auth_service_context.unwrap();
    }

    let access_token = access_token.unwrap_or_default();
    
    // Verify token through AuthService
    // match auth_service.verify_access_token(&access_token).await {
    //     Ok((user_id, token_version)) => {
    //         println!("Token valid for user: {}", user_id);
    //     }
    //     Err(e) => {
    //         println!("Token verification error: {:?}, redirecting to /auth", e);
    //         auth_error = Some("Token verification error".to_string());
    //     }
    // }

    if let Some(error) = auth_error {
        if path.starts_with("/api") {
            return Response::builder()
            .status(StatusCode::FORBIDDEN)
            .body(axum::body::Body::from(error))
            .unwrap();

        }
        println!("Redirecting to /auth");
        return Redirect::to("/auth").into_response();
    }

    return next.run(req).await;
}


// Helper function to extract token from cookies
fn extract_token_from_cookies(req: &Request) -> anyhow::Result<String> {
    let cookies = req.headers()
        .get(axum::http::header::COOKIE)
        .and_then(|v| v.to_str().ok())
        .ok_or_else(|| anyhow::anyhow!("no cookies"))?;
    
    let access_token = cookies.split(';')
        .find_map(|c| {
            let c = c.trim();
            c.strip_prefix("access_token=").map(|v| v.to_string())
        })
        .ok_or_else(|| anyhow::anyhow!("no access token"))?;
    
    Ok(access_token)
}

Part of main.rs:

    let refresh_repo = RefreshRepo::new(db.clone());
    let user_repo = UserRepository::new(db.clone());
    let auth_service = AuthService {
        jwt: jwt_config.clone(),
        refresh_repo: refresh_repo.clone(),
        users: user_repo.clone(),
    };
    
    
    let routes = generate_route_list(App);

    let app = Router::new()
        .route("/img/{*path}", get(image_handler))
        .layer(axum::extract::Extension(auth_service.clone())) 
        .leptos_routes_with_context(
            &leptos_options,
            routes,
            move || {
                provide_context(db.clone());
                provide_context(jwt_config.clone());
                provide_context(refresh_repo.clone());
                provide_context(user_repo.clone());
                provide_context(auth_service.clone());
            },
            {
                let leptos_options = leptos_options.clone();
                move || shell(leptos_options.clone())
            },
        )
        .layer(middleware::from_fn(images::images::deny_images))
        .layer(middleware::from_fn(auth_middleware::auth_required)) 
        .fallback(leptos_axum::file_and_error_handler(shell))
        .with_state(leptos_options);

    // let app = auth_routes.merge(leptos_routes);

    log!("listening on http://{}", &addr);
    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
    axum::serve(listener, app.into_make_service())
        .await
        .unwrap();

Log:
Protected route: /api/protected/get_users15499627470196328238
AuthService not found in Extension, redirecting to /auth
AuthService not found in Context, redirecting to /auth
Protected route: /users
AuthService not found in Extension, redirecting to /auth
AuthService not found in Context, redirecting to /auth
Redirecting to /auth