fix: multiple xdp attach to one iface

xdp/bandwidth/bandwidth.go

@@ -7,12 +7,9 @@ import (
 
 	"github.com/celestiaorg/bittwister/xdp"
 	"github.com/cilium/ebpf"
-	"github.com/cilium/ebpf/link"
 	"go.uber.org/zap"
 )
 
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go bpf kerns/xdp_bandwidth.c -- -I../headers
-
 type Bandwidth struct {
 	NetworkInterface *net.Interface
 	Limit            int64 // Bytes per second
@@ -22,26 +19,15 @@ type Bandwidth struct {
 var _ xdp.XdpLoader = (*Bandwidth)(nil)
 
 func (b *Bandwidth) Start(ctx context.Context, logger *zap.Logger) {
-	// Load pre-compiled programs into the kernel.
-	objs := bpfObjects{}
-	if err := loadBpfObjects(&objs, nil); err != nil {
-		logger.Error(fmt.Sprintf("loading objects: %v", err))
-		return
-	}
-	defer objs.Close()
-
-	l, err := link.AttachXDP(link.XDPOptions{
-		Program:   objs.XdpBandwidthLimit,
-		Interface: b.NetworkInterface.Index,
-	})
+	x, err := xdp.GetPreparedXdpObject(b.NetworkInterface.Index)
 	if err != nil {
-		logger.Error(fmt.Sprintf("could not attach XDP program: %v", err))
+		logger.Error(fmt.Sprintf("Preparing XDP objects: %v", err))
 		return
 	}
-	defer l.Close()
+	defer x.Close()
 
 	key := uint32(0)
-	err = objs.BandwidthLimitMap.Update(key, b.Limit, ebpf.UpdateAny)
+	err = x.BpfObjs.BandwidthLimitMap.Update(key, b.Limit, ebpf.UpdateAny)
 	if err != nil {
 		logger.Error(fmt.Sprintf("could not update bandwidth limit rate: %v", err))
 		return
@@ -57,6 +43,14 @@ func (b *Bandwidth) Start(ctx context.Context, logger *zap.Logger) {
 	b.ready = true
 	<-ctx.Done()
 
+	// Update the map with a rate of 0 to disable the bandwidth limiter.
+	zero := int64(0)
+	err = x.BpfObjs.BandwidthLimitMap.Update(key, zero, ebpf.UpdateAny)
+	if err != nil {
+		logger.Error(fmt.Sprintf("could not update bandwidth limit rate to zero: %v", err))
+		return
+	}
+
 	b.ready = false
 	logger.Info(fmt.Sprintf("Bandwidth limiter stopped on device %q", b.NetworkInterface.Name))
 }

xdp/kerns/main.c

@@ -0,0 +1,17 @@
+// go:build ignore
+
+#include "xdp_bandwidth.c"
+#include "xdp_packetloss.c"
+
+char _license[] SEC("license") = "GPL";
+
+SEC("xdp")
+int xdp_main(struct xdp_md *ctx)
+{
+  int action = xdp_packetloss(ctx);
+  if (action != XDP_PASS)
+  {
+    return action;
+  }
+  return xdp_bandwidth_limit(ctx);
+}

xdp/bandwidth/kerns/xdp_bandwidth.c → xdp/kerns/xdp_bandwidth.c

@@ -32,11 +32,24 @@ struct
   __uint(max_entries, MAX_MAP_ENTRIES);
 } bandwidth_limit_map SEC(".maps");
 
-SEC("xdp")
 int xdp_bandwidth_limit(struct xdp_md *ctx)
 {
-  __u64 current_timestamp = bpf_ktime_get_ns();
   __u32 key = 0;
+  __u64 *bandwidth_limit_ptr = bpf_map_lookup_elem(&bandwidth_limit_map, &key);
+  if (!bandwidth_limit_ptr)
+  {
+    // if it has not set by the user space program,
+    // or the service is not started yet
+    return XDP_PASS;
+  }
+
+  if (*bandwidth_limit_ptr == 0)
+  {
+    // If the service is stopped
+    return XDP_PASS;
+  }
+
+  __u64 current_timestamp = bpf_ktime_get_ns();
   __u64 *last_time_window_start = bpf_map_lookup_elem(&last_packet_timestamp, &key);
   if (!last_time_window_start)
   {
@@ -65,16 +78,9 @@ int xdp_bandwidth_limit(struct xdp_md *ctx)
     bpf_map_update_elem(&last_packet_timestamp, &key, &current_timestamp, BPF_ANY);
   }
 
-  // Look it up only once for performance purposes
-  static __u64 allowed_bytes = 0; // number of bytes per window
-  if (allowed_bytes == 0)
-  {
-    __u64 *bandwidth_limit_ptr = bpf_map_lookup_elem(&bandwidth_limit_map, &key);
-    if (!bandwidth_limit_ptr)
-      return XDP_ABORTED;
-
-    allowed_bytes = (*bandwidth_limit_ptr / 8 * TIME_WINDOW_SEC); // divide by 8 to convert from bits to bytes
-  }
+  // number of bytes per window
+  // divide by 8 to convert from bits to bytes
+  __u64 allowed_bytes = (*bandwidth_limit_ptr / 8 * TIME_WINDOW_SEC);
 
   __u64 *accumulated_bytes = bpf_map_lookup_elem(&byte_counter, &key);
   if (!accumulated_bytes)
@@ -85,5 +91,3 @@ int xdp_bandwidth_limit(struct xdp_md *ctx)
 
   return XDP_PASS;
 }
-
-char _license[] SEC("license") = "GPL";

xdp/kerns/xdp_packetloss.c

@@ -0,0 +1,38 @@
+// go:build ignore
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <string.h>
+
+#define MAX_MAP_ENTRIES 1
+struct
+{
+  __uint(type, BPF_MAP_TYPE_LRU_HASH);
+  __type(key, __u32);
+  __type(value, __s32);
+  __uint(max_entries, MAX_MAP_ENTRIES);
+} packetloss_rate_map SEC(".maps");
+
+int xdp_packetloss(struct xdp_md *ctx)
+{
+  __u32 key = 0;
+  __s32 *drop_rate_ptr = bpf_map_lookup_elem(&packetloss_rate_map, &key);
+  if (!drop_rate_ptr)
+  {
+    // if it has not set by the user space program,
+    // or the service is not started yet
+    return XDP_PASS;
+  }
+
+  if (*drop_rate_ptr == 0)
+  {
+    // If the service is stopped
+    return XDP_PASS;
+  }
+
+  if (bpf_get_prandom_u32() % 100 < *drop_rate_ptr)
+  {
+    return XDP_DROP;
+  }
+
+  return XDP_PASS;
+}