From c3ed425215b67514afa6a46f90adaccfa2074cd4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20M=C3=BCller?= <deso@posteo.net>
Date: Fri, 8 Nov 2024 11:09:00 -0800
Subject: [PATCH] Add a few changes for #996
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add a few fixes on top of pull request #996. Among others, remove an
'as' cast which potentially could be the source of unsoundness, remove
some other unnecessary ones, and don't repeat code needlessly.

Signed-off-by: Daniel Müller <deso@posteo.net>
---
 libbpf-rs/CHANGELOG.md |  6 +++++-
 libbpf-rs/src/map.rs   | 38 +++++++++++++++++---------------------
 2 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/libbpf-rs/CHANGELOG.md b/libbpf-rs/CHANGELOG.md
index 056a1269..9d41763e 100644
--- a/libbpf-rs/CHANGELOG.md
+++ b/libbpf-rs/CHANGELOG.md
@@ -1,8 +1,12 @@
+Unreleased
+----------
+- Added `Map::lookup_batch` and `Map::lookup_and_delete_batch` method
+
+
 0.24.8
 ------
 - Added `Program::attach_netfilter_with_opts` for attaching to netfilter
   hooks
-- Added `Map::lookup_batch` and `Map::lookup_and_delete_batch` method
 
 
 0.24.5
diff --git a/libbpf-rs/src/map.rs b/libbpf-rs/src/map.rs
index 65f59e11..3d87120c 100644
--- a/libbpf-rs/src/map.rs
+++ b/libbpf-rs/src/map.rs
@@ -1397,31 +1397,27 @@ impl<'map> BatchedMapIter<'map> {
     }
 
     fn lookup_next_batch(&mut self) {
-        let prev = self.prev.as_ref().map_or(ptr::null(), |p| p.as_ptr());
+        let prev = self
+            .prev
+            .as_mut()
+            .map_or(ptr::null_mut(), |p| p.as_mut_ptr());
         let mut count = self.count as u32;
 
         let ret = unsafe {
-            if self.delete {
-                libbpf_sys::bpf_map_lookup_and_delete_batch(
-                    self.map_fd.as_raw_fd(),
-                    prev as _,
-                    self.next.as_mut_ptr().cast(),
-                    self.keys.as_mut_ptr().cast(),
-                    self.values.as_mut_ptr().cast(),
-                    (&mut count) as *mut u32,
-                    &self.batch_opts as *const libbpf_sys::bpf_map_batch_opts,
-                )
+            let lookup_fn = if self.delete {
+                libbpf_sys::bpf_map_lookup_and_delete_batch
             } else {
-                libbpf_sys::bpf_map_lookup_batch(
-                    self.map_fd.as_raw_fd(),
-                    prev as _,
-                    self.next.as_mut_ptr().cast(),
-                    self.keys.as_mut_ptr().cast(),
-                    self.values.as_mut_ptr().cast(),
-                    (&mut count) as *mut u32,
-                    &self.batch_opts as *const libbpf_sys::bpf_map_batch_opts,
-                )
-            }
+                libbpf_sys::bpf_map_lookup_batch
+            };
+            lookup_fn(
+                self.map_fd.as_raw_fd(),
+                prev.cast(),
+                self.next.as_mut_ptr().cast(),
+                self.keys.as_mut_ptr().cast(),
+                self.values.as_mut_ptr().cast(),
+                &mut count,
+                &self.batch_opts,
+            )
         };
 
         if let Err(e) = util::parse_ret(ret) {