From bc94e58cb164cbe11ea34c56ec93f5866a9b17c7 Mon Sep 17 00:00:00 2001
From: springzfx <springzfx@gmail.com>
Date: Fri, 12 Jun 2020 00:20:57 +0800
Subject: [PATCH] check again after small period(100ms) to avoid kde cgroup
 override

---
 src/cgproxyd.hpp      | 18 ++++++++++++++----
 src/cgroup_attach.cpp | 17 ++++++++++++++---
 src/cgroup_attach.h   |  1 +
 src/common.cpp        |  2 +-
 4 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/cgproxyd.hpp b/src/cgproxyd.hpp
index 134bdc5..a925b78 100644
--- a/src/cgproxyd.hpp
+++ b/src/cgproxyd.hpp
@@ -103,8 +103,13 @@ class cgproxyd {
         return 0;
       }
       if (!belongToCgroup(cg, config.cgroup_noproxy)) {
-        info("execsnoop; noproxy: %d %s", pid, path.get());
-        return attach(pid, config.cgroup_noproxy_preserved);
+        int res = attach(pid, config.cgroup_noproxy_preserved);
+        if (res == 0) {
+          info("execsnoop; noproxy: %d %s", pid, path.get());
+        } else {
+          info("execsnoop; noproxy failed: %d %s", pid, path.get());
+        }
+        return res;
       }
     }
 
@@ -122,8 +127,13 @@ class cgproxyd {
         return 0;
       }
       if (!belongToCgroup(cg, config.cgroup_proxy)) {
-        info("execsnoop: proxied: %d %s", pid, path.get());
-        return attach(pid, config.cgroup_proxy_preserved);
+        int res = attach(pid, config.cgroup_proxy_preserved);
+        if (res == 0) {
+          info("execsnoop: proxied: %d %s", pid, path.get());
+        } else {
+          info("execsnoop: proxied failed: %d %s", pid, path.get());
+        }
+        return res;
       }
     }
     return 0;
diff --git a/src/cgroup_attach.cpp b/src/cgroup_attach.cpp
index 5c0f5c4..4eb5a79 100644
--- a/src/cgroup_attach.cpp
+++ b/src/cgroup_attach.cpp
@@ -9,6 +9,7 @@
 #include <string>
 #include <sys/stat.h>
 #include <sys/types.h>
+#include <thread>
 #include <unistd.h>
 
 namespace CGPROXY::CGROUP {
@@ -68,9 +69,19 @@ int attach(const string pid, const string cgroup_target) {
   }
 
   // put pid to target cgroup
-  ofstream procs(cgroup_target_procs, ofstream::app);
+  if (write2procs(pid, cgroup_target_procs) != 0) return_error;
+
+  // wait for small period and check again
+  this_thread::sleep_for(std::chrono::milliseconds(100));
+  if (getCgroup(pid) != cgroup_target && write2procs(pid, cgroup_target_procs) != 0)
+    return_error;
+  return_success;
+}
+
+int write2procs(string pid, string procspath) {
+  ofstream procs(procspath, ofstream::app);
   if (!procs.is_open()) {
-    error("open file %s failed", cgroup_target_procs.c_str());
+    error("open file %s failed", procspath.c_str());
     return_error;
   }
   procs << pid.c_str() << endl;
@@ -79,7 +90,7 @@ int attach(const string pid, const string cgroup_target) {
   // maybe there some write error, for example process pid may not exist
   if (!procs) {
     error("write %s to %s failed, maybe process %s not exist", pid.c_str(),
-          cgroup_target_procs.c_str(), pid.c_str());
+          procspath.c_str(), pid.c_str());
     return_error;
   }
   return_success;
diff --git a/src/cgroup_attach.h b/src/cgroup_attach.h
index 43623eb..896edd0 100644
--- a/src/cgroup_attach.h
+++ b/src/cgroup_attach.h
@@ -11,6 +11,7 @@ bool validate(string pid, string cgroup);
 string get_cgroup2_mount_point();
 int attach(const string pid, const string cgroup_target);
 int attach(const int pid, const string cgroup_target);
+int write2procs(string pid, string procspath);
 
 } // namespace CGPROXY::CGROUP
 
diff --git a/src/common.cpp b/src/common.cpp
index 1e9f576..18dd4bb 100644
--- a/src/common.cpp
+++ b/src/common.cpp
@@ -112,7 +112,7 @@ string getCgroup(const string &pid) {
   ifstream ifs(cgroup_f);
   debug("prcessing file %s", cgroup_f.c_str());
   while (ifs.good() && getline(ifs, line)) {
-    debug("process line: %s", line.c_str());
+    // debug("process line: %s", line.c_str());
     if (line[0] == '0') {
       cgroup = line.substr(3);
       debug("get cgroup of %s: %s", pid.c_str(), cgroup.c_str());