[backport v1.0] pkg/sensors: reduce ratelimit map memory footprint

bpf/process/types/basic.h

@@ -1792,6 +1792,8 @@ do_action_signal(int signal)
  */
 #define KEY_BYTES_PER_ARG 40
 
+#ifdef __LARGE_BPF_PROG
+
 struct ratelimit_key {
 	__u64 func_id;
 	__u64 retprobe_id;
@@ -1806,7 +1808,7 @@ struct ratelimit_value {
 
 struct {
 	__uint(type, BPF_MAP_TYPE_LRU_HASH);
-	__uint(max_entries, 32768);
+	__uint(max_entries, 1); // Agent is resizing this if the feature is needed during kprobe load
 	__type(key, struct ratelimit_key);
 	__type(value, struct ratelimit_value);
 } ratelimit_map SEC(".maps");
@@ -1829,7 +1831,6 @@ struct {
 	__type(value, __u8[sizeof(struct ratelimit_key) + 128]);
 } ratelimit_ro_heap SEC(".maps");
 
-#ifdef __LARGE_BPF_PROG
 static inline __attribute__((always_inline)) bool
 rate_limit(__u64 ratelimit_interval, struct msg_generic_kprobe *e)
 {

pkg/sensors/tracing/generickprobe.go

@@ -66,7 +66,10 @@ const (
 	CharBufErrorTooLarge    = -3
 	CharBufSavedForRetprobe = -4
 
-	stackTraceMapMaxEntries = 32768 // this value could be fine tuned
+	// The following values could be fine tuned if either those feature use too
+	// much kernel memory when enabled.
+	stackTraceMapMaxEntries = 32768
+	ratelimitMapMaxEntries  = 32768
 )
 
 func kprobeCharBufErrorToString(e int32) string {
@@ -138,6 +141,9 @@ type genericKprobe struct {
 	// is there stacktrace defined in the kprobe
 	hasStackTrace bool
 
+	// is there ratelimit defined in the kprobe
+	hasRatelimit bool
+
 	customHandler eventhandler.Handler
 }
 
@@ -217,13 +223,15 @@ func createMultiKprobeSensor(sensorPath string, multiIDs, multiRetIDs []idtable.
 	var maps []*program.Map
 
 	oneKprobeHasStackTrace := false
+	oneKprobeHasRatelimit := false
 	for _, id := range multiIDs {
 		gk, err := genericKprobeTableGet(id)
 		if err != nil {
-			logger.GetLogger().WithField("id", id).WithError(err).Warn("createMultiKprobeSensor: failed to retrieve generic kprobe from table, stacktrace could malfunction")
+			logger.GetLogger().WithField("id", id).WithError(err).Warn("createMultiKprobeSensor: failed to retrieve generic kprobe from table, stacktrace or ratelimit could malfunction")
 			continue
 		}
 		oneKprobeHasStackTrace = oneKprobeHasStackTrace || gk.hasStackTrace
+		oneKprobeHasRatelimit = oneKprobeHasRatelimit || gk.hasRatelimit
 	}
 
 	loadProgName := "bpf_multi_kprobe_v53.o"
@@ -310,6 +318,14 @@ func createMultiKprobeSensor(sensorPath string, multiIDs, multiRetIDs []idtable.
 		maps = append(maps, socktrack)
 	}
 
+	if kernels.EnableLargeProgs() {
+		ratelimitMap := program.MapBuilderPin("ratelimit_map", sensors.PathJoin(pinPath, "ratelimit_map"), load)
+		if oneKprobeHasRatelimit {
+			ratelimitMap.SetMaxEntries(ratelimitMapMaxEntries)
+		}
+		maps = append(maps, ratelimitMap)
+	}
+
 	filterMap.SetMaxEntries(len(multiIDs))
 	configMap.SetMaxEntries(len(multiIDs))
 
@@ -806,6 +822,7 @@ func addKprobe(funcName string, f *v1alpha1.KProbeSpec, in *addKprobeIn, selMaps
 		hasOverride:       selectors.HasOverride(f),
 		customHandler:     in.customHandler,
 		hasStackTrace:     selectorsHaveStackTrace(f.Selectors),
+		hasRatelimit:      selectorsHaveRateLimit(f.Selectors),
 	}
 
 	// Parse Filters into kernel filter logic
@@ -950,6 +967,16 @@ func addKprobe(funcName string, f *v1alpha1.KProbeSpec, in *addKprobeIn, selMaps
 		out.maps = append(out.maps, socktrack)
 	}
 
+	if kernels.EnableLargeProgs() {
+		ratelimitMap := program.MapBuilderPin("ratelimit_map", sensors.PathJoin(pinPath, "ratelimit_map"), load)
+		if kprobeEntry.hasRatelimit {
+			// similarly as for stacktrace, we expand the max size only if
+			// needed to reduce the memory footprint when unused
+			ratelimitMap.SetMaxEntries(ratelimitMapMaxEntries)
+		}
+		out.maps = append(out.maps, ratelimitMap)
+	}
+
 	if setRetprobe {
 		pinRetProg := sensors.PathJoin(pinPath, fmt.Sprintf("%s_ret_prog", kprobeEntry.funcName))
 		loadret := program.Builder(
@@ -1854,3 +1881,14 @@ func selectorsHaveStackTrace(selectors []v1alpha1.KProbeSelector) bool {
 	}
 	return false
 }
+
+func selectorsHaveRateLimit(selectors []v1alpha1.KProbeSelector) bool {
+	for _, selector := range selectors {
+		for _, matchAction := range selector.MatchActions {
+			if len(matchAction.RateLimit) > 0 {
+				return true
+			}
+		}
+	}
+	return false
+}

pkg/sensors/tracing/kprobe_test.go

@@ -19,6 +19,7 @@ import (
 	"sync"
 	"syscall"
 	"testing"
+	"time"
 	"unsafe"
 
 	"github.com/cilium/ebpf"
@@ -5368,6 +5369,125 @@ spec:
 	assert.NoError(t, err)
 }
 
+func testKprobeRateLimit(t *testing.T, rateLimit bool) {
+	hook := `apiVersion: cilium.io/v1alpha1
+kind: TracingPolicy
+metadata:
+  name: "datagram"
+spec:
+  kprobes:
+  - call: "ip_send_skb"
+    syscall: false
+    args:
+    - index: 1
+      type: "skb"
+      label: "datagram"
+    selectors:
+    - matchArgs:
+      - index: 1
+        operator: "DAddr"
+        values:
+        - "127.0.0.1"
+      - index: 1
+        operator: "DPort"
+        values:
+        - "9468"
+      - index: 1
+        operator: "Protocol"
+        values:
+        - "IPPROTO_UDP"
+`
+
+	if rateLimit {
+		hook += `
+      matchActions:
+      - action: Post
+        rateLimit: "5"
+`
+	}
+
+	var doneWG, readyWG sync.WaitGroup
+	defer doneWG.Wait()
+
+	ctx, cancel := context.WithTimeout(context.Background(), tus.Conf().CmdWaitTime)
+	defer cancel()
+
+	createCrdFile(t, hook)
+	obs, err := observertesthelper.GetDefaultObserverWithFile(t, ctx, testConfigFile, tus.Conf().TetragonLib)
+	if err != nil {
+		t.Fatalf("GetDefaultObserverWithFile error: %s", err)
+	}
+	observertesthelper.LoopEvents(ctx, t, &doneWG, &readyWG, obs)
+	readyWG.Wait()
+
+	server := "nc.openbsd"
+	cmdServer := exec.Command(server, "-unvlp", "9468", "-s", "127.0.0.1")
+	assert.NoError(t, cmdServer.Start())
+	time.Sleep(1 * time.Second)
+
+	// Generate 5 datagrams
+	socket, err := net.Dial("udp", "127.0.0.1:9468")
+	if err != nil {
+		fmt.Printf("ERROR dialing socket\n")
+		panic(err)
+	}
+
+	for i := 0; i < 5; i++ {
+		_, err := socket.Write([]byte("data"))
+		if err != nil {
+			fmt.Printf("ERROR writing to socket\n")
+			panic(err)
+		}
+	}
+
+	kpChecker := ec.NewProcessKprobeChecker("datagram-checker").
+		WithFunctionName(sm.Full("ip_send_skb")).
+		WithArgs(ec.NewKprobeArgumentListMatcher().
+			WithOperator(lc.Ordered).
+			WithValues(
+				ec.NewKprobeArgumentChecker().WithLabel(sm.Full("datagram")).
+					WithSkbArg(ec.NewKprobeSkbChecker().
+						WithDaddr(sm.Full("127.0.0.1")).
+						WithDport(9468).
+						WithProtocol(sm.Full("IPPROTO_UDP")),
+					),
+			))
+
+	var checkerSuccess *ec.UnorderedEventChecker
+	var checkerFailure *ec.UnorderedEventChecker
+	if rateLimit {
+		// Rate limit. We should have 1. We shouldn't have 2 (or more)
+		checkerSuccess = ec.NewUnorderedEventChecker(kpChecker)
+		checkerFailure = ec.NewUnorderedEventChecker(kpChecker, kpChecker)
+	} else {
+		// No rate limit. We should have 5. We shouldn't have 6.
+		checkerSuccess = ec.NewUnorderedEventChecker(kpChecker, kpChecker, kpChecker, kpChecker, kpChecker)
+		checkerFailure = ec.NewUnorderedEventChecker(kpChecker, kpChecker, kpChecker, kpChecker, kpChecker, kpChecker)
+	}
+	cmdServer.Process.Kill()
+
+	err = jsonchecker.JsonTestCheck(t, checkerSuccess)
+	assert.NoError(t, err)
+	err = jsonchecker.JsonTestCheckExpect(t, checkerFailure, true)
+	assert.NoError(t, err)
+}
+
+func TestKprobeNoRateLimit(t *testing.T) {
+	if !kernels.EnableLargeProgs() {
+		t.Skip("Test requires kernel 5.4")
+	}
+
+	testKprobeRateLimit(t, false)
+}
+
+func TestKprobeRateLimit(t *testing.T) {
+	if !kernels.EnableLargeProgs() {
+		t.Skip("Test requires kernel 5.4")
+	}
+
+	testKprobeRateLimit(t, true)
+}
+
 func TestKprobeListSyscallDupsRange(t *testing.T) {
 	if !kernels.MinKernelVersion("5.3.0") {
 		t.Skip("TestCopyFd requires at least 5.3.0 version")

tests/vmtests/fetch-data.sh

@@ -4,7 +4,7 @@ set -eu -o pipefail
 
 
 OCIORG=quay.io/lvh-images
-ROOTIMG=$OCIORG/root-images
+ROOTIMG=$OCIORG/root-images:20240717.161638@sha256:62a9890111ab39749792fda4f59c8f736fa350ecaedb0667e3eecbbe790d82ed
 KERNIMG=$OCIORG/kernel-images
 CONTAINER_ENGINE=${CONTAINER_ENGINE:-docker}
 KERNEL_VERS="$@"