Skip to content

Commit fc3e193

Browse files
johnstcnjohnstcn
committed
feat(cli): add experimental rpty command
1 parent 097fd85 commit fc3e193

File tree

2 files changed

+211
-0
lines changed

2 files changed

+211
-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: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
package cli
2+
3+
import (
4+
"bufio"
5+
"context"
6+
"encoding/json"
7+
"io"
8+
"os"
9+
"strings"
10+
11+
"github.com/google/uuid"
12+
"github.com/mattn/go-isatty"
13+
"golang.org/x/term"
14+
"golang.org/x/xerrors"
15+
16+
"github.com/coder/coder/v2/cli/cliui"
17+
"github.com/coder/coder/v2/codersdk"
18+
"github.com/coder/coder/v2/codersdk/workspacesdk"
19+
"github.com/coder/coder/v2/pty"
20+
"github.com/coder/serpent"
21+
)
22+
23+
func (r *RootCmd) rptyCommand() *serpent.Command {
24+
var (
25+
client = new(codersdk.Client)
26+
args handleRPTYArgs
27+
)
28+
29+
cmd := &serpent.Command{
30+
Handler: func(inv *serpent.Invocation) error {
31+
if r.disableDirect {
32+
return xerrors.New("direct connections are disabled, but you can try websocat ;-)")
33+
}
34+
args.NamedWorkspace = inv.Args[0]
35+
args.Command = inv.Args[1:]
36+
return handleRPTY(inv, client, args)
37+
},
38+
Long: "Establish an RPTY session with a workspace/agent. This uses the same mechanism as the Web Terminal.",
39+
Middleware: serpent.Chain(
40+
serpent.RequireRangeArgs(1, -1),
41+
r.InitClient(client),
42+
),
43+
Options: []serpent.Option{
44+
{
45+
Name: "container",
46+
Description: "The container name or ID to connect to.",
47+
Flag: "container",
48+
FlagShorthand: "c",
49+
Default: "",
50+
Value: serpent.StringOf(&args.Container),
51+
},
52+
{
53+
Name: "container-user",
54+
Description: "The user to connect as.",
55+
Flag: "container-user",
56+
FlagShorthand: "u",
57+
Default: "",
58+
Value: serpent.StringOf(&args.ContainerUser),
59+
},
60+
{
61+
Name: "reconnect",
62+
Description: "The reconnect ID to use.",
63+
Flag: "reconnect",
64+
FlagShorthand: "r",
65+
Default: "",
66+
Value: serpent.StringOf(&args.ReconnectID),
67+
},
68+
},
69+
Short: "Establish an RPTY session with a workspace/agent.",
70+
Use: "rpty",
71+
}
72+
73+
return cmd
74+
}
75+
76+
type handleRPTYArgs struct {
77+
Command []string
78+
Container string
79+
ContainerUser string
80+
NamedWorkspace string
81+
ReconnectID string
82+
}
83+
84+
func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPTYArgs) error {
85+
ctx, cancel := context.WithCancel(inv.Context())
86+
defer cancel()
87+
88+
var reconnectID uuid.UUID
89+
if args.ReconnectID != "" {
90+
rid, err := uuid.Parse(args.ReconnectID)
91+
if err != nil {
92+
return xerrors.Errorf("invalid reconnect ID: %w", err)
93+
}
94+
reconnectID = rid
95+
} else {
96+
reconnectID = uuid.New()
97+
}
98+
ws, agt, err := getWorkspaceAndAgent(ctx, inv, client, true, args.NamedWorkspace)
99+
if err != nil {
100+
return err
101+
}
102+
103+
var ctID string
104+
if args.Container != "" {
105+
cts, err := client.WorkspaceAgentListContainers(ctx, agt.ID, nil)
106+
if err != nil {
107+
return err
108+
}
109+
for _, ct := range cts.Containers {
110+
if ct.FriendlyName == args.Container || ct.ID == args.Container {
111+
ctID = ct.ID
112+
break
113+
}
114+
}
115+
if ctID == "" {
116+
return xerrors.Errorf("container %q not found", args.Container)
117+
}
118+
}
119+
120+
if err := cliui.Agent(ctx, inv.Stderr, agt.ID, cliui.AgentOptions{
121+
FetchInterval: 0,
122+
Fetch: client.WorkspaceAgent,
123+
FetchLogs: client.WorkspaceAgentLogsAfter,
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+
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, ws.ID, codersdk.PostWorkspaceUsageRequest{
168+
AgentID: agt.ID,
169+
AppName: codersdk.UsageAppNameReconnectingPty,
170+
})
171+
defer closeUsage()
172+
173+
stdinDone := make(chan struct{})
174+
stdoutDone := make(chan struct{})
175+
stderrDone := make(chan struct{})
176+
done := make(chan struct{})
177+
178+
go func() {
179+
defer close(stdinDone)
180+
// This is how we send commands to the agent.
181+
br := bufio.NewScanner(inv.Stdin)
182+
// Split on bytes, otherwise you have to send a newline to flush the buffer.
183+
br.Split(bufio.ScanBytes)
184+
je := json.NewEncoder(conn)
185+
for br.Scan() {
186+
if err := je.Encode(map[string]string{
187+
"data": br.Text(),
188+
}); err != nil {
189+
return
190+
}
191+
}
192+
}()
193+
go func() {
194+
defer close(stdoutDone)
195+
_, _ = io.Copy(inv.Stdout, conn)
196+
}()
197+
go func() {
198+
defer close(stderrDone)
199+
_, _ = io.Copy(inv.Stderr, conn)
200+
}()
201+
go func() {
202+
defer close(done)
203+
<-inv.Context().Done()
204+
_ = conn.Close()
205+
}()
206+
207+
<-done
208+
209+
return nil
210+
}

0 commit comments

Comments
 (0)