diff --git a/api/openapi.yaml b/api/openapi.yaml index e36986a..ba84512 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -415,6 +415,85 @@ paths: description: VM resource with the given name doesn't exist '503': description: Failed to establish connection with the worker responsible for the specified VM + /vms/{name}/exec: + parameters: + - in: path + name: name + required: true + schema: + type: string + get: + summary: "Execute a command inside a VM using WebSocket protocol" + description: | + Upgrades to a WebSocket connection and exchanges JSON text frames. + + Frame schema: + + * Client -> Server uses `ExecClientFrame` + * Server -> Client uses `ExecServerFrame` + + `data` is a base64-encoded byte payload in JSON (`format: byte`). + + Client-side end of stdin is signaled by sending a `stdin` frame with an empty `data` payload. + + Example client frame: + `{"type":"stdin","data":"aGVsbG8K"}` + + Example server frames: + `{"type":"stdout","data":"aGVsbG8K"}` + `{"type":"exit","exit":{"code":0}}` + tags: + - vms + parameters: + - in: query + name: command + description: Command to execute in the guest + schema: + type: string + required: true + - in: query + name: arg + description: Command argument, can be provided multiple times + style: form + explode: true + schema: + type: array + items: + type: string + required: false + - in: query + name: wait + description: Duration in seconds to wait for the VM to transition into "running" state if not already running. + schema: + type: integer + minimum: 0 + maximum: 65535 + default: 10 + required: false + - in: header + name: Connection + description: WebSocket protocol required header + required: true + schema: + type: string + - in: header + name: Upgrade + description: WebSocket protocol required header + required: true + schema: + type: string + responses: + '101': + description: | + WebSocket protocol upgrade succeeded. + + After upgrade, messages follow `ExecClientFrame` and `ExecServerFrame` schemas. + '400': + description: Invalid query parameter specified + '404': + description: VM resource with the given name doesn't exist + '503': + description: Failed to establish SSH session to the specified VM /vms/{name}/ip: parameters: - in: path @@ -693,6 +772,74 @@ components: ip: type: string description: The resolved IP address + ExecTerminalSize: + title: VM Exec terminal size + type: object + properties: + rows: + type: integer + minimum: 1 + description: Terminal row count + cols: + type: integer + minimum: 1 + description: Terminal column count + ExecExit: + title: VM Exec exit payload + type: object + required: + - code + properties: + code: + type: integer + format: int32 + description: Process exit code + ExecClientFrame: + title: VM Exec client WebSocket frame + type: object + required: + - type + properties: + type: + type: string + enum: [ stdin, resize ] + description: | + Frame type sent by the client. + `stdin` sends process input bytes. + `resize` is accepted but ignored by the current SSH-backed implementation. + data: + type: string + format: byte + description: | + Base64-encoded stdin payload for `type=stdin`. + Empty payload indicates stdin EOF. + terminal: + $ref: '#/components/schemas/ExecTerminalSize' + example: + type: stdin + data: aGVsbG8K + ExecServerFrame: + title: VM Exec server WebSocket frame + type: object + required: + - type + properties: + type: + type: string + enum: [ stdout, stderr, exit, error ] + description: Frame type sent by the server + data: + type: string + format: byte + description: Base64-encoded output payload for `type=stdout` and `type=stderr` + exit: + $ref: '#/components/schemas/ExecExit' + error: + type: string + description: Error message for `type=error` + example: + type: stdout + data: aGVsbG8K Event: title: Generic Resource Event type: object diff --git a/internal/controller/api.go b/internal/controller/api.go index 4cbf3d5..3c31b24 100644 --- a/internal/controller/api.go +++ b/internal/controller/api.go @@ -171,6 +171,9 @@ func (controller *Controller) initAPI() *gin.Engine { v1.GET("/vms/:name/port-forward", func(c *gin.Context) { controller.portForwardVM(c).Respond(c) }) + v1.GET("/vms/:name/exec", func(c *gin.Context) { + controller.execVM(c).Respond(c) + }) v1.GET("/vms/:name/ip", func(c *gin.Context) { controller.ip(c).Respond(c) }) diff --git a/internal/controller/api_vms_exec.go b/internal/controller/api_vms_exec.go new file mode 100644 index 0000000..c59e9b1 --- /dev/null +++ b/internal/controller/api_vms_exec.go @@ -0,0 +1,484 @@ +package controller + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "strconv" + "strings" + "time" + + "github.com/cirruslabs/orchard/internal/execstream" + "github.com/cirruslabs/orchard/internal/netconncancel" + "github.com/cirruslabs/orchard/internal/responder" + v1 "github.com/cirruslabs/orchard/pkg/resource/v1" + "github.com/cirruslabs/orchard/rpc" + "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "golang.org/x/crypto/ssh" +) + +const execSessionRendezvousTimeout = 15 * time.Second + +type sshExec struct { + session *ssh.Session + stdout io.Reader + stderr io.Reader + stdin io.WriteCloser +} + +func (controller *Controller) execVM(ctx *gin.Context) responder.Responder { + if responder := controller.authorizeAny(ctx, v1.ServiceAccountRoleComputeWrite, + v1.ServiceAccountRoleComputeConnect); responder != nil { + return responder + } + + name := ctx.Param("name") + + command := ctx.Query("command") + if command == "" { + return responder.Code(http.StatusBadRequest) + } + + args := ctx.QueryArray("arg") + + waitRaw := ctx.DefaultQuery("wait", "10") + wait, err := strconv.ParseUint(waitRaw, 10, 16) + if err != nil { + return responder.Code(http.StatusBadRequest) + } + + waitCtx, waitCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second) + defer waitCancel() + + vm, responderImpl := controller.waitForVM(waitCtx, name) + if responderImpl != nil { + return responderImpl + } + + rvCtx, rvCancel := context.WithCancel(ctx) + defer rvCancel() + + session := uuid.New().String() + + connCh, cancel := controller.connRendezvous.Request(rvCtx, session) + defer cancel() + + err = controller.workerNotifier.Notify(waitCtx, vm.Worker, &rpc.WatchInstruction{ + Action: &rpc.WatchInstruction_PortForwardAction{ + PortForwardAction: &rpc.WatchInstruction_PortForward{ + Session: session, + VmUid: vm.UID, + Port: 22, + }, + }, + }) + if err != nil { + controller.logger.Warnf("failed to request VM SSH port-forwarding from the worker %s: %v", + vm.Worker, err) + + return responder.Code(http.StatusServiceUnavailable) + } + + timeoutTimer := time.NewTimer(execSessionRendezvousTimeout) + defer timeoutTimer.Stop() + + select { + case rvResp := <-connCh: + if rvResp.ErrorMessage != "" { + return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse( + "failed to establish SSH connection to the VM on the worker: %s", rvResp.ErrorMessage)) + } + + if rvResp.Result == nil { + return responder.Code(http.StatusServiceUnavailable) + } + + ws, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{ + OriginPatterns: []string{"*"}, + }) + if err != nil { + _ = rvResp.Result.Close() + + return responder.Error(err) + } + defer func() { + _ = ws.CloseNow() + }() + + tunnel := netconncancel.New(rvResp.Result, rvCancel) + defer func() { + _ = tunnel.Close() + }() + + return controller.execVMViaSSHTunnel(ctx, tunnel, ws, vm, command, args) + case <-timeoutTimer.C: + return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse( + "timed out waiting for worker %s to establish SSH tunnel", vm.Worker)) + case <-ctx.Done(): + return responder.Error(ctx.Err()) + } +} + +func (controller *Controller) execVMViaSSHTunnel(ctx context.Context, tunnel net.Conn, ws *websocket.Conn, vm *v1.VM, cmd string, args []string) responder.Responder { + sshClient, err := newSSHClient(tunnel, vm) + if err != nil { + controller.closeExecWithFrameError(ctx, ws, nil, + fmt.Sprintf("SSH handshake with the VM failed: %v", err)) + + return responder.Empty() + } + defer func() { + _ = sshClient.Close() + }() + + exec, err := startSSHExec(sshClient, cmd, args) + if err != nil { + controller.closeExecWithFrameError(ctx, ws, nil, err.Error()) + + return responder.Empty() + } + defer func() { + _ = exec.session.Close() + }() + + return controller.pumpExecFrames(ctx, ws, exec) +} + +func newSSHClient(conn net.Conn, vm *v1.VM) (*ssh.Client, error) { + sshUser := vm.Username + sshPassword := vm.Password + if sshUser == "" && sshPassword == "" { + sshUser = "admin" + sshPassword = "admin" + } + + sshConn, chans, reqs, err := ssh.NewClientConn(conn, "", &ssh.ClientConfig{ + HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error { + return nil + }, + User: sshUser, + Auth: []ssh.AuthMethod{ + ssh.Password(sshPassword), + }, + Timeout: 10 * time.Second, + }) + if err != nil { + return nil, err + } + + return ssh.NewClient(sshConn, chans, reqs), nil +} + +func startSSHExec(client *ssh.Client, cmd string, args []string) (*sshExec, error) { + session, err := client.NewSession() + if err != nil { + return nil, fmt.Errorf("failed to open SSH session: %v", err) + } + + stdoutPipe, err := session.StdoutPipe() + if err != nil { + _ = session.Close() + + return nil, fmt.Errorf("failed to get SSH stdout pipe: %v", err) + } + + stderrPipe, err := session.StderrPipe() + if err != nil { + _ = session.Close() + + return nil, fmt.Errorf("failed to get SSH stderr pipe: %v", err) + } + + stdinPipe, err := session.StdinPipe() + if err != nil { + _ = session.Close() + + return nil, fmt.Errorf("failed to get SSH stdin pipe: %v", err) + } + + sshCommand := buildSSHCommand(cmd, args) + if err := session.Start(sshCommand); err != nil { + _ = session.Close() + + return nil, fmt.Errorf("failed to start SSH command: %v", err) + } + + return &sshExec{ + session: session, + stdout: stdoutPipe, + stderr: stderrPipe, + stdin: stdinPipe, + }, nil +} + +func (controller *Controller) pumpExecFrames( + ctx context.Context, + ws *websocket.Conn, + exec *sshExec, +) responder.Responder { + wsNetConn := websocket.NetConn(ctx, ws, websocket.MessageText) + defer func() { + _ = wsNetConn.Close() + }() + + encoder := execstream.NewEncoder(wsNetConn) + decoder := execstream.NewDecoder(wsNetConn) + + outCh := make(chan execstream.Frame, 16) + outDoneCh := make(chan struct{}, 2) + outErrCh := make(chan error, 1) + stdinErrCh := make(chan error, 1) + exitCh := make(chan int32, 1) + exitErrCh := make(chan error, 1) + + go streamExecOutput(exec.stdout, execstream.FrameTypeStdout, outCh, outDoneCh, outErrCh) + go streamExecOutput(exec.stderr, execstream.FrameTypeStderr, outCh, outDoneCh, outErrCh) + go streamExecClientFrames(decoder, exec.stdin, stdinErrCh) + go waitForSSHExecExit(exec.session, exitCh, exitErrCh) + + pingTicker := time.NewTicker(controller.pingInterval) + defer pingTicker.Stop() + + readersDone := 0 + exitObserved := false + var exitCode int32 + + for { + if exitObserved && readersDone >= 2 { + for len(outCh) > 0 { + frame := <-outCh + if err := execstream.WriteFrame(encoder, &frame); err != nil { + return controller.wsError(ws, websocket.StatusInternalError, "exec session", + "failed to stream exec output to the client", err) + } + } + + if err := execstream.WriteFrame(encoder, &execstream.Frame{ + Type: execstream.FrameTypeExit, + Exit: &execstream.Exit{Code: exitCode}, + }); err != nil { + return controller.wsError(ws, websocket.StatusInternalError, "exec session", + "failed to send exec exit status to the client", err) + } + + if err := ws.Close(websocket.StatusNormalClosure, + fmt.Sprintf("command exited with code %d", exitCode)); err != nil { + controller.logger.Warnf("exec session: failed to close WebSocket connection: %v", err) + } + + return responder.Empty() + } + + select { + case frame := <-outCh: + if err := execstream.WriteFrame(encoder, &frame); err != nil { + return controller.wsError(ws, websocket.StatusInternalError, "exec session", + "failed to stream exec output to the client", err) + } + case <-outDoneCh: + readersDone++ + case err := <-outErrCh: + if err == nil { + continue + } + + controller.closeExecWithFrameError(ctx, ws, encoder, + fmt.Sprintf("failed while streaming command output: %v", err)) + + return responder.Empty() + case err := <-stdinErrCh: + if err == nil || errors.Is(err, context.Canceled) { + stdinErrCh = nil + continue + } + + if errors.Is(err, io.EOF) { + return responder.Empty() + } + + controller.closeExecWithFrameError(ctx, ws, encoder, + fmt.Sprintf("failed while reading command stdin stream: %v", err)) + + return responder.Empty() + case code := <-exitCh: + exitObserved = true + exitCode = code + case err := <-exitErrCh: + controller.closeExecWithFrameError(ctx, ws, encoder, + fmt.Sprintf("failed while waiting for command completion: %v", err)) + + return responder.Empty() + case <-pingTicker.C: + pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second) + + if err := ws.Ping(pingCtx); err != nil { + controller.logger.Warnf("exec session: failed to ping the client, "+ + "connection might time out: %v", err) + } + + pingCtxCancel() + case <-ctx.Done(): + return responder.Error(ctx.Err()) + } + } +} + +func waitForSSHExecExit(sshSession *ssh.Session, exitCodeCh chan<- int32, exitErrCh chan<- error) { + if err := sshSession.Wait(); err != nil { + var exitError *ssh.ExitError + if errors.As(err, &exitError) { + exitCodeCh <- int32(exitError.ExitStatus()) + + return + } + + exitErrCh <- err + + return + } + + exitCodeCh <- 0 +} + +func streamExecClientFrames( + decoder *json.Decoder, + stdin io.WriteCloser, + errCh chan<- error, +) { + stdinClosed := false + + for { + var frame execstream.Frame + + if err := execstream.ReadFrame(decoder, &frame); err != nil { + if !stdinClosed { + if closeErr := stdin.Close(); closeErr != nil { + errCh <- closeErr + + return + } + } + + errCh <- err + + return + } + + switch frame.Type { + case execstream.FrameTypeStdin: + if len(frame.Data) == 0 { + if !stdinClosed { + if err := stdin.Close(); err != nil { + errCh <- err + + return + } + + stdinClosed = true + } + + errCh <- nil + + return + } + + if stdinClosed { + errCh <- errors.New("stdin is already closed") + + return + } + + if _, err := stdin.Write(frame.Data); err != nil { + errCh <- err + + return + } + case execstream.FrameTypeResize: + // No-op for SSH backend without TTY support. + default: + errCh <- fmt.Errorf("unsupported frame type %q received from client", frame.Type) + + return + } + } +} + +func streamExecOutput( + reader io.Reader, + frameType execstream.FrameType, + outputCh chan<- execstream.Frame, + doneCh chan<- struct{}, + errCh chan<- error, +) { + defer func() { + doneCh <- struct{}{} + }() + + for { + buffer := make([]byte, 4096) + n, err := reader.Read(buffer) + if n > 0 { + outputCh <- execstream.Frame{ + Type: frameType, + Data: append([]byte(nil), buffer[:n]...), + } + } + + if errors.Is(err, io.EOF) { + return + } + + if err != nil { + select { + case errCh <- err: + default: + } + + return + } + } +} + +func buildSSHCommand(command string, args []string) string { + parts := make([]string, 0, 1+len(args)) + parts = append(parts, shellQuoteArg(command)) + for _, arg := range args { + parts = append(parts, shellQuoteArg(arg)) + } + + return strings.Join(parts, " ") +} + +func shellQuoteArg(arg string) string { + if arg == "" { + return "''" + } + + return "'" + strings.ReplaceAll(arg, "'", "'\\''") + "'" +} + +func (controller *Controller) closeExecWithFrameError( + ctx context.Context, + wsConn *websocket.Conn, + encoder *json.Encoder, + message string, +) { + if encoder != nil { + if err := execstream.WriteFrame(encoder, &execstream.Frame{ + Type: execstream.FrameTypeError, + Error: message, + }); err != nil { + controller.logger.Warnf("exec session: failed to send error frame: %v", err) + } + } + + if err := wsConn.Close(websocket.StatusInternalError, message); err != nil { + controller.logger.Warnf("exec session: failed to close WebSocket connection: %v", err) + } +} diff --git a/internal/controller/api_vms_exec_test.go b/internal/controller/api_vms_exec_test.go new file mode 100644 index 0000000..9980192 --- /dev/null +++ b/internal/controller/api_vms_exec_test.go @@ -0,0 +1,114 @@ +//nolint:testpackage // we need to have access to unexported helpers +package controller + +import ( + "bytes" + "io" + "testing" + + "github.com/cirruslabs/orchard/internal/execstream" + "github.com/stretchr/testify/require" +) + +type recordingWriteCloser struct { + bytes.Buffer + closed bool +} + +func (writer *recordingWriteCloser) Close() error { + writer.closed = true + + return nil +} + +func TestStreamExecClientFramesWritesInputAndClosesOnEOFFrame(t *testing.T) { + var input bytes.Buffer + encoder := execstream.NewEncoder(&input) + + require.NoError(t, execstream.WriteFrame(encoder, &execstream.Frame{ + Type: execstream.FrameTypeStdin, + Data: []byte("hello"), + })) + require.NoError(t, execstream.WriteFrame(encoder, &execstream.Frame{ + Type: execstream.FrameTypeResize, + })) + require.NoError(t, execstream.WriteFrame(encoder, &execstream.Frame{ + Type: execstream.FrameTypeStdin, + Data: []byte{}, + })) + + decoder := execstream.NewDecoder(&input) + stdin := &recordingWriteCloser{} + errCh := make(chan error, 1) + + streamExecClientFrames(decoder, stdin, errCh) + + require.NoError(t, <-errCh) + require.True(t, stdin.closed) + require.Equal(t, "hello", stdin.String()) +} + +func TestStreamExecClientFramesUnsupportedType(t *testing.T) { + var input bytes.Buffer + encoder := execstream.NewEncoder(&input) + + require.NoError(t, execstream.WriteFrame(encoder, &execstream.Frame{ + Type: execstream.FrameTypeStdout, + Data: []byte("output"), + })) + + decoder := execstream.NewDecoder(&input) + stdin := &recordingWriteCloser{} + errCh := make(chan error, 1) + + streamExecClientFrames(decoder, stdin, errCh) + + require.EqualError(t, <-errCh, "unsupported frame type \"stdout\" received from client") + require.False(t, stdin.closed) +} + +func TestStreamExecClientFramesClosesStdinOnDecodeError(t *testing.T) { + decoder := execstream.NewDecoder(bytes.NewBuffer(nil)) + stdin := &recordingWriteCloser{} + errCh := make(chan error, 1) + + streamExecClientFrames(decoder, stdin, errCh) + + require.ErrorIs(t, <-errCh, io.EOF) + require.True(t, stdin.closed) +} + +func TestStreamExecOutputEmitsFrameAndSignalsDone(t *testing.T) { + outputCh := make(chan execstream.Frame, 1) + doneCh := make(chan struct{}, 1) + errCh := make(chan error, 1) + + streamExecOutput(bytes.NewBufferString("payload"), + execstream.FrameTypeStderr, outputCh, doneCh, errCh) + + select { + case frame := <-outputCh: + require.Equal(t, execstream.FrameTypeStderr, frame.Type) + require.Equal(t, []byte("payload"), frame.Data) + default: + t.Fatal("expected frame") + } + + select { + case <-doneCh: + default: + t.Fatal("expected done signal") + } + + select { + case err := <-errCh: + t.Fatalf("unexpected error: %v", err) + default: + } +} + +func TestBuildSSHCommandQuotesArguments(t *testing.T) { + result := buildSSHCommand("echo", []string{"hello world", "a'b", ""}) + + require.Equal(t, "'echo' 'hello world' 'a'\\''b' ''", result) +} diff --git a/internal/execstream/frame.go b/internal/execstream/frame.go new file mode 100644 index 0000000..a8e0a44 --- /dev/null +++ b/internal/execstream/frame.go @@ -0,0 +1,68 @@ +package execstream + +import ( + "encoding/json" + "io" +) + +type FrameType string + +const ( + FrameTypeCommand FrameType = "command" + FrameTypeStdin FrameType = "stdin" + FrameTypeStdout FrameType = "stdout" + FrameTypeStderr FrameType = "stderr" + FrameTypeResize FrameType = "resize" + FrameTypeExit FrameType = "exit" + FrameTypeError FrameType = "error" +) + +// Frame captures a single event flowing between controller, worker and clients. +// +// The payload is encoded as JSON where binary blobs (stdin/stdout/stderr data) are +// automatically base64-encoded by the JSON encoder. +type Frame struct { + Type FrameType `json:"type"` + + Command *Command `json:"command,omitempty"` + Data []byte `json:"data,omitempty"` + Terminal *TerminalSize `json:"terminal,omitempty"` + Exit *Exit `json:"exit,omitempty"` + Error string `json:"error,omitempty"` +} + +type Command struct { + Name string `json:"name"` + Args []string `json:"args,omitempty"` + Interactive bool `json:"interactive,omitempty"` + TTY bool `json:"tty,omitempty"` + Terminal *TerminalSize `json:"terminal,omitempty"` +} + +type TerminalSize struct { + Rows uint32 `json:"rows"` + Cols uint32 `json:"cols"` +} + +type Exit struct { + Code int32 `json:"code"` +} + +func NewEncoder(w io.Writer) *json.Encoder { + encoder := json.NewEncoder(w) + encoder.SetEscapeHTML(false) + + return encoder +} + +func NewDecoder(r io.Reader) *json.Decoder { + return json.NewDecoder(r) +} + +func WriteFrame(encoder *json.Encoder, frame *Frame) error { + return encoder.Encode(frame) +} + +func ReadFrame(decoder *json.Decoder, frame *Frame) error { + return decoder.Decode(frame) +} diff --git a/internal/tests/exec_test.go b/internal/tests/exec_test.go new file mode 100644 index 0000000..ac9212e --- /dev/null +++ b/internal/tests/exec_test.go @@ -0,0 +1,151 @@ +package tests_test + +import ( + "bytes" + "context" + "fmt" + "net/http" + "net/url" + "testing" + "time" + + "github.com/cirruslabs/orchard/internal/execstream" + "github.com/cirruslabs/orchard/internal/imageconstant" + "github.com/cirruslabs/orchard/internal/tests/devcontroller" + "github.com/cirruslabs/orchard/internal/tests/wait" + v1 "github.com/cirruslabs/orchard/pkg/resource/v1" + "github.com/coder/websocket" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestVMExec(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + devClient, devController, _ := devcontroller.StartIntegrationTestEnvironment(t) + + vmName := "test-vm-exec-" + uuid.NewString() + + err := devClient.VMs().Create(ctx, &v1.VM{ + Meta: v1.Meta{ + Name: vmName, + }, + Image: imageconstant.DefaultMacosImage, + CPU: 4, + Memory: 8 * 1024, + Headless: true, + }) + require.NoError(t, err) + + require.True(t, wait.Wait(2*time.Minute, func() bool { + vm, getErr := devClient.VMs().Get(ctx, vmName) + require.NoError(t, getErr) + + t.Logf("Waiting for the VM to start. Current status: %s", vm.Status) + + return vm.Status == v1.VMStatusRunning || vm.Status == v1.VMStatusFailed + }), "failed to start a VM") + + vm, err := devClient.VMs().Get(ctx, vmName) + require.NoError(t, err) + require.Equal(t, v1.VMStatusRunning, vm.Status) + + execConn, err := dialExec(ctx, devController.Address(), vmName, "sh", []string{ + "-c", + "echo stdout-line; echo stderr-line >&2; IFS= read -r line; echo stdin:$line; exit 7", + }) + require.NoError(t, err) + t.Cleanup(func() { + _ = execConn.Close() + }) + + encoder := execstream.NewEncoder(execConn) + decoder := execstream.NewDecoder(execConn) + + require.NoError(t, execstream.WriteFrame(encoder, &execstream.Frame{ + Type: execstream.FrameTypeStdin, + Data: []byte("hello-from-test\\n"), + })) + require.NoError(t, execstream.WriteFrame(encoder, &execstream.Frame{ + Type: execstream.FrameTypeStdin, + Data: []byte{}, + })) + + var stdout bytes.Buffer + var stderr bytes.Buffer + var exitFrame *execstream.Exit + + for { + var frame execstream.Frame + require.NoError(t, execstream.ReadFrame(decoder, &frame)) + + switch frame.Type { + case execstream.FrameTypeStdout: + stdout.Write(frame.Data) + case execstream.FrameTypeStderr: + stderr.Write(frame.Data) + case execstream.FrameTypeExit: + require.NotNil(t, frame.Exit) + exitFrame = frame.Exit + case execstream.FrameTypeError: + t.Fatalf("unexpected error frame: %s", frame.Error) + default: + t.Fatalf("unexpected frame type: %q", frame.Type) + } + + if exitFrame != nil { + break + } + } + + require.EqualValues(t, 7, exitFrame.Code) + require.Contains(t, stdout.String(), "stdout-line") + require.Contains(t, stdout.String(), "stdin:hello-from-test") + require.Contains(t, stderr.String(), "stderr-line") +} + +func dialExec( + ctx context.Context, + controllerAddress string, + vmName string, + command string, + args []string, +) (interface { + Read([]byte) (int, error) + Write([]byte) (int, error) + Close() error +}, error) { + endpointURL, err := url.Parse(controllerAddress) + if err != nil { + return nil, fmt.Errorf("failed to parse controller address: %w", err) + } + + endpointURL = endpointURL.JoinPath("v1", "vms", vmName, "exec") + if endpointURL.Scheme == "http" { + endpointURL.Scheme = "ws" + } else { + endpointURL.Scheme = "wss" + } + + query := endpointURL.Query() + query.Set("command", command) + for _, arg := range args { + query.Add("arg", arg) + } + query.Set("wait", "120") + endpointURL.RawQuery = query.Encode() + + wsConn, resp, err := websocket.Dial(ctx, endpointURL.String(), &websocket.DialOptions{ + HTTPClient: http.DefaultClient, + }) + if err != nil { + if resp != nil { + _ = resp.Body.Close() + } + + return nil, fmt.Errorf("failed to establish exec websocket: %w", err) + } + + return websocket.NetConn(ctx, wsConn, websocket.MessageText), nil +}