path: root/src/set.rs

use crate::nlmsg::NlMsg;
use crate::sys::{self, libc};
use crate::{table::Table, MsgType};
use std::{
    cell::Cell,
    ffi::{c_void, CStr, CString},
    fmt::Debug,
    net::{Ipv4Addr, Ipv6Addr},
    os::raw::c_char,
    rc::Rc,
};

#[macro_export]
macro_rules! nft_set {
    ($name:expr, $id:expr, $table:expr) => {
        $crate::set::Set::new(Some($name), $id, $table, $family)
    };
    ($name:expr, $id:expr, $table:expr; [ ]) => {
        nft_set!(Some($name), $id, $table)
    };
    ($name:expr, $id:expr, $table:expr; [ $($value:expr,)* ]) => {{
        let mut set = nft_set!(Some($name), $id, $table).expect("Set allocation failed");
        $(
            set.add($value).expect(stringify!(Unable to add $value to set $name));
        )*
        set
    }};
}

pub struct Set<K> {
    pub(crate) set: *mut sys::nftnl_set,
    pub(crate) table: Rc<Table>,
    _marker: ::std::marker::PhantomData<K>,
}

impl<K> Set<K> {
    pub fn new(name: &CStr, id: u32, table: Rc<Table>) -> Self
    where
        K: SetKey,
    {
        unsafe {
            let set = try_alloc!(sys::nftnl_set_alloc());

            sys::nftnl_set_set_u32(set, sys::NFTNL_SET_FAMILY as u16, table.get_family() as u32);
            sys::nftnl_set_set_str(set, sys::NFTNL_SET_TABLE as u16, table.get_name().as_ptr());
            sys::nftnl_set_set_str(set, sys::NFTNL_SET_NAME as u16, name.as_ptr());
            sys::nftnl_set_set_u32(set, sys::NFTNL_SET_ID as u16, id);

            sys::nftnl_set_set_u32(
                set,
                sys::NFTNL_SET_FLAGS as u16,
                (libc::NFT_SET_ANONYMOUS | libc::NFT_SET_CONSTANT) as u32,
            );
            sys::nftnl_set_set_u32(set, sys::NFTNL_SET_KEY_TYPE as u16, K::TYPE);
            sys::nftnl_set_set_u32(set, sys::NFTNL_SET_KEY_LEN as u16, K::LEN);

            Set {
                set,
                table,
                _marker: ::std::marker::PhantomData,
            }
        }
    }

    pub unsafe fn from_raw(set: *mut sys::nftnl_set, table: Rc<Table>) -> Self
    where
        K: SetKey,
    {
        Set {
            set,
            table,
            _marker: ::std::marker::PhantomData,
        }
    }

    pub fn add(&mut self, key: &K)
    where
        K: SetKey,
    {
        unsafe {
            let elem = try_alloc!(sys::nftnl_set_elem_alloc());

            let data = key.data();
            let data_len = data.len() as u32;
            trace!("Adding key {:?} with len {}", data, data_len);
            sys::nftnl_set_elem_set(
                elem,
                sys::NFTNL_SET_ELEM_KEY as u16,
                data.as_ref() as *const _ as *const c_void,
                data_len,
            );
            sys::nftnl_set_elem_add(self.set, elem);
        }
    }

    pub fn elems_iter(&self) -> SetElemsIter<K> {
        SetElemsIter::new(self)
    }

    #[cfg(feature = "unsafe-raw-handles")]
    /// Returns the raw handle.
    pub fn as_ptr(&self) -> *const sys::nftnl_set {
        self.set as *const sys::nftnl_set
    }

    #[cfg(feature = "unsafe-raw-handles")]
    /// Returns a mutable version of the raw handle.
    pub fn as_mut_ptr(&self) -> *mut sys::nftnl_set {
        self.set
    }

    /// Returns a textual description of the set.
    pub fn get_str(&self) -> CString {
        let mut descr_buf = vec![0i8; 4096];
        unsafe {
            sys::nftnl_set_snprintf(
                descr_buf.as_mut_ptr() as *mut c_char,
                (descr_buf.len() - 1) as u64,
                self.set,
                sys::NFTNL_OUTPUT_DEFAULT,
                0,
            );
            CStr::from_ptr(descr_buf.as_ptr() as *mut c_char).to_owned()
        }
    }

    pub fn get_name(&self) -> Option<&CStr> {
        unsafe {
            let ptr = sys::nftnl_set_get_str(self.set, sys::NFTNL_SET_NAME as u16);
            if !ptr.is_null() {
                Some(CStr::from_ptr(ptr))
            } else {
                None
            }
        }
    }

    pub fn get_id(&self) -> u32 {
        unsafe { sys::nftnl_set_get_u32(self.set, sys::NFTNL_SET_ID as u16) }
    }
}

impl<K> Debug for Set<K> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self.get_str())
    }
}

unsafe impl<K> NlMsg for Set<K> {
    unsafe fn write(&self, buf: &mut Vec<u8>, seq: u32, msg_type: MsgType) {
        let type_ = match msg_type {
            MsgType::Add => libc::NFT_MSG_NEWSET,
            MsgType::Del => libc::NFT_MSG_DELSET,
        };
        /*
        let header = sys::nftnl_nlmsg_build_hdr(
            buf as *mut c_char,
            type_ as u16,
            self.table.get_family() as u16,
            (libc::NLM_F_APPEND | libc::NLM_F_CREATE | libc::NLM_F_ACK) as u16,
            seq,
        );
        sys::nftnl_set_nlmsg_build_payload(header, self.set);
        */
    }
}

impl<K> Drop for Set<K> {
    fn drop(&mut self) {
        unsafe { sys::nftnl_set_free(self.set) };
    }
}

pub struct SetElemsIter<'a, K> {
    set: &'a Set<K>,
    iter: *mut sys::nftnl_set_elems_iter,
    ret: Rc<Cell<i32>>,
}

impl<'a, K> SetElemsIter<'a, K> {
    fn new(set: &'a Set<K>) -> Self {
        let iter = try_alloc!(unsafe {
            sys::nftnl_set_elems_iter_create(set.set as *const sys::nftnl_set)
        });
        SetElemsIter {
            set,
            iter,
            ret: Rc::new(Cell::new(1)),
        }
    }
}

impl<'a, K> Iterator for SetElemsIter<'a, K> {
    type Item = SetElemsMsg<'a, K>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.ret.get() <= 0 || unsafe { sys::nftnl_set_elems_iter_cur(self.iter).is_null() } {
            trace!("SetElemsIter iterator ending");
            None
        } else {
            trace!("SetElemsIter returning new SetElemsMsg");
            Some(SetElemsMsg {
                set: self.set,
                iter: self.iter,
                ret: self.ret.clone(),
            })
        }
    }
}

impl<'a, K> Drop for SetElemsIter<'a, K> {
    fn drop(&mut self) {
        unsafe { sys::nftnl_set_elems_iter_destroy(self.iter) };
    }
}

pub struct SetElemsMsg<'a, K> {
    set: &'a Set<K>,
    iter: *mut sys::nftnl_set_elems_iter,
    ret: Rc<Cell<i32>>,
}

unsafe impl<'a, K> NlMsg for SetElemsMsg<'a, K> {
    unsafe fn write(&self, buf: *mut c_void, seq: u32, msg_type: MsgType) {
        trace!("Writing SetElemsMsg to NlMsg");
        let (type_, flags) = match msg_type {
            MsgType::Add => (
                libc::NFT_MSG_NEWSETELEM,
                libc::NLM_F_CREATE | libc::NLM_F_EXCL | libc::NLM_F_ACK,
            ),
            MsgType::Del => (libc::NFT_MSG_DELSETELEM, libc::NLM_F_ACK),
        };
        let header = sys::nftnl_nlmsg_build_hdr(
            buf as *mut c_char,
            type_ as u16,
            self.set.table.get_family() as u16,
            flags as u16,
            seq,
        );
        self.ret.set(sys::nftnl_set_elems_nlmsg_build_payload_iter(
            header, self.iter,
        ));
    }
}

pub trait SetKey {
    const TYPE: u32;
    const LEN: u32;

    fn data(&self) -> Box<[u8]>;
}

impl SetKey for Ipv4Addr {
    const TYPE: u32 = 7;
    const LEN: u32 = 4;

    fn data(&self) -> Box<[u8]> {
        self.octets().to_vec().into_boxed_slice()
    }
}

impl SetKey for Ipv6Addr {
    const TYPE: u32 = 8;
    const LEN: u32 = 16;

    fn data(&self) -> Box<[u8]> {
        self.octets().to_vec().into_boxed_slice()
    }
}

impl<const N: usize> SetKey for [u8; N] {
    const TYPE: u32 = 5;
    const LEN: u32 = N as u32;

    fn data(&self) -> Box<[u8]> {
        Box::new(*self)
    }
}