3 files changed, 109 insertions, 22 deletions
diff --git a/src/setup/middleware.rs b/src/setup/middleware.rs
deleted file mode 100644
index 5f9996b..0000000
--- a/src/setup/middleware.rs
+++ /dev/null
@@ -1,20 +0,0 @@
-use axum::{
-    extract::{Request, State},
-    http::StatusCode,
-    middleware::Next,
-    response::{IntoResponse, Response},
-};
-
-use crate::{app::App, error::Internal};
-
-pub async fn setup_required(State(app): State<App>, request: Request, next: Next) -> Response {
-    match app.setup().completed().await {
-        Ok(true) => next.run(request).await,
-        Ok(false) => (
-            StatusCode::SERVICE_UNAVAILABLE,
-            "initial setup not completed",
-        )
-            .into_response(),
-        Err(error) => Internal::from(error).into_response(),
-    }
-}
diff --git a/src/setup/mod.rs b/src/setup/mod.rs
index 5a8fa37..a4b821c 100644
--- a/src/setup/mod.rs
+++ b/src/setup/mod.rs
@@ -1,6 +1,6 @@
 pub mod app;
-pub mod middleware;
 pub mod repo;
+mod required;
 mod routes;
 
-pub use self::routes::router;
+pub use self::{required::Required, routes::router};
diff --git a/src/setup/required.rs b/src/setup/required.rs
new file mode 100644
index 0000000..2112e4b
--- /dev/null
+++ b/src/setup/required.rs
@@ -0,0 +1,107 @@
+use axum::{
+    extract::Request,
+    http::StatusCode,
+    response::{IntoResponse, Response},
+};
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tower::{Layer, Service};
+
+use crate::{app::App, error::Internal};
+
+#[derive(Clone)]
+pub struct Required(pub App);
+
+impl Required {
+    pub fn with_fallback<F>(self, fallback: F) -> WithFallback<F> {
+        let Self(app) = self;
+        WithFallback { app, fallback }
+    }
+}
+
+impl<S> Layer<S> for Required {
+    type Service = Middleware<S, Unavailable>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        let Self(app) = self.clone();
+        Middleware {
+            inner,
+            app,
+            fallback: Unavailable,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct WithFallback<F> {
+    app: App,
+    fallback: F,
+}
+
+impl<S, F> Layer<S> for WithFallback<F>
+where
+    Self: Clone,
+{
+    type Service = Middleware<S, F>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        let Self { app, fallback } = self.clone();
+        Middleware {
+            inner,
+            app,
+            fallback,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct Middleware<S, F> {
+    inner: S,
+    app: App,
+    fallback: F,
+}
+
+impl<S, F> Service<Request> for Middleware<S, F>
+where
+    Self: Clone,
+    S: Service<Request, Response = Response> + Send + 'static,
+    S::Future: Send,
+    F: IntoResponse + Clone + Send + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+
+    fn poll_ready(&mut self, ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        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()),
+            }
+        })
+    }
+}
+
+#[derive(Clone)]
+pub struct Unavailable;
+
+impl IntoResponse for Unavailable {
+    fn into_response(self) -> Response {
+        (
+            StatusCode::SERVICE_UNAVAILABLE,
+            "initial setup not completed",
+        )
+            .into_response()
+    }
+}