Optimise map and correct delta type

fix

Author: theop-dd

pkg/security/ebpf/c/include/maps.h

@@ -98,7 +98,6 @@ BPF_LRU_MAP(dns_responses_sent_to_userspace, u16, struct dns_responses_sent_to_u
 BPF_LRU_MAP(capabilities_usage, struct capabilities_usage_key_t, struct capabilities_usage_entry_t, 1) // max entries will be overridden at runtime
 BPF_LRU_MAP(sock_cookie_pid, u64, u32, 1); // max entries will be overridden at runtime
 BPF_LRU_MAP(memfd_tracking, struct memfd_key_t, u32, 1024)
-BPF_LRU_MAP(dropped_packets, u64, u64, 512)
 
 BPF_LRU_MAP_FLAGS(tasks_in_coredump, u64, u8, 64, BPF_F_NO_COMMON_LRU)
 BPF_LRU_MAP_FLAGS(syscalls, u64, struct syscall_cache_t, 1, BPF_F_NO_COMMON_LRU) // max entries will be overridden at runtime
@@ -140,6 +139,7 @@ BPF_PERCPU_ARRAY_MAP(raw_packet_enabled, u32, 1)
 BPF_PERCPU_ARRAY_MAP(sysctl_event_gen, struct sysctl_event_t, 1)
 BPF_PERCPU_ARRAY_MAP(on_demand_event_gen, struct on_demand_event_t, 1)
 BPF_PERCPU_ARRAY_MAP(setsockopt_event, struct setsockopt_event_t, 1)
+BPF_PERCPU_ARRAY_MAP(dropped_packets, u32, 256)
 
 BPF_PROG_ARRAY(args_envs_progs, 3)
 BPF_PROG_ARRAY(dentry_resolver_kprobe_or_fentry_callbacks, EVENT_MAX)

pkg/security/ebpf/probes/rawpacket/bpffilter.go

@@ -38,7 +38,7 @@ const (
 )
 
 // MaxDropActionFilters is the maximum number of network drop action filters tracked in kernel.
-const MaxDropActionFilters = 512
+const MaxDropActionFilters = 256
 
 // ToTCAct converts a policy to a TCAct
 func (p Policy) ToTCAct() TCAct {

pkg/security/ebpf/probes/rawpacket/pcap.go

@@ -41,7 +41,6 @@ const (
 	structRawPacketEventDataSize = 256
 
 	dropStatsKeyStackOffset = int16(-8)
-	dropStatsValStackOffset = int16(-16)
 )
 
 // ProgOpts defines options
@@ -109,38 +108,25 @@ func (opts *ProgOpts) WithDropStatsMapFd(fd int) *ProgOpts {
 
 func dropStatsIncrementInsts(filterIndex int, dropStatsMapFd int, nextLabel string) asm.Instructions {
 	incLabel := fmt.Sprintf("inc_drop_stat_%d", filterIndex)
-	initLabel := fmt.Sprintf("init_drop_stat_%d", filterIndex)
 
 	return asm.Instructions{
 		// Put the key on the stack
 		asm.Mov.Reg(asm.R1, asm.RFP).WithSymbol(incLabel),
 		asm.Add.Imm(asm.R1, int32(dropStatsKeyStackOffset)),
 		asm.Mov.Imm(asm.R2, int32(filterIndex)),
-		asm.StoreMem(asm.R1, 0, asm.R2, asm.DWord),
+		asm.StoreMem(asm.R1, 0, asm.R2, asm.Word),
 		// Lookup in the map
 		asm.LoadMapPtr(asm.R1, dropStatsMapFd),
 		asm.Mov.Reg(asm.R2, asm.RFP),
 		asm.Add.Imm(asm.R2, int32(dropStatsKeyStackOffset)),
 		asm.FnMapLookupElem.Call(),
-		asm.JEq.Imm(asm.R0, 0, initLabel),
-		// Increment if it exists
+		// should never happen
+		asm.JEq.Imm(asm.R0, 0, nextLabel),
+		// Increment
 		asm.Mov.Reg(asm.R5, asm.R0),
-		asm.LoadMem(asm.R6, asm.R5, 0, asm.DWord),
+		asm.LoadMem(asm.R6, asm.R5, 0, asm.Word),
 		asm.Add.Imm(asm.R6, 1),
-		asm.StoreMem(asm.R5, 0, asm.R6, asm.DWord),
-		asm.Ja.Label(nextLabel),
-		// Otherwise create the key and insert it
-		asm.Mov.Reg(asm.R3, asm.RFP).WithSymbol(initLabel),
-		asm.Add.Imm(asm.R3, int32(dropStatsValStackOffset)),
-		asm.Mov.Imm(asm.R4, 1),
-		asm.StoreMem(asm.R3, 0, asm.R4, asm.DWord),
-		asm.LoadMapPtr(asm.R1, dropStatsMapFd),
-		asm.Mov.Reg(asm.R2, asm.RFP),
-		asm.Add.Imm(asm.R2, int32(dropStatsKeyStackOffset)),
-		asm.Mov.Reg(asm.R3, asm.RFP),
-		asm.Add.Imm(asm.R3, int32(dropStatsValStackOffset)),
-		asm.Mov.Imm(asm.R4, 0),
-		asm.FnMapUpdateElem.Call(),
+		asm.StoreMem(asm.R5, 0, asm.R6, asm.Word),
 		asm.Ja.Label(nextLabel),
 	}
 }
@@ -270,6 +256,10 @@ func filtersToProgs(filters []Filter, opts ProgOpts, headerInsts, footerInsts as
 	)
 
 	for i, filter := range filters {
+		if i >= MaxDropActionFilters {
+			mErr = multierror.Append(mErr, fmt.Errorf("too many filters, stop adding filters, max is %d", MaxDropActionFilters))
+			break
+		}
 		filterInsts, err := FilterToInsts(i, filter, opts)
 		if err != nil {
 			mErr = multierror.Append(mErr, fmt.Errorf("unable to generate eBPF bytecode for rule `%s`: %s", filter.RuleID, err))

pkg/security/probe/monitors/rawpacketdrop/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
             "//pkg/security/metrics",
             "//pkg/security/probe/managerhelper",
             "//pkg/security/seclog",
+            "//pkg/security/utils",
             "@com_github_cilium_ebpf//:ebpf",
             "@com_github_datadog_datadog_go_v5//statsd",
             "@com_github_datadog_ebpf_manager//:ebpf-manager",
@@ -18,6 +19,7 @@ go_library(
             "//pkg/security/metrics",
             "//pkg/security/probe/managerhelper",
             "//pkg/security/seclog",
+            "//pkg/security/utils",
             "@com_github_cilium_ebpf//:ebpf",
             "@com_github_datadog_datadog_go_v5//statsd",
             "@com_github_datadog_ebpf_manager//:ebpf-manager",

pkg/security/probe/monitors/rawpacketdrop/rawpacketdrop_monitor.go

@@ -10,6 +10,7 @@ package rawpacketdrop
 
 import (
 	"fmt"
+	"sync"
 
 	manager "github.com/DataDog/ebpf-manager"
 	lib "github.com/cilium/ebpf"
@@ -19,6 +20,7 @@ import (
 	"github.com/DataDog/datadog-agent/pkg/security/metrics"
 	"github.com/DataDog/datadog-agent/pkg/security/probe/managerhelper"
 	"github.com/DataDog/datadog-agent/pkg/security/seclog"
+	"github.com/DataDog/datadog-agent/pkg/security/utils"
 )
 
 // RuleIDsProvider returns the current filter index to rule_id mapping.
@@ -30,6 +32,8 @@ type Monitor struct {
 	droppedMap   *lib.Map
 	ruleIDs      RuleIDsProvider
 	lastCounts   map[string]uint64
+	mu           sync.Mutex
+	numCPU       int
 }
 
 // NewMonitor returns a new Monitor.
@@ -38,48 +42,67 @@ func NewMonitor(manager *manager.Manager, statsdClient statsd.ClientInterface, r
 	if err != nil {
 		return nil, err
 	}
+	numCPU, err := utils.NumCPU()
+	if err != nil {
+		return nil, fmt.Errorf("couldn't fetch the host CPU count: %w", err)
+	}
 
 	return &Monitor{
 		statsdClient: statsdClient,
 		droppedMap:   droppedMap,
 		ruleIDs:      ruleIDs,
 		lastCounts:   make(map[string]uint64),
+		numCPU:       numCPU,
 	}, nil
 }
 
 // ResetCounters clears user-space counters after the kernel map is reset.
 func (m *Monitor) ResetCounters() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
 	m.lastCounts = make(map[string]uint64)
 }
 
 // SendStats emits deltas from the kernel dropped_packets map grouped by rule_id.
 func (m *Monitor) SendStats() error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	// get the up to date corresponding rule IDs for each filter
 	ruleIDs := m.ruleIDs()
 	if len(ruleIDs) == 0 {
 		m.lastCounts = make(map[string]uint64)
 		return nil
 	}
-
 	currentCounts := make(map[string]uint64, len(ruleIDs))
-	iterator := m.droppedMap.Iterate()
 
-	var filterIndex uint64
+	perCPU := make([]uint32, m.numCPU)
+
 	var count uint64
-	for iterator.Next(&filterIndex, &count) {
-		ruleID, ok := ruleIDs[uint32(filterIndex)]
-		if !ok || ruleID == "" {
+	// get the current counts from the kernel map
+	for filterIndex, ruleID := range ruleIDs {
+		if ruleID == "" {
+			continue
+		}
+		if err := m.droppedMap.Lookup(filterIndex, &perCPU); err != nil {
+			seclog.Warnf("failed to lookup dropped_packets map: %s", err)
 			continue
 		}
+		count = 0
+		for _, value := range perCPU {
+			count += uint64(value)
+		}
 		currentCounts[ruleID] += count
 	}
+
 	for ruleID, count := range currentCounts {
 		last := m.lastCounts[ruleID]
-		delta := count - last
-		if delta < 0 {
-			seclog.Errorf("incorrect mapping leading to a negative delta for rule_id %s: %d", ruleID, delta)
+
+		if count < last {
+			seclog.Errorf("incorrect mapping leading to a negative delta for rule_id %s", ruleID)
 			continue
 		}
+		delta := count - last
 		if delta == 0 {
 			continue
 		}

pkg/security/probe/probe_ebpf.go

@@ -852,7 +852,7 @@ func (p *EBPFProbe) rebuildDropActionRuleIDs() {
 	ruleIDs := make(map[uint32]string, len(p.rawPacketActionFilters))
 	for i, filter := range p.rawPacketActionFilters {
 		if i >= rawpacket.MaxDropActionFilters {
-			seclog.Errorf("too many drop action filters, indexes might be incorrect because of the kernel LRU, max is %d", rawpacket.MaxDropActionFilters)
+			seclog.Errorf("too many drop action filters, stop adding them is ruleID mapping, max is %d", rawpacket.MaxDropActionFilters)
 			break
 		}
 		ruleIDs[uint32(i)] = string(filter.RuleID)
@@ -876,11 +876,9 @@ func (p *EBPFProbe) getDropActionRuleIDs() map[uint32]string {
 }
 
 func (p *EBPFProbe) clearDroppedPacketsMap(droppedPacketsMap *lib.Map) {
-	iterator := droppedPacketsMap.Iterate()
-	var key uint64
-	var value uint64
-	for iterator.Next(&key, &value) {
-		_ = droppedPacketsMap.Delete(key)
+	zero := make([]uint32, p.numCPU)
+	for i := uint32(0); i < rawpacket.MaxDropActionFilters; i++ {
+		_ = droppedPacketsMap.Put(i, zero)
 	}
 }