micromegas/servers/
axum_utils.rs

1//! axum-utils : observability middleware
2
3// crate-specific lint exceptions:
4#![allow(clippy::missing_errors_doc)]
5
6use anyhow::Result;
7use async_stream::stream;
8use axum::response::Response;
9use axum::{extract::Request, middleware::Next};
10use micromegas_analytics::response_writer::ResponseWriter;
11use micromegas_tracing::prelude::*;
12use std::sync::Arc;
13use tokio::sync::mpsc::Receiver;
14
15use super::http_utils::get_client_ip;
16
17/// Shared body for `observability_middleware` and `auth_observability_middleware`. Logs only
18/// the request path (never the query string) when `log_query_string` is `false` -- used by
19/// `auth_observability_middleware` because `/auth/callback`'s query carries the OAuth
20/// authorization code and the PKCE verifier (embedded in the signed `state` param) and must
21/// never be written to the telemetry log. The logged field is always named `uri=`, not `path=`
22/// -- even when the value is path-only -- so queries that grep `uri=` to reconstruct HTTP traffic
23/// across the whole `log_entries` stream still match `/auth/*` lines. `client_ip` reuses the
24/// shared `get_client_ip` (rightmost `X-Forwarded-For` entry, then `X-Real-IP`, then the socket
25/// address), same as every other route in this codebase.
26async fn observability_middleware_impl(
27    request: Request,
28    next: Next,
29    log_query_string: bool,
30) -> Response {
31    let (parts, body) = request.into_parts();
32    let uri = if log_query_string {
33        parts.uri.to_string()
34    } else {
35        parts.uri.path().to_string()
36    };
37    let client_ip = get_client_ip(&parts.headers, &parts.extensions);
38    info!(
39        "request method={} uri={uri} client_ip={client_ip}",
40        parts.method
41    );
42    let begin_ticks = now();
43    let response = next.run(Request::from_parts(parts, body)).await;
44    let end_ticks = now();
45    let duration = end_ticks - begin_ticks;
46    imetric!("request_duration", "ticks", duration as u64);
47    info!(
48        "response status={} uri={uri} client_ip={client_ip}",
49        response.status()
50    );
51    response
52}
53
54/// Logs http requests, their duration and status code, including the query string.
55pub async fn observability_middleware(request: Request, next: Next) -> Response {
56    observability_middleware_impl(request, next, true).await
57}
58
59/// Like `observability_middleware`, but never logs the query string -- see
60/// `observability_middleware_impl`'s doc comment for why.
61pub async fn auth_observability_middleware(request: Request, next: Next) -> Response {
62    observability_middleware_impl(request, next, false).await
63}
64
65/// Makes a streaming body from a Tokio MPSC receiver.
66pub fn make_body_from_channel_receiver(mut rx: Receiver<bytes::Bytes>) -> axum::body::Body {
67    let read_stream = stream! {
68        while let Some(value) = rx.recv().await{
69                yield Result::<bytes::Bytes>::Ok(value);
70        }
71    };
72    axum::body::Body::from_stream(read_stream)
73}
74
75/// Streams a response by executing a callback that writes to a `ResponseWriter`.
76///
77/// This function creates a channel and a `ResponseWriter` that writes to this channel.
78/// The `callback` is then executed in a separate Tokio task, allowing it to stream data
79/// back to the client as it becomes available.
80pub fn stream_request<F, Fut>(callback: F) -> Response
81where
82    F: FnOnce(Arc<ResponseWriter>) -> Fut + 'static + Send,
83    Fut: std::future::Future<Output = Result<()>> + Send,
84{
85    let (tx, rx) = tokio::sync::mpsc::channel(10);
86    let writer = Arc::new(ResponseWriter::new(Some(tx)));
87    let response_body = make_body_from_channel_receiver(rx);
88    spawn_with_context(async move {
89        let service_call = callback(writer.clone());
90        if let Err(e) = service_call.await {
91            if writer.is_closed() {
92                info!("Error happened, but connection is closed: {e:?}");
93            } else {
94                // the connection is live, this looks like a real error
95                error!("{e:?}");
96                if let Err(e) = writer.write_string(format!("{e:?}")).await {
97                    //error writing can happen, probably not a big deal
98                    info!("{e:?}");
99                }
100            }
101        }
102    });
103
104    Response::builder().status(200).body(response_body).unwrap()
105}
micromegas/servers/axum_utils.rs

micromegas/servers/
axum_utils.rs
