Add reconnectable SSH exec sessions
This commit is contained in:
Fedor Korotkov
parent
88506b1adb
commit
8f5e55bcca
121
api/openapi.yaml
121
api/openapi.yaml
|
|
@ -429,11 +429,33 @@ paths:
|
|||
parameters:
|
||||
- in: query
|
||||
name: command
|
||||
description: Command to execute
|
||||
description: |
|
||||
Command to execute.
|
||||
|
||||
Required when starting a new exec session. May be omitted when reconnecting to an
|
||||
existing session identified by the `session` parameter.
|
||||
schema:
|
||||
type: string
|
||||
minLength: 1
|
||||
required: true
|
||||
required: false
|
||||
- in: query
|
||||
name: session
|
||||
description: |
|
||||
Optional stable exec session identifier. When present, websocket disconnects detach
|
||||
from the command instead of terminating it, and later requests with the same VM name
|
||||
and session id may reconnect and request buffered history.
|
||||
schema:
|
||||
type: string
|
||||
minLength: 1
|
||||
required: false
|
||||
- in: query
|
||||
name: cmux_session_id
|
||||
description: |
|
||||
Compatibility alias for `session`. Prefer `session` for new Orchard clients.
|
||||
schema:
|
||||
type: string
|
||||
minLength: 1
|
||||
required: false
|
||||
- in: query
|
||||
name: stdin
|
||||
description: |
|
||||
|
|
@ -483,7 +505,9 @@ paths:
|
|||
'400':
|
||||
description: Invalid parameters were supplied
|
||||
'404':
|
||||
description: VM resource with the given name doesn't exist
|
||||
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
|
||||
'503':
|
||||
description: Controller failed to establish a connection with the VM
|
||||
/vms/{name}/ip:
|
||||
|
|
@ -800,10 +824,18 @@ components:
|
|||
description: WebSocket frame from Orchard Client to the Orchard Controller
|
||||
oneOf:
|
||||
- $ref: '#/components/schemas/ExecClientFrameStdin'
|
||||
- $ref: '#/components/schemas/ExecClientFrameHistory'
|
||||
- $ref: '#/components/schemas/ExecClientFrameAck'
|
||||
- $ref: '#/components/schemas/ExecClientFrameDetach'
|
||||
- $ref: '#/components/schemas/ExecClientFrameClose'
|
||||
discriminator:
|
||||
propertyName: type
|
||||
mapping:
|
||||
stdin: '#/components/schemas/ExecClientFrameStdin'
|
||||
history: '#/components/schemas/ExecClientFrameHistory'
|
||||
ack: '#/components/schemas/ExecClientFrameAck'
|
||||
detach: '#/components/schemas/ExecClientFrameDetach'
|
||||
close: '#/components/schemas/ExecClientFrameClose'
|
||||
ExecClientFrameStdin:
|
||||
description: Send bytes to the process standard input
|
||||
type: object
|
||||
|
|
@ -822,6 +854,56 @@ components:
|
|||
example:
|
||||
type: stdin
|
||||
data: aGVsbG8K
|
||||
ExecClientFrameHistory:
|
||||
description: Request buffered output strictly newer than the supplied watermark
|
||||
type: object
|
||||
required: [ type, watermark ]
|
||||
properties:
|
||||
type:
|
||||
type: string
|
||||
enum: [ history ]
|
||||
watermark:
|
||||
type: integer
|
||||
format: int64
|
||||
minimum: 0
|
||||
example:
|
||||
type: history
|
||||
watermark: 42
|
||||
ExecClientFrameAck:
|
||||
description: Acknowledge that output has been durably consumed through the supplied watermark
|
||||
type: object
|
||||
required: [ type, watermark ]
|
||||
properties:
|
||||
type:
|
||||
type: string
|
||||
enum: [ ack ]
|
||||
watermark:
|
||||
type: integer
|
||||
format: int64
|
||||
minimum: 0
|
||||
example:
|
||||
type: ack
|
||||
watermark: 42
|
||||
ExecClientFrameDetach:
|
||||
description: Detach this websocket while leaving the remote command running
|
||||
type: object
|
||||
required: [ type ]
|
||||
properties:
|
||||
type:
|
||||
type: string
|
||||
enum: [ detach ]
|
||||
example:
|
||||
type: detach
|
||||
ExecClientFrameClose:
|
||||
description: Close the reconnectable exec session and terminate the remote command
|
||||
type: object
|
||||
required: [ type ]
|
||||
properties:
|
||||
type:
|
||||
type: string
|
||||
enum: [ close ]
|
||||
example:
|
||||
type: close
|
||||
ExecControllerFrame:
|
||||
description: WebSocket frame from Orchard Controller to the Orchard Client
|
||||
oneOf:
|
||||
|
|
@ -829,6 +911,7 @@ components:
|
|||
- $ref: '#/components/schemas/ExecControllerFrameStderr'
|
||||
- $ref: '#/components/schemas/ExecControllerFrameExit'
|
||||
- $ref: '#/components/schemas/ExecControllerFrameError'
|
||||
- $ref: '#/components/schemas/ExecControllerFrameNoMoreHistory'
|
||||
discriminator:
|
||||
propertyName: type
|
||||
mapping:
|
||||
|
|
@ -836,6 +919,7 @@ components:
|
|||
stderr: '#/components/schemas/ExecControllerFrameStderr'
|
||||
exit: '#/components/schemas/ExecControllerFrameExit'
|
||||
error: '#/components/schemas/ExecControllerFrameError'
|
||||
no_more_history: '#/components/schemas/ExecControllerFrameNoMoreHistory'
|
||||
ExecControllerFrameStdout:
|
||||
description: Standard output from the process
|
||||
type: object
|
||||
|
|
@ -848,6 +932,10 @@ components:
|
|||
type: string
|
||||
format: byte
|
||||
description: Base64-encoded standard output bytes from the process
|
||||
watermark:
|
||||
type: integer
|
||||
format: int64
|
||||
description: Monotonic output watermark present on reconnectable sessions
|
||||
example:
|
||||
type: stdout
|
||||
data: aGVsbG8K
|
||||
|
|
@ -863,6 +951,10 @@ components:
|
|||
type: string
|
||||
format: byte
|
||||
description: Base64-encoded standard error bytes from the process
|
||||
watermark:
|
||||
type: integer
|
||||
format: int64
|
||||
description: Monotonic output watermark present on reconnectable sessions
|
||||
example:
|
||||
type: stderr
|
||||
data: aGVsbG8K
|
||||
|
|
@ -882,6 +974,10 @@ components:
|
|||
type: integer
|
||||
format: int32
|
||||
description: Process exit code
|
||||
watermark:
|
||||
type: integer
|
||||
format: int64
|
||||
description: Monotonic output watermark present on reconnectable sessions
|
||||
example:
|
||||
type: exit
|
||||
exit:
|
||||
|
|
@ -897,9 +993,28 @@ components:
|
|||
error:
|
||||
type: string
|
||||
description: Error message text
|
||||
watermark:
|
||||
type: integer
|
||||
format: int64
|
||||
description: Monotonic output watermark present on reconnectable sessions
|
||||
example:
|
||||
type: error
|
||||
error: Failed to establish SSH connection to a VM
|
||||
ExecControllerFrameNoMoreHistory:
|
||||
description: Marker indicating that the requested replay range has been fully sent
|
||||
type: object
|
||||
required: [ type, watermark ]
|
||||
properties:
|
||||
type:
|
||||
type: string
|
||||
enum: [ no_more_history ]
|
||||
watermark:
|
||||
type: integer
|
||||
format: int64
|
||||
description: Highest watermark known to the exec session at the time of replay
|
||||
example:
|
||||
type: no_more_history
|
||||
watermark: 42
|
||||
Event:
|
||||
title: Generic Resource Event
|
||||
type: object
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ var noExperimentalRPCV2 bool
|
|||
var experimentalPingInterval time.Duration
|
||||
var experimentalDisableDBCompression bool
|
||||
var workerOfflineTimeout time.Duration
|
||||
var execSessionExitTTL time.Duration
|
||||
var synthetic bool
|
||||
|
||||
func newRunCommand() *cobra.Command {
|
||||
|
|
@ -87,6 +88,8 @@ func newRunCommand() *cobra.Command {
|
|||
"duration (e.g. 60s or 5m30s) after which a worker is considered offline for the purposes "+
|
||||
"of scheduling (no new VMs will be scheduled on such worker and already assigned VMs will be "+
|
||||
"marked as failed)")
|
||||
cmd.Flags().DurationVar(&execSessionExitTTL, "exec-session-exit-ttl", 10*time.Minute,
|
||||
"duration to retain reconnectable exec session history after the command exits")
|
||||
|
||||
// Hidden flags
|
||||
cmd.Flags().BoolVar(&synthetic, "synthetic", false, "")
|
||||
|
|
@ -147,6 +150,7 @@ func runController(cmd *cobra.Command, args []string) (err error) {
|
|||
controller.WithListenAddr(address),
|
||||
controller.WithDataDir(dataDir),
|
||||
controller.WithWorkerOfflineTimeout(workerOfflineTimeout),
|
||||
controller.WithExecSessionExitTTL(execSessionExitTTL),
|
||||
controller.WithLogger(logger),
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -28,9 +28,13 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder {
|
|||
|
||||
// Retrieve and parse path and query parameters
|
||||
name := ctx.Param("name")
|
||||
sessionID := ctx.Query("session")
|
||||
if sessionID == "" {
|
||||
sessionID = ctx.Query("cmux_session_id")
|
||||
}
|
||||
|
||||
command := ctx.Query("command")
|
||||
if command == "" {
|
||||
if sessionID == "" && command == "" {
|
||||
return responder.JSON(http.StatusBadRequest,
|
||||
NewErrorResponse("\"command\" parameter cannot be empty"))
|
||||
}
|
||||
|
|
@ -43,6 +47,20 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder {
|
|||
return responder.Code(http.StatusBadRequest)
|
||||
}
|
||||
|
||||
if sessionID != "" {
|
||||
return controller.execVMReconnectable(ctx, name, sessionID, command, stdin, wait)
|
||||
}
|
||||
|
||||
return controller.execVMLegacy(ctx, name, command, stdin, wait)
|
||||
}
|
||||
|
||||
func (controller *Controller) execVMLegacy(
|
||||
ctx *gin.Context,
|
||||
name string,
|
||||
command string,
|
||||
stdin bool,
|
||||
wait uint64,
|
||||
) responder.Responder {
|
||||
// Look-up the VM
|
||||
waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second)
|
||||
defer waitContextCancel()
|
||||
|
|
@ -151,6 +169,198 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder {
|
|||
}
|
||||
}
|
||||
|
||||
func (controller *Controller) execVMReconnectable(
|
||||
ctx *gin.Context,
|
||||
name string,
|
||||
sessionID string,
|
||||
command string,
|
||||
stdin bool,
|
||||
wait uint64,
|
||||
) responder.Responder {
|
||||
key := execSessionKey{
|
||||
vmName: name,
|
||||
sessionID: sessionID,
|
||||
}
|
||||
|
||||
session, ok := controller.execSessions.get(key)
|
||||
if ok {
|
||||
if !session.commandMatches(command) {
|
||||
return responder.JSON(http.StatusConflict,
|
||||
NewErrorResponse("exec session %q is already running a different command", sessionID))
|
||||
}
|
||||
} else {
|
||||
if command == "" {
|
||||
return responder.JSON(http.StatusNotFound,
|
||||
NewErrorResponse("exec session %q does not exist", sessionID))
|
||||
}
|
||||
|
||||
waitContext, waitContextCancel := context.WithTimeout(ctx, time.Duration(wait)*time.Second)
|
||||
defer waitContextCancel()
|
||||
|
||||
vm, responderImpl := controller.waitForVM(waitContext, name)
|
||||
if responderImpl != nil {
|
||||
return responderImpl
|
||||
}
|
||||
|
||||
var err error
|
||||
session, _, err = controller.execSessions.getOrCreate(waitContext, key, func() (*execSession, error) {
|
||||
portForwardConn, err := retry.NewWithData[net.Conn](
|
||||
retry.Context(waitContext),
|
||||
retry.DelayType(retry.FixedDelay),
|
||||
retry.Delay(time.Second),
|
||||
retry.Attempts(0),
|
||||
retry.LastErrorOnly(true),
|
||||
).Do(func() (net.Conn, error) {
|
||||
return controller.portForwardConnection(ctx, waitContext, vm.Worker, vm.UID, 22)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), stdin)
|
||||
if err != nil {
|
||||
_ = portForwardConn.Close()
|
||||
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return newExecSession(
|
||||
key,
|
||||
command,
|
||||
exec,
|
||||
portForwardConn,
|
||||
controller.execSessions,
|
||||
controller.execSessionExitTTL,
|
||||
), nil
|
||||
})
|
||||
if err != nil {
|
||||
return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err))
|
||||
}
|
||||
|
||||
if !session.commandMatches(command) {
|
||||
return responder.JSON(http.StatusConflict,
|
||||
NewErrorResponse("exec session %q is already running a different command", sessionID))
|
||||
}
|
||||
}
|
||||
|
||||
wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{
|
||||
OriginPatterns: []string{"*"},
|
||||
})
|
||||
if err != nil {
|
||||
return responder.Error(err)
|
||||
}
|
||||
defer func() {
|
||||
_ = wsConn.CloseNow()
|
||||
}()
|
||||
|
||||
subscriber, err := session.attach()
|
||||
if err != nil {
|
||||
_ = wsConn.Close(websocket.StatusNormalClosure, err.Error())
|
||||
|
||||
return responder.Empty()
|
||||
}
|
||||
defer session.detach(subscriber)
|
||||
|
||||
readFramesErrCh := make(chan error, 1)
|
||||
go func() {
|
||||
readFramesErrCh <- controller.readReconnectableFrames(ctx, wsConn, session, subscriber)
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case readFramesErr := <-readFramesErrCh:
|
||||
if readFramesErr != nil &&
|
||||
!errors.Is(readFramesErr, errExecSessionDetached) &&
|
||||
!errors.Is(readFramesErr, errExecSessionClosed) {
|
||||
controller.logger.Warnf("failed to read and process reconnectable exec frames from WebSocket: %v",
|
||||
readFramesErr)
|
||||
}
|
||||
|
||||
return responder.Empty()
|
||||
case outgoingFrame, ok := <-subscriber.frames:
|
||||
if !ok {
|
||||
if err := wsConn.Close(websocket.StatusNormalClosure, "Command finished"); err != nil {
|
||||
controller.logger.Warnf("exec: failed to close WebSocket cleanly: %v", err)
|
||||
}
|
||||
|
||||
return responder.Empty()
|
||||
}
|
||||
|
||||
if err := execstream.WriteFrame(ctx, wsConn, outgoingFrame); err != nil {
|
||||
controller.logger.Warnf("failed to write reconnectable exec frame to the client: %v", err)
|
||||
|
||||
return responder.Empty()
|
||||
}
|
||||
case <-time.After(controller.pingInterval):
|
||||
pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
|
||||
if err := wsConn.Ping(pingCtx); err != nil {
|
||||
controller.logger.Warnf("reconnectable exec: failed to ping the client, "+
|
||||
"connection might time out: %v", err)
|
||||
}
|
||||
|
||||
pingCtxCancel()
|
||||
case <-ctx.Done():
|
||||
controller.logger.Warnf("client disconnected prematurely")
|
||||
|
||||
return responder.Empty()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
errExecSessionDetached = errors.New("exec session detached")
|
||||
errExecSessionClosed = errors.New("exec session closed")
|
||||
)
|
||||
|
||||
func (controller *Controller) readReconnectableFrames(
|
||||
ctx context.Context,
|
||||
wsConn *websocket.Conn,
|
||||
session *execSession,
|
||||
subscriber *execSessionSubscriber,
|
||||
) error {
|
||||
for {
|
||||
var frame execstream.Frame
|
||||
|
||||
messageType, payloadBytes, err := wsConn.Read(ctx)
|
||||
if err != nil {
|
||||
var closeErr websocket.CloseError
|
||||
if errors.As(err, &closeErr) && closeErr.Code == websocket.StatusNormalClosure {
|
||||
return errExecSessionDetached
|
||||
}
|
||||
|
||||
return fmt.Errorf("failed to read next frame from WebSocket: %w", err)
|
||||
}
|
||||
|
||||
if messageType != websocket.MessageText {
|
||||
continue
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(payloadBytes, &frame); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
switch frame.Type {
|
||||
case execstream.FrameTypeStdin:
|
||||
if err := session.writeStdin(frame.Data); err != nil {
|
||||
return fmt.Errorf("failed to handle %q frame: %w", frame.Type, err)
|
||||
}
|
||||
case execstream.FrameTypeHistory:
|
||||
session.sendHistory(subscriber, frame.Watermark)
|
||||
case execstream.FrameTypeAck:
|
||||
session.ack(frame.Watermark)
|
||||
case execstream.FrameTypeDetach:
|
||||
return errExecSessionDetached
|
||||
case execstream.FrameTypeClose:
|
||||
session.close()
|
||||
|
||||
return errExecSessionClosed
|
||||
default:
|
||||
return fmt.Errorf("unexpected frame type received: %q", frame.Type)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (controller *Controller) readFrames(
|
||||
ctx context.Context,
|
||||
wsConn *websocket.Conn,
|
||||
|
|
|
|||
|
|
@ -55,6 +55,7 @@ type Controller struct {
|
|||
ipRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[string]]
|
||||
enableSwaggerDocs bool
|
||||
workerOfflineTimeout time.Duration
|
||||
execSessionExitTTL time.Duration
|
||||
experimentalRPCV2 bool
|
||||
disableDBCompression bool
|
||||
pingInterval time.Duration
|
||||
|
|
@ -64,6 +65,7 @@ type Controller struct {
|
|||
sshSigner ssh.Signer
|
||||
sshNoClientAuth bool
|
||||
sshServer *sshserver.SSHServer
|
||||
execSessions *execSessionRegistry
|
||||
|
||||
single singleflight.Group
|
||||
|
||||
|
|
@ -75,7 +77,9 @@ func New(opts ...Option) (*Controller, error) {
|
|||
connRendezvous: rendezvous.New[rendezvous.ResultWithErrorMessage[net.Conn]](),
|
||||
ipRendezvous: rendezvous.New[rendezvous.ResultWithErrorMessage[string]](),
|
||||
workerOfflineTimeout: 3 * time.Minute,
|
||||
execSessionExitTTL: 10 * time.Minute,
|
||||
pingInterval: 30 * time.Second,
|
||||
execSessions: newExecSessionRegistry(),
|
||||
single: singleflight.Group{},
|
||||
}
|
||||
|
||||
|
|
@ -308,6 +312,8 @@ func (controller *Controller) Run(ctx context.Context) error {
|
|||
go func() {
|
||||
<-ctx.Done()
|
||||
|
||||
controller.execSessions.closeAll()
|
||||
|
||||
if err := controller.httpServer.Shutdown(ctx); err != nil {
|
||||
controller.logger.Errorf("failed to cleanly shutdown the HTTP server: %v", err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,462 @@
|
|||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/cirruslabs/orchard/internal/execstream"
|
||||
)
|
||||
|
||||
const execSessionReplayBufferBytes = 4 * 1024 * 1024
|
||||
|
||||
type sshExecRunner interface {
|
||||
Stdin() io.WriteCloser
|
||||
Run(ctx context.Context, command string, outgoingFrames chan<- *execstream.Frame) error
|
||||
Close() error
|
||||
}
|
||||
|
||||
type execSessionKey struct {
|
||||
vmName string
|
||||
sessionID string
|
||||
}
|
||||
|
||||
type execSessionCreation struct {
|
||||
done chan struct{}
|
||||
session *execSession
|
||||
err error
|
||||
}
|
||||
|
||||
type execSessionRegistry struct {
|
||||
mu sync.Mutex
|
||||
sessions map[execSessionKey]*execSession
|
||||
creating map[execSessionKey]*execSessionCreation
|
||||
}
|
||||
|
||||
func newExecSessionRegistry() *execSessionRegistry {
|
||||
return &execSessionRegistry{
|
||||
sessions: map[execSessionKey]*execSession{},
|
||||
creating: map[execSessionKey]*execSessionCreation{},
|
||||
}
|
||||
}
|
||||
|
||||
func (registry *execSessionRegistry) get(key execSessionKey) (*execSession, bool) {
|
||||
registry.mu.Lock()
|
||||
defer registry.mu.Unlock()
|
||||
|
||||
session, ok := registry.sessions[key]
|
||||
|
||||
return session, ok
|
||||
}
|
||||
|
||||
func (registry *execSessionRegistry) getOrCreate(
|
||||
ctx context.Context,
|
||||
key execSessionKey,
|
||||
create func() (*execSession, error),
|
||||
) (*execSession, bool, error) {
|
||||
for {
|
||||
registry.mu.Lock()
|
||||
|
||||
if session, ok := registry.sessions[key]; ok {
|
||||
registry.mu.Unlock()
|
||||
|
||||
return session, false, nil
|
||||
}
|
||||
|
||||
if creation, ok := registry.creating[key]; ok {
|
||||
registry.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, false, ctx.Err()
|
||||
case <-creation.done:
|
||||
if creation.err != nil {
|
||||
return nil, false, creation.err
|
||||
}
|
||||
|
||||
return creation.session, false, nil
|
||||
}
|
||||
}
|
||||
|
||||
creation := &execSessionCreation{done: make(chan struct{})}
|
||||
registry.creating[key] = creation
|
||||
registry.mu.Unlock()
|
||||
|
||||
session, err := create()
|
||||
|
||||
registry.mu.Lock()
|
||||
delete(registry.creating, key)
|
||||
if err == nil {
|
||||
registry.sessions[key] = session
|
||||
}
|
||||
creation.session = session
|
||||
creation.err = err
|
||||
close(creation.done)
|
||||
registry.mu.Unlock()
|
||||
|
||||
return session, true, err
|
||||
}
|
||||
}
|
||||
|
||||
func (registry *execSessionRegistry) remove(key execSessionKey, expected *execSession) {
|
||||
registry.mu.Lock()
|
||||
defer registry.mu.Unlock()
|
||||
|
||||
if registry.sessions[key] == expected {
|
||||
delete(registry.sessions, key)
|
||||
}
|
||||
}
|
||||
|
||||
func (registry *execSessionRegistry) closeAll() {
|
||||
registry.mu.Lock()
|
||||
sessions := make([]*execSession, 0, len(registry.sessions))
|
||||
for _, session := range registry.sessions {
|
||||
sessions = append(sessions, session)
|
||||
}
|
||||
registry.mu.Unlock()
|
||||
|
||||
for _, session := range sessions {
|
||||
session.close()
|
||||
}
|
||||
}
|
||||
|
||||
type execReplayFrame struct {
|
||||
frame *execstream.Frame
|
||||
size int
|
||||
}
|
||||
|
||||
type execSessionSubscriber struct {
|
||||
frames chan *execstream.Frame
|
||||
}
|
||||
|
||||
func newExecSessionSubscriber() *execSessionSubscriber {
|
||||
return &execSessionSubscriber{
|
||||
frames: make(chan *execstream.Frame, 128),
|
||||
}
|
||||
}
|
||||
|
||||
func (subscriber *execSessionSubscriber) enqueue(frame *execstream.Frame) bool {
|
||||
select {
|
||||
case subscriber.frames <- cloneExecFrame(frame):
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
type execSession struct {
|
||||
key execSessionKey
|
||||
command string
|
||||
exec sshExecRunner
|
||||
transport net.Conn
|
||||
registry *execSessionRegistry
|
||||
exitTTL time.Duration
|
||||
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
|
||||
mu sync.Mutex
|
||||
stdin io.WriteCloser
|
||||
stdinClosed bool
|
||||
subscribers map[*execSessionSubscriber]struct{}
|
||||
frames []execReplayFrame
|
||||
bufferBytes int
|
||||
nextWatermark uint64
|
||||
ackedWatermark uint64
|
||||
finished bool
|
||||
closed bool
|
||||
expiryTimer *time.Timer
|
||||
|
||||
done chan struct{}
|
||||
doneOnce sync.Once
|
||||
}
|
||||
|
||||
func newExecSession(
|
||||
key execSessionKey,
|
||||
command string,
|
||||
exec sshExecRunner,
|
||||
transport net.Conn,
|
||||
registry *execSessionRegistry,
|
||||
exitTTL time.Duration,
|
||||
) *execSession {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
session := &execSession{
|
||||
key: key,
|
||||
command: command,
|
||||
exec: exec,
|
||||
transport: transport,
|
||||
registry: registry,
|
||||
exitTTL: exitTTL,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
stdin: exec.Stdin(),
|
||||
subscribers: map[*execSessionSubscriber]struct{}{},
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
|
||||
go session.run()
|
||||
|
||||
return session
|
||||
}
|
||||
|
||||
func (session *execSession) commandMatches(command string) bool {
|
||||
return command == "" || session.command == command
|
||||
}
|
||||
|
||||
func (session *execSession) attach() (*execSessionSubscriber, error) {
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
if session.closed {
|
||||
return nil, errors.New("exec session is closed")
|
||||
}
|
||||
|
||||
subscriber := newExecSessionSubscriber()
|
||||
session.subscribers[subscriber] = struct{}{}
|
||||
|
||||
return subscriber, nil
|
||||
}
|
||||
|
||||
func (session *execSession) detach(subscriber *execSessionSubscriber) {
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
session.detachLocked(subscriber)
|
||||
}
|
||||
|
||||
func (session *execSession) detachLocked(subscriber *execSessionSubscriber) {
|
||||
if _, ok := session.subscribers[subscriber]; !ok {
|
||||
return
|
||||
}
|
||||
|
||||
delete(session.subscribers, subscriber)
|
||||
close(subscriber.frames)
|
||||
}
|
||||
|
||||
func (session *execSession) writeStdin(data []byte) error {
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
if session.stdin == nil || session.stdinClosed {
|
||||
return errors.New("this exec session has no stdin enabled or it is already closed")
|
||||
}
|
||||
|
||||
if len(data) == 0 {
|
||||
if err := session.stdin.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
session.stdinClosed = true
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
_, err := session.stdin.Write(data)
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
func (session *execSession) ack(watermark uint64) {
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
if watermark <= session.ackedWatermark {
|
||||
return
|
||||
}
|
||||
|
||||
session.ackedWatermark = watermark
|
||||
session.trimAcknowledgedLocked()
|
||||
}
|
||||
|
||||
func (session *execSession) sendHistory(
|
||||
subscriber *execSessionSubscriber,
|
||||
watermark uint64,
|
||||
) {
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
if _, ok := session.subscribers[subscriber]; !ok {
|
||||
return
|
||||
}
|
||||
|
||||
for _, record := range session.frames {
|
||||
if record.frame.Watermark <= watermark {
|
||||
continue
|
||||
}
|
||||
|
||||
if !subscriber.enqueue(record.frame) {
|
||||
session.detachLocked(subscriber)
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if !subscriber.enqueue(&execstream.Frame{
|
||||
Type: execstream.FrameTypeNoMoreHistory,
|
||||
Watermark: session.nextWatermark,
|
||||
}) {
|
||||
session.detachLocked(subscriber)
|
||||
}
|
||||
}
|
||||
|
||||
func (session *execSession) close() {
|
||||
session.mu.Lock()
|
||||
if session.closed {
|
||||
session.mu.Unlock()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
session.closed = true
|
||||
if session.expiryTimer != nil {
|
||||
session.expiryTimer.Stop()
|
||||
session.expiryTimer = nil
|
||||
}
|
||||
|
||||
subscribers := make([]*execSessionSubscriber, 0, len(session.subscribers))
|
||||
for subscriber := range session.subscribers {
|
||||
subscribers = append(subscribers, subscriber)
|
||||
}
|
||||
session.subscribers = map[*execSessionSubscriber]struct{}{}
|
||||
session.mu.Unlock()
|
||||
|
||||
for _, subscriber := range subscribers {
|
||||
close(subscriber.frames)
|
||||
}
|
||||
|
||||
session.cancel()
|
||||
_ = session.exec.Close()
|
||||
if session.transport != nil {
|
||||
_ = session.transport.Close()
|
||||
}
|
||||
session.registry.remove(session.key, session)
|
||||
}
|
||||
|
||||
func (session *execSession) run() {
|
||||
outgoingFrames := make(chan *execstream.Frame)
|
||||
runErrCh := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
runErrCh <- session.exec.Run(session.ctx, session.command, outgoingFrames)
|
||||
close(outgoingFrames)
|
||||
}()
|
||||
|
||||
for frame := range outgoingFrames {
|
||||
session.recordFrame(frame)
|
||||
}
|
||||
|
||||
runErr := <-runErrCh
|
||||
if runErr != nil && !errors.Is(runErr, context.Canceled) {
|
||||
session.recordFrame(&execstream.Frame{
|
||||
Type: execstream.FrameTypeError,
|
||||
Error: runErr.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
session.markFinished()
|
||||
}
|
||||
|
||||
func (session *execSession) recordFrame(frame *execstream.Frame) {
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
if session.closed {
|
||||
return
|
||||
}
|
||||
|
||||
session.nextWatermark++
|
||||
frame = cloneExecFrame(frame)
|
||||
frame.Watermark = session.nextWatermark
|
||||
|
||||
session.frames = append(session.frames, execReplayFrame{
|
||||
frame: frame,
|
||||
size: execFrameSize(frame),
|
||||
})
|
||||
session.bufferBytes += execFrameSize(frame)
|
||||
session.trimAcknowledgedLocked()
|
||||
session.trimToLimitLocked()
|
||||
|
||||
for subscriber := range session.subscribers {
|
||||
if subscriber.enqueue(frame) {
|
||||
continue
|
||||
}
|
||||
|
||||
session.detachLocked(subscriber)
|
||||
}
|
||||
}
|
||||
|
||||
func (session *execSession) markFinished() {
|
||||
session.mu.Lock()
|
||||
if session.finished {
|
||||
session.mu.Unlock()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
session.finished = true
|
||||
if !session.closed {
|
||||
session.expiryTimer = time.AfterFunc(session.exitTTL, session.expire)
|
||||
}
|
||||
|
||||
subscribers := make([]*execSessionSubscriber, 0, len(session.subscribers))
|
||||
for subscriber := range session.subscribers {
|
||||
subscribers = append(subscribers, subscriber)
|
||||
}
|
||||
session.subscribers = map[*execSessionSubscriber]struct{}{}
|
||||
session.mu.Unlock()
|
||||
|
||||
for _, subscriber := range subscribers {
|
||||
close(subscriber.frames)
|
||||
}
|
||||
|
||||
session.doneOnce.Do(func() {
|
||||
close(session.done)
|
||||
})
|
||||
}
|
||||
|
||||
func (session *execSession) expire() {
|
||||
session.close()
|
||||
}
|
||||
|
||||
func (session *execSession) trimAcknowledgedLocked() {
|
||||
for len(session.frames) > 0 && session.frames[0].frame.Watermark <= session.ackedWatermark {
|
||||
session.bufferBytes -= session.frames[0].size
|
||||
session.frames = session.frames[1:]
|
||||
}
|
||||
}
|
||||
|
||||
func (session *execSession) trimToLimitLocked() {
|
||||
for session.bufferBytes > execSessionReplayBufferBytes && len(session.frames) > 0 {
|
||||
session.bufferBytes -= session.frames[0].size
|
||||
session.frames = session.frames[1:]
|
||||
}
|
||||
}
|
||||
|
||||
func cloneExecFrame(frame *execstream.Frame) *execstream.Frame {
|
||||
if frame == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
clone := *frame
|
||||
if frame.Data != nil {
|
||||
clone.Data = append([]byte(nil), frame.Data...)
|
||||
}
|
||||
if frame.Exit != nil {
|
||||
exit := *frame.Exit
|
||||
clone.Exit = &exit
|
||||
}
|
||||
|
||||
return &clone
|
||||
}
|
||||
|
||||
func execFrameSize(frame *execstream.Frame) int {
|
||||
if frame == nil {
|
||||
return 0
|
||||
}
|
||||
|
||||
return len(frame.Data) + len(frame.Error) + 16
|
||||
}
|
||||
|
|
@ -0,0 +1,172 @@
|
|||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/cirruslabs/orchard/internal/execstream"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type fakeExec struct {
|
||||
stdin io.WriteCloser
|
||||
run func(context.Context, string, chan<- *execstream.Frame) error
|
||||
closeCalls atomic.Int32
|
||||
}
|
||||
|
||||
func (exec *fakeExec) Stdin() io.WriteCloser {
|
||||
return exec.stdin
|
||||
}
|
||||
|
||||
func (exec *fakeExec) Run(
|
||||
ctx context.Context,
|
||||
command string,
|
||||
outgoingFrames chan<- *execstream.Frame,
|
||||
) error {
|
||||
if exec.run != nil {
|
||||
return exec.run(ctx, command, outgoingFrames)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (exec *fakeExec) Close() error {
|
||||
exec.closeCalls.Add(1)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func newManualExecSessionForTest(
|
||||
key execSessionKey,
|
||||
registry *execSessionRegistry,
|
||||
) *execSession {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
return &execSession{
|
||||
key: key,
|
||||
command: "echo test",
|
||||
exec: &fakeExec{},
|
||||
registry: registry,
|
||||
exitTTL: time.Minute,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
subscribers: map[*execSessionSubscriber]struct{}{},
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecSessionRegistryGetOrCreateReusesInflightCreation(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
key := execSessionKey{vmName: "vm", sessionID: "session"}
|
||||
|
||||
createStarted := make(chan struct{})
|
||||
releaseCreate := make(chan struct{})
|
||||
var createCalls atomic.Int32
|
||||
|
||||
create := func() (*execSession, error) {
|
||||
createCalls.Add(1)
|
||||
close(createStarted)
|
||||
<-releaseCreate
|
||||
|
||||
return newManualExecSessionForTest(key, registry), nil
|
||||
}
|
||||
|
||||
firstDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(firstDone)
|
||||
_, _, err := registry.getOrCreate(context.Background(), key, create)
|
||||
require.NoError(t, err)
|
||||
}()
|
||||
|
||||
<-createStarted
|
||||
|
||||
secondDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(secondDone)
|
||||
_, created, err := registry.getOrCreate(context.Background(), key, create)
|
||||
require.NoError(t, err)
|
||||
require.False(t, created)
|
||||
}()
|
||||
|
||||
close(releaseCreate)
|
||||
|
||||
<-firstDone
|
||||
<-secondDone
|
||||
require.EqualValues(t, 1, createCalls.Load())
|
||||
}
|
||||
|
||||
func TestExecSessionHistoryReplayAndAck(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
|
||||
|
||||
session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStdout, Data: []byte("out")})
|
||||
session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStderr, Data: []byte("err")})
|
||||
session.recordFrame(&execstream.Frame{
|
||||
Type: execstream.FrameTypeExit,
|
||||
Exit: &execstream.Exit{Code: 7},
|
||||
})
|
||||
|
||||
subscriber, err := session.attach()
|
||||
require.NoError(t, err)
|
||||
|
||||
session.sendHistory(subscriber, 0)
|
||||
|
||||
require.Equal(t, execstream.FrameTypeStdout, (<-subscriber.frames).Type)
|
||||
require.Equal(t, execstream.FrameTypeStderr, (<-subscriber.frames).Type)
|
||||
require.Equal(t, execstream.FrameTypeExit, (<-subscriber.frames).Type)
|
||||
noMoreHistory := <-subscriber.frames
|
||||
require.Equal(t, execstream.FrameTypeNoMoreHistory, noMoreHistory.Type)
|
||||
require.EqualValues(t, 3, noMoreHistory.Watermark)
|
||||
|
||||
session.ack(2)
|
||||
require.Len(t, session.frames, 1)
|
||||
require.EqualValues(t, 3, session.frames[0].frame.Watermark)
|
||||
}
|
||||
|
||||
func TestExecSessionDetachKeepsProcessAlive(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
|
||||
exec := session.exec.(*fakeExec)
|
||||
|
||||
subscriber, err := session.attach()
|
||||
require.NoError(t, err)
|
||||
|
||||
session.detach(subscriber)
|
||||
|
||||
require.False(t, session.closed)
|
||||
require.EqualValues(t, 0, exec.closeCalls.Load())
|
||||
}
|
||||
|
||||
func TestExecSessionCloseStopsProcessAndRemovesRegistryEntry(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
key := execSessionKey{vmName: "vm", sessionID: "session"}
|
||||
session := newManualExecSessionForTest(key, registry)
|
||||
exec := session.exec.(*fakeExec)
|
||||
registry.sessions[key] = session
|
||||
|
||||
session.close()
|
||||
|
||||
require.True(t, session.closed)
|
||||
require.EqualValues(t, 1, exec.closeCalls.Load())
|
||||
_, ok := registry.get(key)
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
func TestExecSessionFinishedEntryExpiresAfterTTL(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
key := execSessionKey{vmName: "vm", sessionID: "session"}
|
||||
session := newManualExecSessionForTest(key, registry)
|
||||
session.exitTTL = 10 * time.Millisecond
|
||||
registry.sessions[key] = session
|
||||
|
||||
session.markFinished()
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
_, ok := registry.get(key)
|
||||
|
||||
return !ok
|
||||
}, time.Second, 10*time.Millisecond)
|
||||
}
|
||||
|
|
@ -60,6 +60,12 @@ func WithWorkerOfflineTimeout(workerOfflineTimeout time.Duration) Option {
|
|||
}
|
||||
}
|
||||
|
||||
func WithExecSessionExitTTL(execSessionExitTTL time.Duration) Option {
|
||||
return func(controller *Controller) {
|
||||
controller.execSessionExitTTL = execSessionExitTTL
|
||||
}
|
||||
}
|
||||
|
||||
func WithExperimentalRPCV2() Option {
|
||||
return func(controller *Controller) {
|
||||
controller.experimentalRPCV2 = true
|
||||
|
|
|
|||
|
|
@ -15,6 +15,11 @@ const (
|
|||
FrameTypeStderr FrameType = "stderr"
|
||||
FrameTypeExit FrameType = "exit"
|
||||
FrameTypeError FrameType = "error"
|
||||
FrameTypeHistory FrameType = "history"
|
||||
FrameTypeNoMoreHistory FrameType = "no_more_history"
|
||||
FrameTypeAck FrameType = "ack"
|
||||
FrameTypeDetach FrameType = "detach"
|
||||
FrameTypeClose FrameType = "close"
|
||||
)
|
||||
|
||||
type Frame struct {
|
||||
|
|
@ -22,6 +27,7 @@ type Frame struct {
|
|||
Data []byte `json:"data,omitempty"`
|
||||
Exit *Exit `json:"exit,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Watermark uint64 `json:"watermark,omitempty"`
|
||||
}
|
||||
|
||||
type Exit struct {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,23 @@
|
|||
package execstream
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestFrameRoundTripsWatermark(t *testing.T) {
|
||||
frame := Frame{
|
||||
Type: FrameTypeHistory,
|
||||
Watermark: 42,
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
|
@ -132,6 +132,167 @@ func TestVMExecScript(t *testing.T) {
|
|||
require.Equal(t, websocket.StatusNormalClosure, closeError.Code)
|
||||
}
|
||||
|
||||
func TestVMExecSessionReconnectHistory(t *testing.T) {
|
||||
devClient, vmName := prepareForExec(t)
|
||||
sessionID := uuid.NewString()
|
||||
|
||||
wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
|
||||
Command: "sh -c 'echo first; sleep 1; echo second'",
|
||||
WaitSeconds: 30,
|
||||
Session: sessionID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
firstFrame := readFrame(t, wsConn)
|
||||
require.Equal(t, execstream.FrameTypeStdout, firstFrame.Type)
|
||||
require.Equal(t, "first\n", string(firstFrame.Data))
|
||||
require.EqualValues(t, 1, firstFrame.Watermark)
|
||||
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach})
|
||||
require.NoError(t, err)
|
||||
_ = wsConn.CloseNow()
|
||||
|
||||
wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
|
||||
WaitSeconds: 30,
|
||||
Session: sessionID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer wsConn.CloseNow()
|
||||
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{
|
||||
Type: execstream.FrameTypeHistory,
|
||||
Watermark: firstFrame.Watermark,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
frames := readFramesUntilExit(t, wsConn)
|
||||
require.Len(t, framesByType(frames, execstream.FrameTypeStdout), 1)
|
||||
require.Equal(t, "second\n", string(framesByType(frames, execstream.FrameTypeStdout)[0].Data))
|
||||
require.EqualValues(t, 0, framesByType(frames, execstream.FrameTypeExit)[0].Exit.Code)
|
||||
}
|
||||
|
||||
func TestVMExecSessionReconnectAfterExit(t *testing.T) {
|
||||
devClient, vmName := prepareForExec(t)
|
||||
sessionID := uuid.NewString()
|
||||
|
||||
wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
|
||||
Command: "sh -c 'echo replay-me'",
|
||||
WaitSeconds: 30,
|
||||
Session: sessionID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach})
|
||||
require.NoError(t, err)
|
||||
_ = wsConn.CloseNow()
|
||||
|
||||
time.Sleep(time.Second)
|
||||
|
||||
wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
|
||||
WaitSeconds: 30,
|
||||
Session: sessionID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer wsConn.CloseNow()
|
||||
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeHistory})
|
||||
require.NoError(t, err)
|
||||
|
||||
frames := readFramesUntilExit(t, wsConn)
|
||||
require.Equal(t, "replay-me\n", string(framesByType(frames, execstream.FrameTypeStdout)[0].Data))
|
||||
require.EqualValues(t, 0, framesByType(frames, execstream.FrameTypeExit)[0].Exit.Code)
|
||||
}
|
||||
|
||||
func TestVMExecSessionReplayPreservesStreams(t *testing.T) {
|
||||
devClient, vmName := prepareForExec(t)
|
||||
sessionID := uuid.NewString()
|
||||
|
||||
wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
|
||||
Command: "sh -c 'echo out1; sleep 1; echo err1 >&2; sleep 1; echo out2; sleep 1; echo err2 >&2'",
|
||||
WaitSeconds: 30,
|
||||
Session: sessionID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach})
|
||||
require.NoError(t, err)
|
||||
_ = wsConn.CloseNow()
|
||||
|
||||
time.Sleep(4 * time.Second)
|
||||
|
||||
wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
|
||||
WaitSeconds: 30,
|
||||
Session: sessionID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer wsConn.CloseNow()
|
||||
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeHistory})
|
||||
require.NoError(t, err)
|
||||
|
||||
frames := readFramesUntilExit(t, wsConn)
|
||||
require.Equal(t, []execstream.FrameType{
|
||||
execstream.FrameTypeStdout,
|
||||
execstream.FrameTypeStderr,
|
||||
execstream.FrameTypeStdout,
|
||||
execstream.FrameTypeStderr,
|
||||
execstream.FrameTypeExit,
|
||||
}, frameTypes(frames))
|
||||
require.Equal(t, "out1\n", string(frames[0].Data))
|
||||
require.Equal(t, "err1\n", string(frames[1].Data))
|
||||
require.Equal(t, "out2\n", string(frames[2].Data))
|
||||
require.Equal(t, "err2\n", string(frames[3].Data))
|
||||
}
|
||||
|
||||
func TestVMExecSessionStdinSurvivesReconnect(t *testing.T) {
|
||||
devClient, vmName := prepareForExec(t)
|
||||
sessionID := uuid.NewString()
|
||||
|
||||
wsConn, err := devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
|
||||
Command: "/bin/cat",
|
||||
Stdin: true,
|
||||
WaitSeconds: 30,
|
||||
Session: sessionID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{
|
||||
Type: execstream.FrameTypeStdin,
|
||||
Data: []byte("one\n"),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeDetach})
|
||||
require.NoError(t, err)
|
||||
_ = wsConn.CloseNow()
|
||||
|
||||
wsConn, err = devClient.VMs().ExecSession(t.Context(), vmName, client.ExecSessionOptions{
|
||||
WaitSeconds: 30,
|
||||
Session: sessionID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer wsConn.CloseNow()
|
||||
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{
|
||||
Type: execstream.FrameTypeStdin,
|
||||
Data: []byte("two\n"),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{
|
||||
Type: execstream.FrameTypeStdin,
|
||||
Data: []byte{},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
err = execstream.WriteFrame(t.Context(), wsConn, &execstream.Frame{Type: execstream.FrameTypeHistory})
|
||||
require.NoError(t, err)
|
||||
|
||||
frames := readFramesUntilExit(t, wsConn)
|
||||
stdoutFrames := framesByType(frames, execstream.FrameTypeStdout)
|
||||
require.Len(t, stdoutFrames, 2)
|
||||
require.Equal(t, "one\n", string(stdoutFrames[0].Data))
|
||||
require.Equal(t, "two\n", string(stdoutFrames[1].Data))
|
||||
require.EqualValues(t, 0, framesByType(frames, execstream.FrameTypeExit)[0].Exit.Code)
|
||||
}
|
||||
|
||||
func prepareForExec(t *testing.T) (*client.Client, string) {
|
||||
devClient, _, _ := devcontroller.StartIntegrationTestEnvironment(t)
|
||||
|
||||
|
|
@ -164,3 +325,43 @@ func readFrame(t *testing.T, wsConn *websocket.Conn) *execstream.Frame {
|
|||
|
||||
return &frame
|
||||
}
|
||||
|
||||
func readFramesUntilExit(t *testing.T, wsConn *websocket.Conn) []*execstream.Frame {
|
||||
t.Helper()
|
||||
|
||||
var frames []*execstream.Frame
|
||||
|
||||
for {
|
||||
frame := readFrame(t, wsConn)
|
||||
if frame.Type == execstream.FrameTypeNoMoreHistory {
|
||||
continue
|
||||
}
|
||||
|
||||
frames = append(frames, frame)
|
||||
if frame.Type == execstream.FrameTypeExit {
|
||||
return frames
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func framesByType(frames []*execstream.Frame, frameType execstream.FrameType) []*execstream.Frame {
|
||||
var result []*execstream.Frame
|
||||
|
||||
for _, frame := range frames {
|
||||
if frame.Type == frameType {
|
||||
result = append(result, frame)
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func frameTypes(frames []*execstream.Frame) []execstream.FrameType {
|
||||
var result []execstream.FrameType
|
||||
|
||||
for _, frame := range frames {
|
||||
result = append(result, frame.Type)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
|
|
|||
|
|
@ -47,6 +47,13 @@ type EventsPageOptions struct {
|
|||
Cursor string
|
||||
}
|
||||
|
||||
type ExecSessionOptions struct {
|
||||
Command string
|
||||
Stdin bool
|
||||
WaitSeconds uint16
|
||||
Session string
|
||||
}
|
||||
|
||||
func (service *VMsService) Create(ctx context.Context, vm *v1.VM) error {
|
||||
err := service.client.request(ctx, http.MethodPost, "vms",
|
||||
vm, nil, nil)
|
||||
|
|
@ -164,14 +171,35 @@ func (service *VMsService) Exec(
|
|||
stdin bool,
|
||||
waitSeconds uint16,
|
||||
) (*websocket.Conn, error) {
|
||||
return service.client.wsRequestRaw(ctx, fmt.Sprintf("vms/%s/exec", url.PathEscape(name)),
|
||||
map[string]string{
|
||||
"command": command,
|
||||
"stdin": strconv.FormatBool(stdin),
|
||||
"wait": strconv.FormatUint(uint64(waitSeconds), 10),
|
||||
return service.ExecSession(ctx, name, ExecSessionOptions{
|
||||
Command: command,
|
||||
Stdin: stdin,
|
||||
WaitSeconds: waitSeconds,
|
||||
})
|
||||
}
|
||||
|
||||
func (service *VMsService) ExecSession(
|
||||
ctx context.Context,
|
||||
name string,
|
||||
options ExecSessionOptions,
|
||||
) (*websocket.Conn, error) {
|
||||
params := map[string]string{
|
||||
"wait": strconv.FormatUint(uint64(options.WaitSeconds), 10),
|
||||
}
|
||||
if options.Command != "" {
|
||||
params["command"] = options.Command
|
||||
}
|
||||
if options.Stdin {
|
||||
params["stdin"] = strconv.FormatBool(true)
|
||||
}
|
||||
if options.Session != "" {
|
||||
params["session"] = options.Session
|
||||
}
|
||||
|
||||
return service.client.wsRequestRaw(ctx, fmt.Sprintf("vms/%s/exec", url.PathEscape(name)),
|
||||
params)
|
||||
}
|
||||
|
||||
func (service *VMsService) IP(ctx context.Context, name string, waitSeconds uint16) (string, error) {
|
||||
result := struct {
|
||||
IP string `json:"ip"`
|
||||
|
|
|
|||
|
|
@ -0,0 +1,40 @@
|
|||
package client
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/coder/websocket"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestExecSessionBuildsReconnectableQuery(t *testing.T) {
|
||||
var query map[string][]string
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
query = request.URL.Query()
|
||||
|
||||
conn, err := websocket.Accept(writer, request, nil)
|
||||
require.NoError(t, err)
|
||||
defer conn.CloseNow()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
devClient, err := New(WithAddress(server.URL))
|
||||
require.NoError(t, err)
|
||||
|
||||
conn, err := devClient.VMs().ExecSession(t.Context(), "vm", ExecSessionOptions{
|
||||
Command: "echo hello",
|
||||
Stdin: true,
|
||||
WaitSeconds: 7,
|
||||
Session: "resume-me",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer conn.CloseNow()
|
||||
|
||||
require.Equal(t, []string{"echo hello"}, query["command"])
|
||||
require.Equal(t, []string{"true"}, query["stdin"])
|
||||
require.Equal(t, []string{"7"}, query["wait"])
|
||||
require.Equal(t, []string{"resume-me"}, query["session"])
|
||||
}
|
||||
Loading…
Reference in New Issue