Merge pull request #26 from beordle/file_improve

improve interaction

src/agent.c

@@ -191,7 +191,7 @@ int agent_process_frame(char *str_data, int data_size) {
   // simple command
   if (data_size == strlen("EXIT") && strcmp(str_data, "EXIT") == 0) {
     agent_restore_stdin();
-    printf("\r\n");  // pretty
+    printf("\r");  // pretty
     log_info("exit");
     exit(EXIT_SUCCESS);
   }

src/entry.c

@@ -76,6 +76,36 @@ static void on_resize(int signum) {
               sizeof(struct winsize));
 }
 
+
+void cli_loop(int in, int out, int argc, const char *argv[]) {
+  repl_init();
+  while (true) {
+    interact_run(in, out);
+    kill(getpid(), SIGWINCH);
+    bool one_shot_mode = false;
+    if (one_shot_mode) {
+      // one shot 仅由 cli 做出解释处理，不影响 server
+      //printf("oneshot!\n");
+    } else {
+      repl_run(in, out);
+    }
+    kill(getpid(), SIGWINCH);
+  }
+}
+
+void cli(int argc, const char *argv[], pid_t child_pid) {
+  set_client_process();
+  // Close write end
+  close(in_fd[1]);
+  close(out_fd[0]);
+
+  cli_loop(in_fd[0], out_fd[1], argc, argv);
+  close(in_fd[0]);
+  close(out_fd[1]);
+  kill(child_pid, SIGTERM);
+  exit(EXIT_SUCCESS);
+}
+
 int main(int argc, const char *argv[]) {
   signal(SIGTTIN, SIG_IGN);
   signal(SIGTTOU, SIG_IGN);
@@ -108,7 +138,7 @@ int main(int argc, const char *argv[]) {
       log_add_fp(f, verbose_level);
     }
   }
-
+  termtunnel_state_init();
   q = queue_create();
   // spt_init(argc, argv);
 

src/fileexchange.c

@@ -12,7 +12,7 @@
 #include <sys/stat.h>
 #include <sys/types.h>
 #include <unistd.h>
-
+#include "state.h"
 #include "config.h"
 #include "intent.h"
 #include "log.h"
@@ -120,6 +120,7 @@ int file_sender_start() {
 }
 
 static int file_send_request(path_exchange_t *pe) {
+  set_running_task_changed(1);
   int nfd = vnet_tcp_connect(receiver_service_port);
   log_debug("open local %s to send file %s", pe->src_path, pe->dst_path);
   vnet_send(nfd, pe->dst_path, strlen(pe->dst_path) + 1);  // with_zero_as_split
@@ -146,10 +147,12 @@ static int file_send_request(path_exchange_t *pe) {
   // check path is exist, file can writeable.
   vnet_close(nfd);
   close(fd);
+  set_running_task_changed(-1);
   return 0;
 }
 
 int file_send_start(char *src_path, char *dst_path) {
+
   path_exchange_t *pe = (path_exchange_t *)malloc(sizeof(path_exchange_t));
   strcpy(pe->src_path, src_path);
   strcpy(pe->dst_path, dst_path);
@@ -173,6 +176,7 @@ int file_send_start(char *src_path, char *dst_path) {
 }
 
 static int file_recv_request(path_exchange_t *pe) {
+  set_running_task_changed(1);
   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);
@@ -201,6 +205,7 @@ static int file_recv_request(path_exchange_t *pe) {
   // TODO(jdz) check path is exist, file can writeable.
   vnet_close(nfd);
   close(fd);
+  set_running_task_changed(-1);
   return 0;
 }
 

src/intent.h

@@ -28,6 +28,7 @@
 #define COMMAND_FILE_EXCHANGE 7
 #define COMMAND_PORT_FORWARD 8
 #define COMMAND_RETURN 9
+#define COMMAND_GET_RUNNING_TASK_COUNT 10
 typedef struct ci {
   int i;
   int mode;

src/pipe.c

@@ -42,6 +42,7 @@
 #include "thirdparty/queue/queue_internal.h"
 #include "thirdparty/setproctitle.h"
 #include "utils.h"
+#include "state.h"
 #include "vnet.h"
 static uv_pipe_t in_pipe;
 static uv_pipe_t out_pipe;
@@ -344,9 +345,6 @@ static void common_read_tty(uv_stream_t *stream, ssize_t nread,
     uv_read_stop(stream);
     return;
   } else {
-    if (!server_see_agent_is_repl) {
-      send_tty_to_client(buf->base, nread);
-    }
     fsm_append_input(global_fsm_context, buf->base, nread);
     fsm_run(global_fsm_context);
     char *dst = (char *)malloc(5 * VIR_MTU);
@@ -359,6 +357,9 @@ static void common_read_tty(uv_stream_t *stream, ssize_t nread,
         break;
       }
     }
+    if (!server_see_agent_is_repl) {
+      send_tty_to_client(buf->base, nread);
+    }
     free(dst);
   }
 
@@ -380,6 +381,8 @@ int push_data() {
 }
 
 void timer_callback() {
+
+  log_info("get_running_task_count %d", get_running_task_count());
   // 如果队列没有清空，那么就提醒一下。因为async_send是一个不可靠的提醒，另外，提醒后，因为没有逻辑锁，会出现：
   // uv_aysnc_send:call queue_pop queue_push 的情况，从而导致残留
 #ifdef FLUASH_QUEUE_ON_TIMER
@@ -397,31 +400,6 @@ void timer_callback() {
   }
 }
 
-void cli_loop(int in, int out, int argc, const char *argv[]) {
-  repl_init();
-
-  // 不要贸然进入模式
-  while (true) {
-    interact_run(in, out);
-    kill(getpid(), SIGWINCH);
-    repl_run(in, out);
-    kill(getpid(), SIGWINCH);
-  }
-}
-
-void cli(int argc, const char *argv[], pid_t child_pid) {
-  set_client_process();
-  // Close write end
-  close(in_fd[1]);
-  close(out_fd[0]);
-
-  cli_loop(in_fd[0], out_fd[1], argc, argv);
-  close(in_fd[0]);
-  close(out_fd[1]);
-  kill(child_pid, SIGTERM);
-  exit(EXIT_SUCCESS);
-}
-
 static int pty_nonblock(int fd) {
   int flags = fcntl(fd, F_GETFL, 0);
   if (flags == -1) return -1;
@@ -572,6 +550,12 @@ void server_handle_client_packet(int64_t type, char *buf, ssize_t len) {
       log_info("start ok");
       break;
     }
+    case COMMAND_GET_RUNNING_TASK_COUNT: {
+
+      get_running_task_count();
+      //comm_write_packet_to_cli(COMMAND_RETURN, strdup("bind done (guess)\n"),
+      //                        sizeof("bind done (guess)\n"));
+    }
     case COMMAND_PORT_FORWARD: {
       port_forward_intent_t *a = (port_forward_intent_t *)buf;
       log_info("portforward %s:%hu <-> %s:%hu", a->src_host, a->src_port,

src/pipe.h

@@ -18,7 +18,6 @@ extern int in_fd[2];
 extern int out_fd[2];
 extern void agent_write_data_to_server(char *buf, size_t s, bool autofree);
 extern void send_base64binary_to_agent(const char *buf, size_t size);
-void cli();
 void server();
 int libuv_add_vnet_notify();
 extern int vnet_notify_to_libuv(char *buf, size_t size);

src/state.c

@@ -6,6 +6,7 @@
  */
 
 #include "state.h"
+#include "utils.h"
 static int state_mode;
 
 int get_state_mode() { return state_mode; }
@@ -24,4 +25,20 @@ int set_client_process() {
 int set_agent_process() {
   state_mode = MODE_AGENT_PROCESS;
   return 0;
-}
+}
+
+static struct counter_t task_counter;
+
+void set_running_task_changed(int value) {
+  utils_counter_increment_by(&task_counter, value);
+  return;
+}
+
+int get_running_task_count() {
+  return utils_counter_get(&task_counter);
+}
+
+int termtunnel_state_init() {
+  utils_counter_init(&task_counter);
+  return 0;
+}

src/state.h

@@ -19,4 +19,7 @@ extern int get_state_mode();
 extern int set_server_process();
 extern int set_client_process();
 extern int set_agent_process();
+extern int termtunnel_state_init();
+extern int get_running_task_count();
+extern void set_running_task_changed(int value);
 #endif

src/utils.c

@@ -5,9 +5,6 @@
  * https://opensource.org/licenses/MIT
  */
 
-#ifndef TERMTUNNEL_UTILS_H
-
-#define TERMTUNNEL_UTILS_H
 #include "utils.h"
 #include "log.h"
 #include <arpa/inet.h>
@@ -23,6 +20,7 @@
 #include <termios.h>
 #include <unistd.h>
 #include <errno.h>
+
 //如果仅仅 ~ICANON; ~ECHO
 //那么，我们还是可以通过键盘触发SIGINT，SIGSTOP等信号。此时可以手动选择转发给远端。
 //但是依据ssh，这样设置可以直接将信号转发给远程终端，从而如果我们的程序收到了键盘产生的信号的话，一定是从另一个console中手动kill过来的。
@@ -99,6 +97,27 @@ void set_stdin_raw() {
   _stdin_is_raw = true;
 }
 
+void utils_counter_init(struct counter_t *c) {
+  c->value = 0;
+  pthread_mutex_init(&c->lock, NULL);
+}
+
+void utils_counter_increment_by(struct counter_t *c, int by) {
+  pthread_mutex_lock(&c->lock);
+  c->value += by;
+  pthread_mutex_unlock(&c->lock);
+}
+
+void utils_counter_increment(struct counter_t *c) {
+  utils_counter_increment_by(c, 1);
+}
+
+int utils_counter_get(struct counter_t *c) {
+  pthread_mutex_lock(&c->lock);
+  int rc = c->value;
+  pthread_mutex_unlock(&c->lock);
+  return rc;
+}
 
 char* safe_gethostbyname(char *host, uint16_t port) {
     log_info("gethostbyname %s:%hu", host, port);
@@ -127,5 +146,3 @@ char* safe_gethostbyname(char *host, uint16_t port) {
     return NULL;
 }
 
-
-#endif

src/utils.h

@@ -9,6 +9,7 @@
 
 #define TERMTUNNEL_UTILS_H
 #include <assert.h>
+#include <pthread.h>
 #include <stdint.h>
 #include "log.h"
 #define CHECK(x, ...)       \
@@ -17,10 +18,23 @@
     exit(0);                \
   }
 
+struct counter_t {
+  int value;
+  pthread_mutex_t lock;
+};
+
+
 extern int writen(int fd, void *buf, int n);
 extern void set_stdin_raw();
 extern void restore_stdin();
 extern void *memdup(const void *src, size_t n);
 extern const char *green_encode(const char *buf, int len, int *result_len);
+
+
+extern void utils_counter_init(struct counter_t *c);
+extern void utils_counter_increment_by(struct counter_t *c, int by);
+extern void utils_counter_increment(struct counter_t *c);
+extern int utils_counter_get(struct counter_t *c);
+
 extern char* safe_gethostbyname(char *buf, uint16_t port);
 #endif