Implement cgroup resource limit management and add tests for invocation cgroup

Author: tayyebi

CMakeLists.txt

@@ -132,6 +132,14 @@ if(UNIX)
   target_include_directories(cgroups_test PRIVATE src)
   add_test(NAME cgroups_test COMMAND cgroups_test)
   set_tests_properties(cgroups_test PROPERTIES LABELS "smoke;cgroups")
+  add_executable(cgroups_limits_test tests/cgroups_limits_test.cpp src/sandbox/cgroups.cpp)
+  target_include_directories(cgroups_limits_test PRIVATE src)
+  add_test(NAME cgroups_limits_test COMMAND cgroups_limits_test)
+  set_tests_properties(cgroups_limits_test PROPERTIES LABELS "smoke;cgroups;limits")
+  add_executable(invocation_cgroup_test tests/invocation_cgroup_test.cpp src/sandbox/invocation_cgroup.cpp src/sandbox/cgroups.cpp)
+  target_include_directories(invocation_cgroup_test PRIVATE src)
+  add_test(NAME invocation_cgroup_test COMMAND invocation_cgroup_test)
+  set_tests_properties(invocation_cgroup_test PROPERTIES LABELS "smoke;cgroups;invocation")
 endif()
 
 

src/sandbox/cgroups.cpp

@@ -60,4 +60,39 @@ bool remove_transient_cgroup(const std::string& cgroup_path) {
     return true;
 }
 
+bool set_cgroup_cpu_max(const std::string& cgroup_path, const std::string& cpu_max) {
+    std::string file = cgroup_path + "/cpu.max";
+    if (!write_file(file, cpu_max + "\n")) {
+        std::cerr << "[cgroups] failed to write cpu.max" << std::endl;
+        return false;
+    }
+    return true;
+}
+
+bool set_cgroup_memory_max(const std::string& cgroup_path, const std::string& memory_max) {
+    std::string file = cgroup_path + "/memory.max";
+    if (!write_file(file, memory_max + "\n")) {
+        std::cerr << "[cgroups] failed to write memory.max" << std::endl;
+        return false;
+    }
+    return true;
+}
+
+bool set_cgroup_pids_max(const std::string& cgroup_path, const std::string& pids_max) {
+    std::string file = cgroup_path + "/pids.max";
+    if (!write_file(file, pids_max + "\n")) {
+        std::cerr << "[cgroups] failed to write pids.max" << std::endl;
+        return false;
+    }
+    return true;
+}
+
+std::string read_cgroup_file(const std::string& path) {
+    std::ifstream ifs(path);
+    if (!ifs.is_open()) return std::string();
+    std::string content;
+    std::getline(ifs, content);
+    return content;
+}
+
 } // namespace sandbox

src/sandbox/cgroups.h

@@ -18,4 +18,16 @@ bool add_pid_to_cgroup(const std::string& cgroup_path, pid_t pid = 0);
 // Remove the transient cgroup; best effort.
 bool remove_transient_cgroup(const std::string& cgroup_path);
 
+// Resource limits (cgroup v2 controller files)
+// Set CPU max (value format as in cgroup v2, e.g. "100000 100000" or "max")
+bool set_cgroup_cpu_max(const std::string& cgroup_path, const std::string& cpu_max);
+
+// Set memory max in bytes (or "max" for no limit)
+bool set_cgroup_memory_max(const std::string& cgroup_path, const std::string& memory_max);
+
+// Set pids max (integer or "max")
+bool set_cgroup_pids_max(const std::string& cgroup_path, const std::string& pids_max);
+
+// Read back controller files for verification
+std::string read_cgroup_file(const std::string& path);
 } // namespace sandbox

src/sandbox/invocation_cgroup.cpp

@@ -0,0 +1,74 @@
+#include "invocation_cgroup.h"
+#include "cgroups.h"
+#include <unistd.h>
+#include <chrono>
+#include <sstream>
+#include <random>
+#include <iostream>
+
+namespace sandbox {
+
+static std::string make_unique_name() {
+    pid_t pid = getpid();
+    auto now = std::chrono::steady_clock::now().time_since_epoch().count();
+    std::mt19937_64 rng(static_cast<unsigned long>(now ^ pid));
+    uint64_t r = rng();
+    std::ostringstream ss;
+    ss << "native_node_inv_" << pid << "_" << now << "_" << (r & 0xffffffffULL);
+    return ss.str();
+}
+
+InvocationCgroup::InvocationCgroup(const CgroupLimits& limits) {
+    if (!is_cgroup_v2_available()) {
+        created_ = false;
+        return;
+    }
+
+    std::string name = make_unique_name();
+    std::string p = create_transient_cgroup(name);
+    if (p.empty()) {
+        created_ = false;
+        return;
+    }
+    // Apply limits if provided
+    if (!limits.cpu_max.empty()) set_cgroup_cpu_max(p, limits.cpu_max);
+    if (!limits.memory_max.empty()) set_cgroup_memory_max(p, limits.memory_max);
+    if (!limits.pids_max.empty()) set_cgroup_pids_max(p, limits.pids_max);
+
+    path_ = p;
+    created_ = true;
+}
+
+InvocationCgroup::~InvocationCgroup() {
+    if (created_) {
+        // best-effort: try to remove
+        remove_transient_cgroup(path_);
+    }
+}
+
+InvocationCgroup::InvocationCgroup(InvocationCgroup&& o) noexcept {
+    path_ = std::move(o.path_);
+    created_ = o.created_;
+    o.created_ = false;
+}
+
+InvocationCgroup& InvocationCgroup::operator=(InvocationCgroup&& o) noexcept {
+    if (this != &o) {
+        if (created_) remove_transient_cgroup(path_);
+        path_ = std::move(o.path_);
+        created_ = o.created_;
+        o.created_ = false;
+    }
+    return *this;
+}
+
+bool InvocationCgroup::valid() const { return created_ && !path_.empty(); }
+
+const std::string& InvocationCgroup::path() const { return path_; }
+
+bool InvocationCgroup::add_pid(pid_t pid) {
+    if (!created_) return false;
+    return add_pid_to_cgroup(path_, pid);
+}
+
+} // namespace sandbox

src/sandbox/invocation_cgroup.h

@@ -0,0 +1,40 @@
+#pragma once
+
+#include <string>
+#include <optional>
+
+namespace sandbox {
+
+struct CgroupLimits {
+    std::string cpu_max;    // e.g., "100000 100000" or "max"
+    std::string memory_max; // bytes or "max"
+    std::string pids_max;   // integer or "max"
+};
+
+// RAII helper for per-invocation transient cgroups.
+// Creates a uniquely named cgroup and applies optional resource limits.
+class InvocationCgroup {
+public:
+    explicit InvocationCgroup(const CgroupLimits& limits);
+    ~InvocationCgroup();
+
+    // Non-copyable
+    InvocationCgroup(const InvocationCgroup&) = delete;
+    InvocationCgroup& operator=(const InvocationCgroup&) = delete;
+
+    // Moveable
+    InvocationCgroup(InvocationCgroup&&) noexcept;
+    InvocationCgroup& operator=(InvocationCgroup&&) noexcept;
+
+    bool valid() const;
+    const std::string& path() const;
+
+    // Add a pid to this invocation cgroup (defaults to current process)
+    bool add_pid(pid_t pid = 0);
+
+private:
+    std::string path_;
+    bool created_ = false;
+};
+
+} // namespace sandbox

tests/cgroups_limits_test.cpp

@@ -0,0 +1,58 @@
+#include <iostream>
+#include "sandbox/cgroups.h"
+
+int main() {
+    std::cout << "cgroups_limits_test: starting" << std::endl;
+    if (!sandbox::is_cgroup_v2_available()) {
+        std::cout << "cgroup v2 not available; skipping test" << std::endl;
+        return 0;
+    }
+
+    const std::string name = "native_node_limits_test_12345";
+    std::string path = sandbox::create_transient_cgroup(name);
+    if (path.empty()) {
+        std::cerr << "failed to create transient cgroup" << std::endl;
+        return 2;
+    }
+
+    // Set limits
+    if (!sandbox::set_cgroup_cpu_max(path, "100000 100000")) {
+        std::cerr << "failed to set cpu.max" << std::endl;
+        sandbox::remove_transient_cgroup(path);
+        return 2;
+    }
+
+    if (!sandbox::set_cgroup_memory_max(path, "33554432")) { // 32MB
+        std::cerr << "failed to set memory.max" << std::endl;
+        sandbox::remove_transient_cgroup(path);
+        return 2;
+    }
+
+    if (!sandbox::set_cgroup_pids_max(path, "10")) {
+        std::cerr << "failed to set pids.max" << std::endl;
+        sandbox::remove_transient_cgroup(path);
+        return 2;
+    }
+
+    // Read back and verify (best-effort)
+    auto cpu = sandbox::read_cgroup_file(path + "/cpu.max");
+    auto mem = sandbox::read_cgroup_file(path + "/memory.max");
+    auto pids = sandbox::read_cgroup_file(path + "/pids.max");
+
+    std::cout << "cpu.max='" << cpu << "' memory.max='" << mem << "' pids.max='" << pids << "'" << std::endl;
+
+    if (cpu.empty() || mem.empty() || pids.empty()) {
+        std::cerr << "failed to read back one or more limits" << std::endl;
+        sandbox::remove_transient_cgroup(path);
+        return 2;
+    }
+
+    // Cleanup
+    if (!sandbox::remove_transient_cgroup(path)) {
+        std::cerr << "failed to remove cgroup" << std::endl;
+        return 2;
+    }
+
+    std::cout << "cgroups_limits_test: succeeded" << std::endl;
+    return 0;
+}

tests/invocation_cgroup_test.cpp

@@ -0,0 +1,61 @@
+#include <iostream>
+#include <sys/types.h>
+#include <sys/wait.h>
+#include <unistd.h>
+#include <string>
+#include "sandbox/invocation_cgroup.h"
+#include "sandbox/cgroups.h"
+
+int main() {
+    std::cout << "invocation_cgroup_test: starting" << std::endl;
+    if (!sandbox::is_cgroup_v2_available()) {
+        std::cout << "cgroup v2 not available; skipping test" << std::endl;
+        return 0;
+    }
+
+    sandbox::CgroupLimits limits;
+    limits.cpu_max = "100000 100000";
+    limits.memory_max = "33554432"; // 32MB
+    limits.pids_max = "10";
+
+    sandbox::InvocationCgroup ig(limits);
+    if (!ig.valid()) {
+        std::cerr << "failed to create invocation cgroup (privileges?)" << std::endl;
+        return 2;
+    }
+
+    pid_t pid = fork();
+    if (pid < 0) {
+        std::cerr << "fork failed" << std::endl;
+        return 2;
+    }
+    if (pid == 0) {
+        // child: sleep a bit and exit; pause to allow parent to add pid
+        sleep(2);
+        _exit(0);
+    }
+
+    // Parent: add child pid to cgroup
+    if (!ig.add_pid(pid)) {
+        std::cerr << "failed to add pid to invocation cgroup" << std::endl;
+        kill(pid, SIGKILL);
+        waitpid(pid, nullptr, 0);
+        return 2;
+    }
+
+    // Read cgroup.procs and check presence
+    std::string procs_path = ig.path() + "/cgroup.procs";
+    std::string content = sandbox::read_cgroup_file(procs_path);
+    if (content.find(std::to_string(pid)) == std::string::npos) {
+        std::cerr << "pid not found in cgroup.procs: '" << content << "'" << std::endl;
+        kill(pid, SIGKILL);
+        waitpid(pid, nullptr, 0);
+        return 2;
+    }
+
+    // cleanup: wait for child
+    waitpid(pid, nullptr, 0);
+
+    std::cout << "invocation_cgroup_test: succeeded" << std::endl;
+    return 0;
+}