aljaz/pkdns
1
0
Fork 0
mirror of https://github.com/aljazceru/pkdns.git synced 2025-12-17 05:54:21 +01:00
Code Issues Packages Projects Releases Wiki Activity

optimize code imports

Browse Source
This commit is contained in:
Severin Buhler
2024-02-06 19:40:36 +01:00
parent cfe6b4dbef
commit aee562e375
5 changed files with 134 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
src/main.rs
View File

@@ -1,14 +1,13 @@
use any_dns::{Builder, CustomHandler, CustomHandlerError, DnsSocket};
use async_trait::async_trait;
use ctrlc;
use pkarr::dns::Packet;
use pknames_resolver::PknamesResolver;
use std::{error::Error, net::SocketAddr, sync::mpsc::channel, time::Instant};
use pknames_resolver::PknamesResolver;
use std::{error::Error, net::SocketAddr};
mod packet_lookup;
mod pkarr_cache;
mod pkarr_resolver;
mod pknames_resolver;
mod packet_lookup;
#[derive(Clone)]
struct MyHandler {
@@ -24,15 +23,18 @@ impl MyHandler {
}
#[async_trait]
impl CustomHandler for MyHandler {
async fn lookup(&mut self, query: &Vec<u8>, _socket: DnsSocket) -> std::prelude::v1::Result<Vec<u8>, CustomHandlerError> {
async fn lookup(
&mut self,
query: &Vec<u8>,
_socket: DnsSocket,
) -> std::prelude::v1::Result<Vec<u8>, CustomHandlerError> {
match self.pkarr.resolve(query) {
Ok(reply) => Ok(reply),
Err(_) => Err(CustomHandlerError::Unhandled)
Err(_) => Err(CustomHandlerError::Unhandled),
}
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
const VERSION: &str = env!("CARGO_PKG_VERSION");
@@ -76,7 +78,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
.required(false)
.default_value("4")
.help("Number of threads to process dns queries."),
)
)
.arg(
clap::Arg::new("directory")
.short('d')
@@ -90,7 +92,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
let verbose: bool = *matches.get_one("verbose").unwrap();
let default_cache_ttl = "60".to_string();
let cache_ttl: &String = matches.get_one("cache-ttl").unwrap_or(&default_cache_ttl);
let cache_ttl: u64 = cache_ttl.parse().expect("cache-ttl should be a valid valid positive integer (u64).");
let cache_ttl: u64 = cache_ttl
.parse()
.expect("cache-ttl should be a valid valid positive integer (u64).");
let directory: &String = matches.get_one("directory").unwrap();
let threads: &String = matches.get_one("threads").unwrap();
let threads: u8 = threads.parse().expect("threads should be valid positive integer.");
@@ -133,7 +137,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
.verbose(verbose)
.icann_resolver(forward)
.listen(socket)
.build().await?;
.build()
.await?;
println!("Listening on {}. Waiting for Ctrl-C...", socket);
anydns.wait_on_ctrl_c().await;

123
src/packet_lookup.rs
View File

@@ -1,4 +1,7 @@
use simple_dns::{rdata::{self, RData}, Name, Packet, Question, ResourceRecord, QTYPE, TYPE};
use simple_dns::{
rdata::{self, RData},
Name, Packet, Question, ResourceRecord, QTYPE, TYPE,
};
/**
* Handles all possible ways on how to resolve a query into a reply.
@@ -51,15 +54,18 @@ fn resolve_question<'a>(pkarr_packet: &Packet<'a>, question: &Question<'a>) -> V
fn resolve_cname_for<'a>(pkarr_packet: &Packet<'a>, question: &Question<'a>) -> Vec<ResourceRecord<'a>> {
let cname_matches = direct_matches(pkarr_packet, &question.qname, &QTYPE::TYPE(TYPE::CNAME));
let additional_data: Vec<ResourceRecord<'_>> = cname_matches.iter().flat_map(|cname| {
let cname_content = if let RData::CNAME(rdata::CNAME(cname_pointer)) = &cname.rdata {
cname_pointer
} else {
panic!("Should be cname");
};
let matches = direct_matches(pkarr_packet, &cname_content, &question.qtype);
matches
}).collect();
let additional_data: Vec<ResourceRecord<'_>> = cname_matches
.iter()
.flat_map(|cname| {
let cname_content = if let RData::CNAME(rdata::CNAME(cname_pointer)) = &cname.rdata {
cname_pointer
} else {
panic!("Should be cname");
};
let matches = direct_matches(pkarr_packet, &cname_content, &question.qtype);
matches
})
.collect();
let mut result = vec![];
result.extend(cname_matches);
@@ -72,12 +78,12 @@ fn resolve_cname_for<'a>(pkarr_packet: &Packet<'a>, question: &Question<'a>) ->
* Resolve direct qname and qtype record matches.
*/
fn direct_matches<'a>(pkarr_packet: &Packet<'a>, qname: &Name<'a>, qtype: &QTYPE) -> Vec<ResourceRecord<'a>> {
let matches: Vec<ResourceRecord<'_>> = pkarr_packet.answers.iter()
.filter(|record| {
record.name == *qname && record.match_qtype(*qtype)
})
.map(|record| record.clone())
.collect();
let matches: Vec<ResourceRecord<'_>> = pkarr_packet
.answers
.iter()
.filter(|record| record.name == *qname && record.match_qtype(*qtype))
.map(|record| record.clone())
.collect();
matches
}
@@ -85,12 +91,14 @@ fn direct_matches<'a>(pkarr_packet: &Packet<'a>, qname: &Name<'a>, qtype: &QTYPE
* Find nameserver for given qname.
*/
fn find_nameserver<'a>(pkarr_packet: &Packet<'a>, qname: &Name<'a>) -> Vec<ResourceRecord<'a>> {
let matches: Vec<ResourceRecord<'_>> = pkarr_packet.answers.iter()
.filter(|record| {
record.match_qtype(QTYPE::TYPE(TYPE::NS)) && (qname.is_subdomain_of(&record.name) || record.name == *qname)
})
.map(|record| record.clone())
.collect();
let matches: Vec<ResourceRecord<'_>> = pkarr_packet
.answers
.iter()
.filter(|record| {
record.match_qtype(QTYPE::TYPE(TYPE::NS)) && (qname.is_subdomain_of(&record.name) || record.name == *qname)
})
.map(|record| record.clone())
.collect();
matches
}
@@ -99,7 +107,8 @@ mod tests {
use std::net::Ipv4Addr;
use pkarr::{
dns::{Name, Packet, ResourceRecord}, Keypair, PublicKey
dns::{Name, Packet, ResourceRecord},
Keypair, PublicKey,
};
use simple_dns::{rdata::RData, Question};
@@ -116,18 +125,14 @@ mod tests {
let mut packet = Packet::new_reply(0);
let name = Name::new(&pubkey_z32).unwrap();
let ip: Ipv4Addr = "127.0.0.1".parse().unwrap();
let answer1 = ResourceRecord::new(
name.clone(), simple_dns::CLASS::IN, 100, RData::A(ip.into())
);
let ip: Ipv4Addr = "127.0.0.1".parse().unwrap();
let answer1 = ResourceRecord::new(name.clone(), simple_dns::CLASS::IN, 100, RData::A(ip.into()));
packet.answers.push(answer1);
let name = format!("pknames.p2p.{pubkey_z32}");
let name = Name::new(&name).unwrap();
let ip: Ipv4Addr = "127.0.0.1".parse().unwrap();
let answer1 = ResourceRecord::new(
name.clone(), simple_dns::CLASS::IN, 100, RData::A(ip.into())
);
let ip: Ipv4Addr = "127.0.0.1".parse().unwrap();
let answer1 = ResourceRecord::new(name.clone(), simple_dns::CLASS::IN, 100, RData::A(ip.into()));
packet.answers.push(answer1);
let name = format!("www.pknames.p2p.{pubkey_z32}");
@@ -135,7 +140,10 @@ mod tests {
let data = format!("pknames.p2p.{pubkey_z32}");
let data = Name::new(&data).unwrap();
let answer3 = ResourceRecord::new(
name.clone(), simple_dns::CLASS::IN, 100, RData::CNAME(simple_dns::rdata::CNAME(data))
name.clone(),
simple_dns::CLASS::IN,
100,
RData::CNAME(simple_dns::rdata::CNAME(data)),
);
packet.answers.push(answer3);
@@ -144,16 +152,16 @@ mod tests {
let data = format!("my.ns.example.com");
let data = Name::new(&data).unwrap();
let answer4 = ResourceRecord::new(
name.clone(), simple_dns::CLASS::IN, 100, RData::NS(simple_dns::rdata::NS(data))
name.clone(),
simple_dns::CLASS::IN,
100,
RData::NS(simple_dns::rdata::NS(data)),
);
packet.answers.push(answer4);
(packet.build_bytes_vec_compressed().unwrap(), pubkey)
}
#[test]
fn simple_a_question() {
let (pkarr_packet, pubkey) = example_pkarr_reply();
@@ -163,7 +171,12 @@ mod tests {
let name = format!("pknames.p2p.{pubkey_z32}");
let name = Name::new(&name).unwrap();
let qtype = simple_dns::QTYPE::TYPE(simple_dns::TYPE::A);
let question = Question::new(name.clone(), qtype, simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN), false);
let question = Question::new(
name.clone(),
qtype,
simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN),
false,
);
let reply = resolve_question(&pkarr_packet, &question);
let reply = Packet::parse(&reply).unwrap();
@@ -184,7 +197,12 @@ mod tests {
let name = format!("www.pknames.p2p.{pubkey_z32}");
let name = Name::new(&name).unwrap();
let qtype = simple_dns::QTYPE::TYPE(simple_dns::TYPE::A);
let question = Question::new(name.clone(), qtype, simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN), false);
let question = Question::new(
name.clone(),
qtype,
simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN),
false,
);
let reply = resolve_question(&pkarr_packet, &question);
let reply = Packet::parse(&reply).unwrap();
@@ -210,7 +228,12 @@ mod tests {
let name = format!("other.{pubkey_z32}");
let name = Name::new(&name).unwrap();
let qtype = simple_dns::QTYPE::TYPE(simple_dns::TYPE::A);
let question = Question::new(name.clone(), qtype, simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN), false);
let question = Question::new(
name.clone(),
qtype,
simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN),
false,
);
let reply = resolve_question(&pkarr_packet, &question);
let reply = Packet::parse(&reply).unwrap();
@@ -232,7 +255,12 @@ mod tests {
let name = format!("sub.other.{pubkey_z32}");
let name = Name::new(&name).unwrap();
let qtype = simple_dns::QTYPE::TYPE(simple_dns::TYPE::A);
let question = Question::new(name.clone(), qtype, simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN), false);
let question = Question::new(
name.clone(),
qtype,
simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN),
false,
);
let reply = resolve_question(&pkarr_packet, &question);
let reply = Packet::parse(&reply).unwrap();
@@ -247,16 +275,17 @@ mod tests {
#[test]
fn simple_a_query() {
let (pkarr_packet, pubkey) = example_pkarr_reply();
let (pkarr_packet, _pubkey) = example_pkarr_reply();
let pkarr_packet = Packet::parse(&pkarr_packet).unwrap();
let mut query = Packet::new_query(0);
query.questions = vec![
Question::new(Name::new("pknames.p2p").unwrap(), simple_dns::QTYPE::TYPE(simple_dns::TYPE::A), simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN), false)
];
query.questions = vec![Question::new(
Name::new("pknames.p2p").unwrap(),
simple_dns::QTYPE::TYPE(simple_dns::TYPE::A),
simple_dns::QCLASS::CLASS(simple_dns::CLASS::IN),
false,
)];
let _reply = resolve_query(&pkarr_packet, &query);
}
}
}

19
src/pkarr_cache.rs
View File

@@ -3,22 +3,19 @@ use std::time::Duration;
use pkarr::{dns::Packet, PublicKey};
use ttl_cache::TtlCache;
/**
* Pkarr record ttl cache
*/
pub struct PkarrPacketTtlCache{
pub struct PkarrPacketTtlCache {
cache: TtlCache<String, Vec<u8>>,
max_cache_ttl: u64,
}
impl PkarrPacketTtlCache {
pub fn new(max_cache_ttl: u64) -> Self {
PkarrPacketTtlCache{
PkarrPacketTtlCache {
cache: TtlCache::new(100),
max_cache_ttl
max_cache_ttl,
}
}
@@ -28,7 +25,12 @@ impl PkarrPacketTtlCache {
pub fn add(&mut self, pubkey: PublicKey, reply: Vec<u8>) {
let default_ttl = 1200;
let packet = Packet::parse(&reply).unwrap();
let min_ttl = packet.answers.iter().map(|answer| answer.ttl).min().unwrap_or(default_ttl) as u64;
let min_ttl = packet
.answers
.iter()
.map(|answer| answer.ttl)
.min()
.unwrap_or(default_ttl) as u64;
let ttl = 60.max(min_ttl); // At least 1min
let ttl = ttl.min(self.max_cache_ttl);
@@ -41,5 +43,4 @@ impl PkarrPacketTtlCache {
let z32 = pubkey.to_z32();
self.cache.get(&z32).map(|value| value.clone())
}
}
}

39
src/pkarr_resolver.rs
View File

@@ -1,19 +1,19 @@
use std::{error::Error, sync::{Arc, Mutex}};
use pkarr::{
dns::{Packet, Question, ResourceRecord, QTYPE, TYPE}, PkarrClient, PublicKey, SignedPacket
use std::{
error::Error,
sync::{Arc, Mutex},
};
use chrono::{DateTime, Utc};
use crate::{packet_lookup::resolve_query, pkarr_cache::PkarrPacketTtlCache};
use crate::{packet_lookup::resolve_query, pkarr_cache::PkarrPacketTtlCache};
use chrono::{DateTime, Utc};
use pkarr::{dns::Packet, PkarrClient, PublicKey, SignedPacket};
trait SignedPacketTimestamp {
fn chrono_timestamp(&self) -> DateTime<Utc>;
}
impl SignedPacketTimestamp for SignedPacket {
fn chrono_timestamp(&self) -> DateTime<Utc> {
let timestamp = self.timestamp()/1_000_000;
fn chrono_timestamp(&self) -> DateTime<Utc> {
let timestamp = self.timestamp() / 1_000_000;
let timestamp = DateTime::from_timestamp((timestamp as u32).into(), 0).unwrap();
timestamp
}
@@ -54,10 +54,9 @@ impl PkarrResolver {
let cached_opt = cache.get(pubkey);
if cached_opt.is_some() {
let reply_bytes = cached_opt.unwrap();
return Some(reply_bytes)
return Some(reply_bytes);
};
let packet_option = self.client.resolve(pubkey.clone());
if packet_option.is_none() {
return None;
@@ -71,10 +70,7 @@ impl PkarrResolver {
/**
* Resolves a domain with pkarr.
*/
pub fn resolve(
&mut self,
query: &Vec<u8>
) -> std::prelude::v1::Result<Vec<u8>, Box<dyn Error>> {
pub fn resolve(&mut self, query: &Vec<u8>) -> std::prelude::v1::Result<Vec<u8>, Box<dyn Error>> {
let request = Packet::parse(query)?;
let question_opt = request.questions.first();
@@ -112,10 +108,10 @@ mod tests {
dns::{Name, Packet, Question, ResourceRecord},
Keypair, SignedPacket,
};
use simple_dns::rdata::A;
// use simple_dns::{Name, Question, Packet};
use super::*;
use std::{fmt::format, net::Ipv4Addr};
use std::net::Ipv4Addr;
use zbase32;
fn get_test_keypair() -> Keypair {
@@ -170,7 +166,7 @@ mod tests {
true,
);
query.questions.push(question);
let mut resolver = PkarrResolver::new(0);
let result = resolver.resolve(&query.build_bytes_vec_compressed().unwrap());
assert!(result.is_ok());
@@ -248,7 +244,7 @@ mod tests {
let pubkey = PkarrResolver::parse_pkarr_uri("7fmjpcuuzf54hw18bsgi3zihzyh4awseeuq5tmojefaezjbd64cy").unwrap();
let mut resolver = PkarrResolver::new(0);
let result = resolver.resolve_pubkey_respect_cache(&pubkey);
let _result = resolver.resolve_pubkey_respect_cache(&pubkey);
// assert!(result.is_some());
}
@@ -274,7 +270,10 @@ mod tests {
let data = format!("pknames.p2p.{pubkey_z32}");
let data = Name::new(&data).unwrap();
let answer3 = ResourceRecord::new(
name.clone(), simple_dns::CLASS::IN, 100, simple_dns::rdata::RData::CNAME(simple_dns::rdata::CNAME(data))
name.clone(),
simple_dns::CLASS::IN,
100,
simple_dns::rdata::RData::CNAME(simple_dns::rdata::CNAME(data)),
);
packet.answers.push(answer3);
@@ -285,6 +284,4 @@ mod tests {
let reply_bytes = signed_packet.packet().build_bytes_vec().unwrap();
Packet::parse(&reply_bytes).unwrap(); // Fail
}
}

26
src/pknames_resolver.rs
View File

@@ -1,7 +1,6 @@
use crate::pkarr_resolver::PkarrResolver;
use pkarr::dns::{Name, Packet};
use pknames_core::resolve::resolve_standalone;
use crate::pkarr_resolver::PkarrResolver;
#[derive(Clone)]
pub struct PknamesResolver {
@@ -13,7 +12,7 @@ impl PknamesResolver {
pub fn new(max_cache_ttl: u64, config_dir_path: &str) -> Self {
PknamesResolver {
pkarr: PkarrResolver::new(max_cache_ttl),
config_dir_path: config_dir_path.to_string()
config_dir_path: config_dir_path.to_string(),
}
}
@@ -37,7 +36,6 @@ impl PknamesResolver {
Ok(full_domain)
}
pub fn resolve(&mut self, query: &Vec<u8>) -> std::prelude::v1::Result<Vec<u8>, Box<dyn std::error::Error>> {
let original_query = Packet::parse(query)?;
@@ -46,7 +44,10 @@ impl PknamesResolver {
return pkarr_result; // It was a pkarr hostname
}
let question = original_query.questions.first().ok_or("Query does not include a question.")?;
let question = original_query
.questions
.first()
.ok_or("Query does not include a question.")?;
let domain = question.qname.to_string();
let pkarr_domain = self.predict_pknames_domain(&domain)?;
@@ -62,13 +63,11 @@ impl PknamesResolver {
let mut answer = answer.clone();
answer.name = question.qname.clone();
reply.answers.push(answer);
};
}
Ok(reply.build_bytes_vec_compressed().unwrap())
}
}
#[cfg(test)]
mod tests {
use pkarr::dns::{Name, Packet, Question};
@@ -81,11 +80,16 @@ mod tests {
let mut query = Packet::new_query(0);
let name = Name::new("pknames.p2p").unwrap();
let question = Question::new(name, pkarr::dns::QTYPE::TYPE(pkarr::dns::TYPE::A), pkarr::dns::QCLASS::CLASS(pkarr::dns::CLASS::IN), false);
let question = Question::new(
name,
pkarr::dns::QTYPE::TYPE(pkarr::dns::TYPE::A),
pkarr::dns::QCLASS::CLASS(pkarr::dns::CLASS::IN),
false,
);
query.questions.push(question);
let query_bytes = query.build_bytes_vec_compressed().unwrap();
let result = pknames.resolve(&query_bytes);
let result = pknames.resolve(&query_bytes);
if result.is_err() {
eprintln!("{:?}", result.unwrap_err());
assert!(false);
@@ -96,7 +100,5 @@ mod tests {
let reply = result.unwrap();
let reply = Packet::parse(&reply).unwrap();
println!("{:?}", reply);
}
}