aya: add support for Stack and Queue maps

aya/src/maps/mod.rs

@@ -48,16 +48,19 @@ mod map_lock;
 pub mod array;
 pub mod hash_map;
 pub mod perf;
+pub mod queue;
 pub mod sock;
+pub mod stack;
 pub mod stack_trace;
 
 pub use array::{Array, PerCpuArray, ProgramArray};
 pub use hash_map::{HashMap, PerCpuHashMap};
 pub use map_lock::*;
 pub use perf::PerfEventArray;
+pub use queue::Queue;
 pub use sock::{SockHash, SockMap};
+pub use stack::Stack;
 pub use stack_trace::StackTraceMap;
-
 #[derive(Error, Debug)]
 pub enum MapError {
     #[error("map `{name}` not found ")]
@@ -94,6 +97,9 @@ pub enum MapError {
     #[error("key not found")]
     KeyNotFound,
 
+    #[error("element not found")]
+    ElementNotFound,
+
     #[error("the program is not loaded")]
     ProgramNotLoaded,
 

aya/src/maps/queue.rs

@@ -0,0 +1,121 @@
+use std::{
+    convert::TryFrom,
+    marker::PhantomData,
+    mem,
+    ops::{Deref, DerefMut},
+};
+
+use crate::{
+    generated::bpf_map_type::BPF_MAP_TYPE_QUEUE,
+    maps::{Map, MapError, MapRef, MapRefMut},
+    sys::{bpf_map_lookup_and_delete_elem, bpf_map_push_elem},
+    Pod,
+};
+
+/// A FIFO queue.
+pub struct Queue<T: Deref<Target = Map>, V: Pod> {
+    inner: T,
+    _v: PhantomData<V>,
+}
+
+impl<T: Deref<Target = Map>, V: Pod> Queue<T, V> {
+    fn new(map: T) -> Result<Queue<T, V>, MapError> {
+        let map_type = map.obj.def.map_type;
+        if map_type != BPF_MAP_TYPE_QUEUE as u32 {
+            return Err(MapError::InvalidMapType {
+                map_type: map_type as u32,
+            })?;
+        }
+        let expected = 0;
+        let size = map.obj.def.key_size as usize;
+        if size != expected {
+            return Err(MapError::InvalidKeySize { size, expected });
+        }
+
+        let expected = mem::size_of::<V>();
+        let size = map.obj.def.value_size as usize;
+        if size != expected {
+            return Err(MapError::InvalidValueSize { size, expected });
+        }
+        let _fd = map.fd_or_err()?;
+
+        Ok(Queue {
+            inner: map,
+            _v: PhantomData,
+        })
+    }
+
+    /// Returns the number of elements the queue can hold.
+    ///
+    /// This corresponds to the value of `bpf_map_def::max_entries` on the eBPF side.
+    pub fn capacity(&self) -> u32 {
+        self.inner.obj.def.max_entries
+    }
+}
+
+impl<T: Deref<Target = Map> + DerefMut<Target = Map>, V: Pod> Queue<T, V> {
+    /// Removes the first element and returns it.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`MapError::ElementNotFound`] if the queue is empty, [`MapError::SyscallError`]
+    /// if `bpf_map_lookup_and_delete_elem` fails.
+    pub fn pop(&mut self, flags: u64) -> Result<V, MapError> {
+        let fd = self.inner.fd_or_err()?;
+
+        let value = bpf_map_lookup_and_delete_elem::<u32, _>(fd, None, flags).map_err(
+            |(code, io_error)| MapError::SyscallError {
+                call: "bpf_map_lookup_and_delete_elem".to_owned(),
+                code,
+                io_error,
+            },
+        )?;
+        value.ok_or(MapError::ElementNotFound)
+    }
+
+    /// Appends an element at the end of the queue.
+    ///
+    /// # Errors
+    ///
+    /// [`MapError::SyscallError`] if `bpf_map_update_elem` fails.
+    ///
+    /// # Example
+    /// ```no_run
+    /// # let bpf = aya::Bpf::load(&[], None)?;
+    /// use aya::maps::Queue;
+    /// use std::convert::TryFrom;
+    ///
+    /// let mut queue = Queue::try_from(bpf.map_mut("ARRAY")?)?;
+    /// queue.push(42, 0)?;
+    /// queue.push(43, 0)?;
+    /// assert_eq!(queue.pop(0)?, 42);
+    /// # Ok::<(), aya::BpfError>(())
+    /// ```
+    pub fn push(&mut self, value: V, flags: u64) -> Result<(), MapError> {
+        let fd = self.inner.fd_or_err()?;
+        bpf_map_push_elem(fd, &value, flags).map_err(|(code, io_error)| {
+            MapError::SyscallError {
+                call: "bpf_map_push_elem".to_owned(),
+                code,
+                io_error,
+            }
+        })?;
+        Ok(())
+    }
+}
+
+impl<V: Pod> TryFrom<MapRef> for Queue<MapRef, V> {
+    type Error = MapError;
+
+    fn try_from(a: MapRef) -> Result<Queue<MapRef, V>, MapError> {
+        Queue::new(a)
+    }
+}
+
+impl<V: Pod> TryFrom<MapRefMut> for Queue<MapRefMut, V> {
+    type Error = MapError;
+
+    fn try_from(a: MapRefMut) -> Result<Queue<MapRefMut, V>, MapError> {
+        Queue::new(a)
+    }
+}

aya/src/maps/stack.rs

@@ -0,0 +1,121 @@
+use std::{
+    convert::TryFrom,
+    marker::PhantomData,
+    mem,
+    ops::{Deref, DerefMut},
+};
+
+use crate::{
+    generated::bpf_map_type::BPF_MAP_TYPE_QUEUE,
+    maps::{Map, MapError, MapRef, MapRefMut},
+    sys::{bpf_map_lookup_and_delete_elem, bpf_map_update_elem},
+    Pod,
+};
+
+/// A LIFO stack.
+pub struct Stack<T: Deref<Target = Map>, V: Pod> {
+    inner: T,
+    _v: PhantomData<V>,
+}
+
+impl<T: Deref<Target = Map>, V: Pod> Stack<T, V> {
+    fn new(map: T) -> Result<Stack<T, V>, MapError> {
+        let map_type = map.obj.def.map_type;
+        if map_type != BPF_MAP_TYPE_QUEUE as u32 {
+            return Err(MapError::InvalidMapType {
+                map_type: map_type as u32,
+            })?;
+        }
+        let expected = 0;
+        let size = map.obj.def.key_size as usize;
+        if size != expected {
+            return Err(MapError::InvalidKeySize { size, expected });
+        }
+
+        let expected = mem::size_of::<V>();
+        let size = map.obj.def.value_size as usize;
+        if size != expected {
+            return Err(MapError::InvalidValueSize { size, expected });
+        }
+        let _fd = map.fd_or_err()?;
+
+        Ok(Stack {
+            inner: map,
+            _v: PhantomData,
+        })
+    }
+
+    /// Returns the number of elements the stack can hold.
+    ///
+    /// This corresponds to the value of `bpf_map_def::max_entries` on the eBPF side.
+    pub fn capacity(&self) -> u32 {
+        self.inner.obj.def.max_entries
+    }
+}
+
+impl<T: Deref<Target = Map> + DerefMut<Target = Map>, V: Pod> Stack<T, V> {
+    /// Removes the last element and returns it.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`MapError::ElementNotFound`] if the stack is empty, [`MapError::SyscallError`]
+    /// if `bpf_map_lookup_and_delete_elem` fails.
+    pub fn pop(&mut self, flags: u64) -> Result<V, MapError> {
+        let fd = self.inner.fd_or_err()?;
+
+        let value = bpf_map_lookup_and_delete_elem::<u32, _>(fd, None, flags).map_err(
+            |(code, io_error)| MapError::SyscallError {
+                call: "bpf_map_lookup_and_delete_elem".to_owned(),
+                code,
+                io_error,
+            },
+        )?;
+        value.ok_or(MapError::ElementNotFound)
+    }
+
+    /// Pushes an element on the stack.
+    ///
+    /// # Errors
+    ///
+    /// [`MapError::SyscallError`] if `bpf_map_update_elem` fails.
+    ///
+    /// # Example
+    /// ```no_run
+    /// # let bpf = aya::Bpf::load(&[], None)?;
+    /// use aya::maps::Stack;
+    /// use std::convert::TryFrom;
+    ///
+    /// let mut stack = Stack::try_from(bpf.map_mut("ARRAY")?)?;
+    /// stack.push(42, 0)?;
+    /// stack.push(43, 0)?;
+    /// assert_eq!(stack.pop(0)?, 43);
+    /// # Ok::<(), aya::BpfError>(())
+    /// ```
+    pub fn push(&mut self, value: V, flags: u64) -> Result<(), MapError> {
+        let fd = self.inner.fd_or_err()?;
+        bpf_map_update_elem(fd, &0, &value, flags).map_err(|(code, io_error)| {
+            MapError::SyscallError {
+                call: "bpf_map_update_elem".to_owned(),
+                code,
+                io_error,
+            }
+        })?;
+        Ok(())
+    }
+}
+
+impl<V: Pod> TryFrom<MapRef> for Stack<MapRef, V> {
+    type Error = MapError;
+
+    fn try_from(a: MapRef) -> Result<Stack<MapRef, V>, MapError> {
+        Stack::new(a)
+    }
+}
+
+impl<V: Pod> TryFrom<MapRefMut> for Stack<MapRefMut, V> {
+    type Error = MapError;
+
+    fn try_from(a: MapRefMut) -> Result<Stack<MapRefMut, V>, MapError> {
+        Stack::new(a)
+    }
+}

aya/src/sys/bpf.rs

@@ -66,7 +66,7 @@ pub(crate) fn bpf_load_program(
 
 fn lookup<K: Pod, V: Pod>(
     fd: RawFd,
-    key: &K,
+    key: Option<&K>,
     flags: u64,
     cmd: bpf_cmd,
 ) -> Result<Option<V>, (c_long, io::Error)> {
@@ -75,7 +75,9 @@ fn lookup<K: Pod, V: Pod>(
 
     let u = unsafe { &mut attr.__bindgen_anon_2 };
     u.map_fd = fd as u32;
-    u.key = key as *const _ as u64;
+    if let Some(key) = key {
+        u.key = key as *const _ as u64;
+    }
     u.__bindgen_anon_1.value = &mut value as *mut _ as u64;
     u.flags = flags;
 
@@ -91,7 +93,15 @@ pub(crate) fn bpf_map_lookup_elem<K: Pod, V: Pod>(
     key: &K,
     flags: u64,
 ) -> Result<Option<V>, (c_long, io::Error)> {
-    lookup(fd, key, flags, bpf_cmd::BPF_MAP_LOOKUP_ELEM)
+    lookup(fd, Some(key), flags, bpf_cmd::BPF_MAP_LOOKUP_ELEM)
+}
+
+pub(crate) fn bpf_map_lookup_and_delete_elem<K: Pod, V: Pod>(
+    fd: RawFd,
+    key: Option<&K>,
+    flags: u64,
+) -> Result<Option<V>, (c_long, io::Error)> {
+    lookup(fd, key, flags, bpf_cmd::BPF_MAP_LOOKUP_AND_DELETE_ELEM)
 }
 
 pub(crate) fn bpf_map_lookup_elem_per_cpu<K: Pod, V: Pod>(
@@ -140,6 +150,17 @@ pub(crate) fn bpf_map_update_elem<K, V>(fd: RawFd, key: &K, value: &V, flags: u6
     sys_bpf(bpf_cmd::BPF_MAP_UPDATE_ELEM, &attr)
 }
 
+pub(crate) fn bpf_map_push_elem<V>(fd: RawFd, value: &V, flags: u64) -> SysResult {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+
+    let u = unsafe { &mut attr.__bindgen_anon_2 };
+    u.map_fd = fd as u32;
+    u.__bindgen_anon_1.value = value as *const _ as u64;
+    u.flags = flags;
+
+    sys_bpf(bpf_cmd::BPF_MAP_UPDATE_ELEM, &attr)
+}
+
 pub(crate) fn bpf_map_update_elem_ptr<K, V>(
     fd: RawFd,
     key: *const K,