Only initialize process tree related events if enabled

castai · Jul 19, 2024 · 9eabd05 · 9eabd05
1 parent 5c56bb3
commit 9eabd05
Show file tree

Hide file tree

Showing 6 changed files with 81 additions and 52 deletions.
diff --git a/cmd/agent/daemon/app/app.go b/cmd/agent/daemon/app/app.go
@@ -265,14 +265,18 @@ func (a *App) Run(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
-	processTreeCollector, err := processtree.New(log, procHandler, containersClient)
-	if err != nil {
-		return fmt.Errorf("process tree: %w", err)
-	}
-	err = processTreeCollector.Init(ctx)
-	if err != nil {
-		return fmt.Errorf("process tree: %w", err)
+
+	var processTreeCollector processtree.ProcessTreeCollector
+
+	if cfg.ProcessTree.Enabled {
+		processTreeCollector, err = initializeProcessTree(ctx, log, procHandler, containersClient)
+		if err != nil {
+			return fmt.Errorf("initialize process tree: %w", err)
+		}
+	} else {
+		processTreeCollector = processtree.NewNoop()
 	}
+
 	ct, err := conntrack.NewClient(log)
 	if err != nil {
 		return fmt.Errorf("conntrack: %w", err)
@@ -326,9 +330,6 @@ func (a *App) Run(ctx context.Context) error {
 		SystemEvents: []events.ID{
 			events.CgroupMkdir,
 			events.CgroupRmdir,
-			events.SchedProcessExec,
-			events.SchedProcessExit,
-			events.SchedProcessFork,
 		},
 		Events: []*ebpftracer.EventPolicy{},
 	}
@@ -341,6 +342,14 @@ func (a *App) Run(ctx context.Context) error {
 		),
 	}
 
+	if cfg.ProcessTree.Enabled {
+		policy.SystemEvents = append(policy.SystemEvents, []events.ID{
+			events.SchedProcessExec,
+			events.SchedProcessExit,
+			events.SchedProcessFork,
+		}...)
+	}
+
 	if len(exporters.Events) > 0 {
 		policy.SignatureEvents = signatureEngine.TargetEvents()
 		policy.Events = append(policy.Events, []*ebpftracer.EventPolicy{
@@ -439,6 +448,18 @@ func (a *App) Run(ctx context.Context) error {
 	}
 }
 
+func initializeProcessTree(ctx context.Context, log *logging.Logger, procHandler *proc.Proc, containersClient *containers.Client) (*processtree.ProcessTreeCollectorImpl, error) {
+	processTreeCollector, err := processtree.New(log, procHandler, containersClient)
+	if err != nil {
+		return nil, err
+	}
+	err = processTreeCollector.Init(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return processTreeCollector, nil
+}
+
 func (a *App) syncRemoteConfig(ctx context.Context, client *castai.Client) error {
 	for {
 		select {

diff --git a/pkg/ebpftracer/tracer.go b/pkg/ebpftracer/tracer.go
@@ -47,7 +47,7 @@ type CgroupClient interface {
 	IsDefaultHierarchy(uint32) bool
 }
 
-type ProcessTreeCollector interface {
+type processTreeCollector interface {
 	ProcessStarted(eventTime time.Time, containerID string, p processtree.Process)
 	ProcessForked(eventTime time.Time, containerID string, parent processtree.ProcessKey, processKey processtree.ProcessKey)
 	ProcessExited(eventTime time.Time, containerID string, processKey processtree.ProcessKey, exitTime uint64)
@@ -70,7 +70,7 @@ type Config struct {
 	NetflowSampleSubmitIntervalSeconds uint64
 	NetflowGrouping                    NetflowGrouping
 	TrackSyscallStats                  bool
-	ProcessTreeCollector               ProcessTreeCollector
+	ProcessTreeCollector               processTreeCollector
 }
 
 type cgroupCleanupRequest struct {

diff --git a/pkg/ebpftracer/tracer_filter_test.go b/pkg/ebpftracer/tracer_filter_test.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"log/slog"
 	"testing"
-	"time"
 
 	"github.com/castai/kvisor/pkg/cgroup"
 	"github.com/castai/kvisor/pkg/containers"
@@ -144,19 +143,6 @@ func (c *MockContainerClient) CleanupCgroup(cgroup uint64) {
 	c.CgroupCleaner(cgroup)
 }
 
-type mockProcessTreeCollector struct{}
-
-func (d *mockProcessTreeCollector) ProcessExited(eventTime time.Time, containerID string, processKey processtree.ProcessKey, exitTime uint64) {
-}
-
-func (d *mockProcessTreeCollector) ProcessForked(eventTime time.Time, containerID string, parent processtree.ProcessKey, processKey processtree.ProcessKey) {
-}
-
-func (d *mockProcessTreeCollector) ProcessStarted(eventTime time.Time, containerID string, p processtree.Process) {
-}
-
-var _ ProcessTreeCollector = (*mockProcessTreeCollector)(nil)
-
 type tracerOption func(*Tracer)
 
 func buildTestTracer(options ...tracerOption) *Tracer {
@@ -179,7 +165,7 @@ func buildTestTracer(options ...tracerOption) *Tracer {
 				},
 			},
 			CgroupClient:         &MockCgroupClient{},
-			ProcessTreeCollector: &mockProcessTreeCollector{},
+			ProcessTreeCollector: processtree.NewNoop(),
 		},
 		eventsChan:        make(chan *types.Event, 10),
 		eventPoliciesMap:  map[events.ID]*EventPolicy{},

diff --git a/pkg/ebpftracer/tracer_test.go b/pkg/ebpftracer/tracer_test.go
@@ -78,7 +78,7 @@ func TestTracer(t *testing.T) {
 		HomePIDNS:                          pidNS,
 		NetflowSampleSubmitIntervalSeconds: 0,
 		NetflowGrouping:                    ebpftracer.NetflowGroupingDropSrcPort,
-		ProcessTreeCollector:               &dummyProcessTreeCollector{},
+		ProcessTreeCollector:               processtree.NewNoop(),
 	})
 	defer tr.Close()
 
@@ -326,16 +326,3 @@ func getLayer4TCPFromPacket(packet gopacket.Packet) (*layers.TCP, error) {
 	}
 	return tcp, nil
 }
-
-type dummyProcessTreeCollector struct{}
-
-func (d *dummyProcessTreeCollector) ProcessExited(eventTime time.Time, containerID string, processKey processtree.ProcessKey, exitTime uint64) {
-}
-
-func (d *dummyProcessTreeCollector) ProcessForked(eventTime time.Time, containerID string, parent processtree.ProcessKey, processKey processtree.ProcessKey) {
-}
-
-func (d *dummyProcessTreeCollector) ProcessStarted(eventTime time.Time, containerID string, p processtree.Process) {
-}
-
-var _ ebpftracer.ProcessTreeCollector = (*dummyProcessTreeCollector)(nil)
diff --git a/pkg/processtree/noop.go b/pkg/processtree/noop.go
@@ -0,0 +1,28 @@
+package processtree
+
+import "time"
+
+type NoopProcessTreeCollector struct {
+	dummyChan chan ProcessTreeEvent
+}
+
+func NewNoop() *NoopProcessTreeCollector {
+	return &NoopProcessTreeCollector{
+		dummyChan: make(chan ProcessTreeEvent),
+	}
+}
+
+func (d *NoopProcessTreeCollector) Events() <-chan ProcessTreeEvent {
+	return d.dummyChan
+}
+
+func (d *NoopProcessTreeCollector) ProcessExited(eventTime time.Time, containerID string, processKey ProcessKey, exitTime uint64) {
+}
+
+func (d *NoopProcessTreeCollector) ProcessForked(eventTime time.Time, containerID string, parent ProcessKey, processKey ProcessKey) {
+}
+
+func (d *NoopProcessTreeCollector) ProcessStarted(eventTime time.Time, containerID string, p Process) {
+}
+
+var _ ProcessTreeCollector = (*NoopProcessTreeCollector)(nil)
diff --git a/pkg/processtree/processtree.go b/pkg/processtree/processtree.go
@@ -65,7 +65,14 @@ func (p Process) Exited() bool {
 	return p.ExitTime > 0
 }
 
-type ProcessTreeCollector struct {
+type ProcessTreeCollector interface {
+	ProcessStarted(eventTime time.Time, containerID string, p Process)
+	ProcessForked(eventTime time.Time, containerID string, parent ProcessKey, processKey ProcessKey)
+	ProcessExited(eventTime time.Time, containerID string, processKey ProcessKey, exitTime uint64)
+	Events() <-chan ProcessTreeEvent
+}
+
+type ProcessTreeCollectorImpl struct {
 	log              *logging.Logger
 	proc             *proc.Proc
 	containersClient *containers.Client
@@ -74,8 +81,8 @@ type ProcessTreeCollector struct {
 	eventSink        chan ProcessTreeEvent
 }
 
-func New(log *logging.Logger, p *proc.Proc, containersClient *containers.Client) (*ProcessTreeCollector, error) {
-	return &ProcessTreeCollector{
+func New(log *logging.Logger, p *proc.Proc, containersClient *containers.Client) (*ProcessTreeCollectorImpl, error) {
+	return &ProcessTreeCollectorImpl{
 		log:              log,
 		proc:             p,
 		containersClient: containersClient,
@@ -114,14 +121,14 @@ func ToProcessKeyNs(pid proc.PID, startTimeNs uint64) ProcessKey {
 }
 
 // NOTE: We do not defer cleanup of process trees here, since the container cleanup is already defered.
-func (c *ProcessTreeCollector) onContainerDelete(container *containers.Container) {
+func (c *ProcessTreeCollectorImpl) onContainerDelete(container *containers.Container) {
 	c.processTreesMu.Lock()
 	defer c.processTreesMu.Unlock()
 
 	delete(c.processTrees, container.ID)
 }
 
-func (c *ProcessTreeCollector) Init(ctx context.Context) error {
+func (c *ProcessTreeCollectorImpl) Init(ctx context.Context) error {
 	c.containersClient.RegisterContainerDeletedListener(c.onContainerDelete)
 
 	processes, err := c.containersClient.LoadContainerTasks(ctx)
@@ -194,7 +201,7 @@ func (c *ProcessTreeCollector) Init(ctx context.Context) error {
 	return nil
 }
 
-func (c *ProcessTreeCollector) ProcessStarted(eventTime time.Time, containerID string, p Process) {
+func (c *ProcessTreeCollectorImpl) ProcessStarted(eventTime time.Time, containerID string, p Process) {
 	c.processTreesMu.Lock()
 	defer c.processTreesMu.Unlock()
 
@@ -215,7 +222,7 @@ func (c *ProcessTreeCollector) ProcessStarted(eventTime time.Time, containerID s
 	})
 }
 
-func (c *ProcessTreeCollector) ProcessForked(eventTime time.Time, containerID string, parent ProcessKey, processKey ProcessKey) {
+func (c *ProcessTreeCollectorImpl) ProcessForked(eventTime time.Time, containerID string, parent ProcessKey, processKey ProcessKey) {
 	c.processTreesMu.Lock()
 	defer c.processTreesMu.Unlock()
 
@@ -254,7 +261,7 @@ func (c *ProcessTreeCollector) ProcessForked(eventTime time.Time, containerID st
 	})
 }
 
-func (c *ProcessTreeCollector) ProcessExited(eventTime time.Time, containerID string, processKey ProcessKey, exitTime uint64) {
+func (c *ProcessTreeCollectorImpl) ProcessExited(eventTime time.Time, containerID string, processKey ProcessKey, exitTime uint64) {
 	c.processTreesMu.Lock()
 	defer c.processTreesMu.Unlock()
 
@@ -282,11 +289,11 @@ func (c *ProcessTreeCollector) ProcessExited(eventTime time.Time, containerID st
 	}
 }
 
-func (c *ProcessTreeCollector) Events() <-chan ProcessTreeEvent {
+func (c *ProcessTreeCollectorImpl) Events() <-chan ProcessTreeEvent {
 	return c.eventSink
 }
 
-func (c *ProcessTreeCollector) fireEvent(e ProcessEvent) {
+func (c *ProcessTreeCollectorImpl) fireEvent(e ProcessEvent) {
 	c.log.Debugf("fire process tree event %s (%s): pid: %d, startTime: %s, ppid: %d, parentStartTime: %s, containerID: %s", e.Action.String(),
 		e.Timestamp.Format(time.RFC3339), e.Process.PID, e.Process.StartTime, e.Process.PPID, e.Process.ParentStartTime, e.ContainerID)
 	select {
@@ -299,7 +306,7 @@ func (c *ProcessTreeCollector) fireEvent(e ProcessEvent) {
 	}
 }
 
-func (c *ProcessTreeCollector) fireEvents(event ProcessTreeEvent) {
+func (c *ProcessTreeCollectorImpl) fireEvents(event ProcessTreeEvent) {
 	if c.log.IsEnabled(slog.LevelDebug) {
 		c.log.Debugf("fire process tree event (initial %t) ---", event.Initial)
 		for _, e := range event.Events {