tetragon

1
// SPDX-License-Identifier: Apache-2.0
2
// Copyright Authors of Tetragon
3

4
package sensors
5

6
import (
7
	"context"
8
	"testing"
9

10
	"github.com/cilium/tetragon/pkg/bpf"
11
	"github.com/cilium/tetragon/pkg/observer"
12
	"github.com/cilium/tetragon/pkg/sensors"
13
)
14

15
// LoadSensor is a helper for loading a sensor in tests
16
func LoadSensor(t *testing.T, sensori sensors.SensorIface) {
17
	sensor, ok := sensori.(*sensors.Sensor)
18
	if !ok {
19
		t.Fatalf("Cannot call LoadSensor on type %T", sensori)
20
	}
21

22
	if err := sensor.FindPrograms(); err != nil {
23
		t.Fatalf("ObserverFindProgs error: %s", err)
24
	}
25
	bpfDir := bpf.MapPrefixPath()
26
	if err := sensor.Load(bpfDir); err != nil {
27
		t.Fatalf("observerLoadSensor error: %s", err)
28
	}
29

30
	t.Cleanup(func() {
31
		sensor.Unload()
32
	})
33
}
34

35
// TestSensorManager sensor manager used in tests
36
type TestSensorManager struct {
37
	Manager *sensors.Manager
38
}
39

40
// GetTestSensorManager returns a new test sensor manager.
41
// Some tests require an observer running, some do not. To support both, the function checks if a
42
// sensor manager has already been setup in the observer, and uses it if so.
43
// Otherwise, it creates a new one. If it creates a new one it will use the test name to create a
44
// unqique directory for maps/etc, and will also register the necessary cleanup functions using
45
// t.Cleanup()
46
func GetTestSensorManager(ctx context.Context, t *testing.T) *TestSensorManager {
47
	if mgr := observer.GetSensorManager(); mgr != nil {
48
		return &TestSensorManager{
49
			Manager: mgr,
50
		}
51
	}
52

53
	path := bpf.MapPrefixPath()
54
	mgr, err := sensors.StartSensorManager(path, nil)
55
	if err != nil {
56
		t.Fatalf("startSensorController failed: %s", err)
57
	}
58
	t.Cleanup(func() {
59
		err := mgr.StopSensorManager(ctx)
60
		if err != nil {
61
			t.Logf("stopSensorController failed: %s\n", err)
62
		}
63
	})
64

65
	return &TestSensorManager{
66
		Manager: mgr,
67
	}
68
}
69

70
// AddAndEnableSensor is a helper function that adds and enables a new sensor
71
func (tsm *TestSensorManager) AddAndEnableSensor(
72
	ctx context.Context,
73
	t *testing.T,
74
	sensor *sensors.Sensor,
75
	sensorName string,
76
) {
77
	if err := tsm.Manager.AddSensor(ctx, sensorName, sensor); err != nil {
78
		t.Fatalf("failed to add generic tracepoint sensor: %s", err)
79
	}
80
	t.Cleanup(func() {
81
		tsm.Manager.RemoveSensor(ctx, sensorName)
82
	})
83
	if err := tsm.Manager.EnableSensor(ctx, sensorName); err != nil {
84
		t.Fatalf("EnableSensor error: %s", err)
85
	}
86
	t.Cleanup(func() {
87
		tsm.Manager.DisableSensor(ctx, sensorName)
88
	})
89
}
90

91
// EnableSensors is a helper function that enables a list of sensors
92
func (tsm *TestSensorManager) EnableSensors(
93
	ctx context.Context,
94
	t *testing.T,
95
	targets []*sensors.Sensor,
96
) {
97
	for _, s := range targets {
98
		if err := tsm.Manager.EnableSensor(ctx, s.Name); err != nil {
99
			t.Fatalf("EnableSensor error: %s", err)
100
		}
101
	}
102
}
103

104
// AddAndEnableSensor is a helper function that adds and enables a new sensor
105
func (tsm *TestSensorManager) AddAndEnableSensors(
106
	ctx context.Context,
107
	t *testing.T,
108
	targets []*sensors.Sensor,
109
) {
110
	for i := range targets {
111
		sensor := targets[i]
112
		tsm.AddAndEnableSensor(ctx, t, sensor, sensor.Name)
113
	}
114
}
115

116
// EnableSensors is a helper function that enables a list of sensors
117
func (tsm *TestSensorManager) DisableSensors(
118
	ctx context.Context,
119
	t *testing.T,
120
	targets []*sensors.Sensor,
121
) {
122
	for _, s := range targets {
123
		err := tsm.Manager.DisableSensor(ctx, s.Name)
124
		if err != nil {
125
			t.Logf("DisableSensor failed: %s", err)
126
		}
127
	}
128
}
129