From 62ab2cfd17947c6843ad015eb81c0c675ec22630 Mon Sep 17 00:00:00 2001
From: Thomas O'Dowd <tpodowd@cloudian.com>
Date: Fri, 3 Mar 2023 18:01:34 +0900
Subject: [PATCH] Fix race condition between writing and prompting.

Introduce a new operation.isPrompting field which is true from when a
prompt has been written until data is returned to the caller.
When Write is called on the wrapWriter to write to stdout or stderr,
check if we are currently prompting the user for input and if so
clean up the prompt and write the data before redrawing the
prompt at its new location after the written data.

Previously terminal.IsReading() was used for this, but this had various
race conditions and it was not correct to check this field to make
prompt and buffer redrawing decisions. In turn, I removed all the
isReading code also. The old isReading() check was actually checking
if the terminal goroutine was actively waiting for more input.
---
 operation.go | 65 +++++++++++++++++++++++++++++++++++++++++-----------
 terminal.go  |  8 -------
 utils.go     |  6 ++++-
 3 files changed, 56 insertions(+), 23 deletions(-)

diff --git a/operation.go b/operation.go
index 3641a9e..4bf6fa0 100644
--- a/operation.go
+++ b/operation.go
@@ -27,6 +27,8 @@ type Operation struct {
 	errchan chan error
 	w       io.Writer
 
+	isPrompting bool       // true when prompt written and waiting for input
+
 	history *opHistory
 	*opSearch
 	*opCompleter
@@ -39,29 +41,43 @@ func (o *Operation) SetBuffer(what string) {
 }
 
 type wrapWriter struct {
-	r      *Operation
-	t      *Terminal
+	o      *Operation
 	target io.Writer
 }
 
 func (w *wrapWriter) Write(b []byte) (int, error) {
-	if !w.t.IsReading() {
-		return w.target.Write(b)
+	return w.o.write(w.target, b)
+}
+
+func (o *Operation) write(target io.Writer, b []byte) (int, error) {
+	o.m.Lock()
+	defer o.m.Unlock()
+
+	if !o.isPrompting {
+		return target.Write(b)
 	}
 
 	var (
 		n   int
 		err error
 	)
-	w.r.buf.Refresh(func() {
-		n, err = w.target.Write(b)
+	o.buf.Refresh(func() {
+		n, err = target.Write(b)
+		// Adjust the prompt start position by b
+		rout := runes.ColorFilter([]rune(string(b[:])))
+		sp := SplitByLine(rout, []rune{}, o.buf.ppos, o.buf.width, 1)
+		if len(sp) > 1 {
+			o.buf.ppos = len(sp[len(sp)-1])
+		} else {
+			o.buf.ppos += len(rout)
+		}
 	})
 
-	if w.r.IsSearchMode() {
-		w.r.SearchRefresh(-1)
+	if o.IsSearchMode() {
+		o.SearchRefresh(-1)
 	}
-	if w.r.IsInCompleteMode() {
-		w.r.CompleteRefresh()
+	if o.IsInCompleteMode() {
+		o.CompleteRefresh()
 	}
 	return n, err
 }
@@ -352,7 +368,7 @@ func (o *Operation) ioloop() {
 		} else if o.IsInCompleteMode() {
 			if !keepInCompleteMode {
 				o.ExitCompleteMode(false)
-				o.Refresh()
+				o.refresh()
 			} else {
 				o.buf.Refresh(nil)
 				o.CompleteRefresh()
@@ -367,11 +383,11 @@ func (o *Operation) ioloop() {
 }
 
 func (o *Operation) Stderr() io.Writer {
-	return &wrapWriter{target: o.GetConfig().Stderr, r: o, t: o.t}
+	return &wrapWriter{target: o.GetConfig().Stderr, o: o}
 }
 
 func (o *Operation) Stdout() io.Writer {
-	return &wrapWriter{target: o.GetConfig().Stdout, r: o, t: o.t}
+	return &wrapWriter{target: o.GetConfig().Stdout, o: o}
 }
 
 func (o *Operation) String() (string, error) {
@@ -388,6 +404,11 @@ func (o *Operation) Runes() ([]rune, error) {
 		listener.OnChange(nil, 0, 0)
 	}
 
+	// Before writing the prompt and starting to read, get a lock
+	// so we don't race with wrapWriter trying to write and refresh.
+	o.m.Lock()
+	o.isPrompting = true
+
 	// Query cursor position before printing the prompt as there
 	// maybe existing text on the same line that ideally we don't
 	// want to overwrite and cause prompt to jump left. Note that
@@ -396,6 +417,16 @@ func (o *Operation) Runes() ([]rune, error) {
 	o.buf.Print() // print prompt & buffer contents
 	o.t.KickRead()
 
+	// Prompt written safely, unlock until read completes and then
+	// lock again to unset.
+	o.m.Unlock()
+	defer func() {
+		o.m.Lock()
+		o.isPrompting = false
+		o.buf.SetOffset("1;1")
+		o.m.Unlock()
+	}()
+
 	select {
 	case r := <-o.outchan:
 		return r, nil
@@ -508,7 +539,13 @@ func (o *Operation) SaveHistory(content string) error {
 }
 
 func (o *Operation) Refresh() {
-	if o.t.IsReading() {
+	o.m.Lock()
+	defer o.m.Unlock()
+	o.refresh()
+}
+
+func (o *Operation) refresh() {
+	if o.isPrompting {
 		o.buf.Refresh(nil)
 	}
 }
diff --git a/terminal.go b/terminal.go
index 01f268e..d1faba1 100644
--- a/terminal.go
+++ b/terminal.go
@@ -17,7 +17,6 @@ type Terminal struct {
 	stopChan  chan struct{}
 	kickChan  chan struct{}
 	wg        sync.WaitGroup
-	isReading int32
 	sleeping  int32
 
 	sizeChan chan string
@@ -104,10 +103,6 @@ func (t *Terminal) ReadRune() rune {
 	return ch
 }
 
-func (t *Terminal) IsReading() bool {
-	return atomic.LoadInt32(&t.isReading) == 1
-}
-
 func (t *Terminal) KickRead() {
 	select {
 	case t.kickChan <- struct{}{}:
@@ -132,10 +127,8 @@ func (t *Terminal) ioloop() {
 	buf := bufio.NewReader(t.getStdin())
 	for {
 		if !expectNextChar {
-			atomic.StoreInt32(&t.isReading, 0)
 			select {
 			case <-t.kickChan:
-				atomic.StoreInt32(&t.isReading, 1)
 			case <-t.stopChan:
 				return
 			}
@@ -210,7 +203,6 @@ func (t *Terminal) ioloop() {
 			t.outchan <- r
 		}
 	}
-
 }
 
 func (t *Terminal) Bell() {
diff --git a/utils.go b/utils.go
index 9451a1c..ea73428 100644
--- a/utils.go
+++ b/utils.go
@@ -224,7 +224,11 @@ func SplitByLine(prompt, rs []rune, offset, screenWidth, nextWidth int) [][]rune
 	currentWidth := offset
 	for i, r := range prs {
 		w := runes.Width(r)
-		if currentWidth + w > screenWidth {
+		if r == '\n' {
+			ret = append(ret, prs[si:i+1])
+			si = i + 1
+			currentWidth = 0
+		} else if currentWidth + w > screenWidth {
 			ret = append(ret, prs[si:i])
 			si = i
 			currentWidth = 0