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