From 25ba8d7076a7563ac7b758ede44f4b6be2970362 Mon Sep 17 00:00:00 2001
From: Paul Pan <i@0x7f.app>
Date: Wed, 13 Mar 2024 16:35:47 +0800
Subject: [PATCH] feat: add dump_stats support

---
 cgroup2.cc   | 29 ++++++++++++++++++++++++++++-
 cmdline.cc   | 19 +++++++++++++++++++
 config.cc    |  3 +++
 config.proto |  4 ++++
 nsjail.h     |  2 ++
 5 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/cgroup2.cc b/cgroup2.cc
index e3a73e1..bb71d59 100644
--- a/cgroup2.cc
+++ b/cgroup2.cc
@@ -271,10 +271,37 @@ bool initNsFromParent(nsjconf_t *nsjconf, pid_t pid) {
 	return initNsFromParentCpu(nsjconf, pid);
 }
 
+void dumpCgroupStats(nsjconf_t *nsjconf, const std::string &cgroup_path) {
+	LOG_D("Dump stats to '%s'", nsjconf->dump_stats_file.c_str());
+
+	std::ofstream dump_file(nsjconf->dump_stats_file);
+	if (!dump_file.is_open()) {
+		PLOG_W("Failed to open dump stats file '%s'", nsjconf->dump_stats_file.c_str());
+		return;
+	}
+
+	const auto resources = {"cpu.stat", "io.stat", "memory.stat"};
+	for (const auto &resource : resources) {
+		std::ifstream resource_file(cgroup_path + "/" + resource);
+		if (!resource_file.is_open()) {
+			PLOG_W("Failed to open cgroup resource file '%s'",
+			    (cgroup_path + "/" + resource).c_str());
+			continue;
+		}
+		dump_file << resource_file.rdbuf();
+		dump_file << std::endl;
+		resource_file.close();
+	}
+
+	dump_file.close();
+}
+
 void finishFromParent(nsjconf_t *nsjconf, pid_t pid) {
 	if (nsjconf->cgroup_mem_max != (size_t)0 || nsjconf->cgroup_pids_max != 0U ||
 	    nsjconf->cgroup_cpu_ms_per_sec != 0U) {
-		removeCgroup(getCgroupPath(nsjconf, pid));
+		auto cgroup_path = getCgroupPath(nsjconf, pid);
+		dumpCgroupStats(nsjconf, cgroup_path);
+		removeCgroup(cgroup_path);
 	}
 }
 
diff --git a/cmdline.cc b/cmdline.cc
index acab963..8cdbe65 100644
--- a/cmdline.cc
+++ b/cmdline.cc
@@ -170,6 +170,8 @@ static const struct custom_option custom_opts[] = {
     { { "macvlan_vs_mo", required_argument, nullptr, 0x706 }, "Mode of the 'vs' interface. Can be either 'private', 'vepa', 'bridge' or 'passthru' (default: 'private')" },
     { { "disable_tsc", no_argument, nullptr, 0x707 }, "Disable rdtsc and rdtscp instructions. WARNING: To make it effective, you also need to forbid `prctl(PR_SET_TSC, PR_TSC_ENABLE, ...)` in seccomp rules! (x86 and x86_64 only). Dynamic binaries produced by GCC seem to rely on RDTSC, but static ones should work." },
     { { "forward_signals", no_argument, nullptr, 0x708 }, "Forward fatal signals to the child process instead of always using SIGKILL." },
+    { { "dump_stats", no_argument, nullptr, 0x1001 }, "Dump cgroups stat when process terminated. (Only work with MODE_STANDALONE_ONCE and cgroupv2)" },
+    { { "dump_stats_file", required_argument, nullptr, 0x1002 }, "Location of dump_stats file." },
 };
 // clang-format on
 
@@ -540,6 +542,8 @@ std::unique_ptr<nsjconf_t> parseArgs(int argc, char *argv[]) {
 	nsjconf->seccomp_fprog.len = 0;
 	nsjconf->seccomp_log = false;
 	nsjconf->nice_level = 19;
+	nsjconf->dump_stats = false;
+	nsjconf->dump_stats_file = "/tmp/nsjail.log";
 
 	nsjconf->openfds.push_back(STDIN_FILENO);
 	nsjconf->openfds.push_back(STDOUT_FILENO);
@@ -985,6 +989,12 @@ std::unique_ptr<nsjconf_t> parseArgs(int argc, char *argv[]) {
 		case 0x903:
 			nsjconf->nice_level = (int)strtol(optarg, NULL, 0);
 			break;
+		case 0x1001:
+			nsjconf->dump_stats = true;
+			break;
+		case 0x1002:
+			nsjconf->dump_stats_file = optarg;
+			break;
 		default:
 			cmdlineUsage(argv[0]);
 			return nullptr;
@@ -1008,6 +1018,15 @@ std::unique_ptr<nsjconf_t> parseArgs(int argc, char *argv[]) {
 		LOG_F("cannot set both cgroup_mem_memsw_max and cgroup_mem_swap_max");
 	}
 
+	if (nsjconf->dump_stats) {
+		if (nsjconf->mode != MODE_STANDALONE_ONCE) {
+			LOG_F("dump_stats can only be used with MODE_STANDALONE_ONCE");
+		}
+		if (!nsjconf->use_cgroupv2) {
+			LOG_E("dump_stats can only be used with cgroup v2");
+		}
+	}
+
 	return nsjconf;
 }
 
diff --git a/config.cc b/config.cc
index 7f0b4d7..bef1338 100644
--- a/config.cc
+++ b/config.cc
@@ -299,6 +299,9 @@ static bool parseInternal(nsjconf_t* nsjconf, const nsjail::NsJailConfig& njc) {
 		nsjconf->use_execveat = njc.exec_bin().exec_fd();
 	}
 
+	nsjconf->dump_stats = njc.dump_stats();
+	nsjconf->dump_stats_file = njc.dump_stats_file();
+
 	return true;
 }
 
diff --git a/config.proto b/config.proto
index 7310b14..0295a47 100644
--- a/config.proto
+++ b/config.proto
@@ -278,4 +278,8 @@ message NsJailConfig {
 
 	/* Check whether cgroupv2 is available, and use it if available. */
 	optional bool detect_cgroupv2 = 95 [default = false];
+
+	/* Dump statistics to a file */
+	optional bool dump_stats = 96 [default = false];
+	optional string dump_stats_file = 97 [default = "/tmp/nsjail.log"];
 }
diff --git a/nsjail.h b/nsjail.h
index 38c9851..4de13dc 100644
--- a/nsjail.h
+++ b/nsjail.h
@@ -180,6 +180,8 @@ struct nsjconf_t {
 	std::vector<int> caps;
 	std::vector<std::string> ifaces;
 	std::vector<pipemap_t> pipes;
+	bool dump_stats;
+	std::string dump_stats_file;
 };
 
 #endif /* _NSJAIL_H */