From a4c5d8f433ef23b586d64b49a000e9b3f6e4451f Mon Sep 17 00:00:00 2001
From: fancy <sciencezfx@tju.edu.cn>
Date: Wed, 22 Apr 2020 16:12:06 +0800
Subject: [PATCH] add run_in_cgroup script for other easy use

---
 CMakeLists.txt   |  4 ++--
 cgproxy.sh       | 10 +++++++---
 run_in_cgroup.sh | 50 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 5 deletions(-)
 create mode 100644 run_in_cgroup.sh

diff --git a/CMakeLists.txt b/CMakeLists.txt
index cb13add..42be9e5 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -7,8 +7,8 @@ install(TARGETS cgattach DESTINATION /usr/bin
         PERMISSIONS OWNER_READ OWNER_WRITE OWNER_EXECUTE GROUP_READ GROUP_EXECUTE WORLD_READ WORLD_EXECUTE SETUID)
 install(FILES cgproxy.sh DESTINATION /usr/bin RENAME cgproxy
         PERMISSIONS OWNER_READ OWNER_WRITE OWNER_EXECUTE GROUP_READ GROUP_EXECUTE WORLD_READ WORLD_EXECUTE)
-#install(FILES cgproxy2.sh DESTINATION /usr/bin RENAME cgproxy2
-#        PERMISSIONS OWNER_READ OWNER_WRITE OWNER_EXECUTE GROUP_READ GROUP_EXECUTE WORLD_READ WORLD_EXECUTE)
+install(FILES run_in_cgroup.sh DESTINATION /usr/bin RENAME run_in_cgroup
+       PERMISSIONS OWNER_READ OWNER_WRITE OWNER_EXECUTE GROUP_READ GROUP_EXECUTE WORLD_READ WORLD_EXECUTE)
 install(FILES cgproxy.service DESTINATION /usr/lib/systemd/system/)
 install(FILES cgproxy.conf DESTINATION /etc/)
 install(FILES cgroup-tproxy.sh DESTINATION /usr/share/cgproxy/scripts/)
diff --git a/cgproxy.sh b/cgproxy.sh
index 93ebb30..9d47942 100644
--- a/cgproxy.sh
+++ b/cgproxy.sh
@@ -5,8 +5,12 @@ source $config
 
 # test suid bit
 if [ -u "$(which cgattach)" ]; then 
-    cgattach $$ $proxy_cgroup
+    cgattach $$ $proxy_cgroup && attached=1
 else
-    sudo cgattach $$ $proxy_cgroup
+    sudo cgattach $$ $proxy_cgroup && attached=1
 fi
-$@
\ No newline at end of file
+
+# test attach success or not
+[[ -z "$attached" ]] && echo "config error" && exit 1
+
+exec "$@"
\ No newline at end of file
diff --git a/run_in_cgroup.sh b/run_in_cgroup.sh
new file mode 100644
index 0000000..d1a9530
--- /dev/null
+++ b/run_in_cgroup.sh
@@ -0,0 +1,50 @@
+#!/bin/bash
+
+print_help(){
+cat << 'DOC'
+usage:
+    run_in_cgroup --cggroup=CGROUP <COMMAND>
+    run_in_cgroup --help
+note:    
+    CGROUP must start will slash '/' , and no special character
+example:
+    run_in_cgroup --cggroup=/mycgroup.slice ping 127.0.0.1
+DOC
+}
+
+## parse parameter
+for i in "$@"
+do
+case $i in
+    --cgroup=*)
+        cgroup=${i#*=}
+        shift
+        ;;
+    --help)
+        print_help
+        exit 0
+        shift
+        ;;
+    -*)
+        shift
+        ;;
+    *)
+        break
+        ;;
+esac
+done
+
+[[ -z "$cgroup" ]] && print_help && exit 1
+[[ -z "$@" ]] && print_help && exit 1
+
+# test suid bit
+if [ -u "$(which cgattach)" ]; then 
+    cgattach $$ $cgroup && attached=1
+else
+    sudo cgattach $$ $cgroup && attached=1
+fi
+
+# test attach success or not
+[[ -z "$attached" ]] && print_help && exit 1
+
+exec "$@"