Extract real_ip middleware (#7432)

Author: Turbo87

src/middleware.rs

@@ -6,6 +6,7 @@ mod debug;
 mod ember_html;
 pub mod log_request;
 pub mod normalize_path;
+pub mod real_ip;
 mod require_user_agent;
 pub mod session;
 mod static_or_continue;
@@ -44,6 +45,7 @@ pub fn apply_axum_middleware(state: AppState, router: Router<(), TimeoutBody<Bod
         .layer(TimeoutLayer::new(Duration::from_secs(30)))
         .layer(sentry_tower::NewSentryLayer::new_from_top())
         .layer(sentry_tower::SentryHttpLayer::with_transaction())
+        .layer(from_fn(self::real_ip::middleware))
         .layer(from_fn(log_request::log_requests))
         .layer(CatchPanicLayer::new())
         .layer(from_fn_with_state(

src/middleware/log_request.rs

@@ -4,15 +4,14 @@
 use crate::controllers::util::RequestPartsExt;
 use crate::headers::XRequestId;
 use crate::middleware::normalize_path::OriginalPath;
-use crate::real_ip::process_xff_headers;
+use crate::middleware::real_ip::RealIp;
 use axum::headers::UserAgent;
 use axum::middleware::Next;
 use axum::response::IntoResponse;
 use axum::{Extension, TypedHeader};
 use http::{Method, Request, StatusCode, Uri};
 use parking_lot::Mutex;
 use std::fmt::{self, Display, Formatter};
-use std::net::IpAddr;
 use std::ops::Deref;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
@@ -30,6 +29,7 @@ pub struct RequestMetadata {
     method: Method,
     uri: Uri,
     original_path: Option<Extension<OriginalPath>>,
+    real_ip: Extension<RealIp>,
     user_agent: TypedHeader<UserAgent>,
     request_id: Option<TypedHeader<XRequestId>>,
 }
@@ -40,7 +40,6 @@ pub struct Metadata<'a> {
     cause: Option<&'a CauseField>,
     error: Option<&'a ErrorField>,
     duration: Duration,
-    real_ip: Option<IpAddr>,
     custom_metadata: RequestLog,
 }
 
@@ -73,8 +72,7 @@ impl Display for Metadata<'_> {
             };
         }
 
-        let real_ip = self.real_ip.map(|ip| ip.to_string()).unwrap_or_default();
-        line.add_quoted_field("ip", &real_ip)?;
+        line.add_quoted_field("ip", **self.request.real_ip)?;
 
         let response_time_in_ms = self.duration.as_millis();
         if !is_download_redirect || response_time_in_ms > 0 {
@@ -122,8 +120,6 @@ pub async fn log_requests<B>(
     let custom_metadata = RequestLog::default();
     req.extensions_mut().insert(custom_metadata.clone());
 
-    let real_ip = process_xff_headers(req.headers());
-
     let response = next.run(req).await;
 
     let metadata = Metadata {
@@ -132,7 +128,6 @@ pub async fn log_requests<B>(
         cause: response.extensions().get(),
         error: response.extensions().get(),
         duration: start_instant.elapsed(),
-        real_ip,
         custom_metadata,
     };
 

src/middleware/real_ip.rs

@@ -0,0 +1,22 @@
+use crate::real_ip::process_xff_headers;
+use axum::extract::ConnectInfo;
+use axum::middleware::Next;
+use axum::response::IntoResponse;
+use http::Request;
+use std::net::{IpAddr, SocketAddr};
+
+#[derive(Copy, Clone, Debug, Deref)]
+pub struct RealIp(IpAddr);
+
+pub async fn middleware<B>(
+    ConnectInfo(socket_addr): ConnectInfo<SocketAddr>,
+    mut req: Request<B>,
+    next: Next<B>,
+) -> impl IntoResponse {
+    let xff_ip = process_xff_headers(req.headers());
+    let real_ip = xff_ip.unwrap_or_else(|| socket_addr.ip());
+
+    req.extensions_mut().insert(RealIp(real_ip));
+
+    next.run(req).await
+}

src/tests/util.rs

@@ -29,13 +29,15 @@ use crates_io::models::{ApiToken, CreatedApiToken, User};
 use http::{Method, Request};
 
 use axum::body::Bytes;
+use axum::extract::connect_info::MockConnectInfo;
 use chrono::NaiveDateTime;
 use cookie::Cookie;
 use crates_io::models::token::{CrateScope, EndpointScope};
 use crates_io::util::token::PlainToken;
 use http::header;
 use secrecy::ExposeSecret;
 use std::collections::HashMap;
+use std::net::SocketAddr;
 use tower_service::Service;
 
 mod chaosproxy;
@@ -92,7 +94,12 @@ pub trait RequestHelper {
     /// Run a request that is expected to succeed
     #[track_caller]
     fn run<T>(&self, request: MockRequest) -> Response<T> {
-        let mut router = self.app().router().clone();
+        let router = self.app().router().clone();
+
+        // Add a mock `SocketAddr` to the requests so that the `ConnectInfo`
+        // extractor has something to extract.
+        let mocket_addr = SocketAddr::from(([127, 0, 0, 1], 52381));
+        let mut router = router.layer(MockConnectInfo(mocket_addr));
 
         let rt = tokio::runtime::Builder::new_current_thread()
             .enable_all()