1 files changed, 273 insertions, 0 deletions
diff --git a/src/set.rs b/src/set.rs
new file mode 100644
index 0000000..d6b9514
--- /dev/null
+++ b/src/set.rs
@@ -0,0 +1,273 @@
+use crate::{table::Table, MsgType, ProtoFamily};
+use crate::sys::{self, libc};
+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, $family:expr) => {
+        $crate::set::Set::new($name, $id, $table, $family)
+    };
+    ($name:expr, $id:expr, $table:expr, $family:expr; [ ]) => {
+        nft_set!($name, $id, $table, $family)
+    };
+    ($name:expr, $id:expr, $table:expr, $family:expr; [ $($value:expr,)* ]) => {{
+        let mut set = nft_set!($name, $id, $table, $family).expect("Set allocation failed");
+        $(
+            set.add($value).expect(stringify!(Unable to add $value to set $name));
+        )*
+        set
+    }};
+}
+
+pub struct Set<'a, K> {
+    pub(crate) set: *mut sys::nftnl_set,
+    pub(crate) table: &'a Table,
+    pub(crate) family: ProtoFamily,
+    _marker: ::std::marker::PhantomData<K>,
+}
+
+impl<'a, K> Set<'a, K> {
+    pub fn new(name: &CStr, id: u32, table: &'a Table, family: ProtoFamily) -> 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, 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,
+                family,
+                _marker: ::std::marker::PhantomData,
+            }
+        }
+    }
+
+    pub unsafe fn from_raw(set: *mut sys::nftnl_set, table: &'a Table, family: ProtoFamily) -> Self
+    where
+        K: SetKey,
+    {
+        Set {
+            set,
+            table,
+            family,
+            _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(&'a self) -> SetElemsIter<'a, 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
+    }
+
+    pub fn get_family(&self) -> ProtoFamily {
+        self.family
+    }
+
+    /// 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(),
+                (descr_buf.len() - 1) as u64,
+                self.set,
+                sys::NFTNL_OUTPUT_DEFAULT,
+                0,
+            );
+            CStr::from_ptr(descr_buf.as_ptr()).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<'a, K> Debug for Set<'a, K> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{:?}", self.get_str())
+    }
+}
+
+unsafe impl<'a, K> crate::NlMsg for Set<'a, K> {
+    unsafe fn write(&self, buf: *mut c_void, 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<'a, K> Drop for Set<'a, K> {
+    fn drop(&mut self) {
+        unsafe { sys::nftnl_set_free(self.set) };
+    }
+}
+
+pub struct SetElemsIter<'a, K> {
+    set: &'a Set<'a, K>,
+    iter: *mut sys::nftnl_set_elems_iter,
+    ret: Rc<Cell<i32>>,
+}
+
+impl<'a, K> SetElemsIter<'a, K> {
+    fn new(set: &'a Set<'a, 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: 'a> 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<'a, K>,
+    iter: *mut sys::nftnl_set_elems_iter,
+    ret: Rc<Cell<i32>>,
+}
+
+unsafe impl<'a, K> crate::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.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()
+    }
+}