From 29036b5942d1ac2cc182300d8dc81b5ff4b2cf72 Mon Sep 17 00:00:00 2001 From: Fedor Korotkov Date: Mon, 4 May 2026 15:51:18 -0400 Subject: [PATCH] Add SSH exec session options --- api/openapi.yaml | 87 ++++++++++- internal/controller/api_vms_exec.go | 164 +++++++++++++++++--- internal/controller/api_vms_exec_test.go | 97 ++++++++++++ internal/controller/exec_sessions.go | 78 +++++++++- internal/controller/exec_sessions_test.go | 68 ++++++++ internal/controller/sshexec/sshexec.go | 91 ++++++++++- internal/controller/sshexec/sshexec_test.go | 57 ++++++- internal/execstream/frame.go | 17 +- internal/execstream/frame_test.go | 18 +++ internal/tests/exec_test.go | 90 ++++++++++- pkg/client/vms.go | 28 +++- pkg/client/vms_test.go | 14 +- 12 files changed, 768 insertions(+), 41 deletions(-) create mode 100644 internal/controller/api_vms_exec_test.go diff --git a/api/openapi.yaml b/api/openapi.yaml index 10f4d89..fd668e3 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -457,7 +457,7 @@ paths: minLength: 1 required: false - in: query - name: stdin + name: interactive description: | Whether to allocate an interactive standard input for the command @@ -468,6 +468,59 @@ paths: type: boolean default: false required: false + - in: query + name: stdin + deprecated: true + description: | + Deprecated alias for `interactive`. + + If both `interactive` and `stdin` are provided, their values must match. + schema: + type: boolean + default: false + required: false + - in: query + name: tty + description: Whether to allocate a pseudo-terminal for the command + schema: + type: boolean + default: false + required: false + - in: query + name: rows + description: Initial terminal row count when `tty=true` + schema: + type: integer + minimum: 0 + maximum: 4294967295 + required: false + - in: query + name: cols + description: Initial terminal column count when `tty=true` + schema: + type: integer + minimum: 0 + maximum: 4294967295 + required: false + - in: query + name: env + description: | + Environment variables to expose to the command. + + Use deep object query syntax, for example `env[FOO]=bar&env[BAZ]=qux`. + style: deepObject + explode: true + schema: + type: object + additionalProperties: + type: string + required: false + - in: query + name: workdir + description: Working directory to switch to before starting the command + schema: + type: string + required: false - in: query name: wait description: Duration in seconds for the VM to become available if it's not available already @@ -507,7 +560,7 @@ paths: '404': description: VM resource with the given name or reconnectable exec session doesn't exist '409': - description: Reconnectable exec session already exists with a different command + description: Reconnectable exec session already exists with different options '503': description: Controller failed to establish a connection with the VM /vms/{name}/ip: @@ -824,6 +877,7 @@ components: description: WebSocket frame from Orchard Client to the Orchard Controller oneOf: - $ref: '#/components/schemas/ExecClientFrameStdin' + - $ref: '#/components/schemas/ExecClientFrameResize' - $ref: '#/components/schemas/ExecClientFrameHistory' - $ref: '#/components/schemas/ExecClientFrameAck' - $ref: '#/components/schemas/ExecClientFrameDetach' @@ -832,6 +886,7 @@ components: propertyName: type mapping: stdin: '#/components/schemas/ExecClientFrameStdin' + resize: '#/components/schemas/ExecClientFrameResize' history: '#/components/schemas/ExecClientFrameHistory' ack: '#/components/schemas/ExecClientFrameAck' detach: '#/components/schemas/ExecClientFrameDetach' @@ -854,6 +909,21 @@ components: example: type: stdin data: aGVsbG8K + ExecClientFrameResize: + description: Resize the pseudo-terminal for a TTY exec session + type: object + required: [ type, terminal ] + properties: + type: + type: string + enum: [ resize ] + terminal: + $ref: '#/components/schemas/ExecTerminalSize' + example: + type: resize + terminal: + rows: 40 + cols: 120 ExecClientFrameHistory: description: Request buffered output strictly newer than the supplied watermark type: object @@ -1015,6 +1085,19 @@ components: example: type: no_more_history watermark: 42 + ExecTerminalSize: + description: Pseudo-terminal size + type: object + required: [ rows, cols ] + properties: + rows: + type: integer + format: int64 + minimum: 0 + cols: + type: integer + format: int64 + minimum: 0 Event: title: Generic Resource Event type: object diff --git a/internal/controller/api_vms_exec.go b/internal/controller/api_vms_exec.go index dafb25c..7b72613 100644 --- a/internal/controller/api_vms_exec.go +++ b/internal/controller/api_vms_exec.go @@ -38,7 +38,10 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder { NewErrorResponse("\"command\" parameter cannot be empty")) } - stdin := ctx.Query("stdin") == "true" + spec, runCommand, err := parseExecSessionSpec(ctx, command) + if err != nil { + return responder.JSON(http.StatusBadRequest, NewErrorResponse("%v", err)) + } waitRaw := ctx.DefaultQuery("wait", "10") wait, err := strconv.ParseUint(waitRaw, 10, 16) @@ -47,17 +50,17 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder { } if sessionID != "" { - return controller.execVMReconnectable(ctx, name, sessionID, command, stdin, wait) + return controller.execVMReconnectable(ctx, name, sessionID, spec, runCommand, wait) } - return controller.execVMLegacy(ctx, name, command, stdin, wait) + return controller.execVMLegacy(ctx, name, spec, runCommand, wait) } func (controller *Controller) execVMLegacy( ctx *gin.Context, name string, - command string, - stdin bool, + spec execSessionSpec, + runCommand string, wait uint64, ) responder.Responder { // Look-up the VM @@ -74,8 +77,8 @@ func (controller *Controller) execVMLegacy( waitContext, vm, execSessionKey{vmName: name}, - command, - stdin, + spec, + runCommand, nil, legacyExecSessionPolicy, ) @@ -107,8 +110,8 @@ func (controller *Controller) execVMReconnectable( ctx *gin.Context, name string, sessionID string, - command string, - stdin bool, + spec execSessionSpec, + runCommand string, wait uint64, ) responder.Responder { key := execSessionKey{ @@ -118,12 +121,12 @@ func (controller *Controller) execVMReconnectable( session, ok := controller.execSessions.get(key) if ok { - if !session.commandMatches(command) { + if !session.specMatches(spec) { return responder.JSON(http.StatusConflict, - NewErrorResponse("exec session %q is already running a different command", sessionID)) + NewErrorResponse("exec session %q is already running with different options", sessionID)) } } else { - if command == "" { + if spec.command == "" { return responder.JSON(http.StatusNotFound, NewErrorResponse("exec session %q does not exist", sessionID)) } @@ -143,8 +146,8 @@ func (controller *Controller) execVMReconnectable( waitContext, vm, key, - command, - stdin, + spec, + runCommand, controller.execSessions, reconnectableExecSessionPolicy, ) @@ -153,9 +156,9 @@ func (controller *Controller) execVMReconnectable( return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err)) } - if !session.commandMatches(command) { + if !session.specMatches(spec) { return responder.JSON(http.StatusConflict, - NewErrorResponse("exec session %q is already running a different command", sessionID)) + NewErrorResponse("exec session %q is already running with different options", sessionID)) } } @@ -179,8 +182,8 @@ func (controller *Controller) newSSHExecSession( waitContext context.Context, vm *v1.VM, key execSessionKey, - command string, - stdin bool, + spec execSessionSpec, + runCommand string, registry *execSessionRegistry, policy execSessionPolicy, ) (*execSession, error) { @@ -201,7 +204,12 @@ func (controller *Controller) newSSHExecSession( return nil, err } - exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), stdin) + exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), sshexec.Options{ + Interactive: spec.interactive, + TTY: spec.tty, + Rows: spec.rows, + Cols: spec.cols, + }) if err != nil { sessionContextCancel() _ = portForwardConn.Close() @@ -209,11 +217,12 @@ func (controller *Controller) newSSHExecSession( return nil, fmt.Errorf("failed to establish SSH connection to a VM: %w", err) } - return newExecSessionWithContext( + return newExecSessionWithContextAndSpec( sessionContext, sessionContextCancel, key, - command, + spec, + runCommand, exec, portForwardConn, registry, @@ -288,6 +297,111 @@ var ( errExecSessionClosed = errors.New("exec session closed") ) +func parseExecSessionSpec(ctx *gin.Context, command string) (execSessionSpec, string, error) { + interactive, err := parseExecInteractive(ctx) + if err != nil { + return execSessionSpec{}, "", err + } + + tty, err := parseExecBool(ctx, "tty") + if err != nil { + return execSessionSpec{}, "", err + } + if tty { + interactive = true + } + + rows, err := parseExecUint32(ctx.Query("rows"), "rows") + if err != nil { + return execSessionSpec{}, "", err + } + cols, err := parseExecUint32(ctx.Query("cols"), "cols") + if err != nil { + return execSessionSpec{}, "", err + } + if (rows == 0) != (cols == 0) { + return execSessionSpec{}, "", errors.New("\"rows\" and \"cols\" must be provided together") + } + + spec := execSessionSpec{ + command: command, + interactive: interactive, + tty: tty, + rows: rows, + cols: cols, + env: ctx.QueryMap("env"), + workdir: ctx.Query("workdir"), + } + + runCommand, err := sshexec.CommandWithOptions(command, sshexec.Options{ + Env: spec.env, + Workdir: spec.workdir, + }) + if err != nil { + return execSessionSpec{}, "", err + } + + return spec, runCommand, nil +} + +func parseExecInteractive(ctx *gin.Context) (bool, error) { + interactive, err := parseExecBool(ctx, "interactive") + if err != nil { + return false, err + } + + interactiveRaw, interactivePresent := ctx.GetQuery("interactive") + stdinRaw, stdinPresent := ctx.GetQuery("stdin") + if !stdinPresent { + return interactive, nil + } + + stdin, err := strconv.ParseBool(stdinRaw) + if err != nil { + return false, errors.New("\"stdin\" parameter must be a boolean") + } + + if interactivePresent { + parsedInteractive, _ := strconv.ParseBool(interactiveRaw) + if stdin != parsedInteractive { + return false, errors.New("\"interactive\" and \"stdin\" parameters cannot conflict") + } + } + + if !interactivePresent { + interactive = stdin + } + + return interactive, nil +} + +func parseExecBool(ctx *gin.Context, name string) (bool, error) { + raw, present := ctx.GetQuery(name) + if !present { + return false, nil + } + + value, err := strconv.ParseBool(raw) + if err != nil { + return false, fmt.Errorf("%q parameter must be a boolean", name) + } + + return value, nil +} + +func parseExecUint32(raw string, name string) (uint32, error) { + if raw == "" { + return 0, nil + } + + value, err := strconv.ParseUint(raw, 10, 32) + if err != nil { + return 0, fmt.Errorf("%q parameter must be an unsigned integer", name) + } + + return uint32(value), nil +} + func (controller *Controller) readExecSessionFrames( ctx context.Context, wsConn *websocket.Conn, @@ -320,6 +434,14 @@ func (controller *Controller) readExecSessionFrames( if err := session.writeStdin(frame.Data); err != nil { return fmt.Errorf("failed to handle %q frame: %w", frame.Type, err) } + case execstream.FrameTypeResize: + if frame.Terminal == nil { + return fmt.Errorf("failed to handle %q frame: terminal size is required", frame.Type) + } + + if err := session.resize(frame.Terminal.Rows, frame.Terminal.Cols); err != nil { + return fmt.Errorf("failed to handle %q frame: %w", frame.Type, err) + } case execstream.FrameTypeHistory: if !session.policy.replayEnabled { return fmt.Errorf("unexpected frame type received: %q", frame.Type) diff --git a/internal/controller/api_vms_exec_test.go b/internal/controller/api_vms_exec_test.go new file mode 100644 index 0000000..cfdb6ef --- /dev/null +++ b/internal/controller/api_vms_exec_test.go @@ -0,0 +1,97 @@ +//nolint:testpackage // we need to test unexported exec helpers +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestParseExecInteractive(t *testing.T) { + for _, test := range []struct { + name string + query string + interactive bool + errContains string + }{ + { + name: "default false", + interactive: false, + }, + { + name: "interactive true", + query: "interactive=true", + interactive: true, + }, + { + name: "stdin alias true", + query: "stdin=true", + interactive: true, + }, + { + name: "matching values accepted", + query: "interactive=true&stdin=true", + interactive: true, + }, + { + name: "conflicting values rejected", + query: "interactive=true&stdin=false", + errContains: "cannot conflict", + }, + { + name: "invalid interactive rejected", + query: "interactive=maybe", + errContains: "interactive", + }, + { + name: "invalid stdin rejected", + query: "stdin=maybe", + errContains: "stdin", + }, + } { + t.Run(test.name, func(t *testing.T) { + interactive, err := parseExecInteractive(execQueryContext(test.query)) + if test.errContains != "" { + require.ErrorContains(t, err, test.errContains) + + return + } + + require.NoError(t, err) + require.Equal(t, test.interactive, interactive) + }) + } +} + +func TestParseExecSessionSpec(t *testing.T) { + spec, runCommand, err := parseExecSessionSpec( + execQueryContext("interactive=true&tty=true&rows=24&cols=80&env[GREETING]=hello&workdir=/tmp"), + "printf '%s' \"$GREETING\"", + ) + require.NoError(t, err) + require.Equal(t, execSessionSpec{ + command: "printf '%s' \"$GREETING\"", + interactive: true, + tty: true, + rows: 24, + cols: 80, + env: map[string]string{"GREETING": "hello"}, + workdir: "/tmp", + }, spec) + require.Equal(t, "cd '/tmp' || exit $?\nexport GREETING='hello'\nprintf '%s' \"$GREETING\"", runCommand) +} + +func TestParseExecSessionSpecRejectsPartialTTYSize(t *testing.T) { + _, _, err := parseExecSessionSpec(execQueryContext("tty=true&rows=24"), "echo hello") + require.ErrorContains(t, err, "provided together") +} + +func execQueryContext(query string) *gin.Context { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = httptest.NewRequest(http.MethodGet, "/?"+query, nil) + + return ctx +} diff --git a/internal/controller/exec_sessions.go b/internal/controller/exec_sessions.go index 7357cec..a3d342f 100644 --- a/internal/controller/exec_sessions.go +++ b/internal/controller/exec_sessions.go @@ -4,6 +4,7 @@ import ( "context" "errors" "io" + "maps" "net" "sync" "time" @@ -31,10 +32,37 @@ var ( type sshExecRunner interface { Stdin() io.WriteCloser + Resize(rows uint32, cols uint32) error Run(ctx context.Context, command string, outgoingFrames chan<- *execstream.Frame) error Close() error } +type execSessionSpec struct { + command string + interactive bool + tty bool + rows uint32 + cols uint32 + env map[string]string + workdir string +} + +func (spec execSessionSpec) clone() execSessionSpec { + spec.env = maps.Clone(spec.env) + + return spec +} + +func (spec execSessionSpec) equal(other execSessionSpec) bool { + return spec.command == other.command && + spec.interactive == other.interactive && + spec.tty == other.tty && + spec.rows == other.rows && + spec.cols == other.cols && + spec.workdir == other.workdir && + maps.Equal(spec.env, other.env) +} + type execSessionKey struct { vmName string sessionID string @@ -298,6 +326,7 @@ func (subscriber *execSessionSubscriber) close() { type execSession struct { key execSessionKey + spec execSessionSpec command string exec sshExecRunner transport net.Conn @@ -334,10 +363,11 @@ func newExecSession( ) *execSession { ctx, cancel := context.WithCancel(context.Background()) - return newExecSessionWithContext( + return newExecSessionWithContextAndSpec( ctx, cancel, key, + execSessionSpec{command: command}, command, exec, transport, @@ -357,6 +387,32 @@ func newExecSessionWithContext( registry *execSessionRegistry, exitTTL time.Duration, policy execSessionPolicy, +) *execSession { + return newExecSessionWithContextAndSpec( + ctx, + cancel, + key, + execSessionSpec{command: command}, + command, + exec, + transport, + registry, + exitTTL, + policy, + ) +} + +func newExecSessionWithContextAndSpec( + ctx context.Context, + cancel context.CancelFunc, + key execSessionKey, + spec execSessionSpec, + command string, + exec sshExecRunner, + transport net.Conn, + registry *execSessionRegistry, + exitTTL time.Duration, + policy execSessionPolicy, ) *execSession { if ctx == nil || cancel == nil { ctx, cancel = context.WithCancel(context.Background()) @@ -364,6 +420,7 @@ func newExecSessionWithContext( session := &execSession{ key: key, + spec: spec.clone(), command: command, exec: exec, transport: transport, @@ -380,8 +437,8 @@ func newExecSessionWithContext( return session } -func (session *execSession) commandMatches(command string) bool { - return command == "" || session.command == command +func (session *execSession) specMatches(spec execSessionSpec) bool { + return spec.command == "" || session.spec.equal(spec) } func (session *execSession) start() { @@ -468,6 +525,17 @@ func (session *execSession) writeStdin(data []byte) error { return err } +func (session *execSession) resize(rows uint32, cols uint32) error { + session.mu.Lock() + defer session.mu.Unlock() + + if !session.spec.tty { + return errors.New("this exec session has no TTY") + } + + return session.exec.Resize(rows, cols) +} + func (session *execSession) ack(watermark uint64) { if !session.policy.replayEnabled { return @@ -663,6 +731,10 @@ func cloneExecFrame(frame *execstream.Frame) *execstream.Frame { exit := *frame.Exit clone.Exit = &exit } + if frame.Terminal != nil { + terminal := *frame.Terminal + clone.Terminal = &terminal + } return &clone } diff --git a/internal/controller/exec_sessions_test.go b/internal/controller/exec_sessions_test.go index cd19eb6..7164635 100644 --- a/internal/controller/exec_sessions_test.go +++ b/internal/controller/exec_sessions_test.go @@ -14,6 +14,7 @@ import ( type fakeExec struct { stdin io.WriteCloser run func(context.Context, string, chan<- *execstream.Frame) error + resize func(uint32, uint32) error closeCalls atomic.Int32 } @@ -33,6 +34,14 @@ func (exec *fakeExec) Run( return nil } +func (exec *fakeExec) Resize(rows uint32, cols uint32) error { + if exec.resize != nil { + return exec.resize(rows, cols) + } + + return nil +} + func (exec *fakeExec) Close() error { exec.closeCalls.Add(1) @@ -47,6 +56,7 @@ func newManualExecSessionForTest( return &execSession{ key: key, + spec: execSessionSpec{command: "echo test"}, command: "echo test", exec: &fakeExec{}, registry: registry, @@ -129,6 +139,64 @@ func TestExecSessionStartRunsCommandOnlyOnce(t *testing.T) { require.EqualValues(t, 1, runCalls.Load()) } +func TestExecSessionSpecMatchesOptions(t *testing.T) { + session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, nil) + session.spec = execSessionSpec{ + command: "echo test", + interactive: true, + tty: true, + rows: 24, + cols: 80, + env: map[string]string{"GREETING": "hello"}, + workdir: "/tmp", + } + + require.True(t, session.specMatches(execSessionSpec{})) + require.True(t, session.specMatches(execSessionSpec{ + command: "echo test", + interactive: true, + tty: true, + rows: 24, + cols: 80, + env: map[string]string{"GREETING": "hello"}, + workdir: "/tmp", + })) + require.False(t, session.specMatches(execSessionSpec{ + command: "echo test", + interactive: true, + tty: true, + rows: 24, + cols: 80, + env: map[string]string{"GREETING": "goodbye"}, + workdir: "/tmp", + })) +} + +func TestExecSessionResizeRequiresTTY(t *testing.T) { + session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, nil) + + err := session.resize(24, 80) + require.ErrorContains(t, err, "no TTY") +} + +func TestExecSessionResizeDelegatesToRunner(t *testing.T) { + var resizedRows, resizedCols uint32 + session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, nil) + session.spec.tty = true + session.exec = &fakeExec{ + resize: func(rows uint32, cols uint32) error { + resizedRows = rows + resizedCols = cols + + return nil + }, + } + + require.NoError(t, session.resize(24, 80)) + require.EqualValues(t, 24, resizedRows) + require.EqualValues(t, 80, resizedCols) +} + func TestExecSessionHistoryReplayAndAck(t *testing.T) { registry := newExecSessionRegistry() session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry) diff --git a/internal/controller/sshexec/sshexec.go b/internal/controller/sshexec/sshexec.go index bb412c3..91cc8ab 100644 --- a/internal/controller/sshexec/sshexec.go +++ b/internal/controller/sshexec/sshexec.go @@ -6,13 +6,27 @@ import ( "fmt" "io" "net" + "regexp" "slices" + "sort" + "strings" "github.com/cirruslabs/orchard/internal/execstream" "golang.org/x/crypto/ssh" "golang.org/x/sync/errgroup" ) +var envNamePattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`) + +type Options struct { + Interactive bool + TTY bool + Rows uint32 + Cols uint32 + Env map[string]string + Workdir string +} + type Exec struct { sshClient *ssh.Client sshSession *ssh.Session @@ -20,9 +34,10 @@ type Exec struct { stderr io.Reader stdin io.WriteCloser stdinReader *io.PipeReader + tty bool } -func New(netConn net.Conn, user string, password string, stdin bool) (*Exec, error) { +func New(netConn net.Conn, user string, password string, options Options) (*Exec, error) { // Establish an SSH connection sshConn, sshChans, sshReqs, err := ssh.NewClientConn(netConn, "", &ssh.ClientConfig{ HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error { @@ -50,15 +65,30 @@ func New(netConn net.Conn, user string, password string, stdin bool) (*Exec, err exec := &Exec{ sshClient: sshClient, sshSession: sshSession, + tty: options.TTY, } - if stdin { + if options.Interactive || options.TTY { stdinReader, stdinWriter := io.Pipe() sshSession.Stdin = stdinReader exec.stdinReader = stdinReader exec.stdin = stdinWriter } + if options.TTY { + if err := sshSession.RequestPty( + "xterm-256color", + int(options.Rows), + int(options.Cols), + ssh.TerminalModes{}, + ); err != nil { + _ = sshSession.Close() + _ = sshClient.Close() + + return nil, fmt.Errorf("failed to request PTY for the SSH session: %w", err) + } + } + exec.stdout, err = sshSession.StdoutPipe() if err != nil { _ = sshSession.Close() @@ -84,6 +114,63 @@ func (exec *Exec) Stdin() io.WriteCloser { return exec.stdin } +func CommandWithOptions(command string, options Options) (string, error) { + if strings.ContainsRune(options.Workdir, '\x00') { + return "", errors.New("working directory contains NUL byte") + } + + keys := make([]string, 0, len(options.Env)) + for key, value := range options.Env { + if !envNamePattern.MatchString(key) { + return "", fmt.Errorf("invalid environment variable name %q", key) + } + + if strings.ContainsRune(value, '\x00') { + return "", fmt.Errorf("environment variable %q contains NUL byte", key) + } + + keys = append(keys, key) + } + sort.Strings(keys) + + if command == "" { + return command, nil + } + + if options.Workdir == "" && len(keys) == 0 { + return command, nil + } + + var builder strings.Builder + if options.Workdir != "" { + builder.WriteString("cd ") + builder.WriteString(shellQuote(options.Workdir)) + builder.WriteString(" || exit $?\n") + } + for _, key := range keys { + builder.WriteString("export ") + builder.WriteString(key) + builder.WriteByte('=') + builder.WriteString(shellQuote(options.Env[key])) + builder.WriteByte('\n') + } + builder.WriteString(command) + + return builder.String(), nil +} + +func shellQuote(value string) string { + return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'" +} + +func (exec *Exec) Resize(rows uint32, cols uint32) error { + if !exec.tty { + return errors.New("this exec session does not have a TTY") + } + + return exec.sshSession.WindowChange(int(rows), int(cols)) +} + func (exec *Exec) Run( ctx context.Context, command string, diff --git a/internal/controller/sshexec/sshexec_test.go b/internal/controller/sshexec/sshexec_test.go index bd360ea..d1f2c53 100644 --- a/internal/controller/sshexec/sshexec_test.go +++ b/internal/controller/sshexec/sshexec_test.go @@ -2,6 +2,7 @@ package sshexec_test import ( "net" + "strings" "testing" "time" @@ -21,6 +22,60 @@ func TestContextCancellationViaNetConnClose(t *testing.T) { } }() - _, err := sshexec.New(clientConn, "doesn't", "matter", false) + _, err := sshexec.New(clientConn, "doesn't", "matter", sshexec.Options{}) require.Error(t, err) } + +func TestCommandWithOptionsNoOptionsLeaveCommandUnchanged(t *testing.T) { + command, err := sshexec.CommandWithOptions("echo hello", sshexec.Options{}) + require.NoError(t, err) + require.Equal(t, "echo hello", command) +} + +func TestCommandWithOptionsSortsAndQuotes(t *testing.T) { + command, err := sshexec.CommandWithOptions( + "printf '%s|%s|%s' \"$GREETING\" \"$NAME\" \"$MULTILINE\"", + sshexec.Options{ + Workdir: "/tmp/a'b", + Env: map[string]string{ + "NAME": "O'Reilly", + "GREETING": "hello $USER", + "MULTILINE": "line 1\nline 2", + }, + }, + ) + require.NoError(t, err) + require.Equal(t, strings.Join([]string{ + "cd '/tmp/a'\\''b' || exit $?", + "export GREETING='hello $USER'", + "export MULTILINE='line 1", + "line 2'", + "export NAME='O'\\''Reilly'", + "printf '%s|%s|%s' \"$GREETING\" \"$NAME\" \"$MULTILINE\"", + }, "\n"), command) +} + +func TestCommandWithOptionsRejectsInvalidName(t *testing.T) { + _, err := sshexec.CommandWithOptions("echo hello", sshexec.Options{ + Env: map[string]string{ + "1INVALID": "value", + }, + }) + require.ErrorContains(t, err, "invalid environment variable name") +} + +func TestCommandWithOptionsRejectsNULValue(t *testing.T) { + _, err := sshexec.CommandWithOptions("echo hello", sshexec.Options{ + Env: map[string]string{ + "VALID": "bad\x00value", + }, + }) + require.ErrorContains(t, err, "contains NUL byte") +} + +func TestCommandWithOptionsRejectsNULWorkdir(t *testing.T) { + _, err := sshexec.CommandWithOptions("echo hello", sshexec.Options{ + Workdir: "bad\x00dir", + }) + require.ErrorContains(t, err, "working directory contains NUL byte") +} diff --git a/internal/execstream/frame.go b/internal/execstream/frame.go index ab41f0b..9175b50 100644 --- a/internal/execstream/frame.go +++ b/internal/execstream/frame.go @@ -11,6 +11,7 @@ type FrameType string const ( FrameTypeStdin FrameType = "stdin" + FrameTypeResize FrameType = "resize" FrameTypeStdout FrameType = "stdout" FrameTypeStderr FrameType = "stderr" FrameTypeExit FrameType = "exit" @@ -23,17 +24,23 @@ const ( ) type Frame struct { - Type FrameType `json:"type"` - Data []byte `json:"data,omitempty"` - Exit *Exit `json:"exit,omitempty"` - Error string `json:"error,omitempty"` - Watermark uint64 `json:"watermark,omitempty"` + Type FrameType `json:"type"` + Data []byte `json:"data,omitempty"` + Terminal *TerminalSize `json:"terminal,omitempty"` + Exit *Exit `json:"exit,omitempty"` + Error string `json:"error,omitempty"` + Watermark uint64 `json:"watermark,omitempty"` } type Exit struct { Code int32 `json:"code"` } +type TerminalSize struct { + Rows uint32 `json:"rows"` + Cols uint32 `json:"cols"` +} + func WriteFrame(ctx context.Context, wsConn *websocket.Conn, frame *Frame) error { frameBytes, err := json.Marshal(frame) if err != nil { diff --git a/internal/execstream/frame_test.go b/internal/execstream/frame_test.go index 673f0c3..3b08c73 100644 --- a/internal/execstream/frame_test.go +++ b/internal/execstream/frame_test.go @@ -21,3 +21,21 @@ func TestFrameRoundTripsWatermark(t *testing.T) { require.NoError(t, err) require.Equal(t, frame, decoded) } + +func TestFrameRoundTripsTerminalSize(t *testing.T) { + frame := Frame{ + Type: FrameTypeResize, + Terminal: &TerminalSize{ + Rows: 24, + Cols: 80, + }, + } + + payload, err := json.Marshal(frame) + require.NoError(t, err) + + var decoded Frame + err = json.Unmarshal(payload, &decoded) + require.NoError(t, err) + require.Equal(t, frame, decoded) +} diff --git a/internal/tests/exec_test.go b/internal/tests/exec_test.go index 6d2005c..2d80c52 100644 --- a/internal/tests/exec_test.go +++ b/internal/tests/exec_test.go @@ -80,6 +80,94 @@ func TestVMExecWithStdin(t *testing.T) { require.Equal(t, websocket.StatusNormalClosure, closeError.Code) } +func TestVMExecWithOptions(t *testing.T) { + devClient, vmName := prepareForExec(t) + + script := "sh -c 'printf \"%s|%s|%s\" \"$GREETING\" \"$QUOTE\" \"$PWD\"'" + + wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + Command: script, + Env: map[string]string{ + "GREETING": "Hello, World!", + "QUOTE": "O'Reilly", + }, + Workdir: "/tmp", + WaitSeconds: 30, + }) + require.NoError(t, err) + defer wsConn.CloseNow() + + var stdout bytes.Buffer + var exitFrame *execstream.Frame + + for exitFrame == nil { + frame := readFrame(t, wsConn) + + switch frame.Type { + case execstream.FrameTypeStdout: + stdout.Write(frame.Data) + case execstream.FrameTypeExit: + exitFrame = frame + default: + t.Fatalf("unexpected frame type %q", frame.Type) + } + } + + require.EqualValues(t, 0, exitFrame.Exit.Code) + require.Equal(t, "Hello, World!|O'Reilly|/tmp", stdout.String()) +} + +func TestVMExecTTYResize(t *testing.T) { + devClient, vmName := prepareForExec(t) + + wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ + Command: "sh -c 'stty size; read -r line; stty size'", + Interactive: true, + TTY: true, + Rows: 24, + Cols: 80, + WaitSeconds: 30, + }) + require.NoError(t, err) + defer wsConn.CloseNow() + + firstFrame := readFrame(t, wsConn) + require.Equal(t, execstream.FrameTypeStdout, firstFrame.Type) + require.Contains(t, string(firstFrame.Data), "24 80") + + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{ + Type: execstream.FrameTypeResize, + Terminal: &execstream.TerminalSize{ + Rows: 40, + Cols: 120, + }, + }) + require.NoError(t, err) + err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{ + Type: execstream.FrameTypeStdin, + Data: []byte("continue\n"), + }) + require.NoError(t, err) + + var stdout bytes.Buffer + var exitFrame *execstream.Frame + for exitFrame == nil { + frame := readFrame(t, wsConn) + + switch frame.Type { + case execstream.FrameTypeStdout: + stdout.Write(frame.Data) + case execstream.FrameTypeExit: + exitFrame = frame + default: + t.Fatalf("unexpected frame type %q", frame.Type) + } + } + + require.Contains(t, stdout.String(), "40 120") + require.EqualValues(t, 0, exitFrame.Exit.Code) +} + func TestVMExecScript(t *testing.T) { devClient, vmName := prepareForExec(t) @@ -255,7 +343,7 @@ func TestVMExecSessionStdinSurvivesReconnect(t *testing.T) { wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{ Command: "/bin/cat", - Stdin: true, + Interactive: true, WaitSeconds: 30, Session: sessionID, }) diff --git a/pkg/client/vms.go b/pkg/client/vms.go index 7b73544..a8190c9 100644 --- a/pkg/client/vms.go +++ b/pkg/client/vms.go @@ -49,7 +49,12 @@ type EventsPageOptions struct { type ExecSessionOptions struct { Command string - Stdin bool + Interactive bool + TTY bool + Rows uint32 + Cols uint32 + Env map[string]string + Workdir string WaitSeconds uint16 Session string } @@ -173,7 +178,7 @@ func (service *VMsService) Exec( ) (*websocket.Conn, error) { return service.ExecSession(ctx, name, ExecSessionOptions{ Command: command, - Stdin: stdin, + Interactive: stdin, WaitSeconds: waitSeconds, }) } @@ -189,8 +194,23 @@ func (service *VMsService) ExecSession( if options.Command != "" { params["command"] = options.Command } - if options.Stdin { - params["stdin"] = strconv.FormatBool(true) + if options.Interactive { + params["interactive"] = strconv.FormatBool(true) + } + if options.TTY { + params["tty"] = strconv.FormatBool(true) + } + if options.Rows > 0 { + params["rows"] = strconv.FormatUint(uint64(options.Rows), 10) + } + if options.Cols > 0 { + params["cols"] = strconv.FormatUint(uint64(options.Cols), 10) + } + for key, value := range options.Env { + params[fmt.Sprintf("env[%s]", key)] = value + } + if options.Workdir != "" { + params["workdir"] = options.Workdir } if options.Session != "" { params["session"] = options.Session diff --git a/pkg/client/vms_test.go b/pkg/client/vms_test.go index 79f3394..710c30b 100644 --- a/pkg/client/vms_test.go +++ b/pkg/client/vms_test.go @@ -26,7 +26,12 @@ func TestExecSessionBuildsReconnectableQuery(t *testing.T) { conn, err := devClient.VMs().ExecSession(t.Context(), "vm", ExecSessionOptions{ Command: "echo hello", - Stdin: true, + Interactive: true, + TTY: true, + Rows: 24, + Cols: 80, + Env: map[string]string{"GREETING": "hello"}, + Workdir: "/tmp", WaitSeconds: 7, Session: "resume-me", }) @@ -34,7 +39,12 @@ func TestExecSessionBuildsReconnectableQuery(t *testing.T) { defer conn.CloseNow() require.Equal(t, []string{"echo hello"}, query["command"]) - require.Equal(t, []string{"true"}, query["stdin"]) + require.Equal(t, []string{"true"}, query["interactive"]) + require.Equal(t, []string{"true"}, query["tty"]) + require.Equal(t, []string{"24"}, query["rows"]) + require.Equal(t, []string{"80"}, query["cols"]) + require.Equal(t, []string{"hello"}, query["env[GREETING]"]) + require.Equal(t, []string{"/tmp"}, query["workdir"]) require.Equal(t, []string{"7"}, query["wait"]) require.Equal(t, []string{"resume-me"}, query["session"]) }