public-collect/updns
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

add project files

Browse Source
This commit is contained in:
wyhaya
2019-08-24 15:19:48 +08:00
commit 74a2b3f89d
11 changed files with 2239 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

216
src/config.rs Normal file
View File

@@ -0,0 +1,216 @@
use regex::Regex;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::{IpAddr, SocketAddr};
use std::path::Path;
use std::slice::Iter;
lazy_static! {
static ref REG_IGNORE: Regex = Regex::new(r#"^\s*(#.*)?$"#).unwrap();
static ref REG_BIND: Regex = Regex::new(r#"^\s*bind\s+(?P<val>[^\s#]+)"#).unwrap();
static ref REG_PROXY: Regex = Regex::new(r#"^\s*proxy\s+(?P<val>[^\s#]+)"#).unwrap();
// todo
// The path will also contain '#' and ' '
static ref REG_IMPORT: Regex = Regex::new(r#"\s*import\s+(?P<val>(/.*))"#).unwrap();
static ref REG_DOMAIN_IP: Regex = Regex::new(r#"^\s*(?P<val1>[^\s#]+)\s+(?P<val2>[^\s#]+)"#).unwrap();
}
fn cap_socket_addr(reg: &Regex, text: &str) -> Option<Result<SocketAddr, InvalidType>> {
if let Some(cap) = reg.captures(text) {
return match cap.name("val") {
Some(m) => match m.as_str().parse() {
Ok(addr) => Some(Ok(addr)),
Err(_) => Some(Err(InvalidType::SocketAddr)),
},
None => Some(Err(InvalidType::SocketAddr)),
};
}
None
}
fn cap_ip_addr(reg: &Regex, text: &str) -> Option<Result<(Regex, IpAddr), InvalidType>> {
if let Some(cap) = reg.captures(text) {
if let (Some(val1), Some(val2)) = (cap.name("val1"), cap.name("val2")) {
let (val1, val2) = (val1.as_str(), val2.as_str());
if let Ok(ip) = val1.parse() {
return match Regex::new(val2) {
Ok(reg) => Some(Ok((reg, ip))),
Err(_) => Some(Err(InvalidType::Regex)),
};
} else {
let ip = match val2.parse() {
Ok(ip) => ip,
Err(_) => return Some(Err(InvalidType::IpAddr)),
};
let reg = match Regex::new(val1) {
Ok(reg) => reg,
Err(_) => return Some(Err(InvalidType::Regex)),
};
return Some(Ok((reg, ip)));
}
}
return Some(Err(InvalidType::Other));
}
None
}
#[derive(Debug)]
pub struct Config {
file: File,
content: String,
}
#[derive(Debug)]
pub struct Invalid {
pub line: usize,
pub source: String,
pub err: InvalidType,
}
#[derive(Debug)]
pub enum InvalidType {
Regex,
SocketAddr,
IpAddr,
Other,
}
impl Config {
pub fn new<P: AsRef<Path>>(path: P) -> io::Result<Config> {
let mut file = std::fs::OpenOptions::new()
.read(true)
.append(true)
.create(true)
.open(path)?;
let mut content = String::new();
file.read_to_string(&mut content)?;
Ok(Config { file, content })
}
pub fn add(&mut self, domain: &str, ip: &str) -> std::io::Result<()> {
if self.content.ends_with("\n") {
writeln!(self.file, "{} {}", domain, ip)
} else {
writeln!(self.file, "\n{} {}", domain, ip)
}
}
pub fn parse(&mut self) -> io::Result<(Vec<SocketAddr>, Vec<SocketAddr>, Hosts, Vec<Invalid>)> {
let (mut hosts, mut binds, mut proxys, mut errors) =
(Hosts::new(), Vec::new(), Vec::new(), Vec::new());
for (n, line) in self.content.lines().enumerate() {
// ignore
if REG_IGNORE.is_match(&line) {
continue;
}
// bind
if let Some(addr) = cap_socket_addr(&REG_BIND, &line) {
match addr {
Ok(addr) => binds.push(addr),
Err(err) => {
errors.push(Invalid {
line: n + 1,
source: line.to_string(),
err,
});
}
}
continue;
}
// proxy
if let Some(addr) = cap_socket_addr(&REG_PROXY, &line) {
match addr {
Ok(addr) => proxys.push(addr),
Err(err) => {
errors.push(Invalid {
line: n + 1,
source: line.to_string(),
err,
});
}
}
continue;
}
// import
if let Some(cap) = REG_IMPORT.captures(&line) {
if let Some(m) = cap.name("val") {
let (b, p, h, e) = Config::new(m.as_str())?.parse()?;
binds.extend(b);
proxys.extend(p);
hosts.extend(h);
errors.extend(e);
} else {
// todo
}
continue;
}
// host
if let Some(d) = cap_ip_addr(&REG_DOMAIN_IP, &line) {
match d {
Ok((domain, ip)) => hosts.push(domain, ip),
Err(err) => {
errors.push(Invalid {
line: n + 1,
source: line.to_string(),
err,
});
}
}
continue;
}
errors.push(Invalid {
line: n + 1,
source: line.to_string(),
err: InvalidType::Other,
});
}
Ok((binds, proxys, hosts, errors))
}
}
#[derive(Debug)]
pub struct Hosts {
list: Vec<(Regex, IpAddr)>,
}
impl Hosts {
pub fn new() -> Hosts {
Hosts { list: Vec::new() }
}
fn push(&mut self, domain: Regex, ip: IpAddr) {
self.list.push((domain, ip));
}
fn extend(&mut self, hosts: Hosts) {
for item in hosts.list {
self.list.push(item);
}
}
pub fn iter(&mut self) -> Iter<(Regex, IpAddr)> {
self.list.iter()
}
pub fn get(&self, domain: &str) -> Option<&IpAddr> {
for (reg, ip) in &self.list {
if reg.is_match(domain) {
return Some(ip);
}
}
None
}
}

769
src/lib.rs Normal file
View File

@@ -0,0 +1,769 @@
// From : EmilHernvall/dnsguide
// GitHub : https://github.com/EmilHernvall/dnsguide
#![allow(dead_code)]
use std::io::{Error, ErrorKind};
use std::io::{Read, Result};
use std::net::{Ipv4Addr, Ipv6Addr};
pub struct BytePacketBuffer {
pub buf: [u8; 512],
pub pos: usize,
}
impl BytePacketBuffer {
pub fn new() -> BytePacketBuffer {
BytePacketBuffer {
buf: [0; 512],
pos: 0,
}
}
pub fn pos(&self) -> usize {
self.pos
}
fn step(&mut self, steps: usize) -> Result<()> {
self.pos += steps;
Ok(())
}
fn seek(&mut self, pos: usize) -> Result<()> {
self.pos = pos;
Ok(())
}
fn read(&mut self) -> Result<u8> {
if self.pos >= 512 {
return Err(Error::new(ErrorKind::InvalidInput, "End of buffer"));
}
let res = self.buf[self.pos];
self.pos += 1;
Ok(res)
}
fn get(&mut self, pos: usize) -> Result<u8> {
if pos >= 512 {
return Err(Error::new(ErrorKind::InvalidInput, "End of buffer"));
}
Ok(self.buf[pos])
}
pub fn get_range(&mut self, start: usize, len: usize) -> Result<&[u8]> {
if start + len >= 512 {
return Err(Error::new(ErrorKind::InvalidInput, "End of buffer"));
}
Ok(&self.buf[start..start + len as usize])
}
fn read_u16(&mut self) -> Result<u16> {
let res = ((self.read()? as u16) << 8) | (self.read()? as u16);
Ok(res)
}
fn read_u32(&mut self) -> Result<u32> {
let res = ((self.read()? as u32) << 24)
| ((self.read()? as u32) << 16)
| ((self.read()? as u32) << 8)
| ((self.read()? as u32) << 0);
Ok(res)
}
fn read_qname(&mut self, outstr: &mut String) -> Result<()> {
let mut pos = self.pos();
let mut jumped = false;
let mut delim = "";
loop {
let len = self.get(pos)?;
// A two byte sequence, where the two highest bits of the first byte is
// set, represents a offset relative to the start of the buffer. We
// handle this by jumping to the offset, setting a flag to indicate
// that we shouldn't update the shared buffer position once done.
if (len & 0xC0) == 0xC0 {
// When a jump is performed, we only modify the shared buffer
// position once, and avoid making the change later on.
if !jumped {
self.seek(pos + 2)?;
}
let b2 = self.get(pos + 1)? as u16;
let offset = (((len as u16) ^ 0xC0) << 8) | b2;
pos = offset as usize;
jumped = true;
continue;
}
pos += 1;
// Names are terminated by an empty label of length 0
if len == 0 {
break;
}
outstr.push_str(delim);
let str_buffer = self.get_range(pos, len as usize)?;
outstr.push_str(&String::from_utf8_lossy(str_buffer).to_lowercase());
delim = ".";
pos += len as usize;
}
if !jumped {
self.seek(pos)?;
}
Ok(())
}
fn write(&mut self, val: u8) -> Result<()> {
if self.pos >= 512 {
return Err(Error::new(ErrorKind::InvalidInput, "End of buffer"));
}
self.buf[self.pos] = val;
self.pos += 1;
Ok(())
}
fn write_u8(&mut self, val: u8) -> Result<()> {
self.write(val)?;
Ok(())
}
fn write_u16(&mut self, val: u16) -> Result<()> {
self.write((val >> 8) as u8)?;
self.write((val & 0xFF) as u8)?;
Ok(())
}
fn write_u32(&mut self, val: u32) -> Result<()> {
self.write(((val >> 24) & 0xFF) as u8)?;
self.write(((val >> 16) & 0xFF) as u8)?;
self.write(((val >> 8) & 0xFF) as u8)?;
self.write(((val >> 0) & 0xFF) as u8)?;
Ok(())
}
fn write_qname(&mut self, qname: &str) -> Result<()> {
let split_str = qname.split('.').collect::<Vec<&str>>();
for label in split_str {
let len = label.len();
if len > 0x34 {
return Err(Error::new(
ErrorKind::InvalidInput,
"Single label exceeds 63 characters of length",
));
}
self.write_u8(len as u8)?;
for b in label.as_bytes() {
self.write_u8(*b)?;
}
}
self.write_u8(0)?;
Ok(())
}
fn set(&mut self, pos: usize, val: u8) -> Result<()> {
self.buf[pos] = val;
Ok(())
}
fn set_u16(&mut self, pos: usize, val: u16) -> Result<()> {
self.set(pos, (val >> 8) as u8)?;
self.set(pos + 1, (val & 0xFF) as u8)?;
Ok(())
}
}
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum ResultCode {
NOERROR = 0,
FORMERR = 1,
SERVFAIL = 2,
NXDOMAIN = 3,
NOTIMP = 4,
REFUSED = 5,
}
impl ResultCode {
pub fn from_num(num: u8) -> ResultCode {
match num {
1 => ResultCode::FORMERR,
2 => ResultCode::SERVFAIL,
3 => ResultCode::NXDOMAIN,
4 => ResultCode::NOTIMP,
5 => ResultCode::REFUSED,
0 | _ => ResultCode::NOERROR,
}
}
}
#[derive(Clone, Debug)]
pub struct DnsHeader {
pub id: u16, // 16 bits
pub recursion_desired: bool, // 1 bit
pub truncated_message: bool, // 1 bit
pub authoritative_answer: bool, // 1 bit
pub opcode: u8, // 4 bits
pub response: bool, // 1 bit
pub rescode: ResultCode, // 4 bits
pub checking_disabled: bool, // 1 bit
pub authed_data: bool, // 1 bit
pub z: bool, // 1 bit
pub recursion_available: bool, // 1 bit
pub questions: u16, // 16 bits
pub answers: u16, // 16 bits
pub authoritative_entries: u16, // 16 bits
pub resource_entries: u16, // 16 bits
}
impl DnsHeader {
pub fn new() -> DnsHeader {
DnsHeader {
id: 0,
recursion_desired: false,
truncated_message: false,
authoritative_answer: false,
opcode: 0,
response: false,
rescode: ResultCode::NOERROR,
checking_disabled: false,
authed_data: false,
z: false,
recursion_available: false,
questions: 0,
answers: 0,
authoritative_entries: 0,
resource_entries: 0,
}
}
pub fn read(&mut self, buffer: &mut BytePacketBuffer) -> Result<()> {
self.id = buffer.read_u16()?;
let flags = buffer.read_u16()?;
let a = (flags >> 8) as u8;
let b = (flags & 0xFF) as u8;
self.recursion_desired = (a & (1 << 0)) > 0;
self.truncated_message = (a & (1 << 1)) > 0;
self.authoritative_answer = (a & (1 << 2)) > 0;
self.opcode = (a >> 3) & 0x0F;
self.response = (a & (1 << 7)) > 0;
self.rescode = ResultCode::from_num(b & 0x0F);
self.checking_disabled = (b & (1 << 4)) > 0;
self.authed_data = (b & (1 << 5)) > 0;
self.z = (b & (1 << 6)) > 0;
self.recursion_available = (b & (1 << 7)) > 0;
self.questions = buffer.read_u16()?;
self.answers = buffer.read_u16()?;
self.authoritative_entries = buffer.read_u16()?;
self.resource_entries = buffer.read_u16()?;
// Return the constant header size
Ok(())
}
pub fn write(&self, buffer: &mut BytePacketBuffer) -> Result<()> {
buffer.write_u16(self.id)?;
(buffer.write_u8(
(self.recursion_desired as u8)
| ((self.truncated_message as u8) << 1)
| ((self.authoritative_answer as u8) << 2)
| (self.opcode << 3)
| ((self.response as u8) << 7) as u8,
))?;
(buffer.write_u8(
(self.rescode.clone() as u8)
| ((self.checking_disabled as u8) << 4)
| ((self.authed_data as u8) << 5)
| ((self.z as u8) << 6)
| ((self.recursion_available as u8) << 7),
))?;
buffer.write_u16(self.questions)?;
buffer.write_u16(self.answers)?;
buffer.write_u16(self.authoritative_entries)?;
buffer.write_u16(self.resource_entries)?;
Ok(())
}
}
#[derive(PartialEq, Eq, Debug, Clone, Hash, Copy)]
pub enum QueryType {
UNKNOWN(u16),
A, // 1
NS, // 2
CNAME, // 5
MX, // 15
AAAA, // 28
}
impl QueryType {
pub fn to_num(&self) -> u16 {
match *self {
QueryType::UNKNOWN(x) => x,
QueryType::A => 1,
QueryType::NS => 2,
QueryType::CNAME => 5,
QueryType::MX => 15,
QueryType::AAAA => 28,
}
}
pub fn from_num(num: u16) -> QueryType {
match num {
1 => QueryType::A,
2 => QueryType::NS,
5 => QueryType::CNAME,
15 => QueryType::MX,
28 => QueryType::AAAA,
_ => QueryType::UNKNOWN(num),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DnsQuestion {
pub name: String,
pub qtype: QueryType,
}
impl DnsQuestion {
pub fn new(name: String, qtype: QueryType) -> DnsQuestion {
DnsQuestion {
name: name,
qtype: qtype,
}
}
pub fn read(&mut self, buffer: &mut BytePacketBuffer) -> Result<()> {
buffer.read_qname(&mut self.name)?;
self.qtype = QueryType::from_num(buffer.read_u16()?); // qtype
let _ = buffer.read_u16()?; // class
Ok(())
}
pub fn write(&self, buffer: &mut BytePacketBuffer) -> Result<()> {
buffer.write_qname(&self.name)?;
let typenum = self.qtype.to_num();
buffer.write_u16(typenum)?;
buffer.write_u16(1)?;
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[allow(dead_code)]
pub enum DnsRecord {
UNKNOWN {
domain: String,
qtype: u16,
data_len: u16,
ttl: u32,
}, // 0
A {
domain: String,
addr: Ipv4Addr,
ttl: u32,
}, // 1
NS {
domain: String,
host: String,
ttl: u32,
}, // 2
CNAME {
domain: String,
host: String,
ttl: u32,
}, // 5
MX {
domain: String,
priority: u16,
host: String,
ttl: u32,
}, // 15
AAAA {
domain: String,
addr: Ipv6Addr,
ttl: u32,
}, // 28
}
impl DnsRecord {
pub fn read(buffer: &mut BytePacketBuffer) -> Result<DnsRecord> {
let mut domain = String::new();
buffer.read_qname(&mut domain)?;
let qtype_num = buffer.read_u16()?;
let qtype = QueryType::from_num(qtype_num);
let _ = buffer.read_u16()?;
let ttl = buffer.read_u32()?;
let data_len = buffer.read_u16()?;
match qtype {
QueryType::A => {
let raw_addr = buffer.read_u32()?;
let addr = Ipv4Addr::new(
((raw_addr >> 24) & 0xFF) as u8,
((raw_addr >> 16) & 0xFF) as u8,
((raw_addr >> 8) & 0xFF) as u8,
((raw_addr >> 0) & 0xFF) as u8,
);
Ok(DnsRecord::A {
domain: domain,
addr: addr,
ttl: ttl,
})
}
QueryType::AAAA => {
let raw_addr1 = buffer.read_u32()?;
let raw_addr2 = buffer.read_u32()?;
let raw_addr3 = buffer.read_u32()?;
let raw_addr4 = buffer.read_u32()?;
let addr = Ipv6Addr::new(
((raw_addr1 >> 16) & 0xFFFF) as u16,
((raw_addr1 >> 0) & 0xFFFF) as u16,
((raw_addr2 >> 16) & 0xFFFF) as u16,
((raw_addr2 >> 0) & 0xFFFF) as u16,
((raw_addr3 >> 16) & 0xFFFF) as u16,
((raw_addr3 >> 0) & 0xFFFF) as u16,
((raw_addr4 >> 16) & 0xFFFF) as u16,
((raw_addr4 >> 0) & 0xFFFF) as u16,
);
Ok(DnsRecord::AAAA {
domain: domain,
addr: addr,
ttl: ttl,
})
}
QueryType::NS => {
let mut ns = String::new();
buffer.read_qname(&mut ns)?;
Ok(DnsRecord::NS {
domain: domain,
host: ns,
ttl: ttl,
})
}
QueryType::CNAME => {
let mut cname = String::new();
buffer.read_qname(&mut cname)?;
Ok(DnsRecord::CNAME {
domain: domain,
host: cname,
ttl: ttl,
})
}
QueryType::MX => {
let priority = buffer.read_u16()?;
let mut mx = String::new();
buffer.read_qname(&mut mx)?;
Ok(DnsRecord::MX {
domain: domain,
priority: priority,
host: mx,
ttl: ttl,
})
}
QueryType::UNKNOWN(_) => {
buffer.step(data_len as usize)?;
Ok(DnsRecord::UNKNOWN {
domain: domain,
qtype: qtype_num,
data_len: data_len,
ttl: ttl,
})
}
}
}
pub fn write(&self, buffer: &mut BytePacketBuffer) -> Result<usize> {
let start_pos = buffer.pos();
match *self {
DnsRecord::A {
ref domain,
ref addr,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::A.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
buffer.write_u16(4)?;
let octets = addr.octets();
buffer.write_u8(octets[0])?;
buffer.write_u8(octets[1])?;
buffer.write_u8(octets[2])?;
buffer.write_u8(octets[3])?;
}
DnsRecord::NS {
ref domain,
ref host,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::NS.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
let pos = buffer.pos();
buffer.write_u16(0)?;
buffer.write_qname(host)?;
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::CNAME {
ref domain,
ref host,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::CNAME.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
let pos = buffer.pos();
buffer.write_u16(0)?;
buffer.write_qname(host)?;
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::MX {
ref domain,
priority,
ref host,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::MX.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
let pos = buffer.pos();
buffer.write_u16(0)?;
buffer.write_u16(priority)?;
buffer.write_qname(host)?;
let size = buffer.pos() - (pos + 2);
buffer.set_u16(pos, size as u16)?;
}
DnsRecord::AAAA {
ref domain,
ref addr,
ttl,
} => {
buffer.write_qname(domain)?;
buffer.write_u16(QueryType::AAAA.to_num())?;
buffer.write_u16(1)?;
buffer.write_u32(ttl)?;
buffer.write_u16(16)?;
for octet in &addr.segments() {
buffer.write_u16(*octet)?;
}
}
DnsRecord::UNKNOWN { .. } => {
println!("Skipping record: {:?}", self);
}
}
Ok(buffer.pos() - start_pos)
}
}
#[derive(Clone, Debug)]
pub struct DnsPacket {
pub header: DnsHeader,
pub questions: Vec<DnsQuestion>,
pub answers: Vec<DnsRecord>,
pub authorities: 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(),
}
}
pub fn from_buffer(buffer: &mut BytePacketBuffer) -> Result<DnsPacket> {
let mut result = DnsPacket::new();
result.header.read(buffer)?;
for _ in 0..result.header.questions {
let mut question = DnsQuestion::new("".to_string(), QueryType::UNKNOWN(0));
question.read(buffer)?;
result.questions.push(question);
}
for _ in 0..result.header.answers {
let rec = DnsRecord::read(buffer)?;
result.answers.push(rec);
}
for _ in 0..result.header.authoritative_entries {
let rec = DnsRecord::read(buffer)?;
result.authorities.push(rec);
}
for _ in 0..result.header.resource_entries {
let rec = DnsRecord::read(buffer)?;
result.resources.push(rec);
}
Ok(result)
}
pub fn write(&mut self, buffer: &mut BytePacketBuffer) -> Result<()> {
self.header.questions = self.questions.len() as u16;
self.header.answers = self.answers.len() as u16;
self.header.authoritative_entries = self.authorities.len() as u16;
self.header.resource_entries = self.resources.len() as u16;
self.header.write(buffer)?;
for question in &self.questions {
question.write(buffer)?;
}
for rec in &self.answers {
rec.write(buffer)?;
}
for rec in &self.authorities {
rec.write(buffer)?;
}
for rec in &self.resources {
rec.write(buffer)?;
}
Ok(())
}
pub fn get_random_a(&self) -> Option<String> {
if !self.answers.is_empty() {
let a_record = &self.answers[0];
if let DnsRecord::A { ref addr, .. } = *a_record {
return Some(addr.to_string());
}
}
None
}
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 !qname.ends_with(domain) {
continue;
}
for rsrc in &self.resources {
if let DnsRecord::A {
ref domain,
ref addr,
ttl,
} = *rsrc
{
if domain != host {
continue;
}
let rec = DnsRecord::A {
domain: host.clone(),
addr: *addr,
ttl: ttl,
};
new_authorities.push(rec);
}
}
}
}
if !new_authorities.is_empty() {
if let DnsRecord::A { addr, .. } = new_authorities[0] {
return Some(addr.to_string());
}
}
None
}
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 !qname.ends_with(domain) {
continue;
}
new_authorities.push(host);
}
}
if !new_authorities.is_empty() {
return Some(new_authorities[0].clone());
}
None
}
}

348
src/main.rs Normal file
View File

@@ -0,0 +1,348 @@
#![feature(const_vec_new)]
#[macro_use]
extern crate lazy_static;
mod config;
mod lib;
mod watch;
use ace::App;
use async_std::io;
use async_std::net::UdpSocket;
use async_std::task;
use config::{Config, Hosts, Invalid, InvalidType};
use dirs;
use lib::*;
use std::env;
use std::net::{IpAddr, SocketAddr};
use std::path::PathBuf;
use std::process::Command;
use std::time::{Duration, Instant};
use watch::Watch;
const CONFIG_NAME: &str = ".updns";
const DEFAULT_BIND: &str = "0.0.0.0:53";
const DEFAULT_PROXY: [&str; 2] = ["8.8.8.8:53", "114.114.114.114:53"];
const PROXY_TIMEOUT: u64 = 2000;
const WATCH_INTERVAL: u64 = 3000;
static mut PROXY: Vec<SocketAddr> = Vec::new();
static mut HOSTS: Option<Hosts> = None;
macro_rules! log {
($($arg:tt)*) => {
println!($($arg)*);
};
}
macro_rules! warn {
($($arg:tt)*) => {
print!("\x1B[{}m{}\x1B[0m", "1;33", "warning: ");
println!($($arg)*);
};
}
macro_rules! error {
($($arg:tt)*) => {
eprint!("\x1B[{}m{}\x1B[0m", "1;31", "error: ");
eprintln!($($arg)*);
};
}
macro_rules! exit {
($($arg:tt)*) => {
{
eprint!("\x1B[{}m{}\x1B[0m", "1;31", "error: ");
eprintln!($($arg)*);
std::process::exit(1)
}
};
}
fn main() {
let app = App::new(env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"))
.cmd("add", "Add a DNS record")
.cmd("rm", "Remove a DNS record")
.cmd("ls", "Print all configured DNS records")
.cmd("config", "Call vim to edit the configuration file")
.cmd("path", "Print related directories")
.cmd("help", "Print help information")
.cmd("version", "Print version information")
.opt("-c", "Specify a config file");
let config_path = match app.value("-c") {
Some(values) => {
if values.is_empty() {
exit!("'-c' value: [CONFIG]");
}
PathBuf::from(values[0])
}
None => match dirs::home_dir() {
Some(p) => p.join(CONFIG_NAME),
None => exit!("Can't get home directory"),
},
};
if let Some(cmd) = app.command() {
match cmd.as_str() {
"add" => {
let values = app.value("add").unwrap_or(vec![]);
if values.len() != 2 {
exit!("'add' value: [DOMAIN] [IP]");
}
let mut config = match Config::new(&config_path) {
Ok(c) => c,
Err(err) => exit!("Failed to read config file: {:?}\n{:?}", &config_path, err),
};
if let Err(err) = config.add(&values[0], &values[1]) {
exit!("Add record failed\n{:?}", err);
}
}
"rm" => {
if let Some(value) = app.value("rm") {
if value.is_empty() {
exit!("'rm' value: [DOMAIN | IP]");
}
}
}
"ls" => {
let (_, _, _, mut hosts) = config_parse(&config_path);
let mut n = 0;
for (reg, _) in hosts.iter() {
if reg.as_str().len() > n {
n = reg.as_str().len();
}
}
for (domain, ip) in hosts.iter() {
println!("{:domain$} {}", domain.as_str(), ip, domain = n);
}
}
"config" => {
let cmd = Command::new("vim").arg(&config_path).status();
match cmd {
Ok(status) => {
if status.success() {
config_parse(&config_path);
} else {
warn!("Non-zero state exit\n{:?}", status);
}
}
Err(err) => exit!("Call vim command failed\n{:?}", err),
}
}
"path" => {
let binary = match env::current_exe() {
Ok(p) => p.display().to_string(),
Err(err) => exit!("Failed to get directory\n{:?}", err),
};
println!("Binary: {}\nConfig: {:?}", binary, config_path);
}
"help" => {
app.help();
}
"version" => {
app.version();
}
_ => {
app.error_try("help");
}
}
return;
}
let (_, mut binds, proxys, hosts) = config_parse(&config_path);
if binds.is_empty() {
warn!("Will bind the default address '{}'", DEFAULT_BIND);
binds.push(DEFAULT_BIND.parse().unwrap());
}
if proxys.is_empty() {
warn!(
"Will use the default proxy address '{}'",
DEFAULT_PROXY.join(", ")
);
}
update_config(proxys, hosts);
task::spawn(watch_config(config_path));
task::block_on(run_server(binds));
}
fn config_parse(file: &PathBuf) -> (Config, Vec<SocketAddr>, Vec<SocketAddr>, Hosts) {
let mut config = match Config::new(file) {
Ok(c) => c,
Err(err) => exit!("Failed to read config file: {:?}\n{:?}", file, err),
};
let (binds, proxys, hosts, errors) = match config.parse() {
Ok(d) => d,
Err(err) => exit!("Parsing config file failed\n{:?}", err),
};
output_invalid(errors);
(config, binds, proxys, hosts)
}
fn output_invalid(errors: Vec<Invalid>) {
if !errors.is_empty() {
for invalid in errors {
let msg = match invalid.err {
InvalidType::SocketAddr => "Cannot parse socket addr",
InvalidType::IpAddr => "Cannot parse ip addr",
InvalidType::Regex => "Cannot parse Regular expression",
InvalidType::Other => "Invalid line",
};
warn!("{}", msg);
log!("Line {}: {}", invalid.line, invalid.source);
}
}
}
async fn watch_config(p: PathBuf) {
let mut watch = Watch::new(p, WATCH_INTERVAL);
watch
.change(|c| {
log!("Reload the configuration file: {:?}", &c);
if let Ok(mut config) = Config::new(c) {
if let Ok((_, proxy, hosts, errors)) = config.parse() {
update_config(proxy, hosts);
output_invalid(errors);
}
}
})
.await;
}
fn update_config(mut proxy: Vec<SocketAddr>, hosts: Hosts) {
if proxy.is_empty() {
proxy = DEFAULT_PROXY
.iter()
.map(|p| p.parse().unwrap())
.collect::<Vec<SocketAddr>>();
}
unsafe {
PROXY = proxy;
HOSTS = Some(hosts);
};
}
async fn run_server(binds: Vec<SocketAddr>) {
let mut tasks = vec![];
for addr in binds {
let task = task::spawn(async move {
let socket = match UdpSocket::bind(&addr).await {
Ok(socket) => {
log!("Start listening to '{}'", addr);
socket
}
Err(err) => exit!("Binding '{}' failed\n{:?}", addr, err),
};
loop {
let mut req = BytePacketBuffer::new();
match socket.recv_from(&mut req.buf).await {
Ok((len, src)) => {
let res = match handle(req, len).await {
Ok(data) => data,
Err(err) => {
error!("Processing request failed\n{:?}", err);
continue;
}
};
if let Err(err) = socket.send_to(&res, &src).await {
error!("Replying to '{}' failed\n{:?}", &src, err);
}
}
Err(err) => {
error!("Failed to receive message\n{:?}", err);
}
}
}
});
tasks.push(task);
}
for task in tasks {
task.await;
}
}
async fn proxy(buf: &[u8]) -> io::Result<Vec<u8>> {
let proxy = unsafe { &PROXY };
for addr in proxy.iter() {
let socket = UdpSocket::bind(("0.0.0.0", 0)).await?;
let data = io::timeout(Duration::from_millis(PROXY_TIMEOUT), async {
socket.send_to(&buf, addr).await?;
let mut res = [0; 512];
let len = socket.recv(&mut res).await?;
Ok(res[..len].to_vec())
})
.await;
match data {
Ok(data) => {
return Ok(data);
}
Err(err) => {
error!("Agent request to {}\n{:?}", addr, err);
}
}
}
Err(io::Error::new(
io::ErrorKind::Other,
"Proxy server failed to proxy request",
))
}
fn get_answer(domain: &str, query: QueryType) -> Option<DnsRecord> {
let hosts = unsafe { HOSTS.as_ref().unwrap() };
if let Some(ip) = hosts.get(domain) {
match query {
QueryType::A => {
if let IpAddr::V4(addr) = ip {
return Some(DnsRecord::A {
domain: domain.to_string(),
addr: addr.clone(),
ttl: 3600,
});
}
}
QueryType::AAAA => {
if let IpAddr::V6(addr) = ip {
return Some(DnsRecord::AAAA {
domain: domain.to_string(),
addr: addr.clone(),
ttl: 3600,
});
}
}
_ => {}
}
}
None
}
async fn handle(mut req: BytePacketBuffer, len: usize) -> io::Result<Vec<u8>> {
let mut request = DnsPacket::from_buffer(&mut req)?;
let query = match request.questions.get(0) {
Some(q) => q,
None => return proxy(&req.buf[..len]).await,
};
log!("Query: {} Type: {:?}", query.name, query.qtype);
if let Some(answer) = get_answer(&query.name, query.qtype) {
request.header.recursion_desired = true;
request.header.recursion_available = true;
request.header.response = true;
request.answers.push(answer);
let mut res_buffer = BytePacketBuffer::new();
request.write(&mut res_buffer)?;
let len = res_buffer.pos();
let data = res_buffer.get_range(0, len)?;
Ok(data.to_vec())
} else {
proxy(&req.buf[..len]).await
}
}

46
src/watch.rs Normal file
View File

@@ -0,0 +1,46 @@
use async_std::fs;
use async_std::io;
use async_std::stream;
use async_std::stream::Stream;
use async_std::task;
use std::path::PathBuf;
use std::time::{Duration, SystemTime};
pub struct Watch {
path: PathBuf,
interval: u64,
}
impl Watch {
pub fn new(path: PathBuf, interval: u64) -> Watch {
Watch { interval, path }
}
async fn modified(&self) -> io::Result<SystemTime> {
let file = fs::File::open(&self.path).await?;
let modified = file.metadata().await?.modified()?;
Ok(modified)
}
pub async fn change(&mut self, func: fn(path: &PathBuf)) {
let mut repeat = stream::repeat(0);
let mut before = match self.modified().await {
Ok(time) => Some(time),
Err(_) => None,
};
while let Some(_) = repeat.next().await {
task::sleep(Duration::from_millis(self.interval)).await;
let after = match self.modified().await {
Ok(time) => Some(time),
Err(_) => None,
};
if before != after {
before = after;
func(&self.path);
}
}
}
}