Add exec env support

internal/controller/api_vms_exec.go

@@ -35,7 +35,15 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder {
 			NewErrorResponse("\"command\" parameter cannot be empty"))
 	}
 
-	stdin := ctx.Query("stdin") == "true"
+	interactive, err := parseExecInteractive(ctx)
+	if err != nil {
+		return responder.JSON(http.StatusBadRequest, NewErrorResponse("%v", err))
+	}
+
+	command, err = sshexec.CommandWithEnv(command, ctx.QueryMap("env"))
+	if err != nil {
+		return responder.JSON(http.StatusBadRequest, NewErrorResponse("%v", err))
+	}
 
 	waitRaw := ctx.DefaultQuery("wait", "10")
 	wait, err := strconv.ParseUint(waitRaw, 10, 16)
@@ -68,7 +76,7 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder {
 	defer portForwardConn.Close()
 
 	// Establish an SSH connection to a VM
-	exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), stdin)
+	exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), interactive)
 	if err != nil {
 		return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("failed to establish SSH connection to a VM: %v", err))
 	}
@@ -151,6 +159,38 @@ func (controller *Controller) execVM(ctx *gin.Context) responder.Responder {
 	}
 }
 
+func parseExecInteractive(ctx *gin.Context) (bool, error) {
+	interactiveRaw, interactivePresent := ctx.GetQuery("interactive")
+	stdinRaw, stdinPresent := ctx.GetQuery("stdin")
+
+	interactive := false
+	if interactivePresent {
+		parsed, err := strconv.ParseBool(interactiveRaw)
+		if err != nil {
+			return false, fmt.Errorf("\"interactive\" parameter must be a boolean")
+		}
+
+		interactive = parsed
+	}
+
+	if stdinPresent {
+		stdin, err := strconv.ParseBool(stdinRaw)
+		if err != nil {
+			return false, fmt.Errorf("\"stdin\" parameter must be a boolean")
+		}
+
+		if interactivePresent && stdin != interactive {
+			return false, fmt.Errorf("\"interactive\" and \"stdin\" parameters cannot conflict")
+		}
+
+		if !interactivePresent {
+			interactive = stdin
+		}
+	}
+
+	return interactive, nil
+}
+
 func (controller *Controller) readFrames(
 	ctx context.Context,
 	wsConn *websocket.Conn,

internal/controller/api_vms_exec_test.go

@@ -0,0 +1,74 @@
+//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 execQueryContext(query string) *gin.Context {
+	ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
+	ctx.Request = httptest.NewRequest(http.MethodGet, "/?"+query, nil)
+
+	return ctx
+}

internal/controller/sshexec/sshexec.go

@@ -6,13 +6,18 @@ 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 Exec struct {
 	sshClient  *ssh.Client
 	sshSession *ssh.Session
@@ -87,6 +92,43 @@ func (exec *Exec) Stdin() io.WriteCloser {
 	return exec.stdin
 }
 
+func CommandWithEnv(command string, env map[string]string) (string, error) {
+	if len(env) == 0 {
+		return command, nil
+	}
+
+	keys := make([]string, 0, len(env))
+	for key, value := range 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)
+
+	var builder strings.Builder
+	for _, key := range keys {
+		builder.WriteString("export ")
+		builder.WriteString(key)
+		builder.WriteByte('=')
+		builder.WriteString(shellQuote(env[key]))
+		builder.WriteByte('\n')
+	}
+	builder.WriteString(command)
+
+	return builder.String(), nil
+}
+
+func shellQuote(value string) string {
+	return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'"
+}
+
 func (exec *Exec) Run(
 	ctx context.Context,
 	command string,

internal/controller/sshexec/sshexec_test.go

@@ -2,6 +2,7 @@ package sshexec_test
 
 import (
 	"net"
+	"strings"
 	"testing"
 	"time"
 
@@ -24,3 +25,39 @@ func TestContextCancellationViaNetConnClose(t *testing.T) {
 	_, err := sshexec.New(clientConn, "doesn't", "matter", false)
 	require.Error(t, err)
 }
+
+func TestCommandWithEnvNoEnvLeavesCommandUnchanged(t *testing.T) {
+	command, err := sshexec.CommandWithEnv("echo hello", nil)
+	require.NoError(t, err)
+	require.Equal(t, "echo hello", command)
+}
+
+func TestCommandWithEnvSortsAndQuotes(t *testing.T) {
+	command, err := sshexec.CommandWithEnv("printf '%s|%s|%s' \"$GREETING\" \"$NAME\" \"$MULTILINE\"", map[string]string{
+		"NAME":      "O'Reilly",
+		"GREETING":  "hello $USER",
+		"MULTILINE": "line 1\nline 2",
+	})
+	require.NoError(t, err)
+	require.Equal(t, strings.Join([]string{
+		"export GREETING='hello $USER'",
+		"export MULTILINE='line 1",
+		"line 2'",
+		"export NAME='O'\\''Reilly'",
+		"printf '%s|%s|%s' \"$GREETING\" \"$NAME\" \"$MULTILINE\"",
+	}, "\n"), command)
+}
+
+func TestCommandWithEnvRejectsInvalidName(t *testing.T) {
+	_, err := sshexec.CommandWithEnv("echo hello", map[string]string{
+		"1INVALID": "value",
+	})
+	require.ErrorContains(t, err, "invalid environment variable name")
+}
+
+func TestCommandWithEnvRejectsNULValue(t *testing.T) {
+	_, err := sshexec.CommandWithEnv("echo hello", map[string]string{
+		"VALID": "bad\x00value",
+	})
+	require.ErrorContains(t, err, "contains NUL byte")
+}