From a15a52cbe56518616200ebc2f953e3c631aad8ff Mon Sep 17 00:00:00 2001
From: containerscrew <containerscrew@users.noreply.github.com>
Date: Thu, 19 Dec 2024 00:10:39 +0100
Subject: [PATCH] Refactoring ipv4 tcp ingress proto

---
 nflux-ebpf/src/egress.rs               |  2 ++
 nflux-ebpf/src/{ingress.rs => main.rs} | 16 ++++++++--------
 nflux-ebpf/src/maps.rs                 |  6 +++---
 nflux/src/main.rs                      |  6 +++---
 4 files changed, 16 insertions(+), 14 deletions(-)
 rename nflux-ebpf/src/{ingress.rs => main.rs} (93%)

diff --git a/nflux-ebpf/src/egress.rs b/nflux-ebpf/src/egress.rs
index bb183cd..f8bcb7d 100644
--- a/nflux-ebpf/src/egress.rs
+++ b/nflux-ebpf/src/egress.rs
@@ -6,6 +6,8 @@ use network_types::eth::{EthHdr, EtherType};
 use network_types::ip::Ipv4Hdr;
 use nflux_common::EgressEvent;
 
+use crate::maps::{ACTIVE_CONNECTIONS, EGRESS_EVENT};
+
 
 pub fn try_tc_egress(ctx: TcContext) -> Result<i32, ()> {
     let ethhdr: EthHdr = ctx.load(0).map_err(|_| ())?;
diff --git a/nflux-ebpf/src/ingress.rs b/nflux-ebpf/src/main.rs
similarity index 93%
rename from nflux-ebpf/src/ingress.rs
rename to nflux-ebpf/src/main.rs
index 7f977d0..9103bcc 100644
--- a/nflux-ebpf/src/ingress.rs
+++ b/nflux-ebpf/src/main.rs
@@ -12,7 +12,7 @@ use aya_ebpf::{
     macros::xdp,
     programs::XdpContext,
 };
-use maps::{CONNECTION_EVENTS, CONNECTION_TRACKER, ICMP_RULE, IPV4_RULES};
+use maps::{ACTIVE_CONNECTIONS, CONNECTION_EVENTS, CONNECTION_TRACKER, ICMP_RULE, IPV4_RULES};
 use core::mem;
 use aya_ebpf::bindings::TC_ACT_SHOT;
 use aya_ebpf::macros::classifier;
@@ -97,13 +97,6 @@ fn process_ipv4(ctx: &XdpContext) -> Result<u32, ()> {
 
                     let connection_key = ((source_ip as u64) << 32) | (dst_port as u64);
 
-                    // Check if the incoming connection is part of an active egress connection
-                    if let Some(source_ip) = unsafe { ACTIVE_CONNECTIONS.get(&source_ip) } {
-                        info!(&ctx, "active connection: {:i}", source_ip);
-                        log_new_connection(ctx, source_ip, dst_port, IpProto::Tcp as u8, 1);
-                        return Ok(XDP_PASS); // Allow response to active connection
-                    }
-
                     if rule.ports.contains(&dst_port) && rule.action == 1 {
                         // Allow new connection initiation if permitted by rules
                         if let Some(_) = unsafe { CONNECTION_TRACKER.get(&connection_key) } {
@@ -116,6 +109,13 @@ fn process_ipv4(ctx: &XdpContext) -> Result<u32, ()> {
                         }
                     }
 
+                    // Check if the incoming connection is part of an active egress connection
+                    if let Some(&source_ip) = unsafe { ACTIVE_CONNECTIONS.get(&source_ip) } {
+                        //info!(&ctx, "active connection: {:i}", source_ip);
+                        //log_new_connection(ctx, source_ip, dst_port, IpProto::Tcp as u8, 1);
+                        return Ok(XDP_PASS); // Allow response to active connection
+                    }
+
                     // Drop packets that are not part of any known connection
                     log_new_connection(ctx, source_ip, dst_port, IpProto::Tcp as u8, 0);
                     return Ok(XDP_DROP);
diff --git a/nflux-ebpf/src/maps.rs b/nflux-ebpf/src/maps.rs
index dc17b7d..2f3ebda 100644
--- a/nflux-ebpf/src/maps.rs
+++ b/nflux-ebpf/src/maps.rs
@@ -1,5 +1,5 @@
 use aya_ebpf::{macros::map, maps::{Array, LpmTrie, LruHashMap, PerfEventArray}};
-use nflux_common::{ConnectionEvent, IpRule, LpmKeyIpv4, LpmKeyIpv6};
+use nflux_common::{ConnectionEvent, EgressEvent, IpRule, LpmKeyIpv4, LpmKeyIpv6};
 
 #[map]
 pub static IPV4_RULES: LpmTrie<LpmKeyIpv4, IpRule> = LpmTrie::with_max_entries(1024, 0);
@@ -17,7 +17,7 @@ pub static CONNECTION_EVENTS: PerfEventArray<ConnectionEvent> = PerfEventArray::
 pub static CONNECTION_TRACKER: LruHashMap<u64, u64> = LruHashMap::with_max_entries(1024, 0);
 
 #[map]
-static ACTIVE_CONNECTIONS: LruHashMap<u32, u32> = LruHashMap::with_max_entries(4096, 0);
+pub static ACTIVE_CONNECTIONS: LruHashMap<u32, u32> = LruHashMap::with_max_entries(4096, 0);
 
 #[map]
-static EGRESS_EVENT: PerfEventArray<EgressEvent> = PerfEventArray::new(0);
+pub static EGRESS_EVENT: PerfEventArray<EgressEvent> = PerfEventArray::new(0);
diff --git a/nflux/src/main.rs b/nflux/src/main.rs
index a9975c4..dbc1c38 100644
--- a/nflux/src/main.rs
+++ b/nflux/src/main.rs
@@ -47,9 +47,9 @@ async fn main() -> anyhow::Result<()> {
 
     // Necessary to debug something in the ebpf code
     // By the moment
-    if let Err(e) = EbpfLogger::init(&mut bpf) {
-        warn!("failed to initialize eBPF logger: {}", e);
-    }
+    // if let Err(e) = EbpfLogger::init(&mut bpf) {
+    //     warn!("failed to initialize eBPF logger: {}", e);
+    // }
 
     // Populate eBPF maps with configuration data
     populate_ip_rules(&mut bpf, &config.ip_rules)?;