chayleaf/unbound-rust-mod
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: chayleaf:be0359f85ff009a2cb99753d74aad3bd76f1b9f7
Branches Tags
chayleaf:master
..
compare: chayleaf:e9a6f296dfef8b4382029fdb409038e77891bb19
Branches Tags
chayleaf:master

No commits in common. "be0359f85ff009a2cb99753d74aad3bd76f1b9f7" and "e9a6f296dfef8b4382029fdb409038e77891bb19" have entirely different histories.
be0359f85f ... e9a6f296df

7 changed files with 286 additions and 305 deletions
Show stats Expand all files Collapse all files

14
src/combine.rs
View file

@ -14,20 +14,20 @@ macro_rules! impl_tuple {
{ {
type EnvData = A::EnvData; type EnvData = A::EnvData;
type QstateData = A::QstateData; type QstateData = A::QstateData;
fn init(env: &mut crate::unbound::ModuleEnvMut<Self::EnvData>) -> Result<Self, ()> { fn init(env: &mut crate::unbound::ModuleEnv<Self::EnvData>) -> Result<Self, ()> {
Ok((A::init(env)?, $($t::init(env)?, )*)) Ok((A::init(env)?, $($t::init(env)?, )*))
} }
fn clear(&self, qstate: &mut crate::unbound::ModuleQstateMut<Self::QstateData>) { fn clear(&self, qstate: &mut crate::unbound::ModuleQstate<Self::QstateData>) {
self.0.clear(qstate); self.0.clear(qstate);
$(self.$i.clear(qstate);)* $(self.$i.clear(qstate);)*
} }
fn deinit(self, env: &mut crate::unbound::ModuleEnvMut<Self::EnvData>) { fn deinit(self, env: &mut crate::unbound::ModuleEnv<Self::EnvData>) {
self.0.deinit(env); self.0.deinit(env);
$(self.$i.deinit(env);)* $(self.$i.deinit(env);)*
} }
fn operate( fn operate(
&self, &self,
qstate: &mut crate::unbound::ModuleQstateMut<Self::QstateData>, qstate: &mut crate::unbound::ModuleQstate<Self::QstateData>,
event: crate::unbound::ModuleEvent, event: crate::unbound::ModuleEvent,
entry: &mut crate::unbound::OutboundEntryMut, entry: &mut crate::unbound::OutboundEntryMut,
) -> Option<ModuleExtState> { ) -> Option<ModuleExtState> {
@ -40,13 +40,13 @@ macro_rules! impl_tuple {
})* })*
ret ret
} }
fn get_mem(&self, env: &mut crate::unbound::ModuleEnvMut<Self::EnvData>) -> usize { fn get_mem(&self, env: &mut crate::unbound::ModuleEnv<Self::EnvData>) -> usize {
self.0.get_mem(env) $(* self.$i.get_mem(env))* self.0.get_mem(env) $(* self.$i.get_mem(env))*
} }
fn inform_super( fn inform_super(
&self, &self,
qstate: &mut crate::unbound::ModuleQstateMut<Self::QstateData>, qstate: &mut crate::unbound::ModuleQstate<Self::QstateData>,
super_qstate: &mut crate::unbound::ModuleQstateMut<std::ffi::c_void>, super_qstate: &mut crate::unbound::ModuleQstate<std::ffi::c_void>,
) { ) {
self.0.inform_super(qstate, super_qstate); self.0.inform_super(qstate, super_qstate);
$(self.$i.inform_super(qstate, super_qstate);)* $(self.$i.inform_super(qstate, super_qstate);)*

7
src/domain_tree.rs
View file

@ -1,4 +1,4 @@
use std::{borrow::Borrow, collections::HashMap, hash::Hash}; use std::{collections::HashMap, hash::Hash};
use smallvec::{smallvec, SmallVec}; use smallvec::{smallvec, SmallVec};
@ -36,10 +36,9 @@ impl<T: Hash + Eq> PrefixSet<T> {
} }
} }
} }
pub fn contains<'a, Y>(&self, val: impl IntoIterator<Item = &'a Y>) -> bool pub fn contains<'a>(&self, val: impl IntoIterator<Item = &'a T>) -> bool
where where
T: 'a + Borrow<Y>, T: 'a,
Y: 'a + ?Sized + Eq + Hash,
{ {
match self { match self {
Self::Leaf => true, Self::Leaf => true,

332
src/example.rs
View file

@ -99,26 +99,8 @@ struct IpCache<T>(
); );
#[repr(transparent)] #[repr(transparent)]
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct IpCacheKey(Domain); struct IpCacheKey(Domain);
impl IpCacheKey {
fn from_split_domain<T: AsRef<[u8]>>(
split_domain: impl DoubleEndedIterator + Iterator<Item = T>,
) -> Self {
Self::from_split_rev_domain(split_domain.rev())
}
fn from_split_rev_domain<T: AsRef<[u8]>>(split_rev_domain: impl Iterator<Item = T>) -> Self {
let mut first = true;
Self(split_rev_domain.fold(Domain::new(), |mut ret, seg| {
if first {
first = false;
ret.push(b'.');
}
ret.extend_from_slice(seg.as_ref());
ret
}))
}
}
impl radix_trie::TrieKey for IpCacheKey { impl radix_trie::TrieKey for IpCacheKey {
fn encode_bytes(&self) -> Vec<u8> { fn encode_bytes(&self) -> Vec<u8> {
self.0.to_vec() self.0.to_vec()
@ -136,11 +118,8 @@ impl<T> Default for IpCache<T> {
fn ignore<T>(_: &mut smallvec::SmallVec<[T; 4]>) {} fn ignore<T>(_: &mut smallvec::SmallVec<[T; 4]>) {}
impl<T> IpCache<T> { impl<T> IpCache<T> {
fn extend_set_with_domain<J: IpPrefix + From<T>>( fn extend_set_with_domain<J: IpPrefix + From<T>>(&self, ips: &mut RTrieSet<J>, domain_r: Domain)
&self, where
ips: &mut RTrieSet<J>,
domain_r: IpCacheKey,
) where
T: Copy, T: Copy,
{ {
self.get_maybe_update_rev(domain_r, |val| { self.get_maybe_update_rev(domain_r, |val| {
@ -151,14 +130,15 @@ impl<T> IpCache<T> {
let mut val = Some(ignore); let mut val = Some(ignore);
val = None; val = None;
val val
}); })
} }
fn get_maybe_update_rev<F: for<'a> FnOnce(&'a mut smallvec::SmallVec<[T; 4]>)>( fn get_maybe_update_rev<F: for<'a> FnOnce(&'a mut smallvec::SmallVec<[T; 4]>)>(
&self, &self,
domain_r: IpCacheKey, domain_r: Domain,
upd: impl FnOnce(Option<(&smallvec::SmallVec<[T; 4]>, &Mutex<()>, &AtomicBool)>) -> Option<F>, upd: impl FnOnce(Option<(&smallvec::SmallVec<[T; 4]>, &Mutex<()>, &AtomicBool)>) -> Option<F>,
) { ) {
let lock = self.0.read().unwrap(); let lock = self.0.read().unwrap();
let domain_r = IpCacheKey(domain_r);
let key = lock.0.get(&domain_r).copied(); let key = lock.0.get(&domain_r).copied();
if let Some(val) = if let Some(x) = key.and_then(|key| lock.1.get(key)) { if let Some(val) = if let Some(x) = key.and_then(|key| lock.1.get(key)) {
upd(Some((&x.0.read().unwrap(), &x.1, &x.2))) upd(Some((&x.0.read().unwrap(), &x.1, &x.2)))
@ -191,7 +171,7 @@ impl<T: ToString + PartialEq> IpCache<T> {
let ret1 = &mut ret; let ret1 = &mut ret;
let mut path = self.1.clone(); let mut path = self.1.clone();
path.push(domain); path.push(domain);
self.get_maybe_update_rev(domain_r, |ips| { self.get_maybe_update_rev(domain_r.0, |ips| {
if let Some(ips) = ips.as_ref().filter(|x| x.0 == &val) { if let Some(ips) = ips.as_ref().filter(|x| x.0 == &val) {
*ret1 = false; *ret1 = false;
if ips if ips
@ -236,7 +216,7 @@ impl<T: FromStr> IpCache<T> {
let mut lock = self.0.write().unwrap(); let mut lock = self.0.write().unwrap();
assert!(lock.1.is_empty()); assert!(lock.1.is_empty());
let domains = std::fs::read_dir(dir)?; let domains = std::fs::read_dir(dir)?;
for entry in domains.filter_map(Result::ok) { for entry in domains.filter_map(|x| x.ok()) {
let domain = entry.file_name(); let domain = entry.file_name();
let Some(domain) = domain.to_str() else { let Some(domain) = domain.to_str() else {
continue; continue;
@ -254,7 +234,15 @@ impl<T: FromStr> IpCache<T> {
let Ok(reader) = std::fs::File::open(entry.path()) else { let Ok(reader) = std::fs::File::open(entry.path()) else {
continue; continue;
}; };
let domain_r = IpCacheKey::from_split_domain(domain.split('.')); let domain_r = IpCacheKey(
domain
.split('.')
.rev()
.map(|x| x.as_bytes())
.collect::<Vec<_>>()
.join(&b"."[..])
.into(),
);
let mut reader = BufReader::new(reader); let mut reader = BufReader::new(reader);
let mut line = String::new(); let mut line = String::new();
let mut ips = SmallVec::new(); let mut ips = SmallVec::new();
@ -292,7 +280,7 @@ pub(crate) const DATA_PREFIX: &str = "unbound-mod-test-data";
pub(crate) const CONFIG_PREFIX: &str = "unbound-mod-test-config"; pub(crate) const CONFIG_PREFIX: &str = "unbound-mod-test-config";
#[cfg(not(debug_assertions))] #[cfg(not(debug_assertions))]
pub(crate) const DATA_PREFIX: &str = "/var/lib/unbound"; pub(crate) const PATH_PREFIX: &str = "/var/lib/unbound";
#[cfg(not(debug_assertions))] #[cfg(not(debug_assertions))]
pub(crate) const CONFIG_PREFIX: &str = "/etc/unbound"; pub(crate) const CONFIG_PREFIX: &str = "/etc/unbound";
@ -324,7 +312,7 @@ impl ExampleMod {
Ok(ret) Ok(ret)
} }
fn load_json(&mut self, rulesets: &mut [(NftData<Ipv4Net>, NftData<Ipv6Net>)]) { fn load_json(&mut self, rulesets: &mut [(NftData<Ipv4Net>, NftData<Ipv6Net>)]) {
for (k, v) in &mut self.nft_queries { for (k, v) in self.nft_queries.iter_mut() {
let r = &mut rulesets[v.index]; let r = &mut rulesets[v.index];
let mut v_domains = v.domains.write().unwrap(); let mut v_domains = v.domains.write().unwrap();
for base in [CONFIG_PREFIX, DATA_PREFIX] { for base in [CONFIG_PREFIX, DATA_PREFIX] {
@ -333,8 +321,13 @@ impl ExampleMod {
match read_json::<Vec<String>>(file) { match read_json::<Vec<String>>(file) {
Ok(domains) => { Ok(domains) => {
for domain in domains { for domain in domains {
v_domains v_domains.insert(
.insert(domain.split('.').rev().map(|x| x.as_bytes().into())); domain
.split('.')
.rev()
.map(|x| x.as_bytes().into())
.collect::<SmallVec<[DomainSeg; 5]>>(),
);
} }
} }
Err(err) => Self::report2(&self.error_lock, "domains", err), Err(err) => Self::report2(&self.error_lock, "domains", err),
@ -345,8 +338,13 @@ impl ExampleMod {
match read_json::<Vec<DpiInfo>>(file) { match read_json::<Vec<DpiInfo>>(file) {
Ok(dpi_info) => { Ok(dpi_info) => {
for domain in dpi_info.iter().flat_map(|x| &x.domains) { for domain in dpi_info.iter().flat_map(|x| &x.domains) {
v_domains v_domains.insert(
.insert(domain.split('.').rev().map(|x| x.as_bytes().into())); domain
.split('.')
.rev()
.map(|x| x.as_bytes().into())
.collect::<SmallVec<[DomainSeg; 5]>>(),
);
} }
} }
Err(err) => Self::report2(&self.error_lock, "dpi", err), Err(err) => Self::report2(&self.error_lock, "dpi", err),
@ -377,7 +375,11 @@ impl ExampleMod {
} }
println!("loading cached domain ips for {k}"); println!("loading cached domain ips for {k}");
for rev_domain in v_domains.iter() { for rev_domain in v_domains.iter() {
let rev_domain = IpCacheKey::from_split_rev_domain(rev_domain.into_iter()); let rev_domain: SmallVec<_> = rev_domain
.map(|x| x.as_slice())
.collect::<Vec<_>>()
.join(&b"."[..])
.into();
self.caches self.caches
.0 .0
.extend_set_with_domain(r.0.ips_mut(), rev_domain.clone()); .extend_set_with_domain(r.0.ips_mut(), rev_domain.clone());
@ -399,7 +401,7 @@ impl ExampleMod {
if let Some(s) = std::env::var_os("NFT_QUERIES") { if let Some(s) = std::env::var_os("NFT_QUERIES") {
for (i, (name, set4, set6)) in s for (i, (name, set4, set6)) in s
.to_str() .to_str()
.map(ToOwned::to_owned) .map(|x| x.to_owned())
.ok_or(())? .ok_or(())?
.split(';') .split(';')
.filter_map(|x| x.split_once(':')) .filter_map(|x| x.split_once(':'))
@ -408,9 +410,11 @@ impl ExampleMod {
}) })
.enumerate() .enumerate()
{ {
let (name, dynamic) = name let (name, dynamic) = if let Some(name) = name.strip_suffix('!') {
.strip_suffix('!') (name, true)
.map_or((name, false), |name| (name, true)); } else {
(name, false)
};
self.nft_queries.insert( self.nft_queries.insert(
name.to_owned(), name.to_owned(),
NftQuery { NftQuery {
@ -452,7 +456,7 @@ impl ExampleMod {
} }
fn handle_reply_info( fn handle_reply_info(
&self, &self,
split_domain: &[&[u8]], split_rev_domain: SmallVec<[DomainSeg; 5]>,
qnames: SmallVec<[usize; 5]>, qnames: SmallVec<[usize; 5]>,
rep: &ReplyInfo<'_>, rep: &ReplyInfo<'_>,
) -> Result<(), ()> { ) -> Result<(), ()> {
@ -481,80 +485,91 @@ impl ExampleMod {
} }
} }
} }
self.add_ips(ip4, ip6, split_domain, qnames) self.add_ips(ip4, ip6, split_rev_domain, qnames)
} }
fn add_ips( fn add_ips(
&self, &self,
ip4: SmallVec<[Ipv4Addr; 4]>, ip4: SmallVec<[Ipv4Addr; 4]>,
ip6: SmallVec<[Ipv6Addr; 4]>, ip6: SmallVec<[Ipv6Addr; 4]>,
split_domain: &[&[u8]], split_rev_domain: SmallVec<[DomainSeg; 5]>,
qnames: SmallVec<[usize; 5]>, qnames: SmallVec<[usize; 5]>,
) -> Result<(), ()> { ) -> Result<(), ()> {
println!("adding {ip4:?}/{ip6:?} for {split_domain:?} to {qnames:?}");
if !ip4.is_empty() || !ip6.is_empty() { if !ip4.is_empty() || !ip6.is_empty() {
let domain = match split_domain let domain = match split_rev_domain
.iter() .iter()
.copied() .rev()
.map(std::str::from_utf8) .map(|x| String::from_utf8(x.to_vec()).map(|x| x + "."))
.try_fold(String::new(), |mut s, comp| { .collect::<Result<String, _>>()
if !s.is_empty() { {
s.push('.'); Ok(mut x) => {
} x.pop();
s.push_str(comp?); x
Ok::<_, std::str::Utf8Error>(s) }
}) {
Ok(x) => x,
Err(err) => { Err(err) => {
self.report("domain utf-8", err); self.report("domain utf-8", err);
return Err(()); return Err(());
} }
}; };
let key = IpCacheKey::from_split_domain(split_domain.iter()); let mut split_rev_domain = split_rev_domain.into_iter();
let mut to_send: SmallVec<[IpNet; 8]> = SmallVec::new(); if let Some(first) = split_rev_domain.next() {
to_send.extend(ip4.iter().copied().map(Ipv4Net::from).map(IpNet::from)); let first: Domain = first.to_vec().into();
to_send.extend(ip6.iter().copied().map(Ipv6Net::from).map(IpNet::from)); let joined_rev_domain = split_rev_domain.fold(first, |mut res, mut next| {
let keep4 = !ip4.is_empty() && self.caches.0.set(&domain, key.clone(), ip4); res.push(b'.');
let keep6 = !ip6.is_empty() && self.caches.1.set(&domain, key, ip6); res.append(&mut next);
to_send.retain(|x| x.addr().is_ipv4() && keep4 || x.addr().is_ipv6() && keep6); res
if !to_send.is_empty() { });
self.ruleset_queue let mut to_send: SmallVec<[IpNet; 8]> = SmallVec::new();
.as_ref() to_send.extend(ip4.iter().copied().map(Ipv4Net::from).map(IpNet::from));
.unwrap() to_send.extend(ip6.iter().copied().map(Ipv6Net::from).map(IpNet::from));
.send((qnames, to_send)) let keep4 = !ip4.is_empty()
.unwrap(); && self
.caches
.0
.set(&domain, IpCacheKey(joined_rev_domain.clone()), ip4);
let keep6 = !ip6.is_empty()
&& self
.caches
.1
.set(&domain, IpCacheKey(joined_rev_domain.clone()), ip6);
to_send.retain(|x| x.addr().is_ipv4() && keep4 || x.addr().is_ipv6() && keep6);
if !to_send.is_empty() {
self.ruleset_queue
.as_ref()
.unwrap()
.send((qnames, to_send))
.unwrap();
}
} }
} }
Ok(()) Ok(())
} }
fn run_commands(&self, split_domain: &[&[u8]]) -> Option<ModuleExtState> { fn run_commands(&self, rev_domain: &[u8]) -> Option<ModuleExtState> {
if let Some(split_domain) = self.nft_token.as_ref().and_then(|token| { if let Some(rev_domain) = self
split_domain .nft_token
.split_last() .as_ref()
.filter(|(a, _)| **a == token.as_bytes()) .and_then(|token| rev_domain.strip_prefix(token.as_bytes()))
.map(|(_, b)| b) {
}) { for (qname, query) in self.nft_queries.iter() {
for (qname, query) in &self.nft_queries { if query.dynamic && rev_domain.starts_with(qname.as_bytes()) {
if query.dynamic { if let Some(rev_domain) =
if let Some(split_domain) = split_domain rev_domain.strip_prefix((qname.to_owned() + ".").as_bytes())
.split_last()
.filter(|(a, _)| **a == qname.as_bytes())
.map(|(_, b)| b)
{ {
let rev_domain = rev_domain
.split(|x| *x == b'.')
.map(|x| x.into())
.collect::<SmallVec<[_; 5]>>();
let mut domains = query.domains.write().unwrap(); let mut domains = query.domains.write().unwrap();
if domains.insert(split_domain.iter().copied().rev().map(From::from)) { if domains.insert(rev_domain.clone()) {
drop(domains); drop(domains);
let file_name = format!("{DATA_PREFIX}/{qname}_domains.json"); let file_name = format!("{DATA_PREFIX}/{qname}_domains.json");
let domain = match split_domain let domain = match String::from_utf8(
.iter() rev_domain
.copied() .iter()
.map(std::str::from_utf8) .rev()
.try_fold(String::new(), |mut s, comp| { .map(|x| x.as_slice())
if !s.is_empty() { .collect::<Vec<_>>()
s.push('.'); .join(&b"."[..]),
} ) {
s.push_str(comp?);
Ok::<_, std::str::Utf8Error>(s)
}) {
Ok(x) => x, Ok(x) => x,
Err(err) => { Err(err) => {
self.report("domain utf-8", err); self.report("domain utf-8", err);
@ -562,7 +577,6 @@ impl ExampleMod {
} }
}; };
let _lock = self.domains_write_lock.lock().unwrap(); let _lock = self.domains_write_lock.lock().unwrap();
println!("adding {domain} to {qname}");
let mut old: Vec<String> = if let Ok(file) = File::open(&file_name) { let mut old: Vec<String> = if let Ok(file) = File::open(&file_name) {
match read_json(file) { match read_json(file) {
Ok(x) => x, Ok(x) => x,
@ -588,21 +602,22 @@ impl ExampleMod {
} }
} }
return Some(ModuleExtState::Finished); return Some(ModuleExtState::Finished);
} else if let Some(split_domain) = self.tmp_nft_token.as_ref().and_then(|token| { } else if let Some(rev_domain) = self
split_domain .tmp_nft_token
.split_last() .as_ref()
.filter(|(a, _)| **a == token.as_bytes()) .and_then(|token| rev_domain.strip_prefix(token.as_bytes()))
.map(|(_, b)| b) {
}) { for (qname, query) in self.nft_queries.iter() {
for (qname, query) in &self.nft_queries { if query.dynamic && rev_domain.starts_with(qname.as_bytes()) {
if query.dynamic { if let Some(rev_domain) =
if let Some(split_domain) = split_domain rev_domain.strip_prefix((qname.to_owned() + ".").as_bytes())
.split_last()
.filter(|(a, _)| **a == qname.as_bytes())
.map(|(_, b)| b)
{ {
let rev_domain = rev_domain
.split(|x| *x == b'.')
.map(|x| x.into())
.collect::<SmallVec<[_; 5]>>();
let mut domains = query.domains.write().unwrap(); let mut domains = query.domains.write().unwrap();
domains.insert(split_domain.iter().copied().rev().map(From::from)); domains.insert(rev_domain.clone());
} }
} }
} }
@ -610,15 +625,10 @@ impl ExampleMod {
} }
None None
} }
fn get_qnames(&self, split_domain: &[&[u8]]) -> SmallVec<[usize; 5]> { fn get_qnames(&self, split_rev_domain: &SmallVec<[DomainSeg; 5]>) -> SmallVec<[usize; 5]> {
let mut qnames: SmallVec<[usize; 5]> = SmallVec::new(); let mut qnames: SmallVec<[usize; 5]> = SmallVec::new();
for query in self.nft_queries.values() { for query in self.nft_queries.values() {
if query if query.domains.read().unwrap().contains(split_rev_domain) {
.domains
.read()
.unwrap()
.contains(split_domain.iter().copied().rev().map(From::from))
{
qnames.push(query.index); qnames.push(query.index);
} }
} }
@ -670,31 +680,17 @@ fn read_json<T: 'static + for<'a> Deserialize<'a>>(mut f: File) -> Result<T, ser
serde_json::from_slice(&data) serde_json::from_slice(&data)
} }
// \x06google\x03com
fn unwire_domain(domain: &[u8]) -> SmallVec<[&[u8]; 8]> {
let mut i = 0;
let mut ret = SmallVec::new();
while let Some(val) = domain.get(i).map(|x| *x as usize) {
i += 1;
if let Some(val) = domain.get(i..i + val) {
ret.push(val);
}
i += val;
}
ret
}
impl UnboundMod for ExampleMod { impl UnboundMod for ExampleMod {
type EnvData = (); type EnvData = ();
type QstateData = (); type QstateData = ();
fn init(_env: &mut crate::unbound::ModuleEnvMut<Self::EnvData>) -> Result<Self, ()> { fn init(_env: &mut crate::unbound::ModuleEnv<Self::EnvData>) -> Result<Self, ()> {
Self::new() Self::new()
} }
fn operate( fn operate(
&self, &self,
qstate: &mut crate::unbound::ModuleQstateMut<Self::QstateData>, qstate: &mut crate::unbound::ModuleQstate<Self::QstateData>,
event: ModuleEvent, event: ModuleEvent,
_entry: &mut crate::unbound::OutboundEntryMut, _entry: &mut crate::unbound::OutboundEntryMut,
) -> Option<ModuleExtState> { ) -> Option<ModuleExtState> {
@ -707,21 +703,26 @@ impl UnboundMod for ExampleMod {
return Some(ModuleExtState::Error); return Some(ModuleExtState::Error);
} }
} }
let info = qstate.qinfo(); let info = qstate.qinfo_mut();
let name = info.qname().to_bytes(); let name = info.qname().to_bytes();
// let rev_domain = name.strip_suffix(b".").unwrap_or(name); let rev_domain = name.strip_suffix(b".").unwrap_or(name);
let split_domain = unwire_domain(name); if let Some(val) = self.run_commands(rev_domain) {
println!("handling {split_domain:?}");
if let Some(val) = self.run_commands(&split_domain) {
return Some(val); return Some(val);
} }
let qnames = self.get_qnames(&split_domain); let split_rev_domain = rev_domain
.split(|x| *x == b'.')
.map(|x| x.into())
.collect::<SmallVec<[_; 5]>>();
let qnames = self.get_qnames(&split_rev_domain);
if qnames.is_empty() { if qnames.is_empty() {
return Some(ModuleExtState::Finished); return Some(ModuleExtState::Finished);
} }
if let Some(ret) = qstate.return_msg() { if let Some(ret) = qstate.return_msg_mut() {
if let Some(rep) = ret.rep() { if let Some(rep) = ret.rep() {
if self.handle_reply_info(&split_domain, qnames, &rep).is_err() { if self
.handle_reply_info(split_rev_domain, qnames, &rep)
.is_err()
{
return Some(ModuleExtState::Error); return Some(ModuleExtState::Error);
} }
} }
@ -740,9 +741,9 @@ mod test {
use std::{net::Ipv4Addr, os::unix::fs::MetadataExt, path::PathBuf, str::FromStr, sync::mpsc}; use std::{net::Ipv4Addr, os::unix::fs::MetadataExt, path::PathBuf, str::FromStr, sync::mpsc};
use ipnet::IpNet; use ipnet::IpNet;
use smallvec::smallvec; use smallvec::{smallvec, SmallVec};
use crate::example::{ignore, ExampleMod, IpCacheKey, IpNetDeser, DATA_PREFIX}; use crate::example::{ignore, ExampleMod, IpNetDeser, DATA_PREFIX};
#[test] #[test]
fn test() { fn test() {
@ -793,26 +794,24 @@ mod test {
base_path.push("domains6"); base_path.push("domains6");
t.caches.1.load(&base_path).unwrap(); t.caches.1.load(&base_path).unwrap();
t.caches.0.get_maybe_update_rev( t.caches
IpCacheKey::from_split_domain(["a", "com"].into_iter()), .0
|x| { .get_maybe_update_rev("com.a".as_bytes().into(), |x| {
assert!(x.unwrap().0.len() == 2); assert!(x.unwrap().0.len() == 2);
#[allow(unused_assignments)] #[allow(unused_assignments)]
let mut val = Some(ignore); let mut val = Some(ignore);
val = None; val = None;
val val
}, });
); t.caches
t.caches.0.get_maybe_update_rev( .0
IpCacheKey::from_split_domain(["b", "com"].into_iter()), .get_maybe_update_rev("com.b".as_bytes().into(), |x| {
|x| {
assert!(x.unwrap().0.len() == 1); assert!(x.unwrap().0.len() == 1);
#[allow(unused_assignments)] #[allow(unused_assignments)]
let mut val = Some(ignore); let mut val = Some(ignore);
val = None; val = None;
val val
}, });
);
t.load_json(&mut rulesets); t.load_json(&mut rulesets);
@ -839,53 +838,48 @@ mod test {
tx2.send(rulesets).unwrap(); tx2.send(rulesets).unwrap();
}); });
let split_domain = [&b"c"[..], &b"com"[..]]; let split_rev_domain = smallvec![SmallVec::from(&b"com"[..]), SmallVec::from(&b"c"[..])];
let qnames = t.get_qnames(&split_domain); let qnames = t.get_qnames(&split_rev_domain);
assert_eq!(qnames.len(), 2); assert_eq!(qnames.len(), 2);
t.add_ips( t.add_ips(
smallvec![Ipv4Addr::new(7, 7, 7, 7), Ipv4Addr::new(6, 6, 6, 6)], smallvec![Ipv4Addr::new(7, 7, 7, 7), Ipv4Addr::new(6, 6, 6, 6)],
smallvec![], smallvec![],
&split_domain, split_rev_domain,
qnames, qnames,
) )
.unwrap(); .unwrap();
let split_domain = [&b"a"[..], &b"com"[..]]; let split_rev_domain = smallvec![SmallVec::from(&b"com"[..]), SmallVec::from(&b"a"[..])];
let qnames = t.get_qnames(&split_domain); let qnames = t.get_qnames(&split_rev_domain);
t.add_ips( t.add_ips(
smallvec![Ipv4Addr::new(1, 2, 3, 4), Ipv4Addr::new(5, 6, 7, 8)], smallvec![Ipv4Addr::new(1, 2, 3, 4), Ipv4Addr::new(5, 6, 7, 8)],
smallvec![], smallvec![],
&split_domain, split_rev_domain,
qnames, qnames,
) )
.unwrap(); .unwrap();
t.run_commands(&[&b"w"[..], &b"com"[..], &b"q"[..], &b"token"[..]]) t.run_commands(b"token.q.com.w").unwrap();
.unwrap(); t.run_commands(b"tmptoken.q.com.e").unwrap();
t.run_commands(&[&b"e"[..], &b"com"[..], &b"q"[..], &b"tmptoken"[..]])
.unwrap();
assert!(t
.run_commands(&[&b"e"[..], &b"com"[..], &b"w"[..], &b"tmptoken"[..]])
.is_none());
let split_domain = [&b"e"[..], &b"com"[..]]; let split_rev_domain = smallvec![SmallVec::from(&b"com"[..]), SmallVec::from(&b"e"[..])];
let qnames = t.get_qnames(&split_domain); let qnames = t.get_qnames(&split_rev_domain);
assert_eq!(qnames.len(), 1); assert_eq!(qnames.len(), 1);
t.add_ips( t.add_ips(
smallvec![Ipv4Addr::new(8, 8, 8, 8)], smallvec![Ipv4Addr::new(8, 8, 8, 8)],
smallvec![], smallvec![],
&split_domain, split_rev_domain,
qnames, qnames,
) )
.unwrap(); .unwrap();
let split_domain = [&b"w"[..], &b"com"[..]]; let split_rev_domain = smallvec![SmallVec::from(&b"com"[..]), SmallVec::from(&b"w"[..])];
let qnames = t.get_qnames(&split_domain); let qnames = t.get_qnames(&split_rev_domain);
assert_eq!(qnames.len(), 1); assert_eq!(qnames.len(), 1);
t.add_ips( t.add_ips(
smallvec![Ipv4Addr::new(9, 8, 8, 8)], smallvec![Ipv4Addr::new(9, 8, 8, 8)],
smallvec![], smallvec![],
&split_domain, split_rev_domain,
qnames, qnames,
) )
.unwrap(); .unwrap();

18
src/exports.rs
View file

@ -16,7 +16,11 @@ pub unsafe extern "C" fn init(
env: *mut module_env, env: *mut module_env,
id: ::std::os::raw::c_int, id: ::std::os::raw::c_int,
) -> ::std::os::raw::c_int { ) -> ::std::os::raw::c_int {
crate::MODULE_FACTORY.take().map_or(0, |fac| fac(env, id)) if let Some(fac) = crate::MODULE_FACTORY.take() {
fac(env, id)
} else {
0
}
} }
/// Deinitialize module internals. /// Deinitialize module internals.
@ -29,7 +33,7 @@ pub unsafe extern "C" fn deinit(env: *mut module_env, id: ::std::os::raw::c_int)
} }
/// Perform action on pending query. Accepts a new query, or work on pending query. /// Perform action on pending query. Accepts a new query, or work on pending query.
/// You have to set `qstate.ext_state` on exit. /// You have to set qstate.ext_state on exit.
/// The state informs unbound about result and controls the following states. /// The state informs unbound about result and controls the following states.
/// ///
/// # Arguments /// # Arguments
@ -46,7 +50,7 @@ pub unsafe extern "C" fn operate(
entry: *mut outbound_entry, entry: *mut outbound_entry,
) { ) {
if let Some(module) = crate::module() { if let Some(module) = crate::module() {
module.internal_operate(qstate, event, id, entry); module.internal_operate(qstate, event, id, entry)
} }
} }
@ -65,7 +69,7 @@ pub unsafe extern "C" fn inform_super(
super_qstate: *mut module_qstate, super_qstate: *mut module_qstate,
) { ) {
if let Some(module) = crate::module() { if let Some(module) = crate::module() {
module.internal_inform_super(qstate, id, super_qstate); module.internal_inform_super(qstate, id, super_qstate)
} }
} }
@ -74,7 +78,7 @@ pub unsafe extern "C" fn inform_super(
#[no_mangle] #[no_mangle]
pub unsafe extern "C" fn clear(qstate: *mut module_qstate, id: ::std::os::raw::c_int) { pub unsafe extern "C" fn clear(qstate: *mut module_qstate, id: ::std::os::raw::c_int) {
if let Some(module) = crate::module() { if let Some(module) = crate::module() {
module.internal_clear(qstate, id); module.internal_clear(qstate, id)
} }
} }
@ -82,7 +86,9 @@ pub unsafe extern "C" fn clear(qstate: *mut module_qstate, id: ::std::os::raw::c
/// only happens explicitly and is only used to show memory usage to the user. /// only happens explicitly and is only used to show memory usage to the user.
#[no_mangle] #[no_mangle]
pub unsafe extern "C" fn get_mem(env: *mut module_env, id: ::std::os::raw::c_int) -> usize { pub unsafe extern "C" fn get_mem(env: *mut module_env, id: ::std::os::raw::c_int) -> usize {
crate::module().map_or(0, |module| module.internal_get_mem(env, id)) crate::module()
.map(|module| module.internal_get_mem(env, id))
.unwrap_or(0)
} }
// function interface assertions // function interface assertions

84
src/lib.rs
View file

@ -9,9 +9,10 @@ use unbound::ModuleExtState;
non_snake_case, non_snake_case,
non_upper_case_globals, non_upper_case_globals,
unused_imports, unused_imports,
clippy::all, clippy::useless_transmute,
clippy::nursery, clippy::type_complexity,
clippy::pedantic clippy::too_many_arguments,
clippy::upper_case_acronyms
)] )]
mod bindings; mod bindings;
mod combine; mod combine;
@ -22,17 +23,21 @@ mod exports;
mod nftables; mod nftables;
mod unbound; mod unbound;
pub fn add(left: usize, right: usize) -> usize {
left + right
}
pub trait UnboundMod: Send + Sync + Sized + RefUnwindSafe + UnwindSafe { pub trait UnboundMod: Send + Sync + Sized + RefUnwindSafe + UnwindSafe {
type EnvData; type EnvData;
type QstateData; type QstateData;
#[allow(clippy::result_unit_err)] #[allow(clippy::result_unit_err)]
fn init(_env: &mut unbound::ModuleEnvMut<Self::EnvData>) -> Result<Self, ()> { fn init(_env: &mut unbound::ModuleEnv<Self::EnvData>) -> Result<Self, ()> {
Err(()) Err(())
} }
fn deinit(self, _env: &mut unbound::ModuleEnvMut<Self::EnvData>) {} fn deinit(self, _env: &mut unbound::ModuleEnv<Self::EnvData>) {}
fn operate( fn operate(
&self, &self,
_qstate: &mut unbound::ModuleQstateMut<Self::QstateData>, _qstate: &mut unbound::ModuleQstate<Self::QstateData>,
_event: unbound::ModuleEvent, _event: unbound::ModuleEvent,
_entry: &mut unbound::OutboundEntryMut, _entry: &mut unbound::OutboundEntryMut,
) -> Option<ModuleExtState> { ) -> Option<ModuleExtState> {
@ -40,13 +45,13 @@ pub trait UnboundMod: Send + Sync + Sized + RefUnwindSafe + UnwindSafe {
} }
fn inform_super( fn inform_super(
&self, &self,
_qstate: &mut unbound::ModuleQstateMut<Self::QstateData>, _qstate: &mut unbound::ModuleQstate<Self::QstateData>,
_super_qstate: &mut unbound::ModuleQstateMut<::std::ffi::c_void>, _super_qstate: &mut unbound::ModuleQstate<::std::ffi::c_void>,
) { ) {
} }
fn clear(&self, _qstate: &mut unbound::ModuleQstateMut<Self::QstateData>) {} fn clear(&self, _qstate: &mut unbound::ModuleQstate<Self::QstateData>) {}
fn get_mem(&self, _env: &mut unbound::ModuleEnvMut<Self::EnvData>) -> usize { fn get_mem(&self, _env: &mut unbound::ModuleEnv<Self::EnvData>) -> usize {
0 0
} }
} }
@ -92,7 +97,7 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
id: ::std::os::raw::c_int, id: ::std::os::raw::c_int,
) { ) {
std::panic::catch_unwind(|| { std::panic::catch_unwind(|| {
self.deinit(&mut unbound::ModuleEnvMut(env, id, Default::default())); self.deinit(&mut unbound::ModuleEnv(env, id, Default::default()))
}) })
.unwrap_or(()); .unwrap_or(());
} }
@ -104,16 +109,13 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
entry: *mut bindings::outbound_entry, entry: *mut bindings::outbound_entry,
) { ) {
std::panic::catch_unwind(|| { std::panic::catch_unwind(|| {
let mut qstate = unbound::ModuleQstate(qstate, id, Default::default());
if let Some(ext_state) = self.operate( if let Some(ext_state) = self.operate(
&mut unbound::ModuleQstateMut(unbound::ModuleQstate( &mut qstate,
qstate,
id,
Default::default(),
)),
event.into(), event.into(),
&mut unbound::OutboundEntryMut(entry, Default::default()), &mut unbound::OutboundEntryMut(entry, Default::default()),
) { ) {
(*qstate).ext_state[id as usize] = ext_state as bindings::module_ext_state; qstate.set_ext_state(ext_state);
} }
}) })
.unwrap_or(()); .unwrap_or(());
@ -126,17 +128,9 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
) { ) {
std::panic::catch_unwind(|| { std::panic::catch_unwind(|| {
self.inform_super( self.inform_super(
&mut unbound::ModuleQstateMut(unbound::ModuleQstate( &mut unbound::ModuleQstate(qstate, id, Default::default()),
qstate, &mut unbound::ModuleQstate(super_qstate, -1, Default::default()),
id, )
Default::default(),
)),
&mut unbound::ModuleQstateMut(unbound::ModuleQstate(
super_qstate,
-1,
Default::default(),
)),
);
}) })
.unwrap_or(()); .unwrap_or(());
} }
@ -146,11 +140,7 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
id: ::std::os::raw::c_int, id: ::std::os::raw::c_int,
) { ) {
std::panic::catch_unwind(|| { std::panic::catch_unwind(|| {
self.clear(&mut unbound::ModuleQstateMut(unbound::ModuleQstate( self.clear(&mut unbound::ModuleQstate(qstate, id, Default::default()))
qstate,
id,
Default::default(),
)));
}) })
.unwrap_or(()); .unwrap_or(());
} }
@ -160,7 +150,7 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
id: ::std::os::raw::c_int, id: ::std::os::raw::c_int,
) -> usize { ) -> usize {
std::panic::catch_unwind(|| { std::panic::catch_unwind(|| {
self.get_mem(&mut unbound::ModuleEnvMut(env, id, Default::default())) self.get_mem(&mut unbound::ModuleEnv(env, id, Default::default()))
}) })
.unwrap_or(0) .unwrap_or(0)
} }
@ -183,13 +173,14 @@ pub fn set_unbound_mod<T: 'static + UnboundMod>() {
MODULE_FACTORY MODULE_FACTORY
.set(Box::new(|env, id| { .set(Box::new(|env, id| {
std::panic::catch_unwind(|| { std::panic::catch_unwind(|| {
T::init(&mut unbound::ModuleEnvMut(env, id, Default::default())).map_or( if let Ok(module) =
0, T::init(&mut unbound::ModuleEnv(env, id, Default::default()))
|module| { {
MODULE.set(Box::new(module)).map_err(|_| ()).unwrap(); MODULE.set(Box::new(module)).map_err(|_| ()).unwrap();
1 1
}, } else {
) 0
}
}) })
.unwrap_or(0) .unwrap_or(0)
})) }))
@ -197,3 +188,14 @@ pub fn set_unbound_mod<T: 'static + UnboundMod>() {
.unwrap(); .unwrap();
} }
} }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn it_works() {
let result = add(2, 2);
assert_eq!(result, 4);
}
}

44
src/nftables.rs
View file

@ -48,7 +48,7 @@ struct FlushSetMsg<'a> {
unsafe impl<'a> NlMsg for FlushSetMsg<'a> { unsafe impl<'a> NlMsg for FlushSetMsg<'a> {
unsafe fn write(&self, buf: *mut std::ffi::c_void, seq: u32, _msg_type: MsgType) { unsafe fn write(&self, buf: *mut std::ffi::c_void, seq: u32, _msg_type: MsgType) {
let header = nftnl_sys::nftnl_nlmsg_build_hdr( let header = nftnl_sys::nftnl_nlmsg_build_hdr(
buf.cast(), buf as *mut c_char,
libc::NFT_MSG_DELSETELEM as u16, libc::NFT_MSG_DELSETELEM as u16,
self.set.family() as u16, self.set.family() as u16,
0, 0,
@ -68,7 +68,9 @@ pub struct SetElemsIter<'a> {
impl<'a> SetElemsIter<'a> { impl<'a> SetElemsIter<'a> {
fn new(set: &'a Set1) -> Self { fn new(set: &'a Set1) -> Self {
let iter = unsafe { nftnl_sys::nftnl_set_elems_iter_create(set.as_mut_ptr()) }; let iter = unsafe { nftnl_sys::nftnl_set_elems_iter_create(set.as_mut_ptr()) };
assert!(!iter.is_null(), "oom"); if iter.is_null() {
panic!("oom");
}
SetElemsIter { SetElemsIter {
set, set,
iter, iter,
@ -123,7 +125,7 @@ unsafe impl<'a> NlMsg for SetElemsMsg<'a> {
MsgType::Del => (libc::NFT_MSG_DELSETELEM, 0), MsgType::Del => (libc::NFT_MSG_DELSETELEM, 0),
}; };
let header = nftnl_sys::nftnl_nlmsg_build_hdr( let header = nftnl_sys::nftnl_nlmsg_build_hdr(
buf.cast(), buf as *mut c_char,
type_ as u16, type_ as u16,
self.set.family() as u16, self.set.family() as u16,
flags as u16, flags as u16,
@ -169,7 +171,7 @@ impl Set1 {
pub fn new() -> Self { pub fn new() -> Self {
Self(unsafe { nftnl_sys::nftnl_set_alloc() }) Self(unsafe { nftnl_sys::nftnl_set_alloc() })
} }
pub const fn as_mut_ptr(&self) -> *mut nftnl_sys::nftnl_set { pub fn as_mut_ptr(&self) -> *mut nftnl_sys::nftnl_set {
self.0 self.0
} }
pub fn table_name(&self) -> Option<&str> { pub fn table_name(&self) -> Option<&str> {
@ -193,26 +195,30 @@ impl Set1 {
let data1_len = data1.len() as u32; let data1_len = data1.len() as u32;
unsafe { unsafe {
let elem = nftnl_sys::nftnl_set_elem_alloc(); let elem = nftnl_sys::nftnl_set_elem_alloc();
assert!(!elem.is_null(), "oom"); if elem.is_null() {
panic!("oom");
}
nftnl_sys::nftnl_set_elem_set( nftnl_sys::nftnl_set_elem_set(
elem, elem,
nftnl_sys::NFTNL_SET_ELEM_KEY as u16, nftnl_sys::NFTNL_SET_ELEM_KEY as u16,
data1.as_ptr().cast(), data1.as_ptr() as *const c_void,
data1_len, data1_len,
); );
nftnl_sys::nftnl_set_elem_add(self.as_mut_ptr(), elem); nftnl_sys::nftnl_set_elem_add(self.as_mut_ptr(), elem);
let Some(data2) = excl_upper.map(SetKey::data) else { let Some(data2) = excl_upper.map(|key| key.data()) else {
return; return;
}; };
let data2_len = data2.len() as u32; let data2_len = data2.len() as u32;
let elem = nftnl_sys::nftnl_set_elem_alloc(); let elem = nftnl_sys::nftnl_set_elem_alloc();
assert!(!elem.is_null(), "oom"); if elem.is_null() {
panic!("oom");
}
nftnl_sys::nftnl_set_elem_set( nftnl_sys::nftnl_set_elem_set(
elem, elem,
nftnl_sys::NFTNL_SET_ELEM_KEY as u16, nftnl_sys::NFTNL_SET_ELEM_KEY as u16,
data2.as_ptr().cast(), data2.as_ptr() as *const c_void,
data2_len, data2_len,
); );
nftnl_sys::nftnl_set_elem_set_u32( nftnl_sys::nftnl_set_elem_set_u32(
@ -238,7 +244,7 @@ impl Set1 {
count += 1; count += 1;
batch.add(&set.flush_msg(), nftnl::MsgType::Del); batch.add(&set.flush_msg(), nftnl::MsgType::Del);
} }
for net in cidrs { for net in cidrs.into_iter() {
if count + 2 > max_batch_size { if count + 2 > max_batch_size {
batch.add_iter(SetElemsIter::new(&set), MsgType::Add); batch.add_iter(SetElemsIter::new(&set), MsgType::Add);
send_and_process(socket, &batch.finalize())?; send_and_process(socket, &batch.finalize())?;
@ -259,7 +265,7 @@ impl Set1 {
send_and_process(socket, &batch.finalize()) send_and_process(socket, &batch.finalize())
} }
const fn flush_msg(&self) -> FlushSetMsg<'_> { fn flush_msg(&self) -> FlushSetMsg<'_> {
FlushSetMsg { set: self } FlushSetMsg { set: self }
} }
} }
@ -276,7 +282,7 @@ pub fn get_sets(socket: &mnl::Socket) -> io::Result<Vec<Set1>> {
let mut ret = Vec::new(); let mut ret = Vec::new();
unsafe { unsafe {
nftnl_sys::nftnl_nlmsg_build_hdr( nftnl_sys::nftnl_nlmsg_build_hdr(
buffer.as_mut_ptr().cast(), buffer.as_mut_ptr() as *mut c_char,
libc::NFT_MSG_GETSET as u16, libc::NFT_MSG_GETSET as u16,
nftnl::ProtoFamily::Inet as u16, nftnl::ProtoFamily::Inet as u16,
(libc::NLM_F_DUMP | libc::NLM_F_ACK) as u16, (libc::NLM_F_DUMP | libc::NLM_F_ACK) as u16,
@ -329,8 +335,11 @@ fn should_add<T: Helper>(trie: &RTrieSet<T>, elem: &T) -> bool {
fn iter_ip_trie<T: Helper>(trie: &RTrieSet<T>) -> impl '_ + Iterator<Item = T> { fn iter_ip_trie<T: Helper>(trie: &RTrieSet<T>) -> impl '_ + Iterator<Item = T> {
trie.iter().copied().filter(|x| { trie.iter().copied().filter(|x| {
x.direct_parent() if let Some(par) = x.direct_parent() {
.map_or_else(|| *x != T::ZERO, |par| should_add(trie, &par)) should_add(trie, &par)
} else {
*x != T::ZERO
}
}) })
} }
@ -473,10 +482,11 @@ pub(crate) fn nftables_thread(
println!("nftables init done"); println!("nftables init done");
first = false; first = false;
} }
let Ok((rulesets1, ips)) = rx.recv() else { let (rulesets1, ips) = match rx.recv() {
break; Ok(val) => val,
Err(_) => break,
}; };
for i in rulesets1 { for i in rulesets1.into_iter() {
let ruleset = &mut rulesets[i]; let ruleset = &mut rulesets[i];
for ip1 in ips.iter().copied() { for ip1 in ips.iter().copied() {
match ip1 { match ip1 {

92
src/unbound.rs
View file

@ -6,10 +6,7 @@ use crate::bindings::{
rrset_trust, sec_status, slabhash, sldns_enum_ede_code, sockaddr_in, sockaddr_in6, rrset_trust, sec_status, slabhash, sldns_enum_ede_code, sockaddr_in, sockaddr_in6,
sockaddr_storage, ub_packed_rrset_key, AF_INET, AF_INET6, sockaddr_storage, ub_packed_rrset_key, AF_INET, AF_INET6,
}; };
use std::{ use std::{ffi::CStr, marker::PhantomData, net::SocketAddr, os::raw::c_char, ptr, time::Duration};
ffi::CStr, marker::PhantomData, net::SocketAddr, ops::Deref, os::raw::c_char, ptr,
time::Duration,
};
pub struct ConfigFileMut<'a>( pub struct ConfigFileMut<'a>(
pub(crate) *mut config_file, pub(crate) *mut config_file,
@ -25,7 +22,7 @@ pub struct InfraCacheMut<'a>(
PhantomData<&'a mut infra_cache>, PhantomData<&'a mut infra_cache>,
); );
pub struct KeyCacheMut<'a>(pub(crate) *mut key_cache, PhantomData<&'a mut key_cache>); pub struct KeyCacheMut<'a>(pub(crate) *mut key_cache, PhantomData<&'a mut key_cache>);
pub struct ModuleEnvMut<T>( pub struct ModuleEnv<T>(
pub(crate) *mut module_env, pub(crate) *mut module_env,
pub(crate) std::ffi::c_int, pub(crate) std::ffi::c_int,
pub(crate) PhantomData<T>, pub(crate) PhantomData<T>,
@ -35,39 +32,18 @@ pub struct ModuleQstate<'a, T>(
pub(crate) std::ffi::c_int, pub(crate) std::ffi::c_int,
pub(crate) PhantomData<&'a mut T>, pub(crate) PhantomData<&'a mut T>,
); );
pub struct ModuleQstateMut<'a, T>(pub(crate) ModuleQstate<'a, T>);
impl<'a, T> Deref for ModuleQstateMut<'a, T> {
type Target = ModuleQstate<'a, T>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
pub struct OutboundEntryMut<'a>( pub struct OutboundEntryMut<'a>(
pub(crate) *mut outbound_entry, pub(crate) *mut outbound_entry,
pub(crate) PhantomData<&'a mut outbound_entry>, pub(crate) PhantomData<&'a mut outbound_entry>,
); );
pub struct QueryInfo<'a>( pub struct QueryInfoMut<'a>(
pub(crate) *mut query_info, pub(crate) *mut query_info,
pub(crate) PhantomData<&'a mut query_info>, pub(crate) PhantomData<&'a mut query_info>,
); );
pub struct QueryInfoMut<'a>(QueryInfo<'a>); pub struct DnsMsgMut<'a>(
impl<'a> Deref for QueryInfoMut<'a> {
type Target = QueryInfo<'a>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
pub struct DnsMsg<'a>(
pub(crate) *mut dns_msg, pub(crate) *mut dns_msg,
pub(crate) PhantomData<&'a mut dns_msg>, pub(crate) PhantomData<&'a mut dns_msg>,
); );
pub struct DnsMsgMut<'a>(DnsMsg<'a>);
impl<'a> Deref for DnsMsgMut<'a> {
type Target = DnsMsg<'a>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
pub struct ReplyInfo<'a>( pub struct ReplyInfo<'a>(
pub(crate) *mut reply_info, pub(crate) *mut reply_info,
pub(crate) PhantomData<&'a mut reply_info>, pub(crate) PhantomData<&'a mut reply_info>,
@ -89,7 +65,7 @@ pub struct PackedRrsetData<'a>(
pub(crate) PhantomData<&'a mut packed_rrset_data>, pub(crate) PhantomData<&'a mut packed_rrset_data>,
); );
impl<'a> QueryInfo<'a> { impl<'a> QueryInfoMut<'a> {
pub fn qname(&self) -> &CStr { pub fn qname(&self) -> &CStr {
unsafe { CStr::from_ptr((*self.0).qname as *const c_char) } unsafe { CStr::from_ptr((*self.0).qname as *const c_char) }
} }
@ -101,7 +77,7 @@ impl<'a> QueryInfo<'a> {
} }
} }
impl<T> ModuleEnvMut<T> { impl<T> ModuleEnv<T> {
pub fn config_file_mut(&mut self) -> ConfigFileMut<'_> { pub fn config_file_mut(&mut self) -> ConfigFileMut<'_> {
ConfigFileMut(unsafe { (*self.0).cfg }, Default::default()) ConfigFileMut(unsafe { (*self.0).cfg }, Default::default())
} }
@ -157,7 +133,7 @@ impl<T> ModuleEnvMut<T> {
addr4.sin_port = x.port(); addr4.sin_port = x.port();
addr4.sin_addr.s_addr = (*x.ip()).into(); addr4.sin_addr.s_addr = (*x.ip()).into();
( (
std::ptr::addr_of!(addr4).cast::<sockaddr_storage>(), &addr4 as *const _ as *const sockaddr_storage,
std::mem::size_of_val(&addr4), std::mem::size_of_val(&addr4),
) )
} }
@ -166,27 +142,29 @@ impl<T> ModuleEnvMut<T> {
addr6.sin6_flowinfo = x.flowinfo(); addr6.sin6_flowinfo = x.flowinfo();
addr6.sin6_scope_id = x.scope_id(); addr6.sin6_scope_id = x.scope_id();
( (
std::ptr::addr_of!(addr6).cast(), &addr6 as *const _ as *const sockaddr_storage,
std::mem::size_of_val(&addr6), std::mem::size_of_val(&addr6),
) )
} }
}; };
((*self.0).send_query.unwrap_unchecked())( ((*self.0).send_query.unwrap_unchecked())(
qinfo.0 .0, &qinfo.0 as *const _ as *mut _,
flags, flags,
dnssec as i32, dnssec as i32,
want_dnssec.into(), want_dnssec.into(),
nocaps.into(), nocaps.into(),
check_ratelimit.into(), check_ratelimit.into(),
addr.cast_mut(), addr as *mut _,
addr_len as u32, addr_len as u32,
zone.as_ptr().cast_mut(), zone.as_ptr() as *mut _,
zone.len(), zone.len(),
tcp_upstream.into(), tcp_upstream.into(),
ssl_upstream.into(), ssl_upstream.into(),
tls_auth_name.map_or_else(ptr::null_mut, |x| x.as_ptr().cast_mut()), tls_auth_name
.map(|x| x.as_ptr() as *mut _)
.unwrap_or(ptr::null_mut()),
q.0, q.0,
std::ptr::addr_of_mut!(was_ratelimited), &mut was_ratelimited as *mut _,
) )
}; };
if ret.is_null() { if ret.is_null() {
@ -214,7 +192,7 @@ impl<T> ModuleEnvMut<T> {
let res = unsafe { let res = unsafe {
((*self.0).attach_sub.unwrap_unchecked())( ((*self.0).attach_sub.unwrap_unchecked())(
qstate.0, qstate.0,
qinfo.0 .0, &qinfo.0 as *const _ as *mut _,
qflags, qflags,
prime.into(), prime.into(),
valrec.into(), valrec.into(),
@ -241,35 +219,30 @@ impl<T> ModuleEnvMut<T> {
} }
impl<T> ModuleQstate<'_, T> { impl<T> ModuleQstate<'_, T> {
pub fn qinfo(&self) -> QueryInfo<'_> { pub fn qinfo_mut(&mut self) -> QueryInfoMut<'_> {
QueryInfo( QueryInfoMut(
unsafe { std::ptr::addr_of_mut!((*self.0).qinfo) }, unsafe { &mut (*self.0).qinfo as *mut query_info },
Default::default(), Default::default(),
) )
} }
pub fn return_msg(&self) -> Option<DnsMsg<'_>> { pub fn return_msg_mut(&mut self) -> Option<DnsMsgMut<'_>> {
if unsafe { (*self.0).return_msg.is_null() } { if unsafe { (*self.0).return_msg.is_null() } {
None None
} else { } else {
Some(DnsMsg(unsafe { (*self.0).return_msg }, Default::default())) Some(DnsMsgMut(
unsafe { (*self.0).return_msg },
Default::default(),
))
} }
} }
}
impl<T> ModuleQstateMut<'_, T> {
pub fn qinfo_mut(&mut self) -> QueryInfoMut<'_> {
QueryInfoMut(self.qinfo())
}
pub fn return_msg_mut(&mut self) -> Option<DnsMsgMut<'_>> {
self.return_msg().map(DnsMsgMut)
}
pub fn set_ext_state(&mut self, state: ModuleExtState) { pub fn set_ext_state(&mut self, state: ModuleExtState) {
unsafe { unsafe {
(*self.0 .0).ext_state[self.1 as usize] = state as module_ext_state; (*self.0).ext_state[self.1 as usize] = state as module_ext_state;
} }
} }
} }
impl DnsMsg<'_> { impl DnsMsgMut<'_> {
pub fn rep(&self) -> Option<ReplyInfo<'_>> { pub fn rep(&self) -> Option<ReplyInfo<'_>> {
if unsafe { (*self.0).rep.is_null() } { if unsafe { (*self.0).rep.is_null() } {
None None
@ -345,7 +318,7 @@ impl ReplyInfo<'_> {
impl UbPackedRrsetKey<'_> { impl UbPackedRrsetKey<'_> {
pub fn entry(&self) -> LruHashEntry<'_> { pub fn entry(&self) -> LruHashEntry<'_> {
LruHashEntry( LruHashEntry(
unsafe { std::ptr::addr_of_mut!((*self.0).entry) }, unsafe { &mut (*self.0).entry as *mut _ },
Default::default(), Default::default(),
) )
} }
@ -353,10 +326,7 @@ impl UbPackedRrsetKey<'_> {
unsafe { (*self.0).id } unsafe { (*self.0).id }
} }
pub fn rk(&self) -> PackedRrsetKey<'_> { pub fn rk(&self) -> PackedRrsetKey<'_> {
PackedRrsetKey( PackedRrsetKey(unsafe { &mut (*self.0).rk as *mut _ }, Default::default())
unsafe { std::ptr::addr_of_mut!((*self.0).rk) },
Default::default(),
)
} }
} }
@ -382,7 +352,7 @@ impl PackedRrsetKey<'_> {
impl LruHashEntry<'_> { impl LruHashEntry<'_> {
pub fn data(&self) -> PackedRrsetData<'_> { pub fn data(&self) -> PackedRrsetData<'_> {
// FIXME: shouldnt pthread lock be used here? // FIXME: shouldnt pthread lock be used here?
unsafe { PackedRrsetData((*self.0).data.cast(), Default::default()) } unsafe { PackedRrsetData((*self.0).data as *mut packed_rrset_data, Default::default()) }
} }
} }
@ -640,8 +610,8 @@ pub enum ModuleExtState {
} }
impl ModuleExtState { impl ModuleExtState {
pub(crate) const fn importance(self) -> usize { pub(crate) fn importance(&self) -> usize {
match self { match *self {
Self::Unknown => 0, Self::Unknown => 0,
Self::InitialState => 1, Self::InitialState => 1,
Self::Finished => 2, Self::Finished => 2,