Allow flow filtering to coexists with pkt drop and rtt ebpf hooks (#318)

Signed-off-by: Mohamed Mahmoud <[email protected]>

Author: msherif1234

bpf/flows.c

@@ -43,7 +43,6 @@
 #include "flows_filter.h"
 
 static inline int flow_monitor(struct __sk_buff *skb, u8 direction) {
-    filter_action action = ACCEPT;
     // If sampling is defined, will only parse 1 out of "sampling" flows
     if (sampling > 1 && (bpf_get_prandom_u32() % sampling) != 0) {
         return TC_ACT_OK;
@@ -71,54 +70,9 @@ static inline int flow_monitor(struct __sk_buff *skb, u8 direction) {
     id.direction = direction;
 
     // check if this packet need to be filtered if filtering feature is enabled
-    if (enable_flows_filtering) {
-        u32 *filter_counter_p = NULL;
-        u32 initVal = 1, key = 0;
-        if (is_flow_filtered(&id, &action) != 0 && action != MAX_FILTER_ACTIONS) {
-            // we have matching rules follow through the actions to decide if we should accept or reject the flow
-            // and update global counter for both cases
-            u32 reject_key = FILTER_FLOWS_REJECT_KEY, accept_key = FILTER_FLOWS_ACCEPT_KEY;
-            bool skip = false;
-
-            switch (action) {
-            case REJECT:
-                key = reject_key;
-                skip = true;
-                break;
-            case ACCEPT:
-                key = accept_key;
-                break;
-            // should never come here
-            case MAX_FILTER_ACTIONS:
-                return TC_ACT_OK;
-            }
-
-            // update global counter for flows dropped by filter
-            filter_counter_p = bpf_map_lookup_elem(&global_counters, &key);
-            if (!filter_counter_p) {
-                bpf_map_update_elem(&global_counters, &key, &initVal, BPF_ANY);
-            } else {
-                __sync_fetch_and_add(filter_counter_p, 1);
-            }
-            if (skip) {
-                return TC_ACT_OK;
-            }
-        } else {
-            // we have no matching rules so we update global counter for flows that are not matched by any rule
-            key = FILTER_FLOWS_NOMATCH_KEY;
-            filter_counter_p = bpf_map_lookup_elem(&global_counters, &key);
-            if (!filter_counter_p) {
-                bpf_map_update_elem(&global_counters, &key, &initVal, BPF_ANY);
-            } else {
-                __sync_fetch_and_add(filter_counter_p, 1);
-            }
-            // we have accept rule but no match so we can't let mismatched flows in the hashmap table.
-            if (action == ACCEPT || action == MAX_FILTER_ACTIONS) {
-                return TC_ACT_OK;
-            } else {
-                // we have reject rule and no match so we can add the flows to the hashmap table.
-            }
-        }
+    bool skip = check_and_do_flow_filtering(&id);
+    if (skip) {
+        return TC_ACT_OK;
     }
 
     int dns_errno = 0;

bpf/pkt_drops.h

@@ -47,6 +47,12 @@ static inline int trace_pkt_drop(void *ctx, u8 state, struct sk_buff *skb,
         return 0;
     }
 
+    // check if this packet need to be filtered if filtering feature is enabled
+    bool skip = check_and_do_flow_filtering(&id);
+    if (skip) {
+        return 0;
+    }
+
     long ret = 0;
     for (direction dir = INGRESS; dir < MAX_DIRECTION; dir++) {
         id.direction = dir;

bpf/rtt_tracker.h

@@ -124,6 +124,12 @@ static inline int calculate_flow_rtt_tcp(struct sock *sk, struct sk_buff *skb) {
     rtt = BPF_CORE_READ(ts, srtt_us) >> 3;
     rtt *= 1000u;
 
+    // check if this packet need to be filtered if filtering feature is enabled
+    bool skip = check_and_do_flow_filtering(&id);
+    if (skip) {
+        return 0;
+    }
+
     // update flow with rtt info
     id.direction = INGRESS;
     ret = rtt_lookup_and_update_flow(&id, flags, rtt);

bpf/utils.h

@@ -3,6 +3,7 @@
 
 #include "types.h"
 #include "maps_definition.h"
+#include "flows_filter.h"
 
 // sets the TCP header flags for connection information
 static inline void set_flags(struct tcphdr *th, u16 *flags) {
@@ -276,4 +277,62 @@ static inline long pkt_drop_lookup_and_update_flow(struct sk_buff *skb, flow_id
     return -1;
 }
 
+/*
+ * check if flow filter is enabled and if we need to continue processing the packet or not
+ */
+static inline bool check_and_do_flow_filtering(flow_id *id) {
+    // check if this packet need to be filtered if filtering feature is enabled
+    if (enable_flows_filtering) {
+        filter_action action = ACCEPT;
+        u32 *filter_counter_p = NULL;
+        u32 initVal = 1, key = 0;
+        if (is_flow_filtered(id, &action) != 0 && action != MAX_FILTER_ACTIONS) {
+            // we have matching rules follow through the actions to decide if we should accept or reject the flow
+            // and update global counter for both cases
+            u32 reject_key = FILTER_FLOWS_REJECT_KEY, accept_key = FILTER_FLOWS_ACCEPT_KEY;
+            bool skip = false;
+
+            switch (action) {
+            case REJECT:
+                key = reject_key;
+                skip = true;
+                break;
+            case ACCEPT:
+                key = accept_key;
+                break;
+            // should never come here
+            case MAX_FILTER_ACTIONS:
+                return true;
+            }
+
+            // update global counter for flows dropped by filter
+            filter_counter_p = bpf_map_lookup_elem(&global_counters, &key);
+            if (!filter_counter_p) {
+                bpf_map_update_elem(&global_counters, &key, &initVal, BPF_ANY);
+            } else {
+                __sync_fetch_and_add(filter_counter_p, 1);
+            }
+            if (skip) {
+                return true;
+            }
+        } else {
+            // we have no matching rules so we update global counter for flows that are not matched by any rule
+            key = FILTER_FLOWS_NOMATCH_KEY;
+            filter_counter_p = bpf_map_lookup_elem(&global_counters, &key);
+            if (!filter_counter_p) {
+                bpf_map_update_elem(&global_counters, &key, &initVal, BPF_ANY);
+            } else {
+                __sync_fetch_and_add(filter_counter_p, 1);
+            }
+            // we have accept rule but no match so we can't let mismatched flows in the hashmap table.
+            if (action == ACCEPT || action == MAX_FILTER_ACTIONS) {
+                return true;
+            } else {
+                // we have reject rule and no match so we can add the flows to the hashmap table.
+            }
+        }
+    }
+    return false;
+}
+
 #endif // __UTILS_H__