From 06e8ba9f143428d5faeab1c358a67bcfdac2a348 Mon Sep 17 00:00:00 2001
From: afwerar <1296468573@qq.com>
Date: Mon, 27 Jun 2022 20:56:37 +0800
Subject: [PATCH] os: win shell catch ctrl-c

---
 include/os/osSystem.h            | 3 +++
 source/os/src/osSystem.c         | 9 ++++++---
 tools/shell/src/shellArguments.c | 4 ++--
 tools/shell/src/shellCommand.c   | 5 ++++-
 4 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/include/os/osSystem.h b/include/os/osSystem.h
index 6770be6e46..581e688ccb 100644
--- a/include/os/osSystem.h
+++ b/include/os/osSystem.h
@@ -29,6 +29,9 @@ extern "C" {
     #define tcgetattr TCGETATTR_FUNC_TAOS_FORBID
 #endif
 
+#define TAOS_CONSOLE_PROMPT_HEADER "taos> "
+#define TAOS_CONSOLE_PROMPT_CONTINUE "   -> "
+
 typedef struct TdCmd *TdCmdPtr;
 
 TdCmdPtr taosOpenCmd(const char* cmd);
diff --git a/source/os/src/osSystem.c b/source/os/src/osSystem.c
index ba07b6c3dd..9c66ac1df8 100644
--- a/source/os/src/osSystem.c
+++ b/source/os/src/osSystem.c
@@ -18,6 +18,10 @@
 #include "os.h"
 
 #if defined(WINDOWS)
+BOOL WINAPI CtrlHandler(DWORD fdwCtrlType) {
+  printf("\n" TAOS_CONSOLE_PROMPT_HEADER);
+  return TRUE;
+}
 #elif defined(_TD_DARWIN_64)
 #else
 #include <dlfcn.h>
@@ -124,7 +128,7 @@ int taosSetConsoleEcho(bool on) {
 
 void taosSetTerminalMode() {
 #if defined(WINDOWS)
-  // assert(0);
+  SetConsoleCtrlHandler(CtrlHandler, TRUE);
 
 #else
   struct termios newtio;
@@ -158,7 +162,6 @@ void taosSetTerminalMode() {
 
 int32_t taosGetOldTerminalMode() {
 #if defined(WINDOWS)
-  // assert(0);
 #else
   /* Make sure stdin is a terminal. */
   if (!isatty(STDIN_FILENO)) {
@@ -176,7 +179,7 @@ int32_t taosGetOldTerminalMode() {
 
 void taosResetTerminalMode() {
 #if defined(WINDOWS)
-  // assert(0);
+  SetConsoleCtrlHandler(CtrlHandler, FALSE);
 #else
   if (tcsetattr(0, TCSANOW, &oldtio) != 0) {
     fprintf(stderr, "Fail to reset the terminal properties!\n");
diff --git a/tools/shell/src/shellArguments.c b/tools/shell/src/shellArguments.c
index cd6613b17a..2037f67089 100644
--- a/tools/shell/src/shellArguments.c
+++ b/tools/shell/src/shellArguments.c
@@ -339,8 +339,8 @@ int32_t shellParseArgs(int32_t argc, char *argv[]) {
   shell.info.clientVersion =
       "Welcome to the TDengine shell from %s, Client Version:%s\n"
       "Copyright (c) 2022 by TAOS Data, Inc. All rights reserved.\n\n";
-  shell.info.promptHeader = "taos> ";
-  shell.info.promptContinue = "   -> ";
+  shell.info.promptHeader = TAOS_CONSOLE_PROMPT_HEADER;
+  shell.info.promptContinue = TAOS_CONSOLE_PROMPT_CONTINUE;
   shell.info.promptSize = 6;
   snprintf(shell.info.programVersion, sizeof(shell.info.programVersion), "version: %s", version);
 
diff --git a/tools/shell/src/shellCommand.c b/tools/shell/src/shellCommand.c
index f236c1eb88..9cb3541017 100644
--- a/tools/shell/src/shellCommand.c
+++ b/tools/shell/src/shellCommand.c
@@ -411,8 +411,9 @@ char taosGetConsoleChar() {
   static char mbStr[5];
   static unsigned long bufLen = 0;
   static uint16_t bufIndex = 0, mbStrIndex = 0, mbStrLen = 0;
-  if (bufLen == 0) {
+  while (bufLen == 0) {
     ReadConsoleW(console, buf, SHELL_INPUT_MAX_COMMAND_SIZE, &bufLen, NULL);
+    if (bufLen > 0 && buf[0] == 0) bufLen = 0;
     bufIndex = 0;
   }
   if (mbStrLen == 0){
@@ -466,6 +467,8 @@ int32_t shellReadCommand(char *command) {
     } else if (c < '\033') {
       // Ctrl keys.  TODO: Implement ctrl combinations
       switch (c) {
+        case 0:
+          break;
         case 1:  // ctrl A
           shellPositionCursorHome(&cmd);
           break;