refactor: adjust upstream request filter

vicanso · Apr 4, 2024 · 8105dcd · 8105dcd
1 parent 38b0774
commit 8105dcd
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 31 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git 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
@@ -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
@@ -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>
   );