aya/maps: support TryFrom for LRU hash maps

The macro to implement TryFrom for MapData didn't have the ability to specify that more than one variant of MapData can be valid for a single map implementation. Support for new syntax was added to the macro so that the implementation can succeed for both valid variants in the HashMap and PerCpuHashMap impl. Fixes #636
aya-rs · Sep 1, 2023 · 172859c · 172859c
1 parent 792f467
commit 172859c
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 53 deletions.
diff --git a/aya/src/maps/hash_map/hash_map.rs b/aya/src/maps/hash_map/hash_map.rs
@@ -108,43 +108,22 @@ mod tests {
     use libc::{EFAULT, ENOENT};
 
     use crate::{
-        bpf_map_def,
         generated::{
             bpf_attr, bpf_cmd,
             bpf_map_type::{BPF_MAP_TYPE_HASH, BPF_MAP_TYPE_LRU_HASH},
         },
-        maps::{Map, MapData},
-        obj::{self, maps::LegacyMap, BpfSectionKind},
+        maps::Map,
+        obj,
         sys::{override_syscall, SysResult, Syscall},
     };
 
-    use super::*;
+    use super::{
+        super::test_utils::{self, new_map},
+        *,
+    };
 
     fn new_obj_map() -> obj::Map {
-        obj::Map::Legacy(LegacyMap {
-            def: bpf_map_def {
-                map_type: BPF_MAP_TYPE_HASH as u32,
-                key_size: 4,
-                value_size: 4,
-                max_entries: 1024,
-                ..Default::default()
-            },
-            section_index: 0,
-            section_kind: BpfSectionKind::Maps,
-            data: Vec::new(),
-            symbol_index: None,
-        })
-    }
-
-    fn new_map(obj: obj::Map) -> MapData {
-        override_syscall(|call| match call {
-            Syscall::Bpf {
-                cmd: bpf_cmd::BPF_MAP_CREATE,
-                ..
-            } => Ok(1337),
-            call => panic!("unexpected syscall {:?}", call),
-        });
-        MapData::create(obj, "foo", None).unwrap()
+        test_utils::new_obj_map(BPF_MAP_TYPE_HASH)
     }
 
     fn sys_error(value: i32) -> SysResult<c_long> {
@@ -213,21 +192,10 @@ mod tests {
 
     #[test]
     fn test_try_from_ok_lru() {
-        let map = new_map(obj::Map::Legacy(LegacyMap {
-            def: bpf_map_def {
-                map_type: BPF_MAP_TYPE_LRU_HASH as u32,
-                key_size: 4,
-                value_size: 4,
-                max_entries: 1024,
-                ..Default::default()
-            },
-            section_index: 0,
-            section_kind: BpfSectionKind::Maps,
-            symbol_index: None,
-            data: Vec::new(),
-        }));
-        let map = Map::HashMap(map);
-
+        let map_data = || new_map(test_utils::new_obj_map(BPF_MAP_TYPE_LRU_HASH));
+        let map = Map::HashMap(map_data());
+        assert!(HashMap::<_, u32, u32>::try_from(&map).is_ok());
+        let map = Map::LruHashMap(map_data());
         assert!(HashMap::<_, u32, u32>::try_from(&map).is_ok())
     }
 

diff --git a/aya/src/maps/hash_map/mod.rs b/aya/src/maps/hash_map/mod.rs
@@ -41,3 +41,41 @@ pub(crate) fn remove<K: Pod>(map: &MapData, key: &K) -> Result<(), MapError> {
             .into()
         })
 }
+
+#[cfg(test)]
+mod test_utils {
+    use crate::{
+        bpf_map_def,
+        generated::{bpf_cmd, bpf_map_type},
+        maps::MapData,
+        obj::{self, maps::LegacyMap, BpfSectionKind},
+        sys::{override_syscall, Syscall},
+    };
+
+    pub(super) fn new_map(obj: obj::Map) -> MapData {
+        override_syscall(|call| match call {
+            Syscall::Bpf {
+                cmd: bpf_cmd::BPF_MAP_CREATE,
+                ..
+            } => Ok(1337),
+            call => panic!("unexpected syscall {:?}", call),
+        });
+        MapData::create(obj, "foo", None).unwrap()
+    }
+
+    pub(super) fn new_obj_map(map_type: bpf_map_type) -> obj::Map {
+        obj::Map::Legacy(LegacyMap {
+            def: bpf_map_def {
+                map_type: map_type as u32,
+                key_size: 4,
+                value_size: 4,
+                max_entries: 1024,
+                ..Default::default()
+            },
+            section_index: 0,
+            section_kind: BpfSectionKind::Maps,
+            data: Vec::new(),
+            symbol_index: None,
+        })
+    }
+}
diff --git a/aya/src/maps/hash_map/per_cpu_hash_map.rs b/aya/src/maps/hash_map/per_cpu_hash_map.rs
@@ -146,3 +146,30 @@ impl<T: Borrow<MapData>, K: Pod, V: Pod> IterableMap<K, PerCpuValues<V>>
         Self::get(self, key, 0)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        generated::bpf_map_type::{BPF_MAP_TYPE_LRU_PERCPU_HASH, BPF_MAP_TYPE_PERCPU_HASH},
+        maps::Map,
+    };
+
+    use super::{super::test_utils, *};
+
+    #[test]
+    fn test_try_from_ok() {
+        let map = Map::PerCpuHashMap(test_utils::new_map(test_utils::new_obj_map(
+            BPF_MAP_TYPE_PERCPU_HASH,
+        )));
+        assert!(PerCpuHashMap::<_, u32, u32>::try_from(&map).is_ok())
+    }
+    #[test]
+    fn test_try_from_ok_lru() {
+        let map_data =
+            || test_utils::new_map(test_utils::new_obj_map(BPF_MAP_TYPE_LRU_PERCPU_HASH));
+        let map = Map::PerCpuHashMap(map_data());
+        assert!(PerCpuHashMap::<_, u32, u32>::try_from(&map).is_ok());
+        let map = Map::PerCpuLruHashMap(map_data());
+        assert!(PerCpuHashMap::<_, u32, u32>::try_from(&map).is_ok())
+    }
+}
diff --git a/aya/src/maps/mod.rs b/aya/src/maps/mod.rs
@@ -291,27 +291,27 @@ macro_rules! impl_try_from_map {
     // rather than the repeated idents used later because the macro language does not allow one
     // repetition to be pasted inside another.
     ($ty_param:tt {
-        $($ty:ident $(from $variant:ident)?),+ $(,)?
+        $($ty:ident $(from $($variant:ident)|+)?),+ $(,)?
     }) => {
-        $(impl_try_from_map!(<$ty_param> $ty $(from $variant)?);)+
+        $(impl_try_from_map!(<$ty_param> $ty $(from $($variant)|+)?);)+
     };
     // Add the "from $variant" using $ty as the default if it is missing.
     (<$ty_param:tt> $ty:ident) => {
         impl_try_from_map!(<$ty_param> $ty from $ty);
     };
     // Dispatch for each of the lifetimes.
     (
-        <($($ty_param:ident),*)> $ty:ident from $variant:ident
+        <($($ty_param:ident),*)> $ty:ident from $($variant:ident)|+
     ) => {
-        impl_try_from_map!(<'a> ($($ty_param),*) $ty from $variant);
-        impl_try_from_map!(<'a mut> ($($ty_param),*) $ty from $variant);
-        impl_try_from_map!(<> ($($ty_param),*) $ty from $variant);
+        impl_try_from_map!(<'a> ($($ty_param),*) $ty from $($variant)|+);
+        impl_try_from_map!(<'a mut> ($($ty_param),*) $ty from $($variant)|+);
+        impl_try_from_map!(<> ($($ty_param),*) $ty from $($variant)|+);
     };
     // An individual impl.
     (
         <$($l:lifetime $($m:ident)?)?>
         ($($ty_param:ident),*)
-        $ty:ident from $variant:ident
+        $ty:ident from $($variant:ident)|+
     ) => {
         impl<$($l,)? $($ty_param: Pod),*> TryFrom<$(&$l $($m)?)? Map>
             for $ty<$(&$l $($m)?)? MapData, $($ty_param),*>
@@ -320,7 +320,7 @@ macro_rules! impl_try_from_map {
 
             fn try_from(map: $(&$l $($m)?)? Map) -> Result<Self, Self::Error> {
                 match map {
-                    Map::$variant(map_data) => Self::new(map_data),
+                    $(Map::$variant(map_data) => Self::new(map_data),)+
                     map => Err(MapError::InvalidMapType {
                         map_type: map.map_type()
                     }),
@@ -353,8 +353,8 @@ impl_try_from_map!((V) {
 });
 
 impl_try_from_map!((K, V) {
-    HashMap,
-    PerCpuHashMap,
+    HashMap from HashMap|LruHashMap,
+    PerCpuHashMap from PerCpuHashMap|PerCpuLruHashMap,
     LpmTrie,
 });