Add sampling filter action support

Signed-off-by: Mohamed Mahmoud <[email protected]>
netobserv · Dec 10, 2024 · 2ed36dc · 2ed36dc
1 parent 3eef923
commit 2ed36dc
Show file tree

Hide file tree

Showing 20 changed files with 47 additions and 17 deletions.
diff --git a/bpf/flows.c b/bpf/flows.c
@@ -70,12 +70,6 @@ static inline void update_existing_flow(flow_metrics *aggregate_flow, pkt_info *
 }
 
 static inline int flow_monitor(struct __sk_buff *skb, u8 direction) {
-    // If sampling is defined, will only parse 1 out of "sampling" flows
-    if (sampling > 1 && (bpf_get_prandom_u32() % sampling) != 0) {
-        do_sampling = 0;
-        return TC_ACT_OK;
-    }
-    do_sampling = 1;
     pkt_info pkt;
     __builtin_memset(&pkt, 0, sizeof(pkt));
 
@@ -99,11 +93,27 @@ 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
-    bool skip = check_and_do_flow_filtering(&id, pkt.flags, 0);
+    u32 filter_sampling = 0;
+    bool skip = check_and_do_flow_filtering(&id, pkt.flags, 0, &filter_sampling);
     if (skip) {
         return TC_ACT_OK;
     }
 
+    if (filter_sampling != 0) {
+        if (filter_sampling > 1 && (bpf_get_prandom_u32() % filter_sampling) != 0) {
+            do_sampling = 0;
+            return TC_ACT_OK;
+        }
+        do_sampling = 1;
+    } else {
+        // If sampling is defined, will only parse 1 out of "sampling" flows
+        if (sampling > 1 && (bpf_get_prandom_u32() % sampling) != 0) {
+            do_sampling = 0;
+            return TC_ACT_OK;
+        }
+        do_sampling = 1;
+    }
+
     int dns_errno = 0;
     if (enable_dns_tracking) {
         dns_errno = track_dns_packet(skb, &pkt);

diff --git a/bpf/flows_filter.h b/bpf/flows_filter.h
@@ -33,7 +33,7 @@ static __always_inline int is_equal_ip(u8 *ip1, u8 *ip2, u8 len) {
 
 static __always_inline int do_flow_filter_lookup(flow_id *id, struct filter_key_t *key,
                                                  filter_action *action, u8 len, u8 offset,
-                                                 u16 flags, u32 drop_reason) {
+                                                 u16 flags, u32 drop_reason, u32 *sampling) {
     int result = 0;
 
     struct filter_value_t *rule = (struct filter_value_t *)bpf_map_lookup_elem(&filter_map, key);
@@ -195,6 +195,11 @@ static __always_inline int do_flow_filter_lookup(flow_id *id, struct filter_key_
                 goto end;
             }
         }
+        u32 sample = rule->sample;
+        if (sample && sampling != NULL) {
+            BPF_PRINTK("sampling action is set to %d\n", sample);
+            *sampling = sample;
+        }
     }
 end:
     BPF_PRINTK("result: %d action %d\n", result, *action);
@@ -232,7 +237,7 @@ static __always_inline int flow_filter_setup_lookup_key(flow_id *id, struct filt
  * check if the flow match filter rule and return >= 1 if the flow is to be dropped
  */
 static __always_inline int is_flow_filtered(flow_id *id, filter_action *action, u16 flags,
-                                            u32 drop_reason) {
+                                            u32 drop_reason, u32 *sampling) {
     struct filter_key_t key;
     u8 len, offset;
     int result = 0;
@@ -246,7 +251,7 @@ static __always_inline int is_flow_filtered(flow_id *id, filter_action *action,
         return result;
     }
 
-    result = do_flow_filter_lookup(id, &key, action, len, offset, flags, drop_reason);
+    result = do_flow_filter_lookup(id, &key, action, len, offset, flags, drop_reason, sampling);
     // we have a match so return
     if (result > 0) {
         return result;
@@ -258,7 +263,7 @@ static __always_inline int is_flow_filtered(flow_id *id, filter_action *action,
         return result;
     }
 
-    return do_flow_filter_lookup(id, &key, action, len, offset, flags, drop_reason);
+    return do_flow_filter_lookup(id, &key, action, len, offset, flags, drop_reason, sampling);
 }
 
 #endif //__FLOWS_FILTER_H__
diff --git a/bpf/network_events_monitoring.h b/bpf/network_events_monitoring.h
@@ -99,7 +99,7 @@ static inline int trace_network_events(struct sk_buff *skb, struct rh_psample_me
     }
 
     // check if this packet need to be filtered if filtering feature is enabled
-    bool skip = check_and_do_flow_filtering(&id, flags, 0);
+    bool skip = check_and_do_flow_filtering(&id, flags, 0, NULL);
     if (skip) {
         return 0;
     }

diff --git a/bpf/pca.h b/bpf/pca.h
@@ -57,7 +57,7 @@ static inline bool validate_pca_filter(struct __sk_buff *skb, direction dir) {
     id.direction = dir;
 
     // check if this packet need to be filtered if filtering feature is enabled
-    bool skip = check_and_do_flow_filtering(&id, pkt.flags, 0);
+    bool skip = check_and_do_flow_filtering(&id, pkt.flags, 0, NULL);
     if (skip) {
         return false;
     }

diff --git a/bpf/pkt_drops.h b/bpf/pkt_drops.h
@@ -63,7 +63,7 @@ static inline int trace_pkt_drop(void *ctx, u8 state, struct sk_buff *skb,
     }
 
     // check if this packet need to be filtered if filtering feature is enabled
-    bool skip = check_and_do_flow_filtering(&id, flags, reason);
+    bool skip = check_and_do_flow_filtering(&id, flags, reason, NULL);
     if (skip) {
         return 0;
     }

diff --git a/bpf/rtt_tracker.h b/bpf/rtt_tracker.h
@@ -61,7 +61,7 @@ static inline int calculate_flow_rtt_tcp(struct sock *sk, struct sk_buff *skb) {
     rtt *= 1000u;
 
     // check if this packet need to be filtered if filtering feature is enabled
-    bool skip = check_and_do_flow_filtering(&id, flags, 0);
+    bool skip = check_and_do_flow_filtering(&id, flags, 0, NULL);
     if (skip) {
         return 0;
     }

diff --git a/bpf/types.h b/bpf/types.h
@@ -240,6 +240,7 @@ struct filter_value_t {
     filter_action action;
     tcp_flags tcpFlags;
     u8 filter_drops;
+    u32 sample;
     u8 ip[IP_MAX_LEN];
 } __attribute__((packed));
 // Force emitting struct filter_value_t into the ELF.

diff --git a/bpf/utils.h b/bpf/utils.h
@@ -178,11 +178,12 @@ static inline int fill_ethhdr(struct ethhdr *eth, void *data_end, pkt_info *pkt)
 /*
  * 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, u16 flags, u32 drop_reason) {
+static inline bool check_and_do_flow_filtering(flow_id *id, u16 flags, u32 drop_reason,
+                                               u32 *sampling) {
     // check if this packet need to be filtered if filtering feature is enabled
     if (enable_flows_filtering || enable_pca) {
         filter_action action = ACCEPT;
-        if (is_flow_filtered(id, &action, flags, drop_reason) != 0 &&
+        if (is_flow_filtered(id, &action, flags, drop_reason, sampling) != 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

diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go
@@ -219,6 +219,7 @@ func FlowsAgent(cfg *Config) (*Flows, error) {
 				FilterPort:            tracer.ConvertFilterPortsToInstr(r.FilterPort, r.FilterPortRange, r.FilterPorts),
 				FilterTCPFlags:        r.FilterTCPFlags,
 				FilterDrops:           r.FilterDrops,
+				FilterSample:          r.FilterSample,
 			})
 		}
 	}

diff --git a/pkg/agent/config.go b/pkg/agent/config.go
@@ -75,6 +75,8 @@ type FlowFilter struct {
 	FilterTCPFlags string `json:"tcp_flags,omitempty"`
 	// FilterDrops allow filtering flows with packet drops, default is false.
 	FilterDrops bool `json:"drops,omitempty"`
+	// FilterSample is the sample rate this matching flow will use
+	FilterSample uint32 `json:"sample,omitempty"`
 }
 
 type Config struct {

diff --git a/pkg/agent/packets_agent.go b/pkg/agent/packets_agent.go
@@ -96,6 +96,7 @@ func PacketsAgent(cfg *Config) (*Packets, error) {
 				FilterPort:            tracer.ConvertFilterPortsToInstr(r.FilterPort, r.FilterPortRange, r.FilterPorts),
 				FilterTCPFlags:        r.FilterTCPFlags,
 				FilterDrops:           r.FilterDrops,
+				FilterSample:          r.FilterSample,
 			})
 		}
 	}

diff --git a/pkg/ebpf/bpf_arm64_bpfel.go b/pkg/ebpf/bpf_arm64_bpfel.go
diff --git a/pkg/ebpf/bpf_arm64_bpfel.o b/pkg/ebpf/bpf_arm64_bpfel.o
diff --git a/pkg/ebpf/bpf_powerpc_bpfel.go b/pkg/ebpf/bpf_powerpc_bpfel.go
diff --git a/pkg/ebpf/bpf_powerpc_bpfel.o b/pkg/ebpf/bpf_powerpc_bpfel.o
diff --git a/pkg/ebpf/bpf_s390_bpfeb.go b/pkg/ebpf/bpf_s390_bpfeb.go
diff --git a/pkg/ebpf/bpf_s390_bpfeb.o b/pkg/ebpf/bpf_s390_bpfeb.o
diff --git a/pkg/ebpf/bpf_x86_bpfel.go b/pkg/ebpf/bpf_x86_bpfel.go
diff --git a/pkg/ebpf/bpf_x86_bpfel.o b/pkg/ebpf/bpf_x86_bpfel.o
diff --git a/pkg/tracer/flow_filter.go b/pkg/tracer/flow_filter.go
@@ -25,6 +25,7 @@ type FilterConfig struct {
 	FilterAction          string
 	FilterTCPFlags        string
 	FilterDrops           bool
+	FilterSample          uint32
 }
 
 type Filter struct {
@@ -165,6 +166,10 @@ func (f *Filter) getFilterValue(config *FilterConfig) (ebpf.BpfFilterValueT, err
 	if config.FilterDrops {
 		val.FilterDrops = 1
 	}
+
+	if config.FilterSample != 0 {
+		val.Sample = config.FilterSample
+	}
 	return val, nil
 }