From 2edd39e11a406b4884e0de6dc36bfc282b63e8eb Mon Sep 17 00:00:00 2001
From: Alon Zivony <alon.zivony@aquasec.com>
Date: Sun, 7 Apr 2024 22:37:22 +0300
Subject: [PATCH] feat(ebpf): handle probes with watchers

Change the handling of the probes attach and detach to use the dependencies watchers.
This allows to handle correctly failure of events and future events updates.
---
 pkg/ebpf/tracee.go | 108 ++++++++++++++++++++++++++-------------------
 1 file changed, 63 insertions(+), 45 deletions(-)

diff --git a/pkg/ebpf/tracee.go b/pkg/ebpf/tracee.go
index 15d6f0cb597a..7338d4469a54 100644
--- a/pkg/ebpf/tracee.go
+++ b/pkg/ebpf/tracee.go
@@ -207,7 +207,7 @@ func (t *Tracee) addDependencyEventToState(evtID events.ID, dependantEvts []even
 }
 
 func (t *Tracee) removeEventFromState(evtID events.ID) {
-	logger.Debugw("Cancel event", "event", events.Core.GetDefinitionByID(evtID).GetName())
+	logger.Debugw("Remove event from state", "event", events.Core.GetDefinitionByID(evtID).GetName())
 	delete(t.eventsState, evtID)
 	delete(t.eventSignatures, evtID)
 }
@@ -1056,51 +1056,69 @@ func (t *Tracee) populateFilterMaps(newPolicies *policy.Policies, updateProcTree
 	return nil
 }
 
-// attachProbes attaches selected events probes to their respective eBPF progs
-func (t *Tracee) attachProbes() error {
-	var err error
-
-	// Get probe dependencies for a given event ID
-	getProbeDeps := func(id events.ID) []events.Probe {
-		depsNode, _ := t.eventsDependencies.GetEvent(id)
-		deps := depsNode.GetDependencies()
-		return deps.GetProbes()
+// attachEvent attaches the probes of the given event to their respective eBPF programs.
+// Calling attachment of a probes is a supported behavior, and will do nothing
+// if it was already attached.
+// In the case of failure in attaching, cancel the event using the dependencies API.
+func (t *Tracee) attachEvent(id events.ID) {
+	evtName := events.Core.GetDefinitionByID(id).GetName()
+	depsNode, ok := t.eventsDependencies.GetEvent(id)
+	if !ok {
+		logger.Errorw("Missing event in dependencies", "event", evtName)
+		return
 	}
-
-	// Get the list of probes to attach for each event being traced.
-	probesToEvents := make(map[events.Probe][]events.ID)
-	for id := range t.eventsState {
-		if !events.Core.IsDefined(id) {
-			continue
-		}
-		for _, probeDep := range getProbeDeps(id) {
-			probesToEvents[probeDep] = append(probesToEvents[probeDep], id)
+	for _, probe := range depsNode.GetDependencies().GetProbes() {
+		err := t.probes.Attach(probe.GetHandle(), t.cgroups)
+		if err != nil {
+			if probe.IsRequired() {
+				logger.Warnw(
+					"Cancelling event and its dependencies because of a missing probe",
+					"missing probe", probe.GetHandle(), "event", evtName,
+					"error", err,
+				)
+				t.eventsDependencies.RemoveEvent(id)
+			} else {
+				logger.Debugw(
+					"Failed to attach non-required probe for event",
+					"event", evtName,
+					"probe", probe.GetHandle(), "error", err,
+				)
+			}
 		}
 	}
+}
 
-	// Attach probes to their respective eBPF programs or cancel events if a required probe is missing.
-	for probe, evtID := range probesToEvents {
-		err = t.probes.Attach(probe.GetHandle(), t.cgroups) // attach bpf program to probe
-		if err != nil {
-			for _, evtID := range evtID {
-				evtName := events.Core.GetDefinitionByID(evtID).GetName()
-				if probe.IsRequired() {
-					t.eventsDependencies.RemoveEvent(evtID)
-					logger.Warnw(
-						"Cancelling event and its dependencies because of missing probe",
-						"missing probe", probe.GetHandle(), "event", evtName,
-						"error", err,
-					)
-				} else {
-					logger.Debugw(
-						"Failed to attach non-required probe for event",
-						"event", evtName,
-						"probe", probe.GetHandle(), "error", err,
-					)
+// attachProbes attaches selected events probes to their respective eBPF programs.
+func (t *Tracee) attachProbes() error {
+	// Subscribe to all watchers on the dependencies to attach and/or detach
+	// probes upon changes
+	t.eventsDependencies.SubscribeAdd(
+		func(node *dependencies.EventNode) {
+			t.attachEvent(node.GetID())
+		})
+
+	detachUnusedProbes := func(probesToCheck []events.Probe) {
+		for _, probe := range probesToCheck {
+			_, ok := t.eventsDependencies.GetHandle(probe.GetHandle())
+			if !ok {
+				err := t.probes.Detach(probe.GetHandle())
+				if err != nil {
+					logger.Debugw("Failed to detach probe",
+						"probe", probe.GetHandle(),
+						"error", err)
 				}
 			}
 		}
 	}
+	t.eventsDependencies.SubscribeRemove(func(node *dependencies.EventNode) {
+		detachUnusedProbes(node.GetDependencies().GetProbes())
+	})
+
+	// Attach probes to their respective eBPF programs or cancel events if a required probe is missing.
+	for eventID := range t.eventsState {
+		t.attachEvent(eventID)
+	}
+	// TODO: Update the eBPF maps according to events state after failure of attachments
 	return nil
 }
 
@@ -1155,13 +1173,6 @@ func (t *Tracee) initBPF() error {
 		return errfmt.WrapError(err)
 	}
 
-	// Attach eBPF programs to selected event's probes
-
-	err = t.attachProbes()
-	if err != nil {
-		return errfmt.WrapError(err)
-	}
-
 	// Update all ProcessTreeFilters after probes are attached: reduce the
 	// possible race window between the bpf programs updating the maps and
 	// userland reading procfs and also dealing with same maps.
@@ -1229,6 +1240,13 @@ func (t *Tracee) initBPF() error {
 		return errfmt.Errorf("error initializing logs perf map: %v", err)
 	}
 
+	// Attach eBPF programs to selected event's probes
+
+	err = t.attachProbes()
+	if err != nil {
+		return errfmt.WrapError(err)
+	}
+
 	return errfmt.WrapError(err)
 }