diff --git a/src/main.rs b/src/main.rs index 64b9086..bc8b6bf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,4 @@ -use std::{convert::Infallible, future::Future, net::SocketAddr, pin::Pin}; +use std::{convert::Infallible, future::Future, net::{IpAddr, SocketAddr}, pin::Pin}; use hyper::{body::Incoming, header::HeaderValue, server::conn::http1, service::Service as ServiceTrait, Request, Response, StatusCode}; use tokio::net::TcpListener; @@ -13,8 +13,18 @@ impl ServiceTrait> for Service { type Error = Infallible; type Future = Pin> + Send>>; - fn call(&self, _req: Request) -> Self::Future { - if let Some(ip) = self.remote_ip { + fn call(&self, req: Request) -> Self::Future { + let mut remote_ip = self.remote_ip; + + if let Some(ip) = req.headers().get("X-Forwarded-For") { + if let Ok(str_ip) = ip.to_str() { + if let Ok(ip) = str_ip.parse::() { + remote_ip = Some(ip); + } + } + } + + if let Some(ip) = remote_ip { let mut res = Response::new(ip.ip().to_string()); res.headers_mut().append("Content-Type", HeaderValue::from_static("text/plain")); Box::pin(async { Ok(res) })