From 8105dcdd66394cbd526eb66d9bd9521d9bfa6c1b Mon Sep 17 00:00:00 2001
From: vicanso <tree.xie@outlook.com>
Date: Thu, 4 Apr 2024 17:47:33 +0800
Subject: [PATCH] refactor: adjust upstream request filter

---
 Cargo.lock                         |  2 +-
 Cargo.toml                         |  2 +-
 src/proxy/server.rs                | 46 +++++++++++++++++++-----------
 web/src/components/main-header.tsx | 28 +++++++++---------
 4 files changed, 47 insertions(+), 31 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8d330a84..a75c5096 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1559,7 +1559,7 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
 
 [[package]]
 name = "pingap"
-version = "0.1.11"
+version = "0.1.12"
 dependencies = [
  "async-trait",
  "base64 0.22.0",
diff --git a/Cargo.toml b/Cargo.toml
index e3222b18..acc58456 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pingap"
-version = "0.1.11"
+version = "0.1.12"
 authors = ["Tree Xie <tree.xie@outlook.com>"]
 edition = "2021"
 categories = ["network-programming", "web-programming::http-server"]
diff --git a/src/proxy/server.rs b/src/proxy/server.rs
index 978e32f1..b08c0e90 100644
--- a/src/proxy/server.rs
+++ b/src/proxy/server.rs
@@ -14,7 +14,7 @@ use http::StatusCode;
 use log::{error, info};
 use memory_stats::memory_stats;
 use once_cell::sync::Lazy;
-use pingora::http::ResponseHeader;
+use pingora::http::{RequestHeader, ResponseHeader};
 use pingora::listeners::TlsSettings;
 use pingora::protocols::http::error_resp;
 use pingora::protocols::Digest;
@@ -394,7 +394,6 @@ impl ProxyHttp for Server {
             // TODO parse error
             let _ = new_path.parse::<http::Uri>().map(|uri| header.set_uri(uri));
         }
-        lo.insert_proxy_headers(header);
         ctx.location_index = Some(location_index);
         if let Some(dir) = lo.upstream.get_directory() {
             let result = dir.handle(session, ctx).await?;
@@ -427,20 +426,6 @@ impl ProxyHttp for Server {
             .new_http_peer(ctx, session)
             .ok_or(pingora::Error::new_str("Upstream not found"))?;
 
-        // add x-forwarded-for
-        if let Some(addr) = utils::get_remote_addr(session) {
-            let value = if let Some(value) =
-                session.get_header(utils::HTTP_HEADER_X_FORWARDED_FOR.clone())
-            {
-                format!("{}, {}", value.to_str().unwrap_or_default(), addr)
-            } else {
-                addr.to_string()
-            };
-            let _ = session
-                .req_header_mut()
-                .insert_header(utils::HTTP_HEADER_X_FORWARDED_FOR.clone(), value);
-        }
-
         Ok(Box::new(peer))
     }
     async fn connected_to_upstream(
@@ -459,6 +444,35 @@ impl ProxyHttp for Server {
         ctx.upstream_address = peer.address().to_string();
         Ok(())
     }
+    async fn upstream_request_filter(
+        &self,
+        session: &mut Session,
+        header: &mut RequestHeader,
+        ctx: &mut Self::CTX,
+    ) -> pingora::Result<()>
+    where
+        Self::CTX: Send + Sync,
+    {
+        // add x-forwarded-for
+        if let Some(addr) = utils::get_remote_addr(session) {
+            let value = if let Some(value) =
+                session.get_header(utils::HTTP_HEADER_X_FORWARDED_FOR.clone())
+            {
+                format!("{}, {}", value.to_str().unwrap_or_default(), addr)
+            } else {
+                addr.to_string()
+            };
+            let _ = header.insert_header(utils::HTTP_HEADER_X_FORWARDED_FOR.clone(), value);
+        }
+
+        if let Some(index) = ctx.location_index {
+            if let Some(lo) = self.locations.get(index) {
+                lo.insert_proxy_headers(header);
+            }
+        }
+
+        Ok(())
+    }
     fn upstream_response_filter(
         &self,
         _session: &mut Session,
diff --git a/web/src/components/main-header.tsx b/web/src/components/main-header.tsx
index 70d0b554..9ffca557 100644
--- a/web/src/components/main-header.tsx
+++ b/web/src/components/main-header.tsx
@@ -9,6 +9,7 @@ import SettingsSuggestIcon from "@mui/icons-material/SettingsSuggest";
 import SwipeableDrawer from "@mui/material/SwipeableDrawer";
 import CardContent from "@mui/material/CardContent";
 import Divider from "@mui/material/Divider";
+import Card from "@mui/material/Card";
 
 import useBasicStore from "../states/basic";
 
@@ -49,20 +50,21 @@ export default function MainHeader() {
           setShowSetting(true);
         }}
       >
-        <CardContent>
-          <Typography gutterBottom variant="h5" component="div">
-            Informations
-          </Typography>
-          <Divider />
-          <Box pt={2}>
-            <Typography gutterBottom variant="body2">
-              Start time: {startAt}
-            </Typography>
-            <Typography gutterBottom variant="body2">
-              Memory: {memory}
+        <Card sx={{ minWidth: 275 }}>
+          <CardContent>
+            <Typography gutterBottom variant="h5" component="div">
+              Informations
             </Typography>
-          </Box>
-        </CardContent>
+            <Box pt={2}>
+              <Typography gutterBottom variant="body2">
+                Start Time: {startAt}
+              </Typography>
+              <Typography gutterBottom variant="body2">
+                Memory: {memory}
+              </Typography>
+            </Box>
+          </CardContent>
+        </Card>
       </SwipeableDrawer>
     </React.Fragment>
   );