fix: use axum-client-ip to get the real client IP (#240)

WalletConnect · Oct 9, 2023 · 591effe · 591effe
1 parent 46d1797
commit 591effe
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 17 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -26,6 +26,7 @@ wc = { git = "https://github.com/WalletConnect/utils-rs.git", tag = "v0.5.1", fe
 
 tokio = { version = "1", features = ["full"] }
 axum = { version = "0.6", features = ["json", "multipart", "tokio"] }
+axum-client-ip = "0.4"
 tower = "0.4"
 tower-http = { version = "0.4", features = ["trace", "cors", "request-id", "propagate-header", "catch-panic"] }
 hyper = "0.14"

diff --git a/src/handlers/push_message.rs b/src/handlers/push_message.rs
@@ -1,3 +1,5 @@
+#[cfg(feature = "analytics")]
+use axum_client_ip::SecureClientIp;
 use {
     crate::{
         analytics::message_info::MessageInfo,
@@ -23,8 +25,6 @@ use {
     serde::{Deserialize, Serialize},
     std::sync::Arc,
 };
-#[cfg(feature = "analytics")]
-use {axum::extract::ConnectInfo, std::net::SocketAddr};
 
 #[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
 pub struct MessagePayload {
@@ -46,7 +46,7 @@ pub struct PushMessageBody {
 }
 
 pub async fn handler(
-    #[cfg(feature = "analytics")] ConnectInfo(addr): ConnectInfo<SocketAddr>,
+    #[cfg(feature = "analytics")] SecureClientIp(client_ip): SecureClientIp,
     Path((tenant_id, id)): Path<(String, String)>,
     StateExtractor(state): StateExtractor<Arc<AppState>>,
     headers: HeaderMap,
@@ -104,7 +104,7 @@ pub async fn handler(
         tokio::spawn(async move {
             if let Some(analytics) = &state.analytics {
                 let (country, continent, region) = analytics
-                    .lookup_geo_data(addr.ip())
+                    .lookup_geo_data(client_ip)
                     .map_or((None, None, None), |geo| {
                         (geo.country, geo.continent, geo.region)
                     });
@@ -113,7 +113,7 @@ pub async fn handler(
                     %request_id,
                     %tenant_id,
                     client_id = %id,
-                    ip = %addr.ip(),
+                    ip = %client_ip,
                     "loaded geo data"
                 );
 

diff --git a/src/handlers/register_client.rs b/src/handlers/register_client.rs
@@ -1,5 +1,5 @@
 #[cfg(feature = "analytics")]
-use {crate::analytics::client_info::ClientInfo, axum::extract::ConnectInfo, std::net::SocketAddr};
+use {crate::analytics::client_info::ClientInfo, axum_client_ip::SecureClientIp};
 use {
     crate::{
         error::{
@@ -31,7 +31,7 @@ pub struct RegisterBody {
 }
 
 pub async fn handler(
-    #[cfg(feature = "analytics")] ConnectInfo(addr): ConnectInfo<SocketAddr>,
+    #[cfg(feature = "analytics")] SecureClientIp(client_ip): SecureClientIp,
     Path(tenant_id): Path<String>,
     StateExtractor(state): StateExtractor<Arc<AppState>>,
     headers: HeaderMap,
@@ -108,7 +108,7 @@ pub async fn handler(
     tokio::spawn(async move {
         if let Some(analytics) = &state.analytics {
             let (country, continent, region) = analytics
-                .lookup_geo_data(addr.ip())
+                .lookup_geo_data(client_ip)
                 .map_or((None, None, None), |geo| {
                     (geo.country, geo.continent, geo.region)
                 });
@@ -117,7 +117,7 @@ pub async fn handler(
                 %request_id,
                 %tenant_id,
                 %client_id,
-                ip = %addr.ip(),
+                ip = %client_ip,
                 "loaded geo data"
             );
 

diff --git a/src/handlers/single_tenant_wrappers.rs b/src/handlers/single_tenant_wrappers.rs
@@ -1,3 +1,5 @@
+#[cfg(feature = "analytics")]
+use axum_client_ip::SecureClientIp;
 use {
     crate::{
         error::Result,
@@ -13,8 +15,6 @@ use {
     hyper::HeaderMap,
     std::sync::Arc,
 };
-#[cfg(feature = "analytics")]
-use {axum::extract::ConnectInfo, std::net::SocketAddr};
 
 #[cfg(feature = "multitenant")]
 use crate::error::Error::MissingTenantId;
@@ -37,7 +37,7 @@ pub async fn delete_handler(
 }
 
 pub async fn push_handler(
-    #[cfg(feature = "analytics")] addr: ConnectInfo<SocketAddr>,
+    #[cfg(feature = "analytics")] SecureClientIp(client_ip): SecureClientIp,
     Path(id): Path<String>,
     state: StateExtractor<Arc<AppState>>,
     headers: HeaderMap,
@@ -48,7 +48,7 @@ pub async fn push_handler(
 
     #[cfg(all(not(feature = "multitenant"), feature = "analytics"))]
     return crate::handlers::push_message::handler(
-        addr,
+        SecureClientIp(client_ip),
         Path((DEFAULT_TENANT_ID.to_string(), id)),
         state,
         headers,
@@ -67,7 +67,7 @@ pub async fn push_handler(
 }
 
 pub async fn register_handler(
-    #[cfg(feature = "analytics")] addr: ConnectInfo<SocketAddr>,
+    #[cfg(feature = "analytics")] SecureClientIp(client_ip): SecureClientIp,
     state: StateExtractor<Arc<AppState>>,
     headers: HeaderMap,
     body: Json<RegisterBody>,
@@ -77,7 +77,7 @@ pub async fn register_handler(
 
     #[cfg(all(not(feature = "multitenant"), feature = "analytics"))]
     return crate::handlers::register_client::handler(
-        addr,
+        SecureClientIp(client_ip),
         Path(DEFAULT_TENANT_ID.to_string()),
         state,
         headers,

diff --git a/src/lib.rs b/src/lib.rs
@@ -17,6 +17,7 @@ use {
         routing::{delete, get, post},
         Router,
     },
+    axum_client_ip::SecureClientIpSource,
     config::Config,
     hyper::http::Method,
     opentelemetry::{sdk::Resource, KeyValue},
@@ -210,7 +211,8 @@ pub async fn bootstap(mut shutdown: broadcast::Receiver<()>, config: Config) ->
                     hyper::http::header::CONTENT_TYPE,
                     hyper::http::header::AUTHORIZATION,
                 ]),
-        );
+        )
+        .layer(SecureClientIpSource::RightmostXForwardedFor.into_extension());
 
     #[cfg(feature = "multitenant")]
     let app = {