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]] [[package]]
name = "updns" name = "updns"
version = "0.0.7" version = "0.1.0"
dependencies = [ dependencies = [
"ace 0.0.2 (registry+https://github.com/rust-lang/crates.io-index)", "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)", "dirs 2.0.2 (registry+https://github.com/rust-lang/crates.io-index)",

2
Cargo.toml
View File

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

61
src/config.rs
View File

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

93
src/main.rs
View File

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