Skip to content

Commit c5a265f

Browse files
johnstcnjohnstcn
authored
feat(cli): add experimental rpty command (#16700)
Relates to #16419 Builds upon #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 #16638 .
1 parent 38c0e8a commit c5a265f

File tree

7 files changed

+333
-0
lines changed

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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy