From 6beaf7c4cbf5da0b35154a7cb190db226fd15cb3 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Tue, 25 Feb 2025 12:09:51 +0000 Subject: [PATCH 1/4] chore(cli): skip TestDotfiles if GPG_TTY is set to avoid hanging --- cli/dotfiles_test.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cli/dotfiles_test.go b/cli/dotfiles_test.go index 2f16929cc24ff..002f001e04574 100644 --- a/cli/dotfiles_test.go +++ b/cli/dotfiles_test.go @@ -17,6 +17,10 @@ import ( func TestDotfiles(t *testing.T) { t.Parallel() + // This test will time out if the user has commit signing enabled. + if _, gpgTTYFound := os.LookupEnv("GPG_TTY"); gpgTTYFound { + t.Skip("GPG_TTY is set, skipping test to avoid hanging") + } t.Run("MissingArg", func(t *testing.T) { t.Parallel() inv, _ := clitest.New(t, "dotfiles") From 02463f3ced1f2d5d89b7a2f2e0730c2c73735f33 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Tue, 25 Feb 2025 12:10:11 +0000 Subject: [PATCH 2/4] chore(cli): rename exp commands to have exp_ prefix --- cli/{errors.go => exp_errors.go} | 0 cli/{errors_test.go => exp_errors_test.go} | 0 cli/{prompts.go => exp_prompts.go} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename cli/{errors.go => exp_errors.go} (100%) rename cli/{errors_test.go => exp_errors_test.go} (100%) rename cli/{prompts.go => exp_prompts.go} (100%) diff --git a/cli/errors.go b/cli/exp_errors.go similarity index 100% rename from cli/errors.go rename to cli/exp_errors.go diff --git a/cli/errors_test.go b/cli/exp_errors_test.go similarity index 100% rename from cli/errors_test.go rename to cli/exp_errors_test.go diff --git a/cli/prompts.go b/cli/exp_prompts.go similarity index 100% rename from cli/prompts.go rename to cli/exp_prompts.go From a9ff0456991b57e93e0b4e5a7e9b3bc984bc09b6 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Tue, 25 Feb 2025 16:48:56 +0000 Subject: [PATCH 3/4] feat(cli): add experimental rpty command --- cli/exp.go | 1 + cli/exp_rpty.go | 217 +++++++++++++++++++++++++++++++++++++++++++ cli/exp_rpty_test.go | 111 ++++++++++++++++++++++ 3 files changed, 329 insertions(+) create mode 100644 cli/exp_rpty.go create mode 100644 cli/exp_rpty_test.go diff --git a/cli/exp.go b/cli/exp.go index 5c72d0f9fcd20..2339da86313a6 100644 --- a/cli/exp.go +++ b/cli/exp.go @@ -14,6 +14,7 @@ func (r *RootCmd) expCmd() *serpent.Command { r.scaletestCmd(), r.errorExample(), r.promptExample(), + r.rptyCommand(), }, } return cmd diff --git a/cli/exp_rpty.go b/cli/exp_rpty.go new file mode 100644 index 0000000000000..0ab38c1a45fc6 --- /dev/null +++ b/cli/exp_rpty.go @@ -0,0 +1,217 @@ +package cli + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "os" + "strings" + + "github.com/google/uuid" + "github.com/mattn/go-isatty" + "golang.org/x/term" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/cli/cliui" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/workspacesdk" + "github.com/coder/coder/v2/pty" + "github.com/coder/serpent" +) + +func (r *RootCmd) rptyCommand() *serpent.Command { + var ( + client = new(codersdk.Client) + args handleRPTYArgs + ) + + cmd := &serpent.Command{ + Handler: func(inv *serpent.Invocation) error { + if r.disableDirect { + return xerrors.New("direct connections are disabled, but you can try websocat ;-)") + } + args.NamedWorkspace = inv.Args[0] + args.Command = inv.Args[1:] + return handleRPTY(inv, client, args) + }, + Long: "Establish an RPTY session with a workspace/agent. This uses the same mechanism as the Web Terminal.", + Middleware: serpent.Chain( + serpent.RequireRangeArgs(1, -1), + r.InitClient(client), + ), + Options: []serpent.Option{ + { + Name: "container", + Description: "The container name or ID to connect to.", + Flag: "container", + FlagShorthand: "c", + Default: "", + Value: serpent.StringOf(&args.Container), + }, + { + Name: "container-user", + Description: "The user to connect as.", + Flag: "container-user", + FlagShorthand: "u", + Default: "", + Value: serpent.StringOf(&args.ContainerUser), + }, + { + Name: "reconnect", + Description: "The reconnect ID to use.", + Flag: "reconnect", + FlagShorthand: "r", + Default: "", + Value: serpent.StringOf(&args.ReconnectID), + }, + }, + Short: "Establish an RPTY session with a workspace/agent.", + Use: "rpty", + } + + return cmd +} + +type handleRPTYArgs struct { + Command []string + Container string + ContainerUser string + NamedWorkspace string + ReconnectID string +} + +func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPTYArgs) error { + ctx, cancel := context.WithCancel(inv.Context()) + defer cancel() + + var reconnectID uuid.UUID + if args.ReconnectID != "" { + rid, err := uuid.Parse(args.ReconnectID) + if err != nil { + return xerrors.Errorf("invalid reconnect ID: %w", err) + } + reconnectID = rid + } else { + reconnectID = uuid.New() + } + ws, agt, err := getWorkspaceAndAgent(ctx, inv, client, true, args.NamedWorkspace) + if err != nil { + return err + } + + var ctID string + if args.Container != "" { + cts, err := client.WorkspaceAgentListContainers(ctx, agt.ID, nil) + if err != nil { + return err + } + for _, ct := range cts.Containers { + if ct.FriendlyName == args.Container || ct.ID == args.Container { + ctID = ct.ID + break + } + } + if ctID == "" { + return xerrors.Errorf("container %q not found", args.Container) + } + } + + if err := cliui.Agent(ctx, inv.Stderr, agt.ID, cliui.AgentOptions{ + FetchInterval: 0, + Fetch: client.WorkspaceAgent, + FetchLogs: client.WorkspaceAgentLogsAfter, + Wait: false, + }); err != nil { + return err + } + + // Get the width and height of the terminal. + var termWidth, termHeight uint16 + stdoutFile, validOut := inv.Stdout.(*os.File) + if validOut && isatty.IsTerminal(stdoutFile.Fd()) { + w, h, err := term.GetSize(int(stdoutFile.Fd())) + if err == nil { + //nolint: gosec + termWidth, termHeight = uint16(w), uint16(h) + } + } + + // Set stdin to raw mode so that control characters work. + stdinFile, validIn := inv.Stdin.(*os.File) + if validIn && isatty.IsTerminal(stdinFile.Fd()) { + inState, err := pty.MakeInputRaw(stdinFile.Fd()) + if err != nil { + return xerrors.Errorf("failed to set input terminal to raw mode: %w", err) + } + defer func() { + _ = pty.RestoreTerminal(stdinFile.Fd(), inState) + }() + } + + conn, err := workspacesdk.New(client).AgentReconnectingPTY(ctx, workspacesdk.WorkspaceAgentReconnectingPTYOpts{ + AgentID: agt.ID, + Reconnect: reconnectID, + Command: strings.Join(args.Command, " "), + Container: ctID, + ContainerUser: args.ContainerUser, + Width: termWidth, + Height: termHeight, + }) + if err != nil { + return xerrors.Errorf("open reconnecting PTY: %w", err) + } + defer conn.Close() + + cliui.Infof(inv.Stderr, "Connected to %s (agent id: %s)", args.NamedWorkspace, agt.ID) + closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, ws.ID, codersdk.PostWorkspaceUsageRequest{ + AgentID: agt.ID, + AppName: codersdk.UsageAppNameReconnectingPty, + }) + defer closeUsage() + + stdinDone := make(chan struct{}) + stdoutDone := make(chan struct{}) + stderrDone := make(chan struct{}) + done := make(chan struct{}) + + go func() { + defer close(stdinDone) + // This is how we send commands to the agent. + br := bufio.NewScanner(inv.Stdin) + // Split on bytes, otherwise you have to send a newline to flush the buffer. + br.Split(bufio.ScanBytes) + je := json.NewEncoder(conn) + for br.Scan() { + if err := je.Encode(map[string]string{ + "data": br.Text(), + }); err != nil { + return + } + } + }() + go func() { + defer func() { + close(stdoutDone) + }() + _, _ = io.Copy(inv.Stdout, conn) + }() + go func() { + defer func() { + close(stderrDone) + }() + _, _ = io.Copy(inv.Stderr, conn) + }() + go func() { + defer close(done) + <-stdoutDone + <-stderrDone + _ = conn.Close() + _, _ = fmt.Fprintf(inv.Stderr, "Connection closed\n") + }() + + <-done + + return nil +} diff --git a/cli/exp_rpty_test.go b/cli/exp_rpty_test.go new file mode 100644 index 0000000000000..b215ea3a14dd8 --- /dev/null +++ b/cli/exp_rpty_test.go @@ -0,0 +1,111 @@ +package cli_test + +import ( + "fmt" + "runtime" + "testing" + + "github.com/ory/dockertest/v3" + "github.com/ory/dockertest/v3/docker" + + "github.com/coder/coder/v2/agent" + "github.com/coder/coder/v2/agent/agenttest" + "github.com/coder/coder/v2/cli/clitest" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/pty/ptytest" + "github.com/coder/coder/v2/testutil" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestExpRpty(t *testing.T) { + t.Parallel() + + t.Run("OK", func(t *testing.T) { + t.Parallel() + + client, workspace, agentToken := setupWorkspaceForAgent(t) + inv, root := clitest.New(t, "exp", "rpty", workspace.Name) + clitest.SetupConfig(t, client, root) + pty := ptytest.New(t).Attach(inv) + + ctx := testutil.Context(t, testutil.WaitLong) + + cmdDone := tGo(t, func() { + err := inv.WithContext(ctx).Run() + assert.NoError(t, err) + }) + + _ = agenttest.New(t, client.URL, agentToken) + _ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait() + + pty.ExpectMatch(fmt.Sprintf("Connected to %s", workspace.Name)) + pty.WriteLine("exit") + <-cmdDone + }) + + t.Run("NotFound", func(t *testing.T) { + t.Parallel() + + client, _, _ := setupWorkspaceForAgent(t) + inv, root := clitest.New(t, "exp", "rpty", "not-found") + clitest.SetupConfig(t, client, root) + + ctx := testutil.Context(t, testutil.WaitShort) + err := inv.WithContext(ctx).Run() + require.ErrorContains(t, err, "not found") + }) + + t.Run("Container", func(t *testing.T) { + t.Parallel() + // Skip this test on non-Linux platforms since it requires Docker + if runtime.GOOS != "linux" { + t.Skip("Skipping test on non-Linux platform") + } + + client, workspace, agentToken := setupWorkspaceForAgent(t) + ctx := testutil.Context(t, testutil.WaitLong) + pool, err := dockertest.NewPool("") + require.NoError(t, err, "Could not connect to docker") + ct, err := pool.RunWithOptions(&dockertest.RunOptions{ + Repository: "busybox", + Tag: "latest", + Cmd: []string{"sleep", "infnity"}, + }, func(config *docker.HostConfig) { + config.AutoRemove = true + config.RestartPolicy = docker.RestartPolicy{Name: "no"} + }) + require.NoError(t, err, "Could not start container") + // Wait for container to start + require.Eventually(t, func() bool { + ct, ok := pool.ContainerByName(ct.Container.Name) + return ok && ct.Container.State.Running + }, testutil.WaitShort, testutil.IntervalSlow, "Container did not start in time") + t.Cleanup(func() { + err := pool.Purge(ct) + require.NoError(t, err, "Could not stop container") + }) + + inv, root := clitest.New(t, "exp", "rpty", workspace.Name, "-c", ct.Container.ID) + clitest.SetupConfig(t, client, root) + pty := ptytest.New(t).Attach(inv) + + cmdDone := tGo(t, func() { + err := inv.WithContext(ctx).Run() + assert.NoError(t, err) + }) + + _ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) { + o.ExperimentalContainersEnabled = true + }) + _ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait() + + pty.ExpectMatch(fmt.Sprintf("Connected to %s", workspace.Name)) + pty.ExpectMatch(" #") + pty.WriteLine("hostname") + pty.ExpectMatch(ct.Container.Config.Hostname) + pty.WriteLine("exit") + <-cmdDone + }) +} From b98d4130925c60a0671547e7287437218474e54b Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Wed, 26 Feb 2025 12:19:08 +0000 Subject: [PATCH 4/4] address PR feedback --- cli/exp_rpty.go | 57 ++++++++++++++++++++++---------------------- cli/exp_rpty_test.go | 1 + 2 files changed, 29 insertions(+), 29 deletions(-) diff --git a/cli/exp_rpty.go b/cli/exp_rpty.go index 0ab38c1a45fc6..ddfdc15ece58d 100644 --- a/cli/exp_rpty.go +++ b/cli/exp_rpty.go @@ -121,7 +121,6 @@ func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPT if err := cliui.Agent(ctx, inv.Stderr, agt.ID, cliui.AgentOptions{ FetchInterval: 0, Fetch: client.WorkspaceAgent, - FetchLogs: client.WorkspaceAgentLogsAfter, Wait: false, }); err != nil { return err @@ -165,24 +164,19 @@ func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPT defer conn.Close() cliui.Infof(inv.Stderr, "Connected to %s (agent id: %s)", args.NamedWorkspace, agt.ID) + cliui.Infof(inv.Stderr, "Reconnect ID: %s", reconnectID) closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, ws.ID, codersdk.PostWorkspaceUsageRequest{ AgentID: agt.ID, AppName: codersdk.UsageAppNameReconnectingPty, }) defer closeUsage() - stdinDone := make(chan struct{}) - stdoutDone := make(chan struct{}) - stderrDone := make(chan struct{}) - done := make(chan struct{}) + br := bufio.NewScanner(inv.Stdin) + // Split on bytes, otherwise you have to send a newline to flush the buffer. + br.Split(bufio.ScanBytes) + je := json.NewEncoder(conn) go func() { - defer close(stdinDone) - // This is how we send commands to the agent. - br := bufio.NewScanner(inv.Stdin) - // Split on bytes, otherwise you have to send a newline to flush the buffer. - br.Split(bufio.ScanBytes) - je := json.NewEncoder(conn) for br.Scan() { if err := je.Encode(map[string]string{ "data": br.Text(), @@ -191,27 +185,32 @@ func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPT } } }() + + windowChange := listenWindowSize(ctx) go func() { - defer func() { - close(stdoutDone) - }() - _, _ = io.Copy(inv.Stdout, conn) - }() - go func() { - defer func() { - close(stderrDone) - }() - _, _ = io.Copy(inv.Stderr, conn) - }() - go func() { - defer close(done) - <-stdoutDone - <-stderrDone - _ = conn.Close() - _, _ = fmt.Fprintf(inv.Stderr, "Connection closed\n") + for { + select { + case <-ctx.Done(): + return + case <-windowChange: + } + width, height, err := term.GetSize(int(stdoutFile.Fd())) + if err != nil { + continue + } + if err := je.Encode(map[string]int{ + "width": width, + "height": height, + }); err != nil { + cliui.Errorf(inv.Stderr, "Failed to send window size: %v", err) + } + } }() - <-done + _, _ = io.Copy(inv.Stdout, conn) + cancel() + _ = conn.Close() + _, _ = fmt.Fprintf(inv.Stderr, "Connection closed\n") return nil } diff --git a/cli/exp_rpty_test.go b/cli/exp_rpty_test.go index b215ea3a14dd8..2f0a24bf1cf41 100644 --- a/cli/exp_rpty_test.go +++ b/cli/exp_rpty_test.go @@ -102,6 +102,7 @@ func TestExpRpty(t *testing.T) { _ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait() pty.ExpectMatch(fmt.Sprintf("Connected to %s", workspace.Name)) + pty.ExpectMatch("Reconnect ID: ") pty.ExpectMatch(" #") pty.WriteLine("hostname") pty.ExpectMatch(ct.Container.Config.Hostname)