1use 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
26static TRACES: OnceLock<SdkTracerProvider> = OnceLock::new();
28
29pub const REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
32
33pub const ACCOUNT: HeaderName = HeaderName::from_static("x-edda-account");
36
37pub 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
53async 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
81pub 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
112pub 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
141pub 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 #[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 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 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 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}