add source_ip_from_header option

1 files changed, 18 insertions, 2 deletions

diff --git a/src/main.rs b/src/main.rs
index 7fb07ca..9f395e2 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -40,7 +40,12 @@ use log::{debug, error, info, warn, LevelFilter};
 use modules::{NodeContext, MODULES};
 use quinn::crypto::rustls::QuicServerConfig;
 use std::{
-    collections::HashMap, net::SocketAddr, path::PathBuf, process::exit, str::FromStr, sync::Arc,
+    collections::HashMap,
+    net::{IpAddr, SocketAddr},
+    path::PathBuf,
+    process::exit,
+    str::FromStr,
+    sync::Arc,
 };
 use tokio::{
     fs::File,
@@ -370,7 +375,7 @@ fn error_response(addr: SocketAddr, error: ServiceError) -> Response<BoxBody<Byt
 async fn service(
     state: Arc<State>,
     mut request: Request<BoxBody<Bytes, ServiceError>>,
-    addr: SocketAddr,
+    mut addr: SocketAddr,
     secure: bool,
     listen_addr: SocketAddr,
 ) -> Result<hyper::Response<BoxBody<bytes::Bytes, ServiceError>>, ServiceError> {
@@ -389,6 +394,17 @@ async fn service(
         }
     }
 
+    if config.source_ip_from_header {
+        if let Some(x) = request.headers_mut().remove("x-real-ip") {
+            addr = SocketAddr::new(
+                IpAddr::from_str(x.to_str()?).map_err(|_| ServiceError::InvalidHeader)?,
+                0,
+            );
+        } else {
+            return Err(ServiceError::XRealIPMissing);
+        }
+    }
+
     debug!(
         "{addr} ~> {:?} {}",
         request.headers().get(HOST),