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