diff --git a/src/agent.c b/src/agent.c index 543b32f..77fd48e 100644 --- a/src/agent.c +++ b/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); } diff --git a/src/entry.c b/src/entry.c index e903d15..d4d5419 100644 --- a/src/entry.c +++ b/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); diff --git a/src/fileexchange.c b/src/fileexchange.c index 2d2bbd2..bc573b2 100644 --- a/src/fileexchange.c +++ b/src/fileexchange.c @@ -12,7 +12,7 @@ #include #include #include - +#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; } diff --git a/src/intent.h b/src/intent.h index 74c8d78..e850e37 100644 --- a/src/intent.h +++ b/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; diff --git a/src/pipe.c b/src/pipe.c index 9af0adb..8b30699 100644 --- a/src/pipe.c +++ b/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, diff --git a/src/pipe.h b/src/pipe.h index 0d44b2c..b251cd0 100644 --- a/src/pipe.h +++ b/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); diff --git a/src/state.c b/src/state.c index dbcb257..d56421a 100644 --- a/src/state.c +++ b/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; -} \ No newline at end of file +} + +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; +} diff --git a/src/state.h b/src/state.h index 49bbc7c..17d1ac9 100644 --- a/src/state.h +++ b/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 \ No newline at end of file diff --git a/src/utils.c b/src/utils.c index 3a190ce..5165c83 100644 --- a/src/utils.c +++ b/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 @@ -23,6 +20,7 @@ #include #include #include + //如果仅仅 ~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 \ No newline at end of file diff --git a/src/utils.h b/src/utils.h index bc6eb30..18e853b 100644 --- a/src/utils.h +++ b/src/utils.h @@ -9,6 +9,7 @@ #define TERMTUNNEL_UTILS_H #include +#include #include #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 \ No newline at end of file