Skip to content

Commit c5a265f

Browse files
johnstcnjohnstcn
authored
feat(cli): add experimental rpty command (coder#16700)
Relates to coder#16419 Builds upon coder#16638 and adds a command `exp rpty` that allows you to open a ReconnectingPTY session to an agent. This ultimately allows us to add an integration-style CLI test to verify the functionality added in coder#16638 .
1 parent 38c0e8a commit c5a265f

7 files changed

+333
-0
lines changed

‎cli/dotfiles_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@ import (
1717

1818
func TestDotfiles(t *testing.T) {
1919
t.Parallel()
20+
// This test will time out if the user has commit signing enabled.
21+
if _, gpgTTYFound := os.LookupEnv("GPG_TTY"); gpgTTYFound {
22+
t.Skip("GPG_TTY is set, skipping test to avoid hanging")
23+
}
2024
t.Run("MissingArg", func(t *testing.T) {
2125
t.Parallel()
2226
inv, _ := clitest.New(t, "dotfiles")

‎cli/exp.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ func (r *RootCmd) expCmd() *serpent.Command {
1414
r.scaletestCmd(),
1515
r.errorExample(),
1616
r.promptExample(),
17+
r.rptyCommand(),
1718
},
1819
}
1920
return cmd

‎cli/errors.go renamed to ‎cli/exp_errors.go

File renamed without changes.

‎cli/errors_test.go renamed to ‎cli/exp_errors_test.go

File renamed without changes.

‎cli/prompts.go renamed to ‎cli/exp_prompts.go

File renamed without changes.

‎cli/exp_rpty.go

Lines changed: 216 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
1+
package cli
2+
3+
import (
4+
"bufio"
5+
"context"
6+
"encoding/json"
7+
"fmt"
8+
"io"
9+
"os"
10+
"strings"
11+
12+
"github.com/google/uuid"
13+
"github.com/mattn/go-isatty"
14+
"golang.org/x/term"
15+
"golang.org/x/xerrors"
16+
17+
"github.com/coder/coder/v2/cli/cliui"
18+
"github.com/coder/coder/v2/codersdk"
19+
"github.com/coder/coder/v2/codersdk/workspacesdk"
20+
"github.com/coder/coder/v2/pty"
21+
"github.com/coder/serpent"
22+
)
23+
24+
func (r *RootCmd) rptyCommand() *serpent.Command {
25+
var (
26+
client = new(codersdk.Client)
27+
args handleRPTYArgs
28+
)
29+
30+
cmd := &serpent.Command{
31+
Handler: func(inv *serpent.Invocation) error {
32+
if r.disableDirect {
33+
return xerrors.New("direct connections are disabled, but you can try websocat ;-)")
34+
}
35+
args.NamedWorkspace = inv.Args[0]
36+
args.Command = inv.Args[1:]
37+
return handleRPTY(inv, client, args)
38+
},
39+
Long: "Establish an RPTY session with a workspace/agent. This uses the same mechanism as the Web Terminal.",
40+
Middleware: serpent.Chain(
41+
serpent.RequireRangeArgs(1, -1),
42+
r.InitClient(client),
43+
),
44+
Options: []serpent.Option{
45+
{
46+
Name: "container",
47+
Description: "The container name or ID to connect to.",
48+
Flag: "container",
49+
FlagShorthand: "c",
50+
Default: "",
51+
Value: serpent.StringOf(&args.Container),
52+
},
53+
{
54+
Name: "container-user",
55+
Description: "The user to connect as.",
56+
Flag: "container-user",
57+
FlagShorthand: "u",
58+
Default: "",
59+
Value: serpent.StringOf(&args.ContainerUser),
60+
},
61+
{
62+
Name: "reconnect",
63+
Description: "The reconnect ID to use.",
64+
Flag: "reconnect",
65+
FlagShorthand: "r",
66+
Default: "",
67+
Value: serpent.StringOf(&args.ReconnectID),
68+
},
69+
},
70+
Short: "Establish an RPTY session with a workspace/agent.",
71+
Use: "rpty",
72+
}
73+
74+
return cmd
75+
}
76+
77+
type handleRPTYArgs struct {
78+
Command []string
79+
Container string
80+
ContainerUser string
81+
NamedWorkspace string
82+
ReconnectID string
83+
}
84+
85+
func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPTYArgs) error {
86+
ctx, cancel := context.WithCancel(inv.Context())
87+
defer cancel()
88+
89+
var reconnectID uuid.UUID
90+
if args.ReconnectID != "" {
91+
rid, err := uuid.Parse(args.ReconnectID)
92+
if err != nil {
93+
return xerrors.Errorf("invalid reconnect ID: %w", err)
94+
}
95+
reconnectID = rid
96+
} else {
97+
reconnectID = uuid.New()
98+
}
99+
ws, agt, err := getWorkspaceAndAgent(ctx, inv, client, true, args.NamedWorkspace)
100+
if err != nil {
101+
return err
102+
}
103+
104+
var ctID string
105+
if args.Container != "" {
106+
cts, err := client.WorkspaceAgentListContainers(ctx, agt.ID, nil)
107+
if err != nil {
108+
return err
109+
}
110+
for _, ct := range cts.Containers {
111+
if ct.FriendlyName == args.Container || ct.ID == args.Container {
112+
ctID = ct.ID
113+
break
114+
}
115+
}
116+
if ctID == "" {
117+
return xerrors.Errorf("container %q not found", args.Container)
118+
}
119+
}
120+
121+
if err := cliui.Agent(ctx, inv.Stderr, agt.ID, cliui.AgentOptions{
122+
FetchInterval: 0,
123+
Fetch: client.WorkspaceAgent,
124+
Wait: false,
125+
}); err != nil {
126+
return err
127+
}
128+
129+
// Get the width and height of the terminal.
130+
var termWidth, termHeight uint16
131+
stdoutFile, validOut := inv.Stdout.(*os.File)
132+
if validOut && isatty.IsTerminal(stdoutFile.Fd()) {
133+
w, h, err := term.GetSize(int(stdoutFile.Fd()))
134+
if err == nil {
135+
//nolint: gosec
136+
termWidth, termHeight = uint16(w), uint16(h)
137+
}
138+
}
139+
140+
// Set stdin to raw mode so that control characters work.
141+
stdinFile, validIn := inv.Stdin.(*os.File)
142+
if validIn && isatty.IsTerminal(stdinFile.Fd()) {
143+
inState, err := pty.MakeInputRaw(stdinFile.Fd())
144+
if err != nil {
145+
return xerrors.Errorf("failed to set input terminal to raw mode: %w", err)
146+
}
147+
defer func() {
148+
_ = pty.RestoreTerminal(stdinFile.Fd(), inState)
149+
}()
150+
}
151+
152+
conn, err := workspacesdk.New(client).AgentReconnectingPTY(ctx, workspacesdk.WorkspaceAgentReconnectingPTYOpts{
153+
AgentID: agt.ID,
154+
Reconnect: reconnectID,
155+
Command: strings.Join(args.Command, " "),
156+
Container: ctID,
157+
ContainerUser: args.ContainerUser,
158+
Width: termWidth,
159+
Height: termHeight,
160+
})
161+
if err != nil {
162+
return xerrors.Errorf("open reconnecting PTY: %w", err)
163+
}
164+
defer conn.Close()
165+
166+
cliui.Infof(inv.Stderr, "Connected to %s (agent id: %s)", args.NamedWorkspace, agt.ID)
167+
cliui.Infof(inv.Stderr, "Reconnect ID: %s", reconnectID)
168+
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, ws.ID, codersdk.PostWorkspaceUsageRequest{
169+
AgentID: agt.ID,
170+
AppName: codersdk.UsageAppNameReconnectingPty,
171+
})
172+
defer closeUsage()
173+
174+
br := bufio.NewScanner(inv.Stdin)
175+
// Split on bytes, otherwise you have to send a newline to flush the buffer.
176+
br.Split(bufio.ScanBytes)
177+
je := json.NewEncoder(conn)
178+
179+
go func() {
180+
for br.Scan() {
181+
if err := je.Encode(map[string]string{
182+
"data": br.Text(),
183+
}); err != nil {
184+
return
185+
}
186+
}
187+
}()
188+
189+
windowChange := listenWindowSize(ctx)
190+
go func() {
191+
for {
192+
select {
193+
case <-ctx.Done():
194+
return
195+
case <-windowChange:
196+
}
197+
width, height, err := term.GetSize(int(stdoutFile.Fd()))
198+
if err != nil {
199+
continue
200+
}
201+
if err := je.Encode(map[string]int{
202+
"width": width,
203+
"height": height,
204+
}); err != nil {
205+
cliui.Errorf(inv.Stderr, "Failed to send window size: %v", err)
206+
}
207+
}
208+
}()
209+
210+
_, _ = io.Copy(inv.Stdout, conn)
211+
cancel()
212+
_ = conn.Close()
213+
_, _ = fmt.Fprintf(inv.Stderr, "Connection closed\n")
214+
215+
return nil
216+
}

‎cli/exp_rpty_test.go

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
package cli_test
2+
3+
import (
4+
"fmt"
5+
"runtime"
6+
"testing"
7+
8+
"github.com/ory/dockertest/v3"
9+
"github.com/ory/dockertest/v3/docker"
10+
11+
"github.com/coder/coder/v2/agent"
12+
"github.com/coder/coder/v2/agent/agenttest"
13+
"github.com/coder/coder/v2/cli/clitest"
14+
"github.com/coder/coder/v2/coderd/coderdtest"
15+
"github.com/coder/coder/v2/pty/ptytest"
16+
"github.com/coder/coder/v2/testutil"
17+
18+
"github.com/stretchr/testify/assert"
19+
"github.com/stretchr/testify/require"
20+
)
21+
22+
func TestExpRpty(t *testing.T) {
23+
t.Parallel()
24+
25+
t.Run("OK", func(t *testing.T) {
26+
t.Parallel()
27+
28+
client, workspace, agentToken := setupWorkspaceForAgent(t)
29+
inv, root := clitest.New(t, "exp", "rpty", workspace.Name)
30+
clitest.SetupConfig(t, client, root)
31+
pty := ptytest.New(t).Attach(inv)
32+
33+
ctx := testutil.Context(t, testutil.WaitLong)
34+
35+
cmdDone := tGo(t, func() {
36+
err := inv.WithContext(ctx).Run()
37+
assert.NoError(t, err)
38+
})
39+
40+
_ = agenttest.New(t, client.URL, agentToken)
41+
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
42+
43+
pty.ExpectMatch(fmt.Sprintf("Connected to %s", workspace.Name))
44+
pty.WriteLine("exit")
45+
<-cmdDone
46+
})
47+
48+
t.Run("NotFound", func(t *testing.T) {
49+
t.Parallel()
50+
51+
client, _, _ := setupWorkspaceForAgent(t)
52+
inv, root := clitest.New(t, "exp", "rpty", "not-found")
53+
clitest.SetupConfig(t, client, root)
54+
55+
ctx := testutil.Context(t, testutil.WaitShort)
56+
err := inv.WithContext(ctx).Run()
57+
require.ErrorContains(t, err, "not found")
58+
})
59+
60+
t.Run("Container", func(t *testing.T) {
61+
t.Parallel()
62+
// Skip this test on non-Linux platforms since it requires Docker
63+
if runtime.GOOS != "linux" {
64+
t.Skip("Skipping test on non-Linux platform")
65+
}
66+
67+
client, workspace, agentToken := setupWorkspaceForAgent(t)
68+
ctx := testutil.Context(t, testutil.WaitLong)
69+
pool, err := dockertest.NewPool("")
70+
require.NoError(t, err, "Could not connect to docker")
71+
ct, err := pool.RunWithOptions(&dockertest.RunOptions{
72+
Repository: "busybox",
73+
Tag: "latest",
74+
Cmd: []string{"sleep", "infnity"},
75+
}, func(config *docker.HostConfig) {
76+
config.AutoRemove = true
77+
config.RestartPolicy = docker.RestartPolicy{Name: "no"}
78+
})
79+
require.NoError(t, err, "Could not start container")
80+
// Wait for container to start
81+
require.Eventually(t, func() bool {
82+
ct, ok := pool.ContainerByName(ct.Container.Name)
83+
return ok && ct.Container.State.Running
84+
}, testutil.WaitShort, testutil.IntervalSlow, "Container did not start in time")
85+
t.Cleanup(func() {
86+
err := pool.Purge(ct)
87+
require.NoError(t, err, "Could not stop container")
88+
})
89+
90+
inv, root := clitest.New(t, "exp", "rpty", workspace.Name, "-c", ct.Container.ID)
91+
clitest.SetupConfig(t, client, root)
92+
pty := ptytest.New(t).Attach(inv)
93+
94+
cmdDone := tGo(t, func() {
95+
err := inv.WithContext(ctx).Run()
96+
assert.NoError(t, err)
97+
})
98+
99+
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
100+
o.ExperimentalContainersEnabled = true
101+
})
102+
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
103+
104+
pty.ExpectMatch(fmt.Sprintf("Connected to %s", workspace.Name))
105+
pty.ExpectMatch("Reconnect ID: ")
106+
pty.ExpectMatch(" #")
107+
pty.WriteLine("hostname")
108+
pty.ExpectMatch(ct.Container.Config.Hostname)
109+
pty.WriteLine("exit")
110+
<-cmdDone
111+
})
112+
}

0 commit comments

Comments
 (0)