feat(axum): update route logging middleware

This commit is contained in:
xfy
2024-10-21 09:54:02 +08:00
parent a1c2e73575
commit 602c7f0fde
2 changed files with 28 additions and 76 deletions

View File

@ -40,30 +40,32 @@ pub async fn add_version(
/// This middleware will calculate each request latency
/// and add request's information to each info_span.
pub fn logging_route(router: Router) -> Router {
router.layer(
TraceLayer::new_for_http()
.make_span_with(|req: &Request<_>| {
let unknown = &HeaderValue::from_static("Unknown");
let empty = &HeaderValue::from_static("");
let headers = req.headers();
let ua = headers
.get("User-Agent")
.unwrap_or(unknown)
.to_str()
.unwrap_or("Unknown");
let host = headers.get("Host").unwrap_or(empty).to_str().unwrap_or("");
info_span!("HTTP", method = ?req.method(), host, uri = ?req.uri(), ua)
})
.on_request(|_req: &Request<_>, _span: &Span| {})
.on_response(|res: &Response, latency: Duration, _span: &Span| {
info!("{} {}μs", res.status(), latency.as_micros());
})
.on_body_chunk(|_chunk: &Bytes, _latency: Duration, _span: &Span| {})
.on_eos(|_trailers: Option<&HeaderMap>, _stream_duration: Duration, _span: &Span| {})
.on_failure(
|error: ServerErrorsFailureClass, _latency: Duration, _span: &Span| {
error!("{}", error);
},
),
)
let make_span = |req: &Request<_>| {
let unknown = &HeaderValue::from_static("Unknown");
let empty = &HeaderValue::from_static("");
let headers = req.headers();
let ua = headers
.get("User-Agent")
.unwrap_or(unknown)
.to_str()
.unwrap_or("Unknown");
let host = headers.get("Host").unwrap_or(empty).to_str().unwrap_or("");
info_span!("HTTP", method = ?req.method(), host, uri = ?req.uri(), ua)
};
let trace_layer = TraceLayer::new_for_http()
.make_span_with(make_span)
.on_request(|_req: &Request<_>, _span: &Span| {})
.on_response(|res: &Response, latency: Duration, _span: &Span| {
info!("{}", format_latency(latency, res.status()));
})
.on_body_chunk(|_chunk: &Bytes, _latency: Duration, _span: &Span| {})
.on_eos(|_trailers: Option<&HeaderMap>, _stream_duration: Duration, _span: &Span| {})
.on_failure(
|error: ServerErrorsFailureClass, latency: Duration, _span: &Span| {
error!("{}", format_latency(latency, error));
},
);
router.layer(trace_layer)
}

View File

@ -72,53 +72,3 @@ pub fn routes() -> Router {
pub async fn hello() -> String {
format!("hello {}", env!("CARGO_PKG_NAME"))
}
/// Fallback route handler for handling unmatched routes.
///
/// This asynchronous function takes a `Uri` as an argument, representing the unmatched route.
/// It logs a message indicating that the specified route is not found and returns a standard
/// "Not Found" response with a `StatusCode` of `404`.
///
/// # Arguments
///
/// - `uri`: The `Uri` representing the unmatched route.
///
/// # Returns
///
/// Returns a tuple `(StatusCode, &str)` where `StatusCode` is set to `NOT_FOUND` (404),
/// indicating that the route was not found, and the string "Not found" as the response body.
pub async fn fallback(uri: Uri) -> impl IntoResponse {
info!("route {} not found", uri);
(StatusCode::NOT_FOUND, "Not found")
}
#[derive(Debug)]
enum Version {
V1,
V2,
V3,
}
#[async_trait]
impl<S> FromRequestParts<S> for Version
where
S: Send + Sync,
{
type Rejection = Response;
async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
let params: Path<HashMap<String, String>> =
parts.extract().await.map_err(IntoResponse::into_response)?;
let version = params
.get("version")
.ok_or_else(|| (StatusCode::NOT_FOUND, "version param missing").into_response())?;
match version.as_str() {
"v1" => Ok(Version::V1),
"v2" => Ok(Version::V2),
"v3" => Ok(Version::V3),
_ => Err((StatusCode::NOT_FOUND, "unknown version").into_response()),
}
}
}