Cache recently ACL checking results

ACL checking may requires lots of calculation and network I/O (DNS resolution)

Author: zonyitoo

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,19 +172,22 @@ pub struct Context {
 
     // Check for duplicated IV/Nonce, for prevent replay attack
     // https://github.com/shadowsocks/shadowsocks-org/issues/44
-    nonce_ppbloom: Mutex<PingPongBloom>,
+    nonce_ppbloom: SpinMutex<PingPongBloom>,
 
     // For Android's flow stat report
     #[cfg(feature = "local-flow-stat")]
     local_flow_statistic: ServerFlowStatistic,
 
     // For DNS relay's ACL domain name reverse lookup -- whether the IP shall be forwarded
     #[cfg(feature = "local-dns-relay")]
-    reverse_lookup_cache: Mutex<LruCache<IpAddr, bool>>,
+    reverse_lookup_cache: AsyncMutex<LruCache<IpAddr, bool>>,
 
     // For local DNS upstream
     #[cfg(feature = "local-dns-relay")]
     local_dns: LocalUpstream,
+
+    // ACL check result cache
+    acl_check_cache: AsyncMutex<LruCache<Address, bool>>,
 }
 
 /// 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::<IpAddr, bool>::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

src/relay/dnsrelay/mod.rs

@@ -274,8 +274,16 @@ impl<Remote: upstream::Upstream> DnsRelay<Remote> {
                 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
+                        }
                         _ => (),
                     }
                 }

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),

src/relay/socks5.rs

@@ -205,7 +205,7 @@ impl From<Error> 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),