Compare commits

..

No commits in common. "e9a6f296dfef8b4382029fdb409038e77891bb19" and "14346134b555e9b2c315d23d67535779a5fa85ef" have entirely different histories.

8 changed files with 552 additions and 797 deletions

2
.gitignore vendored
View file

@ -1,4 +1,2 @@
/target
/result
/unbound-mod-test-config
/unbound-mod-test-data

View file

@ -17,7 +17,7 @@ iptrie = "0.8.5"
libc = "0.2.155"
mnl = { version = "0.2.2", features = ["mnl-1-0-4"] }
nftnl = { version = "0.6.2", features = ["nftnl-1-1-2"] }
nix = { version = "0.29.0", features = ["poll", "user"] }
nix = { version = "0.29.0", features = ["poll"] }
radix_trie = "0.2.1"
serde = { version = "1.0.205", features = ["derive"] }
serde_json = "1.0.122"

View file

@ -1,6 +1,5 @@
use std::panic::{RefUnwindSafe, UnwindSafe};
use crate::unbound::ModuleExtState;
use crate::UnboundMod;
macro_rules! impl_tuple {
@ -30,15 +29,9 @@ macro_rules! impl_tuple {
qstate: &mut crate::unbound::ModuleQstate<Self::QstateData>,
event: crate::unbound::ModuleEvent,
entry: &mut crate::unbound::OutboundEntryMut,
) -> Option<ModuleExtState> {
#[allow(unused_mut)]
let mut ret = self.0.operate(qstate, event, entry);
$(if let Some(state) = self.$i.operate(qstate, event, entry) {
if !matches!(ret, Some(ret) if ret.importance() >= state.importance()) {
ret = Some(state);
}
})*
ret
) {
self.0.operate(qstate, event, entry);
$(self.$i.operate(qstate, event, entry);)*
}
fn get_mem(&self, env: &mut crate::unbound::ModuleEnv<Self::EnvData>) -> usize {
self.0.get_mem(env) $(* self.$i.get_mem(env))*

View file

@ -2,7 +2,6 @@ use std::{collections::HashMap, hash::Hash};
use smallvec::{smallvec, SmallVec};
#[derive(Debug)]
pub enum PrefixSet<T> {
Map(HashMap<T, PrefixSet<T>>),
Leaf,
@ -64,7 +63,7 @@ impl<T: Hash + Eq> PrefixSet<T> {
}
struct Iter<'a, T>(
SmallVec<[std::collections::hash_map::Iter<'a, T, PrefixSet<T>>; 9]>,
SmallVec<[std::collections::hash_map::Iter<'a, T, PrefixSet<T>>; 8]>,
SmallVec<[&'a T; 8]>,
);
@ -74,7 +73,9 @@ impl<'a, T> Iterator for Iter<'a, T> {
while let Some(it) = self.0.last_mut() {
let Some((k, v)) = it.next() else {
self.0.pop();
self.1.pop()?;
if self.1.pop().is_none() {
return None;
}
continue;
};
self.1.push(k);

File diff suppressed because it is too large Load diff

View file

@ -1,18 +1,11 @@
#![allow(clippy::type_complexity)]
use std::panic::{RefUnwindSafe, UnwindSafe};
use unbound::ModuleExtState;
#[allow(
dead_code,
improper_ctypes,
non_camel_case_types,
non_snake_case,
non_upper_case_globals,
unused_imports,
clippy::useless_transmute,
clippy::type_complexity,
clippy::too_many_arguments,
clippy::upper_case_acronyms
unused_imports
)]
mod bindings;
mod combine;
@ -40,8 +33,7 @@ pub trait UnboundMod: Send + Sync + Sized + RefUnwindSafe + UnwindSafe {
_qstate: &mut unbound::ModuleQstate<Self::QstateData>,
_event: unbound::ModuleEvent,
_entry: &mut unbound::OutboundEntryMut,
) -> Option<ModuleExtState> {
Some(ModuleExtState::Finished)
) {
}
fn inform_super(
&self,
@ -109,14 +101,11 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
entry: *mut bindings::outbound_entry,
) {
std::panic::catch_unwind(|| {
let mut qstate = unbound::ModuleQstate(qstate, id, Default::default());
if let Some(ext_state) = self.operate(
&mut qstate,
self.operate(
&mut unbound::ModuleQstate(qstate, id, Default::default()),
event.into(),
&mut unbound::OutboundEntryMut(entry, Default::default()),
) {
qstate.set_ext_state(ext_state);
}
)
})
.unwrap_or(());
}

View file

@ -1,23 +1,18 @@
use std::{
cell::Cell,
ffi::CStr,
fmt::Display,
io::{self, Write},
io,
net::{Ipv4Addr, Ipv6Addr},
os::{
fd::BorrowedFd,
raw::{c_char, c_void},
},
rc::Rc,
sync::mpsc,
};
use crate::example::{Helper, DATA_PREFIX};
use ipnet::{IpNet, Ipv4Net, Ipv6Net};
use iptrie::RTrieSet;
use mnl::mnl_sys;
use nftnl::{nftnl_sys, set::SetKey, Batch, FinalizedBatch, MsgType, NlMsg};
use smallvec::SmallVec;
fn cidr_bound_ipv4(net: Ipv4Net) -> Option<Ipv4Addr> {
let data = u32::from(net.network());
@ -329,191 +324,16 @@ pub fn get_sets(socket: &mnl::Socket) -> io::Result<Vec<Set1>> {
Ok(ret)
}
fn should_add<T: Helper>(trie: &RTrieSet<T>, elem: &T) -> bool {
*trie.lookup(elem) == T::ZERO
}
fn iter_ip_trie<T: Helper>(trie: &RTrieSet<T>) -> impl '_ + Iterator<Item = T> {
trie.iter().copied().filter(|x| {
if let Some(par) = x.direct_parent() {
should_add(trie, &par)
} else {
*x != T::ZERO
}
})
}
pub(crate) struct NftData<T: Helper> {
ips: RTrieSet<T>,
dirty: bool,
set: Option<Set1>,
name: String,
}
impl<T: Helper> NftData<T> {
pub fn new(name: &str) -> Self {
Self {
set: None,
ips: RTrieSet::new(),
dirty: true,
name: name.to_owned(),
}
}
}
// SAFETY: set is None initially so Set1 is never actually sent
// (and it might be fine to send anyway actually)
unsafe impl<T: Helper + Send> Send for NftData<T> {}
impl<T: Helper> NftData<T>
where
IpNet: From<T>,
{
#[must_use]
pub fn verify(&mut self) -> bool {
if !self.name.is_empty() && self.set.is_none() {
self.ips = RTrieSet::new();
false
} else {
true
}
}
pub fn flush_changes(
&mut self,
socket: &mnl::Socket,
flush_set: bool,
) -> Result<(), io::Error> {
if let Some(set) = self.set.as_mut().filter(|_| self.dirty) {
if flush_set {
println!(
"initializing set {} with ~{} ips (e.g. {:?})",
self.name,
self.ips.len(),
iter_ip_trie(&self.ips).next(),
);
}
set.add_cidrs(socket, flush_set, iter_ip_trie(&self.ips).map(IpNet::from))
} else {
Ok(())
}
}
pub fn extend(&mut self, ips: impl Iterator<Item = T>) {
for ip in ips {
self.insert(ip, true);
}
}
pub fn insert(&mut self, ip: T, allow_empty_set: bool) {
if (if allow_empty_set {
!self.name.is_empty()
} else {
self.set.is_some()
}) && should_add(&self.ips, &ip)
{
self.ips.insert(ip);
self.dirty = true;
}
}
pub fn ips_mut(&mut self) -> &mut RTrieSet<T> {
&mut self.ips
}
#[cfg(test)]
pub fn ip_count(&self) -> usize {
iter_ip_trie(&self.ips).count()
}
pub fn name(&self) -> &str {
&self.name
}
pub fn set_set(&mut self, set: Set1) {
self.set = Some(set);
}
}
pub(crate) fn nftables_thread(
mut rulesets: Vec<(NftData<Ipv4Net>, NftData<Ipv6Net>)>,
rx: mpsc::Receiver<(SmallVec<[usize; 5]>, smallvec::SmallVec<[IpNet; 8]>)>,
) {
fn report(err: impl Display) {
println!("nftables: {err}");
if let Ok(mut file) = std::fs::OpenOptions::new()
.append(true)
.create(true)
.open(format!("{DATA_PREFIX}/nftables.log"))
{
file.write_all((err.to_string() + "\n").as_bytes())
.unwrap_or(());
}
}
let socket = mnl::Socket::new(mnl::Bus::Netfilter).unwrap();
let all_sets = crate::nftables::get_sets(&socket).unwrap();
for set in all_sets {
for ruleset in &mut rulesets {
if set.table_name() == Some("global") && set.family() == libc::NFPROTO_INET as u32 {
if set.name() == Some(ruleset.0.name()) {
println!("found set {}", ruleset.0.name());
ruleset.0.set_set(set);
break;
} else if set.name() == Some(ruleset.1.name()) {
println!("found set {}", ruleset.1.name());
ruleset.1.set_set(set);
break;
}
}
}
}
for ruleset in &mut rulesets {
if !ruleset.0.verify() {
report(format!("set {} not found", ruleset.0.name()));
}
if !ruleset.1.verify() {
report(format!("set {} not found", ruleset.1.name()));
}
}
let mut first = true;
loop {
for ruleset in &mut rulesets {
if let Err(err) = ruleset.0.flush_changes(&socket, first) {
report(err);
}
if let Err(err) = ruleset.1.flush_changes(&socket, first) {
report(err);
}
}
if first {
println!("nftables init done");
first = false;
}
let (rulesets1, ips) = match rx.recv() {
Ok(val) => val,
Err(_) => break,
};
for i in rulesets1.into_iter() {
let ruleset = &mut rulesets[i];
for ip1 in ips.iter().copied() {
match ip1 {
IpNet::V4(ip) => ruleset.0.insert(ip, false),
IpNet::V6(ip) => ruleset.1.insert(ip, false),
}
}
}
}
}
#[cfg(test)]
mod test {
use std::net::{Ipv4Addr, Ipv6Addr};
use std::net::Ipv6Addr;
use ipnet::{Ipv4Net, Ipv6Net};
use iptrie::RTrieSet;
use crate::nftables::{iter_ip_trie, should_add};
use ipnet::Ipv6Net;
use super::get_sets;
#[test]
fn test_nftables() {
if !nix::unistd::Uid::effective().is_root() {
return;
}
let socket = mnl::Socket::new(mnl::Bus::Netfilter).unwrap();
let sets = get_sets(&socket).unwrap();
assert!(!sets.is_empty());
@ -532,24 +352,4 @@ mod test {
}
panic!();
}
#[test]
fn test_set() {
let mut trie = RTrieSet::new();
assert!(should_add(
&trie,
&Ipv4Net::new(Ipv4Addr::new(127, 0, 0, 1), 32).unwrap()
));
trie.insert(Ipv4Net::new(Ipv4Addr::new(127, 0, 0, 1), 32).unwrap());
assert!(!should_add(
&trie,
&Ipv4Net::new(Ipv4Addr::new(127, 0, 0, 1), 32).unwrap()
));
trie.insert(Ipv4Net::new(Ipv4Addr::new(127, 0, 0, 1), 31).unwrap());
assert!(dbg!(iter_ip_trie(&trie).collect::<Vec<_>>()).len() == 1);
// contains 0.0.0.0, etc
assert!(dbg!(trie.iter().collect::<Vec<_>>()).len() == 3);
trie.insert(Ipv4Net::new(Ipv4Addr::new(127, 0, 1, 1), 32).unwrap());
assert!(dbg!(iter_ip_trie(&trie).collect::<Vec<_>>()).len() == 2);
}
}

View file

@ -609,21 +609,6 @@ pub enum ModuleExtState {
Unknown = 99,
}
impl ModuleExtState {
pub(crate) fn importance(&self) -> usize {
match *self {
Self::Unknown => 0,
Self::InitialState => 1,
Self::Finished => 2,
Self::WaitModule => 3,
Self::RestartNext => 4,
Self::WaitSubquery => 5,
Self::WaitReply => 6,
Self::Error => 7,
}
}
}
impl From<module_ext_state> for ModuleExtState {
fn from(value: module_ext_state) -> Self {
match value {