TLA Line data Source code
1 : use crate::auth::{AuthError, Claims, SwappableJwtAuth};
2 : use crate::http::error::{api_error_handler, route_error_handler, ApiError};
3 : use anyhow::Context;
4 : use hyper::header::{HeaderName, AUTHORIZATION};
5 : use hyper::http::HeaderValue;
6 : use hyper::Method;
7 : use hyper::{header::CONTENT_TYPE, Body, Request, Response};
8 : use metrics::{register_int_counter, Encoder, IntCounter, TextEncoder};
9 : use once_cell::sync::Lazy;
10 : use routerify::ext::RequestExt;
11 : use routerify::{Middleware, RequestInfo, Router, RouterBuilder};
12 : use tracing::{self, debug, info, info_span, warn, Instrument};
13 :
14 : use std::future::Future;
15 : use std::str::FromStr;
16 :
17 : use bytes::{Bytes, BytesMut};
18 : use std::io::Write as _;
19 : use tokio::sync::mpsc;
20 : use tokio_stream::wrappers::ReceiverStream;
21 :
22 CBC 367 : static SERVE_METRICS_COUNT: Lazy<IntCounter> = Lazy::new(|| {
23 367 : register_int_counter!(
24 367 : "libmetrics_metric_handler_requests_total",
25 367 : "Number of metric requests made"
26 367 : )
27 367 : .expect("failed to define a metric")
28 367 : });
29 :
30 : static X_REQUEST_ID_HEADER_STR: &str = "x-request-id";
31 :
32 : static X_REQUEST_ID_HEADER: HeaderName = HeaderName::from_static(X_REQUEST_ID_HEADER_STR);
33 26046 : #[derive(Debug, Default, Clone)]
34 : struct RequestId(String);
35 :
36 : /// Adds a tracing info_span! instrumentation around the handler events,
37 : /// logs the request start and end events for non-GET requests and non-200 responses.
38 : ///
39 : /// Usage: Replace `my_handler` with `|r| request_span(r, my_handler)`
40 : ///
41 : /// Use this to distinguish between logs of different HTTP requests: every request handler wrapped
42 : /// with this will get request info logged in the wrapping span, including the unique request ID.
43 : ///
44 : /// This also handles errors, logging them and converting them to an HTTP error response.
45 : ///
46 : /// NB: If the client disconnects, Hyper will drop the Future, without polling it to
47 : /// completion. In other words, the handler must be async cancellation safe! request_span
48 : /// prints a warning to the log when that happens, so that you have some trace of it in
49 : /// the log.
50 : ///
51 : ///
52 : /// There could be other ways to implement similar functionality:
53 : ///
54 : /// * procmacros placed on top of all handler methods
55 : /// With all the drawbacks of procmacros, brings no difference implementation-wise,
56 : /// and little code reduction compared to the existing approach.
57 : ///
58 : /// * Another `TraitExt` with e.g. the `get_with_span`, `post_with_span` methods to do similar logic,
59 : /// implemented for [`RouterBuilder`].
60 : /// Could be simpler, but we don't want to depend on [`routerify`] more, targeting to use other library later.
61 : ///
62 : /// * In theory, a span guard could've been created in a pre-request middleware and placed into a global collection, to be dropped
63 : /// later, in a post-response middleware.
64 : /// Due to suspendable nature of the futures, would give contradictive results which is exactly the opposite of what `tracing-futures`
65 : /// tries to achive with its `.instrument` used in the current approach.
66 : ///
67 : /// If needed, a declarative macro to substitute the |r| ... closure boilerplate could be introduced.
68 13011 : pub async fn request_span<R, H>(request: Request<Body>, handler: H) -> R::Output
69 13011 : where
70 13011 : R: Future<Output = Result<Response<Body>, ApiError>> + Send + 'static,
71 13011 : H: FnOnce(Request<Body>) -> R + Send + Sync + 'static,
72 13011 : {
73 13011 : let request_id = request.context::<RequestId>().unwrap_or_default().0;
74 13011 : let method = request.method();
75 13011 : let path = request.uri().path();
76 13011 : let request_span = info_span!("request", %method, %path, %request_id);
77 :
78 13011 : let log_quietly = method == Method::GET;
79 13011 : async move {
80 13011 : let cancellation_guard = RequestCancelled::warn_when_dropped_without_responding();
81 13011 : if log_quietly {
82 6236 : debug!("Handling request");
83 : } else {
84 6775 : info!("Handling request");
85 : }
86 :
87 : // No special handling for panics here. There's a `tracing_panic_hook` from another
88 : // module to do that globally.
89 16199 : let res = handler(request).await;
90 :
91 13003 : cancellation_guard.disarm();
92 13003 :
93 13003 : // Log the result if needed.
94 13003 : //
95 13003 : // We also convert any errors into an Ok response with HTTP error code here.
96 13003 : // `make_router` sets a last-resort error handler that would do the same, but
97 13003 : // we prefer to do it here, before we exit the request span, so that the error
98 13003 : // is still logged with the span.
99 13003 : //
100 13003 : // (Because we convert errors to Ok response, we never actually return an error,
101 13003 : // and we could declare the function to return the never type (`!`). However,
102 13003 : // using `routerify::RouterBuilder` requires a proper error type.)
103 13003 : match res {
104 12817 : Ok(response) => {
105 12817 : let response_status = response.status();
106 12817 : if log_quietly && response_status.is_success() {
107 6113 : debug!("Request handled, status: {response_status}");
108 : } else {
109 6704 : info!("Request handled, status: {response_status}");
110 : }
111 12817 : Ok(response)
112 : }
113 186 : Err(err) => Ok(api_error_handler(err)),
114 : }
115 13003 : }
116 13011 : .instrument(request_span)
117 16199 : .await
118 13003 : }
119 :
120 : /// Drop guard to WARN in case the request was dropped before completion.
121 : struct RequestCancelled {
122 : warn: Option<tracing::Span>,
123 : }
124 :
125 : impl RequestCancelled {
126 : /// Create the drop guard using the [`tracing::Span::current`] as the span.
127 13011 : fn warn_when_dropped_without_responding() -> Self {
128 13011 : RequestCancelled {
129 13011 : warn: Some(tracing::Span::current()),
130 13011 : }
131 13011 : }
132 :
133 : /// Consume the drop guard without logging anything.
134 13003 : fn disarm(mut self) {
135 13003 : self.warn = None;
136 13003 : }
137 : }
138 :
139 : impl Drop for RequestCancelled {
140 13005 : fn drop(&mut self) {
141 13005 : if std::thread::panicking() {
142 UBC 0 : // we are unwinding due to panicking, assume we are not dropped for cancellation
143 CBC 13005 : } else if let Some(span) = self.warn.take() {
144 : // the span has all of the info already, but the outer `.instrument(span)` has already
145 : // been dropped, so we need to manually re-enter it for this message.
146 : //
147 : // this is what the instrument would do before polling so it is fine.
148 2 : let _g = span.entered();
149 2 : warn!("request was dropped before completing");
150 13003 : }
151 13005 : }
152 : }
153 :
154 : /// An [`std::io::Write`] implementation on top of a channel sending [`bytes::Bytes`] chunks.
155 : pub struct ChannelWriter {
156 : buffer: BytesMut,
157 : pub tx: mpsc::Sender<std::io::Result<Bytes>>,
158 : written: usize,
159 : }
160 :
161 : impl ChannelWriter {
162 1166 : pub fn new(buf_len: usize, tx: mpsc::Sender<std::io::Result<Bytes>>) -> Self {
163 1166 : assert_ne!(buf_len, 0);
164 1166 : ChannelWriter {
165 1166 : // split about half off the buffer from the start, because we flush depending on
166 1166 : // capacity. first flush will come sooner than without this, but now resizes will
167 1166 : // have better chance of picking up the "other" half. not guaranteed of course.
168 1166 : buffer: BytesMut::with_capacity(buf_len).split_off(buf_len / 2),
169 1166 : tx,
170 1166 : written: 0,
171 1166 : }
172 1166 : }
173 :
174 3837 : pub fn flush0(&mut self) -> std::io::Result<usize> {
175 3837 : let n = self.buffer.len();
176 3837 : if n == 0 {
177 UBC 0 : return Ok(0);
178 CBC 3837 : }
179 3837 :
180 3837 : tracing::trace!(n, "flushing");
181 3837 : let ready = self.buffer.split().freeze();
182 3837 :
183 3837 : // not ideal to call from blocking code to block_on, but we are sure that this
184 3837 : // operation does not spawn_blocking other tasks
185 3837 : let res: Result<(), ()> = tokio::runtime::Handle::current().block_on(async {
186 3837 : self.tx.send(Ok(ready)).await.map_err(|_| ())?;
187 :
188 : // throttle sending to allow reuse of our buffer in `write`.
189 3837 : self.tx.reserve().await.map_err(|_| ())?;
190 :
191 : // now the response task has picked up the buffer and hopefully started
192 : // sending it to the client.
193 3837 : Ok(())
194 3837 : });
195 3837 : if res.is_err() {
196 UBC 0 : return Err(std::io::ErrorKind::BrokenPipe.into());
197 CBC 3837 : }
198 3837 : self.written += n;
199 3837 : Ok(n)
200 3837 : }
201 :
202 1166 : pub fn flushed_bytes(&self) -> usize {
203 1166 : self.written
204 1166 : }
205 : }
206 :
207 : impl std::io::Write for ChannelWriter {
208 33781103 : fn write(&mut self, mut buf: &[u8]) -> std::io::Result<usize> {
209 33781103 : let remaining = self.buffer.capacity() - self.buffer.len();
210 33781103 :
211 33781103 : let out_of_space = remaining < buf.len();
212 33781103 :
213 33781103 : let original_len = buf.len();
214 33781103 :
215 33781103 : if out_of_space {
216 2671 : let can_still_fit = buf.len() - remaining;
217 2671 : self.buffer.extend_from_slice(&buf[..can_still_fit]);
218 2671 : buf = &buf[can_still_fit..];
219 2671 : self.flush0()?;
220 33778432 : }
221 :
222 : // assume that this will often under normal operation just move the pointer back to the
223 : // beginning of allocation, because previous split off parts are already sent and
224 : // dropped.
225 33781103 : self.buffer.extend_from_slice(buf);
226 33781103 : Ok(original_len)
227 33781103 : }
228 :
229 1166 : fn flush(&mut self) -> std::io::Result<()> {
230 1166 : self.flush0().map(|_| ())
231 1166 : }
232 : }
233 :
234 1161 : async fn prometheus_metrics_handler(_req: Request<Body>) -> Result<Response<Body>, ApiError> {
235 1161 : SERVE_METRICS_COUNT.inc();
236 1161 :
237 1161 : let started_at = std::time::Instant::now();
238 1161 :
239 1161 : let (tx, rx) = mpsc::channel(1);
240 1161 :
241 1161 : let body = Body::wrap_stream(ReceiverStream::new(rx));
242 1161 :
243 1161 : let mut writer = ChannelWriter::new(128 * 1024, tx);
244 1161 :
245 1161 : let encoder = TextEncoder::new();
246 1161 :
247 1161 : let response = Response::builder()
248 1161 : .status(200)
249 1161 : .header(CONTENT_TYPE, encoder.format_type())
250 1161 : .body(body)
251 1161 : .unwrap();
252 :
253 1161 : let span = info_span!("blocking");
254 1161 : tokio::task::spawn_blocking(move || {
255 1161 : let _span = span.entered();
256 1161 : let metrics = metrics::gather();
257 1161 : let res = encoder
258 1161 : .encode(&metrics, &mut writer)
259 1161 : .and_then(|_| writer.flush().map_err(|e| e.into()));
260 1161 :
261 1161 : match res {
262 : Ok(()) => {
263 1161 : tracing::info!(
264 1161 : bytes = writer.flushed_bytes(),
265 1161 : elapsed_ms = started_at.elapsed().as_millis(),
266 1161 : "responded /metrics"
267 1161 : );
268 : }
269 UBC 0 : Err(e) => {
270 0 : tracing::warn!("failed to write out /metrics response: {e:#}");
271 : // semantics of this error are quite... unclear. we want to error the stream out to
272 : // abort the response to somehow notify the client that we failed.
273 : //
274 : // though, most likely the reason for failure is that the receiver is already gone.
275 0 : drop(
276 0 : writer
277 0 : .tx
278 0 : .blocking_send(Err(std::io::ErrorKind::BrokenPipe.into())),
279 0 : );
280 : }
281 : }
282 CBC 1161 : });
283 1161 :
284 1161 : Ok(response)
285 1161 : }
286 :
287 1402 : pub fn add_request_id_middleware<B: hyper::body::HttpBody + Send + Sync + 'static>(
288 1402 : ) -> Middleware<B, ApiError> {
289 13043 : Middleware::pre(move |req| async move {
290 13043 : let request_id = match req.headers().get(&X_REQUEST_ID_HEADER) {
291 1 : Some(request_id) => request_id
292 1 : .to_str()
293 1 : .expect("extract request id value")
294 1 : .to_owned(),
295 : None => {
296 13042 : let request_id = uuid::Uuid::new_v4();
297 13042 : request_id.to_string()
298 : }
299 : };
300 13043 : req.set_context(RequestId(request_id));
301 13043 :
302 13043 : Ok(req)
303 13043 : })
304 1402 : }
305 :
306 13035 : async fn add_request_id_header_to_response(
307 13035 : mut res: Response<Body>,
308 13035 : req_info: RequestInfo,
309 13035 : ) -> Result<Response<Body>, ApiError> {
310 13035 : if let Some(request_id) = req_info.context::<RequestId>() {
311 13035 : if let Ok(request_header_value) = HeaderValue::from_str(&request_id.0) {
312 13035 : res.headers_mut()
313 13035 : .insert(&X_REQUEST_ID_HEADER, request_header_value);
314 13035 : };
315 UBC 0 : };
316 :
317 CBC 13035 : Ok(res)
318 13035 : }
319 :
320 1402 : pub fn make_router() -> RouterBuilder<hyper::Body, ApiError> {
321 1402 : Router::builder()
322 1402 : .middleware(add_request_id_middleware())
323 1402 : .middleware(Middleware::post_with_info(
324 1402 : add_request_id_header_to_response,
325 1402 : ))
326 1402 : .get("/metrics", |r| request_span(r, prometheus_metrics_handler))
327 1402 : .err_handler(route_error_handler)
328 1402 : }
329 :
330 557 : pub fn attach_openapi_ui(
331 557 : router_builder: RouterBuilder<hyper::Body, ApiError>,
332 557 : spec: &'static [u8],
333 557 : spec_mount_path: &'static str,
334 557 : ui_mount_path: &'static str,
335 557 : ) -> RouterBuilder<hyper::Body, ApiError> {
336 557 : router_builder
337 557 : .get(spec_mount_path,
338 557 : move |r| request_span(r, move |_| async move {
339 UBC 0 : Ok(Response::builder().body(Body::from(spec)).unwrap())
340 CBC 557 : })
341 557 : )
342 557 : .get(ui_mount_path,
343 557 : move |r| request_span(r, move |_| async move {
344 UBC 0 : Ok(Response::builder().body(Body::from(format!(r#"
345 0 : <!DOCTYPE html>
346 0 : <html lang="en">
347 0 : <head>
348 0 : <title>rweb</title>
349 0 : <link href="https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui.css" rel="stylesheet">
350 0 : </head>
351 0 : <body>
352 0 : <div id="swagger-ui"></div>
353 0 : <script src="https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui-bundle.js" charset="UTF-8"> </script>
354 0 : <script>
355 0 : window.onload = function() {{
356 0 : const ui = SwaggerUIBundle({{
357 0 : "dom_id": "\#swagger-ui",
358 0 : presets: [
359 0 : SwaggerUIBundle.presets.apis,
360 0 : SwaggerUIBundle.SwaggerUIStandalonePreset
361 0 : ],
362 0 : layout: "BaseLayout",
363 0 : deepLinking: true,
364 0 : showExtensions: true,
365 0 : showCommonExtensions: true,
366 0 : url: "{}",
367 0 : }})
368 0 : window.ui = ui;
369 0 : }};
370 0 : </script>
371 0 : </body>
372 0 : </html>
373 0 : "#, spec_mount_path))).unwrap())
374 CBC 557 : })
375 557 : )
376 557 : }
377 :
378 134 : fn parse_token(header_value: &str) -> Result<&str, ApiError> {
379 : // header must be in form Bearer <token>
380 134 : let (prefix, token) = header_value
381 134 : .split_once(' ')
382 134 : .ok_or_else(|| ApiError::Unauthorized("malformed authorization header".to_string()))?;
383 134 : if prefix != "Bearer" {
384 UBC 0 : return Err(ApiError::Unauthorized(
385 0 : "malformed authorization header".to_string(),
386 0 : ));
387 CBC 134 : }
388 134 : Ok(token)
389 134 : }
390 :
391 31 : pub fn auth_middleware<B: hyper::body::HttpBody + Send + Sync + 'static>(
392 31 : provide_auth: fn(&Request<Body>) -> Option<&SwappableJwtAuth>,
393 31 : ) -> Middleware<B, ApiError> {
394 217 : Middleware::pre(move |req| async move {
395 217 : if let Some(auth) = provide_auth(&req) {
396 139 : match req.headers().get(AUTHORIZATION) {
397 134 : Some(value) => {
398 134 : let header_value = value.to_str().map_err(|_| {
399 UBC 0 : ApiError::Unauthorized("malformed authorization header".to_string())
400 CBC 134 : })?;
401 134 : let token = parse_token(header_value)?;
402 :
403 134 : let data = auth.decode(token).map_err(|err| {
404 3 : warn!("Authentication error: {err}");
405 : // Rely on From<AuthError> for ApiError impl
406 3 : err
407 134 : })?;
408 131 : req.set_context(data.claims);
409 : }
410 : None => {
411 5 : return Err(ApiError::Unauthorized(
412 5 : "missing authorization header".to_string(),
413 5 : ))
414 : }
415 : }
416 78 : }
417 209 : Ok(req)
418 217 : })
419 31 : }
420 :
421 557 : pub fn add_response_header_middleware<B>(
422 557 : header: &str,
423 557 : value: &str,
424 557 : ) -> anyhow::Result<Middleware<B, ApiError>>
425 557 : where
426 557 : B: hyper::body::HttpBody + Send + Sync + 'static,
427 557 : {
428 557 : let name =
429 557 : HeaderName::from_str(header).with_context(|| format!("invalid header name: {header}"))?;
430 557 : let value =
431 557 : HeaderValue::from_str(value).with_context(|| format!("invalid header value: {value}"))?;
432 557 : Ok(Middleware::post_with_info(
433 10064 : move |mut response, request_info| {
434 10064 : let name = name.clone();
435 10064 : let value = value.clone();
436 10064 : async move {
437 10064 : let headers = response.headers_mut();
438 10064 : if headers.contains_key(&name) {
439 UBC 0 : warn!(
440 0 : "{} response already contains header {:?}",
441 0 : request_info.uri(),
442 0 : &name,
443 0 : );
444 CBC 10064 : } else {
445 10064 : headers.insert(name, value);
446 10064 : }
447 10064 : Ok(response)
448 10064 : }
449 10064 : },
450 557 : ))
451 557 : }
452 :
453 10085 : pub fn check_permission_with(
454 10085 : req: &Request<Body>,
455 10085 : check_permission: impl Fn(&Claims) -> Result<(), AuthError>,
456 10085 : ) -> Result<(), ApiError> {
457 10085 : match req.context::<Claims>() {
458 131 : Some(claims) => Ok(check_permission(&claims)
459 131 : .map_err(|_err| ApiError::Forbidden("JWT authentication error".to_string()))?),
460 9954 : None => Ok(()), // claims is None because auth is disabled
461 : }
462 10085 : }
463 :
464 : #[cfg(test)]
465 : mod tests {
466 : use super::*;
467 : use futures::future::poll_fn;
468 : use hyper::service::Service;
469 : use routerify::RequestServiceBuilder;
470 : use std::net::{IpAddr, SocketAddr};
471 :
472 1 : #[tokio::test]
473 1 : async fn test_request_id_returned() {
474 1 : let builder = RequestServiceBuilder::new(make_router().build().unwrap()).unwrap();
475 1 : let remote_addr = SocketAddr::new(IpAddr::from_str("127.0.0.1").unwrap(), 80);
476 1 : let mut service = builder.build(remote_addr);
477 1 : if let Err(e) = poll_fn(|ctx| service.poll_ready(ctx)).await {
478 UBC 0 : panic!("request service is not ready: {:?}", e);
479 CBC 1 : }
480 1 :
481 1 : let mut req: Request<Body> = Request::default();
482 1 : req.headers_mut()
483 1 : .append(&X_REQUEST_ID_HEADER, HeaderValue::from_str("42").unwrap());
484 :
485 1 : let resp: Response<hyper::body::Body> = service.call(req).await.unwrap();
486 1 :
487 1 : let header_val = resp.headers().get(&X_REQUEST_ID_HEADER).unwrap();
488 :
489 1 : assert!(header_val == "42", "response header mismatch");
490 : }
491 :
492 1 : #[tokio::test]
493 1 : async fn test_request_id_empty() {
494 1 : let builder = RequestServiceBuilder::new(make_router().build().unwrap()).unwrap();
495 1 : let remote_addr = SocketAddr::new(IpAddr::from_str("127.0.0.1").unwrap(), 80);
496 1 : let mut service = builder.build(remote_addr);
497 1 : if let Err(e) = poll_fn(|ctx| service.poll_ready(ctx)).await {
498 UBC 0 : panic!("request service is not ready: {:?}", e);
499 CBC 1 : }
500 1 :
501 1 : let req: Request<Body> = Request::default();
502 1 : let resp: Response<hyper::body::Body> = service.call(req).await.unwrap();
503 1 :
504 1 : let header_val = resp.headers().get(&X_REQUEST_ID_HEADER);
505 1 :
506 1 : assert_ne!(header_val, None, "response header should NOT be empty");
507 : }
508 : }
|
