LCOV - code coverage report
Current view: top level - libs/http-utils/src - endpoint.rs (source / functions) Coverage Total Hit
Test: 1b0a6a0c05cee5a7de360813c8034804e105ce1c.info Lines: 13.5 % 498 67
Test Date: 2025-03-12 00:01:28 Functions: 2.2 % 543 12

            Line data    Source code
       1              : use std::future::Future;
       2              : use std::io::Write as _;
       3              : use std::str::FromStr;
       4              : use std::time::Duration;
       5              : 
       6              : use anyhow::{Context, anyhow};
       7              : use bytes::{Bytes, BytesMut};
       8              : use hyper::header::{AUTHORIZATION, CONTENT_DISPOSITION, CONTENT_TYPE, HeaderName};
       9              : use hyper::http::HeaderValue;
      10              : use hyper::{Body, Method, Request, Response};
      11              : use metrics::{Encoder, IntCounter, TextEncoder, register_int_counter};
      12              : use once_cell::sync::Lazy;
      13              : use pprof::ProfilerGuardBuilder;
      14              : use pprof::protos::Message as _;
      15              : use routerify::ext::RequestExt;
      16              : use routerify::{Middleware, RequestInfo, Router, RouterBuilder};
      17              : use tokio::sync::{Mutex, Notify, mpsc};
      18              : use tokio_stream::wrappers::ReceiverStream;
      19              : use tokio_util::io::ReaderStream;
      20              : use tracing::{Instrument, debug, info, info_span, warn};
      21              : use utils::auth::{AuthError, Claims, SwappableJwtAuth};
      22              : 
      23              : use crate::error::{ApiError, api_error_handler, route_error_handler};
      24              : use crate::request::{get_query_param, parse_query_param};
      25              : 
      26            0 : static SERVE_METRICS_COUNT: Lazy<IntCounter> = Lazy::new(|| {
      27            0 :     register_int_counter!(
      28            0 :         "libmetrics_metric_handler_requests_total",
      29            0 :         "Number of metric requests made"
      30            0 :     )
      31            0 :     .expect("failed to define a metric")
      32            0 : });
      33              : 
      34              : static X_REQUEST_ID_HEADER_STR: &str = "x-request-id";
      35              : 
      36              : static X_REQUEST_ID_HEADER: HeaderName = HeaderName::from_static(X_REQUEST_ID_HEADER_STR);
      37              : #[derive(Debug, Default, Clone)]
      38              : struct RequestId(String);
      39              : 
      40              : /// Adds a tracing info_span! instrumentation around the handler events,
      41              : /// logs the request start and end events for non-GET requests and non-200 responses.
      42              : ///
      43              : /// Usage: Replace `my_handler` with `|r| request_span(r, my_handler)`
      44              : ///
      45              : /// Use this to distinguish between logs of different HTTP requests: every request handler wrapped
      46              : /// with this will get request info logged in the wrapping span, including the unique request ID.
      47              : ///
      48              : /// This also handles errors, logging them and converting them to an HTTP error response.
      49              : ///
      50              : /// NB: If the client disconnects, Hyper will drop the Future, without polling it to
      51              : /// completion. In other words, the handler must be async cancellation safe! request_span
      52              : /// prints a warning to the log when that happens, so that you have some trace of it in
      53              : /// the log.
      54              : ///
      55              : ///
      56              : /// There could be other ways to implement similar functionality:
      57              : ///
      58              : /// * procmacros placed on top of all handler methods
      59              : ///   With all the drawbacks of procmacros, brings no difference implementation-wise,
      60              : ///   and little code reduction compared to the existing approach.
      61              : ///
      62              : /// * Another `TraitExt` with e.g. the `get_with_span`, `post_with_span` methods to do similar logic,
      63              : ///   implemented for [`RouterBuilder`].
      64              : ///   Could be simpler, but we don't want to depend on [`routerify`] more, targeting to use other library later.
      65              : ///
      66              : /// * In theory, a span guard could've been created in a pre-request middleware and placed into a global collection, to be dropped
      67              : ///   later, in a post-response middleware.
      68              : ///   Due to suspendable nature of the futures, would give contradictive results which is exactly the opposite of what `tracing-futures`
      69              : ///   tries to achive with its `.instrument` used in the current approach.
      70              : ///
      71              : /// If needed, a declarative macro to substitute the |r| ... closure boilerplate could be introduced.
      72            0 : pub async fn request_span<R, H>(request: Request<Body>, handler: H) -> R::Output
      73            0 : where
      74            0 :     R: Future<Output = Result<Response<Body>, ApiError>> + Send + 'static,
      75            0 :     H: FnOnce(Request<Body>) -> R + Send + Sync + 'static,
      76            0 : {
      77            0 :     let request_id = request.context::<RequestId>().unwrap_or_default().0;
      78            0 :     let method = request.method();
      79            0 :     let path = request.uri().path();
      80            0 :     let request_span = info_span!("request", %method, %path, %request_id);
      81              : 
      82            0 :     let log_quietly = method == Method::GET;
      83            0 :     async move {
      84            0 :         let cancellation_guard = RequestCancelled::warn_when_dropped_without_responding();
      85            0 :         if log_quietly {
      86            0 :             debug!("Handling request");
      87              :         } else {
      88            0 :             info!("Handling request");
      89              :         }
      90              : 
      91              :         // No special handling for panics here. There's a `tracing_panic_hook` from another
      92              :         // module to do that globally.
      93            0 :         let res = handler(request).await;
      94              : 
      95            0 :         cancellation_guard.disarm();
      96            0 : 
      97            0 :         // Log the result if needed.
      98            0 :         //
      99            0 :         // We also convert any errors into an Ok response with HTTP error code here.
     100            0 :         // `make_router` sets a last-resort error handler that would do the same, but
     101            0 :         // we prefer to do it here, before we exit the request span, so that the error
     102            0 :         // is still logged with the span.
     103            0 :         //
     104            0 :         // (Because we convert errors to Ok response, we never actually return an error,
     105            0 :         // and we could declare the function to return the never type (`!`). However,
     106            0 :         // using `routerify::RouterBuilder` requires a proper error type.)
     107            0 :         match res {
     108            0 :             Ok(response) => {
     109            0 :                 let response_status = response.status();
     110            0 :                 if log_quietly && response_status.is_success() {
     111            0 :                     debug!("Request handled, status: {response_status}");
     112              :                 } else {
     113            0 :                     info!("Request handled, status: {response_status}");
     114              :                 }
     115            0 :                 Ok(response)
     116              :             }
     117            0 :             Err(err) => Ok(api_error_handler(err)),
     118              :         }
     119            0 :     }
     120            0 :     .instrument(request_span)
     121            0 :     .await
     122            0 : }
     123              : 
     124              : /// Drop guard to WARN in case the request was dropped before completion.
     125              : struct RequestCancelled {
     126              :     warn: Option<tracing::Span>,
     127              : }
     128              : 
     129              : impl RequestCancelled {
     130              :     /// Create the drop guard using the [`tracing::Span::current`] as the span.
     131            0 :     fn warn_when_dropped_without_responding() -> Self {
     132            0 :         RequestCancelled {
     133            0 :             warn: Some(tracing::Span::current()),
     134            0 :         }
     135            0 :     }
     136              : 
     137              :     /// Consume the drop guard without logging anything.
     138            0 :     fn disarm(mut self) {
     139            0 :         self.warn = None;
     140            0 :     }
     141              : }
     142              : 
     143              : impl Drop for RequestCancelled {
     144            0 :     fn drop(&mut self) {
     145            0 :         if std::thread::panicking() {
     146            0 :             // we are unwinding due to panicking, assume we are not dropped for cancellation
     147            0 :         } else if let Some(span) = self.warn.take() {
     148              :             // the span has all of the info already, but the outer `.instrument(span)` has already
     149              :             // been dropped, so we need to manually re-enter it for this message.
     150              :             //
     151              :             // this is what the instrument would do before polling so it is fine.
     152            0 :             let _g = span.entered();
     153            0 :             warn!("request was dropped before completing");
     154            0 :         }
     155            0 :     }
     156              : }
     157              : 
     158              : /// An [`std::io::Write`] implementation on top of a channel sending [`bytes::Bytes`] chunks.
     159              : pub struct ChannelWriter {
     160              :     buffer: BytesMut,
     161              :     pub tx: mpsc::Sender<std::io::Result<Bytes>>,
     162              :     written: usize,
     163              :     /// Time spent waiting for the channel to make progress. It is not the same as time to upload a
     164              :     /// buffer because we cannot know anything about that, but this should allow us to understand
     165              :     /// the actual time taken without the time spent `std::thread::park`ed.
     166              :     wait_time: std::time::Duration,
     167              : }
     168              : 
     169              : impl ChannelWriter {
     170            0 :     pub fn new(buf_len: usize, tx: mpsc::Sender<std::io::Result<Bytes>>) -> Self {
     171            0 :         assert_ne!(buf_len, 0);
     172            0 :         ChannelWriter {
     173            0 :             // split about half off the buffer from the start, because we flush depending on
     174            0 :             // capacity. first flush will come sooner than without this, but now resizes will
     175            0 :             // have better chance of picking up the "other" half. not guaranteed of course.
     176            0 :             buffer: BytesMut::with_capacity(buf_len).split_off(buf_len / 2),
     177            0 :             tx,
     178            0 :             written: 0,
     179            0 :             wait_time: std::time::Duration::ZERO,
     180            0 :         }
     181            0 :     }
     182              : 
     183            0 :     pub fn flush0(&mut self) -> std::io::Result<usize> {
     184            0 :         let n = self.buffer.len();
     185            0 :         if n == 0 {
     186            0 :             return Ok(0);
     187            0 :         }
     188            0 : 
     189            0 :         tracing::trace!(n, "flushing");
     190            0 :         let ready = self.buffer.split().freeze();
     191            0 : 
     192            0 :         let wait_started_at = std::time::Instant::now();
     193            0 : 
     194            0 :         // not ideal to call from blocking code to block_on, but we are sure that this
     195            0 :         // operation does not spawn_blocking other tasks
     196            0 :         let res: Result<(), ()> = tokio::runtime::Handle::current().block_on(async {
     197            0 :             self.tx.send(Ok(ready)).await.map_err(|_| ())?;
     198              : 
     199              :             // throttle sending to allow reuse of our buffer in `write`.
     200            0 :             self.tx.reserve().await.map_err(|_| ())?;
     201              : 
     202              :             // now the response task has picked up the buffer and hopefully started
     203              :             // sending it to the client.
     204            0 :             Ok(())
     205            0 :         });
     206            0 : 
     207            0 :         self.wait_time += wait_started_at.elapsed();
     208            0 : 
     209            0 :         if res.is_err() {
     210            0 :             return Err(std::io::ErrorKind::BrokenPipe.into());
     211            0 :         }
     212            0 :         self.written += n;
     213            0 :         Ok(n)
     214            0 :     }
     215              : 
     216            0 :     pub fn flushed_bytes(&self) -> usize {
     217            0 :         self.written
     218            0 :     }
     219              : 
     220            0 :     pub fn wait_time(&self) -> std::time::Duration {
     221            0 :         self.wait_time
     222            0 :     }
     223              : }
     224              : 
     225              : impl std::io::Write for ChannelWriter {
     226            0 :     fn write(&mut self, mut buf: &[u8]) -> std::io::Result<usize> {
     227            0 :         let remaining = self.buffer.capacity() - self.buffer.len();
     228            0 : 
     229            0 :         let out_of_space = remaining < buf.len();
     230            0 : 
     231            0 :         let original_len = buf.len();
     232            0 : 
     233            0 :         if out_of_space {
     234            0 :             let can_still_fit = buf.len() - remaining;
     235            0 :             self.buffer.extend_from_slice(&buf[..can_still_fit]);
     236            0 :             buf = &buf[can_still_fit..];
     237            0 :             self.flush0()?;
     238            0 :         }
     239              : 
     240              :         // assume that this will often under normal operation just move the pointer back to the
     241              :         // beginning of allocation, because previous split off parts are already sent and
     242              :         // dropped.
     243            0 :         self.buffer.extend_from_slice(buf);
     244            0 :         Ok(original_len)
     245            0 :     }
     246              : 
     247            0 :     fn flush(&mut self) -> std::io::Result<()> {
     248            0 :         self.flush0().map(|_| ())
     249            0 :     }
     250              : }
     251              : 
     252            0 : pub async fn prometheus_metrics_handler(_req: Request<Body>) -> Result<Response<Body>, ApiError> {
     253            0 :     SERVE_METRICS_COUNT.inc();
     254            0 : 
     255            0 :     let started_at = std::time::Instant::now();
     256            0 : 
     257            0 :     let (tx, rx) = mpsc::channel(1);
     258            0 : 
     259            0 :     let body = Body::wrap_stream(ReceiverStream::new(rx));
     260            0 : 
     261            0 :     let mut writer = ChannelWriter::new(128 * 1024, tx);
     262            0 : 
     263            0 :     let encoder = TextEncoder::new();
     264            0 : 
     265            0 :     let response = Response::builder()
     266            0 :         .status(200)
     267            0 :         .header(CONTENT_TYPE, encoder.format_type())
     268            0 :         .body(body)
     269            0 :         .unwrap();
     270              : 
     271            0 :     let span = info_span!("blocking");
     272            0 :     tokio::task::spawn_blocking(move || {
     273            0 :         // there are situations where we lose scraped metrics under load, try to gather some clues
     274            0 :         // since all nodes are queried this, keep the message count low.
     275            0 :         let spawned_at = std::time::Instant::now();
     276            0 : 
     277            0 :         let _span = span.entered();
     278            0 : 
     279            0 :         let metrics = metrics::gather();
     280            0 : 
     281            0 :         let gathered_at = std::time::Instant::now();
     282            0 : 
     283            0 :         let res = encoder
     284            0 :             .encode(&metrics, &mut writer)
     285            0 :             .and_then(|_| writer.flush().map_err(|e| e.into()));
     286            0 : 
     287            0 :         // this instant is not when we finally got the full response sent, sending is done by hyper
     288            0 :         // in another task.
     289            0 :         let encoded_at = std::time::Instant::now();
     290            0 : 
     291            0 :         let spawned_in = spawned_at - started_at;
     292            0 :         let collected_in = gathered_at - spawned_at;
     293            0 :         // remove the wait time here in case the tcp connection was clogged
     294            0 :         let encoded_in = encoded_at - gathered_at - writer.wait_time();
     295            0 :         let total = encoded_at - started_at;
     296            0 : 
     297            0 :         match res {
     298              :             Ok(()) => {
     299            0 :                 tracing::info!(
     300            0 :                     bytes = writer.flushed_bytes(),
     301            0 :                     total_ms = total.as_millis(),
     302            0 :                     spawning_ms = spawned_in.as_millis(),
     303            0 :                     collection_ms = collected_in.as_millis(),
     304            0 :                     encoding_ms = encoded_in.as_millis(),
     305            0 :                     "responded /metrics"
     306              :                 );
     307              :             }
     308            0 :             Err(e) => {
     309            0 :                 // there is a chance that this error is not the BrokenPipe we generate in the writer
     310            0 :                 // for "closed connection", but it is highly unlikely.
     311            0 :                 tracing::warn!(
     312            0 :                     after_bytes = writer.flushed_bytes(),
     313            0 :                     total_ms = total.as_millis(),
     314            0 :                     spawning_ms = spawned_in.as_millis(),
     315            0 :                     collection_ms = collected_in.as_millis(),
     316            0 :                     encoding_ms = encoded_in.as_millis(),
     317            0 :                     "failed to write out /metrics response: {e:?}"
     318              :                 );
     319              :                 // semantics of this error are quite... unclear. we want to error the stream out to
     320              :                 // abort the response to somehow notify the client that we failed.
     321              :                 //
     322              :                 // though, most likely the reason for failure is that the receiver is already gone.
     323            0 :                 drop(
     324            0 :                     writer
     325            0 :                         .tx
     326            0 :                         .blocking_send(Err(std::io::ErrorKind::BrokenPipe.into())),
     327            0 :                 );
     328              :             }
     329              :         }
     330            0 :     });
     331            0 : 
     332            0 :     Ok(response)
     333            0 : }
     334              : 
     335              : /// Generates CPU profiles.
     336            0 : pub async fn profile_cpu_handler(req: Request<Body>) -> Result<Response<Body>, ApiError> {
     337              :     enum Format {
     338              :         Pprof,
     339              :         Svg,
     340              :     }
     341              : 
     342              :     // Parameters.
     343            0 :     let format = match get_query_param(&req, "format")?.as_deref() {
     344            0 :         None => Format::Pprof,
     345            0 :         Some("pprof") => Format::Pprof,
     346            0 :         Some("svg") => Format::Svg,
     347            0 :         Some(format) => return Err(ApiError::BadRequest(anyhow!("invalid format {format}"))),
     348              :     };
     349            0 :     let seconds = match parse_query_param(&req, "seconds")? {
     350            0 :         None => 5,
     351            0 :         Some(seconds @ 1..=60) => seconds,
     352            0 :         Some(_) => return Err(ApiError::BadRequest(anyhow!("duration must be 1-60 secs"))),
     353              :     };
     354            0 :     let frequency_hz = match parse_query_param(&req, "frequency")? {
     355            0 :         None => 99,
     356            0 :         Some(1001..) => return Err(ApiError::BadRequest(anyhow!("frequency must be <=1000 Hz"))),
     357            0 :         Some(frequency) => frequency,
     358              :     };
     359            0 :     let force: bool = parse_query_param(&req, "force")?.unwrap_or_default();
     360              : 
     361              :     // Take the profile.
     362            0 :     static PROFILE_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
     363              :     static PROFILE_CANCEL: Lazy<Notify> = Lazy::new(Notify::new);
     364              : 
     365            0 :     let report = {
     366              :         // Only allow one profiler at a time. If force is true, cancel a running profile (e.g. a
     367              :         // Grafana continuous profile). We use a try_lock() loop when cancelling instead of waiting
     368              :         // for a lock(), to avoid races where the notify isn't currently awaited.
     369            0 :         let _lock = loop {
     370            0 :             match PROFILE_LOCK.try_lock() {
     371            0 :                 Ok(lock) => break lock,
     372            0 :                 Err(_) if force => PROFILE_CANCEL.notify_waiters(),
     373              :                 Err(_) => {
     374            0 :                     return Err(ApiError::Conflict(
     375            0 :                         "profiler already running (use ?force=true to cancel it)".into(),
     376            0 :                     ));
     377              :                 }
     378              :             }
     379            0 :             tokio::time::sleep(Duration::from_millis(1)).await; // don't busy-wait
     380              :         };
     381              : 
     382            0 :         let guard = ProfilerGuardBuilder::default()
     383            0 :             .frequency(frequency_hz)
     384            0 :             .blocklist(&["libc", "libgcc", "pthread", "vdso"])
     385            0 :             .build()
     386            0 :             .map_err(|err| ApiError::InternalServerError(err.into()))?;
     387              : 
     388            0 :         tokio::select! {
     389            0 :             _ = tokio::time::sleep(Duration::from_secs(seconds)) => {},
     390            0 :             _ = PROFILE_CANCEL.notified() => {},
     391              :         };
     392              : 
     393            0 :         guard
     394            0 :             .report()
     395            0 :             .build()
     396            0 :             .map_err(|err| ApiError::InternalServerError(err.into()))?
     397              :     };
     398              : 
     399              :     // Return the report in the requested format.
     400            0 :     match format {
     401              :         Format::Pprof => {
     402            0 :             let body = report
     403            0 :                 .pprof()
     404            0 :                 .map_err(|err| ApiError::InternalServerError(err.into()))?
     405            0 :                 .encode_to_vec();
     406            0 : 
     407            0 :             Response::builder()
     408            0 :                 .status(200)
     409            0 :                 .header(CONTENT_TYPE, "application/octet-stream")
     410            0 :                 .header(CONTENT_DISPOSITION, "attachment; filename=\"profile.pb\"")
     411            0 :                 .body(Body::from(body))
     412            0 :                 .map_err(|err| ApiError::InternalServerError(err.into()))
     413              :         }
     414              : 
     415              :         Format::Svg => {
     416            0 :             let mut body = Vec::new();
     417            0 :             report
     418            0 :                 .flamegraph(&mut body)
     419            0 :                 .map_err(|err| ApiError::InternalServerError(err.into()))?;
     420            0 :             Response::builder()
     421            0 :                 .status(200)
     422            0 :                 .header(CONTENT_TYPE, "image/svg+xml")
     423            0 :                 .body(Body::from(body))
     424            0 :                 .map_err(|err| ApiError::InternalServerError(err.into()))
     425              :         }
     426              :     }
     427            0 : }
     428              : 
     429              : /// Generates heap profiles.
     430              : ///
     431              : /// This only works with jemalloc on Linux.
     432            0 : pub async fn profile_heap_handler(req: Request<Body>) -> Result<Response<Body>, ApiError> {
     433              :     enum Format {
     434              :         Jemalloc,
     435              :         Pprof,
     436              :         Svg,
     437              :     }
     438              : 
     439              :     // Parameters.
     440            0 :     let format = match get_query_param(&req, "format")?.as_deref() {
     441            0 :         None => Format::Pprof,
     442            0 :         Some("jemalloc") => Format::Jemalloc,
     443            0 :         Some("pprof") => Format::Pprof,
     444            0 :         Some("svg") => Format::Svg,
     445            0 :         Some(format) => return Err(ApiError::BadRequest(anyhow!("invalid format {format}"))),
     446              :     };
     447              : 
     448              :     // Obtain profiler handle.
     449            0 :     let mut prof_ctl = jemalloc_pprof::PROF_CTL
     450            0 :         .as_ref()
     451            0 :         .ok_or(ApiError::InternalServerError(anyhow!(
     452            0 :             "heap profiling not enabled"
     453            0 :         )))?
     454            0 :         .lock()
     455            0 :         .await;
     456            0 :     if !prof_ctl.activated() {
     457            0 :         return Err(ApiError::InternalServerError(anyhow!(
     458            0 :             "heap profiling not enabled"
     459            0 :         )));
     460            0 :     }
     461            0 : 
     462            0 :     // Take and return the profile.
     463            0 :     match format {
     464              :         Format::Jemalloc => {
     465              :             // NB: file is an open handle to a tempfile that's already deleted.
     466            0 :             let file = tokio::task::spawn_blocking(move || prof_ctl.dump())
     467            0 :                 .await
     468            0 :                 .map_err(|join_err| ApiError::InternalServerError(join_err.into()))?
     469            0 :                 .map_err(ApiError::InternalServerError)?;
     470            0 :             let stream = ReaderStream::new(tokio::fs::File::from_std(file));
     471            0 :             Response::builder()
     472            0 :                 .status(200)
     473            0 :                 .header(CONTENT_TYPE, "application/octet-stream")
     474            0 :                 .header(CONTENT_DISPOSITION, "attachment; filename=\"heap.dump\"")
     475            0 :                 .body(Body::wrap_stream(stream))
     476            0 :                 .map_err(|err| ApiError::InternalServerError(err.into()))
     477              :         }
     478              : 
     479              :         Format::Pprof => {
     480            0 :             let data = tokio::task::spawn_blocking(move || prof_ctl.dump_pprof())
     481            0 :                 .await
     482            0 :                 .map_err(|join_err| ApiError::InternalServerError(join_err.into()))?
     483            0 :                 .map_err(ApiError::InternalServerError)?;
     484            0 :             Response::builder()
     485            0 :                 .status(200)
     486            0 :                 .header(CONTENT_TYPE, "application/octet-stream")
     487            0 :                 .header(CONTENT_DISPOSITION, "attachment; filename=\"heap.pb.gz\"")
     488            0 :                 .body(Body::from(data))
     489            0 :                 .map_err(|err| ApiError::InternalServerError(err.into()))
     490              :         }
     491              : 
     492              :         Format::Svg => {
     493            0 :             let svg = tokio::task::spawn_blocking(move || prof_ctl.dump_flamegraph())
     494            0 :                 .await
     495            0 :                 .map_err(|join_err| ApiError::InternalServerError(join_err.into()))?
     496            0 :                 .map_err(ApiError::InternalServerError)?;
     497            0 :             Response::builder()
     498            0 :                 .status(200)
     499            0 :                 .header(CONTENT_TYPE, "image/svg+xml")
     500            0 :                 .body(Body::from(svg))
     501            0 :                 .map_err(|err| ApiError::InternalServerError(err.into()))
     502              :         }
     503              :     }
     504            0 : }
     505              : 
     506            2 : pub fn add_request_id_middleware<B: hyper::body::HttpBody + Send + Sync + 'static>()
     507            2 : -> Middleware<B, ApiError> {
     508            2 :     Middleware::pre(move |req| async move {
     509            2 :         let request_id = match req.headers().get(&X_REQUEST_ID_HEADER) {
     510            1 :             Some(request_id) => request_id
     511            1 :                 .to_str()
     512            1 :                 .expect("extract request id value")
     513            1 :                 .to_owned(),
     514              :             None => {
     515            1 :                 let request_id = uuid::Uuid::new_v4();
     516            1 :                 request_id.to_string()
     517              :             }
     518              :         };
     519            2 :         req.set_context(RequestId(request_id));
     520            2 : 
     521            2 :         Ok(req)
     522            2 :     })
     523            2 : }
     524              : 
     525            2 : async fn add_request_id_header_to_response(
     526            2 :     mut res: Response<Body>,
     527            2 :     req_info: RequestInfo,
     528            2 : ) -> Result<Response<Body>, ApiError> {
     529            2 :     if let Some(request_id) = req_info.context::<RequestId>() {
     530            2 :         if let Ok(request_header_value) = HeaderValue::from_str(&request_id.0) {
     531            2 :             res.headers_mut()
     532            2 :                 .insert(&X_REQUEST_ID_HEADER, request_header_value);
     533            2 :         };
     534            0 :     };
     535              : 
     536            2 :     Ok(res)
     537            2 : }
     538              : 
     539            2 : pub fn make_router() -> RouterBuilder<hyper::Body, ApiError> {
     540            2 :     Router::builder()
     541            2 :         .middleware(add_request_id_middleware())
     542            2 :         .middleware(Middleware::post_with_info(
     543            2 :             add_request_id_header_to_response,
     544            2 :         ))
     545            2 :         .err_handler(route_error_handler)
     546            2 : }
     547              : 
     548            0 : pub fn attach_openapi_ui(
     549            0 :     router_builder: RouterBuilder<hyper::Body, ApiError>,
     550            0 :     spec: &'static [u8],
     551            0 :     spec_mount_path: &'static str,
     552            0 :     ui_mount_path: &'static str,
     553            0 : ) -> RouterBuilder<hyper::Body, ApiError> {
     554            0 :     router_builder
     555            0 :         .get(spec_mount_path,
     556            0 :             move |r| request_span(r, move |_| async move {
     557            0 :                 Ok(Response::builder().body(Body::from(spec)).unwrap())
     558            0 :             })
     559            0 :         )
     560            0 :         .get(ui_mount_path,
     561            0 :              move |r| request_span(r, move |_| async move {
     562            0 :                  Ok(Response::builder().body(Body::from(format!(r#"
     563            0 :                 <!DOCTYPE html>
     564            0 :                 <html lang="en">
     565            0 :                 <head>
     566            0 :                 <title>rweb</title>
     567            0 :                 <link href="https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui.css" rel="stylesheet">
     568            0 :                 </head>
     569            0 :                 <body>
     570            0 :                     <div id="swagger-ui"></div>
     571            0 :                     <script src="https://cdn.jsdelivr.net/npm/swagger-ui-dist@3/swagger-ui-bundle.js" charset="UTF-8"> </script>
     572            0 :                     <script>
     573            0 :                         window.onload = function() {{
     574            0 :                         const ui = SwaggerUIBundle({{
     575            0 :                             "dom_id": "\#swagger-ui",
     576            0 :                             presets: [
     577            0 :                             SwaggerUIBundle.presets.apis,
     578            0 :                             SwaggerUIBundle.SwaggerUIStandalonePreset
     579            0 :                             ],
     580            0 :                             layout: "BaseLayout",
     581            0 :                             deepLinking: true,
     582            0 :                             showExtensions: true,
     583            0 :                             showCommonExtensions: true,
     584            0 :                             url: "{}",
     585            0 :                         }})
     586            0 :                         window.ui = ui;
     587            0 :                     }};
     588            0 :                 </script>
     589            0 :                 </body>
     590            0 :                 </html>
     591            0 :             "#, spec_mount_path))).unwrap())
     592            0 :              })
     593            0 :         )
     594            0 : }
     595              : 
     596            0 : fn parse_token(header_value: &str) -> Result<&str, ApiError> {
     597              :     // header must be in form Bearer <token>
     598            0 :     let (prefix, token) = header_value
     599            0 :         .split_once(' ')
     600            0 :         .ok_or_else(|| ApiError::Unauthorized("malformed authorization header".to_string()))?;
     601            0 :     if prefix != "Bearer" {
     602            0 :         return Err(ApiError::Unauthorized(
     603            0 :             "malformed authorization header".to_string(),
     604            0 :         ));
     605            0 :     }
     606            0 :     Ok(token)
     607            0 : }
     608              : 
     609            0 : pub fn auth_middleware<B: hyper::body::HttpBody + Send + Sync + 'static>(
     610            0 :     provide_auth: fn(&Request<Body>) -> Option<&SwappableJwtAuth>,
     611            0 : ) -> Middleware<B, ApiError> {
     612            0 :     Middleware::pre(move |req| async move {
     613            0 :         if let Some(auth) = provide_auth(&req) {
     614            0 :             match req.headers().get(AUTHORIZATION) {
     615            0 :                 Some(value) => {
     616            0 :                     let header_value = value.to_str().map_err(|_| {
     617            0 :                         ApiError::Unauthorized("malformed authorization header".to_string())
     618            0 :                     })?;
     619            0 :                     let token = parse_token(header_value)?;
     620              : 
     621            0 :                     let data = auth.decode(token).map_err(|err| {
     622            0 :                         warn!("Authentication error: {err}");
     623              :                         // Rely on From<AuthError> for ApiError impl
     624            0 :                         err
     625            0 :                     })?;
     626            0 :                     req.set_context(data.claims);
     627              :                 }
     628              :                 None => {
     629            0 :                     return Err(ApiError::Unauthorized(
     630            0 :                         "missing authorization header".to_string(),
     631            0 :                     ));
     632              :                 }
     633              :             }
     634            0 :         }
     635            0 :         Ok(req)
     636            0 :     })
     637            0 : }
     638              : 
     639            0 : pub fn add_response_header_middleware<B>(
     640            0 :     header: &str,
     641            0 :     value: &str,
     642            0 : ) -> anyhow::Result<Middleware<B, ApiError>>
     643            0 : where
     644            0 :     B: hyper::body::HttpBody + Send + Sync + 'static,
     645            0 : {
     646            0 :     let name =
     647            0 :         HeaderName::from_str(header).with_context(|| format!("invalid header name: {header}"))?;
     648            0 :     let value =
     649            0 :         HeaderValue::from_str(value).with_context(|| format!("invalid header value: {value}"))?;
     650            0 :     Ok(Middleware::post_with_info(
     651            0 :         move |mut response, request_info| {
     652            0 :             let name = name.clone();
     653            0 :             let value = value.clone();
     654            0 :             async move {
     655            0 :                 let headers = response.headers_mut();
     656            0 :                 if headers.contains_key(&name) {
     657            0 :                     warn!(
     658            0 :                         "{} response already contains header {:?}",
     659            0 :                         request_info.uri(),
     660            0 :                         &name,
     661              :                     );
     662            0 :                 } else {
     663            0 :                     headers.insert(name, value);
     664            0 :                 }
     665            0 :                 Ok(response)
     666            0 :             }
     667            0 :         },
     668            0 :     ))
     669            0 : }
     670              : 
     671            0 : pub fn check_permission_with(
     672            0 :     req: &Request<Body>,
     673            0 :     check_permission: impl Fn(&Claims) -> Result<(), AuthError>,
     674            0 : ) -> Result<(), ApiError> {
     675            0 :     match req.context::<Claims>() {
     676            0 :         Some(claims) => Ok(check_permission(&claims)
     677            0 :             .map_err(|_err| ApiError::Forbidden("JWT authentication error".to_string()))?),
     678            0 :         None => Ok(()), // claims is None because auth is disabled
     679              :     }
     680            0 : }
     681              : 
     682              : #[cfg(test)]
     683              : mod tests {
     684              :     use std::future::poll_fn;
     685              :     use std::net::{IpAddr, SocketAddr};
     686              : 
     687              :     use hyper::service::Service;
     688              :     use routerify::RequestServiceBuilder;
     689              : 
     690              :     use super::*;
     691              : 
     692              :     #[tokio::test]
     693            1 :     async fn test_request_id_returned() {
     694            1 :         let builder = RequestServiceBuilder::new(make_router().build().unwrap()).unwrap();
     695            1 :         let remote_addr = SocketAddr::new(IpAddr::from_str("127.0.0.1").unwrap(), 80);
     696            1 :         let mut service = builder.build(remote_addr);
     697            1 :         if let Err(e) = poll_fn(|ctx| service.poll_ready(ctx)).await {
     698            1 :             panic!("request service is not ready: {:?}", e);
     699            1 :         }
     700            1 : 
     701            1 :         let mut req: Request<Body> = Request::default();
     702            1 :         req.headers_mut()
     703            1 :             .append(&X_REQUEST_ID_HEADER, HeaderValue::from_str("42").unwrap());
     704            1 : 
     705            1 :         let resp: Response<hyper::body::Body> = service.call(req).await.unwrap();
     706            1 : 
     707            1 :         let header_val = resp.headers().get(&X_REQUEST_ID_HEADER).unwrap();
     708            1 : 
     709            1 :         assert!(header_val == "42", "response header mismatch");
     710            1 :     }
     711              : 
     712              :     #[tokio::test]
     713            1 :     async fn test_request_id_empty() {
     714            1 :         let builder = RequestServiceBuilder::new(make_router().build().unwrap()).unwrap();
     715            1 :         let remote_addr = SocketAddr::new(IpAddr::from_str("127.0.0.1").unwrap(), 80);
     716            1 :         let mut service = builder.build(remote_addr);
     717            1 :         if let Err(e) = poll_fn(|ctx| service.poll_ready(ctx)).await {
     718            1 :             panic!("request service is not ready: {:?}", e);
     719            1 :         }
     720            1 : 
     721            1 :         let req: Request<Body> = Request::default();
     722            1 :         let resp: Response<hyper::body::Body> = service.call(req).await.unwrap();
     723            1 : 
     724            1 :         let header_val = resp.headers().get(&X_REQUEST_ID_HEADER);
     725            1 : 
     726            1 :         assert_ne!(header_val, None, "response header should NOT be empty");
     727            1 :     }
     728              : }
        

Generated by: LCOV version 2.1-beta