From ba2cb079a9008febf6843b44eec9113fa24d9aed Mon Sep 17 00:00:00 2001 From: "Y. T. Chung" Date: Sun, 23 Sep 2018 23:13:11 +0800 Subject: [PATCH] Uses spinlock for resolver config --- Cargo.lock | 7 +++++ Cargo.toml | 1 + src/lib.rs | 5 ++-- src/relay/dns_resolver.rs | 60 +++++++++++++++++++-------------------- src/relay/tcprelay/mod.rs | 30 +++++++++++--------- 5 files changed, 56 insertions(+), 47 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 860cb45e..5bd98ffc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1036,6 +1036,7 @@ dependencies = [ "ring 0.13.2 (registry+https://github.com/rust-lang/crates.io-index)", "serde_json 1.0.28 (registry+https://github.com/rust-lang/crates.io-index)", "serde_urlencoded 0.5.3 (registry+https://github.com/rust-lang/crates.io-index)", + "spin 0.4.9 (registry+https://github.com/rust-lang/crates.io-index)", "subprocess 0.1.15 (registry+https://github.com/rust-lang/crates.io-index)", "time 0.1.40 (registry+https://github.com/rust-lang/crates.io-index)", "tokio 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)", @@ -1070,6 +1071,11 @@ dependencies = [ "winapi 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "spin" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "stable_deref_trait" version = "1.1.1" @@ -1760,6 +1766,7 @@ dependencies = [ "checksum slab 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)" = "5f9776d6b986f77b35c6cf846c11ad986ff128fe0b2b63a3628e3755e8d3102d" "checksum smallvec 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)" = "153ffa32fd170e9944f7e0838edf824a754ec4c1fc64746fcc9fe1f8fa602e5d" "checksum socket2 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)" = "c4d11a52082057d87cb5caa31ad812f4504b97ab44732cd8359df2e9ff9f48e7" +"checksum spin 0.4.9 (registry+https://github.com/rust-lang/crates.io-index)" = "37b5646825922b96b5d7d676b5bb3458a54498e96ed7b0ce09dc43a07038fea4" "checksum stable_deref_trait 1.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "dba1a27d3efae4351c8051072d619e3ade2820635c3958d826bfea39d59b54c8" "checksum string 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "00caf261d6f90f588f8450b8e1230fa0d5be49ee6140fdfbcb55335aff350970" "checksum strsim 0.7.0 (registry+https://github.com/rust-lang/crates.io-index)" = "bb4f380125926a99e52bc279241539c018323fab05ad6368b56f93d9369ff550" diff --git a/Cargo.toml b/Cargo.toml index d8b4e0a2..bb1561ea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -64,6 +64,7 @@ miscreant = { version = "0.4.0-beta", optional = true } lru-cache = "0.1" dns-parser = "0.7" trust-dns-resolver = { version = "0.10.0-alpha", features = ["dns-over-rustls", "dns-over-https"] } +spin = "0.4" [target.'cfg(unix)'.dependencies] tokio-signal = "0.2" diff --git a/src/lib.rs b/src/lib.rs index 221aa443..ecf12289 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -65,12 +65,10 @@ //! let config = Config::load_from_file("shadowsocks.json", ConfigType::Server).unwrap(); //! run_server(config); //! ``` -//! #![crate_type = "lib"] #![crate_name = "shadowsocks"] - -#![recursion_limit="128"] +#![recursion_limit = "128"] extern crate base64; extern crate byte_string; @@ -99,6 +97,7 @@ extern crate subprocess; extern crate tokio_io; extern crate dns_parser; extern crate lru_cache; +extern crate spin; extern crate tokio; #[cfg(any(unix, windows))] extern crate tokio_signal; diff --git a/src/relay/dns_resolver.rs b/src/relay/dns_resolver.rs index 51ca8a6f..aa624d90 100644 --- a/src/relay/dns_resolver.rs +++ b/src/relay/dns_resolver.rs @@ -3,9 +3,10 @@ use std::io::{self, ErrorKind}; use std::mem; use std::net::SocketAddr; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use futures::Future; +use spin::Mutex; use tokio::runtime::current_thread::Runtime; use trust_dns_resolver::config::ResolverConfig; use trust_dns_resolver::AsyncResolver; @@ -24,11 +25,11 @@ lazy_static! { /// Set address for global DNS resolver /// Must be called before servers are actually run pub fn set_dns_config(addr: ResolverConfig) { - *(&*GLOBAL_DNS_ADDRESS).lock().unwrap() = Some(addr); + *(&mut *GLOBAL_DNS_ADDRESS.lock()) = Some(addr); } fn get_dns_address() -> Option { - (&*GLOBAL_DNS_ADDRESS).lock().unwrap().clone() + GLOBAL_DNS_ADDRESS.lock().clone() } fn init_resolver() -> AsyncResolver { @@ -108,36 +109,35 @@ fn inner_resolve(config: Arc, let owned_addr = addr.to_owned(); let owned_addr2 = owned_addr.clone(); - let resolver = &*GLOBAL_DNS_RESOLVER; + GLOBAL_DNS_RESOLVER.lookup_ip(addr) + .map_err(move |err| { + error!("Failed to resolve {}, err: {}", owned_addr2, err); + io::Error::new(io::ErrorKind::Other, "dns resolve error") + }) + .and_then(move |lookup_result| { + let mut vaddr = Vec::new(); + for ip in lookup_result.iter() { + if check_forbidden { + let forbidden_ip = &config.forbidden_ip; + if forbidden_ip.contains(&ip) { + debug!("Resolved {} => {}, which is skipped by forbidden_ip", + owned_addr, ip); + continue; + } + } + vaddr.push(SocketAddr::new(ip, port)); + } - resolver.lookup_ip(addr) - .map_err(move |err| { - error!("Failed to resolve {}, err: {}", owned_addr2, err); - io::Error::new(io::ErrorKind::Other, "dns resolve error") - }) - .and_then(move |lookup_result| { - let mut vaddr = Vec::new(); - for ip in lookup_result.iter() { - if check_forbidden { - let forbidden_ip = &config.forbidden_ip; - if forbidden_ip.contains(&ip) { - debug!("Resolved {} => {}, which is skipped by forbidden_ip", owned_addr, ip); - continue; - } - } - vaddr.push(SocketAddr::new(ip, port)); - } - - if vaddr.is_empty() { - let err = io::Error::new(ErrorKind::Other, + if vaddr.is_empty() { + let err = io::Error::new(ErrorKind::Other, format!("resolved {} to empty address, all IPs are filtered", owned_addr)); - Err(err) - } else { - debug!("Resolved {} => {:?}", owned_addr, vaddr); - Ok(vaddr) - } - }) + Err(err) + } else { + debug!("Resolved {} => {:?}", owned_addr, vaddr); + Ok(vaddr) + } + }) } /// Resolve address to IP diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index 306b2017..96a4e37f 100644 --- a/src/relay/tcprelay/mod.rs +++ b/src/relay/tcprelay/mod.rs @@ -209,13 +209,15 @@ impl> Future for TcpStreamConnect { match mem::replace(self, TcpStreamConnect::Empty) { TcpStreamConnect::Empty => unreachable!(), - TcpStreamConnect::Connect { last_err, .. } => match last_err { - None => { - let err = io::Error::new(ErrorKind::Other, "connect TCP without any addresses"); - Err(err) + TcpStreamConnect::Connect { last_err, .. } => { + match last_err { + None => { + let err = io::Error::new(ErrorKind::Other, "connect TCP without any addresses"); + Err(err) + } + Some(err) => Err(err), } - Some(err) => Err(err), - }, + } } } } @@ -250,16 +252,16 @@ pub fn proxy_server_handshake(remote_stream: TcpStream, let timeout = *svr_cfg.timeout(); proxy_handshake(remote_stream, svr_cfg).and_then(move |(r_fut, w_fut)| { let w_fut = w_fut.and_then(move |enc_w| { - // Send relay address to remote - let mut buf = BytesMut::with_capacity(relay_addr.len()); - relay_addr.write_to_buf(&mut buf); + // Send relay address to remote + let mut buf = BytesMut::with_capacity(relay_addr.len()); + relay_addr.write_to_buf(&mut buf); - trace!("Got encrypt stream and going to send addr: {:?}, buf: {:?}", - relay_addr, - buf); + trace!("Got encrypt stream and going to send addr: {:?}, buf: {:?}", + relay_addr, + buf); - try_timeout(enc_w.write_all(buf), timeout).map(|(w, _)| w) - }); + try_timeout(enc_w.write_all(buf), timeout).map(|(w, _)| w) + }); Ok((r_fut, boxed_future(w_fut))) })