ayun-server 0.23.0

use crate::{traits::MiddlewareTrait, Middleware};
use axum::{error_handling::HandleErrorLayer, http::StatusCode, response::IntoResponse};
use ayun_config::config;
use std::time::Duration;
use tower::ServiceBuilder;

pub struct Timeout;

impl MiddlewareTrait for Timeout {
    fn handle() -> Middleware {
        Box::new(|router| {
            let server = config::<crate::config::Server>("server")?;

            Ok(router.layer(
                ServiceBuilder::new()
                    .layer(HandleErrorLayer::new(handle_error))
                    .timeout(Duration::from_millis(server.middleware.timeout)),
            ))
        })
    }
}

async fn handle_error(error: axum::BoxError) -> axum::response::Response {
    let mut status = StatusCode::INTERNAL_SERVER_ERROR;
    let message = format!("Unhandled internal error: {error}");

    if error.is::<tower::timeout::error::Elapsed>() {
        status = StatusCode::REQUEST_TIMEOUT;
    } else if error.is::<axum::extract::rejection::LengthLimitError>() {
        status = StatusCode::PAYLOAD_TOO_LARGE;
    }

    tracing::error!(error = ?message, "error");

    #[cfg(not(feature = "response-json"))]
    {
        (status, message).into_response()
    }

    #[cfg(feature = "response-json")]
    {
        crate::response::json()
            .code(status)
            .message(message)
            .into_response()
    }
}