Merge pull request #821 from 4t145/fix-gateway-get-real-ip

gateway: fix the way to get real ip
ideal-world · Aug 9, 2024 · 0f466e4 · 0f466e4
2 parents cebea5a + 2598a6e
commit 0f466e4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 19 deletions.
diff --git a/backend/gateways/spacegate-plugins/src/plugin/ip_time.rs b/backend/gateways/spacegate-plugins/src/plugin/ip_time.rs
@@ -4,11 +4,11 @@ use std::sync::Arc;
 use ipnet::IpNet;
 use serde::{Deserialize, Serialize};
 use spacegate_shell::hyper::{Request, Response, StatusCode};
-use spacegate_shell::kernel::extension::PeerAddr;
+use spacegate_shell::kernel::extension::OriginalIpAddr;
 use spacegate_shell::kernel::helper_layers::function::Inner;
 use spacegate_shell::plugin::Plugin;
 
-use spacegate_shell::{BoxError, SgBody, SgResponseExt};
+use spacegate_shell::{BoxError, SgBody, SgRequestExt, SgResponseExt};
 
 use tardis::{log, serde_json};
 pub const CODE: &str = "ip-time";
@@ -112,12 +112,9 @@ impl Plugin for IpTimePlugin {
         Ok(plugin)
     }
     async fn call(&self, req: Request<SgBody>, inner: Inner) -> Result<Response<SgBody>, BoxError> {
-        let Some(socket_addr) = req.extensions().get::<PeerAddr>() else {
-            return Err("Cannot get peer address, it's a implementation bug".into());
-        };
-        let socket_addr = socket_addr.0;
-        let passed = self.check_ip(&socket_addr.ip());
-        log::trace!("[{CODE}] Check ip time rule from {socket_addr}, passed {passed}");
+        let original_addr = req.extract::<OriginalIpAddr>().into_inner();
+        let passed = self.check_ip(&original_addr);
+        log::trace!("[{CODE}] Check ip time rule from {original_addr}, passed {passed}");
         if !passed {
             return Ok(Response::with_code_message(StatusCode::FORBIDDEN, "Blocked by ip-time plugin"));
         }

diff --git a/backend/gateways/spacegate-plugins/src/plugin/rewrite_ns_b_ip.rs b/backend/gateways/spacegate-plugins/src/plugin/rewrite_ns_b_ip.rs
@@ -3,11 +3,11 @@ use serde::{Deserialize, Serialize};
 
 use spacegate_shell::extension::k8s_service::K8sService;
 use spacegate_shell::hyper::{http::uri, Response};
-use spacegate_shell::kernel::extension::PeerAddr;
+use spacegate_shell::kernel::extension::OriginalIpAddr;
 use spacegate_shell::kernel::helper_layers::function::Inner;
 use spacegate_shell::kernel::SgRequest;
-use spacegate_shell::plugin::{schemars, Plugin, PluginConfig, PluginError};
-use spacegate_shell::{BoxError, SgBody};
+use spacegate_shell::plugin::{schemars, Plugin, PluginConfig};
+use spacegate_shell::{BoxError, SgBody, SgRequestExt};
 use std::net::IpAddr;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -99,14 +99,9 @@ impl Plugin for RewriteNsPlugin {
 }
 impl RewriteNsPlugin {
     fn req(&self, req: &mut SgRequest) -> Result<(), BoxError> {
-        let ip = req
-            .headers()
-            .get_all("x-forwarded-for")
-            .iter()
-            .next()
-            .and_then(|s| IpAddr::from_str(s.to_str().unwrap_or_default().trim()).ok())
-            .unwrap_or_else(|| req.extensions().get::<PeerAddr>().expect("peer addr should be settled").0.ip());
-        if self.ip_list.iter().any(|ipnet| ipnet.contains(&ip)) {
+        let original_ip = req
+                    .extract::<OriginalIpAddr>().into_inner();
+        if self.ip_list.iter().any(|ipnet| ipnet.contains(&original_ip)) {
             let defer = req.extensions_mut().get_or_insert_default::<spacegate_shell::kernel::extension::Defer>();
             let target_ns = self.target_ns.clone();
             defer.push_back(move |mut req| {