From 664a2143767b0cf67e3fa731e4820e04ed572997 Mon Sep 17 00:00:00 2001
From: Congqi Xia <congqi.xia@zilliz.com>
Date: Tue, 13 Dec 2022 18:23:08 +0800
Subject: [PATCH] Fix go-prompt caused bash hijack

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 main.go          | 31 +++++++++++++++++++++++++------
 states/exit.go   |  4 ++++
 states/states.go | 15 ++++++++++-----
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/main.go b/main.go
index 5eb9c17..bd0d0a9 100644
--- a/main.go
+++ b/main.go
@@ -19,10 +19,18 @@ var (
 )
 
 func main() {
+	defer handleExit()
 	app := states.Start()
 	runPrompt(app)
 }
 
+func handleExit() {
+	rawModeOff := exec.Command("/bin/stty", "-raw", "echo")
+	rawModeOff.Stdin = os.Stdin
+	_ = rawModeOff.Run()
+	rawModeOff.Wait()
+}
+
 // run start BirdWatcher with promptui. (disable suggestion and history)
 func run(app states.State) {
 	for {
@@ -45,13 +53,13 @@ func run(app states.State) {
 
 // promptApp model wraps states to provide function for go-prompt.
 type promptApp struct {
+	exited       bool
 	currentState states.State
 }
 
 // promptExecute actual execution logic entry.
 func (a *promptApp) promptExecute(in string) {
 	in = strings.TrimSpace(in)
-	var err error
 
 	// try to get $PAGER env
 	pager := os.Getenv("PAGER")
@@ -93,10 +101,10 @@ func (a *promptApp) promptExecute(in string) {
 			}()
 		}
 	}
-
-	a.currentState, err = a.currentState.Process(in)
-	if errors.Is(err, states.ExitErr) {
-		os.Exit(0)
+	a.currentState, _ = a.currentState.Process(in)
+	if a.currentState.IsEnding() {
+		fmt.Println("Bye!")
+		a.exited = true
 	}
 }
 
@@ -119,6 +127,9 @@ func (a *promptApp) completeInput(d prompt.Document) []prompt.Suggest {
 
 // livePrefix implements dynamic change prefix.
 func (a *promptApp) livePrefix() (string, bool) {
+	if a.exited {
+		return "", false
+	}
 	return fmt.Sprintf("%s > ", a.currentState.Label()), true
 }
 
@@ -132,6 +143,14 @@ func runPrompt(app states.State) {
 		prompt.OptionPrefixTextColor(prompt.Yellow),
 		prompt.OptionPreviewSuggestionTextColor(prompt.Blue),
 		prompt.OptionSelectedSuggestionBGColor(prompt.LightGray),
-		prompt.OptionSuggestionBGColor(prompt.DarkGray))
+		prompt.OptionSuggestionBGColor(prompt.DarkGray),
+		prompt.OptionSetExitCheckerOnInput(func(in string, breakline bool) bool {
+			// setup exit command
+			if strings.ToLower(in) == "exit" && breakline {
+				return true
+			}
+			return false
+		}),
+	)
 	p.Run()
 }
diff --git a/states/exit.go b/states/exit.go
index f42bf96..1a93c18 100644
--- a/states/exit.go
+++ b/states/exit.go
@@ -21,6 +21,7 @@ func getExitCmd(state State) *cobra.Command {
 		Aliases: []string{"quit"},
 		RunE: func(*cobra.Command, []string) error {
 			state.SetNext(&exitState{})
+			// cannot return ExitErr here to avoid print help message
 			return nil
 		},
 	}
@@ -36,6 +37,9 @@ type exitState struct {
 // also called after each command run to reset flag values.
 func (s *exitState) SetupCommands() {}
 
+// IsEnding returns true for exit State
+func (s *exitState) IsEnding() bool { return true }
+
 // getDisconnectCmd disconnect from current state.
 // will call close method for current state.
 func getDisconnectCmd(state State) *cobra.Command {
diff --git a/states/states.go b/states/states.go
index f8cb51b..cbb6688 100644
--- a/states/states.go
+++ b/states/states.go
@@ -1,6 +1,7 @@
 package states
 
 import (
+	"errors"
 	"fmt"
 	"strings"
 
@@ -16,6 +17,7 @@ type State interface {
 	SetNext(state State)
 	Suggestions(input string) map[string]string
 	SetupCommands()
+	IsEnding() bool
 }
 
 // cmdState is the basic state to process input command.
@@ -133,15 +135,13 @@ func (s *cmdState) Process(cmd string) (State, error) {
 	}
 	s.rootCmd.SetArgs(args)
 	err = s.rootCmd.Execute()
+	if errors.Is(err, ExitErr) {
+		return s.nextState, ExitErr
+	}
 	if err != nil {
 		return s, err
 	}
 	if s.nextState != nil {
-		//defer s.Close()
-		// TODO fix ugly type cast
-		if _, ok := s.nextState.(*exitState); ok {
-			return s.nextState, ExitErr
-		}
 		nextState := s.nextState
 		s.nextState = nil
 		return nextState, nil
@@ -160,6 +160,9 @@ func (s *cmdState) SetNext(state State) {
 // Close empty method to implement State.
 func (s *cmdState) Close() {}
 
+// Check
+func (s *cmdState) IsEnding() bool { return false }
+
 // Start returns the first state - offline.
 func Start() State {
 	root := &cobra.Command{
@@ -177,6 +180,8 @@ func Start() State {
 		getConnectCommand(state),
 		// load-backup
 		getLoadBackupCmd(state),
+		// pulsarctl
+		getPulsarctlCmd(state),
 		// exit
 		getExitCmd(state))