From a16cefbfb2cbd092c2e5ccd648b3a0c3020d61d3 Mon Sep 17 00:00:00 2001
From: fancy <springzfx@gmail.com>
Date: Sat, 16 May 2020 14:01:10 +0800
Subject: [PATCH] handle signal, robust iptables clean

---
 cgproxyd.cpp     | 44 +++++++++++++++++++++++--------------
 cgroup-tproxy.sh | 56 +++++++++++++++++++++++++++---------------------
 common.hpp       |  3 +++
 3 files changed, 63 insertions(+), 40 deletions(-)

diff --git a/cgproxyd.cpp b/cgproxyd.cpp
index 3210d61..65effaf 100644
--- a/cgproxyd.cpp
+++ b/cgproxyd.cpp
@@ -9,6 +9,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string>
+#include <csignal>
 #include "config.hpp"
 #include "cgroup_attach.hpp"
 
@@ -24,7 +25,7 @@ class cgproxyd{
   thread_arg arg_t;
   Config config;
   pthread_t socket_thread_id = -1;
-  
+
   static cgproxyd* instance;
   static int handle_msg_static(char* msg){
     if (!instance) {
@@ -33,12 +34,11 @@ class cgproxyd{
     }
     return instance->handle_msg(msg);
   }
-
-  int applyConfig(Config *c) {
-    system("sh /usr/share/cgproxy/scripts/cgroup-tproxy.sh stop");
-    c->toEnv();
-    system("sh /usr/share/cgproxy/scripts/cgroup-tproxy.sh");
-    return 0;
+  static void signalHandler( int signum ){
+    debug("Signal %d received.", &signum);
+    if (!instance){ error("no cgproxyd instance assigned");}
+    else { instance->stop(); }
+    exit(signum);
   }
 
   int handle_msg(char *msg) {
@@ -99,22 +99,34 @@ class cgproxyd{
 
   public:
   int start(int argc, char* argv[]) {
+    signal(SIGINT, &signalHandler);
+    signal(SIGTERM,&signalHandler);
+    signal(SIGHUP,&signalHandler);
+
     int shift=1;
     processArgs(argc,argv,shift);
 
-    bool enable_socket = true;
-    string config_path = DEFAULT_CONFIG_FILE;
-    config.loadFromFile(config_path);
+    config.loadFromFile(DEFAULT_CONFIG_FILE);
     applyConfig(&config);
-    if (enable_socket) {
-      assignStaticInstance();
-      socket_thread_id = startSocketListeningThread();
-      pthread_join(socket_thread_id, NULL);
-    }
+
+    assignStaticInstance();
+    socket_thread_id = startSocketListeningThread();
+    pthread_join(socket_thread_id, NULL);
     return 0;
   }
+  int applyConfig(Config *c) {
+    system(TPROXY_IPTABLS_CLEAN);
+    c->toEnv();
+    system(TPROXY_IPTABLS_START);
+    // no need to track running status
+    return 0;
+  }
+  void stop(){
+    debug("stopping");
+    system(TPROXY_IPTABLS_CLEAN);
+  }
   ~cgproxyd(){
-
+    stop();
   }
 };
 
diff --git a/cgroup-tproxy.sh b/cgroup-tproxy.sh
index 4fba5d3..575a756 100644
--- a/cgroup-tproxy.sh
+++ b/cgroup-tproxy.sh
@@ -70,35 +70,41 @@ cgroup_mount_point=$(findmnt -t cgroup2 -n -o TARGET)
 cgroup_type="cgroup2"
 cgroup_procs_file="cgroup.procs"
 
+
+stop(){
+    iptables -t mangle -L TPROXY_PRE &> /dev/null || return
+    echo "cleaning tproxy iptables"
+    iptables -t mangle -D PREROUTING -j TPROXY_PRE
+    iptables -t mangle -D OUTPUT -j TPROXY_OUT
+    iptables -t mangle -F TPROXY_PRE
+    iptables -t mangle -F TPROXY_OUT
+    iptables -t mangle -F TPROXY_ENT
+    iptables -t mangle -X TPROXY_PRE
+    iptables -t mangle -X TPROXY_OUT
+    iptables -t mangle -X TPROXY_ENT
+    ip6tables -t mangle -D PREROUTING -j TPROXY_PRE
+    ip6tables -t mangle -D OUTPUT -j TPROXY_OUT
+    ip6tables -t mangle -F TPROXY_PRE
+    ip6tables -t mangle -F TPROXY_OUT
+    ip6tables -t mangle -F TPROXY_ENT
+    ip6tables -t mangle -X TPROXY_PRE
+    ip6tables -t mangle -X TPROXY_OUT
+    ip6tables -t mangle -X TPROXY_ENT
+    ip rule delete fwmark $fwmark lookup $table
+    ip route flush table $table
+    ip -6 rule delete fwmark $fwmark lookup $table
+    ip -6 route flush table $table
+    ## may not exist, just ignore, and tracking their existence is not reliable
+    iptables -t nat -D POSTROUTING -m owner ! --socket-exists -j MASQUERADE &> /dev/null
+    ip6tables -t nat -D POSTROUTING -m owner ! --socket-exists -s fc00::/7 -j MASQUERADE &> /dev/null
+}
+
 ## parse parameter
 for i in "$@"
 do
 case $i in
     stop)
-        echo "stopping tproxy iptables"
-        iptables -t mangle -D PREROUTING -j TPROXY_PRE
-        iptables -t mangle -D OUTPUT -j TPROXY_OUT
-        iptables -t mangle -F TPROXY_PRE
-        iptables -t mangle -F TPROXY_OUT
-        iptables -t mangle -F TPROXY_ENT
-        iptables -t mangle -X TPROXY_PRE
-        iptables -t mangle -X TPROXY_OUT
-        iptables -t mangle -X TPROXY_ENT
-        ip6tables -t mangle -D PREROUTING -j TPROXY_PRE
-        ip6tables -t mangle -D OUTPUT -j TPROXY_OUT
-        ip6tables -t mangle -F TPROXY_PRE
-        ip6tables -t mangle -F TPROXY_OUT
-        ip6tables -t mangle -F TPROXY_ENT
-        ip6tables -t mangle -X TPROXY_PRE
-        ip6tables -t mangle -X TPROXY_OUT
-        ip6tables -t mangle -X TPROXY_ENT
-        ip rule delete fwmark $fwmark lookup $table
-        ip route flush table $table
-        ip -6 rule delete fwmark $fwmark lookup $table
-        ip -6 route flush table $table
-        ## may not exist, just ignore, and tracking their existence is not reliable
-        iptables -t nat -D POSTROUTING -m owner ! --socket-exists -j MASQUERADE &> /dev/null
-        ip6tables -t nat -D POSTROUTING -m owner ! --socket-exists -s fc00::/7 -j MASQUERADE &> /dev/null
+        stop
         exit 0
         ;;
     --config=*)
@@ -117,6 +123,8 @@ done
 test -d $cgroup_mount_point$cgroup_proxy    || mkdir $cgroup_mount_point$cgroup_proxy   || exit -1; 
 test -d $cgroup_mount_point$cgroup_noproxy  || mkdir $cgroup_mount_point$cgroup_noproxy || exit -1; 
 
+
+echo "applying tproxy iptables"
 ## use TPROXY
 #ipv4#
 ip rule add fwmark $fwmark table $table
diff --git a/common.hpp b/common.hpp
index f731194..07fe998 100644
--- a/common.hpp
+++ b/common.hpp
@@ -1,6 +1,9 @@
 #ifndef COMMON_H
 #define COMMON_H 1
 
+#define TPROXY_IPTABLS_START "sh /usr/share/cgproxy/scripts/cgroup-tproxy.sh"
+#define TPROXY_IPTABLS_CLEAN "sh /usr/share/cgproxy/scripts/cgroup-tproxy.sh stop"
+
 #define SOCKET_PATH "/tmp/cgproxy_unix_socket"
 #define LISTEN_BACKLOG 64
 #define DEFAULT_CONFIG_FILE "/etc/cgproxy/config.json"