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)
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