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

Generated by: LCOV version 2.1-beta