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}