114 lines
3.2 KiB
Rust
114 lines
3.2 KiB
Rust
use std::time::{Duration, Instant};
|
||
|
||
use axum::{
|
||
extract::Request,
|
||
http::{HeaderValue, Request as HttpRequest, header::HeaderName},
|
||
middleware::Next,
|
||
response::Response,
|
||
};
|
||
use shared_contracts::api::API_RESPONSE_ENVELOPE_HEADER;
|
||
use uuid::Uuid;
|
||
|
||
pub use shared_contracts::api::X_REQUEST_ID_HEADER;
|
||
|
||
// 当前阶段先把请求级元信息统一挂到 extensions，后续响应头、envelope 与错误处理中间件继续复用。
|
||
#[derive(Clone, Debug)]
|
||
pub struct RequestContext {
|
||
request_id: String,
|
||
operation: String,
|
||
request_started_at: Instant,
|
||
wants_envelope: bool,
|
||
}
|
||
|
||
impl RequestContext {
|
||
pub fn new(
|
||
request_id: String,
|
||
operation: String,
|
||
elapsed_seed: Duration,
|
||
wants_envelope: bool,
|
||
) -> Self {
|
||
Self {
|
||
request_id,
|
||
operation,
|
||
request_started_at: Instant::now()
|
||
.checked_sub(elapsed_seed)
|
||
.unwrap_or_else(Instant::now),
|
||
wants_envelope,
|
||
}
|
||
}
|
||
|
||
pub fn request_id(&self) -> &str {
|
||
&self.request_id
|
||
}
|
||
|
||
pub fn operation(&self) -> &str {
|
||
&self.operation
|
||
}
|
||
|
||
pub fn wants_envelope(&self) -> bool {
|
||
self.wants_envelope
|
||
}
|
||
|
||
pub fn elapsed(&self) -> u64 {
|
||
self.request_started_at
|
||
.elapsed()
|
||
.as_millis()
|
||
.min(u64::MAX as u128) as u64
|
||
}
|
||
}
|
||
|
||
pub async fn attach_request_context(mut request: Request, next: Next) -> Response {
|
||
let wants_envelope = wants_api_envelope(&request);
|
||
let request_id = request
|
||
.headers()
|
||
.get(X_REQUEST_ID_HEADER)
|
||
.and_then(|value| value.to_str().ok())
|
||
.map(str::trim)
|
||
.filter(|value| !value.is_empty())
|
||
.map(ToOwned::to_owned)
|
||
.unwrap_or_else(|| Uuid::new_v4().to_string());
|
||
let operation = format!("{} {}", request.method(), request.uri());
|
||
|
||
request.extensions_mut().insert(RequestContext::new(
|
||
request_id.clone(),
|
||
operation,
|
||
Duration::ZERO,
|
||
wants_envelope,
|
||
));
|
||
|
||
// 统一把 request_id 写回请求头，方便后续 tracing、响应头与 envelope 层读取同一来源。
|
||
if let Ok(header_value) = HeaderValue::from_str(&request_id) {
|
||
request
|
||
.headers_mut()
|
||
.insert(HeaderName::from_static(X_REQUEST_ID_HEADER), header_value);
|
||
}
|
||
|
||
next.run(request).await
|
||
}
|
||
|
||
pub fn resolve_request_id<B>(request: &HttpRequest<B>) -> Option<String> {
|
||
request
|
||
.extensions()
|
||
.get::<RequestContext>()
|
||
.map(|context| context.request_id().to_string())
|
||
.or_else(|| {
|
||
request
|
||
.headers()
|
||
.get(X_REQUEST_ID_HEADER)
|
||
.and_then(|value| value.to_str().ok())
|
||
.map(str::trim)
|
||
.filter(|value| !value.is_empty())
|
||
.map(ToOwned::to_owned)
|
||
})
|
||
}
|
||
|
||
fn wants_api_envelope<B>(request: &HttpRequest<B>) -> bool {
|
||
request
|
||
.headers()
|
||
.get(API_RESPONSE_ENVELOPE_HEADER)
|
||
.and_then(|value| value.to_str().ok())
|
||
.map(str::trim)
|
||
.map(str::to_lowercase)
|
||
.is_some_and(|value| matches!(value.as_str(), "1" | "true" | "v1" | "envelope"))
|
||
}
|