Merge pull request #36 from beordle/one_shot

One shot

CMakeLists.txt

@@ -69,6 +69,7 @@ thirdparty/queue/queue.c
 thirdparty/queue/queue_internal.c
 thirdparty/setproctitle.c
 thirdparty/tinyfiledialogs/tinyfiledialogs.c
+thirdparty/ya_getopt/ya_getopt.c
 )
 
 IF (CYGWIN)  # cygwin use system default libuv without epoll.

src/agent.c

@@ -19,6 +19,7 @@
 #include "pipe.h"
 #include "state.h"
 #include "thirdparty/base64.h"
+#include "thirdparty/ya_getopt/ya_getopt.h"
 #include "utils.h"
 #include "uv.h"
 #include "vnet.h"
@@ -28,6 +29,11 @@ uv_tty_t agent_stdout_tty;
 uv_tty_t agent_stdin_tty;
 static int64_t pending_send = 0;  //记录待转发的字节，用于tty 流控
 int max_suggested_size = 10240;
+
+int32_t g_oneshot_argc;
+char** g_oneshot_argv;
+
+
 void agent_read_stdin(uv_stream_t *stream, ssize_t nread, const uv_buf_t *buf);
 
 static void alloc_buffer(uv_handle_t *handle, size_t suggested_size,
@@ -225,12 +231,10 @@ int agent_process_frame(char *str_data, int data_size) {
 }
 
 int agent_handle_binary(char *buf, int size) {
-  log_info("agent_handle_binary data (%d)\n", size);
   // simple echo
   // block_write_binary_to_server(buf, size);
   vnet_data_income(buf, size);
   // block_write_binary_to_server(buf, size);
-
   return 0;
 }
 
@@ -252,14 +256,37 @@ void agent_write_data_to_server(char *buf, size_t s, bool autofree) {
   }
 }
 
-void agent() {
+static void sigint_handler(int sig) {
+  exit(0);
+  return;
+}
+
+void agent(int argc, char** argv) {
+  bool opt_is_repl = false;
+  if (argc != 0) {
+    // printf("argv %s\n", argv[0]);
+    g_oneshot_argc = argc;
+    g_oneshot_argv = argv;
+  } else {
+    opt_is_repl = true;
+  }
+
   set_agent_process();
   setvbuf(stdin, NULL, _IONBF, 0);
   agent_set_stdin_noecho();
   atexit(agent_restore_stdin);
+  signal(SIGINT, sigint_handler);
   int len_result;
-  char *result = green_encode("MAGIC!", strlen("MAGIC!"), &len_result);
-  printf("%s\r\n", result);
+  char *str_trigger;
+  // 判断是否使用 oneshot 模式
+  if (opt_is_repl) {
+    str_trigger = "MAGIC!";
+  } else {
+    str_trigger = "ONESHOT!";
+  }
+  char *result = green_encode(str_trigger, strlen(str_trigger), &len_result);
+  usleep(2000);  // 防止粘连，优化显示的目的。 
+  writen(STDOUT_FILENO, result, len_result);
   free(result);
 
   static uv_timer_t timer_watcher;

src/agent.h

@@ -11,7 +11,9 @@
 void agent_restore_stdin();
 void agent_set_stdin_noecho();
 
+extern int g_oneshot_argc;
+extern char** g_oneshot_argv;
 // extern void block_write_frame_to_server(char* data, int data_size);
 extern int write_binary_to_server(const char *buf, size_t size);
-void agent();
+void agent(int argc, char** argv);
 #endif

src/agentcall.c

@@ -12,10 +12,11 @@
 #include <sys/stat.h>
 #include <sys/types.h>
 #include <unistd.h>
+#include "log.h"
 #include "agentcall.h"
 #include "config.h"
 #include "intent.h"
-#include "log.h"
+#include "agent.h"
 #include "lwip/api.h"
 #include "lwip/def.h"
 #include "lwip/ip.h"
@@ -25,7 +26,10 @@
 #include "netif/etharp.h"
 #include "utils.h"
 #include "vnet.h"
+#include "pipe.h"
+#include <stdint.h>
 #include "portforward.h"
+
 static uint16_t agentcall_service_port = 300;
 
 typedef struct thread_arg_pass_t {
@@ -38,7 +42,7 @@ static void agentcall_server_request(void *p) {
   vnet_setsocketdefaultopt(sd);
   char recv_buf[READ_CHUNK_SIZE];
   int n, nwrote;
-  log_info("sd: %d", sd);
+  log_info("agentcall_server_request sd: %d", sd);
   int readbytes = 0;
   int32_t method = 0;
   if (vnet_readn(sd, &method, sizeof(int32_t)) == 0) {
@@ -60,7 +64,16 @@ static void agentcall_server_request(void *p) {
              local_host, local_port, remote_host, remote_port);
     portforward_static_start(local_host, local_port, remote_host,
                                  remote_port);
+
+  }
+  if (method == METHOD_GET_ARGS) {
+    lwip_writen(sd, &g_oneshot_argc, sizeof(g_oneshot_argc));
+    for (int i=0; i < g_oneshot_argc; i++) {
+      lwip_writen(sd, g_oneshot_argv[i], strlen(g_oneshot_argv[i])+1);
+    }
+    log_info("METHOD_GET_ARGS done");
   }
+
   lwip_close(sd);
   return;
 }
@@ -83,6 +96,36 @@ static void call_send_request(thread_arg_pass_t *tmp) {
   return;
 }
 
+
+int bootstrap_get_args(void * _) {
+  int nfd = vnet_tcp_connect(agentcall_service_port);
+  if (nfd < 0) {
+    log_debug("connect agentcall_service_port error");
+    return 0;
+  }
+  int32_t method = METHOD_GET_ARGS;
+  vnet_send(nfd, &method, sizeof(int32_t));
+  vnet_send(nfd, "", 1);
+  int32_t argc;
+  vnet_readn(nfd, &argc, sizeof(int32_t));
+  log_debug("count of arg %d", argc);
+  g_oneshot_argc = argc;
+  g_oneshot_argv = malloc(argc * sizeof(char*));
+  char recv_buf[READ_CHUNK_SIZE];
+  for (int i=0; i < argc; i++) {
+    if (vnet_readstring(nfd, recv_buf, READ_CHUNK_SIZE) == 0) {
+      lwip_close(nfd);
+      return 0;
+    }
+    g_oneshot_argv[i] = strdup(recv_buf);
+    log_debug("%s", recv_buf);
+  }
+  termtunnel_notify(_);
+  vnet_close(nfd);
+  return 0;
+}
+
+
 int server_call_agent(int32_t method, char *strbuf) {
   thread_arg_pass_t *tmp = malloc(sizeof(thread_arg_pass_t));
   tmp->method = method;

src/agentcall.h

@@ -8,6 +8,10 @@
 #ifndef TERMTUNNEL_AGENTCALL_H
 #define TERMTUNNEL_AGENTCALL_H
 #define METHOD_CALL_FORWARD_STATIC 1
+#define METHOD_GET_ARGS 2
 int agentcall_server_start();
 int server_call_agent(int32_t method, char *strbuf);
+typedef void (*get_args_callback)(void *handle);
+
+int bootstrap_get_args(void*);
 #endif

src/config.h

@@ -13,5 +13,6 @@
 #define TIMEOUT_MS 1000
 #define REPEAT_MS 100
 #define TTY_WATERMARK 100
+#define REPL_PROMPT "termtunnel> "
 
 #endif

src/entry.c

@@ -1,6 +1,6 @@
 /**
  * Copyright (c) 2022 Jindong Zhang
- * 
+ *
  * This software is released under the MIT License.
  * https://opensource.org/licenses/MIT
  */
@@ -22,7 +22,7 @@
 #include <sys/types.h>
 #include <termios.h>
 #include <unistd.h>
-
+#include "repl.h"
 #include "agent.h"
 #include "config.h"
 #include "fileexchange.h"
@@ -33,9 +33,9 @@
 #include "portforward.h"
 #include "pty.h"
 #include "repl.h"
-#include "thirdparty/setproctitle.h"
 #include "state.h"
 #include "thirdparty/base64.h"
+#include "thirdparty/setproctitle.h"
 #include "utils.h"
 #include "vnet.h"
 
@@ -59,11 +59,6 @@ void handler(int sig) {
 #endif
 
 static void on_resize(int signum) {
-  // if (first){
-  //     first=false;
-  //     return;
-  // }
-  // printf("%d\n",signum);
   struct winsize ttysize;  // The size of our tty
   int ttyfd = open("/dev/tty", O_RDONLY | O_NOCTTY);
   int r = ioctl(ttyfd, TIOCGWINSZ, &ttysize);
@@ -83,10 +78,8 @@ void cli_loop(int in, int out, int argc, const char *argv[]) {
     interact_run(in, out);
     kill(getpid(), SIGWINCH);
     // one shot 需要由 agent设定，然后传输给server，再控制到 client的行为。
-    bool one_shot_mode = false;
-    if (one_shot_mode) {
-      // one shot 仅由 cli 做出解释处理，不影响 server
-      //printf("oneshot!\n");
+    if (g_oneshot_mode) {
+      oneshot_run(in, out);
     } else {
       repl_run(in, out);
     }
@@ -142,11 +135,12 @@ int main(int argc, const char *argv[]) {
   termtunnel_state_init();
   q = queue_create();
   // spt_init(argc, argv);
-
+  // 其实这里仿照 lldb -- 去启动应用可能会更好？
   // if run as agent
-  if (argc == 2 && strcmp(argv[1], "-a") == 0) {
+  if (argc >= 2 && (strcmp(argv[1], "-a") == 0 || strcmp(argv[1], "--") == 0)) {
     log_info("agent pid %d", getpid());
-    agent();
+    argv+= 2;
+    agent(argc - 2, argv);
     return 0;
   }
 

src/fileexchange.c

@@ -128,6 +128,7 @@ static int file_send_request(path_exchange_t *pe) {
   if (fd < 0) {
     log_error("open_error");
     vnet_close(nfd);
+    set_running_task_changed(-1);
     return -1;
   }
   char buf[READ_CHUNK_SIZE];
@@ -152,7 +153,7 @@ static int file_send_request(path_exchange_t *pe) {
 }
 
 int file_send_start(char *src_path, char *dst_path) {
-
+  log_info("file_send_start1");
   path_exchange_t *pe = (path_exchange_t *)malloc(sizeof(path_exchange_t));
   strcpy(pe->src_path, src_path);
   strcpy(pe->dst_path, dst_path);
@@ -170,13 +171,14 @@ int file_send_start(char *src_path, char *dst_path) {
     // pe->exchange_notify->data=(void*)-1;
     // uv_async_send(pe->exchange_notify);
   }
+  log_info("file_send_start2");
   sys_thread_new("file_send", file_send_request, (void *)pe,
                  DEFAULT_THREAD_STACKSIZE, DEFAULT_THREAD_PRIO);
   return 0;
 }
 
 static int file_recv_request(path_exchange_t *pe) {
-  set_running_task_changed(1);
+  set_running_task_changed(1);  // TODO 计数机制问题
   int nfd = vnet_tcp_connect(sender_service_port);
   vnet_send(nfd, pe->src_path, strlen(pe->src_path) + 1);  // with_zero_as_split
   log_info("open %s for write to recv", pe->dst_path);
@@ -185,6 +187,7 @@ static int file_recv_request(path_exchange_t *pe) {
   if (fd < 0) {
     vnet_close(nfd);
     log_error("open error");
+    set_running_task_changed(-1);
     return 0;
   }
   char buf[READ_CHUNK_SIZE];

src/intent.h

@@ -29,6 +29,8 @@
 #define COMMAND_PORT_FORWARD 8
 #define COMMAND_RETURN 9
 #define COMMAND_GET_RUNNING_TASK_COUNT 10
+#define COMMAND_GET_ARGS 11
+#define COMMAND_GET_ARGS_REPLY 12
 typedef struct ci {
   int i;
   int mode;