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

Use log crate

Browse Source
This commit is contained in:
wyhaya
2020-01-13 20:55:30 +08:00
parent d4d796d3b1
commit e78cf2ba82
7 changed files with 112 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/config.rs
View File

@@ -1,5 +1,6 @@
use crate::matcher::Matcher;
use futures::future::{BoxFuture, FutureExt};
use lazy_static::lazy_static;
use regex::Regex;
use std::{
borrow::Cow,
@@ -7,6 +8,7 @@ use std::{
path::{Path, PathBuf},
result,
slice::Iter,
time::Duration,
};
use tokio::{
fs,
@@ -18,6 +20,32 @@ lazy_static! {
static ref COMMENT_REGEX: Regex = Regex::new("#.*$").unwrap();
}
// Parse time format into Duration
pub fn try_parse_duration(text: &str) -> result::Result<Duration, ()> {
lazy_static! {
static ref REGEX_MATCH_TIME: Regex =
Regex::new(r"^(?P<time>\d+(\.\d+)?)(?P<unit>d|h|m|s|ms)$").unwrap();
}
let reg: &Regex = &REGEX_MATCH_TIME;
let cap = reg.captures(text).ok_or_else(|| ())?;
let time = cap.name("time").unwrap();
let unit = cap.name("unit").unwrap();
let n = time.as_str().parse::<f64>().map_err(|_| ())?;
let ms = match unit.as_str() {
"d" => 24_f64 * 60_f64 * 60_f64 * 1000_f64 * n,
"h" => 60_f64 * 60_f64 * 1000_f64 * n,
"m" => 60_f64 * 1000_f64 * n,
"s" => 1000_f64 * n,
"ms" => n,
_ => panic!(),
};
Ok(Duration::from_millis(ms as u64))
}
#[derive(Debug)]
pub struct Invalid {
pub line: usize,
@@ -85,7 +113,7 @@ pub struct Config {
pub bind: Vec<SocketAddr>,
pub proxy: Vec<SocketAddr>,
pub hosts: Hosts,
pub timeout: Option<u64>,
pub timeout: Option<Duration>,
pub invalid: Vec<Invalid>,
}
@@ -228,7 +256,7 @@ impl Parser {
Ok(addr) => config.proxy.push(addr),
Err(_) => invalid!(InvalidType::SocketAddr),
},
"timeout" => match value.parse::<u64>() {
"timeout" => match try_parse_duration(value) {
Ok(timeout) => config.timeout = Some(timeout),
Err(_) => invalid!(InvalidType::Timeout),
},

26
src/logger.rs Normal file
View File

@@ -0,0 +1,26 @@
use log::{LevelFilter, Metadata, Record, SetLoggerError};
static LOGGER: Logger = Logger;
pub fn init() -> Result<(), SetLoggerError> {
log::set_logger(&LOGGER).map(|()| log::set_max_level(LevelFilter::Trace))
}
struct Logger;
impl log::Log for Logger {
fn enabled(&self, _: &Metadata) -> bool {
true
}
fn log(&self, record: &Record) {
println!(
"[{}] {:<5} {}",
time::now().strftime("%F %T").unwrap(),
record.level(),
record.args()
);
}
fn flush(&self) {}
}

66
src/main.rs
View File

@@ -1,8 +1,6 @@
#[macro_use]
extern crate lazy_static;
mod config;
mod lib;
mod logger;
mod matcher;
mod watch;
@@ -11,6 +9,7 @@ use config::{Config, Hosts, Invalid, Parser};
use dirs;
use futures::prelude::*;
use lib::*;
use log::*;
use regex::Regex;
use std::{
env,
@@ -30,13 +29,13 @@ const CONFIG_FILE: [&str; 2] = [".updns", "config"];
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 DEFAULT_TIMEOUT: Duration = Duration::from_millis(2000);
const WATCH_INTERVAL: u64 = 5000;
const WATCH_INTERVAL: Duration = Duration::from_millis(5000);
static mut PROXY: Vec<SocketAddr> = Vec::new();
static mut HOSTS: Option<Hosts> = None;
static mut TIMEOUT: u64 = DEFAULT_TIMEOUT;
static mut TIMEOUT: Duration = DEFAULT_TIMEOUT;
macro_rules! exit {
($($arg:tt)*) => {
@@ -46,27 +45,11 @@ macro_rules! exit {
}
};
}
macro_rules! error {
($($arg:tt)*) => {
eprint!("{} ERROR ", time::now().strftime("[%Y-%m-%d %H:%M:%S]").unwrap());
eprintln!($($arg)*);
};
}
macro_rules! info {
($($arg:tt)*) => {
print!("{} INFO ", time::now().strftime("[%Y-%m-%d %H:%M:%S]").unwrap());
println!($($arg)*);
};
}
macro_rules! warn {
($($arg:tt)*) => {
print!("{} WARN ", time::now().strftime("[%Y-%m-%d %H:%M:%S]").unwrap());
println!($($arg)*);
};
}
#[tokio::main]
async fn main() {
let _ = logger::init();
let app = App::new(env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"))
.cmd("add", "Add a DNS record")
.cmd("ls", "Print all configured DNS records")
@@ -75,12 +58,18 @@ async fn main() {
.cmd("help", "Print help information")
.cmd("version", "Print version information")
.opt("-c", "Specify a config file")
.opt("-w", "Check the interval of the configuration file (ms)");
.opt(
"-i",
vec![
"Check the interval time of the configuration file",
"format: 1ms, 1s, 1m, 1h, 1d",
],
);
let config_path = match app.value("-c") {
Some(values) => {
if values.is_empty() {
exit!("'-c' value: [CONFIG]");
exit!("'-c' missing a value: [FILE]");
}
PathBuf::from(values[0])
}
@@ -91,14 +80,17 @@ async fn main() {
};
// Check profile interval
let watch_interval = match app.value("-w") {
let watch_interval = match app.value("-i") {
Some(values) => {
if values.is_empty() {
exit!("'-w' value: [ms]");
exit!("'-i' missing a value: : [1ms, 1s, 1m, 1h, 1d]");
}
values[0]
.parse::<u64>()
.unwrap_or_else(|_| exit!("Cannot resolve '{}' to number", &values[0]))
config::try_parse_duration(values[0]).unwrap_or_else(|_| {
exit!(
"Cannot resolve '{}' to interval time, format: 1ms, 1s, 1m, 1h, 1d",
&values[0]
)
})
}
None => WATCH_INTERVAL,
};
@@ -166,9 +158,9 @@ async fn main() {
config_path.display()
);
}
"help" => app.help(),
"version" => app.version(),
_ => app.error_try("help"),
"help" => app.print_help(),
"version" => app.print_version(),
_ => app.print_error_try("help"),
}
return;
}
@@ -195,7 +187,7 @@ async fn main() {
watch_config(config_path, watch_interval).await;
}
fn update_config(mut proxy: Vec<SocketAddr>, hosts: Hosts, timeout: Option<u64>) {
fn update_config(mut proxy: Vec<SocketAddr>, hosts: Hosts, timeout: Option<Duration>) {
if proxy.is_empty() {
proxy = DEFAULT_PROXY
.iter()
@@ -234,7 +226,7 @@ fn output_invalid(errors: &[Invalid]) {
}
}
async fn watch_config(p: PathBuf, t: u64) {
async fn watch_config(p: PathBuf, t: Duration) {
let mut watch = Watch::new(&p, t).await;
while let Some(_) = watch.next().await {
@@ -288,7 +280,7 @@ async fn proxy(buf: &[u8]) -> Result<Vec<u8>> {
for addr in proxy.iter() {
let mut socket = UdpSocket::bind(("0.0.0.0", 0)).await?;
let data: Result<Vec<u8>> = timeout(Duration::from_millis(unsafe { TIMEOUT }), async {
let data: Result<Vec<u8>> = timeout(unsafe { TIMEOUT }, async {
socket.send_to(&buf, addr).await?;
let mut res = [0; 512];
let len = socket.recv(&mut res).await?;

4
src/watch.rs
View File

@@ -19,13 +19,13 @@ pub struct Watch {
}
impl Watch {
pub async fn new<P: AsRef<Path>>(path: P, duration: u64) -> Watch {
pub async fn new<P: AsRef<Path>>(path: P, duration: Duration) -> Watch {
let path = path.as_ref().to_path_buf();
Watch {
path: path.clone(),
state: None,
modified: Self::modified(path).await,
timer: interval(Duration::from_millis(duration)),
timer: interval(duration),
}
}