Skip to content

Commit f1dfeb0

Browse files
deansheatherdeansheather
authored
chore: fix flake in apptest reconnecting-pty test (#7281)
1 parent 35b3ed2 commit f1dfeb0

File tree

3 files changed

+88
-41
lines changed

3 files changed

+88
-41
lines changed

‎coderd/workspaceapps/apptest/apptest.go

Lines changed: 41 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@ import (
2222
"github.com/stretchr/testify/assert"
2323
"github.com/stretchr/testify/require"
2424
"golang.org/x/xerrors"
25-
"nhooyr.io/websocket"
2625

2726
"github.com/coder/coder/coderd/coderdtest"
2827
"github.com/coder/coder/coderd/rbac"
@@ -72,7 +71,13 @@ func Run(t *testing.T, factory DeploymentFactory) {
7271
// Run the test against the path app hostname since that's where the
7372
// reconnecting-pty proxy server we want to test is mounted.
7473
client := appDetails.AppClient(t)
75-
conn, err := client.WorkspaceAgentReconnectingPTY(ctx, appDetails.Agent.ID, uuid.New(), 80, 80, "/bin/bash")
74+
conn, err := client.WorkspaceAgentReconnectingPTY(ctx, codersdk.WorkspaceAgentReconnectingPTYOpts{
75+
AgentID: appDetails.Agent.ID,
76+
Reconnect: uuid.New(),
77+
Height: 80,
78+
Width: 80,
79+
Command: "/bin/bash",
80+
})
7681
require.NoError(t, err)
7782
defer conn.Close()
7883

@@ -125,29 +130,42 @@ func Run(t *testing.T, factory DeploymentFactory) {
125130
})
126131
require.NoError(t, err)
127132

128-
// Try to connect to the endpoint with the signed token and no other
129-
// authentication.
130-
q := u.Query()
131-
q.Set("reconnect", uuid.NewString())
132-
q.Set("height", strconv.Itoa(24))
133-
q.Set("width", strconv.Itoa(80))
134-
q.Set("command", `/bin/sh -c "echo test"`)
135-
q.Set(codersdk.SignedAppTokenQueryParameter, issueRes.SignedToken)
136-
u.RawQuery = q.Encode()
137-
138-
//nolint:bodyclose
139-
wsConn, res, err := websocket.Dial(ctx, u.String(), nil)
140-
if !assert.NoError(t, err) {
141-
dump, err := httputil.DumpResponse(res, true)
142-
if err == nil {
143-
t.Log(string(dump))
144-
}
145-
return
146-
}
147-
defer wsConn.Close(websocket.StatusNormalClosure, "")
148-
conn := websocket.NetConn(ctx, wsConn, websocket.MessageBinary)
133+
// Make an unauthenticated client.
134+
unauthedAppClient := codersdk.New(appDetails.AppClient(t).URL)
135+
conn, err := unauthedAppClient.WorkspaceAgentReconnectingPTY(ctx, codersdk.WorkspaceAgentReconnectingPTYOpts{
136+
AgentID: appDetails.Agent.ID,
137+
Reconnect: uuid.New(),
138+
Height: 80,
139+
Width: 80,
140+
Command: "/bin/bash",
141+
SignedToken: issueRes.SignedToken,
142+
})
143+
require.NoError(t, err)
144+
defer conn.Close()
145+
146+
// First attempt to resize the TTY.
147+
// The websocket will close if it fails!
148+
data, err := json.Marshal(codersdk.ReconnectingPTYRequest{
149+
Height: 250,
150+
Width: 250,
151+
})
152+
require.NoError(t, err)
153+
_, err = conn.Write(data)
154+
require.NoError(t, err)
149155
bufRead := bufio.NewReader(conn)
150156

157+
// Brief pause to reduce the likelihood that we send keystrokes while
158+
// the shell is simultaneously sending a prompt.
159+
time.Sleep(100 * time.Millisecond)
160+
161+
data, err = json.Marshal(codersdk.ReconnectingPTYRequest{
162+
Data: "echo test\r\n",
163+
})
164+
require.NoError(t, err)
165+
_, err = conn.Write(data)
166+
require.NoError(t, err)
167+
168+
expectLine(t, bufRead, matchEchoCommand)
151169
expectLine(t, bufRead, matchEchoOutput)
152170
})
153171
})

‎codersdk/workspaceagents.go

Lines changed: 40 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -385,32 +385,55 @@ func (c *Client) IssueReconnectingPTYSignedToken(ctx context.Context, req IssueR
385385
return resp, json.NewDecoder(res.Body).Decode(&resp)
386386
}
387387

388+
// @typescript-ignore:WorkspaceAgentReconnectingPTYOpts
389+
type WorkspaceAgentReconnectingPTYOpts struct {
390+
AgentID uuid.UUID
391+
Reconnect uuid.UUID
392+
Width uint16
393+
Height uint16
394+
Command string
395+
396+
// SignedToken is an optional signed token from the
397+
// issue-reconnecting-pty-signed-token endpoint. If set, the session token
398+
// on the client will not be sent.
399+
SignedToken string
400+
}
401+
388402
// WorkspaceAgentReconnectingPTY spawns a PTY that reconnects using the token provided.
389403
// It communicates using `agent.ReconnectingPTYRequest` marshaled as JSON.
390404
// Responses are PTY output that can be rendered.
391-
func (c *Client) WorkspaceAgentReconnectingPTY(ctx context.Context, agentID, reconnect uuid.UUID, height, width uint16, command string) (net.Conn, error) {
392-
serverURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/pty", agentID))
405+
func (c *Client) WorkspaceAgentReconnectingPTY(ctx context.Context, opts WorkspaceAgentReconnectingPTYOpts) (net.Conn, error) {
406+
serverURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/pty", opts.AgentID))
393407
if err != nil {
394408
return nil, xerrors.Errorf("parse url: %w", err)
395409
}
396410
q := serverURL.Query()
397-
q.Set("reconnect", reconnect.String())
398-
q.Set("height", strconv.Itoa(int(height)))
399-
q.Set("width", strconv.Itoa(int(width)))
400-
q.Set("command", command)
411+
q.Set("reconnect", opts.Reconnect.String())
412+
q.Set("width", strconv.Itoa(int(opts.Width)))
413+
q.Set("height", strconv.Itoa(int(opts.Height)))
414+
q.Set("command", opts.Command)
415+
// If we're using a signed token, set the query parameter.
416+
if opts.SignedToken != "" {
417+
q.Set(SignedAppTokenQueryParameter, opts.SignedToken)
418+
}
401419
serverURL.RawQuery = q.Encode()
402420

403-
jar, err := cookiejar.New(nil)
404-
if err != nil {
405-
return nil, xerrors.Errorf("create cookie jar: %w", err)
406-
}
407-
jar.SetCookies(serverURL, []*http.Cookie{{
408-
Name: SessionTokenCookie,
409-
Value: c.SessionToken(),
410-
}})
411-
httpClient := &http.Client{
412-
Jar: jar,
413-
Transport: c.HTTPClient.Transport,
421+
// If we're not using a signed token, we need to set the session token as a
422+
// cookie.
423+
httpClient := c.HTTPClient
424+
if opts.SignedToken == "" {
425+
jar, err := cookiejar.New(nil)
426+
if err != nil {
427+
return nil, xerrors.Errorf("create cookie jar: %w", err)
428+
}
429+
jar.SetCookies(serverURL, []*http.Cookie{{
430+
Name: SessionTokenCookie,
431+
Value: c.SessionToken(),
432+
}})
433+
httpClient = &http.Client{
434+
Jar: jar,
435+
Transport: c.HTTPClient.Transport,
436+
}
414437
}
415438
conn, res, err := websocket.Dial(ctx, serverURL.String(), &websocket.DialOptions{
416439
HTTPClient: httpClient,

‎scaletest/reconnectingpty/run.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,13 @@ func (r *Runner) Run(ctx context.Context, _ string, logs io.Writer) error {
6464
_, _ = fmt.Fprintf(logs, "\tHeight: %d\n", height)
6565
_, _ = fmt.Fprintf(logs, "\tCommand: %q\n\n", r.cfg.Init.Command)
6666

67-
conn, err := r.client.WorkspaceAgentReconnectingPTY(ctx, r.cfg.AgentID, id, width, height, r.cfg.Init.Command)
67+
conn, err := r.client.WorkspaceAgentReconnectingPTY(ctx, codersdk.WorkspaceAgentReconnectingPTYOpts{
68+
AgentID: r.cfg.AgentID,
69+
Reconnect: id,
70+
Width: width,
71+
Height: height,
72+
Command: r.cfg.Init.Command,
73+
})
6874
if err != nil {
6975
return xerrors.Errorf("open reconnecting PTY: %w", err)
7076
}

0 commit comments

Comments
 (0)