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 { pub(crate) set: *mut sys::nftnl_set, pub(crate) table: Rc, _marker: ::std::marker::PhantomData, } impl Set { pub fn new(name: &CStr, id: u32, table: Rc
) -> 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
) -> 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 { 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 Debug for Set { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{:?}", self.get_str()) } } unsafe impl NlMsg for Set { unsafe fn write(&self, buf: &mut Vec, 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 Drop for Set { fn drop(&mut self) { unsafe { sys::nftnl_set_free(self.set) }; } } pub struct SetElemsIter<'a, K> { set: &'a Set, iter: *mut sys::nftnl_set_elems_iter, ret: Rc>, } impl<'a, K> SetElemsIter<'a, K> { fn new(set: &'a Set) -> 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 { 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, iter: *mut sys::nftnl_set_elems_iter, ret: Rc>, } 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 SetKey for [u8; N] { const TYPE: u32 = 5; const LEN: u32 = N as u32; fn data(&self) -> Box<[u8]> { Box::new(*self) } }