Skip to content

Commit c67db6e

Browse files
spikecurtisspikecurtis
authored
fix: wait for bash prompt before commands (#9882)
Signed-off-by: Spike Curtis <spike@coder.com>
1 parent 399b428 commit c67db6e

File tree

2 files changed

+45
-51
lines changed

2 files changed

+45
-51
lines changed

‎agent/agent_test.go

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1573,26 +1573,21 @@ func TestAgent_ReconnectingPTY(t *testing.T) {
15731573
//nolint:dogsled
15741574
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
15751575
id := uuid.New()
1576-
netConn1, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash")
1576+
// --norc disables executing .bashrc, which is often used to customize the bash prompt
1577+
netConn1, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash --norc")
15771578
require.NoError(t, err)
15781579
defer netConn1.Close()
1580+
tr1 := testutil.NewTerminalReader(t, netConn1)
15791581

15801582
// A second simultaneous connection.
1581-
netConn2, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash")
1583+
netConn2, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash --norc")
15821584
require.NoError(t, err)
15831585
defer netConn2.Close()
1586+
tr2 := testutil.NewTerminalReader(t, netConn2)
15841587

1585-
// Brief pause to reduce the likelihood that we send keystrokes while
1586-
// the shell is simultaneously sending a prompt.
1587-
time.Sleep(100 * time.Millisecond)
1588-
1589-
data, err := json.Marshal(codersdk.ReconnectingPTYRequest{
1590-
Data: "echo test\r\n",
1591-
})
1592-
require.NoError(t, err)
1593-
_, err = netConn1.Write(data)
1594-
require.NoError(t, err)
1595-
1588+
matchPrompt := func(line string) bool {
1589+
return strings.Contains(line, "$ ") || strings.Contains(line, "# ")
1590+
}
15961591
matchEchoCommand := func(line string) bool {
15971592
return strings.Contains(line, "echo test")
15981593
}
@@ -1606,31 +1601,41 @@ func TestAgent_ReconnectingPTY(t *testing.T) {
16061601
return strings.Contains(line, "exit") || strings.Contains(line, "logout")
16071602
}
16081603

1604+
// Wait for the prompt before writing commands. If the command arrives before the prompt is written, screen
1605+
// will sometimes put the command output on the same line as the command and the test will flake
1606+
require.NoError(t, tr1.ReadUntil(ctx, matchPrompt), "find prompt")
1607+
require.NoError(t, tr2.ReadUntil(ctx, matchPrompt), "find prompt")
1608+
1609+
data, err := json.Marshal(codersdk.ReconnectingPTYRequest{
1610+
Data: "echo test\r",
1611+
})
1612+
require.NoError(t, err)
1613+
_, err = netConn1.Write(data)
1614+
require.NoError(t, err)
1615+
16091616
// Once for typing the command...
1610-
tr1 := testutil.NewTerminalReader(t, netConn1)
16111617
require.NoError(t, tr1.ReadUntil(ctx, matchEchoCommand), "find echo command")
16121618
// And another time for the actual output.
16131619
require.NoError(t, tr1.ReadUntil(ctx, matchEchoOutput), "find echo output")
16141620

16151621
// Same for the other connection.
1616-
tr2 := testutil.NewTerminalReader(t, netConn2)
16171622
require.NoError(t, tr2.ReadUntil(ctx, matchEchoCommand), "find echo command")
16181623
require.NoError(t, tr2.ReadUntil(ctx, matchEchoOutput), "find echo output")
16191624

16201625
_ = netConn1.Close()
16211626
_ = netConn2.Close()
1622-
netConn3, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash")
1627+
netConn3, err := conn.ReconnectingPTY(ctx, id, 80, 80, "bash --norc")
16231628
require.NoError(t, err)
16241629
defer netConn3.Close()
1630+
tr3 := testutil.NewTerminalReader(t, netConn3)
16251631

16261632
// Same output again!
1627-
tr3 := testutil.NewTerminalReader(t, netConn3)
16281633
require.NoError(t, tr3.ReadUntil(ctx, matchEchoCommand), "find echo command")
16291634
require.NoError(t, tr3.ReadUntil(ctx, matchEchoOutput), "find echo output")
16301635

16311636
// Exit should cause the connection to close.
16321637
data, err = json.Marshal(codersdk.ReconnectingPTYRequest{
1633-
Data: "exit\r\n",
1638+
Data: "exit\r",
16341639
})
16351640
require.NoError(t, err)
16361641
_, err = netConn3.Write(data)

‎coderd/workspaceapps/apptest/apptest.go

Lines changed: 22 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -62,13 +62,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
6262
// Run the test against the path app hostname since that's where the
6363
// reconnecting-pty proxy server we want to test is mounted.
6464
client := appDetails.AppClient(t)
65-
testReconnectingPTY(ctx, t, client, codersdk.WorkspaceAgentReconnectingPTYOpts{
66-
AgentID: appDetails.Agent.ID,
67-
Reconnect: uuid.New(),
68-
Height: 100,
69-
Width: 100,
70-
Command: "bash",
71-
})
65+
testReconnectingPTY(ctx, t, client, appDetails.Agent.ID, "")
7266
})
7367

7468
t.Run("SignedTokenQueryParameter", func(t *testing.T) {
@@ -96,14 +90,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
9690

9791
// Make an unauthenticated client.
9892
unauthedAppClient := codersdk.New(appDetails.AppClient(t).URL)
99-
testReconnectingPTY(ctx, t, unauthedAppClient, codersdk.WorkspaceAgentReconnectingPTYOpts{
100-
AgentID: appDetails.Agent.ID,
101-
Reconnect: uuid.New(),
102-
Height: 100,
103-
Width: 100,
104-
Command: "bash",
105-
SignedToken: issueRes.SignedToken,
106-
})
93+
testReconnectingPTY(ctx, t, unauthedAppClient, appDetails.Agent.ID, issueRes.SignedToken)
10794
})
10895
})
10996

@@ -1389,7 +1376,19 @@ func (r *fakeStatsReporter) Report(_ context.Context, stats []workspaceapps.Stat
13891376
return nil
13901377
}
13911378

1392-
func testReconnectingPTY(ctx context.Context, t *testing.T, client *codersdk.Client, opts codersdk.WorkspaceAgentReconnectingPTYOpts) {
1379+
func testReconnectingPTY(ctx context.Context, t *testing.T, client *codersdk.Client, agentID uuid.UUID, signedToken string) {
1380+
opts := codersdk.WorkspaceAgentReconnectingPTYOpts{
1381+
AgentID: agentID,
1382+
Reconnect: uuid.New(),
1383+
Width: 80,
1384+
Height: 80,
1385+
// --norc disables executing .bashrc, which is often used to customize the bash prompt
1386+
Command: "bash --norc",
1387+
SignedToken: signedToken,
1388+
}
1389+
matchPrompt := func(line string) bool {
1390+
return strings.Contains(line, "$ ") || strings.Contains(line, "# ")
1391+
}
13931392
matchEchoCommand := func(line string) bool {
13941393
return strings.Contains(line, "echo test")
13951394
}
@@ -1407,34 +1406,24 @@ func testReconnectingPTY(ctx context.Context, t *testing.T, client *codersdk.Cli
14071406
require.NoError(t, err)
14081407
defer conn.Close()
14091408

1410-
// First attempt to resize the TTY.
1411-
// The websocket will close if it fails!
1412-
data, err := json.Marshal(codersdk.ReconnectingPTYRequest{
1413-
Height: 80,
1414-
Width: 80,
1415-
})
1416-
require.NoError(t, err)
1417-
_, err = conn.Write(data)
1418-
require.NoError(t, err)
1419-
1420-
// Brief pause to reduce the likelihood that we send keystrokes while
1421-
// the shell is simultaneously sending a prompt.
1422-
time.Sleep(500 * time.Millisecond)
1409+
tr := testutil.NewTerminalReader(t, conn)
1410+
// Wait for the prompt before writing commands. If the command arrives before the prompt is written, screen
1411+
// will sometimes put the command output on the same line as the command and the test will flake
1412+
require.NoError(t, tr.ReadUntil(ctx, matchPrompt), "find prompt")
14231413

1424-
data, err = json.Marshal(codersdk.ReconnectingPTYRequest{
1425-
Data: "echo test\r\n",
1414+
data, err := json.Marshal(codersdk.ReconnectingPTYRequest{
1415+
Data: "echo test\r",
14261416
})
14271417
require.NoError(t, err)
14281418
_, err = conn.Write(data)
14291419
require.NoError(t, err)
14301420

1431-
tr := testutil.NewTerminalReader(t, conn)
14321421
require.NoError(t, tr.ReadUntil(ctx, matchEchoCommand), "find echo command")
14331422
require.NoError(t, tr.ReadUntil(ctx, matchEchoOutput), "find echo output")
14341423

14351424
// Exit should cause the connection to close.
14361425
data, err = json.Marshal(codersdk.ReconnectingPTYRequest{
1437-
Data: "exit\r\n",
1426+
Data: "exit\r",
14381427
})
14391428
require.NoError(t, err)
14401429
_, err = conn.Write(data)

0 commit comments

Comments
 (0)