From 25d986a26d9c88cd499a8b795054d583f01476b2 Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Mon, 22 Jul 2024 11:08:48 +0100
Subject: [PATCH] fix(aya): Set PerfEventArray max_entries to nCPUs

Both libbpf and cilium/ebpf have will set the max_entries of a
BPF_MAP_TYPE_PERF_EVENT_ARRAY to the number of online CPUs if
it was omitted at map definition time. This adds that same
logic to Aya.

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 aya/src/maps/bloom_filter.rs |  8 +++-----
 aya/src/maps/lpm_trie.rs     |  8 +++-----
 aya/src/sys/bpf.rs           | 22 +++++++++++++++++++++-
 3 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/aya/src/maps/bloom_filter.rs b/aya/src/maps/bloom_filter.rs
index 944f3f76..daeaa46c 100644
--- a/aya/src/maps/bloom_filter.rs
+++ b/aya/src/maps/bloom_filter.rs
@@ -88,7 +88,7 @@ mod tests {
     use crate::{
         generated::{
             bpf_cmd,
-            bpf_map_type::{BPF_MAP_TYPE_BLOOM_FILTER, BPF_MAP_TYPE_PERF_EVENT_ARRAY},
+            bpf_map_type::{BPF_MAP_TYPE_ARRAY, BPF_MAP_TYPE_BLOOM_FILTER},
         },
         maps::{
             test_utils::{self, new_map},
@@ -120,10 +120,8 @@ mod tests {
 
     #[test]
     fn test_try_from_wrong_map() {
-        let map = new_map(test_utils::new_obj_map::<u32>(
-            BPF_MAP_TYPE_PERF_EVENT_ARRAY,
-        ));
-        let map = Map::PerfEventArray(map);
+        let map = new_map(test_utils::new_obj_map::<u32>(BPF_MAP_TYPE_ARRAY));
+        let map = Map::Array(map);
 
         assert_matches!(
             BloomFilter::<_, u32>::try_from(&map),
diff --git a/aya/src/maps/lpm_trie.rs b/aya/src/maps/lpm_trie.rs
index e2271908..37377efa 100644
--- a/aya/src/maps/lpm_trie.rs
+++ b/aya/src/maps/lpm_trie.rs
@@ -205,7 +205,7 @@ mod tests {
     use crate::{
         generated::{
             bpf_cmd,
-            bpf_map_type::{BPF_MAP_TYPE_LPM_TRIE, BPF_MAP_TYPE_PERF_EVENT_ARRAY},
+            bpf_map_type::{BPF_MAP_TYPE_ARRAY, BPF_MAP_TYPE_LPM_TRIE},
         },
         maps::{
             test_utils::{self, new_map},
@@ -249,10 +249,8 @@ mod tests {
 
     #[test]
     fn test_try_from_wrong_map() {
-        let map = new_map(test_utils::new_obj_map::<u32>(
-            BPF_MAP_TYPE_PERF_EVENT_ARRAY,
-        ));
-        let map = Map::PerfEventArray(map);
+        let map = new_map(test_utils::new_obj_map::<u32>(BPF_MAP_TYPE_ARRAY));
+        let map = Map::Array(map);
 
         assert_matches!(
             LpmTrie::<_, u32, u32>::try_from(&map),
diff --git a/aya/src/sys/bpf.rs b/aya/src/sys/bpf.rs
index f327d2bb..f305f56b 100644
--- a/aya/src/sys/bpf.rs
+++ b/aya/src/sys/bpf.rs
@@ -30,7 +30,7 @@ use crate::{
         copy_instructions,
     },
     sys::{syscall, SysResult, Syscall, SyscallError},
-    util::KernelVersion,
+    util::{nr_cpus, KernelVersion},
     Btf, Pod, VerifierLogLevel, BPF_OBJ_NAME_LEN,
 };
 
@@ -46,6 +46,26 @@ pub(crate) fn bpf_create_map(
     u.map_type = def.map_type();
     u.key_size = def.key_size();
     u.value_size = def.value_size();
+    u.max_entries = def.max_entries();
+
+    // BPF_MAP_TYPE_PERF_EVENT_ARRAY's max_entries should not exceed the number of
+    // CPUs.
+    //
+    // By default, the newest versions of Aya, libbpf and cilium/ebpf define `max_entries` of
+    // `PerfEventArray` as `0`, with an intention to get it replaced with a correct value
+    // by the loader.
+    //
+    // We allow custom values (potentially coming either from older versions of aya-ebpf or
+    // programs written in C) as long as they don't exceed the number of CPUs.
+    //
+    // Otherwise, when the value is `0` or too large, we set it to the number of CPUs.
+    if def.map_type() == bpf_map_type::BPF_MAP_TYPE_PERF_EVENT_ARRAY as u32 {
+        let ncpus = nr_cpus().map_err(|e| (-1i64, e))? as u32;
+        if u.max_entries == 0 || u.max_entries > ncpus {
+            u.max_entries = ncpus;
+        }
+    };
+
     u.max_entries = def.max_entries();
     u.map_flags = def.map_flags();