From afa143db5041572edfb9ec3f8b66e2dbaa71df02 Mon Sep 17 00:00:00 2001 From: Y0ba Date: Sat, 18 Sep 2021 09:10:13 +0300 Subject: [PATCH] Add hashset and domains tree support to ACL rules --- README.md | 14 ++- crates/shadowsocks-service/src/acl/mod.rs | 101 ++++++++++++++---- .../src/acl/sub_domains_tree.rs | 63 +++++++++++ 3 files changed, 153 insertions(+), 25 deletions(-) create mode 100644 crates/shadowsocks-service/src/acl/sub_domains_tree.rs diff --git a/README.md b/README.md index dcee9cea..e39c099a 100644 --- a/README.md +++ b/README.md @@ -630,7 +630,17 @@ Example configuration: [outbound_block_list] 127.0.0.1/8 ::1 -(^|\.)baidu.com +# Using regular expression +^[a-z]{5}\.baidu\.com +# Match exactly +|baidu.com +# Match with subdomains +||google.com +# An internationalized domain name should be converted to punycode +# |☃-⌘\.com - WRONG +|xn----dqo34k.com +# ||джpумлатест.bрфa - WRONG +||xn--p-8sbkgc5ag7bhce.xn--ba-lmcq # CLIENTS # For sslocal, ..., bypasses all targets by default @@ -638,7 +648,7 @@ Example configuration: # Proxy these addresses [proxy_list] -(^|\.)google.com +||google.com 8.8.8.8 ``` diff --git a/crates/shadowsocks-service/src/acl/mod.rs b/crates/shadowsocks-service/src/acl/mod.rs index fd990de8..a75a7b1b 100644 --- a/crates/shadowsocks-service/src/acl/mod.rs +++ b/crates/shadowsocks-service/src/acl/mod.rs @@ -3,6 +3,7 @@ //! This is for advance controlling server behaviors in both local and proxy servers. use std::{ + collections::HashSet, fmt, fs::File, io::{self, BufRead, BufReader, Error, ErrorKind}, @@ -16,6 +17,10 @@ use regex::{RegexSet, RegexSetBuilder}; use shadowsocks::{context::Context, relay::socks5::Address}; +use self::sub_domains_tree::SubDomainsTree; + +mod sub_domains_tree; + /// Strategy mode that ACL is running #[derive(Debug, Copy, Clone, Eq, PartialEq)] pub enum Mode { @@ -29,17 +34,19 @@ pub enum Mode { struct Rules { ipv4: IpRange, ipv6: IpRange, - rule: RegexSet, + rule_regex: RegexSet, + rule_set: HashSet, + rule_tree: SubDomainsTree, } impl fmt::Debug for Rules { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "Rules {{ ipv4: {:?}, ipv6: {:?}, rule: [", self.ipv4, self.ipv6)?; + write!(f, "Rules {{ ipv4: {:?}, ipv6: {:?}, rule_regex: [", self.ipv4, self.ipv6)?; let max_len = 2; - let has_more = self.rule.len() > max_len; + let has_more = self.rule_regex.len() > max_len; - for (idx, r) in self.rule.patterns().iter().take(max_len).enumerate() { + for (idx, r) in self.rule_regex.patterns().iter().take(max_len).enumerate() { if idx > 0 { f.write_str(", ")?; } @@ -50,18 +57,30 @@ impl fmt::Debug for Rules { f.write_str(", ...")?; } - f.write_str("] }") + write!(f, "], rule_set: {:?}, rule_tree: {:?} }}", self.rule_set, self.rule_tree) } } impl Rules { /// Create a new rule - fn new(mut ipv4: IpRange, mut ipv6: IpRange, rule: RegexSet) -> Rules { + fn new( + mut ipv4: IpRange, + mut ipv6: IpRange, + rule_regex: RegexSet, + rule_set: HashSet, + rule_tree: SubDomainsTree, + ) -> Rules { // Optimization, merging networks ipv4.simplify(); ipv6.simplify(); - Rules { ipv4, ipv6, rule } + Rules { + ipv4, + ipv6, + rule_regex, + rule_set, + rule_tree, + } } /// Check if the specified address matches these rules @@ -83,7 +102,7 @@ impl Rules { /// Check if the specified host matches any rules fn check_host_matched(&self, host: &str) -> bool { - self.rule.is_match(host) + self.rule_set.contains(host) || self.rule_tree.contains(host) || self.rule_regex.is_match(host) } /// Check if there are no rules for IP addresses @@ -93,7 +112,7 @@ impl Rules { /// Check if there are no rules for domain names fn is_host_empty(&self) -> bool { - self.rule.is_empty() + self.rule_set.is_empty() && self.rule_tree.is_empty() && self.rule_regex.is_empty() } } @@ -162,17 +181,25 @@ impl AccessControl { let mut outbound_block_ipv4 = IpRange::new(); let mut outbound_block_ipv6 = IpRange::new(); - let mut outbound_block_rules = Vec::new(); + let mut outbound_block_rules_regex = Vec::new(); + let mut outbound_block_rules_set = HashSet::new(); + let mut outbound_block_rules_tree = SubDomainsTree::new(); let mut bypass_ipv4 = IpRange::new(); let mut bypass_ipv6 = IpRange::new(); - let mut bypass_rules = Vec::new(); + let mut bypass_rules_regex = Vec::new(); + let mut bypass_rules_set = HashSet::new(); + let mut bypass_rules_tree = SubDomainsTree::new(); let mut proxy_ipv4 = IpRange::new(); let mut proxy_ipv6 = IpRange::new(); - let mut proxy_rules = Vec::new(); + let mut proxy_rules_regex = Vec::new(); + let mut proxy_rules_set = HashSet::new(); + let mut proxy_rules_tree = SubDomainsTree::new(); let mut curr_ipv4 = &mut bypass_ipv4; let mut curr_ipv6 = &mut bypass_ipv6; - let mut curr_rules = &mut bypass_rules; + let mut curr_rules_regex = &mut bypass_rules_regex; + let mut curr_rules_set = &mut bypass_rules_set; + let mut curr_rules_tree = &mut bypass_rules_tree; for line in r.lines() { let line = line?; @@ -185,6 +212,16 @@ impl AccessControl { continue; } + if line.starts_with("||") { + curr_rules_tree.insert(&line[2..]); + continue; + } + + if line.starts_with('|') { + curr_rules_set.insert(line[1..].to_string()); + continue; + } + match line.as_str() { "[reject_all]" | "[bypass_all]" => { mode = Mode::WhiteList; @@ -195,17 +232,23 @@ impl AccessControl { "[outbound_block_list]" => { curr_ipv4 = &mut outbound_block_ipv4; curr_ipv6 = &mut outbound_block_ipv6; - curr_rules = &mut outbound_block_rules; + curr_rules_regex = &mut outbound_block_rules_regex; + curr_rules_set = &mut outbound_block_rules_set; + curr_rules_tree = &mut outbound_block_rules_tree; } "[black_list]" | "[bypass_list]" => { curr_ipv4 = &mut bypass_ipv4; curr_ipv6 = &mut bypass_ipv6; - curr_rules = &mut bypass_rules; + curr_rules_regex = &mut bypass_rules_regex; + curr_rules_set = &mut bypass_rules_set; + curr_rules_tree = &mut bypass_rules_tree; } "[white_list]" | "[proxy_list]" => { curr_ipv4 = &mut proxy_ipv4; curr_ipv6 = &mut proxy_ipv6; - curr_rules = &mut proxy_rules; + curr_rules_regex = &mut proxy_rules_regex; + curr_rules_set = &mut proxy_rules_set; + curr_rules_tree = &mut proxy_rules_tree; } _ => { match line.parse::() { @@ -226,7 +269,7 @@ impl AccessControl { } Err(..) => { // FIXME: If this line is not a valid regex, how can we know without actually compile it? - curr_rules.push(line); + curr_rules_regex.push(line); } } } @@ -237,7 +280,7 @@ impl AccessControl { const REGEX_SIZE_LIMIT: usize = usize::max_value(); - let outbound_block_regex = match RegexSetBuilder::new(outbound_block_rules) + let outbound_block_regex = match RegexSetBuilder::new(outbound_block_rules_regex) .size_limit(REGEX_SIZE_LIMIT) .build() { @@ -248,7 +291,7 @@ impl AccessControl { } }; - let bypass_regex = match RegexSetBuilder::new(bypass_rules) + let bypass_regex = match RegexSetBuilder::new(bypass_rules_regex) .case_insensitive(true) .size_limit(REGEX_SIZE_LIMIT) .build() @@ -263,7 +306,7 @@ impl AccessControl { } }; - let proxy_regex = match RegexSetBuilder::new(proxy_rules) + let proxy_regex = match RegexSetBuilder::new(proxy_rules_regex) .case_insensitive(true) .size_limit(REGEX_SIZE_LIMIT) .build() @@ -279,9 +322,21 @@ impl AccessControl { }; Ok(AccessControl { - outbound_block: Rules::new(outbound_block_ipv4, outbound_block_ipv6, outbound_block_regex), - black_list: Rules::new(bypass_ipv4, bypass_ipv6, bypass_regex), - white_list: Rules::new(proxy_ipv4, proxy_ipv6, proxy_regex), + outbound_block: Rules::new( + outbound_block_ipv4, + outbound_block_ipv6, + outbound_block_regex, + outbound_block_rules_set, + outbound_block_rules_tree, + ), + black_list: Rules::new( + bypass_ipv4, + bypass_ipv6, + bypass_regex, + bypass_rules_set, + bypass_rules_tree, + ), + white_list: Rules::new(proxy_ipv4, proxy_ipv6, proxy_regex, proxy_rules_set, proxy_rules_tree), mode, }) } diff --git a/crates/shadowsocks-service/src/acl/sub_domains_tree.rs b/crates/shadowsocks-service/src/acl/sub_domains_tree.rs new file mode 100644 index 00000000..f06fcdbc --- /dev/null +++ b/crates/shadowsocks-service/src/acl/sub_domains_tree.rs @@ -0,0 +1,63 @@ +use std::collections::HashMap; + +#[derive(Debug, Clone)] +struct DomainPart { + included: bool, + children: HashMap, +} + +impl DomainPart { + fn new() -> Self { + DomainPart { + included: false, + children: HashMap::new(), + } + } +} + +#[derive(Debug, Clone)] +pub struct SubDomainsTree(HashMap); + +impl SubDomainsTree { + pub fn new() -> Self { + SubDomainsTree(HashMap::new()) + } + + pub fn insert(&mut self, value: &str) { + let mut current_map = &mut self.0; + let mut last_included = None; + for part in value.rsplit('.') { + let entry = current_map.entry(part.to_string()).or_insert_with(|| DomainPart::new()); + // We don't need to include `a.b.c` if we already have `b.c` + if entry.included { + return; + } + current_map = &mut entry.children; + last_included = Some(&mut entry.included); + } + if let Some(last_included) = last_included { + *last_included = true; + // Remove all subdomains to free memory. `contains` will stop here anyway. + current_map.clear(); + } + } + + pub fn contains(&self, value: &str) -> bool { + let mut current_map = &self.0; + for part in value.rsplit('.') { + if let Some(el) = current_map.get(part) { + if el.included { + return true; + } + current_map = &el.children; + } else { + break; + } + } + false + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } +}