diff --git a/libbpf-rs/src/map.rs b/libbpf-rs/src/map.rs index 8d2695ec..591b0990 100644 --- a/libbpf-rs/src/map.rs +++ b/libbpf-rs/src/map.rs @@ -340,13 +340,24 @@ where key.as_ptr().cast::() } +/// Internal selector for which underlying lookup operation to perform. +/// +/// `bpf_map_lookup_elem_flags` accepts `MapFlags`, while +/// `bpf_map_lookup_and_delete_elem` does not. Wrapping the choice in an +/// enum keeps the two cases distinct without leaking an unused `flags` +/// argument into the lookup-and-delete path. +enum LookupOp { + Lookup(MapFlags), + LookupAndDelete, +} + /// Internal function to perform a map lookup and write the value into raw pointer. /// Returns `Ok(true)` if the key was found, `Ok(false)` if not found, or an error. fn lookup_raw( map: &M, key: &[u8], value: &mut [mem::MaybeUninit], - flags: MapFlags, + op: LookupOp, ) -> Result where M: MapCore + ?Sized, @@ -370,13 +381,20 @@ where ); let ret = unsafe { - libbpf_sys::bpf_map_lookup_elem_flags( - map.as_fd().as_raw_fd(), - map_key(map, key), - // TODO: Use `MaybeUninit::slice_as_mut_ptr` once stable. - value.as_mut_ptr().cast(), - flags.bits(), - ) + match op { + LookupOp::Lookup(flags) => libbpf_sys::bpf_map_lookup_elem_flags( + map.as_fd().as_raw_fd(), + map_key(map, key), + // TODO: Use `MaybeUninit::slice_as_mut_ptr` once stable. + value.as_mut_ptr().cast(), + flags.bits(), + ), + LookupOp::LookupAndDelete => libbpf_sys::bpf_map_lookup_and_delete_elem( + map.as_fd().as_raw_fd(), + map_key(map, key), + value.as_mut_ptr().cast(), + ), + } }; if ret == 0 { @@ -392,19 +410,14 @@ where } /// Internal function to return a value from a map into a buffer of the given size. -fn lookup_raw_vec( - map: &M, - key: &[u8], - flags: MapFlags, - out_size: usize, -) -> Result>> +fn lookup_raw_vec(map: &M, key: &[u8], op: LookupOp, out_size: usize) -> Result>> where M: MapCore + ?Sized, { // Allocate without initializing (avoiding memset) let mut out = Vec::with_capacity(out_size); - match lookup_raw(map, key, out.spare_capacity_mut(), flags)? { + match lookup_raw(map, key, out.spare_capacity_mut(), op)? { true => { // SAFETY: `lookup_raw` successfully filled the buffer unsafe { @@ -555,7 +568,7 @@ pub trait MapCore: Debug + AsFd + private::Sealed { fn lookup(&self, key: &[u8], flags: MapFlags) -> Result>> { check_not_bloom_or_percpu(self)?; let out_size = self.value_size() as usize; - lookup_raw_vec(self, key, flags, out_size) + lookup_raw_vec(self, key, LookupOp::Lookup(flags), out_size) } /// Looks up a map value into a pre-allocated buffer, avoiding allocation. @@ -588,7 +601,7 @@ pub trait MapCore: Debug + AsFd + private::Sealed { value.len(), ) }; - lookup_raw(self, key, value, flags) + lookup_raw(self, key, value, LookupOp::Lookup(flags)) } /// Returns many elements in batch mode from the map. @@ -656,7 +669,7 @@ pub trait MapCore: Debug + AsFd + private::Sealed { let aligned_val_size = percpu_aligned_value_size(self); let out_size = percpu_buffer_size(self)?; - let raw_res = lookup_raw_vec(self, key, flags, out_size)?; + let raw_res = lookup_raw_vec(self, key, LookupOp::Lookup(flags), out_size)?; if let Some(raw_vals) = raw_res { let mut out = Vec::new(); for chunk in raw_vals.chunks_exact(aligned_val_size) { @@ -728,42 +741,44 @@ pub trait MapCore: Debug + AsFd + private::Sealed { /// Same as [`Self::lookup()`] except this also deletes the key from the map. /// - /// Note that this operation is currently only implemented in the kernel for [`MapType::Queue`] - /// and [`MapType::Stack`]. + /// Implemented in the kernel for [`MapType::Queue`] and [`MapType::Stack`], + /// and (since Linux 5.14) for [`MapType::Hash`] / [`MapType::LruHash`] / + /// [`MapType::PercpuHash`] / [`MapType::LruPercpuHash`]. /// /// `key` must have exactly [`Self::key_size()`] elements. fn lookup_and_delete(&self, key: &[u8]) -> Result>> { - if key.len() != self.key_size() as usize { + let out_size = self.value_size() as usize; + lookup_raw_vec(self, key, LookupOp::LookupAndDelete, out_size) + } + + /// Same as [`Self::lookup_into()`] except this also deletes the key from the map. + /// + /// This method provides a zero-allocation alternative to [`Self::lookup_and_delete()`]. + /// + /// `key` must have exactly [`Self::key_size()`] elements. + /// `value` must have exactly [`Self::value_size()`] elements. + /// + /// Returns `Ok(true)` if the key was found and the buffer was filled, + /// `Ok(false)` if the key was not found, or an error. + /// + /// See [`Self::lookup_and_delete()`] for kernel support details. + fn lookup_into_and_delete(&self, key: &[u8], value: &mut [u8]) -> Result { + if value.len() != self.value_size() as usize { return Err(Error::with_invalid_data(format!( - "key_size {} != {}", - key.len(), - self.key_size() + "value buffer size {} != {}", + value.len(), + self.value_size() ))); - }; - - let mut out: Vec = Vec::with_capacity(self.value_size() as usize); + } - let ret = unsafe { - libbpf_sys::bpf_map_lookup_and_delete_elem( - self.as_fd().as_raw_fd(), - map_key(self, key), - out.as_mut_ptr().cast::(), + // SAFETY: `u8` and `MaybeUninit` have the same in-memory representation. + let value = unsafe { + slice::from_raw_parts_mut::>( + value.as_mut_ptr().cast(), + value.len(), ) }; - - if ret == 0 { - unsafe { - out.set_len(self.value_size() as usize); - } - Ok(Some(out)) - } else { - let err = io::Error::last_os_error(); - if err.kind() == io::ErrorKind::NotFound { - Ok(None) - } else { - Err(Error::from(err)) - } - } + lookup_raw(self, key, value, LookupOp::LookupAndDelete) } /// Update an element. diff --git a/libbpf-rs/tests/test.rs b/libbpf-rs/tests/test.rs index d595799a..999687eb 100644 --- a/libbpf-rs/tests/test.rs +++ b/libbpf-rs/tests/test.rs @@ -777,6 +777,81 @@ fn test_object_map_lookup_into_consistency() { ); } +/// Test `lookup_and_delete` on a regular HASH map. The kernel grew support for +/// `bpf_map_lookup_and_delete_elem` on hash maps in v5.14 (commit 3e87f192b405); +/// before that, only Queue and Stack were covered (already exercised in +/// `test_object_map_queue_crud` / `test_object_map_stack_crud`). +#[tag(root)] +#[test] +fn test_object_map_lookup_and_delete_hash() { + let mut obj = get_test_object("runqslower.bpf.o"); + let start = get_map_mut(&mut obj, "start"); + + start + .update(&[1, 2, 3, 4], &[1, 2, 3, 4, 5, 6, 7, 8], MapFlags::empty()) + .expect("failed to write"); + + let val = start + .lookup_and_delete(&[1, 2, 3, 4]) + .expect("failed to lookup_and_delete") + .expect("key should exist"); + assert_eq!(val, &[1, 2, 3, 4, 5, 6, 7, 8]); + + // The entry must be gone after lookup_and_delete. + assert!(start + .lookup(&[1, 2, 3, 4], MapFlags::empty()) + .expect("failed to lookup") + .is_none()); + + // A second lookup_and_delete reports the missing key as Ok(None). + assert!(start + .lookup_and_delete(&[1, 2, 3, 4]) + .expect("failed to lookup_and_delete on missing key") + .is_none()); +} + +#[tag(root)] +#[test] +fn test_object_map_lookup_into_and_delete() { + let mut obj = get_test_object("runqslower.bpf.o"); + let start = get_map_mut(&mut obj, "start"); + + start + .update(&[1, 2, 3, 4], &[1, 2, 3, 4, 5, 6, 7, 8], MapFlags::empty()) + .expect("failed to write"); + + let mut value = [0u8; 8]; + let found = start + .lookup_into_and_delete(&[1, 2, 3, 4], &mut value) + .expect("failed to lookup_into_and_delete"); + assert!(found, "key should be found"); + assert_eq!(value, [1, 2, 3, 4, 5, 6, 7, 8]); + + // The entry must be gone afterwards. + assert!(start + .lookup(&[1, 2, 3, 4], MapFlags::empty()) + .expect("failed to lookup") + .is_none()); + + // Missing keys return Ok(false) and leave the buffer untouched. + let mut value2 = [42u8; 8]; + let found2 = start + .lookup_into_and_delete(&[1, 2, 3, 4], &mut value2) + .expect("failed to lookup_into_and_delete on missing key"); + assert!(!found2); + assert_eq!(value2, [42u8; 8]); + + // Wrong-sized buffers and wrong-sized keys must error. + let mut value_small = [0u8; 4]; + assert!(start + .lookup_into_and_delete(&[1, 2, 3, 4], &mut value_small) + .is_err()); + let mut value_ok = [0u8; 8]; + assert!(start + .lookup_into_and_delete(&[1, 2, 3, 4, 5], &mut value_ok) + .is_err()); +} + #[tag(root)] #[test] fn test_object_map_lookup_flags() {