Add Client Server role option (#38)

* complete client-server Signed-off-by: andrewmatilde <[email protected]> * complete cs Signed-off-by: andrewmatilde <[email protected]> * fix minor bug Signed-off-by: andrewmatilde <[email protected]> * allow role use multi ip Signed-off-by: andrewmatilde <[email protected]> * remove set_ip_with_interface_name Signed-off-by: andrewmatilde <[email protected]> * fix merge Signed-off-by: andrewmatilde <[email protected]> * fix comment Signed-off-by: andrewmatilde <[email protected]>
chaos-mesh · Jul 18, 2022 · 955c9f1 · 955c9f1
1 parent 0cc23af
commit 955c9f1
Show file tree

Hide file tree

Showing 8 changed files with 84 additions and 33 deletions.
diff --git a/chaos-tproxy-controller/src/proxy/config.rs b/chaos-tproxy-controller/src/proxy/config.rs
@@ -1,9 +1,12 @@
 use std::convert::TryFrom;
+use std::net::Ipv4Addr;
 
 use anyhow::{anyhow, Error};
-use chaos_tproxy_proxy::raw_config::RawConfig as ProxyRawConfig;
+use pnet::ipnetwork::IpNetwork;
+use chaos_tproxy_proxy::raw_config::{RawConfig as ProxyRawConfig, Role};
+use crate::proxy::net::bridge::get_default_interface;
 
-use crate::raw_config::RawConfig;
+use crate::raw_config::{RawConfig, RawRole};
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Config {
@@ -14,6 +17,18 @@ impl TryFrom<RawConfig> for Config {
     type Error = Error;
 
     fn try_from(raw: RawConfig) -> Result<Self, Self::Error> {
+        let ipv4s: Vec<Ipv4Addr> = get_default_interface()?
+            .ips
+            .iter()
+            .filter_map(|ips| match ips {
+                IpNetwork::V4(ipv4) => Some(ipv4),
+                _ => None,
+            })
+            .map(|ipv4| ipv4.ip())
+            .collect();
+        if ipv4s.is_empty() {
+            return Err(anyhow!("no default ipv4"));
+        }
         Ok(Config {
             proxy_config: ProxyRawConfig {
                 proxy_ports: raw.proxy_ports.clone().map(|c| {
@@ -26,12 +41,14 @@ impl TryFrom<RawConfig> for Config {
                     Some(b) => *b,
                     None => false,
                 },
-                interface: raw.interface,
                 listen_port: get_free_port(raw.proxy_ports.clone())?,
-                rules: match raw.rules {
-                    Some(rules) => rules,
-                    None => vec![],
-                },
+                rules: raw.rules.map_or(vec![], |rules| rules),
+                role: raw.role.and_then(|role| {
+                    Option::from(match role {
+                        RawRole::Client => Role::Client(ipv4s),
+                        RawRole::Server => Role::Server(ipv4s),
+                    })
+                }),
                 tls: raw.tls,
             },
         })
@@ -75,10 +92,11 @@ mod tests {
         let config: Config = RawConfig {
             proxy_ports: None,
             safe_mode: None,
-            interface: None,
             rules: None,
             tls: None,
+            role: None,
 
+            interface: None,
             listen_port: None,
             proxy_mark: None,
             ignore_mark: None,
@@ -93,8 +111,8 @@ mod tests {
                     proxy_ports: None,
                     listen_port: get_free_port(None).unwrap(),
                     safe_mode: false,
-                    interface: None,
                     rules: vec![],
+                    role: None,
                     tls: None
                 }
             }
@@ -103,10 +121,11 @@ mod tests {
         let config: Config = RawConfig {
             proxy_ports: Some(vec![1025u16, 1026u16]),
             safe_mode: Some(true),
-            interface: Some("ens33".parse().unwrap()),
             rules: None,
             tls: None,
+            role: None,
 
+            interface: None,
             listen_port: None,
             proxy_mark: None,
             ignore_mark: None,
@@ -121,8 +140,8 @@ mod tests {
                     proxy_ports: Some("1025,1026".parse().unwrap()),
                     listen_port: 1027u16,
                     safe_mode: true,
-                    interface: Some("ens33".parse().unwrap()),
                     rules: vec![],
+                    role: None,
                     tls: None
                 }
             }

diff --git a/chaos-tproxy-controller/src/proxy/exec.rs b/chaos-tproxy-controller/src/proxy/exec.rs
@@ -83,12 +83,6 @@ impl Proxy {
         };
 
         tracing::info!("Network device name {}", self.net_env.device.clone());
-        match config.interface {
-            None => {}
-            Some(interface) => {
-                self.net_env.set_ip_with_interface_name(&interface)?;
-            }
-        }
         set_net(
             &mut self.rtnl_handle,
             &self.net_env,

diff --git a/chaos-tproxy-controller/src/proxy/net/bridge.rs b/chaos-tproxy-controller/src/proxy/net/bridge.rs
@@ -71,17 +71,6 @@ impl NetEnv {
         }
     }
 
-    pub fn set_ip_with_interface_name(&mut self, interface: &str) -> anyhow::Result<()> {
-        for i in pnet::datalink::interfaces() {
-            if i.name == interface {
-                self.device = i.name.clone();
-                self.ip = get_ipv4(&i).unwrap();
-                return Ok(());
-            }
-        }
-        Err(anyhow!("interface : {} not found", interface))
-    }
-
     pub async fn setenv_bridge(&self, handle: &mut Handle) -> Result<()> {
         let Gateway {
             mac_addr: gateway_mac,

diff --git a/chaos-tproxy-controller/src/raw_config.rs b/chaos-tproxy-controller/src/raw_config.rs
@@ -6,13 +6,21 @@ use serde::{Deserialize, Serialize};
 pub struct RawConfig {
     pub proxy_ports: Option<Vec<u16>>,
     pub safe_mode: Option<bool>,
-    pub interface: Option<String>,
     pub rules: Option<Vec<RawRule>>,
     pub tls: Option<TLSRawConfig>,
+    pub role: Option<RawRole>,
 
+    // Useless options now. TODO: complete them
+    pub interface: Option<String>,
     // Useless options now. Keep these options for upward compatible.
     pub listen_port: Option<u16>,
     pub proxy_mark: Option<i32>,
     pub ignore_mark: Option<i32>,
     pub route_table: Option<u8>,
 }
+
+#[derive(Debug, Eq, PartialEq, Clone, Deserialize, Serialize)]
+pub enum RawRole {
+    Client,
+    Server,
+}
diff --git a/chaos-tproxy-proxy/src/handler/http/selector.rs b/chaos-tproxy-proxy/src/handler/http/selector.rs
@@ -1,8 +1,12 @@
+use std::net::IpAddr;
+
 use http::header::HeaderMap;
 use http::{Method, Request, Response, StatusCode, Uri};
 use hyper::Body;
 use wildmatch::WildMatch;
 
+use crate::raw_config::Role;
+
 #[derive(Debug, Clone)]
 pub struct Selector {
     pub port: Option<u16>,
@@ -13,6 +17,22 @@ pub struct Selector {
     pub response_headers: Option<HeaderMap>,
 }
 
+pub fn select_role(src_ip: &IpAddr, dst_ip: &IpAddr, role: &Role) -> bool {
+    let src_ipv4 = match src_ip {
+        IpAddr::V4(ipv4) => ipv4,
+        _ => return false,
+    };
+    let dst_ipv4 = match dst_ip {
+        IpAddr::V4(ipv4) => ipv4,
+        _ => return false,
+    };
+
+    match role {
+        Role::Client(ipv4s) => ipv4s.iter().any(|ipv4| ipv4 == src_ipv4),
+        Role::Server(ipv4s) => ipv4s.iter().any(|ipv4| ipv4 == dst_ipv4),
+    }
+}
+
 pub fn select_request(port: u16, request: &Request<Body>, selector: &Selector) -> bool {
     selector.port.iter().all(|p| port == *p)
         && selector

diff --git a/chaos-tproxy-proxy/src/proxy/http/config.rs b/chaos-tproxy-proxy/src/proxy/http/config.rs
@@ -1,6 +1,7 @@
 use rustls::{ClientConfig, ServerConfig};
 
 use crate::handler::http::rule::Rule;
+use crate::raw_config::Role;
 
 #[derive(Clone)]
 pub struct Config {
@@ -12,6 +13,7 @@ pub struct Config {
 pub struct HTTPConfig {
     pub proxy_port: u16,
     pub rules: Vec<Rule>,
+    pub role: Option<Role>,
 }
 
 #[derive(Clone)]

diff --git a/chaos-tproxy-proxy/src/proxy/http/server.rs b/chaos-tproxy-proxy/src/proxy/http/server.rs
@@ -24,7 +24,7 @@ use tracing::{debug, error, span, trace, Level};
 
 use crate::handler::http::action::{apply_request_action, apply_response_action};
 use crate::handler::http::rule::Target;
-use crate::handler::http::selector::{select_request, select_response};
+use crate::handler::http::selector::{select_request, select_response, select_role};
 use crate::proxy::http::config::{Config, HTTPConfig};
 use crate::proxy::http::connector::HttpConnector;
 use crate::proxy::tcp::listener::TcpListener;
@@ -203,15 +203,26 @@ impl HttpService {
         }
     }
 
+    fn role_ok(&self) -> bool {
+        let role = match &self.config.role {
+            None => return true,
+            Some(r) => r.clone(),
+        };
+
+        select_role(&self.remote.ip(), &self.target.ip(), &role)
+    }
+
     async fn handle(self, mut request: Request<Body>) -> Result<Response<Body>> {
         let log_key = format!("{{remote = {}, target = {} }}", self.remote, self.target);
         debug!("{} : Proxy is handling http request", log_key);
+
+        let role_ok = self.role_ok();
         let request_rules: Vec<_> = self
             .config
             .rules
             .iter()
             .filter(|rule| {
-                matches!(rule.target, Target::Request)
+                role_ok && matches!(rule.target, Target::Request)
                     && select_request(self.target.port(), &request, &rule.selector)
             })
             .collect();
@@ -280,7 +291,7 @@ impl HttpService {
             .rules
             .iter()
             .filter(|rule| {
-                matches!(rule.target, Target::Response)
+                role_ok && matches!(rule.target, Target::Response)
                     && select_response(
                         self.target.port(),
                         &uri,

diff --git a/chaos-tproxy-proxy/src/raw_config.rs b/chaos-tproxy-proxy/src/raw_config.rs
@@ -3,6 +3,7 @@ use std::convert::{TryFrom, TryInto};
 use std::path::PathBuf;
 use std::time::Duration;
 use std::{fs, io};
+use std::net::Ipv4Addr;
 
 use anyhow::{anyhow, Error};
 use http::header::{HeaderMap, HeaderName};
@@ -27,11 +28,17 @@ pub struct RawConfig {
     pub proxy_ports: Option<String>,
     pub listen_port: u16,
     pub safe_mode: bool,
-    pub interface: Option<String>,
     pub rules: Vec<RawRule>,
+    pub role: Option<Role>,
     pub tls: Option<TLSRawConfig>,
 }
 
+#[derive(Debug, Eq, PartialEq, Clone, Deserialize, Serialize)]
+pub enum Role {
+    Client(Vec<Ipv4Addr>),
+    Server(Vec<Ipv4Addr>),
+}
+
 #[derive(Debug, Eq, PartialEq, Clone, Deserialize, Serialize)]
 #[serde(tag = "type", content = "value")]
 pub enum RawFile {
@@ -191,6 +198,7 @@ impl TryFrom<RawConfig> for Config {
         Ok(Self {
             http_config: HTTPConfig {
                 proxy_port: raw.listen_port,
+                role: raw.role,
                 rules: raw
                     .rules
                     .into_iter()