From 9bcea10ce4015f5ba2018af37c23b4ca74a3d761 Mon Sep 17 00:00:00 2001
From: mii443 <mii@sfc.wide.ad.jp>
Date: Thu, 17 Oct 2024 08:30:52 +0000
Subject: [PATCH] add add/remove_subtree_control

---
 src/bin/izoli.rs         |  4 ++--
 src/cgroup/cgroup.rs     | 39 ++++++++++++++++++++++++++++++++++++++-
 src/cgroup/controller.rs | 18 +++++++++++++++++-
 3 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/src/bin/izoli.rs b/src/bin/izoli.rs
index 943a200..8fa6213 100644
--- a/src/bin/izoli.rs
+++ b/src/bin/izoli.rs
@@ -1,7 +1,7 @@
-use izolilib::cgroup::cgroup::CGroup;
+use izolilib::cgroup::{cgroup::CGroup, controller::Controller};
 
 fn main() {
-    let cgroup = CGroup::new("").unwrap();
+    let cgroup = CGroup::new("test").unwrap();
     println!("{:?}", cgroup.get_root_path());
     println!("{}", cgroup.check_status());
     println!("{:?}", cgroup.read("cgroup.type"));
diff --git a/src/cgroup/cgroup.rs b/src/cgroup/cgroup.rs
index 72f502a..1be60d5 100644
--- a/src/cgroup/cgroup.rs
+++ b/src/cgroup/cgroup.rs
@@ -1,6 +1,6 @@
 use std::{
     fs::{self, File},
-    io::Read,
+    io::{Read, Write},
     path::PathBuf,
     str::FromStr,
 };
@@ -38,6 +38,14 @@ impl CGroup {
         Ok(buf)
     }
 
+    pub fn write(&self, name: &str, data: &str) -> Result<(), std::io::Error> {
+        let path = self.get_file_path(name);
+        let mut file = File::options().append(true).open(path)?;
+        file.write_all(data.as_bytes())?;
+
+        Ok(())
+    }
+
     pub fn check_status(&self) -> bool {
         let root = self.get_root_path();
 
@@ -99,6 +107,35 @@ impl CGroup {
         self.get_limit_value("cgroup.max.descendants")
     }
 
+    // cgroup files write
+
+    pub fn add_subtree_control(&self, controllers: Vec<Controller>) -> Result<(), std::io::Error> {
+        let to_write = controllers
+            .iter()
+            .map(|controller| format!("+{}", controller))
+            .collect::<Vec<String>>()
+            .join(" ");
+
+        self.write("cgroup.subtree_control", &to_write)?;
+
+        Ok(())
+    }
+
+    pub fn remove_subtree_control(
+        &self,
+        controllers: Vec<Controller>,
+    ) -> Result<(), std::io::Error> {
+        let to_write = controllers
+            .iter()
+            .map(|controller| format!("-{}", controller))
+            .collect::<Vec<String>>()
+            .join(" ");
+
+        self.write("cgroup.subtree_control", &to_write)?;
+
+        Ok(())
+    }
+
     fn get_u32_list(&self, name: &str) -> Result<Vec<u32>, std::io::Error> {
         let procs = self
             .read(name)?
diff --git a/src/cgroup/controller.rs b/src/cgroup/controller.rs
index ac56bfb..e075a02 100644
--- a/src/cgroup/controller.rs
+++ b/src/cgroup/controller.rs
@@ -1,4 +1,4 @@
-use std::str::FromStr;
+use std::{fmt, str::FromStr};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Controller {
@@ -33,3 +33,19 @@ impl FromStr for Controller {
         }
     }
 }
+
+impl fmt::Display for Controller {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Controller::Cpu => write!(f, "cpu"),
+            Controller::Cpuset => write!(f, "cpuset"),
+            Controller::Memory => write!(f, "memory"),
+            Controller::Io => write!(f, "io"),
+            Controller::Hugetlb => write!(f, "hugetlb"),
+            Controller::Misc => write!(f, "misc"),
+            Controller::Pids => write!(f, "pids"),
+            Controller::Rdma => write!(f, "rdma"),
+            Controller::Unknown => write!(f, "unknown"),
+        }
+    }
+}