Skip to content

Commit a9ff045

Browse files
johnstcnjohnstcn
committed
feat(cli): add experimental rpty command
1 parent 02463f3 commit a9ff045

File tree

3 files changed

+329
-0
lines changed

3 files changed

+329
-0
lines changed

‎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/exp_rpty.go

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
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+
FetchLogs: client.WorkspaceAgentLogsAfter,
125+
Wait: false,
126+
}); err != nil {
127+
return err
128+
}
129+
130+
// Get the width and height of the terminal.
131+
var termWidth, termHeight uint16
132+
stdoutFile, validOut := inv.Stdout.(*os.File)
133+
if validOut && isatty.IsTerminal(stdoutFile.Fd()) {
134+
w, h, err := term.GetSize(int(stdoutFile.Fd()))
135+
if err == nil {
136+
//nolint: gosec
137+
termWidth, termHeight = uint16(w), uint16(h)
138+
}
139+
}
140+
141+
// Set stdin to raw mode so that control characters work.
142+
stdinFile, validIn := inv.Stdin.(*os.File)
143+
if validIn && isatty.IsTerminal(stdinFile.Fd()) {
144+
inState, err := pty.MakeInputRaw(stdinFile.Fd())
145+
if err != nil {
146+
return xerrors.Errorf("failed to set input terminal to raw mode: %w", err)
147+
}
148+
defer func() {
149+
_ = pty.RestoreTerminal(stdinFile.Fd(), inState)
150+
}()
151+
}
152+
153+
conn, err := workspacesdk.New(client).AgentReconnectingPTY(ctx, workspacesdk.WorkspaceAgentReconnectingPTYOpts{
154+
AgentID: agt.ID,
155+
Reconnect: reconnectID,
156+
Command: strings.Join(args.Command, " "),
157+
Container: ctID,
158+
ContainerUser: args.ContainerUser,
159+
Width: termWidth,
160+
Height: termHeight,
161+
})
162+
if err != nil {
163+
return xerrors.Errorf("open reconnecting PTY: %w", err)
164+
}
165+
defer conn.Close()
166+
167+
cliui.Infof(inv.Stderr, "Connected to %s (agent id: %s)", args.NamedWorkspace, agt.ID)
168+
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, ws.ID, codersdk.PostWorkspaceUsageRequest{
169+
AgentID: agt.ID,
170+
AppName: codersdk.UsageAppNameReconnectingPty,
171+
})
172+
defer closeUsage()
173+
174+
stdinDone := make(chan struct{})
175+
stdoutDone := make(chan struct{})
176+
stderrDone := make(chan struct{})
177+
done := make(chan struct{})
178+
179+
go func() {
180+
defer close(stdinDone)
181+
// This is how we send commands to the agent.
182+
br := bufio.NewScanner(inv.Stdin)
183+
// Split on bytes, otherwise you have to send a newline to flush the buffer.
184+
br.Split(bufio.ScanBytes)
185+
je := json.NewEncoder(conn)
186+
for br.Scan() {
187+
if err := je.Encode(map[string]string{
188+
"data": br.Text(),
189+
}); err != nil {
190+
return
191+
}
192+
}
193+
}()
194+
go func() {
195+
defer func() {
196+
close(stdoutDone)
197+
}()
198+
_, _ = io.Copy(inv.Stdout, conn)
199+
}()
200+
go func() {
201+
defer func() {
202+
close(stderrDone)
203+
}()
204+
_, _ = io.Copy(inv.Stderr, conn)
205+
}()
206+
go func() {
207+
defer close(done)
208+
<-stdoutDone
209+
<-stderrDone
210+
_ = conn.Close()
211+
_, _ = fmt.Fprintf(inv.Stderr, "Connection closed\n")
212+
}()
213+
214+
<-done
215+
216+
return nil
217+
}

‎cli/exp_rpty_test.go

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
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(" #")
106+
pty.WriteLine("hostname")
107+
pty.ExpectMatch(ct.Container.Config.Hostname)
108+
pty.WriteLine("exit")
109+
<-cmdDone
110+
})
111+
}

0 commit comments

Comments
 (0)