From 39377a9378204822cc38ae0d458e9d8da6457231 Mon Sep 17 00:00:00 2001
From: Adam Karpowich <adam.karpowich@datadoghq.com>
Date: Tue, 21 Nov 2023 16:06:30 -0500
Subject: [PATCH] various cleanup

---
 pkg/network/ebpf/c/prebuilt/conntrack.c |  4 ++--
 pkg/network/ebpf/c/prebuilt/conntrack.h |  8 ++++----
 pkg/network/ebpf/c/runtime/conntrack.c  | 18 ++++++++++++------
 pkg/network/ebpf/c/runtime/conntrack.h  | 13 ++++---------
 4 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/pkg/network/ebpf/c/prebuilt/conntrack.c b/pkg/network/ebpf/c/prebuilt/conntrack.c
index 31abe93d1c76e..1b5210dd1b72b 100644
--- a/pkg/network/ebpf/c/prebuilt/conntrack.c
+++ b/pkg/network/ebpf/c/prebuilt/conntrack.c
@@ -21,7 +21,7 @@ int kprobe___nf_conntrack_hash_insert(struct pt_regs* ctx) {
     if (nf_conn_to_conntrack_tuples(ct, &orig, &reply) != 0) {
         return 0;
     }
-    RETURN_IF_NOT_NAT(orig, reply);
+    RETURN_IF_NOT_NAT(&orig, &reply);
 
     bpf_map_update_with_telemetry(conntrack, &orig, &reply, BPF_ANY);
     bpf_map_update_with_telemetry(conntrack, &reply, &orig, BPF_ANY);
@@ -47,7 +47,7 @@ int kprobe_ctnetlink_fill_info(struct pt_regs* ctx) {
         return 0;
     }
 
-    RETURN_IF_NOT_NAT(orig, reply);
+    RETURN_IF_NOT_NAT(&orig, &reply);
 
     bpf_map_update_with_telemetry(conntrack, &orig, &reply, BPF_ANY);
     bpf_map_update_with_telemetry(conntrack, &reply, &orig, BPF_ANY);
diff --git a/pkg/network/ebpf/c/prebuilt/conntrack.h b/pkg/network/ebpf/c/prebuilt/conntrack.h
index 1beba2280d122..29cd7870ca95e 100644
--- a/pkg/network/ebpf/c/prebuilt/conntrack.h
+++ b/pkg/network/ebpf/c/prebuilt/conntrack.h
@@ -34,10 +34,10 @@ offset_ct(ino)
         return 0;                                           \
     }
 
-bool is_conn_nat(conntrack_tuple_t orig, conntrack_tuple_t reply) {
-    return orig.daddr_l != reply.saddr_l || orig.dport != reply.sport || 
-        orig.saddr_l != reply.daddr_l || orig.sport != reply.dport || 
-        orig.daddr_h != reply.saddr_h;
+static __always_inline bool is_conn_nat(const conntrack_tuple_t* orig, const conntrack_tuple_t* reply) {
+    return orig->daddr_l != reply->saddr_l || orig->dport != reply->sport || 
+        orig->saddr_l != reply->daddr_l || orig->sport != reply->dport || 
+        orig->daddr_h != reply->saddr_h;
 }
 
 static __always_inline u32 get_netns(struct nf_conn *ct) {
diff --git a/pkg/network/ebpf/c/runtime/conntrack.c b/pkg/network/ebpf/c/runtime/conntrack.c
index 0c2c82dcab72f..8c50e0e78e800 100644
--- a/pkg/network/ebpf/c/runtime/conntrack.c
+++ b/pkg/network/ebpf/c/runtime/conntrack.c
@@ -27,15 +27,18 @@ SEC("kprobe/__nf_conntrack_hash_insert")
 int kprobe___nf_conntrack_hash_insert(struct pt_regs* ctx) {
     struct nf_conn *ct = (struct nf_conn*)PT_REGS_PARM1(ctx);
 
-    log_debug("kprobe/__nf_conntrack_hash_insert: netns: %u\n", get_netns(&ct->ct_net));
+    u32 status = ct_status(ct);
+    if (!(status&IPS_CONFIRMED) || !(status&IPS_NAT_MASK)) {
+        return 0;
+    }
+
+    log_debug("kprobe/__nf_conntrack_hash_insert: netns: %u, status: %x\n", get_netns(&ct->ct_net), status);
 
     conntrack_tuple_t orig = {}, reply = {};
     if (nf_conn_to_conntrack_tuples(ct, &orig, &reply) != 0) {
         return 0;
     }
 
-    RETURN_IF_NOT_NAT(orig, reply);
-
     bpf_map_update_with_telemetry(conntrack, &orig, &reply, BPF_ANY);
     bpf_map_update_with_telemetry(conntrack, &reply, &orig, BPF_ANY);
     increment_telemetry_registers_count();
@@ -53,15 +56,18 @@ int kprobe_ctnetlink_fill_info(struct pt_regs* ctx) {
 
     struct nf_conn *ct = (struct nf_conn*)PT_REGS_PARM5(ctx);
 
-    log_debug("kprobe/ctnetlink_fill_info: netns: %u\n", get_netns(&ct->ct_net));
+    u32 status = ct_status(ct);
+    if (!(status&IPS_CONFIRMED) || !(status&IPS_NAT_MASK)) {
+        return 0;
+    }
+
+    log_debug("kprobe/__nf_conntrack_hash_insert: netns: %u, status: %x\n", get_netns(&ct->ct_net), status);
 
     conntrack_tuple_t orig = {}, reply = {};
     if (nf_conn_to_conntrack_tuples(ct, &orig, &reply) != 0) {
         return 0;
     }
 
-    RETURN_IF_NOT_NAT(orig, reply);
-
     bpf_map_update_with_telemetry(conntrack, &orig, &reply, BPF_ANY);
     bpf_map_update_with_telemetry(conntrack, &reply, &orig, BPF_ANY);
     increment_telemetry_registers_count();
diff --git a/pkg/network/ebpf/c/runtime/conntrack.h b/pkg/network/ebpf/c/runtime/conntrack.h
index 3b79a01b37074..2fc2d47b15cb0 100644
--- a/pkg/network/ebpf/c/runtime/conntrack.h
+++ b/pkg/network/ebpf/c/runtime/conntrack.h
@@ -28,15 +28,10 @@ static __always_inline u32 get_netns(void *p_net) {
     return net_ns_inum;
 }
 
-#define RETURN_IF_NOT_NAT(orig, reply)                      \
-    if (!is_conn_nat(orig, reply)) {                        \
-        return 0;                                           \
-    }
-
-bool is_conn_nat(conntrack_tuple_t orig, conntrack_tuple_t reply) {
-    return orig.daddr_l != reply.saddr_l || orig.dport != reply.sport || 
-        orig.saddr_l != reply.daddr_l || orig.sport != reply.dport || 
-        orig.daddr_h != reply.saddr_h;
+static __always_inline u32 ct_status(const struct nf_conn *ct) {
+    u32 status = 0;
+    bpf_probe_read_kernel_with_telemetry(&status, sizeof(status), (void *)&ct->status);
+    return status;
 }
 
 static __always_inline int nf_conn_to_conntrack_tuples(struct nf_conn* ct, conntrack_tuple_t* orig, conntrack_tuple_t* reply) {