Skip to main content

edda_server/
observe.rs

1//! What a service logs about each request: a span with its method, path,
2//! request id, and whose it is (the campaign, once its token is checked,
3//! and the account, as the proxy says), in plain text or JSON lines, and
4//! as traces to an OpenTelemetry collector if there is one. Never a body,
5//! and never the query string, which may hold what someone typed.
6
7use std::sync::OnceLock;
8
9use axum::{
10    Router,
11    extract::Request,
12    http::{HeaderName, HeaderValue},
13    middleware::{self, Next},
14    response::Response,
15};
16use opentelemetry::trace::TracerProvider as _;
17use opentelemetry_otlp::{ExporterBuildError, SpanExporter, WithExportConfig as _};
18use opentelemetry_sdk::{Resource, trace::SdkTracerProvider};
19use tower_http::{
20    request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer},
21    trace::{DefaultOnResponse, TraceLayer},
22};
23use tracing::{Level, Span, Subscriber, field};
24use tracing_subscriber::{EnvFilter, Layer as _, fmt, fmt::MakeWriter, layer::SubscriberExt as _};
25
26/// Where traces go, once [`traces`] has set it up; flushed by [`flush`].
27static TRACES: OnceLock<SdkTracerProvider> = OnceLock::new();
28
29/// The request's id: the proxy's, or one of the service's own. Answered in
30/// the response too, so a user's report can be found in the logs.
31pub const REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
32
33/// The account a request is for, set by the hosting's proxy (from the
34/// address it came to). Believed only with `TRUST_PROXY`.
35pub const ACCOUNT: HeaderName = HeaderName::from_static("x-edda-account");
36
37/// Logs `router`'s requests. With `trust_proxy`, the proxy's request id
38/// and account are believed; without, a client's are dropped.
39pub fn requests(router: Router, trust_proxy: bool) -> Router {
40    router
41        .layer(
42            TraceLayer::new_for_http()
43                .make_span_with(move |request: &Request| span(request, trust_proxy))
44                .on_response(DefaultOnResponse::new().level(Level::INFO)),
45        )
46        .layer(PropagateRequestIdLayer::new(REQUEST_ID))
47        .layer(SetRequestIdLayer::new(REQUEST_ID, MakeRequestUuid))
48        .layer(middleware::from_fn(move |request, next| {
49            unless_trusted(request, next, trust_proxy)
50        }))
51}
52
53/// Drops the headers only the proxy may set, unless it's trusted.
54async fn unless_trusted(mut request: Request, next: Next, trust_proxy: bool) -> Response {
55    if !trust_proxy {
56        request.headers_mut().remove(REQUEST_ID);
57        request.headers_mut().remove(ACCOUNT);
58    }
59    next.run(request).await
60}
61
62fn span(request: &Request, trust_proxy: bool) -> Span {
63    let header = |name| {
64        request
65            .headers()
66            .get(name)
67            .and_then(|v: &HeaderValue| v.to_str().ok())
68            .map(str::to_owned)
69    };
70    let account = if trust_proxy { header(ACCOUNT) } else { None };
71    tracing::info_span!(
72        "request",
73        method = %request.method(),
74        path = request.uri().path(),
75        request_id = header(REQUEST_ID),
76        account,
77        campaign = field::Empty,
78    )
79}
80
81/// The subscriber `init_logging` installs, writing to `writer`: JSON lines
82/// when `json`, otherwise text for people; and with `traces`, sending each
83/// span on as a trace too.
84pub fn subscriber<W>(
85    json: bool,
86    filter: EnvFilter,
87    writer: W,
88    traces: Option<&SdkTracerProvider>,
89) -> Box<dyn Subscriber + Send + Sync>
90where
91    W: for<'w> MakeWriter<'w> + Send + Sync + 'static,
92{
93    let lines = if json {
94        fmt::layer()
95            .with_writer(writer)
96            .json()
97            .with_current_span(true)
98            .with_span_list(false)
99            .boxed()
100    } else {
101        fmt::layer().with_writer(writer).boxed()
102    };
103    let traces = traces.map(|p| tracing_opentelemetry::layer().with_tracer(p.tracer("edda")));
104    Box::new(
105        tracing_subscriber::registry()
106            .with(filter)
107            .with(lines)
108            .with(traces),
109    )
110}
111
112/// Sends spans as traces, as `service`, to the OpenTelemetry collector at
113/// `OTEL_EXPORTER_OTLP_ENDPOINT` (such as `http://collector:4318`), over
114/// HTTP; `None` without one.
115pub fn traces(service: &str) -> Result<Option<&'static SdkTracerProvider>, ExporterBuildError> {
116    let Some(endpoint) = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT")
117        .ok()
118        .filter(|e| !e.is_empty())
119    else {
120        return Ok(None);
121    };
122    let provider = tracer_provider(service, &endpoint)?;
123    Ok(Some(TRACES.get_or_init(|| provider)))
124}
125
126fn tracer_provider(service: &str, endpoint: &str) -> Result<SdkTracerProvider, ExporterBuildError> {
127    let exporter = SpanExporter::builder()
128        .with_http()
129        .with_endpoint(format!("{}/v1/traces", endpoint.trim_end_matches('/')))
130        .build()?;
131    Ok(SdkTracerProvider::builder()
132        .with_batch_exporter(exporter)
133        .with_resource(
134            Resource::builder()
135                .with_service_name(service.to_owned())
136                .build(),
137        )
138        .build())
139}
140
141/// Sends the traces not sent yet, before the service stops.
142pub fn flush() {
143    if let Some(provider) = TRACES.get()
144        && let Err(e) = provider.shutdown()
145    {
146        tracing::warn!(error = %e, "sending the last traces failed");
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use std::{
153        io,
154        sync::{Arc, Mutex},
155    };
156
157    use axum::{body::Body, routing::get};
158    use serde_json::Value;
159    use tower::ServiceExt;
160
161    use super::*;
162
163    /// What was logged.
164    #[derive(Clone, Default)]
165    struct Logged(Arc<Mutex<Vec<u8>>>);
166
167    impl io::Write for Logged {
168        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
169            self.0.lock().unwrap().extend_from_slice(bytes);
170            Ok(bytes.len())
171        }
172
173        fn flush(&mut self) -> io::Result<()> {
174            Ok(())
175        }
176    }
177
178    impl<'w> MakeWriter<'w> for Logged {
179        type Writer = Logged;
180
181        fn make_writer(&'w self) -> Self::Writer {
182            self.clone()
183        }
184    }
185
186    impl Logged {
187        /// The spans the lines were logged in.
188        fn spans(&self) -> Vec<Value> {
189            let text = String::from_utf8(self.0.lock().unwrap().clone()).unwrap();
190            text.lines()
191                .map(|line| serde_json::from_str::<Value>(line).unwrap()["span"].clone())
192                .filter(|span| !span.is_null())
193                .collect()
194        }
195
196        fn text(&self) -> String {
197            String::from_utf8(self.0.lock().unwrap().clone()).unwrap()
198        }
199    }
200
201    /// Sends `request` through a service whose handler records campaign
202    /// `c-1`, as checking a token does; returns its response's request id.
203    async fn send(request: Request, trust_proxy: bool, logged: &Logged) -> String {
204        let _logging = tracing::subscriber::set_default(subscriber(
205            true,
206            EnvFilter::new("info"),
207            logged.clone(),
208            None,
209        ));
210        let handler = || async {
211            Span::current().record("campaign", "c-1");
212            "ok"
213        };
214        let response = requests(
215            Router::new().route("/characters", get(handler)),
216            trust_proxy,
217        )
218        .oneshot(request)
219        .await
220        .unwrap();
221        response.headers()[REQUEST_ID].to_str().unwrap().to_owned()
222    }
223
224    fn request() -> Request {
225        Request::get("/characters?search=Secret%20Plans")
226            .header(REQUEST_ID, "from-the-proxy")
227            .header(ACCOUNT, "vox")
228            .body(Body::empty())
229            .unwrap()
230    }
231
232    #[tokio::test]
233    async fn requests_are_logged_with_whose_they_are_and_never_what_was_asked() {
234        let logged = Logged::default();
235        let id = send(request(), true, &logged).await;
236        assert_eq!(id, "from-the-proxy");
237        let spans = logged.spans();
238        assert!(!spans.is_empty(), "{}", logged.text());
239        for span in spans {
240            assert_eq!(span["path"], "/characters");
241            assert_eq!(span["request_id"], "from-the-proxy");
242            assert_eq!(span["account"], "vox");
243        }
244        assert!(logged.text().contains(r#""campaign":"c-1""#));
245        assert!(!logged.text().contains("Secret"), "the query stays out");
246    }
247
248    #[tokio::test]
249    async fn without_a_trusted_proxy_clients_cant_say_whose_they_are() {
250        let logged = Logged::default();
251        let id = send(request(), false, &logged).await;
252        assert_ne!(id, "from-the-proxy", "a request id of the service's own");
253        for span in logged.spans() {
254            assert_eq!(span["request_id"], id.as_str());
255            assert!(span.get("account").is_none(), "{span}");
256        }
257    }
258
259    #[tokio::test(flavor = "multi_thread")]
260    async fn spans_go_to_the_collector_as_traces() {
261        // A collector that only counts what it's sent.
262        let received = Arc::new(Mutex::new(Vec::<String>::new()));
263        let seen = received.clone();
264        let collector = Router::new().route(
265            "/v1/traces",
266            axum::routing::post(move |request: Request| {
267                let seen = seen.clone();
268                async move {
269                    let kind = request.headers()["content-type"]
270                        .to_str()
271                        .unwrap()
272                        .to_owned();
273                    seen.lock().unwrap().push(kind);
274                    ""
275                }
276            }),
277        );
278        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
279        let endpoint = format!("http://{}", listener.local_addr().unwrap());
280        tokio::spawn(async move { axum::serve(listener, collector).await });
281
282        let provider = tracer_provider("test service", &endpoint).unwrap();
283        let logging = subscriber(false, EnvFilter::new("info"), io::sink, Some(&provider));
284        tracing::subscriber::with_default(logging, || {
285            tracing::info_span!("request", path = "/characters").in_scope(|| {});
286        });
287        tokio::task::spawn_blocking(move || provider.shutdown())
288            .await
289            .unwrap()
290            .unwrap();
291        assert_eq!(*received.lock().unwrap(), ["application/x-protobuf"]);
292    }
293}