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

migrate to tokio

Browse Source
This commit is contained in:
wyhaya
2019-11-12 17:17:34 +08:00
parent 54c990725a
commit e168c99635
5 changed files with 659 additions and 374 deletions
Download Patch File Download Diff File Expand all files Collapse all files

345
src/config.rs
View File

@@ -1,11 +1,11 @@
use futures::future::{BoxFuture, FutureExt};
use regex::Regex;
use std::fs;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::{IpAddr, SocketAddr};
use std::path::{Path, PathBuf};
use std::result;
use std::slice::Iter;
use tokio::fs::{create_dir_all, File, OpenOptions};
use tokio::io::{AsyncReadExt, AsyncWriteExt, Result};
lazy_static! {
static ref REG_IGNORE: Regex = Regex::new(r#"^\s*(#.*)?$"#).unwrap();
@@ -18,7 +18,7 @@ lazy_static! {
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>> {
fn cap_socket_addr(reg: &Regex, text: &str) -> Option<result::Result<SocketAddr, InvalidType>> {
let cap = match reg.captures(text) {
Some(cap) => cap,
None => return None,
@@ -33,7 +33,7 @@ fn cap_socket_addr(reg: &Regex, text: &str) -> Option<Result<SocketAddr, Invalid
}
}
fn cap_ip_addr(text: &str) -> Option<Result<(Regex, IpAddr), InvalidType>> {
fn cap_ip_addr(text: &str) -> Option<result::Result<(Regex, IpAddr), InvalidType>> {
let cap = match (&REG_DOMAIN_IP as &Regex).captures(text) {
Some(cap) => cap,
None => return None,
@@ -85,7 +85,7 @@ pub enum InvalidType {
}
impl InvalidType {
pub fn as_str(&self) -> &str {
pub fn text(&self) -> &str {
match self {
InvalidType::SocketAddr => "Cannot parse socket addr",
InvalidType::IpAddr => "Cannot parse ip addr",
@@ -96,168 +96,6 @@ impl InvalidType {
}
}
#[derive(Debug)]
pub struct ParseConfig {
pub bind: Vec<SocketAddr>,
pub proxy: Vec<SocketAddr>,
pub hosts: Hosts,
pub timeout: Option<u64>,
pub invalid: Vec<Invalid>,
}
#[derive(Debug)]
pub struct Config {
path: PathBuf,
file: File,
}
// todo
// async
impl Config {
pub fn new<P: AsRef<Path>>(path: P) -> io::Result<Config> {
let path = path.as_ref();
if let Some(dir) = path.parent() {
fs::create_dir_all(dir)?;
}
Ok(Config {
file: fs::OpenOptions::new()
.read(true)
.append(true)
.create(true)
.open(path)?,
path: path.to_path_buf(),
})
}
fn read_to_string(&mut self) -> io::Result<String> {
let mut content = String::new();
self.file.read_to_string(&mut content)?;
Ok(content)
}
pub fn add(&mut self, domain: &str, ip: &str) -> io::Result<()> {
if self.read_to_string()?.ends_with("\n") {
writeln!(self.file, "{} {}", domain, ip)
} else {
writeln!(self.file, "\n{} {}", domain, ip)
}
}
pub fn parse(&mut self) -> io::Result<ParseConfig> {
let mut hosts = Hosts::new();
let mut bind = Vec::new();
let mut proxy = Vec::new();
let mut invalid = Vec::new();
let mut timeout = None;
for (n, line) in self.read_to_string()?.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) => bind.push(addr),
Err(kind) => {
invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind,
});
}
}
continue;
}
// proxy
if let Some(addr) = cap_socket_addr(&REG_PROXY, &line) {
match addr {
Ok(addr) => proxy.push(addr),
Err(kind) => {
invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind,
});
}
}
continue;
}
// timeout
if let Some(cap) = REG_TIMEOUT.captures(&line) {
if let Some(time) = cap.name("val") {
if let Ok(t) = time.as_str().parse::<u64>() {
timeout = Some(t);
continue;
}
}
invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind: InvalidType::Timeout,
});
continue;
}
// import
if let Some(cap) = REG_IMPORT.captures(&line) {
if let Some(m) = cap.name("val") {
let mut p = Path::new(m.as_str()).to_path_buf();
if p.is_relative() {
if let Some(parent) = self.path.parent() {
p = parent.join(p);
}
}
let config = Config::new(p)?.parse()?;
bind.extend(config.bind);
proxy.extend(config.proxy);
hosts.extend(config.hosts);
invalid.extend(config.invalid);
} else {
// todo
}
continue;
}
// host
if let Some(d) = cap_ip_addr(&line) {
match d {
Ok((domain, ip)) => hosts.push(domain, ip),
Err(kind) => {
invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind,
});
}
}
continue;
}
invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind: InvalidType::Other,
});
}
Ok(ParseConfig {
bind,
proxy,
hosts,
timeout,
invalid,
})
}
}
#[derive(Debug)]
pub struct Hosts {
record: Vec<(Regex, IpAddr)>,
@@ -291,3 +129,172 @@ impl Hosts {
None
}
}
#[derive(Debug)]
pub struct ParseConfig {
pub bind: Vec<SocketAddr>,
pub proxy: Vec<SocketAddr>,
pub hosts: Hosts,
pub timeout: Option<u64>,
pub invalid: Vec<Invalid>,
}
impl ParseConfig {
fn new() -> ParseConfig {
ParseConfig {
hosts: Hosts::new(),
bind: Vec::new(),
proxy: Vec::new(),
invalid: Vec::new(),
timeout: None,
}
}
fn extend(&mut self, other: Self) {
self.bind.extend(other.bind);
self.proxy.extend(other.proxy);
self.hosts.extend(other.hosts);
self.invalid.extend(other.invalid);
if other.timeout.is_some() {
self.timeout = other.timeout;
}
}
}
#[derive(Debug)]
pub struct Config {
path: PathBuf,
file: File,
}
impl Config {
pub async fn new<P: AsRef<Path>>(path: P) -> Result<Config> {
let path = path.as_ref();
if let Some(dir) = path.parent() {
create_dir_all(dir).await?;
}
Ok(Config {
file: OpenOptions::new()
.read(true)
.append(true)
.create(true)
.open(path)
.await?,
path: path.to_path_buf(),
})
}
async fn read_to_string(&mut self) -> Result<String> {
let mut content = String::new();
self.file.read_to_string(&mut content).await?;
Ok(content)
}
pub async fn add(&mut self, domain: &str, ip: &str) -> Result<usize> {
if self.read_to_string().await?.ends_with("\n") {
self.file
.write(format!("{} {}", domain, ip).as_bytes())
.await
} else {
self.file
.write(format!("\n{} {}", domain, ip).as_bytes())
.await
}
}
pub fn parse(mut self) -> BoxFuture<'static, Result<ParseConfig>> {
async move {
let mut parse = ParseConfig::new();
for (n, line) in self.read_to_string().await?.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) => parse.bind.push(addr),
Err(kind) => parse.invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind,
}),
}
continue;
}
// proxy
if let Some(addr) = cap_socket_addr(&REG_PROXY, &line) {
match addr {
Ok(addr) => parse.proxy.push(addr),
Err(kind) => parse.invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind,
}),
}
continue;
}
// timeout
if let Some(cap) = REG_TIMEOUT.captures(&line) {
if let Some(time) = cap.name("val") {
if let Ok(t) = time.as_str().parse::<u64>() {
parse.timeout = Some(t);
continue;
}
}
parse.invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind: InvalidType::Timeout,
});
continue;
}
// import
if let Some(cap) = REG_IMPORT.captures(&line) {
if let Some(m) = cap.name("val") {
let mut p = Path::new(m.as_str()).to_path_buf();
if p.is_relative() {
if let Some(parent) = self.path.parent() {
p = parent.join(p);
}
}
parse.extend(Config::new(p).await?.parse().await?);
} else {
// todo
}
continue;
}
// host
if let Some(d) = cap_ip_addr(&line) {
match d {
Ok((domain, ip)) => parse.hosts.push(domain, ip),
Err(kind) => parse.invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind,
}),
}
continue;
}
parse.invalid.push(Invalid {
line: n + 1,
source: line.to_string(),
kind: InvalidType::Other,
});
}
Ok(parse)
}
.boxed()
}
}

138
src/main.rs
View File

@@ -6,22 +6,22 @@ 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, ParseConfig};
use dirs;
use lib::*;
use regex::Regex;
use std::env;
use std::env::current_exe;
use std::net::{IpAddr, SocketAddr};
use std::path::PathBuf;
use std::process::Command;
use std::time::Duration;
use tokio::io::{Error, ErrorKind, Result};
use tokio::net::UdpSocket;
use tokio::prelude::*;
use tokio::timer::Timeout;
use watch::Watch;
const CONFIG_FILE: [&'static str; 2] = [".updns", "config"];
const CONFIG_COMMAND: &'static str = "vim";
const DEFAULT_BIND: &'static str = "0.0.0.0:53";
const DEFAULT_PROXY: [&'static str; 2] = ["8.8.8.8:53", "1.1.1.1:53"];
@@ -60,13 +60,12 @@ macro_rules! warn {
};
}
fn main() {
let cct = format!("Call '{}' to edit the configuration file", CONFIG_COMMAND);
#[tokio::main]
async 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", cct.as_str())
.cmd("config", "Call 'vim' to edit the configuration file")
.cmd("path", "Print related directories")
.cmd("help", "Print help information")
.cmd("version", "Print version information")
@@ -119,27 +118,16 @@ fn main() {
exit!("Cannot resolve '{}' to ip address", values[1]);
}
let mut config = match Config::new(&config_path) {
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]) {
if let Err(err) = config.add(&values[0], &values[1]).await {
exit!("Add record failed\n{:?}", err);
}
}
"rm" => {
if let Some(value) = app.value("rm") {
if value.is_empty() {
exit!("'rm' value: [DOMAIN | IP]");
}
match ask("Confirm delete? Y/N\n") {
Ok(_) => println!("todo"),
Err(err) => exit!("{:?}", err),
}
}
}
"ls" => {
let mut config = config_parse(&config_path);
let mut config = config_parse(&config_path).await;
let mut n = 0;
for (reg, _) in config.hosts.iter() {
if reg.as_str().len() > n {
@@ -151,23 +139,20 @@ fn main() {
}
}
"config" => {
let cmd = Command::new(CONFIG_COMMAND).arg(&config_path).status();
let cmd = Command::new("vim").arg(&config_path).status();
match cmd {
Ok(status) => {
if status.success() {
config_parse(&config_path);
config_parse(&config_path).await;
} else {
println!(
"'{}' exits with a non-zero status code: {:?}",
CONFIG_COMMAND, status
);
println!("'vim' exits with a non-zero status code: {:?}", status);
}
}
Err(err) => exit!("Call '{}' command failed\n{:?}", CONFIG_COMMAND, err),
Err(err) => exit!("Call 'vim' command failed\n{:?}", err),
}
}
"path" => {
let binary = match env::current_exe() {
let binary = match current_exe() {
Ok(p) => p.display().to_string(),
Err(err) => exit!("Failed to get directory\n{:?}", err),
};
@@ -177,20 +162,14 @@ fn main() {
config_path.to_string_lossy()
);
}
"help" => {
app.help();
}
"version" => {
app.version();
}
_ => {
app.error_try("help");
}
"help" => app.help(),
"version" => app.version(),
_ => app.error_try("help"),
}
return;
}
let mut parse = config_parse(&config_path);
let mut parse = config_parse(&config_path).await;
if parse.bind.is_empty() {
warn!("Will bind the default address '{}'", DEFAULT_BIND);
parse.bind.push(DEFAULT_BIND.parse().unwrap());
@@ -206,26 +185,10 @@ fn main() {
// Run server
for addr in parse.bind {
task::spawn(run_server(addr.clone()));
tokio::spawn(run_server(addr.clone()));
}
// watch config
task::block_on(watch_config(config_path, watch_interval));
}
fn ask(text: &str) -> io::Result<bool> {
use std::io;
use std::io::Write;
io::stdout().write(text.as_bytes())?;
io::stdout().flush()?;
let mut s = String::new();
io::stdin().read_line(&mut s)?;
match s.to_uppercase().as_str() {
"Y\n" => Ok(true),
"N\n" => Ok(false),
_ => Ok(ask(&text)?),
}
watch_config(config_path, watch_interval).await;
}
fn update_config(mut proxy: Vec<SocketAddr>, hosts: Hosts, timeout: Option<u64>) {
@@ -245,13 +208,13 @@ fn update_config(mut proxy: Vec<SocketAddr>, hosts: Hosts, timeout: Option<u64>)
};
}
fn config_parse(file: &PathBuf) -> ParseConfig {
let mut config = match Config::new(file) {
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),
};
let parse = match config.parse() {
let parse: ParseConfig = match config.parse().await {
Ok(d) => d,
Err(err) => exit!("Parsing config file failed\n{:?}", err),
};
@@ -265,29 +228,27 @@ fn output_invalid(errors: &Vec<Invalid>) {
error!(
"[line:{}] {} `{}`",
invalid.line,
invalid.kind.as_str(),
invalid.kind.text(),
invalid.source
);
}
}
async fn watch_config(p: PathBuf, t: u64) {
let mut watch = Watch::new(p, t);
watch
.for_each(|c| {
info!("Reload the configuration file: {:?}", &c);
if let Ok(mut config) = Config::new(c) {
if let Ok(parse) = config.parse() {
update_config(parse.proxy, parse.hosts, parse.timeout);
output_invalid(&parse.invalid);
}
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);
}
})
.await;
}
}
}
async fn run_server(addr: SocketAddr) {
let socket = match UdpSocket::bind(&addr).await {
let mut socket = match UdpSocket::bind(&addr).await {
Ok(socket) => {
info!("Start listening to '{}'", addr);
socket
@@ -317,19 +278,22 @@ async fn run_server(addr: SocketAddr) {
}
}
async fn proxy(buf: &[u8]) -> io::Result<Vec<u8>> {
async fn proxy(buf: &[u8]) -> Result<Vec<u8>> {
let proxy = unsafe { &PROXY };
for addr in proxy.iter() {
let socket = UdpSocket::bind(("0.0.0.0", 0)).await?;
let mut socket = UdpSocket::bind(("0.0.0.0", 0)).await?;
let data = io::timeout(Duration::from_millis(unsafe { 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;
let data: Result<Vec<u8>> = Timeout::new(
async {
socket.send_to(&buf, addr).await?;
let mut res = [0; 512];
let len = socket.recv(&mut res).await?;
Ok(res[..len].to_vec())
},
Duration::from_millis(unsafe { TIMEOUT }),
)
.await?;
match data {
Ok(data) => {
@@ -341,8 +305,8 @@ async fn proxy(buf: &[u8]) -> io::Result<Vec<u8>> {
}
}
Err(io::Error::new(
io::ErrorKind::Other,
Err(Error::new(
ErrorKind::Other,
"Proxy server failed to proxy request",
))
}
@@ -375,7 +339,7 @@ fn get_answer(domain: &str, query: QueryType) -> Option<DnsRecord> {
None
}
async fn handle(mut req: BytePacketBuffer, len: usize) -> io::Result<Vec<u8>> {
async fn handle(mut req: BytePacketBuffer, len: usize) -> Result<Vec<u8>> {
let mut request = DnsPacket::from_buffer(&mut req)?;
let query = match request.questions.get(0) {

80
src/watch.rs
View File

@@ -1,44 +1,68 @@
use async_std::fs;
use async_std::io;
use async_std::prelude::*;
use async_std::stream;
use std::path::PathBuf;
use futures::ready;
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, SystemTime};
use tokio::fs::File;
use tokio::io::Result;
use tokio::prelude::*;
use tokio::timer::Interval;
pub struct Watch {
path: PathBuf,
interval: u64,
state: Option<Pin<Box<dyn Future<Output = Result<SystemTime>>>>>,
modified: Result<SystemTime>,
timer: Interval,
}
impl Watch {
pub fn new(path: PathBuf, interval: u64) -> Watch {
Watch { interval, path }
pub async fn new<P: AsRef<Path>>(path: P, duration: u64) -> Watch {
let path = path.as_ref().to_path_buf();
Watch {
path: path.clone(),
state: None,
modified: Self::modified(path).await,
timer: Interval::new_interval(Duration::from_millis(duration)),
}
}
async fn modified(&self) -> io::Result<SystemTime> {
let file = fs::File::open(&self.path).await?;
let modified = file.metadata().await?.modified()?;
Ok(modified)
async fn modified(p: PathBuf) -> Result<SystemTime> {
let file = File::open(p).await?;
file.metadata().await?.modified()
}
// todo
// use Stream
pub async fn for_each(&mut self, func: fn(path: &PathBuf)) {
let mut before = match self.modified().await {
Ok(time) => Some(time),
Err(_) => None,
};
fn eq(a: &Result<SystemTime>, b: &Result<SystemTime>) -> bool {
if a.is_ok() && b.is_ok() {
if a.as_ref().ok() == b.as_ref().ok() {
return true;
}
} else if a.is_err() && b.is_err() {
let left = a.as_ref().err().unwrap();
let right = b.as_ref().err().unwrap();
if left.kind() == right.kind() && left.raw_os_error() == right.raw_os_error() {
return true;
}
}
false
}
}
let mut interval = stream::interval(Duration::from_millis(self.interval));
while let Some(_) = interval.next().await {
let after = match self.modified().await {
Ok(time) => Some(time),
Err(_) => None,
};
impl Stream for Watch {
type Item = ();
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
loop {
if let Some(state) = &mut self.state {
let modified: Result<SystemTime> = ready!(Pin::new(state).poll(cx));
self.state = None;
if before != after {
before = after;
func(&self.path);
if !Self::eq(&self.modified, &modified) {
self.modified = modified;
return Poll::Ready(Some(()));
}
} else {
ready!(self.timer.poll_next_unpin(cx));
self.state = Some(Box::pin(Self::modified(self.path.clone())));
}
}
}