Skip to main content

nftnl/
set.rs

1use crate::{MsgType, ProtoFamily, table::Table};
2use nftnl_sys::{self as sys, libc};
3use std::{
4    cell::Cell,
5    ffi::{CStr, c_void},
6    net::{Ipv4Addr, Ipv6Addr},
7    os::raw::c_char,
8    ptr,
9    rc::Rc,
10};
11
12#[macro_export]
13macro_rules! nft_set {
14    ($name:expr, $id:expr, $table:expr, $family:expr) => {
15        $crate::set::Set::new($name, $id, $table, $family)
16    };
17    ($name:expr, $id:expr, $table:expr, $family:expr; [ ]) => {
18        nft_set!($name, $id, $table, $family)
19    };
20    ($name:expr, $id:expr, $table:expr, $family:expr; [ $($value:expr,)* ]) => {{
21        let mut set = nft_set!($name, $id, $table, $family);
22        $(
23            set.add($value);
24        )*
25        set
26    }};
27}
28
29pub struct Set<'a, K> {
30    set: ptr::NonNull<sys::nftnl_set>,
31    table: &'a Table,
32    family: ProtoFamily,
33    _marker: ::std::marker::PhantomData<K>,
34}
35
36impl<'a, K> Set<'a, K> {
37    pub fn new(name: &CStr, id: u32, table: &'a Table, family: ProtoFamily) -> Self
38    where
39        K: SetKey,
40    {
41        let set = try_alloc!(unsafe { sys::nftnl_set_alloc() });
42
43        unsafe {
44            let set = set.as_ptr();
45            sys::nftnl_set_set_u32(set, sys::NFTNL_SET_FAMILY as u16, family as u32);
46            sys::nftnl_set_set_str(set, sys::NFTNL_SET_TABLE as u16, table.get_name().as_ptr());
47            sys::nftnl_set_set_str(set, sys::NFTNL_SET_NAME as u16, name.as_ptr());
48            sys::nftnl_set_set_u32(set, sys::NFTNL_SET_ID as u16, id);
49
50            sys::nftnl_set_set_u32(
51                set,
52                sys::NFTNL_SET_FLAGS as u16,
53                (libc::NFT_SET_ANONYMOUS | libc::NFT_SET_CONSTANT) as u32,
54            );
55            sys::nftnl_set_set_u32(set, sys::NFTNL_SET_KEY_TYPE as u16, K::TYPE);
56            sys::nftnl_set_set_u32(set, sys::NFTNL_SET_KEY_LEN as u16, K::LEN);
57        }
58
59        Set {
60            set,
61            table,
62            family,
63            _marker: ::std::marker::PhantomData,
64        }
65    }
66
67    pub fn add(&mut self, key: &K)
68    where
69        K: SetKey,
70    {
71        unsafe {
72            let elem = try_alloc!(sys::nftnl_set_elem_alloc());
73
74            let data = key.data();
75            let data_len = data.len() as u32;
76            trace!("Adding key {data:?} with len {data_len}");
77            sys::nftnl_set_elem_set(
78                elem.as_ptr(),
79                sys::NFTNL_SET_ELEM_KEY as u16,
80                data.as_ref() as *const _ as *const c_void,
81                data_len,
82            );
83            sys::nftnl_set_elem_add(self.set.as_ptr(), elem.as_ptr());
84        }
85    }
86
87    pub fn elems_iter(&'a self) -> SetElemsIter<'a, K> {
88        SetElemsIter::new(self)
89    }
90
91    pub fn as_ptr(&self) -> ptr::NonNull<sys::nftnl_set> {
92        self.set
93    }
94
95    pub fn get_family(&self) -> ProtoFamily {
96        self.family
97    }
98
99    pub fn get_name(&self) -> &CStr {
100        unsafe {
101            let ptr = sys::nftnl_set_get_str(self.set.as_ptr(), sys::NFTNL_SET_NAME as u16);
102            CStr::from_ptr(ptr)
103        }
104    }
105
106    pub fn get_id(&self) -> u32 {
107        unsafe { sys::nftnl_set_get_u32(self.set.as_ptr(), sys::NFTNL_SET_ID as u16) }
108    }
109}
110
111unsafe impl<K> crate::NlMsg for Set<'_, K> {
112    unsafe fn write(&self, buf: *mut c_void, seq: u32, msg_type: MsgType) {
113        let type_ = match msg_type {
114            MsgType::Add => libc::NFT_MSG_NEWSET,
115            MsgType::Del => libc::NFT_MSG_DELSET,
116        };
117        let header = unsafe {
118            sys::nftnl_nlmsg_build_hdr(
119                buf.cast::<c_char>(),
120                type_ as u16,
121                self.table.get_family() as u16,
122                (libc::NLM_F_APPEND | libc::NLM_F_CREATE | libc::NLM_F_ACK) as u16,
123                seq,
124            )
125        };
126        unsafe { sys::nftnl_set_nlmsg_build_payload(header, self.set.as_ptr()) };
127    }
128}
129
130impl<K> Drop for Set<'_, K> {
131    fn drop(&mut self) {
132        unsafe { sys::nftnl_set_free(self.set.as_ptr()) };
133    }
134}
135
136pub struct SetElemsIter<'a, K> {
137    set: &'a Set<'a, K>,
138    iter: ptr::NonNull<sys::nftnl_set_elems_iter>,
139    ret: Rc<Cell<i32>>,
140}
141
142impl<'a, K> SetElemsIter<'a, K> {
143    fn new(set: &'a Set<'a, K>) -> Self {
144        let iter = try_alloc!(unsafe { sys::nftnl_set_elems_iter_create(set.set.as_ptr()) });
145        SetElemsIter {
146            set,
147            iter,
148            ret: Rc::new(Cell::new(1)),
149        }
150    }
151}
152
153impl<'a, K: 'a> Iterator for SetElemsIter<'a, K> {
154    type Item = SetElemsMsg<'a, K>;
155
156    fn next(&mut self) -> Option<Self::Item> {
157        if self.ret.get() <= 0
158            || unsafe { sys::nftnl_set_elems_iter_cur(self.iter.as_ptr()).is_null() }
159        {
160            trace!("SetElemsIter iterator ending");
161            None
162        } else {
163            trace!("SetElemsIter returning new SetElemsMsg");
164            Some(SetElemsMsg {
165                set: self.set,
166                iter: self.iter.as_ptr(),
167                ret: self.ret.clone(),
168            })
169        }
170    }
171}
172
173impl<K> Drop for SetElemsIter<'_, K> {
174    fn drop(&mut self) {
175        unsafe { sys::nftnl_set_elems_iter_destroy(self.iter.as_ptr()) };
176    }
177}
178
179pub struct SetElemsMsg<'a, K> {
180    set: &'a Set<'a, K>,
181    iter: *mut sys::nftnl_set_elems_iter,
182    ret: Rc<Cell<i32>>,
183}
184
185unsafe impl<K> crate::NlMsg for SetElemsMsg<'_, K> {
186    unsafe fn write(&self, buf: *mut c_void, seq: u32, msg_type: MsgType) {
187        trace!("Writing SetElemsMsg to NlMsg");
188        let (type_, flags) = match msg_type {
189            MsgType::Add => (
190                libc::NFT_MSG_NEWSETELEM,
191                libc::NLM_F_CREATE | libc::NLM_F_EXCL | libc::NLM_F_ACK,
192            ),
193            MsgType::Del => (libc::NFT_MSG_DELSETELEM, libc::NLM_F_ACK),
194        };
195        let header = unsafe {
196            sys::nftnl_nlmsg_build_hdr(
197                buf.cast::<c_char>(),
198                type_ as u16,
199                self.set.get_family() as u16,
200                flags as u16,
201                seq,
202            )
203        };
204        self.ret
205            .set(unsafe { sys::nftnl_set_elems_nlmsg_build_payload_iter(header, self.iter) });
206    }
207}
208
209pub trait SetKey {
210    const TYPE: u32;
211    const LEN: u32;
212
213    fn data(&self) -> Box<[u8]>;
214}
215
216impl SetKey for Ipv4Addr {
217    const TYPE: u32 = 7;
218    const LEN: u32 = 4;
219
220    fn data(&self) -> Box<[u8]> {
221        self.octets().to_vec().into_boxed_slice()
222    }
223}
224
225impl SetKey for Ipv6Addr {
226    const TYPE: u32 = 8;
227    const LEN: u32 = 16;
228
229    fn data(&self) -> Box<[u8]> {
230        self.octets().to_vec().into_boxed_slice()
231    }
232}