feat: support proxying existing cgroup slices (#135)

.github/workflows/build.yml

@@ -1,10 +1,9 @@
 on: [push, pull_request]
 
-name: CI
+name: check and test
 
 jobs:
-  check:
-    name: Check
+  cargo-check:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
@@ -17,8 +16,7 @@ jobs:
         with:
           command: check
 
-  test:
-    name: Test Suite
+  cargo-test:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
@@ -30,3 +28,38 @@ jobs:
       - uses: actions-rs/cargo@v1
         with:
           command: test
+
+  e2e-test:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash -l {0}
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v3
+
+      - uses: actions-rs/toolchain@v1
+        with:
+          profile: minimal
+          toolchain: stable
+          override: true
+
+      - name: Install Dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y iptables iproute2 socat curl dnsutils cgroup-tools
+
+      - name: Build cproxy
+        run: |
+          cargo build --release
+
+      - name: Prepare Test Environment
+        run: |
+          mkdir -p test/logs
+          chmod +x test/*.sh
+          sudo install -m 755 target/release/cproxy /usr/local/bin
+
+      - name: Run Redirect Mode Tests
+        run: |
+          sudo ./test/run_all_tests.sh

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "cproxy"
-version = "4.2.1"
+version = "4.2.2"
 authors = ["Xiangru Lian <[email protected]>"]
 description = "Transparent proxy built on cgroup net_cls."
 homepage = "https://github.com/NOBLES5E/cproxy"

README.md

@@ -102,6 +102,18 @@ sudo cproxy --mode trace <your-program>
 
 You will be able to see log in `dmesg`. Note that this requires a recent enough kernel and iptables.
 
+### Advanced Usage: Proxy Specific Cgroup Paths
+
+`cproxy` allows you to proxy all processes within specific cgroup paths. This is particularly useful for managing groups of related processes without specifying individual PIDs.
+
+Suppose you have a cgroup at `/sys/fs/cgroup/mygroup` containing several processes you wish to proxy. You can run:
+
+```bash
+sudo cproxy --port 1080 --cgroup-path /sys/fs/cgroup/mygroup --mode tproxy
+```
+
+This command will proxy all TCP and UDP traffic from processes within the `/sys/fs/cgroup/mygroup` cgroup using TPROXY mode on port `1080`.
+
 ## The Secret Sauce
 
 `cproxy` simply creates a unique `cgroup` for the proxied program, and redirect its traffic with packet rules.

src/guards.rs

@@ -5,7 +5,7 @@ use std::time::Duration;
 
 #[allow(unused)]
 pub struct CGroupGuard {
-    pub pid: u32,
+    pub pid: Option<u32>,
     pub cg: Cgroup,
     pub cg_path: String,
     pub class_id: u32,
@@ -26,13 +26,39 @@ impl CGroupGuard {
         cg.add_task_by_tgid(CgroupPid::from(pid as u64))
             .expect("add task failed");
         Ok(Self {
-            pid,
+            pid: Some(pid),
             hier_v2,
             cg,
             cg_path,
             class_id,
         })
     }
+
+    pub fn from_path(path: &str) -> Result<Self> {
+        let hier = cgroups_rs::hierarchies::auto();
+        let hier_v2 = hier.v2();
+        // Use path hash as class_id to avoid conflicts
+        let class_id = {
+            use std::hash::{Hash, Hasher};
+            let mut hasher = std::collections::hash_map::DefaultHasher::new();
+            path.hash(&mut hasher);
+            hasher.finish() as u32
+        };
+
+        let cg = CgroupBuilder::new(path)
+            .network()
+            .class_id(class_id as u64)
+            .done()
+            .build(hier)?;
+
+        Ok(Self {
+            pid: None,
+            hier_v2,
+            cg,
+            cg_path: path.to_string(),
+            class_id,
+        })
+    }
 }
 
 impl Drop for CGroupGuard {
@@ -47,7 +73,9 @@ impl Drop for CGroupGuard {
                 );
             }
         }
-        self.cg.delete().unwrap();
+        if let Err(e) = self.cg.delete() {
+            tracing::warn!("failed to delete cgroup. error: {}", e)
+        }
     }
 }
 

src/main.rs

@@ -4,6 +4,7 @@ use crate::guards::TraceGuard;
 use eyre::Result;
 use guards::{CGroupGuard, RedirectGuard, TProxyGuard};
 use std::os::unix::prelude::CommandExt;
+use std::process::ExitStatus;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
@@ -16,18 +17,27 @@ struct Cli {
     /// Redirect traffic to specific local port.
     #[structopt(long, env = "CPROXY_PORT", default_value = "1080")]
     port: u32,
+
     /// redirect DNS traffic. This option only works with redirect mode
     #[structopt(long)]
     redirect_dns: bool,
+
     /// Proxy mode can be `trace` (use iptables TRACE target to debug program network), `tproxy`, or `redirect`.
     #[structopt(long, default_value = "redirect")]
     mode: String,
+
     /// Override dns server address. This option only works with tproxy mode
     #[structopt(long)]
     override_dns: Option<String>,
+
     /// Proxy an existing process.
     #[structopt(long)]
     pid: Option<u32>,
+
+    /// Proxy specific cgroup paths, can be specified multiple times)
+    #[structopt(long)]
+    cgroup_path: Vec<String>,
+
     #[structopt(subcommand)]
     command: Option<ChildCommand>,
 }
@@ -38,7 +48,7 @@ enum ChildCommand {
     Command(Vec<String>),
 }
 
-fn proxy_new_command(args: &Cli) -> Result<()> {
+fn proxy_new_command(args: &Cli) -> Result<ExitStatus> {
     let pid = std::process::id();
     let ChildCommand::Command(child_command) = &args
         .command
@@ -51,7 +61,7 @@ fn proxy_new_command(args: &Cli) -> Result<()> {
     let cgroup_guard = CGroupGuard::new(pid)?;
     let _guard: Box<dyn Drop> = match args.mode.as_str() {
         "redirect" => {
-            let output_chain_name = format!("cproxy_redirect_out_{}", pid);
+            let output_chain_name = format!("cp_rd_out_{}", pid);
             Box::new(RedirectGuard::new(
                 port,
                 output_chain_name.as_str(),
@@ -60,8 +70,8 @@ fn proxy_new_command(args: &Cli) -> Result<()> {
             )?)
         }
         "tproxy" => {
-            let output_chain_name = format!("cproxy_tproxy_out_{}", pid);
-            let prerouting_chain_name = format!("cproxy_tproxy_pre_{}", pid);
+            let output_chain_name = format!("cp_tp_out_{}", pid);
+            let prerouting_chain_name = format!("cp_tp_pre_{}", pid);
             let mark = pid;
             Box::new(TProxyGuard::new(
                 port,
@@ -73,8 +83,8 @@ fn proxy_new_command(args: &Cli) -> Result<()> {
             )?)
         }
         "trace" => {
-            let prerouting_chain_name = format!("cproxy_trace_pre_{}", pid);
-            let output_chain_name = format!("cproxy_trace_out_{}", pid);
+            let prerouting_chain_name = format!("cp_tr_pre_{}", pid);
+            let output_chain_name = format!("cp_tr_out_{}", pid);
             Box::new(TraceGuard::new(
                 output_chain_name.as_str(),
                 prerouting_chain_name.as_str(),
@@ -111,9 +121,8 @@ fn proxy_new_command(args: &Cli) -> Result<()> {
         println!("received ctrl-c, terminating...");
     })?;
 
-    child.wait()?;
-
-    Ok(())
+    let exit_status = child.wait()?;
+    Ok(exit_status)
 }
 
 fn proxy_existing_pid(pid: u32, args: &Cli) -> Result<()> {
@@ -122,17 +131,17 @@ fn proxy_existing_pid(pid: u32, args: &Cli) -> Result<()> {
     let cgroup_guard = CGroupGuard::new(pid)?;
     let _guard: Box<dyn Drop> = match args.mode.as_str() {
         "redirect" => {
-            let output_chain_name = format!("cproxy_redirect_out_{}", pid);
+            let output_chain_name = format!("cp_rd_out_{}", pid);
             Box::new(RedirectGuard::new(
                 port,
                 output_chain_name.as_str(),
                 cgroup_guard,
-                !args.redirect_dns,
+                args.redirect_dns,
             )?)
         }
         "tproxy" => {
-            let output_chain_name = format!("cproxy_tproxy_out_{}", pid);
-            let prerouting_chain_name = format!("cproxy_tproxy_pre_{}", pid);
+            let output_chain_name = format!("cp_tp_out_{}", pid);
+            let prerouting_chain_name = format!("cp_tp_pre_{}", pid);
             let mark = pid;
             Box::new(TProxyGuard::new(
                 port,
@@ -144,8 +153,8 @@ fn proxy_existing_pid(pid: u32, args: &Cli) -> Result<()> {
             )?)
         }
         "trace" => {
-            let prerouting_chain_name = format!("cproxy_trace_pre_{}", pid);
-            let output_chain_name = format!("cproxy_trace_out_{}", pid);
+            let prerouting_chain_name = format!("cp_tr_pre_{}", pid);
+            let output_chain_name = format!("cp_tr_out_{}", pid);
             Box::new(TraceGuard::new(
                 output_chain_name.as_str(),
                 prerouting_chain_name.as_str(),
@@ -172,6 +181,67 @@ fn proxy_existing_pid(pid: u32, args: &Cli) -> Result<()> {
     Ok(())
 }
 
+fn proxy_cgroup_paths(paths: Vec<String>, args: &Cli) -> Result<()> {
+    let port = args.port;
+
+    let mut guards: Vec<Box<dyn Drop>> = Vec::new();
+
+    for path in paths {
+        let cgroup_guard = CGroupGuard::from_path(&path)?;
+        let guard: Box<dyn Drop> = match args.mode.as_str() {
+            "redirect" => {
+                let output_chain_name = format!("cp_rd_out_{}", cgroup_guard.class_id);
+                Box::new(RedirectGuard::new(
+                    port,
+                    output_chain_name.as_str(),
+                    cgroup_guard,
+                    args.redirect_dns,
+                )?)
+            }
+            "tproxy" => {
+                let output_chain_name = format!("cp_tp_out_{}", cgroup_guard.class_id);
+                let prerouting_chain_name = format!("cp_tp_pre_{}", cgroup_guard.class_id);
+                let mark = cgroup_guard.class_id;
+                Box::new(TProxyGuard::new(
+                    port,
+                    mark,
+                    output_chain_name.as_str(),
+                    prerouting_chain_name.as_str(),
+                    cgroup_guard,
+                    args.override_dns.clone(),
+                )?)
+            }
+            "trace" => {
+                let prerouting_chain_name = format!("cp_tr_pre_{}", cgroup_guard.class_id);
+                let output_chain_name = format!("cp_tr_out_{}", cgroup_guard.class_id);
+                Box::new(TraceGuard::new(
+                    output_chain_name.as_str(),
+                    prerouting_chain_name.as_str(),
+                    cgroup_guard,
+                )?)
+            }
+            _ => {
+                unimplemented!()
+            }
+        };
+        guards.push(guard);
+    }
+
+    let running = Arc::new(AtomicBool::new(true));
+    let r = running.clone();
+
+    ctrlc::set_handler(move || {
+        println!("received ctrl-c, terminating...");
+        r.store(false, Ordering::SeqCst);
+    })?;
+
+    while running.load(Ordering::SeqCst) {
+        std::thread::sleep(Duration::from_millis(100));
+    }
+
+    Ok(())
+}
+
 fn main() -> Result<()> {
     color_eyre::install()?;
     tracing_subscriber::fmt()
@@ -183,12 +253,17 @@ fn main() -> Result<()> {
         .expect("cproxy failed to seteuid, please run as root");
     let args: Cli = Cli::from_args();
 
-    match args.pid {
-        None => {
-            proxy_new_command(&args)?;
-        }
-        Some(existing_pid) => {
-            proxy_existing_pid(existing_pid, &args)?;
+    if args.cgroup_path.len() > 0 {
+        proxy_cgroup_paths(args.cgroup_path.clone(), &args)?;
+    } else {
+        match args.pid {
+            None => {
+                let exit_status = proxy_new_command(&args)?;
+                std::process::exit(exit_status.code().unwrap_or(1));
+            }
+            Some(existing_pid) => {
+                proxy_existing_pid(existing_pid, &args)?;
+            }
         }
     }