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

fix tests

Browse source
This commit is contained in:
chayleaf 2024-08-13 08:38:01 +07:00
parent be0359f85f
commit 671bb514e6
Signed by: chayleaf
GPG key ID: 78171AD46227E68E
2 changed files with 41 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
src/example.rs
View file

@ -112,6 +112,7 @@ impl IpCacheKey {
Self(split_rev_domain.fold(Domain::new(), |mut ret, seg| { Self(split_rev_domain.fold(Domain::new(), |mut ret, seg| {
if first { if first {
first = false; first = false;
} else {
ret.push(b'.'); ret.push(b'.');
} }
ret.extend_from_slice(seg.as_ref()); ret.extend_from_slice(seg.as_ref());
@ -389,10 +390,10 @@ impl ExampleMod {
} }
fn load_env(&mut self) -> Result<Vec<(NftData<Ipv4Net>, NftData<Ipv6Net>)>, ()> { fn load_env(&mut self) -> Result<Vec<(NftData<Ipv4Net>, NftData<Ipv6Net>)>, ()> {
self.nft_token = std::env::var_os("NFT_TOKEN") self.nft_token = std::env::var_os("NFT_TOKEN")
.map(|x| x.to_str().ok_or(()).map(|s| s.to_owned() + ".")) .map(|x| x.to_str().ok_or(()).map(|s| s.to_owned()))
.transpose()?; .transpose()?;
self.tmp_nft_token = std::env::var_os("NFT_TOKEN") self.tmp_nft_token = std::env::var_os("NFT_TOKEN")
.map(|x| x.to_str().ok_or(()).map(|s| format!("tmp{s}."))) .map(|x| x.to_str().ok_or(()).map(|s| format!("tmp{s}")))
.transpose()?; .transpose()?;
let mut rulesets = Vec::new(); let mut rulesets = Vec::new();
assert!(self.nft_queries.is_empty()); assert!(self.nft_queries.is_empty());
@ -492,12 +493,15 @@ impl ExampleMod {
) -> Result<(), ()> { ) -> Result<(), ()> {
println!("adding {ip4:?}/{ip6:?} for {split_domain:?} to {qnames:?}"); println!("adding {ip4:?}/{ip6:?} for {split_domain:?} to {qnames:?}");
if !ip4.is_empty() || !ip6.is_empty() { if !ip4.is_empty() || !ip6.is_empty() {
let mut first = true;
let domain = match split_domain let domain = match split_domain
.iter() .iter()
.copied() .copied()
.map(std::str::from_utf8) .map(std::str::from_utf8)
.try_fold(String::new(), |mut s, comp| { .try_fold(String::new(), |mut s, comp| {
if !s.is_empty() { if first {
first = false;
} else {
s.push('.'); s.push('.');
} }
s.push_str(comp?); s.push_str(comp?);
@ -527,6 +531,7 @@ impl ExampleMod {
Ok(()) Ok(())
} }
fn run_commands(&self, split_domain: &[&[u8]]) -> Option<ModuleExtState> { fn run_commands(&self, split_domain: &[&[u8]]) -> Option<ModuleExtState> {
println!("{split_domain:?} {:?}", self.nft_token);
if let Some(split_domain) = self.nft_token.as_ref().and_then(|token| { if let Some(split_domain) = self.nft_token.as_ref().and_then(|token| {
split_domain split_domain
.split_last() .split_last()
@ -544,12 +549,15 @@ impl ExampleMod {
if domains.insert(split_domain.iter().copied().rev().map(From::from)) { if domains.insert(split_domain.iter().copied().rev().map(From::from)) {
drop(domains); drop(domains);
let file_name = format!("{DATA_PREFIX}/{qname}_domains.json"); let file_name = format!("{DATA_PREFIX}/{qname}_domains.json");
let mut first = false;
let domain = match split_domain let domain = match split_domain
.iter() .iter()
.copied() .copied()
.map(std::str::from_utf8) .map(std::str::from_utf8)
.try_fold(String::new(), |mut s, comp| { .try_fold(String::new(), |mut s, comp| {
if !s.is_empty() { if !first {
first = true;
} else {
s.push('.'); s.push('.');
} }
s.push_str(comp?); s.push_str(comp?);
@ -584,10 +592,11 @@ impl ExampleMod {
Err(err) => self.report("domains create", err), Err(err) => self.report("domains create", err),
} }
} }
return Some(ModuleExtState::Finished);
} }
} }
} }
return Some(ModuleExtState::Finished); return Some(ModuleExtState::Error);
} else if let Some(split_domain) = self.tmp_nft_token.as_ref().and_then(|token| { } else if let Some(split_domain) = self.tmp_nft_token.as_ref().and_then(|token| {
split_domain split_domain
.split_last() .split_last()
@ -603,10 +612,11 @@ impl ExampleMod {
{ {
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(split_domain.iter().copied().rev().map(From::from));
return Some(ModuleExtState::Finished);
} }
} }
} }
return Some(ModuleExtState::Finished); return Some(ModuleExtState::Error);
} }
None None
} }
@ -709,9 +719,7 @@ impl UnboundMod for ExampleMod {
} }
let info = qstate.qinfo(); let info = qstate.qinfo();
let name = info.qname().to_bytes(); let name = info.qname().to_bytes();
// let rev_domain = name.strip_suffix(b".").unwrap_or(name);
let split_domain = unwire_domain(name); let split_domain = unwire_domain(name);
println!("handling {split_domain:?}");
if let Some(val) = self.run_commands(&split_domain) { if let Some(val) = self.run_commands(&split_domain) {
return Some(val); return Some(val);
} }
@ -742,7 +750,10 @@ mod test {
use ipnet::IpNet; use ipnet::IpNet;
use smallvec::smallvec; use smallvec::smallvec;
use crate::example::{ignore, ExampleMod, IpCacheKey, IpNetDeser, DATA_PREFIX}; use crate::{
example::{ignore, ExampleMod, IpCacheKey, IpNetDeser, DATA_PREFIX},
unbound::ModuleExtState,
};
#[test] #[test]
fn test() { fn test() {
@ -860,13 +871,25 @@ mod test {
) )
.unwrap(); .unwrap();
t.run_commands(&[&b"w"[..], &b"com"[..], &b"q"[..], &b"token"[..]]) assert_eq!(
.unwrap(); t.run_commands(&[&b"w"[..], &b"com"[..], &b"q"[..], &b"token"[..]])
t.run_commands(&[&b"e"[..], &b"com"[..], &b"q"[..], &b"tmptoken"[..]]) .unwrap(),
.unwrap(); ModuleExtState::Finished
assert!(t );
.run_commands(&[&b"e"[..], &b"com"[..], &b"w"[..], &b"tmptoken"[..]]) assert_eq!(
.is_none()); t.run_commands(&[&b"w"[..], &b"com"[..], &b"q"[..], &b"wrongtoken"[..]]),
None
);
assert_eq!(
t.run_commands(&[&b"e"[..], &b"com"[..], &b"q"[..], &b"tmptoken"[..]])
.unwrap(),
ModuleExtState::Finished
);
assert_eq!(
t.run_commands(&[&b"e"[..], &b"com"[..], &b"w"[..], &b"tmptoken"[..]])
.unwrap(),
ModuleExtState::Error
);
let split_domain = [&b"e"[..], &b"com"[..]]; let split_domain = [&b"e"[..], &b"com"[..]];
let qnames = t.get_qnames(&split_domain); let qnames = t.get_qnames(&split_domain);

7
src/nftables.rs
View file

@ -4,10 +4,7 @@ use std::{
fmt::Display, fmt::Display,
io::{self, Write}, io::{self, Write},
net::{Ipv4Addr, Ipv6Addr}, net::{Ipv4Addr, Ipv6Addr},
os::{ os::{fd::BorrowedFd, raw::c_void},
fd::BorrowedFd,
raw::{c_char, c_void},
},
rc::Rc, rc::Rc,
sync::mpsc, sync::mpsc,
}; };
@ -46,7 +43,7 @@ struct FlushSetMsg<'a> {
set: &'a Set1, set: &'a Set1,
} }
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 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.cast(),
libc::NFT_MSG_DELSETELEM as u16, libc::NFT_MSG_DELSETELEM as u16,