tetragon: Add cgrouprate processCgroup test

Adding throttle processCgroup test that models possible cases
and checks the throttle stop event is properly sent or not.

Signed-off-by: Jiri Olsa <[email protected]>

pkg/cgrouprate/cgrouprate.go

@@ -97,6 +97,13 @@ func NewCgroupRate(ctx context.Context,
 	go handle.process(ctx)
 }
 
+func NewTestCgroupRate(listener observer.Listener,
+	hash *program.Map,
+	opts *option.CgroupRate) {
+
+	handle = newCgroupRate(listener, hash, opts)
+}
+
 func (r *CgroupRate) notify(msg notify.Message) {
 	if err := r.listener.Notify(msg); err != nil {
 		r.log.WithError(err).Warn("failed to notify listener")

pkg/cgrouprate/cgrouprate_test.go

@@ -0,0 +1,259 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Authors of Tetragon
+
+package cgrouprate
+
+import (
+	"os"
+	"testing"
+	"time"
+	"unsafe"
+
+	"github.com/cilium/ebpf"
+	"github.com/cilium/tetragon/api/v1/tetragon"
+	"github.com/cilium/tetragon/pkg/api/processapi"
+	"github.com/cilium/tetragon/pkg/bpf"
+	"github.com/cilium/tetragon/pkg/option"
+	"github.com/cilium/tetragon/pkg/reader/notify"
+	"github.com/cilium/tetragon/pkg/sensors/program"
+	tus "github.com/cilium/tetragon/pkg/testutils/sensors"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMain(m *testing.M) {
+	ec := tus.TestSensorsRun(m, "SensorExec")
+	os.Exit(ec)
+}
+
+type listener struct {
+	throttle tetragon.ThrottleType
+	cgroup   string
+}
+
+func (l *listener) Notify(msg notify.Message) error {
+	response := msg.HandleMessage()
+	switch response.Event.(type) {
+	case *tetragon.GetEventsResponse_ProcessThrottle:
+		ev := response.GetProcessThrottle()
+		l.throttle = ev.Type
+		l.cgroup = ev.Cgroup
+	}
+	return nil
+}
+
+func (l *listener) Close() error {
+	return nil
+}
+
+type testData struct {
+	opts     option.CgroupRate
+	values   [2]processapi.CgroupRateValue
+	last     uint64
+	throttle tetragon.ThrottleType
+	ret      bool
+}
+
+func TestProcessCgroup(t *testing.T) {
+	key := processapi.CgroupRateKey{
+		Id: 123,
+	}
+
+	cgroup := "cgroup"
+
+	// Test that we get (or not) STOP throttle event, which depends on
+	// wether the cgroup is alive and the rate is below the limit.
+
+	data := []testData{
+		// 0: both rate and last time update are beyond limit on both
+		// cpus - expecting STOP
+		{
+			opts: option.CgroupRate{
+				Events:   10,
+				Interval: uint64(time.Second),
+			},
+			values: [2]processapi.CgroupRateValue{
+				{
+					Rate:      1,
+					Time:      uint64(time.Second),
+					Throttled: uint64(time.Second),
+				},
+				{
+					Rate:      2,
+					Time:      uint64(time.Second),
+					Throttled: uint64(time.Second),
+				},
+			},
+			last: uint64(time.Second) * 8,
+
+			// expecting:
+			throttle: tetragon.ThrottleType_THROTTLE_STOP,
+			ret:      true,
+		},
+		// 1: rate is above limit and last time update is recent enough from
+		// rate time - no event
+		{
+			opts: option.CgroupRate{
+				Events:   10,
+				Interval: uint64(time.Second),
+			},
+			values: [2]processapi.CgroupRateValue{
+				{
+					Rate:      1,
+					Time:      uint64(time.Second),
+					Throttled: uint64(time.Second),
+				},
+				{
+					Rate:      2,
+					Time:      uint64(time.Second),
+					Throttled: uint64(time.Second),
+				},
+			},
+			last: uint64(time.Second),
+
+			// expecting:
+			throttle: tetragon.ThrottleType_THROTTLE_UNKNOWN,
+			ret:      false,
+		},
+		// 2: rate is below limit but last time update is recent enough from
+		// throttle time - expecting no event
+		{
+			opts: option.CgroupRate{
+				Events:   10,
+				Interval: uint64(time.Second),
+			},
+			values: [2]processapi.CgroupRateValue{
+				{
+					Rate:      1,
+					Time:      uint64(time.Second * 3),
+					Throttled: uint64(time.Second),
+				},
+				{
+					Rate:      2,
+					Time:      uint64(time.Second * 3),
+					Throttled: uint64(time.Second),
+				},
+			},
+			last: uint64(time.Second * 3),
+
+			// expecting:
+			throttle: tetragon.ThrottleType_THROTTLE_UNKNOWN,
+			ret:      false,
+		},
+		// 3: rate is below limit but last time update is recent enough from
+		// throttle time on one cpu, the other one is dead - expecting no event
+		{
+			opts: option.CgroupRate{
+				Events:   10,
+				Interval: uint64(time.Second),
+			},
+			values: [2]processapi.CgroupRateValue{
+				{
+					Rate:      0,
+					Time:      uint64(time.Second),
+					Throttled: uint64(time.Second),
+				},
+				{
+					Rate:      2,
+					Time:      uint64(time.Second * 10),
+					Throttled: uint64(time.Second * 5),
+				},
+			},
+			last: uint64(time.Second * 7),
+
+			// expecting:
+			throttle: tetragon.ThrottleType_THROTTLE_UNKNOWN,
+			ret:      false,
+		},
+		// 4: rate is above limit, but the last time update is beyond limit on
+		// both cpus - expecting STOP
+		{
+			opts: option.CgroupRate{
+				Events:   10,
+				Interval: uint64(time.Second),
+			},
+			values: [2]processapi.CgroupRateValue{
+				{
+					Rate:      20,
+					Time:      uint64(time.Second),
+					Throttled: uint64(time.Second),
+				},
+				{
+					Rate:      20,
+					Time:      uint64(time.Second),
+					Throttled: uint64(time.Second),
+				},
+			},
+			last: uint64(time.Second) * 8,
+
+			// expecting:
+			throttle: tetragon.ThrottleType_THROTTLE_STOP,
+			ret:      true,
+		},
+		// 5: rate is below limit and the last time is recent enough on
+		// both cpus - expecting STOP
+		{
+			opts: option.CgroupRate{
+				Events:   10,
+				Interval: uint64(time.Second),
+			},
+			values: [2]processapi.CgroupRateValue{
+				{
+					Rate:      2,
+					Time:      uint64(time.Second * 8),
+					Throttled: uint64(time.Second),
+				},
+				{
+					Rate:      3,
+					Time:      uint64(time.Second * 8),
+					Throttled: uint64(time.Second),
+				},
+			},
+			last: uint64(time.Second * 9),
+
+			// expecting:
+			throttle: tetragon.ThrottleType_THROTTLE_STOP,
+			ret:      true,
+		},
+	}
+
+	values := make([]processapi.CgroupRateValue, bpf.GetNumPossibleCPUs())
+
+	spec := &ebpf.MapSpec{
+		Type:       ebpf.PerCPUHash,
+		KeySize:    uint32(unsafe.Sizeof(key)),
+		ValueSize:  uint32(unsafe.Sizeof(values[0])),
+		MaxEntries: 32768,
+	}
+
+	hash := program.MapBuilder("hash", nil)
+	err := hash.New(spec)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer hash.Close()
+
+	for idx, d := range data {
+		l := &listener{
+			throttle: tetragon.ThrottleType_THROTTLE_UNKNOWN,
+		}
+		NewTestCgroupRate(l, hash, &d.opts)
+
+		// setup cgrouprate cgroup
+		handle.cgroups[key.Id] = cgroup
+		assert.NotEqual(t, nil, handle)
+
+		// store hash values
+		values[0] = d.values[0]
+		values[1] = d.values[1]
+
+		if err := hash.MapHandle.Put(key, values); err != nil {
+			t.Fatal("Can't put:", err)
+		}
+
+		t.Logf("Test %d", idx)
+		ret := handle.processCgroup(key.Id, cgroup, d.last)
+
+		assert.Equal(t, d.ret, ret)
+		assert.Equal(t, d.throttle, l.throttle)
+	}
+}