From b90e6ae7449141da5861f1cdf3b0f2faf2cbe7e0 Mon Sep 17 00:00:00 2001
From: Hamcha <hamcha@crunchy.rocks>
Date: Sat, 18 Nov 2023 16:58:28 +0100
Subject: [PATCH] add error middleware

---
 src/http/error.rs | 82 ++++++++++++++++++++++++++++++++---------------
 src/main.rs       |  8 +++--
 src/route/mod.rs  |  2 +-
 3 files changed, 63 insertions(+), 29 deletions(-)

diff --git a/src/http/error.rs b/src/http/error.rs
index 19c31ea..41a20be 100644
--- a/src/http/error.rs
+++ b/src/http/error.rs
@@ -1,19 +1,21 @@
 use askama::Template;
 use axum::{
     extract::rejection::JsonRejection,
-    http::StatusCode,
+    http::{header::ACCEPT, Request, StatusCode},
+    middleware::Next,
     response::{IntoResponse, Response},
     Json,
 };
+use serde_json::json;
 use thiserror::Error;
 
 #[derive(Error, Debug)]
-pub enum AppError<'a> {
+pub enum AppError {
     #[error("client error: <{code}> {message}")]
     Client {
         status: StatusCode,
-        code: &'a str,
-        message: &'a str,
+        code: &'static str,
+        message: &'static str,
     },
 
     #[error("unexpected internal error: {0}")]
@@ -23,34 +25,39 @@ pub enum AppError<'a> {
     JSONFormat(#[from] JsonRejection),
 }
 
-#[derive(Template)]
-#[template(path = "error.html")]
-struct ErrorTemplate<'a> {
-    code: &'a str,
-    message: &'a str,
+struct ErrorInfo {
+    code: String,
+    message: String,
 }
 
-impl IntoResponse for AppError<'_> {
+#[derive(Template)]
+#[template(path = "error.html")]
+struct ErrorTemplate {
+    code: String,
+    message: String,
+}
+
+impl IntoResponse for AppError {
     fn into_response(self) -> Response {
-        match self {
+        let (status, info) = match self {
             AppError::Internal(err) => (
                 StatusCode::INTERNAL_SERVER_ERROR,
-                ErrorTemplate {
-                    code: "server-error",
-                    message: &err.to_string(),
+                ErrorInfo {
+                    code: "server-error".to_string(),
+                    message: err.to_string(),
                 },
-                // Json(json!({"code":"server-error", "message": err.to_string()})),
-            )
-                .into_response(),
+            ),
             AppError::Client {
                 status,
                 code,
                 message,
             } => (
                 status,
-                ErrorTemplate { code, message }, // Json(json!({"code":code, "message": message}))
-            )
-                .into_response(),
+                ErrorInfo {
+                    code: code.to_string(),
+                    message: message.to_string(),
+                },
+            ),
             AppError::JSONFormat(rejection) => {
                 let status = match rejection {
                     JsonRejection::JsonDataError(_) => StatusCode::UNPROCESSABLE_ENTITY,
@@ -60,13 +67,36 @@ impl IntoResponse for AppError<'_> {
                 };
                 (
                     status,
-                    ErrorTemplate {
-                        code: "invalid-body",
-                        message: &rejection.body_text(),
-                    }, // Json(json!({"code":"invalid-body", "message": rejection.body_text()})),
+                    ErrorInfo {
+                        code: "invalid-body".to_string(),
+                        message: rejection.body_text().clone(),
+                    },
                 )
-                    .into_response()
             }
-        }
+        };
+
+        let mut response = status.into_response();
+        response.extensions_mut().insert(info);
+        response
     }
 }
+
+pub async fn error_interceptor<B>(request: Request<B>, next: Next<B>) -> Response {
+    let accept_header = request
+        .headers()
+        .get(&ACCEPT)
+        .map(|value| value.as_ref().to_owned());
+
+    let mut response = next.run(request).await;
+
+    if let Some(ErrorInfo { code, message }) = response.extensions_mut().remove::<ErrorInfo>() {
+        match accept_header.as_deref() {
+            Some(b"application/json") => {
+                return Json(json!({"code": code, "message": message})).into_response()
+            }
+            _ => return ErrorTemplate { code, message }.into_response(),
+        }
+    }
+
+    response
+}
diff --git a/src/main.rs b/src/main.rs
index e97f75b..b9bca3b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,9 +1,11 @@
 use anyhow::Result;
-use axum::{routing::get, Router};
+use axum::middleware::from_fn;
 use bollard::Docker;
 use clap::Parser;
 use std::net::SocketAddr;
 
+use crate::http::error::error_interceptor;
+
 mod http;
 mod route;
 mod stack;
@@ -48,7 +50,9 @@ async fn main() -> Result<()> {
         docker,
     };
 
-    let app = route::router().with_state(state);
+    let app = route::router()
+        .with_state(state)
+        .layer(from_fn(error_interceptor));
 
     axum::Server::bind(&args.bind)
         .serve(app.into_make_service())
diff --git a/src/route/mod.rs b/src/route/mod.rs
index 83745ea..29cd367 100644
--- a/src/route/mod.rs
+++ b/src/route/mod.rs
@@ -9,7 +9,7 @@ struct HomeTemplate {
     stacks: Vec<String>,
 }
 
-async fn home(State(state): State<AppState>) -> Result<HomeTemplate, AppError<'static>> {
+async fn home(State(state): State<AppState>) -> Result<HomeTemplate, AppError> {
     let list = stack::list(&state.stack_dir)
         .await
         .map_err(AppError::from)?;