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

release 0.1.0

Browse Source
This commit is contained in:
wyhaya
2019-11-25 15:08:52 +08:00
parent dbbc4d0f91
commit 7827521b98
4 changed files with 86 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.lock generated
View File

@@ -891,7 +891,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
[[package]]
name = "updns"
version = "0.0.7"
version = "0.1.0"
dependencies = [
"ace 0.0.2 (registry+https://github.com/rust-lang/crates.io-index)",
"dirs 2.0.2 (registry+https://github.com/rust-lang/crates.io-index)",

2
Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "updns"
version = "0.0.7"
version = "0.1.0"
edition = "2018"
authors = ["wyhaya <wyhaya@gmail.com>"]

67
src/config.rs
View File

@@ -8,7 +8,8 @@ use std::{
slice::Iter,
};
use tokio::{
fs::{create_dir_all, File, OpenOptions},
fs,
fs::{File, OpenOptions},
io::{AsyncReadExt, AsyncWriteExt, Result},
};
@@ -33,7 +34,7 @@ pub enum InvalidType {
}
impl InvalidType {
pub fn text(&self) -> &str {
pub fn as_str(&self) -> &str {
match self {
InvalidType::SocketAddr => "Cannot parse socket address",
InvalidType::IpAddr => "Cannot parse ip address",
@@ -131,7 +132,7 @@ impl Host {
}
#[derive(Debug)]
pub struct ParseConfig {
pub struct Config {
pub bind: Vec<SocketAddr>,
pub proxy: Vec<SocketAddr>,
pub hosts: Hosts,
@@ -139,9 +140,9 @@ pub struct ParseConfig {
pub invalid: Vec<Invalid>,
}
impl ParseConfig {
fn new() -> ParseConfig {
ParseConfig {
impl Config {
fn new() -> Config {
Config {
hosts: Hosts::new(),
bind: Vec::new(),
proxy: Vec::new(),
@@ -162,20 +163,20 @@ impl ParseConfig {
}
#[derive(Debug)]
pub struct Config {
pub struct Parser {
path: PathBuf,
file: File,
}
impl Config {
pub async fn new<P: AsRef<Path>>(path: P) -> Result<Config> {
impl Parser {
pub async fn new<P: AsRef<Path>>(path: P) -> Result<Parser> {
let path = path.as_ref();
if let Some(dir) = path.parent() {
create_dir_all(dir).await?;
fs::create_dir_all(dir).await?;
}
Ok(Config {
Ok(Parser {
file: OpenOptions::new()
.read(true)
.append(true)
@@ -216,21 +217,20 @@ impl Config {
None
}
fn parse_host(key: &str, value: &str) -> result::Result<(Host, IpAddr), InvalidType> {
// match host
// example.com 0.0.0.0
// 0.0.0.0 example.com
// match host
// example.com 0.0.0.0
// 0.0.0.0 example.com
fn record(left: &str, right: &str) -> result::Result<(Host, IpAddr), InvalidType> {
// ip domain
if let Ok(ip) = key.parse() {
return Host::new(value)
if let Ok(ip) = right.parse() {
return Host::new(left)
.map(|host| (host, ip))
.map_err(|_| InvalidType::Regex);
}
// domain ip
if let Ok(ip) = value.parse() {
return Host::new(key)
if let Ok(ip) = left.parse() {
return Host::new(right)
.map(|host| (host, ip))
.map_err(|_| InvalidType::Regex);
}
@@ -238,27 +238,26 @@ impl Config {
Err(InvalidType::IpAddr)
}
pub fn parse(mut self) -> BoxFuture<'static, Result<ParseConfig>> {
pub fn parse(mut self) -> BoxFuture<'static, Result<Config>> {
async move {
let mut parse = ParseConfig::new();
let content = self.read_to_string().await?;
let mut config = Config::new();
for (n, line) in self.read_to_string().await?.lines().enumerate() {
for (i, line) in content.lines().enumerate() {
if line.is_empty() {
continue;
}
// remove comment
// example # ... -> example
let line: Cow<str> = COMMENT_REGEX.replace(line, "");
if line.trim().is_empty() {
continue;
}
macro_rules! invalid {
($type: expr) => {{
parse.invalid.push(Invalid {
line: n + 1,
config.invalid.push(Invalid {
line: i + 1,
source: line.to_string(),
kind: $type,
});
@@ -273,15 +272,15 @@ impl Config {
match key {
"bind" => match value.parse::<SocketAddr>() {
Ok(addr) => parse.bind.push(addr),
Ok(addr) => config.bind.push(addr),
Err(_) => invalid!(InvalidType::SocketAddr),
},
"proxy" => match value.parse::<SocketAddr>() {
Ok(addr) => parse.proxy.push(addr),
Ok(addr) => config.proxy.push(addr),
Err(_) => invalid!(InvalidType::SocketAddr),
},
"timeout" => match value.parse::<u64>() {
Ok(timeout) => parse.timeout = Some(timeout),
Ok(timeout) => config.timeout = Some(timeout),
Err(_) => invalid!(InvalidType::Timeout),
},
"import" => {
@@ -291,16 +290,16 @@ impl Config {
path = parent.join(path);
}
}
parse.extend(Config::new(path).await?.parse().await?);
config.extend(Parser::new(path).await?.parse().await?);
}
_ => match Self::parse_host(key, value) {
Ok(record) => parse.hosts.push(record),
_ => match Self::record(key, value) {
Ok(record) => config.hosts.push(record),
Err(kind) => invalid!(kind),
},
}
}
Ok(parse)
Ok(config)
}
.boxed()
}

101
src/main.rs
View File

@@ -6,7 +6,7 @@ mod lib;
mod watch;
use ace::App;
use config::{Config, Hosts, Invalid, ParseConfig};
use config::{Config, Hosts, Invalid, Parser};
use dirs;
use lib::*;
use regex::Regex;
@@ -31,12 +31,12 @@ const DEFAULT_BIND: &str = "0.0.0.0:53";
const DEFAULT_PROXY: [&str; 2] = ["8.8.8.8:53", "1.1.1.1:53"];
const DEFAULT_TIMEOUT: u64 = 2000;
const WATCH_INTERVAL: u64 = 5000;
static mut PROXY: Vec<SocketAddr> = Vec::new();
static mut HOSTS: Option<Hosts> = None;
static mut TIMEOUT: u64 = DEFAULT_TIMEOUT;
const WATCH_INTERVAL: u64 = 5000;
macro_rules! exit {
($($arg:tt)*) => {
{
@@ -95,10 +95,9 @@ async fn main() {
if values.is_empty() {
exit!("'-w' value: [ms]");
}
match values[0].parse::<u64>() {
Ok(t) => t,
Err(_) => exit!("Cannot resolve '{}' to number", &values[0]),
}
values[0]
.parse::<u64>()
.unwrap_or_else(|_| exit!("Cannot resolve '{}' to number", &values[0]))
}
None => WATCH_INTERVAL,
};
@@ -111,6 +110,7 @@ async fn main() {
exit!("'add' value: [DOMAIN] [IP]");
}
// Check is positive
if let Err(err) = Regex::new(values[0]) {
exit!(
"Cannot resolve '{}' to regular expression\n{:?}",
@@ -122,11 +122,11 @@ async fn main() {
exit!("Cannot resolve '{}' to ip address", values[1]);
}
let mut config = match Config::new(&config_path).await {
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]).await {
let mut parser = Parser::new(&config_path).await.unwrap_or_else(|err| {
exit!("Failed to read config file {:?}\n{:?}", &config_path, err)
});
if let Err(err) = parser.add(values[0], values[1]).await {
exit!("Add record failed\n{:?}", err);
}
}
@@ -144,27 +144,25 @@ async fn main() {
}
}
"config" => {
let cmd = Command::new("vim").arg(&config_path).status();
match cmd {
Ok(status) => {
if status.success() {
config_parse(&config_path).await;
} else {
println!("'vim' exits with a non-zero status code: {:?}", status);
}
}
Err(err) => exit!("Call 'vim' command failed\n{:?}", err),
let status = Command::new("vim")
.arg(&config_path)
.status()
.unwrap_or_else(|err| exit!("Call 'vim' command failed\n{:?}", err));
if status.success() {
config_parse(&config_path).await;
} else {
println!("'vim' exits with a non-zero status code: {:?}", status);
}
}
"path" => {
let binary = match env::current_exe() {
Ok(p) => p.display().to_string(),
Err(err) => exit!("Failed to get directory\n{:?}", err),
};
let binary = env::current_exe()
.unwrap_or_else(|err| exit!("Failed to get directory\n{:?}", err));
println!(
"Binary: {}\nConfig: {}",
binary,
config_path.to_string_lossy()
binary.display(),
config_path.display()
);
}
"help" => app.help(),
@@ -174,22 +172,22 @@ async fn main() {
return;
}
let mut parse = config_parse(&config_path).await;
if parse.bind.is_empty() {
let mut config = config_parse(&config_path).await;
if config.bind.is_empty() {
warn!("Will bind the default address '{}'", DEFAULT_BIND);
parse.bind.push(DEFAULT_BIND.parse().unwrap());
config.bind.push(DEFAULT_BIND.parse().unwrap());
}
if parse.proxy.is_empty() {
if config.proxy.is_empty() {
warn!(
"Will use the default proxy address '{}'",
DEFAULT_PROXY.join(", ")
);
}
update_config(parse.proxy, parse.hosts, parse.timeout);
update_config(config.proxy, config.hosts, config.timeout);
// Run server
for addr in parse.bind {
for addr in config.bind {
tokio::spawn(run_server(addr));
}
// watch config
@@ -210,19 +208,18 @@ fn update_config(mut proxy: Vec<SocketAddr>, hosts: Hosts, timeout: Option<u64>)
};
}
async fn config_parse(file: &PathBuf) -> ParseConfig {
let config = match Config::new(file).await {
Ok(c) => c,
Err(err) => exit!("Failed to read config file {:?}\n{:?}", file, err),
};
async fn config_parse(file: &PathBuf) -> Config {
let parser = Parser::new(file)
.await
.unwrap_or_else(|err| exit!("Failed to read config file {:?}\n{:?}", file, err));
let parse: ParseConfig = match config.parse().await {
Ok(d) => d,
Err(err) => exit!("Parsing config file failed\n{:?}", err),
};
output_invalid(&parse.invalid);
let config: Config = parser
.parse()
.await
.unwrap_or_else(|err| exit!("Parsing config file failed\n{:?}", err));
parse
output_invalid(&config.invalid);
config
}
fn output_invalid(errors: &[Invalid]) {
@@ -230,7 +227,7 @@ fn output_invalid(errors: &[Invalid]) {
error!(
"[line:{}] {} `{}`",
invalid.line,
invalid.kind.text(),
invalid.kind.as_str(),
invalid.source
);
}
@@ -238,12 +235,13 @@ fn output_invalid(errors: &[Invalid]) {
async fn watch_config(p: PathBuf, t: u64) {
let mut watch = Watch::new(&p, t).await;
while let Some(_) = watch.next().await {
info!("Reload the configuration file: {:?}", &p);
if let Ok(config) = Config::new(&p).await {
if let Ok(parse) = config.parse().await {
update_config(parse.proxy, parse.hosts, parse.timeout);
output_invalid(&parse.invalid);
if let Ok(parser) = Parser::new(&p).await {
if let Ok(config) = parser.parse().await {
update_config(config.proxy, config.hosts, config.timeout);
output_invalid(&config.invalid);
}
}
}
@@ -260,6 +258,7 @@ async fn run_server(addr: SocketAddr) {
loop {
let mut req = BytePacketBuffer::new();
let (len, src) = match socket.recv_from(&mut req.buf).await {
Ok(r) => r,
Err(err) => {
@@ -267,6 +266,7 @@ async fn run_server(addr: SocketAddr) {
continue;
}
};
let res = match handle(req, len).await {
Ok(data) => data,
Err(err) => {
@@ -274,6 +274,7 @@ async fn run_server(addr: SocketAddr) {
continue;
}
};
if let Err(err) = socket.send_to(&res, &src).await {
error!("Replying to '{}' failed {:?}", &src, err);
}