use axum::{ extract::Request, http::StatusCode, response::{IntoResponse, Response}, }; use std::pin::Pin; use std::task::{Context, Poll}; use tower::Service; use crate::{app::App, error::Internal}; const UNAVAILABLE: (StatusCode, &str) = ( StatusCode::SERVICE_UNAVAILABLE, "initial setup not completed", ); #[derive(Clone)] pub struct Layer { app: App, fallback: F, } impl Layer<(StatusCode, &'static str)> { pub fn or_unavailable(app: App) -> Self { Self::with_fallback(app, UNAVAILABLE) } } impl Layer { pub fn with_fallback(app: App, fallback: F) -> Self { Layer { app, fallback } } } impl tower::Layer for Layer where Self: Clone, { type Service = Middleware; fn layer(&self, inner: S) -> Self::Service { let Self { app, fallback } = self.clone(); Middleware { inner, app, fallback, } } } #[derive(Clone)] pub struct Middleware { inner: S, app: App, fallback: F, } impl Service for Middleware where Self: Clone, S: Service + Send + 'static, S::Future: Send, F: IntoResponse + Clone + Send + 'static, { type Response = S::Response; type Error = S::Error; type Future = Pin> + Send>>; fn poll_ready(&mut self, ctx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(ctx) } fn call(&mut self, req: Request) -> Self::Future { let Self { mut inner, app, fallback, } = self.clone(); Box::pin(async move { match app.setup().completed().await { Ok(true) => inner.call(req).await, Ok(false) => Ok(fallback.into_response()), Err(error) => Ok(Internal::from(error).into_response()), } }) } }