1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! A LIFO stack.
use std::{
    convert::TryFrom,
    marker::PhantomData,
    mem,
    ops::{Deref, DerefMut},
};

use crate::{
    generated::bpf_map_type::BPF_MAP_TYPE_STACK,
    maps::{Map, MapError, MapRef, MapRefMut},
    sys::{bpf_map_lookup_and_delete_elem, bpf_map_update_elem},
    Pod,
};

/// A LIFO stack.
///
/// # Examples
/// ```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("STACK")?)?;
/// stack.push(42, 0)?;
/// stack.push(43, 0)?;
/// assert_eq!(stack.pop(0)?, 43);
/// # Ok::<(), aya::BpfError>(())
/// ```
#[doc(alias = "BPF_MAP_TYPE_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_STACK 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.
    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)
    }
}