diff --git a/src/context.rs b/src/context.rs index 7c8ad7786725..527717844481 100644 --- a/src/context.rs +++ b/src/context.rs @@ -7,16 +7,17 @@ use std::{ atomic::{AtomicBool, Ordering}, Arc, }, + time::Duration, }; #[cfg(feature = "local-dns-relay")] -use std::{net::IpAddr, time::Duration}; +use std::net::IpAddr; use bloomfilter::Bloom; -use log::{log_enabled, warn}; -#[cfg(feature = "local-dns-relay")] +use log::{log_enabled, trace, warn}; use lru_time_cache::LruCache; -use spin::Mutex; +use spin::Mutex as SpinMutex; +use tokio::sync::Mutex as AsyncMutex; #[cfg(feature = "trust-dns")] use trust_dns_resolver::TokioAsyncResolver; @@ -171,7 +172,7 @@ pub struct Context { // Check for duplicated IV/Nonce, for prevent replay attack // https://github.com/shadowsocks/shadowsocks-org/issues/44 - nonce_ppbloom: Mutex, + nonce_ppbloom: SpinMutex, // For Android's flow stat report #[cfg(feature = "local-flow-stat")] @@ -179,11 +180,14 @@ pub struct Context { // For DNS relay's ACL domain name reverse lookup -- whether the IP shall be forwarded #[cfg(feature = "local-dns-relay")] - reverse_lookup_cache: Mutex>, + reverse_lookup_cache: AsyncMutex>, // For local DNS upstream #[cfg(feature = "local-dns-relay")] local_dns: LocalUpstream, + + // ACL check result cache + acl_check_cache: AsyncMutex>, } /// Unique context thw whole server @@ -224,11 +228,7 @@ impl Context { } } - let nonce_ppbloom = Mutex::new(PingPongBloom::new(config.config_type)); - #[cfg(feature = "local-dns-relay")] - let reverse_lookup_cache = Mutex::new(LruCache::::with_expiry_duration(Duration::from_secs( - 3 * 24 * 60 * 60, - ))); + let nonce_ppbloom = SpinMutex::new(PingPongBloom::new(config.config_type)); #[cfg(feature = "local-dns-relay")] let local_dns = LocalUpstream::new(&config); @@ -240,9 +240,15 @@ impl Context { #[cfg(feature = "local-flow-stat")] local_flow_statistic: ServerFlowStatistic::new(), #[cfg(feature = "local-dns-relay")] - reverse_lookup_cache, + reverse_lookup_cache: AsyncMutex::new(LruCache::with_expiry_duration(Duration::from_secs( + 3 * 24 * 60 * 60, + ))), #[cfg(feature = "local-dns-relay")] local_dns, + acl_check_cache: AsyncMutex::new(LruCache::with_expiry_duration_and_capacity( + Duration::from_secs(24 * 60 * 60), + 512, + )), } } @@ -376,14 +382,14 @@ impl Context { /// Add a record to the reverse lookup cache #[cfg(feature = "local-dns-relay")] - pub fn add_to_reverse_lookup_cache(&self, addr: &IpAddr, forward: bool) { + pub async fn add_to_reverse_lookup_cache(&self, addr: &IpAddr, forward: bool) { let is_exception = forward != match self.acl() { // Proxy everything by default None => true, Some(a) => a.check_ip_in_proxy_list(addr), }; - let mut reverse_lookup_cache = self.reverse_lookup_cache.lock(); + let mut reverse_lookup_cache = self.reverse_lookup_cache.lock().await; match reverse_lookup_cache.get_mut(addr) { Some(value) => { if is_exception { @@ -453,7 +459,23 @@ impl Context { } } } - a.check_target_bypassed(self, target).await + + // ACL checking may need over 500ms (DNS resolving) + if let Some(bypassed) = self.acl_check_cache.lock().await.get(target) { + trace!( + "check bypassing {} cached result: {}", + target, + if *bypassed { "bypassed" } else { "proxied" } + ); + + return *bypassed; + } + + let r = a.check_target_bypassed(self, target).await; + + self.acl_check_cache.lock().await.insert(target.clone(), r); + + r } /// Get client flow statistics diff --git a/src/relay/dnsrelay/mod.rs b/src/relay/dnsrelay/mod.rs index 8d1928128e81..f1a468066fdc 100644 --- a/src/relay/dnsrelay/mod.rs +++ b/src/relay/dnsrelay/mod.rs @@ -274,8 +274,16 @@ impl DnsRelay { for rec in result.answers() { debug!("dns answer: {:?}", rec); match rec.rdata() { - RData::A(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward), - RData::AAAA(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward), + RData::A(ref ip) => { + self.context + .add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward) + .await + } + RData::AAAA(ref ip) => { + self.context + .add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward) + .await + } _ => (), } } diff --git a/src/relay/socks4.rs b/src/relay/socks4.rs index d10a1bb57f38..169d502b41a9 100644 --- a/src/relay/socks4.rs +++ b/src/relay/socks4.rs @@ -111,7 +111,7 @@ impl fmt::Display for ResultCode { } /// SOCKS4 Address type -#[derive(Clone, PartialEq, Eq, Hash)] +#[derive(Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] pub enum Address { /// Socket address (IP Address) SocketAddress(SocketAddrV4), diff --git a/src/relay/socks5.rs b/src/relay/socks5.rs index 7c2d07f5527d..4eff04a2840f 100644 --- a/src/relay/socks5.rs +++ b/src/relay/socks5.rs @@ -205,7 +205,7 @@ impl From for io::Error { } /// SOCKS5 address type -#[derive(Clone, PartialEq, Eq, Hash)] +#[derive(Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] pub enum Address { /// Socket address (IP Address) SocketAddress(SocketAddr),