[WP] Refactor custom events to properly send them (#47586)

### What does this PR do?
- refactor a bit how custom events are handled so we always send them each time an action report is sent.
- add a new status to the kill action to handle early exit processes. When we try to kill one process (scope process), if the process exits before we kill it, we abort the kill instead of sending `kill_queued` status.

Note: This PR only fixes issues for scope process. Other changes will be needed in the future to handle other hedge cases.

### Motivation
Now every time an action report is sent, we have an associated remediation status event sent. Previously, we had situation where one was emitted but not the second one.

Co-authored-by: theo.putegnat <theo.putegnat@datadoghq.com>

Author: theop-dd

pkg/security/module/server.go

@@ -328,6 +328,15 @@ func (a *APIServer) updateCustomEventTags(msg *api.SecurityEventMessage) {
 	}
 }
 
+// SendCustomEventKillAction sends a custom remediation event for each resolved kill action report
+func SendCustomEventKillAction(probe *sprobe.Probe, tags []string, actionReports []model.ActionReport) {
+	for _, report := range actionReports {
+		if _, ok := report.(*sprobe.KillActionReport); ok {
+			probe.SendCustomEventKillAction(report, tags)
+		}
+	}
+}
+
 func (a *APIServer) start(ctx context.Context) {
 	ticker := time.NewTicker(200 * time.Millisecond)
 	defer ticker.Stop()
@@ -359,6 +368,9 @@ func (a *APIServer) start(ctx context.Context) {
 					return false
 				}
 
+				// For kill actions, send a custom remediation event per resolved kill report
+				// If a rule contains multiple kill, a custom event will be sent for each action
+				SendCustomEventKillAction(a.probe, msg.tags, msg.actionReports)
 				if a.containerFilter != nil {
 					containerName, imageName, podNamespace := utils.GetContainerFilterTags(msg.tags)
 					if a.containerFilter.IsExcluded(nil, containerName, imageName, podNamespace) {

pkg/security/probe/actions.go

@@ -33,14 +33,23 @@ const (
 	// KillActionStatusQueued indicates the kill action was queued until the end of the first rule period
 	KillActionStatusQueued KillActionStatus = "kill_queued"
 	// KillActionStatusPartiallyPerformed indicates the kill action was performed on some processes but not all
-	KillActionStatusPartiallyPerformed = "partially_performed"
+	KillActionStatusPartiallyPerformed KillActionStatus = "partially_performed"
+	// KillActionStatusKillAborted indicates the kill action was aborted because the process exited before the kill was performed
+	KillActionStatusKillAborted KillActionStatus = "kill_aborted"
 
 	// maxRetryForMsgWithKillAction is the maximum number of retries for a kill action
 	// - a kill can be queued up to the end of the first disarmer period (1min by default)
 	// - so, we set the server retry period to 1min and 2sec (+2sec to have the time to trigger the kill and wait to catch the process exit)
 	maxRetryForMsgWithKillAction = 62
 )
 
+// RemediationContainerContext represents the container context for remediation events (e.g. container ID and created_at).
+// Defined here so KillActionReport can use it on all platforms; Linux-specific remediation logic lives in remediations_linux.go.
+type RemediationContainerContext struct {
+	CreatedAt uint64 `json:"created_at,omitempty"`
+	ID        string `json:"id,omitempty"`
+}
+
 // KillActionReport defines a kill action reports
 type KillActionReport struct {
 	sync.RWMutex
@@ -55,9 +64,10 @@ type KillActionReport struct {
 	DisarmerType string
 
 	// internal
-	Pid      uint32
-	resolved bool
-	rule     *rules.Rule
+	Pid              uint32
+	resolved         bool
+	rule             *rules.Rule
+	containerContext RemediationContainerContext // This is an internal field needed for remediation status events
 }
 
 // JKillActionReport used to serialize date
@@ -128,3 +138,10 @@ func (k *KillActionReport) IsMatchingRule(ruleID eval.RuleID) bool {
 
 	return k.rule.ID == ruleID
 }
+
+// GetRemediationContainerContext returns the container ID and created_at when the process was in a container, otherwise empty/zero.
+func (k *KillActionReport) GetRemediationContainerContext() RemediationContainerContext {
+	k.RLock()
+	defer k.RUnlock()
+	return k.containerContext
+}

pkg/security/probe/probe.go

@@ -59,6 +59,7 @@ type PlatformProbe interface {
 	GetEventTags(_ containerutils.ContainerID) []string
 	EnableEnforcement(bool)
 	ReplayEvents()
+	SendCustomEventKillAction(_ model.ActionReport, _ []string)
 }
 
 var probeTelemetry = struct {
@@ -357,6 +358,11 @@ func (p *Probe) AddDiscarderPushedCallback(cb DiscarderPushedCallback) {
 	p.PlatformProbe.AddDiscarderPushedCallback(cb)
 }
 
+// SendCustomEventKillAction sends a custom remediation-style event for a resolved kill action report.
+func (p *Probe) SendCustomEventKillAction(report model.ActionReport, tags []string) {
+	p.PlatformProbe.SendCustomEventKillAction(report, tags)
+}
+
 // DispatchCustomEvent sends a custom event to the probe event handler
 func (p *Probe) DispatchCustomEvent(rule *rules.Rule, event *events.CustomEvent) {
 	p.logTraceEvent(event.GetEventType(), event)

pkg/security/probe/probe_ebpf.go

@@ -3366,7 +3366,7 @@ func (p *EBPFProbe) HandleActions(ctx *eval.Context, rule *rules.Rule) {
 				p.probe.onRuleActionPerformed(rule, action.Def)
 			}
 			if report != nil {
-				p.HandleKillRemediation(rule, ev, report, action)
+				p.HandleKillRemediation(rule, ev, action)
 			}
 
 		case action.Def.CoreDump != nil:

pkg/security/probe/probe_ebpfless.go

@@ -720,6 +720,9 @@ func (p *EBPFLessProbe) EnableEnforcement(state bool) {
 	p.processKiller.SetState(state)
 }
 
+// SendCustomEventKillAction is a no-op for EBPFLess (remediation custom events use EBPF probe).
+func (p *EBPFLessProbe) SendCustomEventKillAction(_ model.ActionReport, _ []string) {}
+
 // GetAgentContainerContext returns the agent container context
 func (p *EBPFLessProbe) GetAgentContainerContext() *events.AgentContainerContext {
 	return p.probe.GetAgentContainerContext()

pkg/security/probe/probe_others.go

@@ -133,3 +133,6 @@ func (p *Probe) GetAgentContainerContext() *events.AgentContainerContext {
 // Walk iterates through the entire tree and call the provided callback on each entry
 func (p *Probe) Walk(_ func(*model.ProcessCacheEntry)) {
 }
+
+// SendCustomEventKillAction is a no-op on unsupported platforms (remediation custom events are Linux-only).
+func (p *Probe) SendCustomEventKillAction(_ model.ActionReport, _ []string) {}

pkg/security/probe/probe_windows.go

@@ -1593,6 +1593,9 @@ func (p *WindowsProbe) EnableEnforcement(state bool) {
 	p.processKiller.SetState(state)
 }
 
+// SendCustomEventKillAction is a no-op on Windows (remediation custom events are Linux-only).
+func (p *WindowsProbe) SendCustomEventKillAction(_ model.ActionReport, _ []string) {}
+
 // NewProbe instantiates a new runtime security agent probe
 func NewProbe(config *config.Config, hostname string, opts Opts) (*Probe, error) {
 	opts.normalize()

pkg/security/probe/process_killer.go

@@ -241,6 +241,12 @@ func (p *ProcessKiller) HandleProcessExited(event *model.Event) {
 		defer report.Unlock()
 
 		if report.Pid == event.ProcessContext.Pid {
+			if report.Scope == "process" {
+				if report.Status == KillActionStatusQueued {
+					// The process exited before the kill was performed
+					report.Status = KillActionStatusKillAborted
+				}
+			}
 			report.ExitedAt = event.ProcessContext.ExitTime
 			report.resolved = true
 			return true
@@ -324,6 +330,10 @@ func (p *ProcessKiller) KillAndReport(kill *rules.KillDefinition, rule *rules.Ru
 				Pid:          ev.ProcessContext.Pid,
 				rule:         rule,
 			}
+			if !ev.ProcessContext.Process.ContainerContext.IsNull() {
+				report.containerContext.ID = string(ev.ProcessContext.Process.ContainerContext.ContainerID)
+				report.containerContext.CreatedAt = ev.ProcessContext.Process.ContainerContext.CreatedAt
+			}
 			if dismantled {
 				report.Status = KillActionStatusRuleDismantled
 				seclog.Warnf("skipping kill action of rule `%s` because it has been dismantled", rule.ID)
@@ -423,6 +433,10 @@ func (p *ProcessKiller) KillAndReport(kill *rules.KillDefinition, rule *rules.Ru
 		Pid:        ev.ProcessContext.Pid,
 		rule:       rule,
 	}
+	if !ev.ProcessContext.Process.ContainerContext.IsNull() {
+		report.containerContext.ID = string(ev.ProcessContext.Process.ContainerContext.ContainerID)
+		report.containerContext.CreatedAt = ev.ProcessContext.Process.ContainerContext.CreatedAt
+	}
 
 	if disarmer != nil && p.warmupEnqueued(disarmer, sig, pcs) {
 		log.Warnf("rule %s triggered on first period, putting pids to kill on wait list", rule.ID)

pkg/security/probe/remediations_linux.go

@@ -69,13 +69,6 @@ type RemediationProcessContext struct {
 	PID uint32 `json:"pid,omitempty"`
 }
 
-// RemediationContainerContext represents the container context for remediation events
-// easyjson:json
-type RemediationContainerContext struct {
-	CreatedAt uint64 `json:"created_at,omitempty"`
-	ID        string `json:"id,omitempty"`
-}
-
 // RemediationAgentContext represents the agent context for remediation events
 // easyjson:json
 type RemediationAgentContext struct {
@@ -109,15 +102,6 @@ func (k RemediationEvent) ToJSON() ([]byte, error) {
 	return utils.MarshalEasyJSON(k)
 }
 
-func getAgentEventID(rule *rules.Rule) string {
-	for _, tag := range rule.Tags {
-		if strings.HasPrefix(tag, "agent_event_id:") {
-			return tag[len("agent_event_id:"):]
-		}
-	}
-	return ""
-}
-
 func getRemediationTagBool(rule *rules.Rule) bool {
 	for _, tag := range rule.Tags {
 		if strings.HasPrefix(tag, "remediation_rule:") {
@@ -140,27 +124,30 @@ func generateNetworkIsolationActionKey(ruleID string, filter string) string {
 	return NetworkIsolationKeyPrefix + ruleID + hex.EncodeToString(hash[:])
 
 }
-func generateRemediationActionKey(rule *rules.Rule) string {
+func generateRemediationActionKey(key string) string {
 	// prefix + agent_event_id
-	return RemediationKeyPrefix + getAgentEventID(rule)
+	return RemediationKeyPrefix + key
 }
 
 func getRemediationKeyFromAction(rule *rules.Rule, action *rules.Action) string {
-	if getRemediationTagBool(rule) {
-		// Should not have multiple remediation actions for the same rule that were generated by customers (remediation)
-		return generateRemediationActionKey(rule)
-	}
+	key := ""
+
 	// Having multiple actions in the same rulemeans that they are from a rule that was not dynamically generated for the remediation feature
 	// We assume this combination unique
 	if action.Def.Kill != nil {
 		// ruleID + scope + signal
-		return generateKillActionKey(rule.ID, action.Def.Kill.Scope, action.Def.Kill.Signal)
-	}
-	if action.Def.NetworkFilter != nil {
+		key = generateKillActionKey(rule.ID, action.Def.Kill.Scope, action.Def.Kill.Signal)
+	} else if action.Def.NetworkFilter != nil {
 		// ruleID + bpffilter
-		return generateNetworkIsolationActionKey(rule.ID, action.Def.NetworkFilter.BPFFilter)
+		key = generateNetworkIsolationActionKey(rule.ID, action.Def.NetworkFilter.BPFFilter)
+	}
+
+	if getRemediationTagBool(rule) {
+		// Should not have multiple remediation actions for the same rule that were generated by customers (remediation)
+		return generateRemediationActionKey(key)
 	}
-	return ""
+
+	return key
 }
 
 // NewRemediationEvent creates a new Remediation event from the latest action report
@@ -211,6 +198,17 @@ func getTagsFromRule(rule *rules.Rule) RuleTags {
 	return ruleTags
 }
 
+// tagsToRuleTags converts a slice of "key:value" tags into RuleTags map.
+func tagsToRuleTags(tags []string) RuleTags {
+	ruleTags := make(RuleTags)
+	for _, tag := range tags {
+		if before, after, ok := strings.Cut(tag, ":"); ok {
+			ruleTags[before] = after
+		}
+	}
+	return ruleTags
+}
+
 // HandleRemediationStatus is called when a new ruleset is loaded
 // It cleans up the activeRemediations map from the kill actions and network isolation actions that are not persistent
 func (p *EBPFProbe) HandleRemediationStatus(rs *rules.RuleSet) {
@@ -278,7 +276,33 @@ func (p *EBPFProbe) HandleRemediationStatus(rs *rules.RuleSet) {
 	}
 }
 
-func (p *EBPFProbe) HandleKillRemediation(rule *rules.Rule, ev *model.Event, report *KillActionReport, action *rules.Action) {
+// SendCustomEventKillAction sends a custom remediation event for a resolved kill action report
+func (p *EBPFProbe) SendCustomEventKillAction(report model.ActionReport, tags []string) {
+	killReport, ok := report.(*KillActionReport)
+	if !ok {
+		return
+	}
+	killReport.RLock()
+	status := string(killReport.Status)
+	scope := killReport.Scope
+	pid := killReport.Pid
+	killReport.RUnlock()
+
+	containerContext := killReport.GetRemediationContainerContext()
+
+	remediation := &Remediation{
+		actionType:       RemediationTypeKill,
+		triggered:        true,
+		scope:            scope,
+		containerContext: containerContext,
+		processContext:   RemediationProcessContext{PID: pid},
+		ruleTags:         tagsToRuleTags(tags),
+	}
+	re := NewRemediationEvent(p, remediation, status, RemediationTypeKillStr)
+	p.SendRemediationEvent(re)
+}
+
+func (p *EBPFProbe) HandleKillRemediation(rule *rules.Rule, ev *model.Event, action *rules.Action) {
 	remediationKey := getRemediationKeyFromAction(rule, action)
 	p.activeRemediationsLock.Lock()
 	defer p.activeRemediationsLock.Unlock()
@@ -292,30 +316,9 @@ func (p *EBPFProbe) HandleKillRemediation(rule *rules.Rule, ev *model.Event, rep
 		remediation.policy = ""
 		remediation.ruleTags = getTagsFromRule(rule)
 
-	} else {
-		// Don't create a new entry for kill actions that are not from the remediation feature
-		// It will only be used to send an event
-		remediation = &Remediation{
-			actionType: RemediationTypeKill,
-			triggered:  true,
-			processContext: RemediationProcessContext{
-				PID: ev.ProcessContext.Process.Pid,
-			},
-			containerContext: RemediationContainerContext{
-				ID:        string(ev.ProcessContext.Process.ContainerContext.ContainerID),
-				CreatedAt: ev.ProcessContext.Process.ContainerContext.CreatedAt,
-			},
-			scope: action.Def.Kill.Scope,
-		}
 	}
-
-	// Get kill status
-	report.RLock()
-	status := string(report.Status)
-	report.RUnlock()
-	// Send custom event
-	killActionEvent := NewRemediationEvent(p, remediation, status, RemediationTypeKillStr)
-	p.SendRemediationEvent(killActionEvent)
+	// Don't send an event for kill action here
+	// The event will be sent when the report is resolved to handle the cases where disarmers are used
 }
 
 func (p *EBPFProbe) HandleNetworkRemediation(rule *rules.Rule, ev *model.Event, report *RawPacketActionReport, action *rules.Action) {