Compare commits
2 commits
e9a6f296df
...
be0359f85f
Author | SHA1 | Date | |
---|---|---|---|
chayleaf | be0359f85f | ||
chayleaf | 422001fc71 |
|
@ -14,20 +14,20 @@ macro_rules! impl_tuple {
|
|||
{
|
||||
type EnvData = A::EnvData;
|
||||
type QstateData = A::QstateData;
|
||||
fn init(env: &mut crate::unbound::ModuleEnv<Self::EnvData>) -> Result<Self, ()> {
|
||||
fn init(env: &mut crate::unbound::ModuleEnvMut<Self::EnvData>) -> Result<Self, ()> {
|
||||
Ok((A::init(env)?, $($t::init(env)?, )*))
|
||||
}
|
||||
fn clear(&self, qstate: &mut crate::unbound::ModuleQstate<Self::QstateData>) {
|
||||
fn clear(&self, qstate: &mut crate::unbound::ModuleQstateMut<Self::QstateData>) {
|
||||
self.0.clear(qstate);
|
||||
$(self.$i.clear(qstate);)*
|
||||
}
|
||||
fn deinit(self, env: &mut crate::unbound::ModuleEnv<Self::EnvData>) {
|
||||
fn deinit(self, env: &mut crate::unbound::ModuleEnvMut<Self::EnvData>) {
|
||||
self.0.deinit(env);
|
||||
$(self.$i.deinit(env);)*
|
||||
}
|
||||
fn operate(
|
||||
&self,
|
||||
qstate: &mut crate::unbound::ModuleQstate<Self::QstateData>,
|
||||
qstate: &mut crate::unbound::ModuleQstateMut<Self::QstateData>,
|
||||
event: crate::unbound::ModuleEvent,
|
||||
entry: &mut crate::unbound::OutboundEntryMut,
|
||||
) -> Option<ModuleExtState> {
|
||||
|
@ -40,13 +40,13 @@ macro_rules! impl_tuple {
|
|||
})*
|
||||
ret
|
||||
}
|
||||
fn get_mem(&self, env: &mut crate::unbound::ModuleEnv<Self::EnvData>) -> usize {
|
||||
fn get_mem(&self, env: &mut crate::unbound::ModuleEnvMut<Self::EnvData>) -> usize {
|
||||
self.0.get_mem(env) $(* self.$i.get_mem(env))*
|
||||
}
|
||||
fn inform_super(
|
||||
&self,
|
||||
qstate: &mut crate::unbound::ModuleQstate<Self::QstateData>,
|
||||
super_qstate: &mut crate::unbound::ModuleQstate<std::ffi::c_void>,
|
||||
qstate: &mut crate::unbound::ModuleQstateMut<Self::QstateData>,
|
||||
super_qstate: &mut crate::unbound::ModuleQstateMut<std::ffi::c_void>,
|
||||
) {
|
||||
self.0.inform_super(qstate, super_qstate);
|
||||
$(self.$i.inform_super(qstate, super_qstate);)*
|
||||
|
|
|
@ -1,4 +1,4 @@
|
|||
use std::{collections::HashMap, hash::Hash};
|
||||
use std::{borrow::Borrow, collections::HashMap, hash::Hash};
|
||||
|
||||
use smallvec::{smallvec, SmallVec};
|
||||
|
||||
|
@ -36,9 +36,10 @@ impl<T: Hash + Eq> PrefixSet<T> {
|
|||
}
|
||||
}
|
||||
}
|
||||
pub fn contains<'a>(&self, val: impl IntoIterator<Item = &'a T>) -> bool
|
||||
pub fn contains<'a, Y>(&self, val: impl IntoIterator<Item = &'a Y>) -> bool
|
||||
where
|
||||
T: 'a,
|
||||
T: 'a + Borrow<Y>,
|
||||
Y: 'a + ?Sized + Eq + Hash,
|
||||
{
|
||||
match self {
|
||||
Self::Leaf => true,
|
||||
|
|
332
src/example.rs
332
src/example.rs
|
@ -99,8 +99,26 @@ struct IpCache<T>(
|
|||
);
|
||||
|
||||
#[repr(transparent)]
|
||||
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
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 {
|
||||
fn encode_bytes(&self) -> Vec<u8> {
|
||||
self.0.to_vec()
|
||||
|
@ -118,8 +136,11 @@ impl<T> Default for IpCache<T> {
|
|||
fn ignore<T>(_: &mut smallvec::SmallVec<[T; 4]>) {}
|
||||
|
||||
impl<T> IpCache<T> {
|
||||
fn extend_set_with_domain<J: IpPrefix + From<T>>(&self, ips: &mut RTrieSet<J>, domain_r: Domain)
|
||||
where
|
||||
fn extend_set_with_domain<J: IpPrefix + From<T>>(
|
||||
&self,
|
||||
ips: &mut RTrieSet<J>,
|
||||
domain_r: IpCacheKey,
|
||||
) where
|
||||
T: Copy,
|
||||
{
|
||||
self.get_maybe_update_rev(domain_r, |val| {
|
||||
|
@ -130,15 +151,14 @@ impl<T> IpCache<T> {
|
|||
let mut val = Some(ignore);
|
||||
val = None;
|
||||
val
|
||||
})
|
||||
});
|
||||
}
|
||||
fn get_maybe_update_rev<F: for<'a> FnOnce(&'a mut smallvec::SmallVec<[T; 4]>)>(
|
||||
&self,
|
||||
domain_r: Domain,
|
||||
domain_r: IpCacheKey,
|
||||
upd: impl FnOnce(Option<(&smallvec::SmallVec<[T; 4]>, &Mutex<()>, &AtomicBool)>) -> Option<F>,
|
||||
) {
|
||||
let lock = self.0.read().unwrap();
|
||||
let domain_r = IpCacheKey(domain_r);
|
||||
let key = lock.0.get(&domain_r).copied();
|
||||
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)))
|
||||
|
@ -171,7 +191,7 @@ impl<T: ToString + PartialEq> IpCache<T> {
|
|||
let ret1 = &mut ret;
|
||||
let mut path = self.1.clone();
|
||||
path.push(domain);
|
||||
self.get_maybe_update_rev(domain_r.0, |ips| {
|
||||
self.get_maybe_update_rev(domain_r, |ips| {
|
||||
if let Some(ips) = ips.as_ref().filter(|x| x.0 == &val) {
|
||||
*ret1 = false;
|
||||
if ips
|
||||
|
@ -216,7 +236,7 @@ impl<T: FromStr> IpCache<T> {
|
|||
let mut lock = self.0.write().unwrap();
|
||||
assert!(lock.1.is_empty());
|
||||
let domains = std::fs::read_dir(dir)?;
|
||||
for entry in domains.filter_map(|x| x.ok()) {
|
||||
for entry in domains.filter_map(Result::ok) {
|
||||
let domain = entry.file_name();
|
||||
let Some(domain) = domain.to_str() else {
|
||||
continue;
|
||||
|
@ -234,15 +254,7 @@ impl<T: FromStr> IpCache<T> {
|
|||
let Ok(reader) = std::fs::File::open(entry.path()) else {
|
||||
continue;
|
||||
};
|
||||
let domain_r = IpCacheKey(
|
||||
domain
|
||||
.split('.')
|
||||
.rev()
|
||||
.map(|x| x.as_bytes())
|
||||
.collect::<Vec<_>>()
|
||||
.join(&b"."[..])
|
||||
.into(),
|
||||
);
|
||||
let domain_r = IpCacheKey::from_split_domain(domain.split('.'));
|
||||
let mut reader = BufReader::new(reader);
|
||||
let mut line = String::new();
|
||||
let mut ips = SmallVec::new();
|
||||
|
@ -280,7 +292,7 @@ pub(crate) const DATA_PREFIX: &str = "unbound-mod-test-data";
|
|||
pub(crate) const CONFIG_PREFIX: &str = "unbound-mod-test-config";
|
||||
|
||||
#[cfg(not(debug_assertions))]
|
||||
pub(crate) const PATH_PREFIX: &str = "/var/lib/unbound";
|
||||
pub(crate) const DATA_PREFIX: &str = "/var/lib/unbound";
|
||||
#[cfg(not(debug_assertions))]
|
||||
pub(crate) const CONFIG_PREFIX: &str = "/etc/unbound";
|
||||
|
||||
|
@ -312,7 +324,7 @@ impl ExampleMod {
|
|||
Ok(ret)
|
||||
}
|
||||
fn load_json(&mut self, rulesets: &mut [(NftData<Ipv4Net>, NftData<Ipv6Net>)]) {
|
||||
for (k, v) in self.nft_queries.iter_mut() {
|
||||
for (k, v) in &mut self.nft_queries {
|
||||
let r = &mut rulesets[v.index];
|
||||
let mut v_domains = v.domains.write().unwrap();
|
||||
for base in [CONFIG_PREFIX, DATA_PREFIX] {
|
||||
|
@ -321,13 +333,8 @@ impl ExampleMod {
|
|||
match read_json::<Vec<String>>(file) {
|
||||
Ok(domains) => {
|
||||
for domain in domains {
|
||||
v_domains.insert(
|
||||
domain
|
||||
.split('.')
|
||||
.rev()
|
||||
.map(|x| x.as_bytes().into())
|
||||
.collect::<SmallVec<[DomainSeg; 5]>>(),
|
||||
);
|
||||
v_domains
|
||||
.insert(domain.split('.').rev().map(|x| x.as_bytes().into()));
|
||||
}
|
||||
}
|
||||
Err(err) => Self::report2(&self.error_lock, "domains", err),
|
||||
|
@ -338,13 +345,8 @@ impl ExampleMod {
|
|||
match read_json::<Vec<DpiInfo>>(file) {
|
||||
Ok(dpi_info) => {
|
||||
for domain in dpi_info.iter().flat_map(|x| &x.domains) {
|
||||
v_domains.insert(
|
||||
domain
|
||||
.split('.')
|
||||
.rev()
|
||||
.map(|x| x.as_bytes().into())
|
||||
.collect::<SmallVec<[DomainSeg; 5]>>(),
|
||||
);
|
||||
v_domains
|
||||
.insert(domain.split('.').rev().map(|x| x.as_bytes().into()));
|
||||
}
|
||||
}
|
||||
Err(err) => Self::report2(&self.error_lock, "dpi", err),
|
||||
|
@ -375,11 +377,7 @@ impl ExampleMod {
|
|||
}
|
||||
println!("loading cached domain ips for {k}");
|
||||
for rev_domain in v_domains.iter() {
|
||||
let rev_domain: SmallVec<_> = rev_domain
|
||||
.map(|x| x.as_slice())
|
||||
.collect::<Vec<_>>()
|
||||
.join(&b"."[..])
|
||||
.into();
|
||||
let rev_domain = IpCacheKey::from_split_rev_domain(rev_domain.into_iter());
|
||||
self.caches
|
||||
.0
|
||||
.extend_set_with_domain(r.0.ips_mut(), rev_domain.clone());
|
||||
|
@ -401,7 +399,7 @@ impl ExampleMod {
|
|||
if let Some(s) = std::env::var_os("NFT_QUERIES") {
|
||||
for (i, (name, set4, set6)) in s
|
||||
.to_str()
|
||||
.map(|x| x.to_owned())
|
||||
.map(ToOwned::to_owned)
|
||||
.ok_or(())?
|
||||
.split(';')
|
||||
.filter_map(|x| x.split_once(':'))
|
||||
|
@ -410,11 +408,9 @@ impl ExampleMod {
|
|||
})
|
||||
.enumerate()
|
||||
{
|
||||
let (name, dynamic) = if let Some(name) = name.strip_suffix('!') {
|
||||
(name, true)
|
||||
} else {
|
||||
(name, false)
|
||||
};
|
||||
let (name, dynamic) = name
|
||||
.strip_suffix('!')
|
||||
.map_or((name, false), |name| (name, true));
|
||||
self.nft_queries.insert(
|
||||
name.to_owned(),
|
||||
NftQuery {
|
||||
|
@ -456,7 +452,7 @@ impl ExampleMod {
|
|||
}
|
||||
fn handle_reply_info(
|
||||
&self,
|
||||
split_rev_domain: SmallVec<[DomainSeg; 5]>,
|
||||
split_domain: &[&[u8]],
|
||||
qnames: SmallVec<[usize; 5]>,
|
||||
rep: &ReplyInfo<'_>,
|
||||
) -> Result<(), ()> {
|
||||
|
@ -485,91 +481,80 @@ impl ExampleMod {
|
|||
}
|
||||
}
|
||||
}
|
||||
self.add_ips(ip4, ip6, split_rev_domain, qnames)
|
||||
self.add_ips(ip4, ip6, split_domain, qnames)
|
||||
}
|
||||
fn add_ips(
|
||||
&self,
|
||||
ip4: SmallVec<[Ipv4Addr; 4]>,
|
||||
ip6: SmallVec<[Ipv6Addr; 4]>,
|
||||
split_rev_domain: SmallVec<[DomainSeg; 5]>,
|
||||
split_domain: &[&[u8]],
|
||||
qnames: SmallVec<[usize; 5]>,
|
||||
) -> Result<(), ()> {
|
||||
println!("adding {ip4:?}/{ip6:?} for {split_domain:?} to {qnames:?}");
|
||||
if !ip4.is_empty() || !ip6.is_empty() {
|
||||
let domain = match split_rev_domain
|
||||
let domain = match split_domain
|
||||
.iter()
|
||||
.rev()
|
||||
.map(|x| String::from_utf8(x.to_vec()).map(|x| x + "."))
|
||||
.collect::<Result<String, _>>()
|
||||
{
|
||||
Ok(mut x) => {
|
||||
x.pop();
|
||||
x
|
||||
}
|
||||
.copied()
|
||||
.map(std::str::from_utf8)
|
||||
.try_fold(String::new(), |mut s, comp| {
|
||||
if !s.is_empty() {
|
||||
s.push('.');
|
||||
}
|
||||
s.push_str(comp?);
|
||||
Ok::<_, std::str::Utf8Error>(s)
|
||||
}) {
|
||||
Ok(x) => x,
|
||||
Err(err) => {
|
||||
self.report("domain utf-8", err);
|
||||
return Err(());
|
||||
}
|
||||
};
|
||||
let mut split_rev_domain = split_rev_domain.into_iter();
|
||||
if let Some(first) = split_rev_domain.next() {
|
||||
let first: Domain = first.to_vec().into();
|
||||
let joined_rev_domain = split_rev_domain.fold(first, |mut res, mut next| {
|
||||
res.push(b'.');
|
||||
res.append(&mut next);
|
||||
res
|
||||
});
|
||||
let mut to_send: SmallVec<[IpNet; 8]> = SmallVec::new();
|
||||
to_send.extend(ip4.iter().copied().map(Ipv4Net::from).map(IpNet::from));
|
||||
to_send.extend(ip6.iter().copied().map(Ipv6Net::from).map(IpNet::from));
|
||||
let keep4 = !ip4.is_empty()
|
||||
&& 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();
|
||||
}
|
||||
let key = IpCacheKey::from_split_domain(split_domain.iter());
|
||||
let mut to_send: SmallVec<[IpNet; 8]> = SmallVec::new();
|
||||
to_send.extend(ip4.iter().copied().map(Ipv4Net::from).map(IpNet::from));
|
||||
to_send.extend(ip6.iter().copied().map(Ipv6Net::from).map(IpNet::from));
|
||||
let keep4 = !ip4.is_empty() && self.caches.0.set(&domain, key.clone(), ip4);
|
||||
let keep6 = !ip6.is_empty() && self.caches.1.set(&domain, key, 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(())
|
||||
}
|
||||
fn run_commands(&self, rev_domain: &[u8]) -> Option<ModuleExtState> {
|
||||
if let Some(rev_domain) = self
|
||||
.nft_token
|
||||
.as_ref()
|
||||
.and_then(|token| rev_domain.strip_prefix(token.as_bytes()))
|
||||
{
|
||||
for (qname, query) in self.nft_queries.iter() {
|
||||
if query.dynamic && rev_domain.starts_with(qname.as_bytes()) {
|
||||
if let Some(rev_domain) =
|
||||
rev_domain.strip_prefix((qname.to_owned() + ".").as_bytes())
|
||||
fn run_commands(&self, split_domain: &[&[u8]]) -> Option<ModuleExtState> {
|
||||
if let Some(split_domain) = self.nft_token.as_ref().and_then(|token| {
|
||||
split_domain
|
||||
.split_last()
|
||||
.filter(|(a, _)| **a == token.as_bytes())
|
||||
.map(|(_, b)| b)
|
||||
}) {
|
||||
for (qname, query) in &self.nft_queries {
|
||||
if query.dynamic {
|
||||
if let Some(split_domain) = split_domain
|
||||
.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();
|
||||
if domains.insert(rev_domain.clone()) {
|
||||
if domains.insert(split_domain.iter().copied().rev().map(From::from)) {
|
||||
drop(domains);
|
||||
let file_name = format!("{DATA_PREFIX}/{qname}_domains.json");
|
||||
let domain = match String::from_utf8(
|
||||
rev_domain
|
||||
.iter()
|
||||
.rev()
|
||||
.map(|x| x.as_slice())
|
||||
.collect::<Vec<_>>()
|
||||
.join(&b"."[..]),
|
||||
) {
|
||||
let domain = match split_domain
|
||||
.iter()
|
||||
.copied()
|
||||
.map(std::str::from_utf8)
|
||||
.try_fold(String::new(), |mut s, comp| {
|
||||
if !s.is_empty() {
|
||||
s.push('.');
|
||||
}
|
||||
s.push_str(comp?);
|
||||
Ok::<_, std::str::Utf8Error>(s)
|
||||
}) {
|
||||
Ok(x) => x,
|
||||
Err(err) => {
|
||||
self.report("domain utf-8", err);
|
||||
|
@ -577,6 +562,7 @@ impl ExampleMod {
|
|||
}
|
||||
};
|
||||
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) {
|
||||
match read_json(file) {
|
||||
Ok(x) => x,
|
||||
|
@ -602,22 +588,21 @@ impl ExampleMod {
|
|||
}
|
||||
}
|
||||
return Some(ModuleExtState::Finished);
|
||||
} else if let Some(rev_domain) = self
|
||||
.tmp_nft_token
|
||||
.as_ref()
|
||||
.and_then(|token| rev_domain.strip_prefix(token.as_bytes()))
|
||||
{
|
||||
for (qname, query) in self.nft_queries.iter() {
|
||||
if query.dynamic && rev_domain.starts_with(qname.as_bytes()) {
|
||||
if let Some(rev_domain) =
|
||||
rev_domain.strip_prefix((qname.to_owned() + ".").as_bytes())
|
||||
} else if let Some(split_domain) = self.tmp_nft_token.as_ref().and_then(|token| {
|
||||
split_domain
|
||||
.split_last()
|
||||
.filter(|(a, _)| **a == token.as_bytes())
|
||||
.map(|(_, b)| b)
|
||||
}) {
|
||||
for (qname, query) in &self.nft_queries {
|
||||
if query.dynamic {
|
||||
if let Some(split_domain) = split_domain
|
||||
.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();
|
||||
domains.insert(rev_domain.clone());
|
||||
domains.insert(split_domain.iter().copied().rev().map(From::from));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
@ -625,10 +610,15 @@ impl ExampleMod {
|
|||
}
|
||||
None
|
||||
}
|
||||
fn get_qnames(&self, split_rev_domain: &SmallVec<[DomainSeg; 5]>) -> SmallVec<[usize; 5]> {
|
||||
fn get_qnames(&self, split_domain: &[&[u8]]) -> SmallVec<[usize; 5]> {
|
||||
let mut qnames: SmallVec<[usize; 5]> = SmallVec::new();
|
||||
for query in self.nft_queries.values() {
|
||||
if query.domains.read().unwrap().contains(split_rev_domain) {
|
||||
if query
|
||||
.domains
|
||||
.read()
|
||||
.unwrap()
|
||||
.contains(split_domain.iter().copied().rev().map(From::from))
|
||||
{
|
||||
qnames.push(query.index);
|
||||
}
|
||||
}
|
||||
|
@ -680,17 +670,31 @@ fn read_json<T: 'static + for<'a> Deserialize<'a>>(mut f: File) -> Result<T, ser
|
|||
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 {
|
||||
type EnvData = ();
|
||||
type QstateData = ();
|
||||
|
||||
fn init(_env: &mut crate::unbound::ModuleEnv<Self::EnvData>) -> Result<Self, ()> {
|
||||
fn init(_env: &mut crate::unbound::ModuleEnvMut<Self::EnvData>) -> Result<Self, ()> {
|
||||
Self::new()
|
||||
}
|
||||
|
||||
fn operate(
|
||||
&self,
|
||||
qstate: &mut crate::unbound::ModuleQstate<Self::QstateData>,
|
||||
qstate: &mut crate::unbound::ModuleQstateMut<Self::QstateData>,
|
||||
event: ModuleEvent,
|
||||
_entry: &mut crate::unbound::OutboundEntryMut,
|
||||
) -> Option<ModuleExtState> {
|
||||
|
@ -703,26 +707,21 @@ impl UnboundMod for ExampleMod {
|
|||
return Some(ModuleExtState::Error);
|
||||
}
|
||||
}
|
||||
let info = qstate.qinfo_mut();
|
||||
let info = qstate.qinfo();
|
||||
let name = info.qname().to_bytes();
|
||||
let rev_domain = name.strip_suffix(b".").unwrap_or(name);
|
||||
if let Some(val) = self.run_commands(rev_domain) {
|
||||
// let rev_domain = name.strip_suffix(b".").unwrap_or(name);
|
||||
let split_domain = unwire_domain(name);
|
||||
println!("handling {split_domain:?}");
|
||||
if let Some(val) = self.run_commands(&split_domain) {
|
||||
return Some(val);
|
||||
}
|
||||
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);
|
||||
let qnames = self.get_qnames(&split_domain);
|
||||
if qnames.is_empty() {
|
||||
return Some(ModuleExtState::Finished);
|
||||
}
|
||||
if let Some(ret) = qstate.return_msg_mut() {
|
||||
if let Some(ret) = qstate.return_msg() {
|
||||
if let Some(rep) = ret.rep() {
|
||||
if self
|
||||
.handle_reply_info(split_rev_domain, qnames, &rep)
|
||||
.is_err()
|
||||
{
|
||||
if self.handle_reply_info(&split_domain, qnames, &rep).is_err() {
|
||||
return Some(ModuleExtState::Error);
|
||||
}
|
||||
}
|
||||
|
@ -741,9 +740,9 @@ mod test {
|
|||
use std::{net::Ipv4Addr, os::unix::fs::MetadataExt, path::PathBuf, str::FromStr, sync::mpsc};
|
||||
|
||||
use ipnet::IpNet;
|
||||
use smallvec::{smallvec, SmallVec};
|
||||
use smallvec::smallvec;
|
||||
|
||||
use crate::example::{ignore, ExampleMod, IpNetDeser, DATA_PREFIX};
|
||||
use crate::example::{ignore, ExampleMod, IpCacheKey, IpNetDeser, DATA_PREFIX};
|
||||
|
||||
#[test]
|
||||
fn test() {
|
||||
|
@ -794,24 +793,26 @@ mod test {
|
|||
base_path.push("domains6");
|
||||
t.caches.1.load(&base_path).unwrap();
|
||||
|
||||
t.caches
|
||||
.0
|
||||
.get_maybe_update_rev("com.a".as_bytes().into(), |x| {
|
||||
t.caches.0.get_maybe_update_rev(
|
||||
IpCacheKey::from_split_domain(["a", "com"].into_iter()),
|
||||
|x| {
|
||||
assert!(x.unwrap().0.len() == 2);
|
||||
#[allow(unused_assignments)]
|
||||
let mut val = Some(ignore);
|
||||
val = None;
|
||||
val
|
||||
});
|
||||
t.caches
|
||||
.0
|
||||
.get_maybe_update_rev("com.b".as_bytes().into(), |x| {
|
||||
},
|
||||
);
|
||||
t.caches.0.get_maybe_update_rev(
|
||||
IpCacheKey::from_split_domain(["b", "com"].into_iter()),
|
||||
|x| {
|
||||
assert!(x.unwrap().0.len() == 1);
|
||||
#[allow(unused_assignments)]
|
||||
let mut val = Some(ignore);
|
||||
val = None;
|
||||
val
|
||||
});
|
||||
},
|
||||
);
|
||||
|
||||
t.load_json(&mut rulesets);
|
||||
|
||||
|
@ -838,48 +839,53 @@ mod test {
|
|||
tx2.send(rulesets).unwrap();
|
||||
});
|
||||
|
||||
let split_rev_domain = smallvec![SmallVec::from(&b"com"[..]), SmallVec::from(&b"c"[..])];
|
||||
let qnames = t.get_qnames(&split_rev_domain);
|
||||
let split_domain = [&b"c"[..], &b"com"[..]];
|
||||
let qnames = t.get_qnames(&split_domain);
|
||||
assert_eq!(qnames.len(), 2);
|
||||
t.add_ips(
|
||||
smallvec![Ipv4Addr::new(7, 7, 7, 7), Ipv4Addr::new(6, 6, 6, 6)],
|
||||
smallvec![],
|
||||
split_rev_domain,
|
||||
&split_domain,
|
||||
qnames,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let split_rev_domain = smallvec![SmallVec::from(&b"com"[..]), SmallVec::from(&b"a"[..])];
|
||||
let qnames = t.get_qnames(&split_rev_domain);
|
||||
let split_domain = [&b"a"[..], &b"com"[..]];
|
||||
let qnames = t.get_qnames(&split_domain);
|
||||
t.add_ips(
|
||||
smallvec![Ipv4Addr::new(1, 2, 3, 4), Ipv4Addr::new(5, 6, 7, 8)],
|
||||
smallvec![],
|
||||
split_rev_domain,
|
||||
&split_domain,
|
||||
qnames,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
t.run_commands(b"token.q.com.w").unwrap();
|
||||
t.run_commands(b"tmptoken.q.com.e").unwrap();
|
||||
t.run_commands(&[&b"w"[..], &b"com"[..], &b"q"[..], &b"token"[..]])
|
||||
.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_rev_domain = smallvec![SmallVec::from(&b"com"[..]), SmallVec::from(&b"e"[..])];
|
||||
let qnames = t.get_qnames(&split_rev_domain);
|
||||
let split_domain = [&b"e"[..], &b"com"[..]];
|
||||
let qnames = t.get_qnames(&split_domain);
|
||||
assert_eq!(qnames.len(), 1);
|
||||
t.add_ips(
|
||||
smallvec![Ipv4Addr::new(8, 8, 8, 8)],
|
||||
smallvec![],
|
||||
split_rev_domain,
|
||||
&split_domain,
|
||||
qnames,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let split_rev_domain = smallvec![SmallVec::from(&b"com"[..]), SmallVec::from(&b"w"[..])];
|
||||
let qnames = t.get_qnames(&split_rev_domain);
|
||||
let split_domain = [&b"w"[..], &b"com"[..]];
|
||||
let qnames = t.get_qnames(&split_domain);
|
||||
assert_eq!(qnames.len(), 1);
|
||||
t.add_ips(
|
||||
smallvec![Ipv4Addr::new(9, 8, 8, 8)],
|
||||
smallvec![],
|
||||
split_rev_domain,
|
||||
&split_domain,
|
||||
qnames,
|
||||
)
|
||||
.unwrap();
|
||||
|
|
|
@ -16,11 +16,7 @@ pub unsafe extern "C" fn init(
|
|||
env: *mut module_env,
|
||||
id: ::std::os::raw::c_int,
|
||||
) -> ::std::os::raw::c_int {
|
||||
if let Some(fac) = crate::MODULE_FACTORY.take() {
|
||||
fac(env, id)
|
||||
} else {
|
||||
0
|
||||
}
|
||||
crate::MODULE_FACTORY.take().map_or(0, |fac| fac(env, id))
|
||||
}
|
||||
|
||||
/// Deinitialize module internals.
|
||||
|
@ -33,7 +29,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.
|
||||
/// 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.
|
||||
///
|
||||
/// # Arguments
|
||||
|
@ -50,7 +46,7 @@ pub unsafe extern "C" fn operate(
|
|||
entry: *mut outbound_entry,
|
||||
) {
|
||||
if let Some(module) = crate::module() {
|
||||
module.internal_operate(qstate, event, id, entry)
|
||||
module.internal_operate(qstate, event, id, entry);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -69,7 +65,7 @@ pub unsafe extern "C" fn inform_super(
|
|||
super_qstate: *mut module_qstate,
|
||||
) {
|
||||
if let Some(module) = crate::module() {
|
||||
module.internal_inform_super(qstate, id, super_qstate)
|
||||
module.internal_inform_super(qstate, id, super_qstate);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -78,7 +74,7 @@ pub unsafe extern "C" fn inform_super(
|
|||
#[no_mangle]
|
||||
pub unsafe extern "C" fn clear(qstate: *mut module_qstate, id: ::std::os::raw::c_int) {
|
||||
if let Some(module) = crate::module() {
|
||||
module.internal_clear(qstate, id)
|
||||
module.internal_clear(qstate, id);
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -86,9 +82,7 @@ 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.
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn get_mem(env: *mut module_env, id: ::std::os::raw::c_int) -> usize {
|
||||
crate::module()
|
||||
.map(|module| module.internal_get_mem(env, id))
|
||||
.unwrap_or(0)
|
||||
crate::module().map_or(0, |module| module.internal_get_mem(env, id))
|
||||
}
|
||||
|
||||
// function interface assertions
|
||||
|
|
84
src/lib.rs
84
src/lib.rs
|
@ -9,10 +9,9 @@ use unbound::ModuleExtState;
|
|||
non_snake_case,
|
||||
non_upper_case_globals,
|
||||
unused_imports,
|
||||
clippy::useless_transmute,
|
||||
clippy::type_complexity,
|
||||
clippy::too_many_arguments,
|
||||
clippy::upper_case_acronyms
|
||||
clippy::all,
|
||||
clippy::nursery,
|
||||
clippy::pedantic
|
||||
)]
|
||||
mod bindings;
|
||||
mod combine;
|
||||
|
@ -23,21 +22,17 @@ mod exports;
|
|||
mod nftables;
|
||||
mod unbound;
|
||||
|
||||
pub fn add(left: usize, right: usize) -> usize {
|
||||
left + right
|
||||
}
|
||||
|
||||
pub trait UnboundMod: Send + Sync + Sized + RefUnwindSafe + UnwindSafe {
|
||||
type EnvData;
|
||||
type QstateData;
|
||||
#[allow(clippy::result_unit_err)]
|
||||
fn init(_env: &mut unbound::ModuleEnv<Self::EnvData>) -> Result<Self, ()> {
|
||||
fn init(_env: &mut unbound::ModuleEnvMut<Self::EnvData>) -> Result<Self, ()> {
|
||||
Err(())
|
||||
}
|
||||
fn deinit(self, _env: &mut unbound::ModuleEnv<Self::EnvData>) {}
|
||||
fn deinit(self, _env: &mut unbound::ModuleEnvMut<Self::EnvData>) {}
|
||||
fn operate(
|
||||
&self,
|
||||
_qstate: &mut unbound::ModuleQstate<Self::QstateData>,
|
||||
_qstate: &mut unbound::ModuleQstateMut<Self::QstateData>,
|
||||
_event: unbound::ModuleEvent,
|
||||
_entry: &mut unbound::OutboundEntryMut,
|
||||
) -> Option<ModuleExtState> {
|
||||
|
@ -45,13 +40,13 @@ pub trait UnboundMod: Send + Sync + Sized + RefUnwindSafe + UnwindSafe {
|
|||
}
|
||||
fn inform_super(
|
||||
&self,
|
||||
_qstate: &mut unbound::ModuleQstate<Self::QstateData>,
|
||||
_super_qstate: &mut unbound::ModuleQstate<::std::ffi::c_void>,
|
||||
_qstate: &mut unbound::ModuleQstateMut<Self::QstateData>,
|
||||
_super_qstate: &mut unbound::ModuleQstateMut<::std::ffi::c_void>,
|
||||
) {
|
||||
}
|
||||
fn clear(&self, _qstate: &mut unbound::ModuleQstate<Self::QstateData>) {}
|
||||
fn clear(&self, _qstate: &mut unbound::ModuleQstateMut<Self::QstateData>) {}
|
||||
|
||||
fn get_mem(&self, _env: &mut unbound::ModuleEnv<Self::EnvData>) -> usize {
|
||||
fn get_mem(&self, _env: &mut unbound::ModuleEnvMut<Self::EnvData>) -> usize {
|
||||
0
|
||||
}
|
||||
}
|
||||
|
@ -97,7 +92,7 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
|
|||
id: ::std::os::raw::c_int,
|
||||
) {
|
||||
std::panic::catch_unwind(|| {
|
||||
self.deinit(&mut unbound::ModuleEnv(env, id, Default::default()))
|
||||
self.deinit(&mut unbound::ModuleEnvMut(env, id, Default::default()));
|
||||
})
|
||||
.unwrap_or(());
|
||||
}
|
||||
|
@ -109,13 +104,16 @@ 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,
|
||||
&mut unbound::ModuleQstateMut(unbound::ModuleQstate(
|
||||
qstate,
|
||||
id,
|
||||
Default::default(),
|
||||
)),
|
||||
event.into(),
|
||||
&mut unbound::OutboundEntryMut(entry, Default::default()),
|
||||
) {
|
||||
qstate.set_ext_state(ext_state);
|
||||
(*qstate).ext_state[id as usize] = ext_state as bindings::module_ext_state;
|
||||
}
|
||||
})
|
||||
.unwrap_or(());
|
||||
|
@ -128,9 +126,17 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
|
|||
) {
|
||||
std::panic::catch_unwind(|| {
|
||||
self.inform_super(
|
||||
&mut unbound::ModuleQstate(qstate, id, Default::default()),
|
||||
&mut unbound::ModuleQstate(super_qstate, -1, Default::default()),
|
||||
)
|
||||
&mut unbound::ModuleQstateMut(unbound::ModuleQstate(
|
||||
qstate,
|
||||
id,
|
||||
Default::default(),
|
||||
)),
|
||||
&mut unbound::ModuleQstateMut(unbound::ModuleQstate(
|
||||
super_qstate,
|
||||
-1,
|
||||
Default::default(),
|
||||
)),
|
||||
);
|
||||
})
|
||||
.unwrap_or(());
|
||||
}
|
||||
|
@ -140,7 +146,11 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
|
|||
id: ::std::os::raw::c_int,
|
||||
) {
|
||||
std::panic::catch_unwind(|| {
|
||||
self.clear(&mut unbound::ModuleQstate(qstate, id, Default::default()))
|
||||
self.clear(&mut unbound::ModuleQstateMut(unbound::ModuleQstate(
|
||||
qstate,
|
||||
id,
|
||||
Default::default(),
|
||||
)));
|
||||
})
|
||||
.unwrap_or(());
|
||||
}
|
||||
|
@ -150,7 +160,7 @@ unsafe impl<T: UnboundMod> SealedUnboundMod for T {
|
|||
id: ::std::os::raw::c_int,
|
||||
) -> usize {
|
||||
std::panic::catch_unwind(|| {
|
||||
self.get_mem(&mut unbound::ModuleEnv(env, id, Default::default()))
|
||||
self.get_mem(&mut unbound::ModuleEnvMut(env, id, Default::default()))
|
||||
})
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
@ -173,14 +183,13 @@ pub fn set_unbound_mod<T: 'static + UnboundMod>() {
|
|||
MODULE_FACTORY
|
||||
.set(Box::new(|env, id| {
|
||||
std::panic::catch_unwind(|| {
|
||||
if let Ok(module) =
|
||||
T::init(&mut unbound::ModuleEnv(env, id, Default::default()))
|
||||
{
|
||||
MODULE.set(Box::new(module)).map_err(|_| ()).unwrap();
|
||||
1
|
||||
} else {
|
||||
0
|
||||
}
|
||||
T::init(&mut unbound::ModuleEnvMut(env, id, Default::default())).map_or(
|
||||
0,
|
||||
|module| {
|
||||
MODULE.set(Box::new(module)).map_err(|_| ()).unwrap();
|
||||
1
|
||||
},
|
||||
)
|
||||
})
|
||||
.unwrap_or(0)
|
||||
}))
|
||||
|
@ -188,14 +197,3 @@ pub fn set_unbound_mod<T: 'static + UnboundMod>() {
|
|||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn it_works() {
|
||||
let result = add(2, 2);
|
||||
assert_eq!(result, 4);
|
||||
}
|
||||
}
|
||||
|
|
|
@ -48,7 +48,7 @@ struct FlushSetMsg<'a> {
|
|||
unsafe impl<'a> NlMsg for FlushSetMsg<'a> {
|
||||
unsafe fn write(&self, buf: *mut std::ffi::c_void, seq: u32, _msg_type: MsgType) {
|
||||
let header = nftnl_sys::nftnl_nlmsg_build_hdr(
|
||||
buf as *mut c_char,
|
||||
buf.cast(),
|
||||
libc::NFT_MSG_DELSETELEM as u16,
|
||||
self.set.family() as u16,
|
||||
0,
|
||||
|
@ -68,9 +68,7 @@ pub struct SetElemsIter<'a> {
|
|||
impl<'a> SetElemsIter<'a> {
|
||||
fn new(set: &'a Set1) -> Self {
|
||||
let iter = unsafe { nftnl_sys::nftnl_set_elems_iter_create(set.as_mut_ptr()) };
|
||||
if iter.is_null() {
|
||||
panic!("oom");
|
||||
}
|
||||
assert!(!iter.is_null(), "oom");
|
||||
SetElemsIter {
|
||||
set,
|
||||
iter,
|
||||
|
@ -125,7 +123,7 @@ unsafe impl<'a> NlMsg for SetElemsMsg<'a> {
|
|||
MsgType::Del => (libc::NFT_MSG_DELSETELEM, 0),
|
||||
};
|
||||
let header = nftnl_sys::nftnl_nlmsg_build_hdr(
|
||||
buf as *mut c_char,
|
||||
buf.cast(),
|
||||
type_ as u16,
|
||||
self.set.family() as u16,
|
||||
flags as u16,
|
||||
|
@ -171,7 +169,7 @@ impl Set1 {
|
|||
pub fn new() -> Self {
|
||||
Self(unsafe { nftnl_sys::nftnl_set_alloc() })
|
||||
}
|
||||
pub fn as_mut_ptr(&self) -> *mut nftnl_sys::nftnl_set {
|
||||
pub const fn as_mut_ptr(&self) -> *mut nftnl_sys::nftnl_set {
|
||||
self.0
|
||||
}
|
||||
pub fn table_name(&self) -> Option<&str> {
|
||||
|
@ -195,30 +193,26 @@ impl Set1 {
|
|||
let data1_len = data1.len() as u32;
|
||||
unsafe {
|
||||
let elem = nftnl_sys::nftnl_set_elem_alloc();
|
||||
if elem.is_null() {
|
||||
panic!("oom");
|
||||
}
|
||||
assert!(!elem.is_null(), "oom");
|
||||
nftnl_sys::nftnl_set_elem_set(
|
||||
elem,
|
||||
nftnl_sys::NFTNL_SET_ELEM_KEY as u16,
|
||||
data1.as_ptr() as *const c_void,
|
||||
data1.as_ptr().cast(),
|
||||
data1_len,
|
||||
);
|
||||
nftnl_sys::nftnl_set_elem_add(self.as_mut_ptr(), elem);
|
||||
|
||||
let Some(data2) = excl_upper.map(|key| key.data()) else {
|
||||
let Some(data2) = excl_upper.map(SetKey::data) else {
|
||||
return;
|
||||
};
|
||||
let data2_len = data2.len() as u32;
|
||||
|
||||
let elem = nftnl_sys::nftnl_set_elem_alloc();
|
||||
if elem.is_null() {
|
||||
panic!("oom");
|
||||
}
|
||||
assert!(!elem.is_null(), "oom");
|
||||
nftnl_sys::nftnl_set_elem_set(
|
||||
elem,
|
||||
nftnl_sys::NFTNL_SET_ELEM_KEY as u16,
|
||||
data2.as_ptr() as *const c_void,
|
||||
data2.as_ptr().cast(),
|
||||
data2_len,
|
||||
);
|
||||
nftnl_sys::nftnl_set_elem_set_u32(
|
||||
|
@ -244,7 +238,7 @@ impl Set1 {
|
|||
count += 1;
|
||||
batch.add(&set.flush_msg(), nftnl::MsgType::Del);
|
||||
}
|
||||
for net in cidrs.into_iter() {
|
||||
for net in cidrs {
|
||||
if count + 2 > max_batch_size {
|
||||
batch.add_iter(SetElemsIter::new(&set), MsgType::Add);
|
||||
send_and_process(socket, &batch.finalize())?;
|
||||
|
@ -265,7 +259,7 @@ impl Set1 {
|
|||
send_and_process(socket, &batch.finalize())
|
||||
}
|
||||
|
||||
fn flush_msg(&self) -> FlushSetMsg<'_> {
|
||||
const fn flush_msg(&self) -> FlushSetMsg<'_> {
|
||||
FlushSetMsg { set: self }
|
||||
}
|
||||
}
|
||||
|
@ -282,7 +276,7 @@ pub fn get_sets(socket: &mnl::Socket) -> io::Result<Vec<Set1>> {
|
|||
let mut ret = Vec::new();
|
||||
unsafe {
|
||||
nftnl_sys::nftnl_nlmsg_build_hdr(
|
||||
buffer.as_mut_ptr() as *mut c_char,
|
||||
buffer.as_mut_ptr().cast(),
|
||||
libc::NFT_MSG_GETSET as u16,
|
||||
nftnl::ProtoFamily::Inet as u16,
|
||||
(libc::NLM_F_DUMP | libc::NLM_F_ACK) as u16,
|
||||
|
@ -335,11 +329,8 @@ fn should_add<T: Helper>(trie: &RTrieSet<T>, elem: &T) -> bool {
|
|||
|
||||
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
|
||||
}
|
||||
x.direct_parent()
|
||||
.map_or_else(|| *x != T::ZERO, |par| should_add(trie, &par))
|
||||
})
|
||||
}
|
||||
|
||||
|
@ -482,11 +473,10 @@ pub(crate) fn nftables_thread(
|
|||
println!("nftables init done");
|
||||
first = false;
|
||||
}
|
||||
let (rulesets1, ips) = match rx.recv() {
|
||||
Ok(val) => val,
|
||||
Err(_) => break,
|
||||
let Ok((rulesets1, ips)) = rx.recv() else {
|
||||
break;
|
||||
};
|
||||
for i in rulesets1.into_iter() {
|
||||
for i in rulesets1 {
|
||||
let ruleset = &mut rulesets[i];
|
||||
for ip1 in ips.iter().copied() {
|
||||
match ip1 {
|
||||
|
|
|
@ -6,7 +6,10 @@ use crate::bindings::{
|
|||
rrset_trust, sec_status, slabhash, sldns_enum_ede_code, sockaddr_in, sockaddr_in6,
|
||||
sockaddr_storage, ub_packed_rrset_key, AF_INET, AF_INET6,
|
||||
};
|
||||
use std::{ffi::CStr, marker::PhantomData, net::SocketAddr, os::raw::c_char, ptr, time::Duration};
|
||||
use std::{
|
||||
ffi::CStr, marker::PhantomData, net::SocketAddr, ops::Deref, os::raw::c_char, ptr,
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
pub struct ConfigFileMut<'a>(
|
||||
pub(crate) *mut config_file,
|
||||
|
@ -22,7 +25,7 @@ pub struct InfraCacheMut<'a>(
|
|||
PhantomData<&'a mut infra_cache>,
|
||||
);
|
||||
pub struct KeyCacheMut<'a>(pub(crate) *mut key_cache, PhantomData<&'a mut key_cache>);
|
||||
pub struct ModuleEnv<T>(
|
||||
pub struct ModuleEnvMut<T>(
|
||||
pub(crate) *mut module_env,
|
||||
pub(crate) std::ffi::c_int,
|
||||
pub(crate) PhantomData<T>,
|
||||
|
@ -32,18 +35,39 @@ pub struct ModuleQstate<'a, T>(
|
|||
pub(crate) std::ffi::c_int,
|
||||
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(crate) *mut outbound_entry,
|
||||
pub(crate) PhantomData<&'a mut outbound_entry>,
|
||||
);
|
||||
pub struct QueryInfoMut<'a>(
|
||||
pub struct QueryInfo<'a>(
|
||||
pub(crate) *mut query_info,
|
||||
pub(crate) PhantomData<&'a mut query_info>,
|
||||
);
|
||||
pub struct DnsMsgMut<'a>(
|
||||
pub struct QueryInfoMut<'a>(QueryInfo<'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) 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(crate) *mut reply_info,
|
||||
pub(crate) PhantomData<&'a mut reply_info>,
|
||||
|
@ -65,7 +89,7 @@ pub struct PackedRrsetData<'a>(
|
|||
pub(crate) PhantomData<&'a mut packed_rrset_data>,
|
||||
);
|
||||
|
||||
impl<'a> QueryInfoMut<'a> {
|
||||
impl<'a> QueryInfo<'a> {
|
||||
pub fn qname(&self) -> &CStr {
|
||||
unsafe { CStr::from_ptr((*self.0).qname as *const c_char) }
|
||||
}
|
||||
|
@ -77,7 +101,7 @@ impl<'a> QueryInfoMut<'a> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<T> ModuleEnv<T> {
|
||||
impl<T> ModuleEnvMut<T> {
|
||||
pub fn config_file_mut(&mut self) -> ConfigFileMut<'_> {
|
||||
ConfigFileMut(unsafe { (*self.0).cfg }, Default::default())
|
||||
}
|
||||
|
@ -133,7 +157,7 @@ impl<T> ModuleEnv<T> {
|
|||
addr4.sin_port = x.port();
|
||||
addr4.sin_addr.s_addr = (*x.ip()).into();
|
||||
(
|
||||
&addr4 as *const _ as *const sockaddr_storage,
|
||||
std::ptr::addr_of!(addr4).cast::<sockaddr_storage>(),
|
||||
std::mem::size_of_val(&addr4),
|
||||
)
|
||||
}
|
||||
|
@ -142,29 +166,27 @@ impl<T> ModuleEnv<T> {
|
|||
addr6.sin6_flowinfo = x.flowinfo();
|
||||
addr6.sin6_scope_id = x.scope_id();
|
||||
(
|
||||
&addr6 as *const _ as *const sockaddr_storage,
|
||||
std::ptr::addr_of!(addr6).cast(),
|
||||
std::mem::size_of_val(&addr6),
|
||||
)
|
||||
}
|
||||
};
|
||||
((*self.0).send_query.unwrap_unchecked())(
|
||||
&qinfo.0 as *const _ as *mut _,
|
||||
qinfo.0 .0,
|
||||
flags,
|
||||
dnssec as i32,
|
||||
want_dnssec.into(),
|
||||
nocaps.into(),
|
||||
check_ratelimit.into(),
|
||||
addr as *mut _,
|
||||
addr.cast_mut(),
|
||||
addr_len as u32,
|
||||
zone.as_ptr() as *mut _,
|
||||
zone.as_ptr().cast_mut(),
|
||||
zone.len(),
|
||||
tcp_upstream.into(),
|
||||
ssl_upstream.into(),
|
||||
tls_auth_name
|
||||
.map(|x| x.as_ptr() as *mut _)
|
||||
.unwrap_or(ptr::null_mut()),
|
||||
tls_auth_name.map_or_else(ptr::null_mut, |x| x.as_ptr().cast_mut()),
|
||||
q.0,
|
||||
&mut was_ratelimited as *mut _,
|
||||
std::ptr::addr_of_mut!(was_ratelimited),
|
||||
)
|
||||
};
|
||||
if ret.is_null() {
|
||||
|
@ -192,7 +214,7 @@ impl<T> ModuleEnv<T> {
|
|||
let res = unsafe {
|
||||
((*self.0).attach_sub.unwrap_unchecked())(
|
||||
qstate.0,
|
||||
&qinfo.0 as *const _ as *mut _,
|
||||
qinfo.0 .0,
|
||||
qflags,
|
||||
prime.into(),
|
||||
valrec.into(),
|
||||
|
@ -219,30 +241,35 @@ impl<T> ModuleEnv<T> {
|
|||
}
|
||||
|
||||
impl<T> ModuleQstate<'_, T> {
|
||||
pub fn qinfo_mut(&mut self) -> QueryInfoMut<'_> {
|
||||
QueryInfoMut(
|
||||
unsafe { &mut (*self.0).qinfo as *mut query_info },
|
||||
pub fn qinfo(&self) -> QueryInfo<'_> {
|
||||
QueryInfo(
|
||||
unsafe { std::ptr::addr_of_mut!((*self.0).qinfo) },
|
||||
Default::default(),
|
||||
)
|
||||
}
|
||||
pub fn return_msg_mut(&mut self) -> Option<DnsMsgMut<'_>> {
|
||||
pub fn return_msg(&self) -> Option<DnsMsg<'_>> {
|
||||
if unsafe { (*self.0).return_msg.is_null() } {
|
||||
None
|
||||
} else {
|
||||
Some(DnsMsgMut(
|
||||
unsafe { (*self.0).return_msg },
|
||||
Default::default(),
|
||||
))
|
||||
Some(DnsMsg(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) {
|
||||
unsafe {
|
||||
(*self.0).ext_state[self.1 as usize] = state as module_ext_state;
|
||||
(*self.0 .0).ext_state[self.1 as usize] = state as module_ext_state;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DnsMsgMut<'_> {
|
||||
impl DnsMsg<'_> {
|
||||
pub fn rep(&self) -> Option<ReplyInfo<'_>> {
|
||||
if unsafe { (*self.0).rep.is_null() } {
|
||||
None
|
||||
|
@ -318,7 +345,7 @@ impl ReplyInfo<'_> {
|
|||
impl UbPackedRrsetKey<'_> {
|
||||
pub fn entry(&self) -> LruHashEntry<'_> {
|
||||
LruHashEntry(
|
||||
unsafe { &mut (*self.0).entry as *mut _ },
|
||||
unsafe { std::ptr::addr_of_mut!((*self.0).entry) },
|
||||
Default::default(),
|
||||
)
|
||||
}
|
||||
|
@ -326,7 +353,10 @@ impl UbPackedRrsetKey<'_> {
|
|||
unsafe { (*self.0).id }
|
||||
}
|
||||
pub fn rk(&self) -> PackedRrsetKey<'_> {
|
||||
PackedRrsetKey(unsafe { &mut (*self.0).rk as *mut _ }, Default::default())
|
||||
PackedRrsetKey(
|
||||
unsafe { std::ptr::addr_of_mut!((*self.0).rk) },
|
||||
Default::default(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -352,7 +382,7 @@ impl PackedRrsetKey<'_> {
|
|||
impl LruHashEntry<'_> {
|
||||
pub fn data(&self) -> PackedRrsetData<'_> {
|
||||
// FIXME: shouldnt pthread lock be used here?
|
||||
unsafe { PackedRrsetData((*self.0).data as *mut packed_rrset_data, Default::default()) }
|
||||
unsafe { PackedRrsetData((*self.0).data.cast(), Default::default()) }
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -610,8 +640,8 @@ pub enum ModuleExtState {
|
|||
}
|
||||
|
||||
impl ModuleExtState {
|
||||
pub(crate) fn importance(&self) -> usize {
|
||||
match *self {
|
||||
pub(crate) const fn importance(self) -> usize {
|
||||
match self {
|
||||
Self::Unknown => 0,
|
||||
Self::InitialState => 1,
|
||||
Self::Finished => 2,
|
||||
|
|
Loading…
Reference in a new issue