test: add tests for einat skel

EHfive · Nov 20, 2024 · 9025fb9 · 9025fb9
1 parent 0e4fe71
commit 9025fb9
Show file tree

Hide file tree

Showing 3 changed files with 171 additions and 4 deletions.
diff --git a/src/macros.rs b/src/macros.rs
@@ -13,7 +13,7 @@ macro_rules! derive_pod {
         }
     ) => {
         $( #[$attr] )*
-        #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, ::bytemuck::Zeroable, ::bytemuck::Pod)]
+        #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, ::bytemuck::Zeroable, ::bytemuck::Pod)]
         $vis struct $name {
             $(
                $( #[$attr_f] )?

diff --git a/src/skel/einat/mod.rs b/src/skel/einat/mod.rs
@@ -113,3 +113,170 @@ impl<T> EinatEbpfSkel for T where T: EinatEbpf + EinatEbpfInet<Ipv4Net> {}
 pub trait EinatEbpfSkel: EinatEbpf + EinatEbpfInet<Ipv4Net> + EinatEbpfInet<Ipv6Net> {}
 #[cfg(feature = "ipv6")]
 impl<T> EinatEbpfSkel for T where T: EinatEbpf + EinatEbpfInet<Ipv4Net> + EinatEbpfInet<Ipv6Net> {}
+
+#[cfg(test)]
+mod tests {
+    use std::borrow::Borrow;
+    use std::collections::HashSet;
+    use std::fmt::Debug;
+    use std::hash::RandomState;
+    use std::net;
+
+    use super::*;
+    use crate::skel::{EbpfHashMap, EbpfLpmTrie, EbpfMapFlags};
+    use crate::utils::IpAddress;
+
+    #[cfg(feature = "aya")]
+    #[test]
+    #[ignore = "bpf"]
+    fn test_aya_maps() {
+        test_skel_maps::<aya::EinatAya>();
+        test_skel_inet_maps::<Ipv4Net, aya::EinatAya>();
+        #[cfg(feature = "ipv6")]
+        test_skel_inet_maps::<Ipv6Net, aya::EinatAya>();
+
+        test_skel_attach::<aya::EinatAya>();
+    }
+
+    #[cfg(feature = "libbpf")]
+    #[test]
+    #[ignore = "bpf"]
+    fn test_libbpf_maps() {
+        test_skel_maps::<libbpf::EinatLibbpf>();
+        test_skel_inet_maps::<Ipv4Net, libbpf::EinatLibbpf>();
+        #[cfg(feature = "ipv6")]
+        test_skel_inet_maps::<Ipv6Net, libbpf::EinatLibbpf>();
+
+        test_skel_attach::<libbpf::EinatLibbpf>();
+    }
+
+    #[cfg(feature = "libbpf-skel")]
+    #[test]
+    #[ignore = "bpf"]
+    fn test_libbpf_skel_maps() {
+        test_skel_maps::<libbpf_skel::EinatLibbpfSkel>();
+        test_skel_inet_maps::<Ipv4Net, libbpf_skel::EinatLibbpfSkel>();
+        #[cfg(feature = "ipv6")]
+        test_skel_inet_maps::<Ipv6Net, libbpf_skel::EinatLibbpfSkel>();
+
+        test_skel_attach::<libbpf_skel::EinatLibbpfSkel>();
+    }
+
+    fn test_skel_attach<T: EinatEbpf>() {
+        let mut skel = T::load(EinatConstConfig::default()).unwrap();
+        let links = skel.attach("lo", 1).unwrap();
+        skel.detach(links).unwrap();
+    }
+
+    // test if key and value struct size matches BTF map and general map operations
+    fn test_skel_maps<T: EinatEbpf>() {
+        let mut skel = T::load(EinatConstConfig::default()).unwrap();
+
+        skel.with_updating_wait(|_| {}).unwrap();
+
+        macro_rules! test_map {
+            ($map:ident, $kt:tt) => {{
+                let mut keys: Vec<_> = (0..100)
+                    .map(|i| $kt {
+                        if_index: i,
+                        ..Default::default()
+                    })
+                    .collect();
+
+                for k in keys.iter() {
+                    skel.$map()
+                        .update(k, &Default::default(), EbpfMapFlags::NO_EXIST)
+                        .unwrap();
+                }
+
+                let keys_set = HashSet::<_, RandomState>::from_iter(keys.iter());
+                for k in skel.$map().keys() {
+                    let k = k.unwrap();
+                    assert!(keys_set.contains(k.borrow()));
+                }
+
+                let key = keys.pop().unwrap();
+
+                assert!(skel
+                    .$map()
+                    .lookup(&key, EbpfMapFlags::ANY)
+                    .unwrap()
+                    .is_some());
+
+                skel.$map().delete(&key).unwrap();
+
+                assert!(skel
+                    .$map()
+                    .lookup(&key, EbpfMapFlags::ANY)
+                    .unwrap()
+                    .is_none());
+
+                skel.$map().delete_batch(&keys, EbpfMapFlags::ANY).unwrap();
+
+                for k in keys.iter() {
+                    assert!(skel.$map().lookup(k, EbpfMapFlags::ANY).unwrap().is_none());
+                }
+            }};
+        }
+
+        test_map!(map_binding_mut, MapBindingKey);
+        test_map!(map_ct_mut, MapCtKey);
+    }
+
+    trait InetKeyGen: Sized + IpAddress {
+        fn gen_idx(i: u32) -> Self {
+            Self::gen_idx_len(i, Self::LEN)
+        }
+
+        fn gen_idx_len(i: u32, len: u8) -> Self;
+    }
+
+    impl InetKeyGen for Ipv4Net {
+        fn gen_idx_len(i: u32, len: u8) -> Self {
+            IpNetwork::from(net::Ipv4Addr::from_bits(i as _), len)
+        }
+    }
+
+    #[cfg(feature = "ipv6")]
+    impl InetKeyGen for Ipv6Net {
+        fn gen_idx_len(i: u32, len: u8) -> Self {
+            IpNetwork::from(net::Ipv6Addr::from_bits(i as _), len)
+        }
+    }
+
+    fn test_skel_inet_maps<
+        P: IpNetwork + InetKeyGen + Eq + Debug + Copy,
+        T: EinatEbpf + EinatEbpfInet<P>,
+    >() {
+        let mut skel = T::load(EinatConstConfig::default()).unwrap();
+        let addr = P::gen_idx(1);
+        skel.set_external_addr(addr).unwrap();
+        assert_eq!(addr, skel.external_addr().unwrap());
+
+        macro_rules! test_map {
+            ($map:ident) => {{
+                let keys: Vec<_> = (0..P::LEN).map(|i| P::gen_idx_len(i as _, i)).collect();
+                for k in keys.iter() {
+                    skel.$map()
+                        .update(k, &Default::default(), EbpfMapFlags::NO_EXIST)
+                        .unwrap();
+                }
+
+                for k in keys.iter() {
+                    assert!(skel.$map().lookup(k, EbpfMapFlags::ANY).unwrap().is_some());
+                }
+
+                for k in keys.iter() {
+                    skel.$map().delete(k).unwrap();
+                }
+
+                for k in keys.iter() {
+                    assert!(skel.$map().lookup(k, EbpfMapFlags::ANY).unwrap().is_none());
+                }
+            }};
+        }
+
+        test_map!(map_external_config);
+        test_map!(map_dest_config);
+    }
+}
diff --git a/src/skel/einat/types.rs b/src/skel/einat/types.rs
@@ -124,7 +124,7 @@ derive_pod!(
 );
 
 bitflags! {
-    #[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Zeroable, Pod)]
+    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default, Zeroable, Pod)]
     #[repr(transparent)]
     pub struct ExternalFlags: u8 {
         const IS_INTERNAL = 0b1;
@@ -154,7 +154,7 @@ derive_pod!(
 );
 
 bitflags! {
-    #[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Zeroable, Pod)]
+    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default, Zeroable, Pod)]
     #[repr(transparent)]
     pub struct DestFlags: u8 {
         const HAIRPIN = 0b01;
@@ -170,7 +170,7 @@ derive_pod!(
 );
 
 bitflags! {
-    #[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Zeroable, Pod)]
+    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default, Zeroable, Pod)]
     #[repr(transparent)]
     pub struct BindingFlags: u8 {
         const ORIG_DIR = 0b001;