Add ip or domain filter support

Cargo.toml

@@ -45,6 +45,7 @@ bincode = "1.3.3"
 default-net = "0.2.0"
 system_gateway = {git="https://github.com/aruntomar/system_gateway"}
 base64 = "0.13.0"
+trust-dns-resolver = "0.21.2"
 
 [dev-dependencies]
 test-case = "1.2"

chaos-tproxy-controller/Cargo.toml

@@ -44,4 +44,5 @@ futures-util = { version = "0.3.7", default-features = false, features = ["alloc
 chaos-tproxy-proxy = {path = "../chaos-tproxy-proxy"}
 pnet = "0.28.0"
 default-net = "0.2.0"
-system_gateway = {git="https://github.com/aruntomar/system_gateway"}
+system_gateway = {git="https://github.com/aruntomar/system_gateway"}
+trust-dns-resolver = "0.21.2"

chaos-tproxy-controller/src/proxy/config.rs

@@ -1,7 +1,11 @@
 use std::convert::TryFrom;
+use std::net::{IpAddr, Ipv4Addr};
+use std::sync::mpsc::channel;
 
-use anyhow::{anyhow, Error};
+use anyhow::{anyhow, Error, Result};
 use chaos_tproxy_proxy::raw_config::RawConfig as ProxyRawConfig;
+use trust_dns_resolver::system_conf::read_system_conf;
+use trust_dns_resolver::Resolver;
 
 use crate::raw_config::RawConfig;
 
@@ -22,6 +26,37 @@ impl TryFrom<RawConfig> for Config {
                         .collect::<Vec<_>>()
                         .join(",")
                 }),
+                proxy_ips: raw.proxy_domains.map(|domains| {
+                    let ips: Vec<Ipv4Addr> = domains
+                        .into_iter()
+                        .map(|domain| {
+                            let (config, opt) = read_system_conf()?;
+                            let resolver = Resolver::new(config, opt)?;
+
+                            let (sender, receiver) = channel();
+                            let _ = std::thread::spawn(move || {
+                                let _ = sender.send(resolver.lookup_ip(domain));
+                            })
+                            .join();
+
+                            let rsp = receiver.recv()??;
+                            let ips: Vec<Ipv4Addr> = rsp
+                                .iter()
+                                .filter_map(|ip| match ip {
+                                    IpAddr::V4(ipv4) => Some(ipv4),
+                                    IpAddr::V6(_) => None,
+                                })
+                                .collect();
+                            Ok(ips)
+                        })
+                        .filter_map(|r: Result<Vec<Ipv4Addr>>| {
+                            r.map_err(|e| tracing::error!("resolve domain with error: {}", e))
+                                .ok()
+                        })
+                        .flatten()
+                        .collect();
+                    ips
+                }),
                 safe_mode: match &raw.safe_mode {
                     Some(b) => *b,
                     None => false,
@@ -54,73 +89,3 @@ pub(crate) fn get_free_port(ports: Option<Vec<u16>>) -> anyhow::Result<u16> {
         "never apply all ports in 1025-65535 to be proxy ports"
     ))
 }
-
-#[cfg(test)]
-mod tests {
-    use std::convert::TryInto;
-
-    use chaos_tproxy_proxy::raw_config::RawConfig as ProxyRawConfig;
-
-    use crate::proxy::config::{get_free_port, Config};
-    use crate::raw_config::RawConfig;
-
-    #[test]
-    fn test_get_free_port() {
-        assert!(get_free_port(Some((u16::MIN..u16::MAX).collect())).is_err());
-    }
-
-    #[test]
-    fn test_try_into() {
-        let config: Config = RawConfig {
-            proxy_ports: None,
-            safe_mode: None,
-            interface: None,
-            rules: None,
-
-            listen_port: None,
-            proxy_mark: None,
-            ignore_mark: None,
-            route_table: None,
-        }
-        .try_into()
-        .unwrap();
-        assert_eq!(
-            config,
-            Config {
-                proxy_config: ProxyRawConfig {
-                    proxy_ports: None,
-                    listen_port: get_free_port(None).unwrap(),
-                    safe_mode: false,
-                    interface: None,
-                    rules: vec![]
-                }
-            }
-        );
-
-        let config: Config = RawConfig {
-            proxy_ports: Some(vec![1025u16, 1026u16]),
-            safe_mode: Some(true),
-            interface: Some("ens33".parse().unwrap()),
-            rules: None,
-
-            listen_port: None,
-            proxy_mark: None,
-            ignore_mark: None,
-            route_table: None,
-        }
-        .try_into()
-        .unwrap();
-        assert_eq!(
-            config,
-            Config {
-                proxy_config: ProxyRawConfig {
-                    proxy_ports: Some("1025,1026".parse().unwrap()),
-                    listen_port: 1027u16,
-                    safe_mode: true,
-                    interface: Some("ens33".parse().unwrap()),
-                    rules: vec![]
-                }
-            }
-        );
-    }
-}

chaos-tproxy-controller/src/raw_config.rs

@@ -5,6 +5,7 @@ use serde::{Deserialize, Serialize};
 #[serde(deny_unknown_fields)] // To prevent typos.
 pub struct RawConfig {
     pub proxy_ports: Option<Vec<u16>>,
+    pub proxy_domains: Option<Vec<String>>,
     pub safe_mode: Option<bool>,
     pub interface: Option<String>,
     pub rules: Option<Vec<RawRule>>,

chaos-tproxy-proxy/Cargo.toml

@@ -35,3 +35,5 @@ bincode = "1.3.3"
 tempfile = "3.2.0"
 uuid = { version = "0.8", features = ["serde", "v4"] }
 base64 = "0.13.0"
+trust-dns-resolver = "0.21.2"
+

chaos-tproxy-proxy/src/raw_config.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::convert::{TryFrom, TryInto};
+use std::net::Ipv4Addr;
 use std::time::Duration;
 
 use anyhow::{anyhow, Error};
@@ -19,6 +20,7 @@ use crate::proxy::http::config::Config;
 #[derive(Debug, Eq, PartialEq, Clone, Deserialize, Serialize, Default)]
 pub struct RawConfig {
     pub proxy_ports: Option<String>,
+    pub proxy_ips: Option<Vec<Ipv4Addr>>,
     pub listen_port: u16,
     pub safe_mode: bool,
     pub interface: Option<String>,

tests/Cargo.toml

@@ -42,3 +42,4 @@ chaos-tproxy-controller = {path = "../chaos-tproxy-controller"}
 pnet = "0.28.0"
 default-net = "0.2.0"
 system_gateway = {git="https://github.com/aruntomar/system_gateway"}
+trust-dns-resolver = "0.21.2"

tests/integrations/test_http_action.rs

@@ -1,7 +1,7 @@
+use chaos_tproxy_proxy::handler::http::action::{apply_request_action, Actions, ReplaceAction};
 use http::header::CONTENT_LENGTH;
 use http::HeaderMap;
 use hyper::{Body, Client, Method, Request};
-use chaos_tproxy_proxy::handler::http::action::{apply_request_action, Actions, ReplaceAction};
 
 #[tokio::test]
 #[ignore]