maps: add ProgramArray

ProgramArray is a wrapper around BPF_MAP_TYPE_PROG_ARRAY maps. Can be
used to setup jump tables for bpf_tail_call().

src/maps/mod.rs

@@ -8,10 +8,12 @@ use crate::{
 };
 
 mod hash_map;
-pub use hash_map::*;
-
 mod perf_map;
+mod program_array;
+
+pub use hash_map::*;
 pub use perf_map::*;
+pub use program_array::*;
 
 #[derive(Error, Debug)]
 pub enum MapError {
@@ -40,6 +42,12 @@ pub enum MapError {
     #[error("invalid value size {size}, expected {expected}")]
     InvalidValueSize { size: usize, expected: usize },
 
+    #[error("the index is {index} but `max_entries` is {max_entries}")]
+    OutOfBounds { index: u32, max_entries: u32 },
+
+    #[error("the program is not loaded")]
+    ProgramNotLoaded,
+
     #[error("the BPF_MAP_UPDATE_ELEM syscall failed with code {code} io_error {io_error}")]
     UpdateElementFailed { code: i64, io_error: io::Error },
 

src/maps/program_array.rs

@@ -0,0 +1,128 @@
+use std::{
+    cell::{Ref, RefMut},
+    convert::TryFrom,
+    mem,
+    ops::{Deref, DerefMut},
+    os::unix::prelude::RawFd,
+};
+
+use crate::{
+    generated::bpf_map_type::BPF_MAP_TYPE_PROG_ARRAY,
+    maps::{IterableMap, Map, MapError, MapIter, MapKeys},
+    programs::ProgramFd,
+    syscalls::{
+        bpf_map_delete_elem, bpf_map_lookup_and_delete_elem, bpf_map_lookup_elem,
+        bpf_map_update_elem,
+    },
+};
+
+pub struct ProgramArray<T: Deref<Target = Map>> {
+    inner: T,
+}
+
+impl<T: Deref<Target = Map>> ProgramArray<T> {
+    pub fn new(map: T) -> Result<ProgramArray<T>, MapError> {
+        let inner = map.deref();
+        let map_type = inner.obj.def.map_type;
+        if map_type != BPF_MAP_TYPE_PROG_ARRAY {
+            return Err(MapError::InvalidMapType {
+                map_type: map_type as u32,
+            })?;
+        }
+        let expected = mem::size_of::<RawFd>();
+        let size = inner.obj.def.key_size as usize;
+        if size != expected {
+            return Err(MapError::InvalidKeySize { size, expected });
+        }
+
+        let expected = mem::size_of::<RawFd>();
+        let size = inner.obj.def.value_size as usize;
+        if size != expected {
+            return Err(MapError::InvalidValueSize { size, expected });
+        }
+
+        Ok(ProgramArray { inner: map })
+    }
+
+    pub unsafe fn get(&self, key: &u32, flags: u64) -> Result<Option<RawFd>, MapError> {
+        let fd = self.inner.deref().fd_or_err()?;
+        let fd = bpf_map_lookup_elem(fd, key, flags)
+            .map_err(|(code, io_error)| MapError::LookupElementFailed { code, io_error })?;
+        Ok(fd)
+    }
+
+    pub unsafe fn iter<'coll>(&'coll self) -> MapIter<'coll, u32, RawFd> {
+        MapIter::new(self)
+    }
+
+    pub unsafe fn keys<'coll>(&'coll self) -> MapKeys<'coll, u32, RawFd> {
+        MapKeys::new(self)
+    }
+
+    fn check_bounds(&self, index: u32) -> Result<(), MapError> {
+        let max_entries = self.inner.deref().obj.def.max_entries;
+        if index >= self.inner.deref().obj.def.max_entries {
+            Err(MapError::OutOfBounds { index, max_entries })
+        } else {
+            Ok(())
+        }
+    }
+}
+
+impl<T: Deref<Target = Map> + DerefMut<Target = Map>> ProgramArray<T> {
+    pub fn insert(
+        &mut self,
+        index: u32,
+        program: &dyn ProgramFd,
+        flags: u64,
+    ) -> Result<(), MapError> {
+        let fd = self.inner.deref().fd_or_err()?;
+        self.check_bounds(index)?;
+        let prog_fd = program.fd().ok_or(MapError::ProgramNotLoaded)?;
+
+        bpf_map_update_elem(fd, &index, &prog_fd, flags)
+            .map_err(|(code, io_error)| MapError::UpdateElementFailed { code, io_error })?;
+        Ok(())
+    }
+
+    pub unsafe fn pop(&mut self, index: &u32) -> Result<Option<RawFd>, MapError> {
+        let fd = self.inner.deref().fd_or_err()?;
+        self.check_bounds(*index)?;
+        bpf_map_lookup_and_delete_elem(fd, index)
+            .map_err(|(code, io_error)| MapError::LookupAndDeleteElementFailed { code, io_error })
+    }
+
+    pub fn remove(&mut self, index: &u32) -> Result<(), MapError> {
+        let fd = self.inner.deref().fd_or_err()?;
+        self.check_bounds(*index)?;
+        bpf_map_delete_elem(fd, index)
+            .map(|_| ())
+            .map_err(|(code, io_error)| MapError::DeleteElementFailed { code, io_error })
+    }
+}
+
+impl<T: Deref<Target = Map>> IterableMap<u32, RawFd> for ProgramArray<T> {
+    fn fd(&self) -> Result<RawFd, MapError> {
+        self.inner.deref().fd_or_err()
+    }
+
+    unsafe fn get(&self, index: &u32) -> Result<Option<RawFd>, MapError> {
+        self.get(index, 0)
+    }
+}
+
+impl<'a> TryFrom<Ref<'a, Map>> for ProgramArray<Ref<'a, Map>> {
+    type Error = MapError;
+
+    fn try_from(inner: Ref<'a, Map>) -> Result<ProgramArray<Ref<'a, Map>>, MapError> {
+        ProgramArray::new(inner)
+    }
+}
+
+impl<'a> TryFrom<RefMut<'a, Map>> for ProgramArray<RefMut<'a, Map>> {
+    type Error = MapError;
+
+    fn try_from(inner: RefMut<'a, Map>) -> Result<ProgramArray<RefMut<'a, Map>>, MapError> {
+        ProgramArray::new(inner)
+    }
+}

src/programs/mod.rs

@@ -77,12 +77,8 @@ pub enum ProgramError {
     Other { message: String },
 }
 
-#[derive(Debug)]
-pub(crate) struct ProgramData {
-    pub(crate) name: String,
-    pub(crate) obj: obj::Program,
-    pub(crate) fd: Option<RawFd>,
-    pub(crate) links: Vec<Rc<RefCell<dyn Link>>>,
+pub trait ProgramFd {
+    fn fd(&self) -> Option<RawFd>;
 }
 
 #[derive(Debug)]
@@ -110,6 +106,16 @@ impl Program {
         }
     }
 
+    pub(crate) fn data(&self) -> &ProgramData {
+        match self {
+            Program::KProbe(p) => &p.data,
+            Program::UProbe(p) => &p.data,
+            Program::TracePoint(p) => &p.data,
+            Program::SocketFilter(p) => &p.data,
+            Program::Xdp(p) => &p.data,
+        }
+    }
+
     fn data_mut(&mut self) -> &mut ProgramData {
         match self {
             Program::KProbe(p) => &mut p.data,
@@ -121,6 +127,34 @@ impl Program {
     }
 }
 
+impl ProgramFd for Program {
+    fn fd(&self) -> Option<RawFd> {
+        self.data().fd
+    }
+}
+
+macro_rules! impl_program_fd {
+    ($($struct_name:ident),+ $(,)?) => {
+        $(
+            impl ProgramFd for $struct_name {
+                fn fd(&self) -> Option<RawFd> {
+                    self.data.fd
+                }
+            }
+        )+
+    }
+}
+
+impl_program_fd!(KProbe, UProbe, TracePoint, SocketFilter, Xdp);
+
+#[derive(Debug)]
+pub(crate) struct ProgramData {
+    pub(crate) name: String,
+    pub(crate) obj: obj::Program,
+    pub(crate) fd: Option<RawFd>,
+    pub(crate) links: Vec<Rc<RefCell<dyn Link>>>,
+}
+
 impl ProgramData {
     fn fd_or_err(&self) -> Result<RawFd, ProgramError> {
         self.fd.ok_or(ProgramError::NotLoaded {
@@ -129,10 +163,6 @@ impl ProgramData {
     }
 }
 
-impl Drop for ProgramData {
-    fn drop(&mut self) {}
-}
-
 const MAX_LOG_BUF_SIZE: usize = (std::u32::MAX >> 8) as usize;
 
 pub struct VerifierLog {