diff --git a/src/lib.rs b/src/lib.rs index cf1f7769..23ccbad5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -236,16 +236,25 @@ pub async fn bootstap(mut shutdown: broadcast::Receiver<()>, config: Config) -> ); let app = Router::new() - .route("/health", get(handlers::health::handler)) - .nest("/tenants", tenancy_routes) + .route("/health", get(handlers::health::handler).layer( + axum::middleware::from_fn_with_state(state_arc.clone(), rate_limit_middleware), + )) + .nest("/tenants", tenancy_routes.layer( + axum::middleware::from_fn_with_state(state_arc.clone(), rate_limit_middleware), + )) .route( "/:tenant_id/clients", - post(handlers::register_client::handler), + post(handlers::register_client::handler).layer( + axum::middleware::from_fn_with_state(state_arc.clone(), rate_limit_middleware), + ), ) .route( "/:tenant_id/clients/:id", - delete(handlers::delete_client::handler), + delete(handlers::delete_client::handler).layer( + axum::middleware::from_fn_with_state(state_arc.clone(), rate_limit_middleware), + ), ) + // Rate limiting middleware is not applying to push_handler because it is used by the relay .route( "/:tenant_id/clients/:id", post(handlers::push_message::handler), @@ -266,15 +275,22 @@ pub async fn bootstap(mut shutdown: broadcast::Receiver<()>, config: Config) -> #[cfg(not(feature = "multitenant"))] let app = Router::new() - .route("/health", get(handlers::health::handler)) + .route("/health", get(handlers::health::handler).layer( + axum::middleware::from_fn_with_state(state_arc.clone(), rate_limit_middleware), + )) .route( "/clients", - post(handlers::single_tenant_wrappers::register_handler), + post(handlers::single_tenant_wrappers::register_handler).layer( + axum::middleware::from_fn_with_state(state_arc.clone(), rate_limit_middleware), + ), ) .route( "/clients/:id", - delete(handlers::single_tenant_wrappers::delete_handler), + delete(handlers::single_tenant_wrappers::delete_handler).layer( + axum::middleware::from_fn_with_state(state_arc.clone(), rate_limit_middleware), + ), ) + // Rate limiting middleware is not applying to push_handler because it is used by the relay .route( "/clients/:id", post(handlers::single_tenant_wrappers::push_handler), @@ -285,10 +301,6 @@ pub async fn bootstap(mut shutdown: broadcast::Receiver<()>, config: Config) -> } else { app }; - let app = app.route_layer(axum::middleware::from_fn_with_state( - state_arc.clone(), - rate_limit_middleware, - )); let app = app.with_state(state_arc.clone()); let private_app = Router::new()