Skip to content

Commit 55a3f63

Browse files
ammarioammario
committed
Merge remote-tracking branch 'origin/main' into agent-metadata
2 parents dc631f5 + a8346bd commit 55a3f63

14 files changed

+240
-203
lines changed

‎agent/agent.go

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1273,13 +1273,28 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
12731273
go func() {
12741274
_, _ = io.Copy(ptty.Input(), session)
12751275
}()
1276+
1277+
// In low parallelism scenarios, the command may exit and we may close
1278+
// the pty before the output copy has started. This can result in the
1279+
// output being lost. To avoid this, we wait for the output copy to
1280+
// start before waiting for the command to exit. This ensures that the
1281+
// output copy goroutine will be scheduled before calling close on the
1282+
// pty. There is still a risk of data loss if a command produces a lot
1283+
// of output, see TestAgent_Session_TTY_HugeOutputIsNotLost (skipped).
1284+
outputCopyStarted := make(chan struct{})
1285+
ptyOutput := func() io.Reader {
1286+
defer close(outputCopyStarted)
1287+
return ptty.Output()
1288+
}
12761289
wg.Add(1)
12771290
go func() {
12781291
// Ensure data is flushed to session on command exit, if we
12791292
// close the session too soon, we might lose data.
12801293
defer wg.Done()
1281-
_, _ = io.Copy(session, ptty.Output())
1294+
_, _ = io.Copy(session, ptyOutput())
12821295
}()
1296+
<-outputCopyStarted
1297+
12831298
err = process.Wait()
12841299
var exitErr *exec.ExitError
12851300
// ExitErrors just mean the command we run returned a non-zero exit code, which is normal

‎agent/agent_test.go

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -374,9 +374,12 @@ func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
374374

375375
var stdout bytes.Buffer
376376
// NOTE(mafredri): Increase iterations to increase chance of failure,
377-
// assuming bug is present.
377+
// assuming bug is present. Limiting GOMAXPROCS further
378+
// increases the chance of failure.
378379
// Using 1000 iterations is basically a guaranteed failure (but let's
379380
// not increase test times needlessly).
381+
// Limit GOMAXPROCS (e.g. `export GOMAXPROCS=1`) to further increase
382+
// chance of failure. Also -race helps.
380383
for i := 0; i < 5; i++ {
381384
func() {
382385
stdout.Reset()
@@ -400,6 +403,63 @@ func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
400403
}
401404
}
402405

406+
func TestAgent_Session_TTY_HugeOutputIsNotLost(t *testing.T) {
407+
t.Parallel()
408+
if runtime.GOOS == "windows" {
409+
// This might be our implementation, or ConPTY itself.
410+
// It's difficult to find extensive tests for it, so
411+
// it seems like it could be either.
412+
t.Skip("ConPTY appears to be inconsistent on Windows.")
413+
}
414+
t.Skip("This test proves we have a bug where parts of large output on a PTY can be lost after the command exits, skipped to avoid test failures.")
415+
416+
// This test is here to prevent prove we have a bug where quickly executing
417+
// commands (with TTY) don't flush their output to the SSH session. This is
418+
// due to the pty being closed before all the output has been copied, but
419+
// protecting against this requires a non-trivial rewrite of the output
420+
// processing (or figuring out a way to put the pty in a mode where this
421+
// does not happen).
422+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
423+
defer cancel()
424+
//nolint:dogsled
425+
conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
426+
sshClient, err := conn.SSHClient(ctx)
427+
require.NoError(t, err)
428+
defer sshClient.Close()
429+
430+
ptty := ptytest.New(t)
431+
432+
var stdout bytes.Buffer
433+
// NOTE(mafredri): Increase iterations to increase chance of failure,
434+
// assuming bug is present.
435+
// Using 10 iterations is basically a guaranteed failure (but let's
436+
// not increase test times needlessly). Run with -race and do not
437+
// limit parallelism (`export GOMAXPROCS=10`) to increase the chance
438+
// of failure.
439+
for i := 0; i < 1; i++ {
440+
func() {
441+
stdout.Reset()
442+
443+
session, err := sshClient.NewSession()
444+
require.NoError(t, err)
445+
defer session.Close()
446+
err = session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
447+
require.NoError(t, err)
448+
449+
session.Stdout = &stdout
450+
session.Stderr = ptty.Output()
451+
session.Stdin = ptty.Input()
452+
want := strings.Repeat("wazzup", 1024+1) // ~6KB, +1 because 1024 is a common buffer size.
453+
err = session.Start("echo " + want)
454+
require.NoError(t, err)
455+
456+
err = session.Wait()
457+
require.NoError(t, err)
458+
require.Contains(t, stdout.String(), want, "should output entire greeting")
459+
}()
460+
}
461+
}
462+
403463
//nolint:paralleltest // This test reserves a port.
404464
func TestAgent_TCPLocalForwarding(t *testing.T) {
405465
random, err := net.Listen("tcp", "127.0.0.1:0")

‎cli/clibase/cmd.go

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -293,11 +293,19 @@ func (i *Invocation) run(state *runState) error {
293293
return
294294
}
295295

296+
// If flag was changed, we need to remove the default values.
296297
sv, ok := f.Value.(pflag.SliceValue)
297298
if !ok {
298299
panic("defaulted array option is not a slice value")
299300
}
300-
err := sv.Replace(sv.GetSlice()[i:])
301+
ss := sv.GetSlice()
302+
if len(ss) == 0 {
303+
// Slice likely zeroed by a flag.
304+
// E.g. "--fruit" may default to "apples,oranges" but the user
305+
// provided "--fruit=""".
306+
return
307+
}
308+
err := sv.Replace(ss[i:])
301309
if err != nil {
302310
panic(err)
303311
}

‎cli/clibase/cmd_test.go

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -503,3 +503,35 @@ func TestCommand_SliceFlags(t *testing.T) {
503503
err = cmd("bad", "bad", "bad").Invoke().Run()
504504
require.NoError(t, err)
505505
}
506+
507+
func TestCommand_EmptySlice(t *testing.T) {
508+
t.Parallel()
509+
510+
cmd := func(want ...string) *clibase.Cmd {
511+
var got []string
512+
return &clibase.Cmd{
513+
Use: "root",
514+
Options: clibase.OptionSet{
515+
{
516+
Name: "arr",
517+
Flag: "arr",
518+
Default: "bad,bad,bad",
519+
Env: "ARR",
520+
Value: clibase.StringArrayOf(&got),
521+
},
522+
},
523+
Handler: (func(i *clibase.Invocation) error {
524+
require.Equal(t, want, got)
525+
return nil
526+
}),
527+
}
528+
}
529+
530+
// Base-case
531+
err := cmd("bad", "bad", "bad").Invoke().Run()
532+
require.NoError(t, err)
533+
534+
inv := cmd().Invoke("--arr", "")
535+
err = inv.Run()
536+
require.NoError(t, err)
537+
}

‎cli/clibase/option.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,9 @@ func (s *OptionSet) ParseEnv(vs []EnvVar) error {
126126
// way for a user to change a Default value to an empty string from
127127
// the environment. Unfortunately, we have old configuration files
128128
// that rely on the faulty behavior.
129+
//
130+
// TODO: We should remove this hack in May 2023, when deployments
131+
// have had months to migrate to the new behavior.
129132
if !ok || envVal == "" {
130133
continue
131134
}

‎cli/clibase/values.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,10 @@ func writeAsCSV(vals []string) string {
146146
}
147147

148148
func (s *StringArray) Set(v string) error {
149+
if v == "" {
150+
*s = nil
151+
return nil
152+
}
149153
ss, err := readAsCSV(v)
150154
if err != nil {
151155
return err

‎cli/configssh.go

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ func (o *sshConfigOptions) addOptions(options ...string) error {
6262
}
6363

6464
func (o *sshConfigOptions) addOption(option string) error {
65-
key, _, err := codersdk.ParseSSHConfigOption(option)
65+
key, value, err := codersdk.ParseSSHConfigOption(option)
6666
if err != nil {
6767
return err
6868
}
@@ -77,11 +77,20 @@ func (o *sshConfigOptions) addOption(option string) error {
7777
continue
7878
}
7979
if strings.EqualFold(existingKey, key) {
80-
o.sshOptions[i] = option
80+
if value == "" {
81+
// Delete existing option.
82+
o.sshOptions = append(o.sshOptions[:i], o.sshOptions[i+1:]...)
83+
} else {
84+
// Override existing option.
85+
o.sshOptions[i] = option
86+
}
8187
return nil
8288
}
8389
}
84-
o.sshOptions = append(o.sshOptions, option)
90+
// Only append the option if it is not empty.
91+
if value != "" {
92+
o.sshOptions = append(o.sshOptions, option)
93+
}
8594
return nil
8695
}
8796

‎cli/configssh_test.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,13 +66,15 @@ func TestConfigSSH(t *testing.T) {
6666

6767
const hostname = "test-coder."
6868
const expectedKey = "ConnectionAttempts"
69+
const removeKey = "ConnectionTimeout"
6970
client := coderdtest.New(t, &coderdtest.Options{
7071
IncludeProvisionerDaemon: true,
7172
ConfigSSH: codersdk.SSHConfigResponse{
7273
HostnamePrefix: hostname,
7374
SSHConfigOptions: map[string]string{
7475
// Something we can test for
7576
expectedKey: "3",
77+
removeKey: "",
7678
},
7779
},
7880
})
@@ -176,6 +178,7 @@ func TestConfigSSH(t *testing.T) {
176178
fileContents, err := os.ReadFile(sshConfigFile)
177179
require.NoError(t, err, "read ssh config file")
178180
require.Contains(t, string(fileContents), expectedKey, "ssh config file contains expected key")
181+
require.NotContains(t, string(fileContents), removeKey, "ssh config file should not have removed key")
179182

180183
home := filepath.Dir(filepath.Dir(sshConfigFile))
181184
// #nosec

‎cli/root.go

Lines changed: 40 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package cli
22

33
import (
4-
"bufio"
54
"context"
65
"errors"
76
"flag"
@@ -14,14 +13,11 @@ import (
1413
"os"
1514
"os/signal"
1615
"path/filepath"
17-
"regexp"
1816
"runtime"
1917
"strings"
2018
"syscall"
2119
"time"
22-
"unicode/utf8"
2320

24-
"golang.org/x/crypto/ssh/terminal"
2521
"golang.org/x/exp/slices"
2622
"golang.org/x/xerrors"
2723

@@ -822,89 +818,61 @@ func isConnectionError(err error) bool {
822818
}
823819

824820
type prettyErrorFormatter struct {
825-
level int
826-
w io.Writer
827-
}
828-
829-
func (prettyErrorFormatter) prefixLines(spaces int, s string) string {
830-
twidth, _, err := terminal.GetSize(0)
831-
if err != nil {
832-
twidth = 80
833-
}
834-
835-
s = lipgloss.NewStyle().Width(twidth - spaces).Render(s)
836-
837-
var b strings.Builder
838-
scanner := bufio.NewScanner(strings.NewReader(s))
839-
for i := 0; scanner.Scan(); i++ {
840-
// The first line is already padded.
841-
if i == 0 {
842-
_, _ = fmt.Fprintf(&b, "%s\n", scanner.Text())
843-
continue
844-
}
845-
_, _ = fmt.Fprintf(&b, "%s%s\n", strings.Repeat(" ", spaces), scanner.Text())
846-
}
847-
return strings.TrimSuffix(strings.TrimSuffix(b.String(), "\n"), " ")
821+
w io.Writer
848822
}
849823

850824
func (p *prettyErrorFormatter) format(err error) {
851-
underErr := errors.Unwrap(err)
852-
853-
arrowStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("#515151"))
825+
errTail := errors.Unwrap(err)
854826

855827
//nolint:errorlint
856-
if _, ok := err.(*clibase.RunCommandError); ok && p.level == 0 && underErr != nil {
857-
// We can do a better job now.
858-
p.format(underErr)
828+
if _, ok := err.(*clibase.RunCommandError); ok && errTail != nil {
829+
// Avoid extra nesting.
830+
p.format(errTail)
859831
return
860832
}
861833

862-
var (
863-
padding string
864-
arrowWidth int
865-
)
866-
if p.level > 0 {
867-
const arrow = "┗━ "
868-
arrowWidth = utf8.RuneCount([]byte(arrow))
869-
padding = strings.Repeat(" ", arrowWidth*p.level)
870-
_, _ = fmt.Fprintf(p.w, "%v%v", padding, arrowStyle.Render(arrow))
871-
}
872-
873-
if underErr != nil {
874-
header := strings.TrimSuffix(err.Error(), ": "+underErr.Error())
875-
_, _ = fmt.Fprintf(p.w, "%s\n", p.prefixLines(len(padding)+arrowWidth, header))
876-
p.level++
877-
p.format(underErr)
878-
return
834+
var headErr string
835+
if errTail != nil {
836+
headErr = strings.TrimSuffix(err.Error(), ": "+errTail.Error())
837+
} else {
838+
headErr = err.Error()
879839
}
880840

881-
{
882-
style := lipgloss.NewStyle().Foreground(lipgloss.Color("#D16644")).Background(lipgloss.Color("#000000")).Bold(false)
883-
// This is the last error in a tree.
884-
p.wrappedPrintf(
885-
"%s\n",
886-
p.prefixLines(
887-
len(padding)+arrowWidth,
888-
fmt.Sprintf(
889-
"%s%s%s",
890-
lipgloss.NewStyle().Inherit(style).Underline(true).Render("ERROR"),
891-
lipgloss.NewStyle().Inherit(style).Foreground(arrowStyle.GetForeground()).Render(" ► "),
892-
style.Render(err.Error()),
893-
),
894-
),
895-
)
841+
var msg string
842+
var sdkError *codersdk.Error
843+
if errors.As(err, &sdkError) {
844+
// We don't want to repeat the same error message twice, so we
845+
// only show the SDK error on the top of the stack.
846+
msg = sdkError.Message
847+
if sdkError.Helper != "" {
848+
msg = msg + "\n" + sdkError.Helper
849+
}
850+
// The SDK error is usually good enough, and we don't want to overwhelm
851+
// the user with output.
852+
errTail = nil
853+
} else {
854+
msg = headErr
896855
}
897-
}
898856

899-
func (p *prettyErrorFormatter) wrappedPrintf(format string, a ...interface{}) {
900-
s := lipgloss.NewStyle().Width(ttyWidth()).Render(
901-
fmt.Sprintf(format, a...),
857+
headStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("#D16644"))
858+
p.printf(
859+
headStyle,
860+
"%s",
861+
msg,
902862
)
903863

904-
// Not sure why, but lipgloss is adding extra spaces we need to remove.
905-
excessSpaceRe := regexp.MustCompile(`[[:blank:]]*\n[[:blank:]]*$`)
906-
s = excessSpaceRe.ReplaceAllString(s, "\n")
864+
tailStyle := headStyle.Copy().Foreground(lipgloss.Color("#969696"))
865+
866+
if errTail != nil {
867+
p.printf(headStyle, ": ")
868+
// Grey out the less important, deep errors.
869+
p.printf(tailStyle, "%s", errTail.Error())
870+
}
871+
p.printf(tailStyle, "\n")
872+
}
907873

874+
func (p *prettyErrorFormatter) printf(style lipgloss.Style, format string, a ...interface{}) {
875+
s := style.Render(fmt.Sprintf(format, a...))
908876
_, _ = p.w.Write(
909877
[]byte(
910878
s,

0 commit comments

Comments
 (0)