Summarize API exec changes
This commit is contained in:
Fedor Korotkov
parent
c4b7378883
commit
8cf8c68557
147
api/openapi.yaml
147
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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
Loading…
Reference in New Issue