Add "dns-cache-size" command-line option

crates/shadowsocks-service/src/config.rs

@@ -168,6 +168,9 @@ struct SSConfig {
     #[serde(skip_serializing_if = "Option::is_none")]
     dns: Option<SSDnsConfig>,
 
+    #[serde(skip_serializing_if = "Option::is_none")]
+    dns_cache_size: Option<usize>,
+
     #[serde(skip_serializing_if = "Option::is_none")]
     mode: Option<String>,
 
@@ -1099,6 +1102,7 @@ pub struct Config {
     /// - `cloudflare`, `cloudflare_tls`, `cloudflare_https`
     /// - `quad9`, `quad9_tls`
     pub dns: DnsConfig,
+    pub dns_cache_size: Option<usize>,
     /// Uses IPv6 addresses first
     ///
     /// Set to `true` if you want to query IPv6 addresses before IPv4
@@ -1249,6 +1253,7 @@ impl Config {
             local: Vec::new(),
 
             dns: DnsConfig::default(),
+            dns_cache_size: None,
             ipv6_first: false,
             ipv6_only: false,
 
@@ -1935,6 +1940,7 @@ impl Config {
                 Some(SSDnsConfig::TrustDns(c)) => nconfig.dns = DnsConfig::TrustDns(c),
                 None => nconfig.dns = DnsConfig::System,
             }
+            nconfig.dns_cache_size = config.dns_cache_size;
         }
 
         // TCP nodelay

crates/shadowsocks-service/src/dns/mod.rs

@@ -6,7 +6,7 @@ use shadowsocks::{dns_resolver::DnsResolver, net::ConnectOpts};
 use crate::config::DnsConfig;
 
 #[allow(unused_variables, dead_code)]
-pub async fn build_dns_resolver(dns: DnsConfig, ipv6_first: bool, connect_opts: &ConnectOpts) -> Option<DnsResolver> {
+pub async fn build_dns_resolver(dns: DnsConfig, ipv6_first: bool, dns_cache_size: Option<usize>, connect_opts: &ConnectOpts) -> Option<DnsResolver> {
     match dns {
         DnsConfig::System => {
             #[cfg(feature = "trust-dns")]
@@ -23,7 +23,7 @@ pub async fn build_dns_resolver(dns: DnsConfig, ipv6_first: bool, connect_opts:
                 };
 
                 if !force_system_builtin {
-                    return match DnsResolver::trust_dns_system_resolver(connect_opts.clone()).await {
+                    return match DnsResolver::trust_dns_system_resolver(dns_cache_size, connect_opts.clone()).await {
                         Ok(r) => Some(r),
                         Err(err) => {
                             warn!(
@@ -41,7 +41,7 @@ pub async fn build_dns_resolver(dns: DnsConfig, ipv6_first: bool, connect_opts:
             None
         }
         #[cfg(feature = "trust-dns")]
-        DnsConfig::TrustDns(dns) => match DnsResolver::trust_dns_resolver(dns, connect_opts.clone()).await {
+        DnsConfig::TrustDns(dns) => match DnsResolver::trust_dns_resolver(dns, dns_cache_size, connect_opts.clone()).await {
             Ok(r) => Some(r),
             Err(err) => {
                 use log::warn;

crates/shadowsocks-service/src/local/mod.rs

@@ -165,7 +165,7 @@ impl Server {
         accept_opts.tcp.mptcp = config.mptcp;
         context.set_accept_opts(accept_opts);
 
-        if let Some(resolver) = build_dns_resolver(config.dns, config.ipv6_first, context.connect_opts_ref()).await {
+        if let Some(resolver) = build_dns_resolver(config.dns, config.ipv6_first, config.dns_cache_size, context.connect_opts_ref()).await {
             context.set_dns_resolver(Arc::new(resolver));
         }
 

crates/shadowsocks-service/src/manager/mod.rs

@@ -64,7 +64,7 @@ pub async fn run(config: Config) -> io::Result<()> {
     accept_opts.tcp.keepalive = config.keep_alive.or(Some(SERVER_DEFAULT_KEEPALIVE_TIMEOUT));
     accept_opts.tcp.mptcp = config.mptcp;
 
-    if let Some(resolver) = build_dns_resolver(config.dns, config.ipv6_first, &connect_opts).await {
+    if let Some(resolver) = build_dns_resolver(config.dns, config.ipv6_first, config.dns_cache_size, &connect_opts).await {
         manager_builder.set_dns_resolver(Arc::new(resolver));
     }
     manager_builder.set_ipv6_first(config.ipv6_first);

crates/shadowsocks-service/src/server/mod.rs

@@ -95,7 +95,7 @@ pub async fn run(config: Config) -> io::Result<()> {
     accept_opts.tcp.keepalive = config.keep_alive.or(Some(SERVER_DEFAULT_KEEPALIVE_TIMEOUT));
     accept_opts.tcp.mptcp = config.mptcp;
 
-    let resolver = build_dns_resolver(config.dns, config.ipv6_first, &connect_opts)
+    let resolver = build_dns_resolver(config.dns, config.ipv6_first, config.dns_cache_size, &connect_opts)
         .await
         .map(Arc::new);
 

crates/shadowsocks/src/dns_resolver/resolver.rs

@@ -39,6 +39,7 @@ pub trait DnsResolve {
 pub struct TrustDnsSystemResolver {
     resolver: ArcSwap<TrustDnsResolver>,
     connect_opts: ConnectOpts,
+    dns_cache_size: Option<usize>,
 }
 
 /// Collections of DNS resolver
@@ -195,7 +196,7 @@ async fn trust_dns_notify_update_dns(resolver: Arc<TrustDnsSystemResolver>) -> n
                 // Update once for all those Modify events
                 time::sleep(Duration::from_secs(1)).await;
 
-                match create_resolver(None, resolver.connect_opts.clone()).await {
+                match create_resolver(None, resolver.dns_cache_size.clone(), resolver.connect_opts.clone()).await {
                     Ok(r) => {
                         debug!("auto-reload {DNS_RESOLV_FILE_PATH}");
 
@@ -226,14 +227,15 @@ impl DnsResolver {
     ///
     /// On *nix system, it will try to read configurations from `/etc/resolv.conf`.
     #[cfg(feature = "trust-dns")]
-    pub async fn trust_dns_system_resolver(connect_opts: ConnectOpts) -> io::Result<DnsResolver> {
+    pub async fn trust_dns_system_resolver(dns_cache_size: Option<usize>, connect_opts: ConnectOpts) -> io::Result<DnsResolver> {
         use super::trust_dns_resolver::create_resolver;
 
-        let resolver = create_resolver(None, connect_opts.clone()).await?;
+        let resolver = create_resolver(None, dns_cache_size.clone(), connect_opts.clone()).await?;
 
         let inner = Arc::new(TrustDnsSystemResolver {
             resolver: ArcSwap::from(Arc::new(resolver)),
             connect_opts,
+            dns_cache_size,
         });
 
         cfg_if! {
@@ -256,9 +258,9 @@ impl DnsResolver {
 
     /// Use trust-dns DNS resolver (with DNS cache)
     #[cfg(feature = "trust-dns")]
-    pub async fn trust_dns_resolver(dns: ResolverConfig, connect_opts: ConnectOpts) -> io::Result<DnsResolver> {
+    pub async fn trust_dns_resolver(dns: ResolverConfig, dns_cache_size: Option<usize>, connect_opts: ConnectOpts) -> io::Result<DnsResolver> {
         use super::trust_dns_resolver::create_resolver;
-        Ok(DnsResolver::TrustDns(create_resolver(Some(dns), connect_opts).await?))
+        Ok(DnsResolver::TrustDns(create_resolver(Some(dns), dns_cache_size, connect_opts).await?))
     }
 
     /// Custom DNS resolver

crates/shadowsocks/src/dns_resolver/trust_dns_resolver.rs

@@ -99,15 +99,19 @@ pub type ShadowDnsConnectionProvider = GenericConnector<ShadowDnsRuntimeProvider
 pub type DnsResolver = AsyncResolver<ShadowDnsConnectionProvider>;
 
 /// Create a `trust-dns` asynchronous DNS resolver
-pub async fn create_resolver(dns: Option<ResolverConfig>, connect_opts: ConnectOpts) -> ResolveResult<DnsResolver> {
+pub async fn create_resolver(dns: Option<ResolverConfig>, dns_cache_size: Option<usize>, connect_opts: ConnectOpts) -> ResolveResult<DnsResolver> {
     // Customized dns resolution
     match dns {
         Some(conf) => {
             let mut resolver_opts = ResolverOpts::default();
             // Use Ipv4AndIpv6 strategy. Because Ipv4ThenIpv6 or Ipv6ThenIpv4 will return if the first query returned.
             // Since we want to use Happy Eyeballs to connect to both IPv4 and IPv6 addresses, we need both A and AAAA records.
             resolver_opts.ip_strategy = LookupIpStrategy::Ipv4AndIpv6;
-            resolver_opts.cache_size = 0;
+
+            match dns_cache_size {
+                Some(size) => resolver_opts.cache_size = size,
+                None => {},
+            }
 
             trace!(
                 "initializing DNS resolver with config {:?} opts {:?}",

src/service/local.rs

@@ -195,6 +195,7 @@ pub fn define_command_line_options(mut app: Command) -> Command {
             .help("Path to ACL (Access Control List)"),
     )
     .arg(Arg::new("DNS").long("dns").num_args(1).action(ArgAction::Set).help("DNS nameservers, formatted like [(tcp|udp)://]host[:port][,host[:port]]..., or unix:///path/to/dns, or predefined keys like \"google\", \"cloudflare\""))
+    .arg(Arg::new("DNS_CACHE_SIZE").long("dns-cache-size").num_args(1).action(ArgAction::Set).value_parser(clap::value_parser!(usize)).help("DNS cache size in number of records. Works when trust-dns DNS backend is enabled."))
     .arg(Arg::new("TCP_NO_DELAY").long("tcp-no-delay").alias("no-delay").action(ArgAction::SetTrue).help("Set TCP_NODELAY option for sockets"))
     .arg(Arg::new("TCP_FAST_OPEN").long("tcp-fast-open").alias("fast-open").action(ArgAction::SetTrue).help("Enable TCP Fast Open (TFO)"))
     .arg(Arg::new("TCP_KEEP_ALIVE").long("tcp-keep-alive").num_args(1).action(ArgAction::Set).value_parser(clap::value_parser!(u64)).help("Set TCP keep alive timeout seconds"))
@@ -784,6 +785,10 @@ pub fn main(matches: &ArgMatches) -> ExitCode {
             config.set_dns_formatted(dns).expect("dns");
         }
 
+        if let Some(dns_cache_size) = matches.get_one::<usize>("DNS_CACHE_SIZE") {
+            config.dns_cache_size = Some(*dns_cache_size);
+        }
+
         if matches.get_flag("IPV6_FIRST") {
             config.ipv6_first = true;
         }

src/service/manager.rs

@@ -100,6 +100,7 @@ pub fn define_command_line_options(mut app: Command) -> Command {
                 .help("Default SIP003 plugin options"),
         ).arg(Arg::new("ACL").long("acl").num_args(1).action(ArgAction::Set).value_hint(ValueHint::FilePath).help("Path to ACL (Access Control List)"))
         .arg(Arg::new("DNS").long("dns").num_args(1).action(ArgAction::Set).help("DNS nameservers, formatted like [(tcp|udp)://]host[:port][,host[:port]]..., or unix:///path/to/dns, or predefined keys like \"google\", \"cloudflare\""))
+        .arg(Arg::new("DNS_CACHE_SIZE").long("dns-cache-size").num_args(1).action(ArgAction::Set).value_parser(clap::value_parser!(usize)).help("DNS cache size in number of records. Works when trust-dns DNS backend is used."))
         .arg(Arg::new("TCP_NO_DELAY").long("tcp-no-delay").alias("no-delay").action(ArgAction::SetTrue).help("Set TCP_NODELAY option for sockets"))
         .arg(Arg::new("TCP_FAST_OPEN").long("tcp-fast-open").alias("fast-open").action(ArgAction::SetTrue).help("Enable TCP Fast Open (TFO)"))
         .arg(Arg::new("TCP_KEEP_ALIVE").long("tcp-keep-alive").num_args(1).action(ArgAction::Set).value_parser(clap::value_parser!(u64)).help("Set TCP keep alive timeout seconds"))
@@ -417,6 +418,10 @@ pub fn main(matches: &ArgMatches) -> ExitCode {
             config.set_dns_formatted(dns).expect("dns");
         }
 
+        if let Some(dns_cache_size) = matches.get_one::<usize>("DNS_CACHE_SIZE") {
+            config.dns_cache_size = Some(*dns_cache_size);
+        }
+
         if matches.get_flag("IPV6_FIRST") {
             config.ipv6_first = true;
         }

src/service/server.rs

@@ -130,6 +130,7 @@ pub fn define_command_line_options(mut app: Command) -> Command {
         .arg(Arg::new("MANAGER_ADDR").long("manager-addr").num_args(1).action(ArgAction::Set).value_parser(vparser::parse_manager_addr).alias("manager-address").help("ShadowSocks Manager (ssmgr) address, could be \"IP:Port\", \"Domain:Port\" or \"/path/to/unix.sock\""))
         .arg(Arg::new("ACL").long("acl").num_args(1).action(ArgAction::Set).value_hint(ValueHint::FilePath).help("Path to ACL (Access Control List)"))
         .arg(Arg::new("DNS").long("dns").num_args(1).action(ArgAction::Set).help("DNS nameservers, formatted like [(tcp|udp)://]host[:port][,host[:port]]..., or unix:///path/to/dns, or predefined keys like \"google\", \"cloudflare\""))
+        .arg(Arg::new("DNS_CACHE_SIZE").long("dns-cache-size").num_args(1).action(ArgAction::Set).value_parser(clap::value_parser!(usize)).help("DNS cache size in number of records. Works when trust-dns DNS backend is enabled."))
         .arg(Arg::new("TCP_NO_DELAY").long("tcp-no-delay").alias("no-delay").action(ArgAction::SetTrue).help("Set TCP_NODELAY option for sockets"))
         .arg(Arg::new("TCP_FAST_OPEN").long("tcp-fast-open").alias("fast-open").action(ArgAction::SetTrue).help("Enable TCP Fast Open (TFO)"))
         .arg(Arg::new("TCP_KEEP_ALIVE").long("tcp-keep-alive").num_args(1).action(ArgAction::Set).value_parser(clap::value_parser!(u64)).help("Set TCP keep alive timeout seconds"))
@@ -436,6 +437,10 @@ pub fn main(matches: &ArgMatches) -> ExitCode {
             config.set_dns_formatted(dns).expect("dns");
         }
 
+        if let Some(dns_cache_size) = matches.get_one::<usize>("DNS_CACHE_SIZE") {
+            config.dns_cache_size = Some(*dns_cache_size);
+        }
+
         if matches.get_flag("IPV6_FIRST") {
             config.ipv6_first = true;
         }