From 6e05092055d20af2bded2b5c9fdca9a7b53512cb Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Thu, 18 Aug 2022 10:41:05 -0700
Subject: [PATCH] monitor: add common shell features to dev prompt
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This makes the monitor mode from —invoke more
user friendly by enabling features user might expect
from dev shell.

Interactive mode now has colors, history and keyboard
control for movement (arrows, moving to beginning/end,
by word, emacs controls etc).

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 monitor/monitor.go                            | 102 ++++++++++++++----
 .../c-bata/go-prompt/output_vt100.go          |   8 ++
 2 files changed, 90 insertions(+), 20 deletions(-)

diff --git a/monitor/monitor.go b/monitor/monitor.go
index fbddbdd6..07de8fcd 100644
--- a/monitor/monitor.go
+++ b/monitor/monitor.go
@@ -1,16 +1,15 @@
 package monitor
 
 import (
-	"bufio"
 	"context"
 	"fmt"
 	"io"
 	"sync"
 
+	prompt "github.com/c-bata/go-prompt"
 	"github.com/docker/buildx/build"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
-	"golang.org/x/term"
 )
 
 const helpMessage = `
@@ -69,16 +68,7 @@ func RunMonitor(ctx context.Context, containerConfig build.ContainerConfig, relo
 				<-ctx.Done()
 				in.Close()
 			}()
-			t := term.NewTerminal(readWriter{in.stdin, in.stdout}, "(buildx) ")
-			for {
-				l, err := t.ReadLine()
-				if err != nil {
-					if err != io.EOF {
-						errCh <- err
-						return
-					}
-					return
-				}
+			exec := func(l string) {
 				switch l {
 				case "":
 					// nop
@@ -104,6 +94,36 @@ func RunMonitor(ctx context.Context, containerConfig build.ContainerConfig, relo
 					fmt.Fprint(stdout, helpMessage)
 				}
 			}
+			completer := func(d prompt.Document) []prompt.Suggest {
+				return []prompt.Suggest{}
+			}
+
+			prompt.New(exec, completer, prompt.OptionHistory([]string{
+				"reload", "rollback", "exit", "help",
+			}),
+				prompt.OptionPrefix("(buildx) "),
+				prompt.OptionParser(&customParser{
+					PosixParser: *prompt.NewStandardInputParser(),
+					reader:      in.stdin,
+				}),
+				prompt.OptionWriter(&posixWriter{
+					w: in.stdout,
+				}),
+				prompt.OptionAddASCIICodeBind(
+					prompt.ASCIICodeBind{
+						ASCIICode: []byte{0x1b, 'f'},
+						Fn: func(b *prompt.Buffer) {
+							b.CursorRight(1 + b.Document().FindEndOfCurrentWordWithSpace())
+						},
+					},
+					prompt.ASCIICodeBind{
+						ASCIICode: []byte{0x1b, 'b'},
+						Fn: func(b *prompt.Buffer) {
+							b.CursorLeft(len(b.Document().TextBeforeCursor()) - b.Document().FindStartOfPreviousWordWithSpace())
+						},
+					},
+				),
+			).Run()
 		}()
 		select {
 		case <-doneCh:
@@ -307,12 +327,12 @@ func newMuxIO(in ioSetIn, out []ioSetOutContext, initIdx int, toggleMessage func
 		errToggle := errors.Errorf("toggle IO")
 		for {
 			prevIsControlSequence := false
-			if err := copyToFunc(traceReader(in.stdin, func(r rune) (bool, error) {
+			if err := copyToFunc(traceReader(in.stdin, func(r []byte) (bool, error) {
 				// Toggle IO when it detects C-a-c
 				// TODO: make it configurable if needed
-				if int(r) == 1 {
+				if len(r) == 1 && int(r[0]) == 1 {
 					prevIsControlSequence = true
-					return false, nil
+					return true, nil
 				}
 				defer func() { prevIsControlSequence = false }()
 				if prevIsControlSequence {
@@ -404,12 +424,12 @@ func (m *muxIO) toggleIO() {
 	fmt.Fprint(m.in.stdout, m.toggleMessage(prev, res))
 }
 
-func traceReader(r io.ReadCloser, f func(rune) (bool, error)) io.ReadCloser {
+func traceReader(r io.ReadCloser, f func([]byte) (bool, error)) io.ReadCloser {
+	dt := make([]byte, 32)
 	pr, pw := io.Pipe()
 	go func() {
-		br := bufio.NewReader(r)
 		for {
-			rn, _, err := br.ReadRune()
+			n, err := r.Read(dt)
 			if err != nil {
 				if err == io.EOF {
 					pw.Close()
@@ -418,13 +438,13 @@ func traceReader(r io.ReadCloser, f func(rune) (bool, error)) io.ReadCloser {
 				pw.CloseWithError(err)
 				return
 			}
-			if isWrite, err := f(rn); err != nil {
+			if isWrite, err := f(dt[:n]); err != nil {
 				pw.CloseWithError(err)
 				return
 			} else if !isWrite {
 				continue
 			}
-			if _, err := pw.Write([]byte(string(rn))); err != nil {
+			if _, err := pw.Write(dt[:n]); err != nil {
 				pw.CloseWithError(err)
 				return
 			}
@@ -514,3 +534,45 @@ type readerWithClose struct {
 func (r *readerWithClose) Close() error {
 	return r.closeFunc()
 }
+
+type customParser struct {
+	prompt.PosixParser
+	reader io.Reader
+}
+
+func (cp *customParser) Read() ([]byte, error) {
+	buf := make([]byte, 1024)
+	n, err := cp.reader.Read(buf)
+	if err != nil {
+		return []byte{}, err
+	}
+	return buf[:n], nil
+}
+
+type posixWriter struct {
+	prompt.VT100Writer
+	w io.Writer
+}
+
+func (w *posixWriter) Flush() error {
+	buf := w.Buffer()
+	l := len(buf)
+	offset := 0
+	retry := 0
+	for {
+		n, err := w.w.Write(buf[offset:])
+		if err != nil {
+			if retry < 3 {
+				retry++
+				continue
+			}
+			return err
+		}
+		offset += n
+		if offset == l {
+			break
+		}
+	}
+	w.SetBuffer([]byte{})
+	return nil
+}
diff --git a/vendor/github.com/c-bata/go-prompt/output_vt100.go b/vendor/github.com/c-bata/go-prompt/output_vt100.go
index 20850fea..feab8605 100644
--- a/vendor/github.com/c-bata/go-prompt/output_vt100.go
+++ b/vendor/github.com/c-bata/go-prompt/output_vt100.go
@@ -10,6 +10,14 @@ type VT100Writer struct {
 	buffer []byte
 }
 
+func (w *VT100Writer) Buffer() []byte {
+	return w.buffer
+}
+
+func (w *VT100Writer) SetBuffer(buffer []byte) {
+	w.buffer = buffer
+}
+
 // WriteRaw to write raw byte array
 func (w *VT100Writer) WriteRaw(data []byte) {
 	w.buffer = append(w.buffer, data...)