From 92abcb1851988e019706c52e76414f3796e015b2 Mon Sep 17 00:00:00 2001
From: springzfx <springzfx@gmail.com>
Date: Sat, 23 May 2020 03:30:46 +0800
Subject: [PATCH] add --pid option

---
 src/cgproxy.hpp | 28 +++++++++++++++++++++-------
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/src/cgproxy.hpp b/src/cgproxy.hpp
index 7bb9afc..5243713 100644
--- a/src/cgproxy.hpp
+++ b/src/cgproxy.hpp
@@ -11,6 +11,7 @@ using namespace CGPROXY::CONFIG;
 namespace CGPROXY::CGPROXY {
 
 bool print_help = false, proxy = true;
+bool attach_pid=false; string arg_pid;
 inline void print_usage() {
   if (proxy) {
     cout << "Run program with proxy" << endl;
@@ -22,14 +23,24 @@ inline void print_usage() {
   }
 }
 
-void processArgs(const int argc, char *argv[], int &shift) {
-  for (int i = 1; i < argc; i++) {
+bool processArgs(const int argc, char *argv[], int &shift) {
+  int i;
+  for (i = 1; i < argc; i++) {
+    if (strcmp(argv[i], "--pid") == 0) {
+      attach_pid = true;
+      i++;
+      if (i==argc) return false;
+      arg_pid=argv[i];
+      if (!validPid(arg_pid)) return false;
+      continue;
+    }
     if (strcmp(argv[i], "--noproxy") == 0) { proxy = false; }
     if (strcmp(argv[i], "--debug") == 0) { enable_debug = true; }
     if (strcmp(argv[i], "--help") == 0) { print_help = true; }
     if (argv[i][0] != '-') { break; }
-    shift += 1;
   }
+  shift = i;
+  return true;
 }
 
 void send_pid(const pid_t pid, bool proxy, int &status) {
@@ -40,21 +51,24 @@ void send_pid(const pid_t pid, bool proxy, int &status) {
 }
 
 int main(int argc, char *argv[]) {
-  int shift = 1;
-  processArgs(argc, argv, shift);
+  int shift = -1;
+  if (!processArgs(argc, argv, shift)){
+    error("parameter error");
+    exit(EXIT_FAILURE);
+  }
 
   if (print_help) {
     print_usage();
     exit(0);
   }
 
-  if (argc == shift) {
+  if (!attach_pid && argc == shift) {
     error("no program specified");
     exit(EXIT_FAILURE);
   }
 
   int status = -1;
-  send_pid(getpid(), proxy, status);
+  send_pid(attach_pid?stoi(arg_pid):getpid(), proxy, status);
   if (status != 0) {
     error("attach process failed");
     if (status==1) error("maybe cgproxy.service not running");