From 8ab6725ce55f2cfea09029c878eb53ea64a5748e Mon Sep 17 00:00:00 2001
From: yuanchao <yuanchao@yunshan.net>
Date: Wed, 18 Jun 2025 20:04:13 +0800
Subject: [PATCH] feat: support gre nat ip

---
 agent/plugins/tunnel/src/lib.rs      |  2 +-
 agent/src/common/decapsulate.rs      |  5 +++-
 agent/src/flow_generator/flow_map.rs | 39 +++++++++++++++++++++++++++-
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/agent/plugins/tunnel/src/lib.rs b/agent/plugins/tunnel/src/lib.rs
index 1c5f0f5144b..5861d158b86 100644
--- a/agent/plugins/tunnel/src/lib.rs
+++ b/agent/plugins/tunnel/src/lib.rs
@@ -43,6 +43,6 @@ pub fn decapsulate_tencent_gre(
     _flags: u16,
     _gre_protocol_type: u16,
     _ip_header_size: usize,
-) -> Option<(usize, u32)> {
+) -> Option<(usize, u32, u32)> {
     None
 }
diff --git a/agent/src/common/decapsulate.rs b/agent/src/common/decapsulate.rs
index 1867303b395..193a12e5d50 100644
--- a/agent/src/common/decapsulate.rs
+++ b/agent/src/common/decapsulate.rs
@@ -199,6 +199,7 @@ pub struct TunnelInfo {
     pub mac_src: u32, // lowest 4B
     pub mac_dst: u32, // lowest 4B
     pub id: u32,
+    pub ip: u32,
     pub tunnel_type: TunnelType,
     pub tier: u8,
     pub is_ipv6: bool,
@@ -213,6 +214,7 @@ impl Default for TunnelInfo {
             mac_src: 0,
             mac_dst: 0,
             id: 0,
+            ip: 0,
             tunnel_type: TunnelType::default(),
             tier: 0,
             is_ipv6: false,
@@ -385,7 +387,7 @@ impl TunnelInfo {
         gre_protocol_type: u16,
         ip_header_size: usize,
     ) -> usize {
-        let Some((offset, id)) =
+        let Some((offset, id, ip)) =
             decapsulate_tencent_gre(packet, l2_len, flags, gre_protocol_type, ip_header_size)
         else {
             return 0;
@@ -396,6 +398,7 @@ impl TunnelInfo {
             self.decapsulate_mac(packet);
             self.tunnel_type = TunnelType::TencentGre;
             self.id = id;
+            self.ip = ip;
         }
         self.tier += 1;
 
diff --git a/agent/src/flow_generator/flow_map.rs b/agent/src/flow_generator/flow_map.rs
index 284e8923aef..43e535075a3 100644
--- a/agent/src/flow_generator/flow_map.rs
+++ b/agent/src/flow_generator/flow_map.rs
@@ -19,7 +19,7 @@ use std::{
     cell::RefCell,
     collections::HashSet,
     mem,
-    net::Ipv4Addr,
+    net::{IpAddr, Ipv4Addr},
     num::NonZeroUsize,
     rc::Rc,
     str::FromStr,
@@ -54,6 +54,7 @@ use super::{
 
 use crate::{
     common::{
+        decapsulate::TunnelType,
         ebpf::EbpfType,
         endpoint::{EndpointData, EndpointDataPov, EndpointInfo, EPC_DEEPFLOW, EPC_INTERNET},
         enums::{CaptureNetworkType, EthernetType, HeaderType, IpProtocol, TcpFlags},
@@ -244,6 +245,8 @@ pub struct FlowMap {
     stats_collector: Arc<stats::Collector>,
 
     obfuscate_cache: Option<ObfuscateCache>,
+
+    gre_tunnel_id_maps: LruCache<u32, u8>,
 }
 
 impl FlowMap {
@@ -374,6 +377,9 @@ impl FlowMap {
             stats_collector,
             capacity: config.flow_capacity() as usize,
             size: 0,
+            gre_tunnel_id_maps: LruCache::new(
+                NonZeroUsize::new(config.hash_slots as usize).unwrap(),
+            ),
         }
     }
 
@@ -707,7 +713,38 @@ impl FlowMap {
         (self.policy_getter).lookup(meta_packet, self.id as usize, local_epc_id);
     }
 
+    fn modify_meta_packet(&mut self, meta_packet: &mut MetaPacket) {
+        let Some(tunnel) = meta_packet.tunnel else {
+            return;
+        };
+
+        if tunnel.tunnel_type != TunnelType::TencentGre {
+            return;
+        }
+
+        let lookup_key = &mut meta_packet.lookup_key;
+
+        if tunnel.id > 0 && tunnel.ip > 0 && lookup_key.is_ipv4() {
+            let IpAddr::V4(src_ip) = lookup_key.src_ip else {
+                return;
+            };
+
+            let nat_ip = Ipv4Addr::from(tunnel.ip);
+            if nat_ip == src_ip {
+                self.gre_tunnel_id_maps.put(tunnel.ip, 0);
+            } else {
+                if self.gre_tunnel_id_maps.contains(&tunnel.ip) {
+                    lookup_key.dst_ip = IpAddr::V4(nat_ip);
+                }
+            }
+        }
+    }
+
     pub fn inject_meta_packet(&mut self, config: &Config, meta_packet: &mut MetaPacket) {
+        if config.flow.cloud_gateway_traffic {
+            self.modify_meta_packet(meta_packet);
+        }
+
         if !self.inject_flush_ticker(config, meta_packet.lookup_key.timestamp.into()) {
             self.lookup_without_flow(config, meta_packet);
             return;