reverse proxy ip checking

1 files changed, 51 insertions, 60 deletions

diff --git a/src/types/extract.rs b/src/types/extract.rs
index f21c352..e79aa7a 100644
--- a/src/types/extract.rs
+++ b/src/types/extract.rs
@@ -1,14 +1,15 @@
-use std::{io::Read, net::SocketAddr};
+use std::io::Read;
 
 use axum::{
     async_trait,
     body::HttpBody,
-    extract::{ConnectInfo, FromRequest, FromRequestParts},
+    extract::{FromRequest, FromRequestParts},
     headers::Cookie,
     http::{request::Parts, Request},
     response::Response,
     BoxError, RequestExt, TypedHeader,
 };
+use axum_client_ip::ClientIp;
 use bytes::Bytes;
 use serde::de::DeserializeOwned;
 
@@ -53,6 +54,7 @@ where
 }
 
 pub struct Log;
+
 #[async_trait]
 impl<S, B> FromRequest<S, B> for Log
 where
@@ -63,36 +65,8 @@ where
 {
     type Rejection = Response;
 
-    async fn from_request(mut req: Request<B>, state: &S) -> Result<Self> {
-        let Ok(ConnectInfo(info)) = req.extract_parts::<ConnectInfo<SocketAddr>>().await else {
-            return Ok(Self)
-        };
-        let method = req.method().clone();
-        let path = req
-            .extensions()
-            .get::<RouterURI>()
-            .map_or("", |path| path.0);
-        let uri = req.uri().clone();
-
-        let Ok(bytes) = Bytes::from_request(req, state).await else {
-            console::log(info.ip(), method.clone(), uri.clone(), Some(path.to_string()), None).await;
-            return Ok(Self)
-        };
-
-        let Ok(body) = String::from_utf8(bytes.bytes().flatten().collect()) else {
-            console::log(info.ip(), method.clone(), uri.clone(), Some(path.to_string()), None).await;
-            return Ok(Self)
-        };
-
-        console::log(
-            info.ip(),
-            method.clone(),
-            uri.clone(),
-            Some(path.to_string()),
-            Some(body.to_string()),
-        )
-        .await;
-
+    async fn from_request(req: Request<B>, state: &S) -> Result<Self> {
+        parse_body(req, state).await?;
         Ok(Self)
     }
 }
@@ -110,36 +84,13 @@ where
 {
     type Rejection = Response;
 
-    async fn from_request(mut req: Request<B>, state: &S) -> Result<Self> {
-        let Ok(ConnectInfo(info)) = req.extract_parts::<ConnectInfo<SocketAddr>>().await else {
-            tracing::error!("Failed to read connection info");
-            return Err(ResponseCode::InternalServerError.text("Failed to read connection info"));
-        };
-        let method = req.method().clone();
-        let path = req
-            .extensions()
-            .get::<RouterURI>()
-            .map_or("", |path| path.0);
-        let uri = req.uri().clone();
-
-        let Ok(bytes) = Bytes::from_request(req, state).await else {
-            tracing::error!("Failed to read request body");
-            return Err(ResponseCode::InternalServerError.text("Failed to read request body"));
+    async fn from_request(req: Request<B>, state: &S) -> Result<Self> {
+        
+        let body = match parse_body(req, state).await {
+            Ok(body) => body,
+            Err(err) => return Err(err)
         };
 
-        let Ok(body) = String::from_utf8(bytes.bytes().flatten().collect()) else {
-            return Err(ResponseCode::BadRequest.text("Invalid utf8 body"))
-        };
-
-        console::log(
-            info.ip(),
-            method.clone(),
-            uri.clone(),
-            Some(path.to_string()),
-            Some(body.to_string()),
-        )
-        .await;
-
         let Ok(value) = serde_json::from_str::<T>(&body) else {
             return Err(ResponseCode::BadRequest.text("Invalid request body"))
         };
@@ -172,5 +123,45 @@ pub trait Check {
     }
 }
 
+pub async fn parse_body<S, B>(mut req: Request<B>, state: &S) -> Result<String>
+where
+    B: HttpBody + Sync + Send + 'static,
+    B::Data: Send,
+    B::Error: Into<BoxError>,
+    S: Send + Sync
+{
+    let Ok(ClientIp(ip)) = req.extract_parts::<ClientIp>().await else {
+        tracing::error!("Failed to read client ip");
+        return Err(ResponseCode::InternalServerError.text("Failed to read client ip"));
+    };
+
+    let method = req.method().clone();
+    let uri = req.uri().clone();
+    let path = req
+        .extensions()
+        .get::<RouterURI>()
+        .map_or("", |path| path.0);
+
+    let Ok(bytes) = Bytes::from_request(req, state).await else {
+        tracing::error!("Failed to read request body");
+        return Err(ResponseCode::InternalServerError.text("Failed to read request body"));
+    };
+
+    let Ok(body) = String::from_utf8(bytes.bytes().flatten().collect()) else {
+        return Err(ResponseCode::BadRequest.text("Invalid utf8 body"))
+    };
+
+    console::log(
+        ip,
+        method,
+        uri,
+        Some(path.to_string()),
+        Some(body.to_string()),
+    )
+    .await;
+
+    Ok(body)
+}
+
 #[derive(Clone)]
 pub struct RouterURI(pub &'static str);