read X-Forwarded-For header from trusted proxies

src/http.rs

@@ -1,4 +1,10 @@
-use std::{collections::HashMap, io::{BufRead, BufReader, Error, ErrorKind, Result, Write}, net::{SocketAddr, TcpListener, TcpStream}};
+use std::{
+    collections::HashMap,
+    io::{BufRead, BufReader, Error, ErrorKind, Result, Write},
+    net::{IpAddr, SocketAddr, TcpListener, TcpStream},
+    str::FromStr,
+    sync::Arc,
+};
 
 use chrono::Local;
 
@@ -45,10 +51,11 @@ pub struct Request<'a> {
     headers: HashMap<&'a str, &'a str>,
     query: HashMap<&'a str, &'a str>,
     body: Option<String>,
+    real_address: IpAddr,
 }
 
 impl<'a> Request<'a> {
-    pub fn new(stream: &'a TcpStream, lines: &'a Vec<String>) -> Result<Request<'a>> {
+    pub fn new(stream: &'a TcpStream, lines: &'a Vec<String>, trusted_proxies: Vec<IpAddr>) -> Result<Request<'a>> {
         let request_line = lines[0].as_str();
         let request_line_split: Vec<&str> = request_line.split(" ").collect();
         if request_line_split.len() < 3 {
@@ -88,6 +95,22 @@ impl<'a> Request<'a> {
             }
         }
 
+        let mut real_address = IpAddr::from(stream.peer_addr().unwrap().ip());
+
+        headers.get("X-Forwarded-For").inspect(|address| {
+            match IpAddr::from_str(address) {
+                Ok(address) => {
+                    for proxy in trusted_proxies {
+                        if real_address == proxy {
+                            real_address = address;
+                            break;
+                        }
+                    }
+                }
+                Err(_) => {}
+            }
+        });
+
         let mut body: Option<String> = None;
         if lines.len() > headers.len() + 2 &&
             (method == "POST" || method == "PUT") &&
@@ -104,12 +127,16 @@ impl<'a> Request<'a> {
             headers,
             query,
             body,
+            real_address,
         })
     }
 
     pub fn address(&self) -> Result<SocketAddr> {
         self.stream.peer_addr()
     }
+    pub fn real_address(&self) -> &IpAddr {
+        &self.real_address
+    }
     pub fn path(&self) -> &'a str {
         self.path
     }
@@ -187,11 +214,12 @@ impl<'a> Response<'a> {
 pub struct HttpServer {
     address: String,
     port: u16,
+    trusted_proxies: Arc<Vec<IpAddr>>,
     max_connections: usize,
 }
 
 impl HttpServer {
-    pub fn new(address: String, max_connections: usize) -> HttpServer {
+    pub fn new(address: String, max_connections: usize, trusted_proxies: Vec<IpAddr>) -> HttpServer {
         let mut _address = address.clone();
         let mut _port: u16 = 8080;
         match address.split_once(":") {
@@ -204,6 +232,7 @@ impl HttpServer {
         HttpServer {
             address: _address,
             port: _port,
+            trusted_proxies: Arc::new(trusted_proxies),
             max_connections,
         }
     }
@@ -217,8 +246,9 @@ impl HttpServer {
         for stream in listener.incoming() {
             match stream {
                 Ok(stream) => {
+                    let trusted_proxies = self.trusted_proxies.clone();
                     pool.execute(move || {
-                        HttpServer::handle_client(&stream, handler);
+                        HttpServer::handle_client(&stream, handler, trusted_proxies);
                     });
                 }
                 Err(e) => {
@@ -230,7 +260,7 @@ impl HttpServer {
         Ok(())
     }
 
-    fn handle_client(stream: &TcpStream, handler: HttpHandlerFunc) {
+    fn handle_client(stream: &TcpStream, handler: HttpHandlerFunc, trusted_proxies: Arc<Vec<IpAddr>>) {
         let buf_reader = BufReader::new(stream);
         let http_request: Vec<String> = buf_reader
             .lines()
@@ -238,7 +268,7 @@ impl HttpServer {
             .take_while(|line| !line.is_empty())
             .collect();
 
-        let request = Request::new(stream, &http_request);
+        let request = Request::new(stream, &http_request, trusted_proxies.to_vec());
         if request.is_err() {
             eprintln!("Failed to process request: {}", request.err().unwrap());
             return;
@@ -251,6 +281,8 @@ impl HttpServer {
         match handler(&request, response) {
             Ok(status) => {
                 let end_date = Local::now();
+
+
                 println!(
                     "[{}] {} {} {} - {} {} - {}ms - {} ({})",
                     start_date.format("%Y-%m-%d %H:%M:%S"),
@@ -265,7 +297,7 @@ impl HttpServer {
                     (end_date - start_date).num_milliseconds(),
 
                     request.headers().get("User-Agent").map_or("[]", |v| v),
-                    request.address().unwrap().ip(),
+                    request.real_address().to_string(),
                 );
             }
             Err(e) => {

src/main.rs

@@ -1,5 +1,6 @@
 use std::io::{Result};
-use std::net::{ToSocketAddrs};
+use std::net::{IpAddr, ToSocketAddrs};
+use std::str::FromStr;
 use std::{env};
 
 use mcstatusface::http::{HttpServer, StatusCode};
@@ -46,8 +47,22 @@ env!("CARGO_PKG_VERSION"));
 
     let mut address = "0.0.0.0:8080".to_string();
     if args.len() > 2 { address = args[2].to_string() }
+    let trusted_proxies: Vec<IpAddr> =
+        match env::var("MCSTATUSFACE_TRUSTED_PROXIES") {
+            Ok(envar) => {
+                let mut trusted_proxies: Vec<IpAddr> = Vec::new();
+                for addr in envar.split(",") {
+                    match IpAddr::from_str(addr) {
+                        Ok(addr) => { trusted_proxies.push(addr); }
+                        Err(_) => {}
+                    }
+                }
+                trusted_proxies
+            }
+            Err(_) => { vec![] }
+        };
 
-    HttpServer::new(address, 64).start(|request, mut response| {
+    HttpServer::new(address, 64, trusted_proxies).start(|request, mut response| {
         response.status(StatusCode::OK);
         response.set_header("Content-Type", "text/plain".to_string());
         response.set_header("x-powered-by", "GIRL FUEL".to_string());
@@ -132,7 +147,7 @@ env!("CARGO_PKG_VERSION"));
             // JSON response
             match request.query().get("s") {
                 None => {
-                    response.status(StatusCode::BadRequest);
+                    response.status(StatusCode::OK);
                     response.body("?s=<server address>\n".to_string());
                     return response.send();
                 }