use crate::app::context::AppContext;
use crate::service::http::middleware::Middleware;
use axum::Router;
use axum::http::StatusCode;
use axum_core::extract::FromRef;
use serde_derive::{Deserialize, Serialize};
use serde_with::serde_as;
use std::time::Duration;
use tower_http::timeout::TimeoutLayer;
use validator::Validate;
#[serde_as]
#[derive(Debug, Clone, Validate, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
#[non_exhaustive]
pub struct TimeoutConfig {
#[serde_as(as = "serde_with::DurationMilliSeconds")]
pub timeout: Duration,
#[serde(with = "http_serde::status_code")]
pub status_code: StatusCode,
}
pub struct TimeoutMiddleware;
impl<S> Middleware<S> for TimeoutMiddleware
where
S: 'static + Send + Sync + Clone,
AppContext: FromRef<S>,
{
type Error = crate::error::Error;
fn name(&self) -> String {
"timeout".to_string()
}
fn enabled(&self, state: &S) -> bool {
AppContext::from_ref(state)
.config()
.service
.http
.custom
.middleware
.timeout
.common
.enabled(state)
}
fn priority(&self, state: &S) -> i32 {
AppContext::from_ref(state)
.config()
.service
.http
.custom
.middleware
.timeout
.common
.priority
}
fn install(&self, state: &S, router: Router) -> Result<Router, Self::Error> {
let context = AppContext::from_ref(state);
let config = &context
.config()
.service
.http
.custom
.middleware
.timeout
.custom;
let status_code = config.status_code;
let timeout = config.timeout;
let router = router.layer(TimeoutLayer::with_status_code(status_code, timeout));
Ok(router)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::AppConfig;
use rstest::rstest;
#[rstest]
#[case(false, Some(true), true)]
#[case(false, Some(false), false)]
#[cfg_attr(coverage_nightly, coverage(off))]
fn timeout_enabled(
#[case] default_enable: bool,
#[case] enable: Option<bool>,
#[case] expected_enabled: bool,
) {
let mut config = AppConfig::test(None).unwrap();
config.service.http.custom.middleware.default_enable = default_enable;
config.service.http.custom.middleware.timeout.common.enable = enable;
let context = AppContext::test(Some(config), None, None).unwrap();
let middleware = TimeoutMiddleware;
assert_eq!(middleware.enabled(&context), expected_enabled);
}
#[rstest]
#[case(None, 0)]
#[case(Some(1234), 1234)]
#[cfg_attr(coverage_nightly, coverage(off))]
fn timeout_priority(#[case] override_priority: Option<i32>, #[case] expected_priority: i32) {
let mut config = AppConfig::test(None).unwrap();
if let Some(priority) = override_priority {
config
.service
.http
.custom
.middleware
.timeout
.common
.priority = priority;
}
let context = AppContext::test(Some(config), None, None).unwrap();
let middleware = TimeoutMiddleware;
assert_eq!(middleware.priority(&context), expected_priority);
}
}