Sweetbread/Alfis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Code reformatting.

Browse Source
This commit is contained in:
Revertron
2021-06-09 20:36:36 +02:00
parent 2d12fd0447
commit d513c29cfe
40 changed files with 546 additions and 815 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/dns/authority.rs
+15 -19
View File
@@ -6,9 +6,9 @@ use std::io::Write;
use std::path::Path;
use std::sync::{LockResult, RwLock, RwLockReadGuard, RwLockWriteGuard};
use derive_more::{Display, Error, From};
#[allow(unused_imports)]
use log::{trace, debug, info, warn, error};
use derive_more::{Display, From, Error};
use log::{debug, error, info, trace, warn};
use crate::dns::buffer::{PacketBuffer, StreamPacketBuffer, VectorPacketBuffer};
use crate::dns::protocol::{DnsPacket, DnsRecord, QueryType, ResultCode, TransientTtl};
@@ -18,7 +18,7 @@ pub enum AuthorityError {
Buffer(crate::dns::buffer::BufferError),
Protocol(crate::dns::protocol::ProtocolError),
Io(std::io::Error),
PoisonedLock,
PoisonedLock
}
type Result<T> = std::result::Result<T, AuthorityError>;
@@ -33,7 +33,7 @@ pub struct Zone {
pub retry: u32,
pub expire: u32,
pub minimum: u32,
pub records: BTreeSet<DnsRecord>,
pub records: BTreeSet<DnsRecord>
}
impl Zone {
@@ -62,19 +62,17 @@ impl Zone {
#[derive(Default)]
pub struct Zones {
zones: BTreeMap<String, Zone>,
zones: BTreeMap<String, Zone>
}
impl<'a> Zones {
pub fn new() -> Zones {
Zones {
zones: BTreeMap::new(),
}
Zones { zones: BTreeMap::new() }
}
pub fn load(&mut self) -> Result<()> {
let zones_dir = match Path::new("zones").read_dir() {
Ok(result) => { result }
Ok(result) => result,
Err(_) => {
debug!("Authority dir (zones) not found, skipping.");
return Ok(());
@@ -84,12 +82,12 @@ impl<'a> Zones {
for wrapped_filename in zones_dir {
let filename = match wrapped_filename {
Ok(x) => x,
Err(_) => continue,
Err(_) => continue
};
let mut zone_file = match File::open(filename.path()) {
Ok(x) => x,
Err(_) => continue,
Err(_) => continue
};
let mut buffer = StreamPacketBuffer::new(&mut zone_file);
@@ -171,14 +169,12 @@ impl<'a> Zones {
#[derive(Default)]
pub struct Authority {
zones: RwLock<Zones>,
zones: RwLock<Zones>
}
impl Authority {
pub fn new() -> Authority {
Authority {
zones: RwLock::new(Zones::new()),
}
Authority { zones: RwLock::new(Zones::new()) }
}
pub fn load(&self) -> Result<()> {
@@ -194,7 +190,7 @@ impl Authority {
pub fn query(&self, qname: &str, qtype: QueryType) -> Option<DnsPacket> {
let zones = match self.zones.read().ok() {
Some(x) => x,
None => return None,
None => return None
};
let mut best_match = None;
@@ -214,7 +210,7 @@ impl Authority {
let zone = match best_match {
Some((_, zone)) => zone,
None => return None,
None => return None
};
let mut packet = DnsPacket::new();
@@ -223,7 +219,7 @@ impl Authority {
for rec in &zone.records {
let domain = match rec.get_domain() {
Some(x) => x,
None => continue,
None => continue
};
if &domain != qname {
@@ -248,7 +244,7 @@ impl Authority {
retry: zone.retry,
expire: zone.expire,
minimum: zone.minimum,
ttl: TransientTtl(zone.minimum),
ttl: TransientTtl(zone.minimum)
});
}
src/dns/buffer.rs
+17 -23
View File
@@ -8,7 +8,7 @@ use derive_more::{Display, Error, From};
#[derive(Debug, Display, From, Error)]
pub enum BufferError {
Io(std::io::Error),
EndOfBuffer,
EndOfBuffer
}
type Result<T> = std::result::Result<T, BufferError>;
@@ -155,16 +155,12 @@ pub trait PacketBuffer {
pub struct VectorPacketBuffer {
pub buffer: Vec<u8>,
pub pos: usize,
pub label_lookup: BTreeMap<String, usize>,
pub label_lookup: BTreeMap<String, usize>
}
impl VectorPacketBuffer {
pub fn new() -> VectorPacketBuffer {
VectorPacketBuffer {
buffer: Vec::new(),
pos: 0,
label_lookup: BTreeMap::new(),
}
VectorPacketBuffer { buffer: Vec::new(), pos: 0, label_lookup: BTreeMap::new() }
}
}
@@ -222,10 +218,11 @@ impl PacketBuffer for VectorPacketBuffer {
}
}
pub struct StreamPacketBuffer<'a, T> where T: Read {
pub struct StreamPacketBuffer<'a, T>
where T: Read {
pub stream: &'a mut T,
pub buffer: Vec<u8>,
pub pos: usize,
pub pos: usize
}
impl<'a, T> StreamPacketBuffer<'a, T> where T: Read + 'a {
@@ -305,15 +302,12 @@ impl<'a, T> PacketBuffer for StreamPacketBuffer<'a, T> where T: Read + 'a {
pub struct BytePacketBuffer {
pub buf: [u8; 512],
pub pos: usize,
pub pos: usize
}
impl BytePacketBuffer {
pub fn new() -> BytePacketBuffer {
BytePacketBuffer {
buf: [0; 512],
pos: 0,
}
BytePacketBuffer { buf: [0; 512], pos: 0 }
}
}
@@ -401,7 +395,7 @@ mod tests {
// First write the standard string
match buffer.write_qname(&instr1) {
Ok(_) => {}
Err(_) => panic!(),
Err(_) => panic!()
}
// Then we set up a slight variation with relies on a jump back to the data of
@@ -410,7 +404,7 @@ mod tests {
for b in &crafted_data {
match buffer.write_u8(*b) {
Ok(_) => {}
Err(_) => panic!(),
Err(_) => panic!()
}
}
@@ -421,7 +415,7 @@ mod tests {
let mut outstr1 = String::new();
match buffer.read_qname(&mut outstr1) {
Ok(_) => {}
Err(_) => panic!(),
Err(_) => panic!()
}
assert_eq!(instr1, outstr1);
@@ -430,7 +424,7 @@ mod tests {
let mut outstr2 = String::new();
match buffer.read_qname(&mut outstr2) {
Ok(_) => {}
Err(_) => panic!(),
Err(_) => panic!()
}
assert_eq!(instr2, outstr2);
@@ -445,24 +439,24 @@ mod tests {
match buffer.write_qname(&"ns1.google.com".to_string()) {
Ok(_) => {}
Err(_) => panic!(),
Err(_) => panic!()
}
match buffer.write_qname(&"ns2.google.com".to_string()) {
Ok(_) => {}
Err(_) => panic!(),
Err(_) => panic!()
}
assert_eq!(22, buffer.pos());
match buffer.seek(0) {
Ok(_) => {}
Err(_) => panic!(),
Err(_) => panic!()
}
let mut str1 = String::new();
match buffer.read_qname(&mut str1) {
Ok(_) => {}
Err(_) => panic!(),
Err(_) => panic!()
}
assert_eq!("ns1.google.com", str1);
@@ -470,7 +464,7 @@ mod tests {
let mut str2 = String::new();
match buffer.read_qname(&mut str2) {
Ok(_) => {}
Err(_) => panic!(),
Err(_) => panic!()
}
assert_eq!("ns2.google.com", str2);
src/dns/cache.rs
+24 -67
View File
@@ -15,7 +15,7 @@ use crate::dns::protocol::{DnsPacket, DnsRecord, QueryType, ResultCode};
#[derive(Debug, Display, From, Error)]
pub enum CacheError {
Io(std::io::Error),
PoisonedLock,
PoisonedLock
}
type Result<T> = std::result::Result<T, CacheError>;
@@ -23,13 +23,13 @@ type Result<T> = std::result::Result<T, CacheError>;
pub enum CacheState {
PositiveCache,
NegativeCache,
NotCached,
NotCached
}
#[derive(Clone, Eq, Debug, Serialize, Deserialize)]
pub struct RecordEntry {
pub record: DnsRecord,
pub timestamp: DateTime<Local>,
pub timestamp: DateTime<Local>
}
impl PartialEq<RecordEntry> for RecordEntry {
@@ -47,7 +47,7 @@ impl Hash for RecordEntry {
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum RecordSet {
NoRecords { qtype: QueryType, ttl: u32, timestamp: DateTime<Local> },
Records { qtype: QueryType, records: HashSet<RecordEntry> },
Records { qtype: QueryType, records: HashSet<RecordEntry> }
}
#[derive(Clone, Debug)]
@@ -55,7 +55,7 @@ pub struct DomainEntry {
pub domain: String,
pub record_types: HashMap<QueryType, RecordSet>,
pub hits: u32,
pub updates: u32,
pub updates: u32
}
impl DomainEntry {
@@ -128,7 +128,7 @@ impl DomainEntry {
CacheState::NegativeCache
}
}
None => CacheState::NotCached,
None => CacheState::NotCached
}
}
@@ -137,7 +137,7 @@ impl DomainEntry {
let current_set = match self.record_types.get(&qtype) {
Some(x) => x,
None => return,
None => return
};
if let RecordSet::Records { ref records, .. } = *current_set {
@@ -158,7 +158,7 @@ impl DomainEntry {
#[derive(Default)]
pub struct Cache {
domain_entries: BTreeMap<String, Arc<DomainEntry>>,
domain_entries: BTreeMap<String, Arc<DomainEntry>>
}
impl Cache {
@@ -169,11 +169,11 @@ impl Cache {
fn get_cache_state(&mut self, qname: &str, qtype: QueryType) -> CacheState {
match self.domain_entries.get(qname) {
Some(x) => x.get_cache_state(qtype),
None => CacheState::NotCached,
None => CacheState::NotCached
}
}
fn fill_queryresult(&mut self,qname: &str, qtype: QueryType, result_vec: &mut Vec<DnsRecord>, increment_stats: bool) {
fn fill_queryresult(&mut self, qname: &str, qtype: QueryType, result_vec: &mut Vec<DnsRecord>, increment_stats: bool) {
if let Some(domain_entry) = self.domain_entries.get_mut(qname).and_then(Arc::get_mut) {
if increment_stats {
domain_entry.hits += 1
@@ -198,7 +198,7 @@ impl Cache {
Some(qr)
}
CacheState::NotCached => None,
CacheState::NotCached => None
}
}
@@ -206,7 +206,7 @@ impl Cache {
for rec in records {
let domain = match rec.get_domain() {
Some(x) => x,
None => continue,
None => continue
};
if let Some(ref mut rs) = self.domain_entries.get_mut(&domain).and_then(Arc::get_mut) {
@@ -234,7 +234,7 @@ impl Cache {
#[derive(Default)]
pub struct SynchronizedCache {
pub cache: RwLock<Cache>,
pub cache: RwLock<Cache>
}
impl SynchronizedCache {
@@ -257,7 +257,7 @@ impl SynchronizedCache {
pub fn lookup(&self, qname: &str, qtype: QueryType) -> Option<DnsPacket> {
let mut cache = match self.cache.write() {
Ok(x) => x,
Err(_) => return None,
Err(_) => return None
};
cache.lookup(qname, qtype)
@@ -284,7 +284,6 @@ impl SynchronizedCache {
mod tests {
use super::*;
use crate::dns::protocol::{DnsRecord, QueryType, ResultCode, TransientTtl};
#[test]
@@ -318,17 +317,17 @@ mod tests {
records.push(DnsRecord::A {
domain: "www.google.com".to_string(),
addr: "127.0.0.1".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
records.push(DnsRecord::A {
domain: "www.yahoo.com".to_string(),
addr: "127.0.0.2".parse().unwrap(),
ttl: TransientTtl(0),
ttl: TransientTtl(0)
});
records.push(DnsRecord::CNAME {
domain: "www.microsoft.com".to_string(),
host: "www.somecdn.com".to_string(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
cache.store(&records);
@@ -361,7 +360,7 @@ mod tests {
records2.push(DnsRecord::A {
domain: "www.yahoo.com".to_string(),
addr: "127.0.0.2".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
cache.store(&records2);
@@ -373,53 +372,11 @@ mod tests {
// Check stat counter behavior
assert_eq!(3, cache.domain_entries.len());
assert_eq!(
1,
cache
.domain_entries
.get(&"www.google.com".to_string())
.unwrap()
.hits
);
assert_eq!(
2,
cache
.domain_entries
.get(&"www.google.com".to_string())
.unwrap()
.updates
);
assert_eq!(
1,
cache
.domain_entries
.get(&"www.yahoo.com".to_string())
.unwrap()
.hits
);
assert_eq!(
3,
cache
.domain_entries
.get(&"www.yahoo.com".to_string())
.unwrap()
.updates
);
assert_eq!(
1,
cache
.domain_entries
.get(&"www.microsoft.com".to_string())
.unwrap()
.updates
);
assert_eq!(
1,
cache
.domain_entries
.get(&"www.microsoft.com".to_string())
.unwrap()
.hits
);
assert_eq!(1, cache.domain_entries.get(&"www.google.com".to_string()).unwrap().hits);
assert_eq!(2, cache.domain_entries.get(&"www.google.com".to_string()).unwrap().updates);
assert_eq!(1, cache.domain_entries.get(&"www.yahoo.com".to_string()).unwrap().hits);
assert_eq!(3, cache.domain_entries.get(&"www.yahoo.com".to_string()).unwrap().updates);
assert_eq!(1, cache.domain_entries.get(&"www.microsoft.com".to_string()).unwrap().updates);
assert_eq!(1, cache.domain_entries.get(&"www.microsoft.com".to_string()).unwrap().hits);
}
}
src/dns/client.rs
+12 -16
View File
@@ -2,7 +2,7 @@
use std::io::Write;
use std::marker::{Send, Sync};
use std::net::{TcpStream, UdpSocket, ToSocketAddrs, SocketAddr};
use std::net::{SocketAddr, TcpStream, ToSocketAddrs, UdpSocket};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Sender};
use std::sync::{Arc, Mutex};
@@ -22,7 +22,7 @@ pub enum ClientError {
Io(std::io::Error),
PoisonedLock,
LookupFailed,
TimeOut,
TimeOut
}
type Result<T> = std::result::Result<T, ClientError>;
@@ -56,7 +56,7 @@ pub struct DnsNetworkClient {
socket_ipv6: UdpSocket,
/// Queries in progress
pending_queries: Arc<Mutex<Vec<PendingQuery>>>,
pending_queries: Arc<Mutex<Vec<PendingQuery>>>
}
/// A query in progress. This struct holds the `id` if the request, and a channel
@@ -65,7 +65,7 @@ pub struct DnsNetworkClient {
struct PendingQuery {
seq: u16,
timestamp: DateTime<Local>,
tx: Sender<Option<DnsPacket>>,
tx: Sender<Option<DnsPacket>>
}
unsafe impl Send for DnsNetworkClient {}
@@ -80,7 +80,7 @@ impl DnsNetworkClient {
seq: AtomicUsize::new(0),
socket_ipv4: UdpSocket::bind(format!("0.0.0.0:{}", port)).expect("Error binding IPv4"),
socket_ipv6: UdpSocket::bind(format!("[::]:{}", port + 1)).expect("Error binding IPv6"),
pending_queries: Arc::new(Mutex::new(Vec::new())),
pending_queries: Arc::new(Mutex::new(Vec::new()))
}
}
@@ -331,7 +331,7 @@ impl DnsClient for DnsNetworkClient {
Ok(())
}
fn send_query(&self,qname: &str, qtype: QueryType, server: &str, recursive: bool) -> Result<DnsPacket> {
fn send_query(&self, qname: &str, qtype: QueryType, server: &str, recursive: bool) -> Result<DnsPacket> {
let packet = self.send_udp_query(qname, qtype, server, recursive)?;
if !packet.header.truncated_message {
return Ok(packet);
@@ -350,7 +350,7 @@ pub mod tests {
pub type StubCallback = dyn Fn(&str, QueryType, &str, bool) -> Result<DnsPacket>;
pub struct DnsStubClient {
callback: Box<StubCallback>,
callback: Box<StubCallback>
}
impl<'a> DnsStubClient {
@@ -376,7 +376,7 @@ pub mod tests {
Ok(())
}
fn send_query(&self,qname: &str, qtype: QueryType, server: &str, recursive: bool) -> Result<DnsPacket> {
fn send_query(&self, qname: &str, qtype: QueryType, server: &str, recursive: bool) -> Result<DnsPacket> {
(self.callback)(qname, qtype, server, recursive)
}
}
@@ -386,9 +386,7 @@ pub mod tests {
let client = DnsNetworkClient::new(31456);
client.run().unwrap();
let res = client
.send_udp_query("google.com", QueryType::A, ("8.8.8.8", 53), true)
.unwrap();
let res = client.send_udp_query("google.com", QueryType::A, ("8.8.8.8", 53), true).unwrap();
assert_eq!(res.questions[0].name, "google.com");
assert!(res.answers.len() > 0);
@@ -397,16 +395,14 @@ pub mod tests {
DnsRecord::A { ref domain, .. } => {
assert_eq!("google.com", domain);
}
_ => panic!(),
_ => panic!()
}
}
#[test]
pub fn test_tcp_client() {
let client = DnsNetworkClient::new(31458);
let res = client
.send_tcp_query("google.com", QueryType::A, ("8.8.8.8", 53), true)
.unwrap();
let res = client.send_tcp_query("google.com", QueryType::A, ("8.8.8.8", 53), true).unwrap();
assert_eq!(res.questions[0].name, "google.com");
assert!(res.answers.len() > 0);
@@ -415,7 +411,7 @@ pub mod tests {
DnsRecord::A { ref domain, .. } => {
assert_eq!("google.com", domain);
}
_ => panic!(),
_ => panic!()
}
}
}
src/dns/context.rs
+9 -17
View File
@@ -8,21 +8,21 @@ use derive_more::{Display, Error, From};
use crate::dns::authority::Authority;
use crate::dns::cache::SynchronizedCache;
use crate::dns::client::{DnsClient, DnsNetworkClient};
use crate::dns::resolve::{DnsResolver, ForwardingDnsResolver, RecursiveDnsResolver};
use crate::dns::filter::DnsFilter;
use crate::dns::resolve::{DnsResolver, ForwardingDnsResolver, RecursiveDnsResolver};
#[derive(Debug, Display, From, Error)]
pub enum ContextError {
Authority(crate::dns::authority::AuthorityError),
Client(crate::dns::client::ClientError),
Io(std::io::Error),
Io(std::io::Error)
}
type Result<T> = std::result::Result<T, ContextError>;
pub struct ServerStatistics {
pub tcp_query_count: AtomicUsize,
pub udp_query_count: AtomicUsize,
pub udp_query_count: AtomicUsize
}
impl ServerStatistics {
@@ -37,7 +37,7 @@ impl ServerStatistics {
pub enum ResolveStrategy {
Recursive,
Forward { upstreams: Vec<String> },
Forward { upstreams: Vec<String> }
}
pub struct ServerContext {
@@ -76,11 +76,8 @@ impl ServerContext {
enable_udp: true,
enable_tcp: true,
enable_api: false,
statistics: ServerStatistics {
tcp_query_count: AtomicUsize::new(0),
udp_query_count: AtomicUsize::new(0),
},
zones_dir: "zones",
statistics: ServerStatistics { tcp_query_count: AtomicUsize::new(0), udp_query_count: AtomicUsize::new(0) },
zones_dir: "zones"
}
}
@@ -110,13 +107,11 @@ pub mod tests {
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use super::*;
use crate::dns::authority::Authority;
use crate::dns::cache::SynchronizedCache;
use crate::dns::client::tests::{DnsStubClient, StubCallback};
use super::*;
pub fn create_test_context(callback: Box<StubCallback>) -> Arc<ServerContext> {
Arc::new(ServerContext {
authority: Authority::new(),
@@ -130,11 +125,8 @@ pub mod tests {
enable_udp: true,
enable_tcp: true,
enable_api: false,
statistics: ServerStatistics {
tcp_query_count: AtomicUsize::new(0),
udp_query_count: AtomicUsize::new(0),
},
zones_dir: "zones",
statistics: ServerStatistics { tcp_query_count: AtomicUsize::new(0), udp_query_count: AtomicUsize::new(0) },
zones_dir: "zones"
})
}
}
src/dns/filter.rs
+2 -4
View File
@@ -1,12 +1,10 @@
use crate::dns::protocol::{QueryType, DnsPacket};
use crate::dns::protocol::{DnsPacket, QueryType};
pub trait DnsFilter {
fn lookup(&self, qname: &str, qtype: QueryType) -> Option<DnsPacket>;
}
pub struct DummyFilter {
}
pub struct DummyFilter {}
#[allow(unused_variables)]
impl DnsFilter for DummyFilter {
src/dns/hosts.rs
+14 -10
View File
@@ -1,11 +1,12 @@
use std::net::IpAddr;
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use crate::dns::filter::DnsFilter;
use crate::dns::protocol::{DnsPacket, QueryType, DnsRecord, TransientTtl, DnsQuestion};
use std::net::IpAddr;
const NAME_SERVER: & str = "hosts";
use crate::dns::filter::DnsFilter;
use crate::dns::protocol::{DnsPacket, DnsQuestion, DnsRecord, QueryType, TransientTtl};
const NAME_SERVER: &str = "hosts";
pub struct HostsFilter {
hosts: HashMap<String, Vec<IpAddr>>
@@ -33,16 +34,14 @@ impl HostsFilter {
let domain = parts[1].trim().to_owned();
if let Ok(addr) = ip.parse::<IpAddr>() {
if !domain.is_empty() {
map.entry(domain).or_insert(vec!(addr));
map.entry(domain).or_insert(vec![addr]);
}
}
}
map
}
Err(..) => {
HashMap::new()
}
Err(..) => HashMap::new()
};
HostsFilter { hosts }
}
@@ -70,7 +69,11 @@ impl DnsFilter for HostsFilter {
packet.header.authoritative_answer = true;
packet.questions.push(DnsQuestion::new(String::from(qname), qtype));
packet.authorities.push(DnsRecord::NS { domain: String::from("hosts"), host: String::from(NAME_SERVER), ttl: TransientTtl(600) });
packet.authorities.push(DnsRecord::NS {
domain: String::from("hosts"),
host: String::from(NAME_SERVER),
ttl: TransientTtl(600)
});
return Some(packet);
}
@@ -80,9 +83,10 @@ impl DnsFilter for HostsFilter {
#[cfg(test)]
mod tests {
use crate::dns::hosts::HostsFilter;
use std::env;
use crate::dns::hosts::HostsFilter;
#[test]
#[ignore]
pub fn load_hosts() {
src/dns/mod.rs
+2 -2
View File
@@ -19,10 +19,10 @@ pub mod buffer;
pub mod cache;
pub mod client;
pub mod context;
pub mod filter;
pub mod hosts;
pub mod protocol;
pub mod resolve;
pub mod server;
pub mod filter;
pub mod hosts;
mod netutil;
src/dns/protocol.rs
+70 -254
View File
@@ -15,7 +15,7 @@ use crate::dns::buffer::{PacketBuffer, VectorPacketBuffer};
#[derive(Debug, Display, From, Error)]
pub enum ProtocolError {
Buffer(crate::dns::buffer::BufferError),
Io(std::io::Error),
Io(std::io::Error)
}
type Result<T> = std::result::Result<T, ProtocolError>;
@@ -91,10 +91,7 @@ impl PartialOrd<TransientTtl> for TransientTtl {
}
impl Hash for TransientTtl {
fn hash<H>(&self, _: &mut H)
where
H: Hasher,
{
fn hash<H>(&self, _: &mut H) where H: Hasher {
// purposely left empty
}
}
@@ -110,22 +107,22 @@ pub enum DnsRecord {
domain: String,
qtype: u16,
data_len: u16,
ttl: TransientTtl,
ttl: TransientTtl
}, // 0
A {
domain: String,
addr: Ipv4Addr,
ttl: TransientTtl,
ttl: TransientTtl
}, // 1
NS {
domain: String,
host: String,
ttl: TransientTtl,
ttl: TransientTtl
}, // 2
CNAME {
domain: String,
host: String,
ttl: TransientTtl,
ttl: TransientTtl
}, // 5
SOA {
domain: String,
@@ -136,28 +133,28 @@ pub enum DnsRecord {
retry: u32,
expire: u32,
minimum: u32,
ttl: TransientTtl,
ttl: TransientTtl
}, // 6
PTR {
domain: String,
data: String,
ttl: TransientTtl,
ttl: TransientTtl
}, // 12
MX {
domain: String,
priority: u16,
host: String,
ttl: TransientTtl,
ttl: TransientTtl
}, // 15
TXT {
domain: String,
data: String,
ttl: TransientTtl,
ttl: TransientTtl
}, // 16
AAAA {
domain: String,
addr: Ipv6Addr,
ttl: TransientTtl,
ttl: TransientTtl
}, // 28
SRV {
domain: String,
@@ -165,12 +162,12 @@ pub enum DnsRecord {
weight: u16,
port: u16,
host: String,
ttl: TransientTtl,
ttl: TransientTtl
}, // 33
OPT {
packet_len: u16,
flags: u32,
data: String,
data: String
}, // 41
}
@@ -192,14 +189,10 @@ impl DnsRecord {
((raw_addr >> 24) & 0xFF) as u8,
((raw_addr >> 16) & 0xFF) as u8,
((raw_addr >> 8) & 0xFF) as u8,
((raw_addr >> 0) & 0xFF) as u8,
((raw_addr >> 0) & 0xFF) as u8
);
Ok(DnsRecord::A {
domain,
addr,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::A { domain, addr, ttl: TransientTtl(ttl) })
}
QueryType::AAAA => {
let raw_addr1 = buffer.read_u32()?;
@@ -214,34 +207,22 @@ impl DnsRecord {
((raw_addr3 >> 16) & 0xFFFF) as u16,
((raw_addr3 >> 0) & 0xFFFF) as u16,
((raw_addr4 >> 16) & 0xFFFF) as u16,
((raw_addr4 >> 0) & 0xFFFF) as u16,
((raw_addr4 >> 0) & 0xFFFF) as u16
);
Ok(DnsRecord::AAAA {
domain,
addr,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::AAAA { domain, addr, ttl: TransientTtl(ttl) })
}
QueryType::NS => {
let mut ns = String::new();
buffer.read_qname(&mut ns)?;
Ok(DnsRecord::NS {
domain,
host: ns,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::NS { domain, host: ns, ttl: TransientTtl(ttl) })
}
QueryType::CNAME => {
let mut cname = String::new();
buffer.read_qname(&mut cname)?;
Ok(DnsRecord::CNAME {
domain,
host: cname,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::CNAME { domain, host: cname, ttl: TransientTtl(ttl) })
}
QueryType::SRV => {
let priority = buffer.read_u16()?;
@@ -251,36 +232,20 @@ impl DnsRecord {
let mut srv = String::new();
buffer.read_qname(&mut srv)?;
Ok(DnsRecord::SRV {
domain,
priority,
weight,
port,
host: srv,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::SRV { domain, priority, weight, port, host: srv, ttl: TransientTtl(ttl) })
}
QueryType::PTR => {
let mut ptr = String::new();
buffer.read_qname(&mut ptr)?;
Ok(DnsRecord::PTR {
domain,
data: ptr,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::PTR { domain, data: ptr, ttl: TransientTtl(ttl) })
}
QueryType::MX => {
let priority = buffer.read_u16()?;
let mut mx = String::new();
buffer.read_qname(&mut mx)?;
Ok(DnsRecord::MX {
domain,
priority,
host: mx,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::MX { domain, priority, host: mx, ttl: TransientTtl(ttl) })
}
QueryType::SOA => {
let mut m_name = String::new();
@@ -295,58 +260,31 @@ impl DnsRecord {
let expire = buffer.read_u32()?;
let minimum = buffer.read_u32()?;
Ok(DnsRecord::SOA {
domain,
m_name,
r_name,
serial,
refresh,
retry,
expire,
minimum,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::SOA { domain, m_name, r_name, serial, refresh, retry, expire, minimum, ttl: TransientTtl(ttl) })
}
QueryType::TXT => {
let mut txt = String::new();
let cur_pos = buffer.pos();
txt.push_str(&String::from_utf8_lossy(
buffer.get_range(cur_pos, data_len as usize)?,
));
txt.push_str(&String::from_utf8_lossy(buffer.get_range(cur_pos, data_len as usize)?));
buffer.step(data_len as usize)?;
Ok(DnsRecord::TXT {
domain,
data: txt,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::TXT { domain, data: txt, ttl: TransientTtl(ttl) })
}
QueryType::OPT => {
let mut data = String::new();
let cur_pos = buffer.pos();
data.push_str(&String::from_utf8_lossy(
buffer.get_range(cur_pos, data_len as usize)?,
));
data.push_str(&String::from_utf8_lossy(buffer.get_range(cur_pos, data_len as usize)?));
buffer.step(data_len as usize)?;
Ok(DnsRecord::OPT {
packet_len: class,
flags: ttl,
data,
})
Ok(DnsRecord::OPT { packet_len: class, flags: ttl, data })
}
QueryType::UNKNOWN(_) => {
buffer.step(data_len as usize)?;
Ok(DnsRecord::UNKNOWN {
domain,
qtype: qtype_num,
data_len,
ttl: TransientTtl(ttl),
})
Ok(DnsRecord::UNKNOWN { domain, qtype: qtype_num, data_len, ttl: TransientTtl(ttl) })
}
}
}
@@ -355,11 +293,7 @@ impl DnsRecord {
let start_pos = buffer.pos();
match *self {
DnsRecord::A {
ref domain,
ref addr,
ttl: TransientTtl(ttl),
} => {
DnsRecord::A { ref domain, ref addr, ttl: TransientTtl(ttl) } => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::A.to_num())?;
buffer.write_u16(1)?;
@@ -372,11 +306,7 @@ impl DnsRecord {
buffer.write_u8(octets[2])?;
buffer.write_u8(octets[3])?;
}
DnsRecord::AAAA {
ref domain,
ref addr,
ttl: TransientTtl(ttl),
} => {
DnsRecord::AAAA { ref domain, ref addr, ttl: TransientTtl(ttl) } => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::AAAA.to_num())?;
buffer.write_u16(1)?;
@@ -387,11 +317,7 @@ impl DnsRecord {
buffer.write_u16(*octet)?;
}
}
DnsRecord::NS {
ref domain,
ref host,
ttl: TransientTtl(ttl),
} => {
DnsRecord::NS { ref domain, ref host, ttl: TransientTtl(ttl) } => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::NS.to_num())?;
buffer.write_u16(1)?;
@@ -405,11 +331,7 @@ impl DnsRecord {
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::CNAME {
ref domain,
ref host,
ttl: TransientTtl(ttl),
} => {
DnsRecord::CNAME { ref domain, ref host, ttl: TransientTtl(ttl) } => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::CNAME.to_num())?;
buffer.write_u16(1)?;
@@ -423,14 +345,7 @@ impl DnsRecord {
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::SRV {
ref domain,
priority,
weight,
port,
ref host,
ttl: TransientTtl(ttl),
} => {
DnsRecord::SRV { ref domain, priority, weight, port, ref host, ttl: TransientTtl(ttl) } => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::SRV.to_num())?;
buffer.write_u16(1)?;
@@ -461,12 +376,7 @@ impl DnsRecord {
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::MX {
ref domain,
priority,
ref host,
ttl: TransientTtl(ttl),
} => {
DnsRecord::MX { ref domain, priority, ref host, ttl: TransientTtl(ttl) } => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::MX.to_num())?;
buffer.write_u16(1)?;
@@ -481,17 +391,7 @@ impl DnsRecord {
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::SOA {
ref domain,
ref m_name,
ref r_name,
serial,
refresh,
retry,
expire,
minimum,
ttl: TransientTtl(ttl),
} => {
DnsRecord::SOA { ref domain, ref m_name, ref r_name, serial, refresh, retry, expire, minimum, ttl: TransientTtl(ttl) } => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::SOA.to_num())?;
buffer.write_u16(1)?;
@@ -511,11 +411,7 @@ impl DnsRecord {
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::TXT {
ref domain,
ref data,
ttl: TransientTtl(ttl),
} => {
DnsRecord::TXT { ref domain, ref data, ttl: TransientTtl(ttl) } => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::TXT.to_num())?;
buffer.write_u16(1)?;
@@ -590,47 +486,17 @@ impl DnsRecord {
pub fn get_ttl(&self) -> u32 {
match *self {
DnsRecord::A {
ttl: TransientTtl(ttl),
..
}
| DnsRecord::AAAA {
ttl: TransientTtl(ttl),
..
}
| DnsRecord::NS {
ttl: TransientTtl(ttl),
..
}
| DnsRecord::CNAME {
ttl: TransientTtl(ttl),
..
}
| DnsRecord::SRV {
ttl: TransientTtl(ttl),
..
}
| DnsRecord::PTR {
ttl: TransientTtl(ttl),
..
}
| DnsRecord::MX {
ttl: TransientTtl(ttl),
..
}
| DnsRecord::UNKNOWN {
ttl: TransientTtl(ttl),
..
}
| DnsRecord::SOA {
ttl: TransientTtl(ttl),
..
}
| DnsRecord::TXT {
ttl: TransientTtl(ttl),
..
} => ttl,
DnsRecord::OPT { .. } => 0,
DnsRecord::A { ttl: TransientTtl(ttl), .. }
| DnsRecord::AAAA { ttl: TransientTtl(ttl), .. }
| DnsRecord::NS { ttl: TransientTtl(ttl), .. }
| DnsRecord::CNAME { ttl: TransientTtl(ttl), .. }
| DnsRecord::SRV { ttl: TransientTtl(ttl), .. }
| DnsRecord::PTR { ttl: TransientTtl(ttl), .. }
| DnsRecord::MX { ttl: TransientTtl(ttl), .. }
| DnsRecord::UNKNOWN { ttl: TransientTtl(ttl), .. }
| DnsRecord::SOA { ttl: TransientTtl(ttl), .. }
| DnsRecord::TXT { ttl: TransientTtl(ttl), .. } => ttl,
DnsRecord::OPT { .. } => 0
}
}
}
@@ -643,7 +509,7 @@ pub enum ResultCode {
SERVFAIL = 2,
NXDOMAIN = 3,
NOTIMP = 4,
REFUSED = 5,
REFUSED = 5
}
impl Default for ResultCode {
@@ -660,7 +526,7 @@ impl ResultCode {
3 => ResultCode::NXDOMAIN,
4 => ResultCode::NOTIMP,
5 => ResultCode::REFUSED,
0 | _ => ResultCode::NOERROR,
0 | _ => ResultCode::NOERROR
}
}
}
@@ -685,7 +551,7 @@ pub struct DnsHeader {
pub questions: u16, // 16 bits
pub answers: u16, // 16 bits
pub authoritative_entries: u16, // 16 bits
pub resource_entries: u16, // 16 bits
pub resource_entries: u16 // 16 bits
}
impl DnsHeader {
@@ -708,7 +574,7 @@ impl DnsHeader {
questions: 0,
answers: 0,
authoritative_entries: 0,
resource_entries: 0,
resource_entries: 0
}
}
@@ -720,7 +586,7 @@ impl DnsHeader {
| ((self.truncated_message as u8) << 1)
| ((self.authoritative_answer as u8) << 2)
| (self.opcode << 3)
| ((self.response as u8) << 7) as u8,
| ((self.response as u8) << 7) as u8
)?;
buffer.write_u8(
@@ -728,7 +594,7 @@ impl DnsHeader {
| ((self.checking_disabled as u8) << 4)
| ((self.authed_data as u8) << 5)
| ((self.z as u8) << 6)
| ((self.recursion_available as u8) << 7),
| ((self.recursion_available as u8) << 7)
)?;
buffer.write_u16(self.questions)?;
@@ -778,11 +644,7 @@ impl fmt::Display for DnsHeader {
write!(f, "\trecursion_desired: {0}\n", self.recursion_desired)?;
write!(f, "\ttruncated_message: {0}\n", self.truncated_message)?;
write!(
f,
"\tauthoritative_answer: {0}\n",
self.authoritative_answer
)?;
write!(f, "\tauthoritative_answer: {0}\n", self.authoritative_answer)?;
write!(f, "\topcode: {0}\n", self.opcode)?;
write!(f, "\tresponse: {0}\n", self.response)?;
@@ -794,11 +656,7 @@ impl fmt::Display for DnsHeader {
write!(f, "\tquestions: {0}\n", self.questions)?;
write!(f, "\tanswers: {0}\n", self.answers)?;
write!(
f,
"\tauthoritative_entries: {0}\n",
self.authoritative_entries
)?;
write!(f, "\tauthoritative_entries: {0}\n", self.authoritative_entries)?;
write!(f, "\tresource_entries: {0}\n", self.resource_entries)?;
Ok(())
@@ -809,7 +667,7 @@ impl fmt::Display for DnsHeader {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DnsQuestion {
pub name: String,
pub qtype: QueryType,
pub qtype: QueryType
}
impl DnsQuestion {
@@ -818,10 +676,7 @@ impl DnsQuestion {
}
pub fn binary_len(&self) -> usize {
self.name
.split('.')
.map(|x| x.len() + 1)
.fold(1, |x, y| x + y)
self.name.split('.').map(|x| x.len() + 1).fold(1, |x, y| x + y)
}
pub fn write<T: PacketBuffer>(&self, buffer: &mut T) -> Result<()> {
@@ -864,18 +719,12 @@ pub struct DnsPacket {
pub questions: Vec<DnsQuestion>,
pub answers: Vec<DnsRecord>,
pub authorities: Vec<DnsRecord>,
pub resources: Vec<DnsRecord>,
pub resources: Vec<DnsRecord>
}
impl DnsPacket {
pub fn new() -> DnsPacket {
DnsPacket {
header: DnsHeader::new(),
questions: Vec::new(),
answers: Vec::new(),
authorities: Vec::new(),
resources: Vec::new(),
}
DnsPacket { header: DnsHeader::new(), questions: Vec::new(), answers: Vec::new(), authorities: Vec::new(), resources: Vec::new() }
}
pub fn from_buffer<T: PacketBuffer>(buffer: &mut T) -> Result<DnsPacket> {
@@ -977,32 +826,18 @@ impl DnsPacket {
pub fn get_resolved_ns(&self, qname: &str) -> Option<String> {
let mut new_authorities = Vec::new();
for auth in &self.authorities {
if let DnsRecord::NS {
ref domain,
ref host,
..
} = *auth
{
if let DnsRecord::NS { ref domain, ref host, .. } = *auth {
if !qname.ends_with(domain) {
continue;
}
for rsrc in &self.resources {
if let DnsRecord::A {
ref domain,
ref addr,
ttl: TransientTtl(ttl),
} = *rsrc
{
if let DnsRecord::A { ref domain, ref addr, ttl: TransientTtl(ttl) } = *rsrc {
if domain != host {
continue;
}
let rec = DnsRecord::A {
domain: host.clone(),
addr: *addr,
ttl: TransientTtl(ttl),
};
let rec = DnsRecord::A { domain: host.clone(), addr: *addr, ttl: TransientTtl(ttl) };
new_authorities.push(rec);
}
@@ -1023,12 +858,7 @@ impl DnsPacket {
pub fn get_unresolved_ns(&self, qname: &str) -> Option<String> {
let mut new_authorities = Vec::new();
for auth in &self.authorities {
if let DnsRecord::NS {
ref domain,
ref host,
..
} = *auth
{
if let DnsRecord::NS { ref domain, ref host, .. } = *auth {
if !qname.ends_with(domain) {
continue;
}
@@ -1056,13 +886,7 @@ impl DnsPacket {
let mut record_count = self.answers.len() + self.authorities.len() + self.resources.len();
for (i, rec) in self
.answers
.iter()
.chain(self.authorities.iter())
.chain(self.resources.iter())
.enumerate()
{
for (i, rec) in self.answers.iter().chain(self.authorities.iter()).chain(self.resources.iter()).enumerate() {
size += rec.write(&mut test_buffer)?;
if size > max_size {
record_count = i;
@@ -1085,13 +909,7 @@ impl DnsPacket {
question.write(buffer)?;
}
for rec in self
.answers
.iter()
.chain(self.authorities.iter())
.chain(self.resources.iter())
.take(record_count)
{
for rec in self.answers.iter().chain(self.authorities.iter()).chain(self.resources.iter()).take(record_count) {
rec.write(buffer)?;
}
@@ -1111,29 +929,27 @@ mod tests {
packet.header.id = 1337;
packet.header.response = true;
packet
.questions
.push(DnsQuestion::new("google.com".to_string(), QueryType::NS));
packet.questions.push(DnsQuestion::new("google.com".to_string(), QueryType::NS));
//packet.answers.push(DnsRecord::A("ns1.google.com".to_string(), "127.0.0.1".parse::<Ipv4Addr>().unwrap(), 3600));
packet.answers.push(DnsRecord::NS {
domain: "google.com".to_string(),
host: "ns1.google.com".to_string(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
packet.answers.push(DnsRecord::NS {
domain: "google.com".to_string(),
host: "ns2.google.com".to_string(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
packet.answers.push(DnsRecord::NS {
domain: "google.com".to_string(),
host: "ns3.google.com".to_string(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
packet.answers.push(DnsRecord::NS {
domain: "google.com".to_string(),
host: "ns4.google.com".to_string(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
let mut buffer = VectorPacketBuffer::new();
src/dns/resolve.rs
+39 -64
View File
@@ -5,17 +5,17 @@ use std::sync::Arc;
use std::vec::Vec;
use derive_more::{Display, Error, From};
use rand::seq::IteratorRandom;
use crate::dns::context::ServerContext;
use crate::dns::protocol::{DnsPacket, QueryType, ResultCode};
use rand::seq::IteratorRandom;
#[derive(Debug, Display, From, Error)]
pub enum ResolveError {
Client(crate::dns::client::ClientError),
Cache(crate::dns::cache::CacheError),
Io(std::io::Error),
NoServerFound,
NoServerFound
}
type Result<T> = std::result::Result<T, ResolveError>;
@@ -69,7 +69,7 @@ pub trait DnsResolver {
/// This resolver uses an external DNS server to service a query
pub struct ForwardingDnsResolver {
context: Arc<ServerContext>,
upstreams: Vec<String>,
upstreams: Vec<String>
}
impl ForwardingDnsResolver {
@@ -87,9 +87,7 @@ impl DnsResolver for ForwardingDnsResolver {
let mut random = rand::thread_rng();
let upstream = self.upstreams.iter().choose(&mut random).unwrap();
let result = match self.context.cache.lookup(qname, qtype) {
None => {
self.context.client.send_query(qname, qtype, upstream, true)?
}
None => self.context.client.send_query(qname, qtype, upstream, true)?,
Some(packet) => packet
};
@@ -103,7 +101,7 @@ impl DnsResolver for ForwardingDnsResolver {
///
/// This resolver can answer any request using the root servers of the internet
pub struct RecursiveDnsResolver {
context: Arc<ServerContext>,
context: Arc<ServerContext>
}
impl RecursiveDnsResolver {
@@ -139,7 +137,7 @@ impl DnsResolver for RecursiveDnsResolver {
tentative_ns = Some(addr);
break;
}
None => continue,
None => continue
}
}
@@ -152,10 +150,7 @@ impl DnsResolver for RecursiveDnsResolver {
let ns_copy = ns.clone();
let server = format!("{}:{}", ns_copy.as_str(), 53);
let response = self
.context
.client
.send_query(qname, qtype.clone(), &server, false)?;
let response = self.context.client.send_query(qname, qtype.clone(), &server, false)?;
// If we've got an actual answer, we're done!
if !response.answers.is_empty() && response.header.rescode == ResultCode::NOERROR {
@@ -187,7 +182,7 @@ impl DnsResolver for RecursiveDnsResolver {
// If not, we'll have to resolve the ip of a NS record
let new_ns_name = match response.get_unresolved_ns(qname) {
Some(x) => x,
None => return Ok(response.clone()),
None => return Ok(response.clone())
};
// Recursively resolve the NS
@@ -208,12 +203,10 @@ mod tests {
use std::sync::Arc;
use crate::dns::protocol::{DnsPacket, DnsRecord, QueryType, ResultCode, TransientTtl};
use super::*;
use crate::dns::context::tests::create_test_context;
use crate::dns::context::ResolveStrategy;
use crate::dns::protocol::{DnsPacket, DnsRecord, QueryType, ResultCode, TransientTtl};
#[test]
fn test_forwarding_resolver() {
@@ -224,7 +217,7 @@ mod tests {
packet.answers.push(DnsRecord::A {
domain: "google.com".to_string(),
addr: "127.0.0.1".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
} else {
packet.header.rescode = ResultCode::NXDOMAIN;
@@ -235,11 +228,9 @@ mod tests {
match Arc::get_mut(&mut context) {
Some(mut ctx) => {
ctx.resolve_strategy = ResolveStrategy::Forward {
upstreams: vec![String::from("127.0.0.1:53")]
};
ctx.resolve_strategy = ResolveStrategy::Forward { upstreams: vec![String::from("127.0.0.1:53")] };
}
None => panic!(),
None => panic!()
}
let mut resolver = context.create_resolver(Arc::clone(&context));
@@ -248,7 +239,7 @@ mod tests {
{
let res = match resolver.resolve("google.com", QueryType::A, true) {
Ok(x) => x,
Err(_) => panic!(),
Err(_) => panic!()
};
assert_eq!(1, res.answers.len());
@@ -257,7 +248,7 @@ mod tests {
DnsRecord::A { ref domain, .. } => {
assert_eq!("google.com", domain);
}
_ => panic!(),
_ => panic!()
}
};
@@ -266,14 +257,14 @@ mod tests {
{
let res = match resolver.resolve("google.com", QueryType::A, true) {
Ok(x) => x,
Err(_) => panic!(),
Err(_) => panic!()
};
assert_eq!(1, res.answers.len());
let list = match context.cache.list() {
Ok(x) => x,
Err(_) => panic!(),
Err(_) => panic!()
};
assert_eq!(1, list.len());
@@ -287,7 +278,7 @@ mod tests {
{
let res = match resolver.resolve("yahoo.com", QueryType::A, true) {
Ok(x) => x,
Err(_) => panic!(),
Err(_) => panic!()
};
assert_eq!(0, res.answers.len());
@@ -328,11 +319,7 @@ mod tests {
// Insert name server, but no corresponding A record
let mut nameservers = Vec::new();
nameservers.push(DnsRecord::NS {
domain: "".to_string(),
host: "a.myroot.net".to_string(),
ttl: TransientTtl(3600),
});
nameservers.push(DnsRecord::NS { domain: "".to_string(), host: "a.myroot.net".to_string(), ttl: TransientTtl(3600) });
let _ = context.cache.store(&nameservers);
@@ -352,7 +339,7 @@ mod tests {
packet.answers.push(DnsRecord::A {
domain: "a.google.com".to_string(),
addr: "127.0.0.1".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
return Ok(packet);
@@ -362,7 +349,7 @@ mod tests {
packet.answers.push(DnsRecord::A {
domain: "b.google.com".to_string(),
addr: "127.0.0.1".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
return Ok(packet);
@@ -372,7 +359,7 @@ mod tests {
packet.answers.push(DnsRecord::A {
domain: "c.google.com".to_string(),
addr: "127.0.0.1".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
return Ok(packet);
@@ -393,15 +380,11 @@ mod tests {
// Insert root servers
{
let mut nameservers = Vec::new();
nameservers.push(DnsRecord::NS {
domain: "".to_string(),
host: "a.myroot.net".to_string(),
ttl: TransientTtl(3600),
});
nameservers.push(DnsRecord::NS { domain: "".to_string(), host: "a.myroot.net".to_string(), ttl: TransientTtl(3600) });
nameservers.push(DnsRecord::A {
domain: "a.myroot.net".to_string(),
addr: "127.0.0.1".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
let _ = context.cache.store(&nameservers);
@@ -411,21 +394,17 @@ mod tests {
Ok(packet) => {
assert_eq!(1, packet.header.id);
}
Err(_) => panic!(),
Err(_) => panic!()
}
// Insert TLD servers
{
let mut nameservers = Vec::new();
nameservers.push(DnsRecord::NS {
domain: "com".to_string(),
host: "a.mytld.net".to_string(),
ttl: TransientTtl(3600),
});
nameservers.push(DnsRecord::NS { domain: "com".to_string(), host: "a.mytld.net".to_string(), ttl: TransientTtl(3600) });
nameservers.push(DnsRecord::A {
domain: "a.mytld.net".to_string(),
addr: "127.0.0.2".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
let _ = context.cache.store(&nameservers);
@@ -435,7 +414,7 @@ mod tests {
Ok(packet) => {
assert_eq!(2, packet.header.id);
}
Err(_) => panic!(),
Err(_) => panic!()
}
// Insert authoritative servers
@@ -444,12 +423,12 @@ mod tests {
nameservers.push(DnsRecord::NS {
domain: "google.com".to_string(),
host: "ns1.google.com".to_string(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
nameservers.push(DnsRecord::A {
domain: "ns1.google.com".to_string(),
addr: "127.0.0.3".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
let _ = context.cache.store(&nameservers);
@@ -459,7 +438,7 @@ mod tests {
Ok(packet) => {
assert_eq!(3, packet.header.id);
}
Err(_) => panic!(),
Err(_) => panic!()
}
}
@@ -472,7 +451,7 @@ mod tests {
packet.answers.push(DnsRecord::A {
domain: "google.com".to_string(),
addr: "127.0.0.1".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
} else {
packet.header.rescode = ResultCode::NXDOMAIN;
@@ -486,7 +465,7 @@ mod tests {
retry: 3600,
expire: 3600,
minimum: 3600,
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
}
@@ -497,15 +476,11 @@ mod tests {
// Insert name servers
let mut nameservers = Vec::new();
nameservers.push(DnsRecord::NS {
domain: "google.com".to_string(),
host: "ns1.google.com".to_string(),
ttl: TransientTtl(3600),
});
nameservers.push(DnsRecord::NS { domain: "google.com".to_string(), host: "ns1.google.com".to_string(), ttl: TransientTtl(3600) });
nameservers.push(DnsRecord::A {
domain: "ns1.google.com".to_string(),
addr: "127.0.0.1".parse().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
let _ = context.cache.store(&nameservers);
@@ -514,7 +489,7 @@ mod tests {
{
let res = match resolver.resolve("google.com", QueryType::A, true) {
Ok(x) => x,
Err(_) => panic!(),
Err(_) => panic!()
};
assert_eq!(1, res.answers.len());
@@ -523,7 +498,7 @@ mod tests {
DnsRecord::A { ref domain, .. } => {
assert_eq!("google.com", domain);
}
_ => panic!(),
_ => panic!()
}
};
@@ -531,7 +506,7 @@ mod tests {
{
let res = match resolver.resolve("foobar.google.com", QueryType::A, true) {
Ok(x) => x,
Err(_) => panic!(),
Err(_) => panic!()
};
assert_eq!(ResultCode::NXDOMAIN, res.header.rescode);
@@ -542,7 +517,7 @@ mod tests {
{
let res = match resolver.resolve("google.com", QueryType::A, true) {
Ok(x) => x,
Err(_) => panic!(),
Err(_) => panic!()
};
assert_eq!(1, res.answers.len());
@@ -552,7 +527,7 @@ mod tests {
{
let list = match context.cache.list() {
Ok(x) => x,
Err(_) => panic!(),
Err(_) => panic!()
};
assert_eq!(3, list.len());
src/dns/server.rs
+28 -53
View File
@@ -2,16 +2,15 @@
use std::collections::VecDeque;
use std::io::Write;
use std::net::SocketAddr;
use std::net::{Shutdown, TcpListener, TcpStream, UdpSocket};
use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream, UdpSocket};
use std::sync::atomic::Ordering;
use std::sync::mpsc::{channel, Sender};
use std::sync::{Arc, Condvar, Mutex};
use std::thread::Builder;
use derive_more::{Display, Error, From};
use log::{debug, error, warn};
use rand::random;
use log::{error, warn, debug};
use crate::dns::buffer::{BytePacketBuffer, PacketBuffer, StreamPacketBuffer, VectorPacketBuffer};
use crate::dns::context::ServerContext;
@@ -21,7 +20,7 @@ use crate::dns::resolve::DnsResolver;
#[derive(Debug, Display, From, Error)]
pub enum ServerError {
Io(std::io::Error),
Io(std::io::Error)
}
type Result<T> = std::result::Result<T, ServerError>;
@@ -61,7 +60,7 @@ pub trait DnsServer {
/// Utility function for resolving domains referenced in for example CNAME or SRV
/// records. This usually spares the client from having to perform additional lookups.
fn resolve_cnames(lookup_list: &[DnsRecord], results: &mut Vec<DnsPacket>, resolver: &mut Box<dyn DnsResolver>, depth: u16,) {
fn resolve_cnames(lookup_list: &[DnsRecord], results: &mut Vec<DnsPacket>, resolver: &mut Box<dyn DnsResolver>, depth: u16) {
if depth > 10 {
return;
}
@@ -161,17 +160,12 @@ pub struct DnsUdpServer {
context: Arc<ServerContext>,
request_queue: Arc<Mutex<VecDeque<(SocketAddr, DnsPacket)>>>,
request_cond: Arc<Condvar>,
thread_count: usize,
thread_count: usize
}
impl DnsUdpServer {
pub fn new(context: Arc<ServerContext>, thread_count: usize) -> DnsUdpServer {
DnsUdpServer {
context,
request_queue: Arc::new(Mutex::new(VecDeque::new())),
request_cond: Arc::new(Condvar::new()),
thread_count,
}
DnsUdpServer { context, request_queue: Arc::new(Mutex::new(VecDeque::new())), request_cond: Arc::new(Condvar::new()), thread_count }
}
}
@@ -292,7 +286,7 @@ impl DnsServer for DnsUdpServer {
pub struct DnsTcpServer {
context: Arc<ServerContext>,
senders: Vec<Sender<TcpStream>>,
thread_count: usize,
thread_count: usize
}
impl DnsTcpServer {
@@ -318,7 +312,7 @@ impl DnsServer for DnsTcpServer {
loop {
let mut stream = match rx.recv() {
Ok(x) => x,
Err(_) => continue,
Err(_) => continue
};
let _ = context.statistics.tcp_query_count.fetch_add(1, Ordering::Release);
@@ -392,22 +386,16 @@ mod tests {
use std::net::Ipv4Addr;
use std::sync::Arc;
use crate::dns::protocol::{
DnsPacket, DnsQuestion, DnsRecord, QueryType, ResultCode, TransientTtl,
};
use super::*;
use crate::dns::context::tests::create_test_context;
use crate::dns::context::ResolveStrategy;
use crate::dns::protocol::{DnsPacket, DnsQuestion, DnsRecord, QueryType, ResultCode, TransientTtl};
fn build_query(qname: &str, qtype: QueryType) -> DnsPacket {
let mut query_packet = DnsPacket::new();
query_packet.header.recursion_desired = true;
query_packet
.questions
.push(DnsQuestion::new(qname.into(), qtype));
query_packet.questions.push(DnsQuestion::new(qname.into(), qtype));
query_packet
}
@@ -422,30 +410,30 @@ mod tests {
packet.answers.push(DnsRecord::A {
domain: "google.com".to_string(),
addr: "127.0.0.1".parse::<Ipv4Addr>().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
} else if qname == "www.facebook.com" && qtype == QueryType::CNAME {
packet.answers.push(DnsRecord::CNAME {
domain: "www.facebook.com".to_string(),
host: "cdn.facebook.com".to_string(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
packet.answers.push(DnsRecord::A {
domain: "cdn.facebook.com".to_string(),
addr: "127.0.0.1".parse::<Ipv4Addr>().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
} else if qname == "www.microsoft.com" && qtype == QueryType::CNAME {
packet.answers.push(DnsRecord::CNAME {
domain: "www.microsoft.com".to_string(),
host: "cdn.microsoft.com".to_string(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
} else if qname == "cdn.microsoft.com" && qtype == QueryType::A {
packet.answers.push(DnsRecord::A {
domain: "cdn.microsoft.com".to_string(),
addr: "127.0.0.1".parse::<Ipv4Addr>().unwrap(),
ttl: TransientTtl(3600),
ttl: TransientTtl(3600)
});
} else {
packet.header.rescode = ResultCode::NXDOMAIN;
@@ -456,11 +444,9 @@ mod tests {
match Arc::get_mut(&mut context) {
Some(mut ctx) => {
ctx.resolve_strategy = ResolveStrategy::Forward {
upstreams: vec![String::from("127.0.0.1:53")]
};
ctx.resolve_strategy = ResolveStrategy::Forward { upstreams: vec![String::from("127.0.0.1:53")] };
}
None => panic!(),
None => panic!()
}
// A successful resolve
@@ -472,53 +458,47 @@ mod tests {
DnsRecord::A { ref domain, .. } => {
assert_eq!("google.com", domain);
}
_ => panic!(),
_ => panic!()
}
};
// A successful resolve, that also resolves a CNAME without recursive lookup
{
let res = execute_query(
Arc::clone(&context),
&build_query("www.facebook.com", QueryType::CNAME),
);
let res = execute_query(Arc::clone(&context), &build_query("www.facebook.com", QueryType::CNAME));
assert_eq!(2, res.answers.len());
match res.answers[0] {
DnsRecord::CNAME { ref domain, .. } => {
assert_eq!("www.facebook.com", domain);
}
_ => panic!(),
_ => panic!()
}
match res.answers[1] {
DnsRecord::A { ref domain, .. } => {
assert_eq!("cdn.facebook.com", domain);
}
_ => panic!(),
_ => panic!()
}
};
// A successful resolve, that also resolves a CNAME through recursive lookup
{
let res = execute_query(
Arc::clone(&context),
&build_query("www.microsoft.com", QueryType::CNAME),
);
let res = execute_query(Arc::clone(&context), &build_query("www.microsoft.com", QueryType::CNAME));
assert_eq!(2, res.answers.len());
match res.answers[0] {
DnsRecord::CNAME { ref domain, .. } => {
assert_eq!("www.microsoft.com", domain);
}
_ => panic!(),
_ => panic!()
}
match res.answers[1] {
DnsRecord::A { ref domain, .. } => {
assert_eq!("cdn.microsoft.com", domain);
}
_ => panic!(),
_ => panic!()
}
};
@@ -534,7 +514,7 @@ mod tests {
Some(mut ctx) => {
ctx.allow_recursive = false;
}
None => panic!(),
None => panic!()
}
// This should generate an error code, since recursive resolves are
@@ -555,19 +535,14 @@ mod tests {
// Now construct a context where the dns client will return a failure
let mut context2 = create_test_context(Box::new(|_, _, _, _| {
Err(crate::dns::client::ClientError::Io(std::io::Error::new(
std::io::ErrorKind::NotFound,
"Fail",
)))
Err(crate::dns::client::ClientError::Io(std::io::Error::new(std::io::ErrorKind::NotFound, "Fail")))
}));
match Arc::get_mut(&mut context2) {
Some(mut ctx) => {
ctx.resolve_strategy = ResolveStrategy::Forward {
upstreams: vec![String::from("127.0.0.1:53")]
};
ctx.resolve_strategy = ResolveStrategy::Forward { upstreams: vec![String::from("127.0.0.1:53")] };
}
None => panic!(),
None => panic!()
}
// We expect this to set the server failure rescode