Thanks to visit codestin.com
Credit goes to github.com

Skip to content

fix: fix loss of buffered input on cliui.Prompt #15421

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 7, 2024
Merged

fix: fix loss of buffered input on cliui.Prompt #15421

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 31 additions & 6 deletions cli/cliui/prompt.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
package cliui

import (
"bufio"
"bytes"
"encoding/json"
"fmt"
"io"
"os"
"os/signal"
"strings"
Expand Down Expand Up @@ -96,14 +96,13 @@ func Prompt(inv *serpent.Invocation, opts PromptOptions) (string, error) {
signal.Notify(interrupt, os.Interrupt)
defer signal.Stop(interrupt)

reader := bufio.NewReader(inv.Stdin)
line, err = reader.ReadString('\n')
line, err = readUntil(inv.Stdin, '\n')

// Check if the first line beings with JSON object or array chars.
// This enables multiline JSON to be pasted into an input, and have
// it parse properly.
if err == nil && (strings.HasPrefix(line, "{") || strings.HasPrefix(line, "[")) {
line, err = promptJSON(reader, line)
line, err = promptJSON(inv.Stdin, line)
}
}
if err != nil {
Expand Down Expand Up @@ -144,7 +143,7 @@ func Prompt(inv *serpent.Invocation, opts PromptOptions) (string, error) {
}
}

func promptJSON(reader *bufio.Reader, line string) (string, error) {
func promptJSON(reader io.Reader, line string) (string, error) {
var data bytes.Buffer
for {
_, _ = data.WriteString(line)
Expand All @@ -162,7 +161,7 @@ func promptJSON(reader *bufio.Reader, line string) (string, error) {
// Read line-by-line. We can't use a JSON decoder
// here because it doesn't work by newline, so
// reads will block.
line, err = reader.ReadString('\n')
line, err = readUntil(reader, '\n')
if err != nil {
break
}
Expand All @@ -179,3 +178,29 @@ func promptJSON(reader *bufio.Reader, line string) (string, error) {
}
return line, nil
}

// readUntil the first occurrence of delim in the input, returning a string containing the data up
// to and including the delimiter. Unlike `bufio`, it only reads until the delimiter and no further
// bytes. If readUntil encounters an error before finding a delimiter, it returns the data read
// before the error and the error itself (often io.EOF). readUntil returns err != nil if and only if
// the returned data does not end in delim.
func readUntil(r io.Reader, delim byte) (string, error) {
var (
have []byte
b = make([]byte, 1)
)
for {
n, err := r.Read(b)
if n > 0 {
have = append(have, b[0])
if b[0] == delim {
// match `bufio` in that we only return non-nil if we didn't find the delimiter,
// regardless of whether we also erred.
return string(have), nil
}
}
if err != nil {
return string(have), err
}
}
}
68 changes: 53 additions & 15 deletions cli/cliui/prompt_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/xerrors"

"github.com/coder/coder/v2/cli/cliui"
"github.com/coder/coder/v2/pty"
Expand All @@ -22,26 +23,29 @@ func TestPrompt(t *testing.T) {
t.Parallel()
t.Run("Success", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
ptty := ptytest.New(t)
msgChan := make(chan string)
go func() {
resp, err := newPrompt(ptty, cliui.PromptOptions{
resp, err := newPrompt(ctx, ptty, cliui.PromptOptions{
Text: "Example",
}, nil)
assert.NoError(t, err)
msgChan <- resp
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("hello")
require.Equal(t, "hello", <-msgChan)
resp := testutil.RequireRecvCtx(ctx, t, msgChan)
require.Equal(t, "hello", resp)
})

t.Run("Confirm", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
ptty := ptytest.New(t)
doneChan := make(chan string)
go func() {
resp, err := newPrompt(ptty, cliui.PromptOptions{
resp, err := newPrompt(ctx, ptty, cliui.PromptOptions{
Text: "Example",
IsConfirm: true,
}, nil)
Expand All @@ -50,18 +54,20 @@ func TestPrompt(t *testing.T) {
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("yes")
require.Equal(t, "yes", <-doneChan)
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
require.Equal(t, "yes", resp)
})

t.Run("Skip", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
ptty := ptytest.New(t)
var buf bytes.Buffer

// Copy all data written out to a buffer. When we close the ptty, we can
// no longer read from the ptty.Output(), but we can read what was
// written to the buffer.
dataRead, doneReading := context.WithTimeout(context.Background(), testutil.WaitShort)
dataRead, doneReading := context.WithCancel(ctx)
go func() {
// This will throw an error sometimes. The underlying ptty
// has its own cleanup routines in t.Cleanup. Instead of
Expand All @@ -74,7 +80,7 @@ func TestPrompt(t *testing.T) {

doneChan := make(chan string)
go func() {
resp, err := newPrompt(ptty, cliui.PromptOptions{
resp, err := newPrompt(ctx, ptty, cliui.PromptOptions{
Text: "ShouldNotSeeThis",
IsConfirm: true,
}, func(inv *serpent.Invocation) {
Expand All @@ -85,7 +91,8 @@ func TestPrompt(t *testing.T) {
doneChan <- resp
}()

require.Equal(t, "yes", <-doneChan)
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
require.Equal(t, "yes", resp)
// Close the reader to end the io.Copy
require.NoError(t, ptty.Close(), "close eof reader")
// Wait for the IO copy to finish
Expand All @@ -96,42 +103,47 @@ func TestPrompt(t *testing.T) {
})
t.Run("JSON", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
ptty := ptytest.New(t)
doneChan := make(chan string)
go func() {
resp, err := newPrompt(ptty, cliui.PromptOptions{
resp, err := newPrompt(ctx, ptty, cliui.PromptOptions{
Text: "Example",
}, nil)
assert.NoError(t, err)
doneChan <- resp
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("{}")
require.Equal(t, "{}", <-doneChan)
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
require.Equal(t, "{}", resp)
})

t.Run("BadJSON", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
ptty := ptytest.New(t)
doneChan := make(chan string)
go func() {
resp, err := newPrompt(ptty, cliui.PromptOptions{
resp, err := newPrompt(ctx, ptty, cliui.PromptOptions{
Text: "Example",
}, nil)
assert.NoError(t, err)
doneChan <- resp
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("{a")
require.Equal(t, "{a", <-doneChan)
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
require.Equal(t, "{a", resp)
})

t.Run("MultilineJSON", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
ptty := ptytest.New(t)
doneChan := make(chan string)
go func() {
resp, err := newPrompt(ptty, cliui.PromptOptions{
resp, err := newPrompt(ctx, ptty, cliui.PromptOptions{
Text: "Example",
}, nil)
assert.NoError(t, err)
Expand All @@ -141,11 +153,37 @@ func TestPrompt(t *testing.T) {
ptty.WriteLine(`{
"test": "wow"
}`)
require.Equal(t, `{"test":"wow"}`, <-doneChan)
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
require.Equal(t, `{"test":"wow"}`, resp)
})

t.Run("InvalidValid", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
ptty := ptytest.New(t)
doneChan := make(chan string)
go func() {
resp, err := newPrompt(ctx, ptty, cliui.PromptOptions{
Text: "Example",
Validate: func(s string) error {
t.Logf("validate: %q", s)
if s != "valid" {
return xerrors.New("invalid")
}
return nil
},
}, nil)
assert.NoError(t, err)
doneChan <- resp
}()
ptty.ExpectMatch("Example")
ptty.WriteLine("foo\nbar\nbaz\n\n\nvalid\n")
resp := testutil.RequireRecvCtx(ctx, t, doneChan)
require.Equal(t, "valid", resp)
})
}

func newPrompt(ptty *ptytest.PTY, opts cliui.PromptOptions, invOpt func(inv *serpent.Invocation)) (string, error) {
func newPrompt(ctx context.Context, ptty *ptytest.PTY, opts cliui.PromptOptions, invOpt func(inv *serpent.Invocation)) (string, error) {
value := ""
cmd := &serpent.Command{
Handler: func(inv *serpent.Invocation) error {
Expand All @@ -163,7 +201,7 @@ func newPrompt(ptty *ptytest.PTY, opts cliui.PromptOptions, invOpt func(inv *ser
inv.Stdout = ptty.Output()
inv.Stderr = ptty.Output()
inv.Stdin = ptty.Input()
return value, inv.WithContext(context.Background()).Run()
return value, inv.WithContext(ctx).Run()
}

func TestPasswordTerminalState(t *testing.T) {
Expand Down
Loading