v0.3.1

CHANGELOG.md

@@ -1,11 +1,15 @@
+# 0.3.1
+
+* Optimize memory usage
+
 # 0.2.0
 
 * add response soa record
-* fixed bugs
+* fix bugs
 
 # 0.1.1
 
-* fixed mips build error
+* fix mips build error
 
 # 0.1.0
 

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "yadns"
-version = "0.2.0"
+version = "0.3.1"
 authors = ["zyj"]
 edition = "2021"
 
@@ -36,6 +36,7 @@ anyhow = "1.0.75"
 futures = "0.3.28"
 async-recursion = "1.0.5"
 publicsuffix = "2.2.3"
+once_cell = "1.18.0"
 
 [target.'cfg(any(target_arch = "mips", target_arch = "mips64"))'.dependencies]
 trust-dns-resolver = { version = "0.23.0", default-features = false, features = ["dns-over-openssl"], optional = true }

src/config.rs

@@ -1,5 +1,4 @@
 use crate::ip::IpRange;
-use crate::Transpose;
 use failure::{err_msg, Error};
 use ipnet::IpNet;
 use serde_derive::Deserialize;
@@ -315,3 +314,20 @@ pub enum RuleAction {
     #[serde(rename = "drop")]
     Drop,
 }
+
+trait Transpose {
+    type Output;
+    fn transpose(self) -> Self::Output;
+}
+
+impl<T, E> Transpose for Option<Result<T, E>> {
+    type Output = Result<Option<T>, E>;
+
+    fn transpose(self) -> Self::Output {
+        match self {
+            Some(Ok(x)) => Ok(Some(x)),
+            Some(Err(e)) => Err(e),
+            None => Ok(None),
+        }
+    }
+}

src/filter.rs

@@ -1,12 +1,18 @@
 use crate::{
     config::{RequestRule, ResponseRule, RuleAction},
-    CONFIG, STDERR,
+    handler_config::HandlerConfig,
+    loger::STDERR,
 };
 use slog::debug;
 use trust_dns_proto::{op::LowerQuery, rr::RecordType};
 use trust_dns_resolver::lookup::Lookup;
 
-pub fn check_response(domain: &str, upstream_name: &str, resp: &Lookup) -> RuleAction {
+pub fn check_response(
+    cfg: &HandlerConfig,
+    domain: &str,
+    upstream_name: &str,
+    resp: &Lookup,
+) -> RuleAction {
     let answers = resp.records();
 
     // Drop empty response
@@ -31,7 +37,7 @@ pub fn check_response(domain: &str, upstream_name: &str, resp: &Lookup) -> RuleA
                     let toggle = (range_pattern.len() - range_name.len()) % 2 == 1;
 
                     // See if the range contains the IP
-                    let range = CONFIG.app_config.ranges.get(range_name);
+                    let range = cfg.ranges.get(range_name);
                     range
                         .map(|range| {
                             answers
@@ -57,18 +63,16 @@ pub fn check_response(domain: &str, upstream_name: &str, resp: &Lookup) -> RuleA
             .unwrap_or(true) // No ranges field means matching all ranges
     };
 
-    CONFIG
-        .app_config
-        .response_rules
+    cfg.response_rules
         .iter()
         .find(|rule| {
-            check_upstream(rule) && check_ranges(rule) && check_domains(domain, &rule.domains)
+            check_upstream(rule) && check_ranges(rule) && check_domains(cfg, domain, &rule.domains)
         })
         .map(|rule| rule.action)
         .unwrap_or(RuleAction::Accept)
 }
 
-pub fn resolvers(query: &LowerQuery) -> Vec<&str> {
+pub fn resolvers<'a>(cfg: &'a HandlerConfig, query: &LowerQuery) -> Vec<&'a str> {
     let name = query.name().to_string();
 
     let check_type = |rule: &RequestRule| {
@@ -78,28 +82,22 @@ pub fn resolvers(query: &LowerQuery) -> Vec<&str> {
             .unwrap_or(true)
     };
 
-    let rule = CONFIG
-        .app_config
+    let rule = cfg
         .request_rules
         .iter()
-        .find(|r| check_domains(&name, &r.domains) && check_type(r));
+        .find(|r| check_domains(cfg, &name, &r.domains) && check_type(r));
 
     if let Some(rule) = rule {
         debug!(STDERR, "Query {} matches rule {:?}", name, rule);
         rule.upstreams.iter().map(String::as_str).collect()
     } else {
         debug!(STDERR, "No rule matches for {}. Use defaults.", name);
         // If no rule matches, use defaults
-        CONFIG
-            .app_config
-            .defaults
-            .iter()
-            .map(String::as_str)
-            .collect()
+        cfg.defaults.iter().map(String::as_str).collect()
     }
 }
 
-fn check_domains(domain: &str, domains: &Option<Vec<String>>) -> bool {
+fn check_domains(cfg: &HandlerConfig, domain: &str, domains: &Option<Vec<String>>) -> bool {
     let name = domain.trim_end_matches(".");
     domains
         .as_ref()
@@ -108,7 +106,7 @@ fn check_domains(domain: &str, domains: &Option<Vec<String>>) -> bool {
                 // Process the leading `!`
                 let domains_tag = domains_pattern.trim_start_matches('!');
                 let toggle = (domains_pattern.len() - domains_tag.len()) % 2 == 1;
-                let domains = CONFIG.app_config.domains.get(domains_tag);
+                let domains = cfg.domains.get(domains_tag);
                 domains
                     .map(|domains| {
                         (domains.regex_set.is_match(&name) || domains.suffix.contains(&name))

src/handler.rs

@@ -1,10 +1,11 @@
-use crate::{config::RuleAction, filter, CONFIG, STDERR};
+use crate::{config::RuleAction, filter, handler_config::HandlerConfig, loger::STDERR};
 use anyhow::Error;
 use async_recursion::async_recursion;
 use futures::{
     future::{self, MapErr, MapOk},
     Future, FutureExt, TryFutureExt,
 };
+use once_cell::sync::OnceCell;
 use slog::{debug, error};
 use std::pin::Pin;
 use trust_dns_proto::op::Query;
@@ -18,17 +19,27 @@ use trust_dns_server::{
     server::{Request, RequestHandler, ResponseHandler, ResponseInfo},
 };
 
+static HANDLER_CONFIG: OnceCell<HandlerConfig> = OnceCell::new();
+
+fn handler_config() -> &'static HandlerConfig {
+    HANDLER_CONFIG
+        .get()
+        .expect("HandlerConfig is not initialized")
+}
+
 /// DNS Request Handler
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct Handler {
     //pub counter: Arc<AtomicU64>,
 }
 
 impl Handler {
-    /// Create default handler.
-    pub fn default() -> Self {
-        Handler {
-            // counter: Arc::new(AtomicU64::new(0)),
+    /// Create handler from app config.
+    pub fn new(cfg: HandlerConfig) -> Self {
+        match HANDLER_CONFIG.set(cfg) {
+            _ => Handler {
+                // counter: Arc::new(AtomicU64::new(0)),
+            },
         }
     }
 
@@ -39,7 +50,7 @@ impl Handler {
         mut responder: R,
     ) -> Result<ResponseInfo, Error> {
         //self.counter.fetch_add(1, Ordering::SeqCst);
-        let resolvers = filter::resolvers(request.query());
+        let resolvers = filter::resolvers(handler_config(), request.query());
         let tasks: Vec<_> = resolvers
             .into_iter()
             .map(|name| {
@@ -48,7 +59,7 @@ impl Handler {
                 let query_type = request.query().query_type().to_owned();
                 let name1 = name.to_owned();
                 let name2 = name.to_owned();
-                let rs = CONFIG.app_config.resolvers.get(name);
+                let rs = handler_config().resolvers.get(name);
                 rs.unwrap()
                     .resolve(domain1, query_type)
                     .boxed()
@@ -89,7 +100,7 @@ impl Handler {
                 match future::select_all(tasks).await {
                     (Ok((domain, name, resp)), _index, remaining) => {
                         //debug!(STDERR, "DNS {} result {:?}", name, resp);
-                        match filter::check_response(&domain, &name, &resp) {
+                        match filter::check_response(handler_config(), &domain, &name, &resp) {
                             RuleAction::Accept => {
                                 // Ignore the remaining future
                                 tokio::spawn(future::join_all(remaining).map(|_| ()));

src/app_config.rs → src/handler_config.rs

@@ -9,7 +9,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 #[derive(Clone, Debug)]
-pub struct AppConfig {
+pub struct HandlerConfig {
     pub defaults: Arc<Vec<String>>,
     pub resolvers: Arc<HashMap<String, Arc<RecursiveResolver>>>,
     pub domains: Arc<HashMap<String, Domains>>,
@@ -24,8 +24,8 @@ pub struct Domains {
     pub suffix: DomainSuffix,
 }
 
-impl AppConfig {
-    pub fn new(config: Config) -> Self {
+impl From<Config> for HandlerConfig {
+    fn from(config: Config) -> Self {
         // debug!(STDERR, "{:#?}", config);
         let resolvers: HashMap<_, _> = config
             .upstreams
@@ -67,7 +67,7 @@ impl AppConfig {
             })
             .collect();
 
-        AppConfig {
+        HandlerConfig {
             defaults: Arc::new(config.default_upstreams),
             resolvers: Arc::new(resolvers),
             domains: Arc::new(domains),

src/loger.rs

@@ -0,0 +1,23 @@
+use lazy_static::lazy_static;
+use slog::{o, Drain, Logger};
+
+lazy_static! {
+    pub static ref STDOUT: Logger = stdout_logger();
+    pub static ref STDERR: Logger = stderr_logger();
+}
+
+fn stdout_logger() -> Logger {
+    let decorator = slog_term::TermDecorator::new().build();
+    let drain = slog_term::CompactFormat::new(decorator).build().fuse();
+    let drain = slog_async::Async::new(drain).build().fuse();
+
+    Logger::root(drain, o!())
+}
+
+fn stderr_logger() -> Logger {
+    let decorator = slog_term::TermDecorator::new().build();
+    let drain = slog_term::CompactFormat::new(decorator).build();
+    let drain = std::sync::Mutex::new(drain).fuse();
+
+    Logger::root(drain, o!())
+}