From ff70e09aa090a140729483783076e7938fad0d39 Mon Sep 17 00:00:00 2001
From: Jonathan Davies <jpds@protonmail.com>
Date: Thu, 9 Feb 2023 13:48:48 +0000
Subject: [PATCH] util/forwarded_headers.rs: Generalized
 handle_forwarded_for_headers() here.

---
 src/util/forwarded_headers.rs | 63 +++++++++++++++++++++++++++++++++++
 src/util/lib.rs               |  1 +
 2 files changed, 64 insertions(+)
 create mode 100644 src/util/forwarded_headers.rs

diff --git a/src/util/forwarded_headers.rs b/src/util/forwarded_headers.rs
new file mode 100644
index 00000000..6ae275aa
--- /dev/null
+++ b/src/util/forwarded_headers.rs
@@ -0,0 +1,63 @@
+use http::{HeaderMap, HeaderValue};
+use std::net::IpAddr;
+use std::str::FromStr;
+
+use crate::error::{Error, OkOrMessage};
+
+pub fn handle_forwarded_for_headers(headers: &HeaderMap<HeaderValue>) -> Result<String, Error> {
+	let forwarded_for_header = headers
+		.get("x-forwarded-for")
+		.ok_or_message("X-Forwarded-For header not provided")?;
+
+	let forwarded_for_ip_str = forwarded_for_header
+		.to_str()
+		.ok_or_message("Error parsing X-Forwarded-For header")?;
+
+	let client_ip = IpAddr::from_str(&forwarded_for_ip_str)
+		.ok_or_message("Valid IP address not found in X-Forwarded-For header")?;
+
+	Ok(client_ip.to_string())
+}
+
+#[cfg(test)]
+mod test {
+	use super::*;
+
+	#[test]
+	fn test_handle_forwarded_for_headers_ipv4_client() {
+		let mut test_headers = HeaderMap::new();
+		test_headers.insert("X-Forwarded-For", "192.0.2.100".parse().unwrap());
+
+		if let Ok(forwarded_ip) = handle_forwarded_for_headers(&test_headers) {
+			assert_eq!(forwarded_ip, "192.0.2.100");
+		}
+	}
+
+	#[test]
+	fn test_handle_forwarded_for_headers_ipv6_client() {
+		let mut test_headers = HeaderMap::new();
+		test_headers.insert("X-Forwarded-For", "2001:db8::f00d:cafe".parse().unwrap());
+
+		if let Ok(forwarded_ip) = handle_forwarded_for_headers(&test_headers) {
+			assert_eq!(forwarded_ip, "2001:db8::f00d:cafe");
+		}
+	}
+
+	#[test]
+	fn test_handle_forwarded_for_headers_invalid_ip() {
+		let mut test_headers = HeaderMap::new();
+		test_headers.insert("X-Forwarded-For", "www.example.com".parse().unwrap());
+
+		let result = handle_forwarded_for_headers(&test_headers);
+		assert!(result.is_err());
+	}
+
+	#[test]
+	fn test_handle_forwarded_for_headers_missing() {
+		let mut test_headers = HeaderMap::new();
+		test_headers.insert("Host", "www.deuxfleurs.fr".parse().unwrap());
+
+		let result = handle_forwarded_for_headers(&test_headers);
+		assert!(result.is_err());
+	}
+}
diff --git a/src/util/lib.rs b/src/util/lib.rs
index d35ca72f..c9110fb2 100644
--- a/src/util/lib.rs
+++ b/src/util/lib.rs
@@ -11,6 +11,7 @@ pub mod data;
 pub mod encode;
 pub mod error;
 pub mod formater;
+pub mod forwarded_headers;
 pub mod metrics;
 pub mod migrate;
 pub mod persister;