Merge pull request #5697 from Turbo87/metadata-trait

middleware/log_requests: Extract `CustomMetadataRequestExt` trait

src/auth.rs

@@ -1,6 +1,6 @@
 use crate::controllers;
 use crate::db::RequestTransaction;
-use crate::middleware::log_request;
+use crate::middleware::log_request::CustomMetadataRequestExt;
 use crate::middleware::session::RequestSession;
 use crate::models::token::{CrateScope, EndpointScope};
 use crate::models::{ApiToken, User};
@@ -71,9 +71,9 @@ impl AuthCheck {
             }
         }
 
-        log_request::add_custom_metadata(request, "uid", auth.user_id());
+        request.add_custom_metadata("uid", auth.user_id());
         if let Some(id) = auth.api_token_id() {
-            log_request::add_custom_metadata(request, "tokenid", id);
+            request.add_custom_metadata("tokenid", id);
         }
 
         if let Some(ref token) = auth.token {

src/controllers/helpers/pagination.rs

@@ -1,6 +1,6 @@
 use crate::config::Server;
 use crate::controllers::prelude::*;
-use crate::middleware::log_request::add_custom_metadata;
+use crate::middleware::log_request::CustomMetadataRequestExt;
 use crate::models::helpers::with_count::*;
 use crate::util::errors::{bad_request, AppResult};
 use crate::util::request_header;
@@ -95,7 +95,7 @@ impl PaginationOptionsBuilder {
                 }
 
                 if numeric_page > MAX_PAGE_BEFORE_SUSPECTED_BOT {
-                    add_custom_metadata(req, "bot", "suspected");
+                    req.add_custom_metadata("bot", "suspected");
                 }
 
                 // Block large offsets for known violators of the crawler policy
@@ -104,7 +104,7 @@ impl PaginationOptionsBuilder {
                     if numeric_page > config.max_allowed_page_offset
                         && is_useragent_or_ip_blocked(config, req)
                     {
-                        add_custom_metadata(req, "cause", "large page offset");
+                        req.add_custom_metadata("cause", "large page offset");
                         return Err(bad_request("requested page offset is too large"));
                     }
                 }

src/controllers/krate/publish.rs

@@ -17,7 +17,7 @@ use crate::models::{
 };
 use crate::worker;
 
-use crate::middleware::log_request::add_custom_metadata;
+use crate::middleware::log_request::CustomMetadataRequestExt;
 use crate::models::token::EndpointScope;
 use crate::schema::*;
 use crate::util::errors::{cargo_err, AppResult};
@@ -62,8 +62,8 @@ pub fn publish(req: &mut dyn RequestExt) -> EndpointResult {
 
     let new_crate = parse_new_headers(req)?;
 
-    add_custom_metadata(req, "crate_name", new_crate.name.to_string());
-    add_custom_metadata(req, "crate_version", new_crate.vers.to_string());
+    req.add_custom_metadata("crate_name", new_crate.name.to_string());
+    req.add_custom_metadata("crate_version", new_crate.vers.to_string());
 
     let conn = app.primary_database.get()?;
 
@@ -307,7 +307,7 @@ fn count_versions_published_today(krate_id: i32, conn: &PgConnection) -> QueryRe
 fn parse_new_headers(req: &mut dyn RequestExt) -> AppResult<EncodableCrateUpload> {
     // Read the json upload request
     let metadata_length = u64::from(read_le_u32(req.body())?);
-    add_custom_metadata(req, "metadata_length", metadata_length);
+    req.add_custom_metadata("metadata_length", metadata_length);
 
     let max = req.app().config.max_upload_size;
     if metadata_length > max {

src/controllers/version/downloads.rs

@@ -5,7 +5,7 @@
 use super::{extract_crate_name_and_semver, version_and_crate};
 use crate::controllers::prelude::*;
 use crate::db::PoolError;
-use crate::middleware::log_request::add_custom_metadata;
+use crate::middleware::log_request::CustomMetadataRequestExt;
 use crate::models::{Crate, VersionDownload};
 use crate::schema::*;
 use crate::views::EncodableVersionDownload;
@@ -109,7 +109,7 @@ pub fn download(req: &mut dyn RequestExt) -> EndpointResult {
         .crate_location(&crate_name, version);
 
     if let Some((key, value)) = log_metadata {
-        add_custom_metadata(req, key, value);
+        req.add_custom_metadata(key, value);
     }
 
     if req.wants_json() {

src/middleware.rs

@@ -1,5 +1,5 @@
 mod prelude {
-    pub use super::log_request::add_custom_metadata;
+    pub use crate::middleware::log_request::CustomMetadataRequestExt;
     pub use conduit::{box_error, Body, Handler, RequestExt};
     pub use conduit_middleware::{AfterResult, AroundMiddleware, BeforeResult, Middleware};
     pub use http::{header, Response, StatusCode};

src/middleware/balance_capacity.rs

@@ -57,11 +57,11 @@ impl BalanceCapacity {
     fn handle_high_load(&self, request: &mut dyn RequestExt, note: &str) -> AfterResult {
         if self.report_only {
             // In report-only mode we serve all requests but add log metadata
-            add_custom_metadata(request, "would_reject", note);
+            request.add_custom_metadata("would_reject", note);
             self.handle(request)
         } else {
             // Reject the request
-            add_custom_metadata(request, "cause", note);
+            request.add_custom_metadata("cause", note);
             let body = "Service temporarily unavailable";
             Response::builder()
                 .status(StatusCode::SERVICE_UNAVAILABLE)
@@ -85,7 +85,7 @@ impl Handler for BalanceCapacity {
 
         // Begin logging total request count so early stages of load increase can be located
         if in_flight_total >= self.log_total_at_count {
-            add_custom_metadata(request, "in_flight_total", in_flight_total);
+            request.add_custom_metadata("in_flight_total", in_flight_total);
         }
 
         // Download requests are always accepted and do not affect the capacity tracking
@@ -99,7 +99,7 @@ impl Handler for BalanceCapacity {
 
         // Begin logging non-download request count so early stages of non-download load increase can be located
         if load >= self.log_at_percentage {
-            add_custom_metadata(request, "in_flight_non_dl_requests", count);
+            request.add_custom_metadata("in_flight_non_dl_requests", count);
         }
 
         // Reject read-only requests as load nears capacity. Bots are likely to send only safe

src/middleware/block_traffic.rs

@@ -48,7 +48,7 @@ impl Handler for BlockTraffic {
             .any(|value| self.blocked_values.iter().any(|v| v == value));
         if has_blocked_value {
             let cause = format!("blocked due to contents of header {}", self.header_name);
-            add_custom_metadata(req, "cause", cause);
+            req.add_custom_metadata("cause", cause);
             let body = format!(
                 "We are unable to process your request at this time. \
                  This usually means that you are in violation of our crawler \

src/middleware/log_request.rs

@@ -26,7 +26,7 @@ impl Middleware for LogRequests {
     fn after(&self, req: &mut dyn RequestExt, res: AfterResult) -> AfterResult {
         if let Err(error) = &res {
             // Move handler error into custom metadata for axum traffic logging
-            add_custom_metadata(req, "error", error.to_string());
+            req.add_custom_metadata("error", error.to_string())
         }
 
         res
@@ -144,14 +144,30 @@ pub async fn log_requests<B>(
 #[derive(Clone, Debug, Deref, Default)]
 pub struct CustomMetadata(Arc<Mutex<Vec<(&'static str, String)>>>);
 
-pub fn add_custom_metadata<V: Display>(req: &dyn RequestExt, key: &'static str, value: V) {
-    if let Some(metadata) = req.extensions().get::<CustomMetadata>() {
-        if let Ok(mut metadata) = metadata.lock() {
-            metadata.push((key, value.to_string()));
+pub trait CustomMetadataRequestExt {
+    fn add_custom_metadata<V: Display>(&self, key: &'static str, value: V) {
+        if let Some(metadata) = self.metadata_extension() {
+            if let Ok(mut metadata) = metadata.lock() {
+                metadata.push((key, value.to_string()));
+            }
         }
+
+        sentry::configure_scope(|scope| scope.set_extra(key, value.to_string().into()));
     }
 
-    sentry::configure_scope(|scope| scope.set_extra(key, value.to_string().into()));
+    fn metadata_extension(&self) -> Option<&CustomMetadata>;
+}
+
+impl CustomMetadataRequestExt for dyn RequestExt + '_ {
+    fn metadata_extension(&self) -> Option<&CustomMetadata> {
+        self.extensions().get::<CustomMetadata>()
+    }
+}
+
+impl<B> CustomMetadataRequestExt for Request<B> {
+    fn metadata_extension(&self) -> Option<&CustomMetadata> {
+        self.extensions().get::<CustomMetadata>()
+    }
 }
 
 #[cfg(test)]

src/middleware/require_user_agent.rs

@@ -31,7 +31,7 @@ impl Handler for RequireUserAgent {
         let has_user_agent = !agent.is_empty() && agent != cdn_user_agent;
         let is_download = req.path().ends_with("download");
         if !has_user_agent && !is_download {
-            add_custom_metadata(req, "cause", "no user agent");
+            req.add_custom_metadata("cause", "no user agent");
             let body = format!(
                 include_str!("no_user_agent_message.txt"),
                 request_header(req, "x-request-id"),

src/router.rs

@@ -5,7 +5,7 @@ use conduit_router::{RequestParams, RouteBuilder, RoutePattern};
 
 use crate::controllers::*;
 use crate::middleware::app::RequestApp;
-use crate::middleware::log_request::add_custom_metadata;
+use crate::middleware::log_request::CustomMetadataRequestExt;
 use crate::util::errors::{std_error, AppError, RouteBlocked};
 use crate::util::EndpointResult;
 use crate::{App, Env};
@@ -200,7 +200,7 @@ impl Handler for C {
             Ok(resp) => Ok(resp),
             Err(e) => {
                 if let Some(cause) = e.cause() {
-                    add_custom_metadata(req, "cause", cause.to_string())
+                    req.add_custom_metadata("cause", cause.to_string())
                 };
                 match e.response() {
                     Some(response) => Ok(response),