Merge pull request #152 from dfinity/igor/fix-ip

fix(BOUN-1497): fix client IP addr resolution

Author: blind-oracle

Cargo.lock

@@ -3,6 +3,7 @@ name = "ic-gateway"
 version = "0.2.0"
 description = "HTTP-to-IC gateway"
 edition = "2024"
+default-run = "ic-gateway"
 
 [features]
 default = []

Cargo.toml

@@ -33,7 +33,10 @@ use prometheus::{
 use tower_http::compression::CompressionLayer;
 use tracing::info;
 
-use crate::core::{ENV, HOSTNAME};
+use crate::{
+    core::{ENV, HOSTNAME},
+    routing::RemoteAddr,
+};
 
 use crate::routing::{
     CanisterId, RequestCtx,
@@ -167,10 +170,16 @@ pub async fn middleware(
     State(state): State<Arc<HttpMetrics>>,
     Extension(conn_info): Extension<Arc<ConnInfo>>,
     Extension(request_id): Extension<RequestId>,
-    request: Request,
+    mut request: Request,
     next: Next,
 ) -> impl IntoResponse {
+    let remote_addr = request.extensions_mut().remove::<RemoteAddr>();
     let tls_info = request.extensions().get::<Arc<TlsInfo>>().cloned();
+    let country_code = request
+        .extensions_mut()
+        .remove::<CountryCode>()
+        .map(|x| x.0)
+        .unwrap_or_default();
 
     // Prepare to execute the request and count its body size
     let (parts, body) = request.into_parts();
@@ -220,11 +229,6 @@ pub async fn middleware(
     let canister_id = response.extensions_mut().get::<CanisterId>().copied();
     let error_cause = response.extensions_mut().remove::<ErrorCause>();
     let ic_status = response.extensions_mut().remove::<IcResponseStatus>();
-    let country_code = response
-        .extensions_mut()
-        .remove::<CountryCode>()
-        .map(|x| x.0)
-        .unwrap_or_default();
     let status = response.status().as_u16();
 
     // IC request metadata
@@ -338,7 +342,7 @@ pub async fn middleware(
         let conn_rcvd = conn_info.traffic.rcvd();
         let conn_sent = conn_info.traffic.sent();
         let conn_reqs = conn_info.req_count();
-        let remote_addr = conn_info.remote_addr.ip().to_string();
+        let remote_addr = remote_addr.map(|x| x.to_string()).unwrap_or_default();
         let request_id_str = request_id.to_string();
 
         let (ic_http_streaming, ic_http_upgrade) = ic_status.as_ref().map_or((false, false), |x| {

src/metrics/mod.rs

@@ -6,10 +6,11 @@ use axum::{
     middleware::Next,
     response::Response,
 };
-use ic_bn_lib::http::middleware::extract_ip_from_request;
 use maxminddb::geoip2;
 use tracing::warn;
 
+use crate::routing::RemoteAddr;
+
 #[derive(Clone, Debug)]
 pub struct CountryCode(pub String);
 
@@ -44,17 +45,19 @@ pub async fn middleware(
     mut request: Request,
     next: Next,
 ) -> Response {
-    let ip = extract_ip_from_request(&request);
+    let remote_addr = request.extensions().get::<RemoteAddr>().copied();
 
     // Lookup code
-    let country_code = ip.map(|x| geoip.lookup(x));
+    let country_code = remote_addr.and_then(|x| geoip.lookup(*x));
 
     if let Some(v) = &country_code {
         request.extensions_mut().insert(v.clone());
     }
 
+    #[allow(unused_mut)]
     let mut response = next.run(request).await;
 
+    #[cfg(test)]
     if let Some(v) = country_code {
         response.extensions_mut().insert(v);
     }

src/routing/middleware/geoip.rs

@@ -24,7 +24,7 @@ const HEADERS_REMOVE: [HeaderName; 12] = [
     X_IC_COUNTRY_CODE,
 ];
 
-// Add various headers
+// Add various headers to the response
 pub async fn middleware(request: Request, next: Next) -> Response {
     let mut response = next.run(request).await;
 

src/routing/middleware/headers.rs

@@ -8,9 +8,11 @@ use axum::{
 use bytes::Bytes;
 use derive_new::new;
 use http::header::HeaderValue;
-use ic_bn_lib::http::headers::X_REQUEST_ID;
+use ic_bn_lib::http::{headers::X_REQUEST_ID, middleware::extract_ip_from_request};
 use uuid::Uuid;
 
+use crate::routing::RemoteAddr;
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct RequestId(pub Uuid);
 
@@ -42,7 +44,7 @@ pub async fn middleware(
     mut request: Request,
     next: Next,
 ) -> Response {
-    // Try to get & parse incoming UUID if it's there
+    // Try to get & parse incoming request UUID if it's there
     let request_id = if let Some(v) = extract_request_id(&request)
         && state.trust_incoming
     {
@@ -51,16 +53,28 @@ pub async fn middleware(
         RequestId(Uuid::now_v7())
     };
 
+    // Extract client's IP
+    let remote_addr = extract_ip_from_request(&request).map(RemoteAddr);
+    if let Some(v) = remote_addr {
+        request.extensions_mut().insert(v);
+    }
+
     let hdr = HeaderValue::from_maybe_shared(Bytes::from(request_id.to_string())).unwrap();
 
     request.extensions_mut().insert(request_id);
     request.headers_mut().insert(X_REQUEST_ID, hdr.clone());
 
     let mut response = next.run(request).await;
-
-    response.extensions_mut().insert(request_id);
     response.headers_mut().insert(X_REQUEST_ID, hdr);
 
+    #[cfg(test)]
+    {
+        response.extensions_mut().insert(request_id);
+        if let Some(v) = remote_addr {
+            response.extensions_mut().insert(v);
+        }
+    }
+
     response
 }
 

src/routing/middleware/request_id.rs

@@ -4,7 +4,7 @@ pub mod ic;
 pub mod middleware;
 pub mod proxy;
 
-use std::{ops::Deref, str::FromStr, sync::Arc, time::Duration};
+use std::{net::IpAddr, ops::Deref, str::FromStr, sync::Arc, time::Duration};
 
 use anyhow::{Context, Error};
 use axum::{
@@ -181,6 +181,18 @@ impl TypeExtractor for RequestTypeExtractor {
     }
 }
 
+/// Client address
+#[derive(Debug, Clone, Copy)]
+pub struct RemoteAddr(pub IpAddr);
+
+impl Deref for RemoteAddr {
+    type Target = IpAddr;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
 // TODO: make it less horrible by using maybe builder pattern or just a struct
 #[allow(clippy::too_many_arguments)]
 #[allow(clippy::cognitive_complexity)]
@@ -534,12 +546,12 @@ pub async fn setup_router(
             request_type_state,
             request_type::middleware,
         ))
+        .layer(geoip_mw)
         .layer(metrics_mw)
         .layer(option_layer(waf_layer))
         .layer(load_shedder_system_mw)
         .layer(from_fn_with_state(validate_state, validate::middleware))
         .layer(concurrency_limit_mw)
-        .layer(geoip_mw)
         .layer(load_shedder_latency_mw);
 
     let api_hostname = cli.api.api_hostname.clone().map(|x| x.to_string());
@@ -634,15 +646,20 @@ pub async fn setup_router(
 
 #[cfg(test)]
 mod test {
-    use crate::test::setup_test_router;
+    use crate::{
+        routing::middleware::{geoip::CountryCode, request_id::RequestId},
+        test::setup_test_router,
+    };
 
     use super::*;
     use axum::body::{Body, to_bytes};
-    use http::Uri;
+    use http::{HeaderValue, Uri};
+    use ic_bn_lib::http::headers::{X_REAL_IP, X_REQUEST_ID};
     use ic_bn_lib_common::types::http::ConnInfo;
     use rand::{seq::SliceRandom, thread_rng};
     use std::str::FromStr;
     use tower::Service;
+    use uuid::Uuid;
 
     #[test]
     fn test_request_type() {
@@ -704,11 +721,30 @@ mod test {
         let mut req = axum::extract::Request::new(Body::from(""));
         *req.uri_mut() = Uri::try_from(format!("http://{domain}")).unwrap();
         let conn_info = Arc::new(ConnInfo::default());
-        (*req.extensions_mut()).insert(conn_info);
+        req.extensions_mut().insert(conn_info);
+        // Some Swiss IP
+        let remote_addr = IpAddr::from_str("77.109.180.4").unwrap();
+        req.headers_mut().insert(
+            X_REAL_IP,
+            HeaderValue::from_str(&remote_addr.to_string()).unwrap(),
+        );
+
+        // Send request id
+        let request_id = Uuid::from_str("7373F02E-9560-4E16-AC6D-4974300C827B").unwrap();
+        req.headers_mut().insert(
+            X_REQUEST_ID,
+            HeaderValue::from_str(&request_id.to_string()).unwrap(),
+        );
 
         // Make sure that we get right answer
         let resp = router.call(req).await.unwrap();
         assert_eq!(resp.status(), StatusCode::OK);
+        assert_eq!(
+            resp.extensions().get::<RemoteAddr>().unwrap().0,
+            remote_addr,
+        );
+        assert_eq!(resp.extensions().get::<CountryCode>().unwrap().0, "CH");
+        assert_eq!(resp.extensions().get::<RequestId>().unwrap().0, request_id);
 
         let body = resp.into_body();
         let body = to_bytes(body, 1024).await.unwrap();

src/routing/mod.rs

@@ -103,6 +103,7 @@ pub async fn setup_test_router(tasks: &mut TaskManager) -> (Router, Vec<String>)
     let args = vec![
         "",
         "--ic-unsafe-disable-response-verification",
+        "--network-trust-x-request-id",
         "--cache-size",
         "2gb",
         "--domain",
@@ -113,6 +114,8 @@ pub async fn setup_test_router(tasks: &mut TaskManager) -> (Router, Vec<String>)
         "test_data/denylist.json",
         "--policy-denylist-allowlist",
         "test_data/allowlist.txt",
+        "--geoip-db",
+        "test_data/GeoLite2-Country.mmdb",
         "--log-vector-url",
         "http://127.0.0.1/vector",
     ];