From 8009361694a7f8967a31734d109f79a6b26516dc Mon Sep 17 00:00:00 2001
From: Andrew Stoycos <astoycos@redhat.com>
Date: Tue, 4 Oct 2022 14:09:34 -0400
Subject: [PATCH] Fixups

Respond to review comments, specifically:

- Remove Map::map_type()

- Update some comments

- remove `docs` from feature macros

- generalize check_bounds, check_kv_size,
and check_v_size functions to remove
duplicate code

Signed-off-by: Andrew Stoycos <astoycos@redhat.com>
---
 aya/src/bpf.rs                            | 16 ++++++---
 aya/src/maps/array/array.rs               | 18 +++-------
 aya/src/maps/array/mod.rs                 | 11 -------
 aya/src/maps/array/per_cpu_array.rs       | 18 +++-------
 aya/src/maps/array/program_array.rs       | 18 +++-------
 aya/src/maps/bloom_filter.rs              | 10 ++----
 aya/src/maps/hash_map/hash_map.rs         |  4 +--
 aya/src/maps/hash_map/mod.rs              | 16 ---------
 aya/src/maps/hash_map/per_cpu_hash_map.rs | 14 +++-----
 aya/src/maps/lpm_trie.rs                  | 13 ++------
 aya/src/maps/mod.rs                       | 40 +++++++++++++++++++----
 aya/src/maps/queue.rs                     | 14 ++------
 aya/src/maps/sock/sock_hash.rs            |  6 ++--
 aya/src/maps/sock/sock_map.rs             | 33 +++++--------------
 aya/src/maps/stack.rs                     | 14 ++------
 15 files changed, 83 insertions(+), 162 deletions(-)

diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs
index e473802f..6b54fa51 100644
--- a/aya/src/bpf.rs
+++ b/aya/src/bpf.rs
@@ -13,7 +13,7 @@ use thiserror::Error;
 
 use crate::{
     generated::{
-        bpf_map_type::*, AYA_PERF_EVENT_IOC_DISABLE, AYA_PERF_EVENT_IOC_ENABLE,
+        bpf_map_type, bpf_map_type::*, AYA_PERF_EVENT_IOC_DISABLE, AYA_PERF_EVENT_IOC_ENABLE,
         AYA_PERF_EVENT_IOC_SET_BPF,
     },
     maps::{Map, MapData, MapError},
@@ -650,14 +650,16 @@ impl<'a> BpfLoader<'a> {
 fn parse_map(data: (String, MapData)) -> Result<(String, Map), BpfError> {
     let name = data.0;
     let map = data.1;
-    let map_type = map.map_type().map_err(BpfError::MapError)?;
+    let map_type = bpf_map_type::try_from(map.obj.map_type())?;
     let map = match map_type {
         BPF_MAP_TYPE_ARRAY => Ok(Map::Array(map)),
         BPF_MAP_TYPE_PERCPU_ARRAY => Ok(Map::PerCpuArray(map)),
         BPF_MAP_TYPE_PROG_ARRAY => Ok(Map::ProgramArray(map)),
         BPF_MAP_TYPE_HASH => Ok(Map::HashMap(map)),
         BPF_MAP_TYPE_PERCPU_HASH => Ok(Map::PerCpuHashMap(map)),
-        BPF_MAP_TYPE_PERF_EVENT_ARRAY => Ok(Map::PerfEventArray(map)),
+        BPF_MAP_TYPE_PERF_EVENT_ARRAY | BPF_MAP_TYPE_LRU_PERCPU_HASH => {
+            Ok(Map::PerfEventArray(map))
+        }
         BPF_MAP_TYPE_SOCKHASH => Ok(Map::SockHash(map)),
         BPF_MAP_TYPE_SOCKMAP => Ok(Map::SockMap(map)),
         BPF_MAP_TYPE_BLOOM_FILTER => Ok(Map::BloomFilter(map)),
@@ -770,9 +772,13 @@ impl Bpf {
             })
     }
 
-    /// Returns a map with the given name.
+    /// Takes ownership of a map with the given name.
     ///
-    /// WARNING: This transfers ownership of the map to the user.
+    /// This API is intended for cases where the map must be moved into spawned
+    /// task. For example, when using an [`AsyncPerfEventArray`]. For map interactions
+    /// without taking ownership, see `map` or `map_mut`. An owned map will be
+    /// closed on `Drop`, therefore the the caller is now responsible for managing
+    /// its lifetime.
     ///
     /// The returned type is mostly opaque. In order to do anything useful with it you need to
     /// convert it to a [typed map](crate::maps).
diff --git a/aya/src/maps/array/array.rs b/aya/src/maps/array/array.rs
index 8b60b8a5..d661be18 100644
--- a/aya/src/maps/array/array.rs
+++ b/aya/src/maps/array/array.rs
@@ -1,11 +1,10 @@
 use std::{
     convert::{AsMut, AsRef},
     marker::PhantomData,
-    mem,
 };
 
 use crate::{
-    maps::{array, IterableMap, MapData, MapError},
+    maps::{check_bounds, check_kv_size, IterableMap, MapData, MapError},
     sys::{bpf_map_lookup_elem, bpf_map_update_elem},
     Pod,
 };
@@ -38,17 +37,8 @@ pub struct Array<T, V: Pod> {
 impl<T: AsRef<MapData>, V: Pod> Array<T, V> {
     pub(crate) fn new(map: T) -> Result<Array<T, V>, MapError> {
         let data = map.as_ref();
-        let expected = mem::size_of::<u32>();
-        let size = data.obj.key_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidKeySize { size, expected });
-        }
+        check_kv_size::<u32, V>(data)?;
 
-        let expected = mem::size_of::<V>();
-        let size = data.obj.value_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidValueSize { size, expected });
-        }
         let _fd = data.fd_or_err()?;
 
         Ok(Array {
@@ -72,7 +62,7 @@ impl<T: AsRef<MapData>, V: Pod> Array<T, V> {
     /// if `bpf_map_lookup_elem` fails.
     pub fn get(&self, index: &u32, flags: u64) -> Result<V, MapError> {
         let data = self.inner.as_ref();
-        array::check_bounds(data, *index)?;
+        check_bounds(data, *index)?;
         let fd = data.fd_or_err()?;
 
         let value = bpf_map_lookup_elem(fd, index, flags).map_err(|(_, io_error)| {
@@ -100,7 +90,7 @@ impl<T: AsMut<MapData>, V: Pod> Array<T, V> {
     /// if `bpf_map_update_elem` fails.
     pub fn set(&mut self, index: u32, value: V, flags: u64) -> Result<(), MapError> {
         let data = self.inner.as_mut();
-        array::check_bounds(data, index)?;
+        check_bounds(data, index)?;
         let fd = data.fd_or_err()?;
         bpf_map_update_elem(fd, Some(&index), &value, flags).map_err(|(_, io_error)| {
             MapError::SyscallError {
diff --git a/aya/src/maps/array/mod.rs b/aya/src/maps/array/mod.rs
index c33f91cf..728e4851 100644
--- a/aya/src/maps/array/mod.rs
+++ b/aya/src/maps/array/mod.rs
@@ -7,14 +7,3 @@ mod program_array;
 pub use array::*;
 pub use per_cpu_array::PerCpuArray;
 pub use program_array::ProgramArray;
-
-use crate::maps::{MapData, MapError};
-
-pub(crate) fn check_bounds(map: &MapData, index: u32) -> Result<(), MapError> {
-    let max_entries = map.obj.max_entries();
-    if index >= map.obj.max_entries() {
-        Err(MapError::OutOfBounds { index, max_entries })
-    } else {
-        Ok(())
-    }
-}
diff --git a/aya/src/maps/array/per_cpu_array.rs b/aya/src/maps/array/per_cpu_array.rs
index 5589e306..d9a4f73c 100644
--- a/aya/src/maps/array/per_cpu_array.rs
+++ b/aya/src/maps/array/per_cpu_array.rs
@@ -1,11 +1,10 @@
 use std::{
     convert::{AsMut, AsRef},
     marker::PhantomData,
-    mem,
 };
 
 use crate::{
-    maps::{array, IterableMap, MapData, MapError, PerCpuValues},
+    maps::{check_bounds, check_kv_size, IterableMap, MapData, MapError, PerCpuValues},
     sys::{bpf_map_lookup_elem_per_cpu, bpf_map_update_elem_per_cpu},
     Pod,
 };
@@ -57,17 +56,8 @@ pub struct PerCpuArray<T, V: Pod> {
 impl<T: AsRef<MapData>, V: Pod> PerCpuArray<T, V> {
     pub(crate) fn new(map: T) -> Result<PerCpuArray<T, V>, MapError> {
         let data = map.as_ref();
-        let expected = mem::size_of::<u32>();
-        let size = data.obj.key_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidKeySize { size, expected });
-        }
+        check_kv_size::<u32, V>(data)?;
 
-        let expected = mem::size_of::<V>();
-        let size = data.obj.value_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidValueSize { size, expected });
-        }
         let _fd = data.fd_or_err()?;
 
         Ok(PerCpuArray {
@@ -91,7 +81,7 @@ impl<T: AsRef<MapData>, V: Pod> PerCpuArray<T, V> {
     /// if `bpf_map_lookup_elem` fails.
     pub fn get(&self, index: &u32, flags: u64) -> Result<PerCpuValues<V>, MapError> {
         let data = self.inner.as_ref();
-        array::check_bounds(data, *index)?;
+        check_bounds(data, *index)?;
         let fd = data.fd_or_err()?;
 
         let value = bpf_map_lookup_elem_per_cpu(fd, index, flags).map_err(|(_, io_error)| {
@@ -119,7 +109,7 @@ impl<T: AsMut<MapData>, V: Pod> PerCpuArray<T, V> {
     /// if `bpf_map_update_elem` fails.
     pub fn set(&mut self, index: u32, values: PerCpuValues<V>, flags: u64) -> Result<(), MapError> {
         let data = self.inner.as_mut();
-        array::check_bounds(data, index)?;
+        check_bounds(data, index)?;
         let fd = data.fd_or_err()?;
 
         bpf_map_update_elem_per_cpu(fd, &index, &values, flags).map_err(|(_, io_error)| {
diff --git a/aya/src/maps/array/program_array.rs b/aya/src/maps/array/program_array.rs
index c51cfbb7..5245d001 100644
--- a/aya/src/maps/array/program_array.rs
+++ b/aya/src/maps/array/program_array.rs
@@ -2,12 +2,11 @@
 
 use std::{
     convert::{AsMut, AsRef},
-    mem,
     os::unix::prelude::{AsRawFd, RawFd},
 };
 
 use crate::{
-    maps::{array, MapData, MapError, MapKeys},
+    maps::{check_bounds, check_kv_size, MapData, MapError, MapKeys},
     programs::ProgramFd,
     sys::{bpf_map_delete_elem, bpf_map_update_elem},
 };
@@ -55,17 +54,8 @@ pub struct ProgramArray<T> {
 impl<T: AsRef<MapData>> ProgramArray<T> {
     pub(crate) fn new(map: T) -> Result<ProgramArray<T>, MapError> {
         let data = map.as_ref();
-        let expected = mem::size_of::<u32>();
-        let size = data.obj.key_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidKeySize { size, expected });
-        }
+        check_kv_size::<u32, RawFd>(data)?;
 
-        let expected = mem::size_of::<RawFd>();
-        let size = data.obj.value_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidValueSize { size, expected });
-        }
         let _fd = data.fd_or_err()?;
 
         Ok(ProgramArray { inner: map })
@@ -85,7 +75,7 @@ impl<T: AsMut<MapData>> ProgramArray<T> {
     /// flow will jump to `program`.
     pub fn set(&mut self, index: u32, program: ProgramFd, flags: u64) -> Result<(), MapError> {
         let data = self.inner.as_mut();
-        array::check_bounds(data, index)?;
+        check_bounds(data, index)?;
         let fd = data.fd_or_err()?;
         let prog_fd = program.as_raw_fd();
 
@@ -104,7 +94,7 @@ impl<T: AsMut<MapData>> ProgramArray<T> {
     /// error.
     pub fn clear_index(&mut self, index: &u32) -> Result<(), MapError> {
         let data = self.inner.as_mut();
-        array::check_bounds(data, *index)?;
+        check_bounds(data, *index)?;
         let fd = self.inner.as_mut().fd_or_err()?;
 
         bpf_map_delete_elem(fd, index)
diff --git a/aya/src/maps/bloom_filter.rs b/aya/src/maps/bloom_filter.rs
index 5e49749a..7a540367 100644
--- a/aya/src/maps/bloom_filter.rs
+++ b/aya/src/maps/bloom_filter.rs
@@ -1,10 +1,8 @@
 //! A Bloom Filter.
 use std::{convert::AsRef, marker::PhantomData};
 
-use core::mem;
-
 use crate::{
-    maps::{MapData, MapError},
+    maps::{check_v_size, MapData, MapError},
     sys::{bpf_map_lookup_elem_ptr, bpf_map_push_elem},
     Pod,
 };
@@ -40,11 +38,7 @@ pub struct BloomFilter<T, V: Pod> {
 impl<T: AsRef<MapData>, V: Pod> BloomFilter<T, V> {
     pub(crate) fn new(map: T) -> Result<BloomFilter<T, V>, MapError> {
         let data = map.as_ref();
-        let size = mem::size_of::<V>();
-        let expected = data.obj.value_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidValueSize { size, expected });
-        };
+        check_v_size::<V>(data)?;
 
         let _ = data.fd_or_err()?;
 
diff --git a/aya/src/maps/hash_map/hash_map.rs b/aya/src/maps/hash_map/hash_map.rs
index 8a120201..474337bd 100644
--- a/aya/src/maps/hash_map/hash_map.rs
+++ b/aya/src/maps/hash_map/hash_map.rs
@@ -4,7 +4,7 @@ use std::{
 };
 
 use crate::{
-    maps::{hash_map, IterableMap, MapData, MapError, MapIter, MapKeys},
+    maps::{check_kv_size, hash_map, IterableMap, MapData, MapError, MapIter, MapKeys},
     sys::bpf_map_lookup_elem,
     Pod,
 };
@@ -41,7 +41,7 @@ pub struct HashMap<T, K, V> {
 impl<T: AsRef<MapData>, K: Pod, V: Pod> HashMap<T, K, V> {
     pub(crate) fn new(map: T) -> Result<HashMap<T, K, V>, MapError> {
         let data = map.as_ref();
-        hash_map::check_kv_size::<K, V>(data)?;
+        check_kv_size::<K, V>(data)?;
         let _ = data.fd_or_err()?;
 
         Ok(HashMap {
diff --git a/aya/src/maps/hash_map/mod.rs b/aya/src/maps/hash_map/mod.rs
index 92333700..5be375c9 100644
--- a/aya/src/maps/hash_map/mod.rs
+++ b/aya/src/maps/hash_map/mod.rs
@@ -1,6 +1,4 @@
 //! Hash map types.
-use std::mem;
-
 use crate::{
     maps::MapError,
     sys::{bpf_map_delete_elem, bpf_map_update_elem},
@@ -15,20 +13,6 @@ pub use per_cpu_hash_map::*;
 
 use super::MapData;
 
-pub(crate) fn check_kv_size<K, V>(map: &MapData) -> Result<(), MapError> {
-    let size = mem::size_of::<K>();
-    let expected = map.obj.key_size() as usize;
-    if size != expected {
-        return Err(MapError::InvalidKeySize { size, expected });
-    }
-    let size = mem::size_of::<V>();
-    let expected = map.obj.value_size() as usize;
-    if size != expected {
-        return Err(MapError::InvalidValueSize { size, expected });
-    };
-    Ok(())
-}
-
 pub(crate) fn insert<K, V>(
     map: &mut MapData,
     key: K,
diff --git a/aya/src/maps/hash_map/per_cpu_hash_map.rs b/aya/src/maps/hash_map/per_cpu_hash_map.rs
index 329a51fc..946ab63a 100644
--- a/aya/src/maps/hash_map/per_cpu_hash_map.rs
+++ b/aya/src/maps/hash_map/per_cpu_hash_map.rs
@@ -5,8 +5,9 @@ use std::{
 };
 
 use crate::{
-    generated::bpf_map_type::{BPF_MAP_TYPE_LRU_PERCPU_HASH, BPF_MAP_TYPE_PERCPU_HASH},
-    maps::{hash_map, IterableMap, MapData, MapError, MapIter, MapKeys, PerCpuValues},
+    maps::{
+        check_kv_size, hash_map, IterableMap, MapData, MapError, MapIter, MapKeys, PerCpuValues,
+    },
     sys::{bpf_map_lookup_elem_per_cpu, bpf_map_update_elem_per_cpu},
     Pod,
 };
@@ -49,15 +50,8 @@ pub struct PerCpuHashMap<T, K: Pod, V: Pod> {
 impl<T: AsRef<MapData>, K: Pod, V: Pod> PerCpuHashMap<T, K, V> {
     pub(crate) fn new(map: T) -> Result<PerCpuHashMap<T, K, V>, MapError> {
         let data = map.as_ref();
-        let map_type = data.obj.map_type();
+        check_kv_size::<K, V>(data)?;
 
-        // validate the map definition
-        if map_type != BPF_MAP_TYPE_PERCPU_HASH as u32
-            && map_type != BPF_MAP_TYPE_LRU_PERCPU_HASH as u32
-        {
-            return Err(MapError::InvalidMapType { map_type });
-        }
-        hash_map::check_kv_size::<K, V>(data)?;
         let _ = data.fd_or_err()?;
 
         Ok(PerCpuHashMap {
diff --git a/aya/src/maps/lpm_trie.rs b/aya/src/maps/lpm_trie.rs
index 4cbc6ec7..ae479529 100644
--- a/aya/src/maps/lpm_trie.rs
+++ b/aya/src/maps/lpm_trie.rs
@@ -6,7 +6,7 @@ use std::{
 };
 
 use crate::{
-    maps::{IterableMap, MapData, MapError},
+    maps::{check_kv_size, IterableMap, MapData, MapError},
     sys::{bpf_map_delete_elem, bpf_map_lookup_elem, bpf_map_update_elem},
     Pod,
 };
@@ -102,16 +102,7 @@ unsafe impl<K: Pod> Pod for Key<K> {}
 impl<T: AsRef<MapData>, K: Pod, V: Pod> LpmTrie<T, K, V> {
     pub(crate) fn new(map: T) -> Result<LpmTrie<T, K, V>, MapError> {
         let data = map.as_ref();
-        let size = mem::size_of::<Key<K>>();
-        let expected = data.obj.key_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidKeySize { size, expected });
-        }
-        let size = mem::size_of::<V>();
-        let expected = data.obj.value_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidValueSize { size, expected });
-        };
+        check_kv_size::<Key<K>, V>(data)?;
 
         let _ = data.fd_or_err()?;
 
diff --git a/aya/src/maps/mod.rs b/aya/src/maps/mod.rs
index 4abcf3af..29bf0658 100644
--- a/aya/src/maps/mod.rs
+++ b/aya/src/maps/mod.rs
@@ -78,7 +78,8 @@ pub use array::{Array, PerCpuArray, ProgramArray};
 pub use bloom_filter::BloomFilter;
 pub use hash_map::{HashMap, PerCpuHashMap};
 pub use lpm_trie::LpmTrie;
-#[cfg(any(feature = "async", doc))]
+#[cfg(feature = "async")]
+#[cfg_attr(docsrs, doc(cfg(feature = "async")))]
 pub use perf::AsyncPerfEventArray;
 pub use perf::PerfEventArray;
 pub use queue::Queue;
@@ -410,6 +411,38 @@ macro_rules! impl_try_from_map_generic_key_and_value {
 
 impl_try_from_map_generic_key_and_value!(HashMap, PerCpuHashMap, LpmTrie);
 
+pub(crate) fn check_bounds(map: &MapData, index: u32) -> Result<(), MapError> {
+    let max_entries = map.obj.max_entries();
+    if index >= max_entries {
+        Err(MapError::OutOfBounds { index, max_entries })
+    } else {
+        Ok(())
+    }
+}
+
+pub(crate) fn check_kv_size<K, V>(map: &MapData) -> Result<(), MapError> {
+    let size = mem::size_of::<K>();
+    let expected = map.obj.key_size() as usize;
+    if size != expected {
+        return Err(MapError::InvalidKeySize { size, expected });
+    }
+    let size = mem::size_of::<V>();
+    let expected = map.obj.value_size() as usize;
+    if size != expected {
+        return Err(MapError::InvalidValueSize { size, expected });
+    };
+    Ok(())
+}
+
+pub(crate) fn check_v_size<V>(map: &MapData) -> Result<(), MapError> {
+    let size = mem::size_of::<V>();
+    let expected = map.obj.value_size() as usize;
+    if size != expected {
+        return Err(MapError::InvalidValueSize { size, expected });
+    };
+    Ok(())
+}
+
 /// A generic handle to a BPF map.
 ///
 /// You should never need to use this unless you're implementing a new map type.
@@ -530,11 +563,6 @@ impl MapData {
         })
     }
 
-    /// Returns the [`bpf_map_type`] of this map
-    pub fn map_type(&self) -> Result<bpf_map_type, MapError> {
-        bpf_map_type::try_from(self.obj.map_type())
-    }
-
     pub(crate) fn fd_or_err(&self) -> Result<RawFd, MapError> {
         self.fd.ok_or(MapError::NotCreated)
     }
diff --git a/aya/src/maps/queue.rs b/aya/src/maps/queue.rs
index 10dbd96e..8484f2df 100644
--- a/aya/src/maps/queue.rs
+++ b/aya/src/maps/queue.rs
@@ -2,11 +2,10 @@
 use std::{
     convert::{AsMut, AsRef},
     marker::PhantomData,
-    mem,
 };
 
 use crate::{
-    maps::{MapData, MapError},
+    maps::{check_kv_size, MapData, MapError},
     sys::{bpf_map_lookup_and_delete_elem, bpf_map_push_elem},
     Pod,
 };
@@ -37,17 +36,8 @@ pub struct Queue<T, V: Pod> {
 impl<T: AsRef<MapData>, V: Pod> Queue<T, V> {
     pub(crate) fn new(map: T) -> Result<Queue<T, V>, MapError> {
         let data = map.as_ref();
-        let expected = 0;
-        let size = data.obj.key_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidKeySize { size, expected });
-        }
+        check_kv_size::<(), V>(data)?;
 
-        let expected = mem::size_of::<V>();
-        let size = data.obj.value_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidValueSize { size, expected });
-        }
         let _fd = data.fd_or_err()?;
 
         Ok(Queue {
diff --git a/aya/src/maps/sock/sock_hash.rs b/aya/src/maps/sock/sock_hash.rs
index c7ceb752..58ba8944 100644
--- a/aya/src/maps/sock/sock_hash.rs
+++ b/aya/src/maps/sock/sock_hash.rs
@@ -5,7 +5,9 @@ use std::{
 };
 
 use crate::{
-    maps::{hash_map, sock::SockMapFd, IterableMap, MapData, MapError, MapIter, MapKeys},
+    maps::{
+        check_kv_size, hash_map, sock::SockMapFd, IterableMap, MapData, MapError, MapIter, MapKeys,
+    },
     sys::bpf_map_lookup_elem,
     Pod,
 };
@@ -69,7 +71,7 @@ pub struct SockHash<T, K> {
 impl<T: AsRef<MapData>, K: Pod> SockHash<T, K> {
     pub(crate) fn new(map: T) -> Result<SockHash<T, K>, MapError> {
         let data = map.as_ref();
-        hash_map::check_kv_size::<K, u32>(data)?;
+        check_kv_size::<K, u32>(data)?;
         let _ = data.fd_or_err()?;
 
         Ok(SockHash {
diff --git a/aya/src/maps/sock/sock_map.rs b/aya/src/maps/sock/sock_map.rs
index b13291a9..e4b7131e 100644
--- a/aya/src/maps/sock/sock_map.rs
+++ b/aya/src/maps/sock/sock_map.rs
@@ -2,12 +2,11 @@
 
 use std::{
     convert::{AsMut, AsRef},
-    mem,
     os::unix::{io::AsRawFd, prelude::RawFd},
 };
 
 use crate::{
-    maps::{sock::SockMapFd, MapData, MapError, MapKeys},
+    maps::{check_bounds, check_kv_size, sock::SockMapFd, MapData, MapError, MapKeys},
     sys::{bpf_map_delete_elem, bpf_map_update_elem},
 };
 
@@ -48,17 +47,8 @@ pub struct SockMap<T> {
 impl<T: AsRef<MapData>> SockMap<T> {
     pub(crate) fn new(map: T) -> Result<SockMap<T>, MapError> {
         let data = map.as_ref();
-        let expected = mem::size_of::<u32>();
-        let size = data.obj.key_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidKeySize { size, expected });
-        }
+        check_kv_size::<u32, RawFd>(data)?;
 
-        let expected = mem::size_of::<RawFd>();
-        let size = data.obj.value_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidValueSize { size, expected });
-        }
         let _fd = data.fd_or_err()?;
 
         Ok(SockMap { inner: map })
@@ -80,8 +70,9 @@ impl<T: AsRef<MapData>> SockMap<T> {
 impl<T: AsMut<MapData>> SockMap<T> {
     /// Stores a socket into the map.
     pub fn set<I: AsRawFd>(&mut self, index: u32, socket: &I, flags: u64) -> Result<(), MapError> {
-        let fd = self.inner.as_mut().fd_or_err()?;
-        self.check_bounds(index)?;
+        let data = self.inner.as_mut();
+        let fd = data.fd_or_err()?;
+        check_bounds(data, index)?;
         bpf_map_update_elem(fd, Some(&index), &socket.as_raw_fd(), flags).map_err(
             |(_, io_error)| MapError::SyscallError {
                 call: "bpf_map_update_elem".to_owned(),
@@ -93,8 +84,9 @@ impl<T: AsMut<MapData>> SockMap<T> {
 
     /// Removes the socket stored at `index` from the map.
     pub fn clear_index(&mut self, index: &u32) -> Result<(), MapError> {
-        let fd = self.inner.as_mut().fd_or_err()?;
-        self.check_bounds(*index)?;
+        let data = self.inner.as_mut();
+        let fd = data.fd_or_err()?;
+        check_bounds(data, *index)?;
         bpf_map_delete_elem(fd, index)
             .map(|_| ())
             .map_err(|(_, io_error)| MapError::SyscallError {
@@ -102,13 +94,4 @@ impl<T: AsMut<MapData>> SockMap<T> {
                 io_error,
             })
     }
-
-    fn check_bounds(&mut self, index: u32) -> Result<(), MapError> {
-        let max_entries = self.inner.as_mut().obj.max_entries();
-        if index >= self.inner.as_mut().obj.max_entries() {
-            Err(MapError::OutOfBounds { index, max_entries })
-        } else {
-            Ok(())
-        }
-    }
 }
diff --git a/aya/src/maps/stack.rs b/aya/src/maps/stack.rs
index dcff4832..135c18d3 100644
--- a/aya/src/maps/stack.rs
+++ b/aya/src/maps/stack.rs
@@ -2,11 +2,10 @@
 use std::{
     convert::{AsMut, AsRef},
     marker::PhantomData,
-    mem,
 };
 
 use crate::{
-    maps::{MapData, MapError},
+    maps::{check_kv_size, MapData, MapError},
     sys::{bpf_map_lookup_and_delete_elem, bpf_map_update_elem},
     Pod,
 };
@@ -37,17 +36,8 @@ pub struct Stack<T, V: Pod> {
 impl<T: AsRef<MapData>, V: Pod> Stack<T, V> {
     pub(crate) fn new(map: T) -> Result<Stack<T, V>, MapError> {
         let data = map.as_ref();
-        let expected = 0;
-        let size = data.obj.key_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidKeySize { size, expected });
-        }
+        check_kv_size::<(), V>(data)?;
 
-        let expected = mem::size_of::<V>();
-        let size = data.obj.value_size() as usize;
-        if size != expected {
-            return Err(MapError::InvalidValueSize { size, expected });
-        }
         let _fd = data.fd_or_err()?;
 
         Ok(Stack {