cgrouprate: refactor base registration code

Instead of using base objects directly, it's cleaner to operate on the
contents of the sensor object passed at registration time. Add some
helper functions to identify these objects and use them instead.

Also, refactor the global state so that everything is under a single
object for easier access.

Move everything into a single file for clarity.

Signed-off-by: Kornilios Kourtis <[email protected]>

cmd/tetragon/main.go

@@ -492,7 +492,9 @@ func tetragonExecuteCtx(ctx context.Context, cancel context.CancelFunc, ready fu
 		return err
 	}
 
-	cgrouprate.NewCgroupRate(ctx, pm, &option.Config.CgroupRate)
+	if err := cgrouprate.NewCgroupRate(ctx, pm, &option.Config.CgroupRate); err != nil {
+		return err
+	}
 	cgrouprate.Config()
 
 	err = loadTpFromDir(ctx, option.Config.TracingPolicyDir)

pkg/cgrouprate/cgrouprate.go

@@ -23,6 +23,7 @@ package cgrouprate
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"time"
 
@@ -36,6 +37,8 @@ import (
 	"github.com/cilium/tetragon/pkg/observer"
 	"github.com/cilium/tetragon/pkg/option"
 	"github.com/cilium/tetragon/pkg/reader/notify"
+	"github.com/cilium/tetragon/pkg/sensors"
+	"github.com/cilium/tetragon/pkg/sensors/base"
 	"github.com/cilium/tetragon/pkg/sensors/program"
 	"github.com/sirupsen/logrus"
 )
@@ -44,12 +47,18 @@ const (
 	aliveCnt            = 5
 	cleanupInterval     = time.Minute
 	cleanupInactiveTime = time.Minute
+	cgRateMaxEntries    = 32768 // this value could be fine tuned
 )
 
-var (
-	handle     *CgroupRate
-	handleLock sync.RWMutex
-)
+type globalState struct {
+	cgRateMap        *program.Map
+	cgRateOptionsMap *program.Map
+
+	handle *CgroupRate
+	mu     sync.RWMutex
+}
+
+var glSt globalState
 
 type cgroupQueue struct {
 	id    uint64
@@ -84,26 +93,30 @@ func newCgroupRate(
 
 func NewCgroupRate(ctx context.Context,
 	listener observer.Listener,
-	opts *option.CgroupRate) {
+	opts *option.CgroupRate) error {
 
 	if opts.Events == 0 || opts.Interval == 0 {
 		logger.GetLogger().Infof("Cgroup rate disabled (%d/%s)",
 			opts.Events, time.Duration(opts.Interval).String())
-		return
+		return nil
 	}
 
-	handleLock.Lock()
-	defer handleLock.Unlock()
+	glSt.mu.Lock()
+	defer glSt.mu.Unlock()
+	if glSt.cgRateMap == nil {
+		return fmt.Errorf("cgrouprate has not been registered to base sensor")
+	}
 
-	handle = newCgroupRate(listener, cgroupRateMap, opts)
-	go handle.process(ctx)
+	glSt.handle = newCgroupRate(listener, glSt.cgRateMap, opts)
+	go glSt.handle.process(ctx)
+	return nil
 }
 
 func NewTestCgroupRate(listener observer.Listener,
 	hash *program.Map,
 	opts *option.CgroupRate) {
 
-	handle = newCgroupRate(listener, hash, opts)
+	glSt.handle = newCgroupRate(listener, hash, opts)
 }
 
 func (r *CgroupRate) notify(msg notify.Message) {
@@ -117,6 +130,11 @@ func (r *CgroupRate) process(ctx context.Context) {
 	r.log.Infof("Cgroup rate started (%d/%s)",
 		r.opts.Events, time.Duration(r.opts.Interval).String())
 
+	defer func() {
+		// cleanup
+		glSt.handle = nil
+	}()
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -263,14 +281,14 @@ func (r *CgroupRate) processCgroup(id uint64, cgroup string, last uint64) bool {
 // Called from event handlers to kick off the cgroup rate
 // periodical check for event's cgroup.
 func Check(kube *processapi.MsgK8s, ktime uint64) {
-	if handle == nil {
+	if glSt.handle == nil {
 		return
 	}
 
-	handleLock.RLock()
-	defer handleLock.RUnlock()
+	glSt.mu.RLock()
+	defer glSt.mu.RUnlock()
 
-	if handle == nil {
+	if glSt.handle == nil {
 		return
 	}
 
@@ -280,27 +298,76 @@ func Check(kube *processapi.MsgK8s, ktime uint64) {
 		name:  string(kube.Docker[:]),
 	}
 
-	handle.ch <- cq
+	glSt.handle.ch <- cq
 	cgroupratemetrics.CgroupRateTotalInc(cgroupratemetrics.Check)
 }
 
 func Config() {
-	if handle == nil {
+	if glSt.handle == nil {
 		return
 	}
 
-	if cgroupRateOptionsMap.MapHandle == nil {
-		handle.log.Warn("failed to update cgroup rate options map")
+	if glSt.cgRateOptionsMap.MapHandle == nil {
+		glSt.handle.log.Warn("failed to update cgroup rate options map")
 		return
 	}
 
 	key := uint32(0)
 	opts := processapi.CgroupRateOptions{
-		Events:   handle.opts.Events,
-		Interval: handle.opts.Interval,
+		Events:   glSt.handle.opts.Events,
+		Interval: glSt.handle.opts.Interval,
 	}
 
-	if err := cgroupRateOptionsMap.MapHandle.Put(key, opts); err != nil {
+	if err := glSt.cgRateOptionsMap.MapHandle.Put(key, opts); err != nil {
 		cgroupratemetrics.CgroupRateTotalInc(cgroupratemetrics.UpdateFail)
 	}
 }
+
+func RegisterCgroupRate(sensor *sensors.Sensor) (*sensors.Sensor, error) {
+	if !option.CgroupRateEnabled() {
+		return sensor, nil
+	}
+
+	glSt.mu.Lock()
+	defer glSt.mu.Unlock()
+
+	if glSt.handle != nil {
+		return nil, fmt.Errorf("cgrouprate: handle is already set, need to cleanup first")
+	}
+
+	var rateProgs []*program.Program
+	var optsProgs []*program.Program
+	for _, p := range sensor.Progs {
+		if base.IsExecve(p) || base.IsFork(p) || base.IsExit(p) {
+			rateProgs = append(rateProgs, p)
+		}
+		if base.IsExecve(p) {
+			optsProgs = append(optsProgs, p)
+		}
+	}
+
+	if len(optsProgs) == 0 || len(rateProgs) == 0 {
+		return nil, fmt.Errorf("failed to find base programs")
+	}
+
+	cgRmdirProg := program.Builder(
+		"bpf_cgroup.o",
+		"cgroup/cgroup_rmdir",
+		"raw_tracepoint/cgroup_rmdir",
+		"tg_cgroup_rmdir",
+		"raw_tracepoint",
+	).SetPolicy(optsProgs[0].Policy)
+	rateProgs = append(rateProgs, cgRmdirProg)
+
+	glSt.cgRateMap = program.MapBuilder("cgroup_rate_map", rateProgs...)
+	glSt.cgRateMap.SetMaxEntries(cgRateMaxEntries)
+	glSt.cgRateOptionsMap = program.MapBuilder("cgroup_rate_options_map", optsProgs...)
+
+	sensor.Progs = append(sensor.Progs, cgRmdirProg)
+	sensor.Maps = append(sensor.Maps, glSt.cgRateMap, glSt.cgRateOptionsMap)
+	return sensor, nil
+}
+
+func init() {
+	base.RegisterExtensionAtInit("cgroup_rate", RegisterCgroupRate)
+}

pkg/cgrouprate/cgrouprate_test.go

@@ -240,8 +240,8 @@ func TestProcessCgroup(t *testing.T) {
 		NewTestCgroupRate(l, hash, &d.opts)
 
 		// setup cgrouprate cgroup
-		handle.cgroups[key.Id] = cgroup
-		assert.NotEqual(t, nil, handle)
+		glSt.handle.cgroups[key.Id] = cgroup
+		assert.NotEqual(t, nil, glSt.handle)
 
 		// store hash values
 		values[0] = d.values[0]
@@ -252,7 +252,7 @@ func TestProcessCgroup(t *testing.T) {
 		}
 
 		t.Logf("Test %d", idx)
-		ret := handle.processCgroup(key.Id, cgroup, d.last)
+		ret := glSt.handle.processCgroup(key.Id, cgroup, d.last)
 
 		assert.Equal(t, d.ret, ret)
 		assert.Equal(t, d.throttle, l.throttle)

pkg/observer/observertesthelper/observer_test_helper.go

@@ -423,8 +423,7 @@ func loadExporter(tb testing.TB, ctx context.Context, obs *observer.Observer, op
 		obs.RemoveListener(processManager)
 	})
 
-	cgrouprate.NewCgroupRate(ctx, processManager, &option.Config.CgroupRate)
-	return nil
+	return cgrouprate.NewCgroupRate(ctx, processManager, &option.Config.CgroupRate)
 }
 
 func loadObserver(tb testing.TB, ctx context.Context, base *sensors.Sensor,

pkg/sensors/base/base.go

@@ -139,7 +139,7 @@ func initBaseSensor() *sensors.Sensor {
 	setupPrograms()
 	sensor.Progs = GetDefaultPrograms()
 	sensor.Maps = GetDefaultMaps()
-	return applyExtensions(&sensor)
+	return ApplyExtensions(&sensor)
 }
 
 func initBaseSensorFn() func(tb testing.TB) *sensors.Sensor {

pkg/sensors/base/extensions.go

@@ -27,7 +27,7 @@ func RegisterExtensionAtInit(name string, fn ExtensionFn) {
 	})
 }
 
-func applyExtensions(s *sensors.Sensor) *sensors.Sensor {
+func ApplyExtensions(s *sensors.Sensor) *sensors.Sensor {
 	for _, ext := range extensions {
 		newS, err := ext.fn(s)
 		if err != nil {

pkg/sensors/base/utils.go

@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Authors of Tetragon
+
+package base
+
+import "github.com/cilium/tetragon/pkg/sensors/program"
+
+// IsExecve returns true if this is a base execve program
+func IsExecve(p *program.Program) bool {
+	return p.PinName == "event_execve" && p.Policy == basePolicy
+}
+
+func IsFork(p *program.Program) bool {
+	return p.PinName == "kprobe_pid_clear" && p.Policy == basePolicy
+}
+
+func IsExit(p *program.Program) bool {
+	return p.PinName == "event_exit" && p.Policy == basePolicy
+}