use axum::Json;
use axum::http::{HeaderValue, StatusCode, header};
use axum::response::{IntoResponse, Response};
use serde::{Deserialize, Serialize};
#[derive(Debug, Default, Serialize, Deserialize, Clone)]
struct ErrorData {
#[serde(skip_serializing_if = "Option::is_none")]
sub_category: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
code: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
exception: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
extra: Option<Vec<String>>,
#[serde(skip)]
redact: Option<bool>,
}
#[derive(Serialize)]
struct ErrorSerialize<'a> {
#[serde(skip_serializing_if = "str::is_empty")]
category: &'a str,
message: &'a str,
#[serde(flatten)]
data: &'a ErrorData,
}
#[derive(Deserialize)]
struct ErrorDeserialize {
#[serde(default)]
category: String,
#[serde(default)]
message: String,
#[serde(flatten)]
data: ErrorData,
}
#[derive(Debug, Clone, Default)]
pub struct Error {
status: u16,
category: String,
message: String,
data: Box<ErrorData>,
}
impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.message)
}
}
impl std::error::Error for Error {}
impl Serialize for Error {
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
ErrorSerialize {
category: &self.category,
message: &self.message,
data: &self.data,
}
.serialize(serializer)
}
}
impl<'de> Deserialize<'de> for Error {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let d = ErrorDeserialize::deserialize(deserializer)?;
Ok(Self {
status: 0,
category: d.category,
message: d.message,
data: Box::new(d.data),
})
}
}
impl Error {
#[must_use]
pub fn new(message: impl ToString) -> Self {
Self {
message: message.to_string(),
..Default::default()
}
}
#[must_use]
pub fn with_category(mut self, category: impl Into<String>) -> Self {
self.category = category.into();
self
}
#[must_use]
pub fn with_sub_category(mut self, sub_category: impl Into<String>) -> Self {
self.data.sub_category = Some(sub_category.into());
self
}
#[must_use]
pub fn with_code(mut self, code: impl Into<String>) -> Self {
self.data.code = Some(code.into());
self
}
#[must_use]
pub fn with_status(mut self, status: u16) -> Self {
self.status = status;
self
}
#[must_use]
pub fn with_exception(mut self, exception: bool) -> Self {
self.data.exception = Some(exception);
self
}
#[must_use]
pub fn with_redact(mut self, redact: bool) -> Self {
self.data.redact = Some(redact);
self
}
#[must_use]
pub fn add_extra(mut self, value: impl Into<String>) -> Self {
self.data
.extra
.get_or_insert_with(Vec::new)
.push(value.into());
self
}
pub fn status(&self) -> u16 {
self.status
}
pub fn category(&self) -> &str {
&self.category
}
pub fn message(&self) -> &str {
&self.message
}
pub fn sub_category(&self) -> Option<&str> {
self.data.sub_category.as_deref()
}
pub fn code(&self) -> Option<&str> {
self.data.code.as_deref()
}
pub fn is_exception(&self) -> bool {
self.data.exception.unwrap_or(false)
}
pub fn redact(&self) -> Option<bool> {
self.data.redact
}
pub fn extra(&self) -> &[String] {
self.data.extra.as_deref().unwrap_or(&[])
}
#[must_use]
pub fn from_upstream(status: u16, body: &[u8]) -> Self {
match serde_json::from_slice::<Self>(body) {
Ok(err) if !err.message.is_empty() => err.with_status(status),
_ => Self::new(truncate_upstream_body(body)).with_status(status),
}
}
}
pub const MAX_UPSTREAM_MESSAGE_LEN: usize = 512;
fn truncate_upstream_body(body: &[u8]) -> String {
if body.is_empty() {
return "upstream returned an empty body".to_string();
}
let text = String::from_utf8_lossy(body);
let trimmed = text.trim();
if trimmed.len() <= MAX_UPSTREAM_MESSAGE_LEN {
return trimmed.to_string();
}
let mut end = MAX_UPSTREAM_MESSAGE_LEN;
while end > 0 && !trimmed.is_char_boundary(end) {
end -= 1;
}
format!("{}…", &trimmed[..end])
}
fn response_status(status: u16) -> StatusCode {
match StatusCode::from_u16(status) {
Ok(code) if code.is_client_error() || code.is_server_error() => code,
_ => StatusCode::INTERNAL_SERVER_ERROR,
}
}
fn redacted_message(status: StatusCode) -> String {
match status.canonical_reason() {
Some(reason) => reason.to_ascii_lowercase(),
None if status.is_client_error() => "client error".to_string(),
None => "internal server error".to_string(),
}
}
fn should_redact(status: StatusCode, explicit: Option<bool>) -> bool {
explicit.unwrap_or_else(|| status.is_server_error())
}
impl IntoResponse for Error {
fn into_response(self) -> Response {
let status = response_status(self.status);
let mut res = if should_redact(status, self.data.redact) {
let redacted = ErrorData {
sub_category: self.data.sub_category.clone(),
code: self.data.code.clone(),
exception: None,
extra: None,
redact: None,
};
let message = redacted_message(status);
(
status,
Json(ErrorSerialize {
category: &self.category,
message: &message,
data: &redacted,
}),
)
.into_response()
} else {
(status, Json(&self)).into_response()
};
res.extensions_mut().insert(self);
res.headers_mut()
.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
res
}
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
#[test]
fn empty_category_is_omitted_and_round_trips() {
let json = serde_json::to_string(&Error::new("m")).expect("序列化");
assert_eq!(json, r#"{"message":"m"}"#);
let back: Error = serde_json::from_str(&json).expect("反序列化");
assert_eq!(back.category(), "");
assert_eq!(back.message(), "m");
}
async fn body_json(res: Response) -> serde_json::Value {
let bytes = axum::body::to_bytes(res.into_body(), usize::MAX)
.await
.expect("读取响应体");
serde_json::from_slice(&bytes).expect("响应体应是合法 JSON")
}
#[test]
fn redact_defaults_to_status_class_only() {
assert!(should_redact(StatusCode::INTERNAL_SERVER_ERROR, None));
assert!(should_redact(StatusCode::SERVICE_UNAVAILABLE, None));
assert!(!should_redact(StatusCode::BAD_REQUEST, None));
assert!(!should_redact(StatusCode::NOT_FOUND, None));
assert!(!should_redact(StatusCode::OK, None));
assert!(should_redact(StatusCode::BAD_REQUEST, Some(true)));
assert!(!should_redact(
StatusCode::INTERNAL_SERVER_ERROR,
Some(false)
));
}
#[tokio::test]
async fn exception_on_4xx_keeps_client_facing_message() {
let err = Error::new("email format invalid")
.with_category("params")
.with_status(400)
.with_exception(true);
let res = err.into_response();
assert_eq!(res.status(), StatusCode::BAD_REQUEST);
let body = body_json(res).await;
assert_eq!(body["message"], "email format invalid");
assert_eq!(body["category"], "params");
}
#[tokio::test]
async fn server_error_message_is_redacted() {
let err = Error::new("relation \"users\" does not exist")
.with_category("db")
.with_sub_category("sqlx")
.with_status(500)
.add_extra("connection=primary");
let res = err.into_response();
let body = body_json(res).await;
assert_eq!(body["message"], "internal server error");
assert_eq!(body["category"], "db");
assert_eq!(body["sub_category"], "sqlx");
assert!(body["extra"].is_null());
assert!(!body.to_string().contains("does not exist"));
}
#[tokio::test]
async fn unset_status_falls_back_to_redacted_500() {
let err = Error::new("raw internal detail").with_category("cache");
let res = err.into_response();
assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR);
assert_eq!(body_json(res).await["message"], "internal server error");
}
#[tokio::test]
async fn explicit_redact_overrides_both_directions() {
let res = Error::new("internal policy rule #42 denied")
.with_status(403)
.with_redact(true)
.into_response();
assert_eq!(res.status(), StatusCode::FORBIDDEN);
assert_eq!(body_json(res).await["message"], "forbidden");
let res = Error::new("upstream is warming up, retry later")
.with_status(503)
.with_redact(false)
.into_response();
assert_eq!(
body_json(res).await["message"],
"upstream is warming up, retry later"
);
}
#[test]
fn full_error_is_preserved_in_extensions() {
let res = Error::new("secret detail").with_status(500).into_response();
let stored = res
.extensions()
.get::<Error>()
.expect("Error 应存入 extensions");
assert_eq!(stored.message(), "secret detail");
}
#[test]
fn no_store_header_is_always_set() {
let res = Error::new("x").with_status(400).into_response();
assert_eq!(
res.headers().get(header::CACHE_CONTROL).unwrap(),
"no-store"
);
}
#[test]
fn non_error_status_falls_back_to_500() {
for bogus in [0, 99, 100, 200, 204, 302, 399, 600, 999] {
let res = Error::new("x").with_status(bogus).into_response();
assert_eq!(
res.status(),
StatusCode::INTERNAL_SERVER_ERROR,
"status={bogus} 应回退为 500"
);
}
for ok in [400, 404, 418, 429, 499, 500, 502, 503, 599] {
let res = Error::new("x").with_status(ok).into_response();
assert_eq!(res.status().as_u16(), ok);
}
}
#[test]
fn redacted_message_matches_status_class() {
assert_eq!(
redacted_message(StatusCode::INTERNAL_SERVER_ERROR),
"internal server error"
);
assert_eq!(
redacted_message(StatusCode::SERVICE_UNAVAILABLE),
"service unavailable"
);
assert_eq!(redacted_message(StatusCode::FORBIDDEN), "forbidden");
let unknown_4xx = StatusCode::from_u16(499).expect("合法状态码");
assert_eq!(redacted_message(unknown_4xx), "client error");
let unknown_5xx = StatusCode::from_u16(599).expect("合法状态码");
assert_eq!(redacted_message(unknown_5xx), "internal server error");
}
#[test]
fn redact_flag_is_not_serialized() {
let json = serde_json::to_string(&Error::new("m").with_redact(true)).unwrap();
assert!(!json.contains("redact"), "序列化结果不应含 redact: {json}");
}
#[tokio::test]
async fn unset_optional_fields_are_omitted() {
let res = Error::new("bad email")
.with_category("params")
.with_status(400)
.into_response();
let body = body_json(res).await;
let obj = body.as_object().expect("响应体应是 JSON 对象");
assert_eq!(obj.len(), 2, "只应剩 category / message,实际: {obj:?}");
assert_eq!(obj["category"], "params");
assert!(!obj.contains_key("code"));
assert!(!obj.contains_key("exception"));
assert!(!obj.contains_key("extra"));
assert!(!obj.contains_key("sub_category"));
}
#[tokio::test]
async fn redacted_response_hides_exception_flag() {
let res = Error::new("sqlx: connection refused")
.with_category("db")
.with_status(500)
.with_exception(true)
.into_response();
let body = body_json(res).await;
assert_eq!(body["message"], "internal server error");
assert!(
body.get("exception").is_none(),
"脱敏响应仍暴露了告警分级: {body}"
);
}
#[test]
fn redacted_response_still_carries_exception_in_extensions() {
let res = Error::new("boom")
.with_status(500)
.with_exception(true)
.into_response();
let stored = res
.extensions()
.get::<Error>()
.expect("Error 应存入 extensions");
assert!(stored.is_exception());
}
#[test]
fn from_upstream_restores_structured_fields() {
let upstream = Error::new("account is locked")
.with_category("user")
.with_sub_category("locked")
.with_code("E4031")
.add_extra("retry_after=600")
.with_status(403);
let body = serde_json::to_vec(&upstream).expect("序列化上游错误");
let restored = Error::from_upstream(403, &body);
assert_eq!(restored.message(), "account is locked");
assert_eq!(restored.category(), "user");
assert_eq!(restored.sub_category(), Some("locked"));
assert_eq!(restored.code(), Some("E4031"));
assert_eq!(restored.extra(), ["retry_after=600".to_string()]);
assert_eq!(restored.status(), 403);
}
#[test]
fn from_upstream_falls_back_to_plain_text() {
let err = Error::from_upstream(502, b" <html>Bad Gateway</html> ");
assert_eq!(err.message(), "<html>Bad Gateway</html>");
assert_eq!(err.status(), 502);
let err = Error::from_upstream(400, br#"{"error":"invalid_grant"}"#);
assert_eq!(err.status(), 400);
assert!(err.message().contains("invalid_grant"));
let err = Error::from_upstream(504, b"");
assert_eq!(err.status(), 504);
assert!(!err.message().is_empty(), "空 body 也要有可读的 message");
}
#[test]
fn from_upstream_truncates_oversized_body() {
let body = "x".repeat(MAX_UPSTREAM_MESSAGE_LEN * 4);
let err = Error::from_upstream(500, body.as_bytes());
assert!(
err.message().len() <= MAX_UPSTREAM_MESSAGE_LEN + 4,
"实际长度 {}",
err.message().len()
);
assert!(err.message().ends_with('…'));
}
#[test]
fn from_upstream_truncation_respects_char_boundaries() {
let body = "错".repeat(MAX_UPSTREAM_MESSAGE_LEN);
let err = Error::from_upstream(500, body.as_bytes());
assert!(err.message().chars().all(|c| c == '错' || c == '…'));
assert!(err.message().len() <= MAX_UPSTREAM_MESSAGE_LEN + 4);
}
#[test]
fn serde_round_trip_preserves_public_fields() {
let err = Error::new("boom")
.with_category("db")
.with_sub_category("sqlx")
.with_code("E1001")
.with_exception(true)
.add_extra("a")
.add_extra("b");
let json = serde_json::to_string(&err).unwrap();
let back: Error = serde_json::from_str(&json).unwrap();
assert_eq!(back.message(), "boom");
assert_eq!(back.category(), "db");
assert_eq!(back.sub_category(), Some("sqlx"));
assert_eq!(back.code(), Some("E1001"));
assert!(back.is_exception());
assert_eq!(back.extra(), ["a".to_string(), "b".to_string()]);
assert_eq!(back.status(), 0);
}
}