core socket feature

main.cpp

@@ -8,103 +8,19 @@
 #include <sstream>
 #include <stdio.h>
 #include <stdlib.h>
-#include <string.h>
+#include <string>
+#include "config.h"
+#include "cgroup_attach.h"
 
 using namespace std;
 using json = nlohmann::json;
+using namespace CGPROXY::SOCKET;
+using namespace CGPROXY::CONFIG;
+using namespace CGPROXY::CGROUP;
 
-struct Config {
-  string cgroup_proxy = "/proxy.slice";
-  string cgroup_noproxy = "/noproxy.slice";
-  bool enable_gateway = false;
-  int port = 12345;
-  bool enable_dns = true;
-  bool enable_tcp = true;
-  bool enable_udp = true;
-  bool enable_ipv4 = true;
-  bool enable_ipv6 = true;
-
-  void toEnv() {
-    #define env(v) setenv(#v, to_str(v).c_str(), 1)
-    env(cgroup_proxy);
-    env(cgroup_noproxy);
-    env(enable_gateway);
-    env(port);
-    env(enable_dns);
-    env(enable_tcp);
-    env(enable_udp);
-    env(enable_ipv4);
-    env(enable_ipv6);
-    #undef env
-  }
-
-  int safeLoadFromFile(const string path){
-    Config tmp=*this;
-    int flag=tmp.loadFromFile(path);
-    if (flag!=0) return flag;
-    if (tmp.isValid()){
-      loadFromFile(path);
-      return 0;
-    }else{
-      return PARAM_ERROR;
-    }
-  }
-
-  int safeLoadFromJson(const json& j){
-    Config tmp=*this;
-    int flag=tmp.loadFromJson(j);
-    if (flag!=0) return flag;
-    if (tmp.isValid()){
-      loadFromJson(j);
-      return 0;
-    }else{
-      return PARAM_ERROR;
-    }
-  }
-  
-  private:
-  int loadFromFile(const string f) {
-    debug("loading config: %s", f.c_str());
-    libconfig::Config config_f;
-    try { config_f.readFile(f.c_str()); } catch (exception &e) { return PARSE_ERROR; }
-    #define assign(v, t) if (config_f.exists(#v)) {v = (t)config_f.lookup(#v);}
-    assign(cgroup_proxy, string);
-    assign(cgroup_noproxy, string);
-    assign(enable_gateway, bool);
-    assign(port, int);
-    assign(enable_dns, bool);
-    assign(enable_tcp, bool);
-    assign(enable_udp, bool);
-    assign(enable_ipv4, bool);
-    assign(enable_ipv6, bool);
-    #undef assign
-    return 0;
-  }
-
-  int loadFromJson(const json &j) {
-    #define get_to(v) try {j.at(#v).get_to(v); } catch (exception& e) {}
-    get_to(cgroup_proxy);
-    get_to(cgroup_noproxy);
-    get_to(enable_gateway);
-    get_to(port);
-    get_to(enable_dns);
-    get_to(enable_tcp);
-    get_to(enable_udp);
-    get_to(enable_ipv4);
-    get_to(enable_ipv6);
-    #undef get_to
-    return 0;
-  }
-
-  bool isValid(){
-    // TODO
-    return true;
-  }
-};
-
-SocketControl sc;
+SocketServer sc;
 thread_arg arg_t;
-Config config_tproxy;
+Config config;
 pthread_t socket_thread_id = -1;
 
 int applyConfig(Config *c) {
@@ -118,21 +34,42 @@ int handle_msg(char *msg) {
   debug("received msg: %s", msg);
   json j;
   try{ j = json::parse(msg); }catch(exception& e){debug("msg paser error");return MSG_ERROR;}
-  int type = -1, status;
+
+  int type, status;
+  string pid, cgroup_target;
   try {
     type = j.at("type").get<int>();
-    if (type == MSG_TYPE_JSON) { // json data
-      status=config_tproxy.safeLoadFromJson(j.at("data"));
-    } else if (type == MSG_TYPE_CONFIG_PATH) { // config file
-      status=config_tproxy.safeLoadFromFile(j.at("data").get<string>());
-    }
+    switch (type)
+    {
+    case MSG_TYPE_JSON:
+      status=config.loadFromJson(j.at("data"));
+      if (status==SUCCESS) status=applyConfig(&config);
+      return status;
+      break;
+    case MSG_TYPE_CONFIG_PATH:
+      status=config.loadFromFile(j.at("data").get<string>());
+      if (status==SUCCESS) status=applyConfig(&config);
+      return status;
+      break;
+    case MSG_TYPE_PROXY_PID:
+      pid=j.at("data").get<string>();
+      status=attach(pid, config.cgroup_proxy_preserved);
+      return status;
+      break;
+    case MSG_TYPE_NOPROXY_PID:
+      pid=j.at("data").get<string>();
+      status=attach(pid, config.cgroup_noproxy_preserved);
+      return status;
+      break;
+    default:
+        return MSG_ERROR;
+      break;
+    };
   } catch (out_of_range &e) {
     return MSG_ERROR;
+  } catch (exception &e){
+    return ERROR;
   }
-  if (status==0){
-    return applyConfig(&config_tproxy);
-  }
-  return status;
 }
 
 pthread_t startSocketListeningThread() {
@@ -140,7 +77,7 @@ pthread_t startSocketListeningThread() {
   arg_t.handle_msg = &handle_msg;
   pthread_t thread_id;
   int status =
-      pthread_create(&thread_id, NULL, &SocketControl::startThread, &arg_t);
+      pthread_create(&thread_id, NULL, &SocketServer::startThread, &arg_t);
   if (status != 0)
     error("socket thread create failed");
   return thread_id;
@@ -149,13 +86,11 @@ pthread_t startSocketListeningThread() {
 int main() {
   bool enable_socket = true;
   string config_path = DEFAULT_CONFIG_FILE;
-  config_tproxy.safeLoadFromFile(config_path);
-  applyConfig(&config_tproxy);
+  config.loadFromFile(config_path);
+  applyConfig(&config);
   if (enable_socket) {
     socket_thread_id = startSocketListeningThread();
     pthread_join(socket_thread_id, NULL);
   }
   return 0;
-}
-
-// TODO handle attch pid
+}