From ac17b66a8de4f70d8ce153486f57d67e3ef74488 Mon Sep 17 00:00:00 2001
From: Timo Ley <timo.ley@tilera.dev>
Date: Sat, 8 Jan 2022 19:01:19 +0100
Subject: [PATCH] Implement IP header extraction

---
 Cargo.toml          |  3 ++-
 src/lib/error.rs    | 28 ++++++++++++++++++++++++++++
 src/lib/ipheader.rs | 32 ++++++++++++++++++++++++++++++++
 src/lib/mod.rs      |  4 ++++
 src/main.rs         |  1 +
 5 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 src/lib/error.rs
 create mode 100644 src/lib/ipheader.rs
 create mode 100644 src/lib/mod.rs

diff --git a/Cargo.toml b/Cargo.toml
index 97db21b..1ccfe1f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -24,4 +24,5 @@ headers = "0.3.5"
 url = {version = "2.2.2", features = ["serde"]}
 askama = "0.10"
 urlencoding = "2.1.0"
-thiserror = "1.0.30"
\ No newline at end of file
+thiserror = "1.0.30"
+async-trait = "0.1.51"
\ No newline at end of file
diff --git a/src/lib/error.rs b/src/lib/error.rs
new file mode 100644
index 0000000..7dd016b
--- /dev/null
+++ b/src/lib/error.rs
@@ -0,0 +1,28 @@
+use std::{convert::Infallible, net::AddrParseError};
+
+use axum::{
+    body::{Bytes, Full},
+    response::IntoResponse,
+};
+use hyper::{header::ToStrError, StatusCode};
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum ExtractError {
+    #[error("`{0}` header is missing")]
+    HeaderMissing(String),
+    #[error("Header value contains illegal character: {0}")]
+    ToStr(#[from] ToStrError),
+    #[error("Header value is not a valid IP address: {0}")]
+    IPParse(#[from] AddrParseError),
+}
+
+impl IntoResponse for ExtractError {
+    type Body = Full<Bytes>;
+
+    type BodyError = Infallible;
+
+    fn into_response(self) -> hyper::Response<Self::Body> {
+        (StatusCode::BAD_REQUEST, self.to_string()).into_response()
+    }
+}
diff --git a/src/lib/ipheader.rs b/src/lib/ipheader.rs
new file mode 100644
index 0000000..af24c1d
--- /dev/null
+++ b/src/lib/ipheader.rs
@@ -0,0 +1,32 @@
+use std::{net::IpAddr, str::FromStr};
+
+use async_trait::async_trait;
+use axum::extract::{FromRequest, RequestParts};
+
+use super::error::ExtractError;
+
+pub struct ExtractIP(pub IpAddr);
+
+#[async_trait]
+impl<B> FromRequest<B> for ExtractIP
+where
+    B: Send,
+{
+    type Rejection = ExtractError;
+
+    async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
+        let header = req
+            .headers()
+            .and_then(|headers| headers.get("x-forwarded-for"));
+        let header = header.ok_or(ExtractError::HeaderMissing("X-Forwarded-For".to_string()))?;
+        let mut value = header.to_str()?;
+        let pos = value.chars().position(|r| r == ',');
+        value = match pos {
+            Some(p) => &value[0..p],
+            None => value,
+        };
+        let ip = IpAddr::from_str(value)?;
+
+        Ok(Self(ip))
+    }
+}
diff --git a/src/lib/mod.rs b/src/lib/mod.rs
new file mode 100644
index 0000000..2432b6e
--- /dev/null
+++ b/src/lib/mod.rs
@@ -0,0 +1,4 @@
+mod error;
+mod ipheader;
+
+pub use ipheader::ExtractIP;
diff --git a/src/main.rs b/src/main.rs
index af19ae0..ac1d573 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -14,6 +14,7 @@ mod cdn;
 mod config;
 mod error;
 mod ipfs;
+mod lib;
 mod v1;
 
 #[derive(StructOpt)]