Add support for ip-alias

src/config/mod.rs

@@ -261,6 +261,14 @@ pub struct Config {
 
     /// ip set
     pub ip_sets: HashMap<String, Vec<IpNet>>,
+
+    pub ip_alias: Vec<IpAlias>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct IpAlias {
+    pub ip: IpOrSet,
+    pub to: Arc<[IpAddr]>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]

src/config/parser/ip_alias.rs

@@ -0,0 +1,43 @@
+use super::*;
+
+impl NomParser for IpAlias {
+    fn parse(input: &str) -> IResult<&str, Self> {
+        let ip_list = separated_list1(tuple((space0, char(','), space0)), nom_recipes::ip);
+        map(
+            separated_pair(IpOrSet::parse, space1, ip_list),
+            |(ip, to)| IpAlias { ip, to: to.into() },
+        )(input)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse() {
+        assert_eq!(
+            IpAlias::parse("0.1.2.3 4.5.6.7,::89AB:CDEF"),
+            Ok((
+                "",
+                IpAlias {
+                    ip: IpOrSet::Net("0.1.2.3/32".parse().unwrap()),
+                    to: ["4.5.6.7", "::89AB:CDEF"]
+                        .map(|x| x.parse().unwrap())
+                        .into(),
+                }
+            ))
+        );
+
+        assert_eq!(
+            IpAlias::parse("ip-set:name  ::  ,  1.2.3.4"),
+            Ok((
+                "",
+                IpAlias {
+                    ip: IpOrSet::Set("name".to_string()),
+                    to: ["::", "1.2.3.4"].map(|x| x.parse().unwrap()).into(),
+                }
+            ))
+        );
+    }
+}

src/config/parser/mod.rs

@@ -15,6 +15,7 @@ mod file_mode;
 mod forward_rule;
 mod glob_pattern;
 mod https_record;
+mod ip_alias;
 mod ip_set;
 mod ipnet;
 mod iporset;
@@ -146,6 +147,7 @@ pub enum OneConfig {
     WhitelistIp(IpOrSet),
     User(String),
     IpSetProvider(IpSetProvider),
+    IpAlias(IpAlias),
 }
 
 pub fn parse_config(input: &str) -> IResult<&str, OneConfig> {
@@ -256,6 +258,7 @@ pub fn parse_config(input: &str) -> IResult<&str, OneConfig> {
         map(parse_item("user"), OneConfig::User),
         map(parse_item("whitelist-ip"), OneConfig::WhitelistIp),
         map(parse_item("ip-set"), OneConfig::IpSetProvider),
+        map(parse_item("ip-alias"), OneConfig::IpAlias),
         map(NomParser::parse, OneConfig::Listener),
         map(NomParser::parse, OneConfig::Server),
     ));

src/dns_conf.rs

@@ -11,6 +11,7 @@ use std::sync::Arc;
 
 pub use crate::config::*;
 use crate::dns::DomainRuleGetter;
+use crate::infra::ipset::IpMap;
 use crate::log;
 use crate::{
     dns_rule::{DomainRuleMap, DomainRuleTreeNode},
@@ -41,6 +42,8 @@ pub struct RuntimeConfig {
 
     /// List of IPs that will be ignored
     ignore_ip: Arc<IpSet>,
+
+    ip_alias: Arc<IpMap<Arc<[IpAddr]>>>,
 }
 
 impl RuntimeConfig {
@@ -337,6 +340,10 @@ impl RuntimeConfig {
         &self.ignore_ip
     }
 
+    pub fn ip_alias(&self) -> &Arc<IpMap<Arc<[IpAddr]>>> {
+        &self.ip_alias
+    }
+
     /// speed check mode
     #[inline]
     pub fn speed_check_mode(&self) -> Option<&SpeedCheckModeList> {
@@ -609,17 +616,21 @@ impl RuntimeConfigBuilder {
             cfg.listeners.push(UdpListenerConfig::default().into())
         }
 
-        let make_ip_set = |set: &[IpOrSet]| {
-            let iter = set.iter().flat_map(|ip| match ip {
+        fn get_ip_set<'a>(ip: &'a IpOrSet, cfg: &'a Config) -> &'a [IpNet] {
+            match ip {
                 IpOrSet::Net(net) => std::slice::from_ref(net),
                 IpOrSet::Set(name) => match cfg.ip_sets.get(name) {
-                    Some(net) => &**net,
+                    Some(net) => net,
                     None => {
                         warn!("unknown ip-set:{name}");
                         &[]
                     }
                 },
-            });
+            }
+        }
+
+        let make_ip_set = |set: &[IpOrSet]| {
+            let iter = set.iter().flat_map(|ip| get_ip_set(ip, &cfg));
             Arc::new(IpSet::new(iter.copied()))
         };
 
@@ -628,6 +639,12 @@ impl RuntimeConfigBuilder {
         let whitelist_ip = make_ip_set(&cfg.whitelist_ip);
         let ignore_ip = make_ip_set(&cfg.ignore_ip);
 
+        let ip_alias = cfg.ip_alias.iter().flat_map(|alias| {
+            let to = std::iter::repeat(alias.to.clone());
+            get_ip_set(&alias.ip, &cfg).iter().copied().zip(to)
+        });
+        let ip_alias = Arc::new(IpMap::from_iter(ip_alias));
+
         if !cfg.cnames.is_empty() {
             cfg.cnames.dedup_by(|a, b| a.domain == b.domain);
         }
@@ -752,6 +769,7 @@ impl RuntimeConfigBuilder {
             blacklist_ip,
             whitelist_ip,
             ignore_ip,
+            ip_alias,
             proxy_servers: Arc::new(proxy_servers),
         }
     }
@@ -908,6 +926,7 @@ impl RuntimeConfigBuilder {
                     }
                 }
                 MdnsLookup(enable) => self.mdns_lookup = Some(enable),
+                IpAlias(alias) => self.ip_alias.push(alias),
                 // #[allow(unreachable_patterns)]
                 // c => log::warn!("unhandled config {:?}", c),
             },
@@ -1622,4 +1641,15 @@ mod tests {
         assert_eq!(cfg.ip_sets["cf-ipv6"], v6);
         assert_eq!(*cfg.whitelist_ip, all);
     }
+
+    #[test]
+    fn test_ip_alias() {
+        let cfg = RuntimeConfig::load_from_file("tests/test_data/b_main.conf");
+        let addr = |s: &str| s.parse::<IpAddr>().unwrap();
+        let get_alias = |s: &str| &**cfg.ip_alias.get(&addr(s)).unwrap();
+
+        assert_eq!(get_alias("104.16.0.0"), [addr("1.2.3.4"), addr("::5678")]);
+        assert_eq!(get_alias("2400:cb00::"), [addr("::1234"), addr("5.6.7.8")]);
+        assert_eq!(get_alias("172.64.0.0"), [addr("90AB::CDEF")]);
+    }
 }

src/dns_mw_ns.rs

@@ -5,7 +5,7 @@ use std::{borrow::Borrow, net::IpAddr, time::Duration};
 
 use crate::dns_client::{LookupOptions, NameServer};
 
-use crate::infra::ipset::IpSet;
+use crate::infra::ipset::{IpMap, IpSet};
 use crate::infra::ping::{PingError, PingOutput};
 use crate::third_ext::FutureTimeoutExt;
 use crate::{
@@ -123,6 +123,7 @@ impl Middleware<DnsContext, DnsRequest, DnsResponse, DnsError> for NameServerMid
                     ignore_ip: cfg.ignore_ip().clone(),
                     blacklist_ip: cfg.blacklist_ip().clone(),
                     whitelist_ip: cfg.whitelist_ip().clone(),
+                    ip_alias: cfg.ip_alias().clone(),
                     lookup_options,
                 },
                 None => LookupIpOptions {
@@ -132,6 +133,7 @@ impl Middleware<DnsContext, DnsRequest, DnsResponse, DnsError> for NameServerMid
                     ignore_ip: cfg.ignore_ip().clone(),
                     blacklist_ip: cfg.blacklist_ip().clone(),
                     whitelist_ip: cfg.whitelist_ip().clone(),
+                    ip_alias: cfg.ip_alias().clone(),
                     lookup_options,
                 },
             };
@@ -155,6 +157,7 @@ struct LookupIpOptions {
     ignore_ip: Arc<IpSet>,
     whitelist_ip: Arc<IpSet>,
     blacklist_ip: Arc<IpSet>,
+    ip_alias: Arc<IpMap<Arc<[IpAddr]>>>,
     lookup_options: LookupOptions,
 }
 
@@ -553,11 +556,12 @@ async fn per_nameserver_lookup_ip(
     let LookupIpOptions {
         whitelist_ip,
         blacklist_ip,
+        ip_alias,
         ignore_ip,
         ..
     } = options;
 
-    if !whitelist_on && !blacklist_on && ignore_ip.is_empty() {
+    if !whitelist_on && !blacklist_on && ignore_ip.is_empty() && ip_alias.is_empty() {
         return res;
     }
 
@@ -578,15 +582,29 @@ async fn per_nameserver_lookup_ip(
         Ok(mut lookup) => {
             let query = lookup.query().clone();
 
-            let answers = lookup
-                .answers()
-                .iter()
-                .filter(|record| match record.data().ip_addr() {
-                    Some(ip) => ip_filter(&ip),
-                    _ => false,
-                })
-                .cloned()
-                .collect::<Vec<_>>();
+            let answers = lookup.take_answers();
+            let answers = {
+                let mut new_ans = Vec::new();
+                let mut alias_set = Vec::new(); // dedup
+                for record in answers {
+                    let Some(ip) = record.data().ip_addr().filter(ip_filter) else {
+                        continue;
+                    };
+                    match ip_alias.get(&ip) {
+                        None => new_ans.push(record),
+                        Some(ip) if !alias_set.contains(&ip.as_ptr()) => {
+                            alias_set.push(ip.as_ptr());
+                            new_ans.extend(ip.iter().map(|&ip| {
+                                let mut record = record.clone();
+                                record.set_data(ip.into());
+                                record
+                            }));
+                        }
+                        Some(_) => continue,
+                    }
+                }
+                new_ans
+            };
 
             if answers.is_empty() {
                 return Err(ProtoErrorKind::NoRecordsFound {

tests/test_data/b_main.conf

@@ -12,4 +12,8 @@ ip-set -n cf-ipv6 -f cf-ipv6.txt
 
 whitelist-ip 1.1.1.1
 whitelist-ip ip-set:cf-ipv4
-whitelist-ip ip-set:cf-ipv6
+whitelist-ip ip-set:cf-ipv6
+
+ip-alias ip-set:cf-ipv4 1.2.3.4,::5678
+ip-alias ip-set:cf-ipv6 ::1234,5.6.7.8
+ip-alias 172.64.0.0/13 90AB::CDEF