From acda5afc54350c90902622f4ffc2a172d765807e Mon Sep 17 00:00:00 2001
From: Yonatan Goldschmidt <yonatan.goldschmidt@wiz.io>
Date: Sat, 18 Jan 2025 18:14:15 +0200
Subject: [PATCH] map: Add max_entries() for Map objects

Signed-off-by: Yonatan Goldschmidt <yonatan.goldschmidt@wiz.io>
---
 libbpf-rs/src/map.rs    | 23 +++++++++++++++++++++++
 libbpf-rs/tests/test.rs | 41 +++++++++++++++++++++++++++++++++++++++--
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/libbpf-rs/src/map.rs b/libbpf-rs/src/map.rs
index 3d87120c..6dac55ba 100644
--- a/libbpf-rs/src/map.rs
+++ b/libbpf-rs/src/map.rs
@@ -103,6 +103,11 @@ impl<'obj> OpenMap<'obj> {
             Some(data)
         }
     }
+
+    /// Retrieve max_entries of the map.
+    pub fn max_entries(&self) -> u32 {
+        unsafe { libbpf_sys::bpf_map__max_entries(self.ptr.as_ptr()) }
+    }
 }
 
 impl<'obj> OpenMapMut<'obj> {
@@ -438,6 +443,9 @@ pub trait MapCore: Debug + AsFd + private::Sealed {
     /// Retrieve the size of the map's values.
     fn value_size(&self) -> u32;
 
+    /// Retrieve max_entries of the map.
+    fn max_entries(&self) -> u32;
+
     /// Fetch extra map information
     #[inline]
     fn info(&self) -> Result<MapInfo> {
@@ -934,6 +942,11 @@ where
     fn value_size(&self) -> u32 {
         unsafe { libbpf_sys::bpf_map__value_size(self.ptr.as_ptr()) }
     }
+
+    #[inline]
+    fn max_entries(&self) -> u32 {
+        unsafe { libbpf_sys::bpf_map__max_entries(self.ptr.as_ptr()) }
+    }
 }
 
 impl AsRawLibbpf for Map<'_> {
@@ -967,6 +980,7 @@ pub struct MapHandle {
     ty: MapType,
     key_size: u32,
     value_size: u32,
+    max_entries: u32,
 }
 
 impl MapHandle {
@@ -1014,6 +1028,7 @@ impl MapHandle {
             ty: map_type,
             key_size,
             value_size,
+            max_entries,
         })
     }
 
@@ -1064,6 +1079,7 @@ impl MapHandle {
             ty: info.map_type(),
             key_size: info.info.key_size,
             value_size: info.info.value_size,
+            max_entries: info.info.max_entries,
         })
     }
 
@@ -1116,6 +1132,11 @@ impl MapCore for MapHandle {
     fn value_size(&self) -> u32 {
         self.value_size
     }
+
+    #[inline]
+    fn max_entries(&self) -> u32 {
+        self.max_entries
+    }
 }
 
 impl AsFd for MapHandle {
@@ -1141,6 +1162,7 @@ where
             ty: other.map_type(),
             key_size: other.key_size(),
             value_size: other.value_size(),
+            max_entries: other.max_entries(),
         })
     }
 }
@@ -1158,6 +1180,7 @@ impl TryFrom<&MapHandle> for MapHandle {
             ty: other.map_type(),
             key_size: other.key_size(),
             value_size: other.value_size(),
+            max_entries: other.max_entries(),
         })
     }
 }
diff --git a/libbpf-rs/tests/test.rs b/libbpf-rs/tests/test.rs
index c2dc6249..14f3750e 100644
--- a/libbpf-rs/tests/test.rs
+++ b/libbpf-rs/tests/test.rs
@@ -1576,12 +1576,14 @@ fn test_object_map_handle_clone() {
     assert_eq!(map.map_type(), handle1.map_type());
     assert_eq!(map.key_size(), handle1.key_size());
     assert_eq!(map.value_size(), handle1.value_size());
+    assert_eq!(map.max_entries(), handle1.max_entries());
 
     let handle2 = MapHandle::try_from(&handle1).expect("failed to duplicate existing handle");
     assert_eq!(handle1.name(), handle2.name());
     assert_eq!(handle1.map_type(), handle2.map_type());
     assert_eq!(handle1.key_size(), handle2.key_size());
     assert_eq!(handle1.value_size(), handle2.value_size());
+    assert_eq!(handle1.max_entries(), handle2.max_entries());
 
     let info1 = map.info().expect("failed to get map info from map");
     let info2 = handle2.info().expect("failed to get map info from handle");
@@ -1999,10 +2001,10 @@ fn test_map_autocreate_disable() {
     open_obj.load().expect("failed to load object");
 }
 
-/// Check that we can resize a map.
+/// Check that we can resize a map value.
 #[tag(root)]
 #[test]
-fn test_map_resize() {
+fn test_map_resize_value() {
     bump_rlimit_mlock();
 
     let mut open_obj = open_test_object("map_auto_pin.bpf.o");
@@ -2021,6 +2023,41 @@ fn test_map_resize() {
     assert_eq!(new_len, len * 2);
 }
 
+/// Check that we can resize map max entries.
+#[tag(root)]
+#[test]
+fn test_object_map_max_entries() {
+    bump_rlimit_mlock();
+
+    let mut obj = open_test_object("runqslower.bpf.o");
+
+    // resize the map to have twice the number of entries
+    let mut start = obj
+        .maps_mut()
+        .find(|map| map.name() == OsStr::new("start"))
+        .expect("failed to find `start` map");
+    let initial_max_entries = start.max_entries();
+    let new_max_entries = initial_max_entries * 2;
+    start
+        .set_max_entries(new_max_entries)
+        .expect("failed to set max entries");
+    // check that it reflects on the open map
+    assert_eq!(start.max_entries(), new_max_entries);
+
+    // check that it reflects after loading the map
+    let obj = obj.load().expect("failed to load object");
+    let start = obj
+        .maps()
+        .find(|map| map.name() == OsStr::new("start"))
+        .expect("failed to find `start` map");
+    assert_eq!(start.max_entries(), new_max_entries);
+
+    // check that it reflects after recreating the map handle from map id
+    let start = MapHandle::from_map_id(start.info().expect("failed to get map info").info.id)
+        .expect("failed to get map handle from id");
+    assert!(start.max_entries() == new_max_entries);
+}
+
 /// Check that we are able to attach using ksyscall
 #[tag(root)]
 #[test]