From a761d04ccf7a5a06cface27f5364227f77dbc908 Mon Sep 17 00:00:00 2001
From: aorhant <aorhant@scaleway.com>
Date: Wed, 26 Feb 2025 18:54:13 +0100
Subject: [PATCH] aya: Fix PerCpuHashMap NotFound

PerCpuHashMap was never returning MapError::KeyNotFound because
bpf_map_lookup_elem_per_cpu was replacing Ok(None) with
Ok(Some(zeroed_value)).

Update bpf_map_lookup_elem_per_cpu to map the Option value.
---
 aya/src/maps/hash_map/per_cpu_hash_map.rs | 24 ++++++++++++++++++++++-
 aya/src/sys/bpf.rs                        |  2 +-
 aya/src/util.rs                           | 23 ++++++++++++++--------
 3 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/aya/src/maps/hash_map/per_cpu_hash_map.rs b/aya/src/maps/hash_map/per_cpu_hash_map.rs
index d9194668..afb76620 100644
--- a/aya/src/maps/hash_map/per_cpu_hash_map.rs
+++ b/aya/src/maps/hash_map/per_cpu_hash_map.rs
@@ -151,12 +151,23 @@ impl<T: Borrow<MapData>, K: Pod, V: Pod> IterableMap<K, PerCpuValues<V>>
 
 #[cfg(test)]
 mod tests {
+    use std::{ffi::c_long, io};
+
+    use assert_matches::assert_matches;
     use aya_obj::generated::bpf_map_type::{
         BPF_MAP_TYPE_LRU_PERCPU_HASH, BPF_MAP_TYPE_PERCPU_HASH,
     };
+    use libc::ENOENT;
 
     use super::*;
-    use crate::maps::{test_utils, Map};
+    use crate::{
+        maps::{test_utils, Map},
+        sys::{override_syscall, SysResult},
+    };
+
+    fn sys_error(value: i32) -> SysResult<c_long> {
+        Err((-1, io::Error::from_raw_os_error(value)))
+    }
 
     #[test]
     fn test_try_from_ok() {
@@ -174,4 +185,15 @@ mod tests {
         let map = Map::PerCpuLruHashMap(map_data());
         assert!(PerCpuHashMap::<_, u32, u32>::try_from(&map).is_ok())
     }
+    #[test]
+    fn test_get_not_found() {
+        let map_data =
+            || test_utils::new_map(test_utils::new_obj_map::<u32>(BPF_MAP_TYPE_LRU_PERCPU_HASH));
+        let map = Map::PerCpuHashMap(map_data());
+        let map = PerCpuHashMap::<_, u32, u32>::try_from(&map).unwrap();
+
+        override_syscall(|_| sys_error(ENOENT));
+
+        assert_matches!(map.get(&1, 0), Err(MapError::KeyNotFound));
+    }
 }
diff --git a/aya/src/sys/bpf.rs b/aya/src/sys/bpf.rs
index 791959f1..51ca7a97 100644
--- a/aya/src/sys/bpf.rs
+++ b/aya/src/sys/bpf.rs
@@ -249,7 +249,7 @@ pub(crate) fn bpf_map_lookup_elem_per_cpu<K: Pod, V: Pod>(
 ) -> SysResult<Option<PerCpuValues<V>>> {
     let mut mem = PerCpuValues::<V>::alloc_kernel_mem().map_err(|io_error| (-1, io_error))?;
     match bpf_map_lookup_elem_ptr(fd, Some(key), mem.as_mut_ptr(), flags) {
-        Ok(_) => Ok(Some(unsafe { PerCpuValues::from_kernel_mem(mem) })),
+        Ok(v) => Ok(v.map(|()| unsafe { PerCpuValues::from_kernel_mem(mem) })),
         Err((_, io_error)) if io_error.raw_os_error() == Some(ENOENT) => Ok(None),
         Err(e) => Err(e),
     }
diff --git a/aya/src/util.rs b/aya/src/util.rs
index 377ccc59..f94e4540 100644
--- a/aya/src/util.rs
+++ b/aya/src/util.rs
@@ -194,15 +194,22 @@ pub fn online_cpus() -> Result<Vec<u32>, (&'static str, io::Error)> {
 ///
 /// See `/sys/devices/system/cpu/possible`.
 pub fn nr_cpus() -> Result<usize, (&'static str, io::Error)> {
-    thread_local! {
-        // TODO(https://github.com/rust-lang/rust/issues/109737): Use
-        // `std::cell::OnceCell` when `get_or_try_init` is stabilized.
-        static CACHE: once_cell::unsync::OnceCell<usize> = const { once_cell::unsync::OnceCell::new() };
+    #[cfg(miri)]
+    {
+        Ok(1)
+    }
+    #[cfg(not(miri))]
+    {
+        thread_local! {
+            // TODO(https://github.com/rust-lang/rust/issues/109737): Use
+            // `std::cell::OnceCell` when `get_or_try_init` is stabilized.
+            static CACHE: once_cell::unsync::OnceCell<usize> = const { once_cell::unsync::OnceCell::new() };
+        }
+        CACHE.with(|cell| {
+            cell.get_or_try_init(|| read_cpu_ranges(POSSIBLE_CPUS).map(|cpus| cpus.len()))
+                .copied()
+        })
     }
-    CACHE.with(|cell| {
-        cell.get_or_try_init(|| read_cpu_ranges(POSSIBLE_CPUS).map(|cpus| cpus.len()))
-            .copied()
-    })
 }
 
 fn read_cpu_ranges(path: &'static str) -> Result<Vec<u32>, (&'static str, io::Error)> {