Add exec env support

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")
+}