Merge pull request 'feat: runtime env var forwarding via SSH SetEnv' (#6) from worktree-env-forwarding into main

Reviewed-on: https://forgejo.tail9a847c.ts.net/sh/pixels/pulls/6

Author: deevus

cmd/console.go

@@ -72,5 +72,5 @@ func runConsole(cmd *cobra.Command, args []string) error {
 	}
 
 	// Console replaces the process — does not return on success.
-	return ssh.Console(ip, cfg.SSH.User, cfg.SSH.Key)
+	return ssh.Console(ip, cfg.SSH.User, cfg.SSH.Key, cfg.EnvForward)
 }

cmd/create.go

@@ -297,7 +297,7 @@ func runCreate(cmd *cobra.Command, args []string) error {
 				go func() {
 					defer close(done)
 					ssh.Exec(journalCtx, ip, "root", cfg.SSH.Key,
-						[]string{"journalctl", "-fu", "pixels-devtools", "--no-pager", "-o", "cat"})
+						[]string{"journalctl", "-fu", "pixels-devtools", "--no-pager", "-o", "cat"}, nil)
 				}()
 			}
 
@@ -309,7 +309,7 @@ func runCreate(cmd *cobra.Command, args []string) error {
 				<-done
 			}
 		}
-		return ssh.Console(ip, cfg.SSH.User, cfg.SSH.Key)
+		return ssh.Console(ip, cfg.SSH.User, cfg.SSH.Key, cfg.EnvForward)
 	}
 
 	return nil

cmd/exec.go

@@ -65,7 +65,7 @@ func runExec(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	exitCode, err := ssh.Exec(ctx, ip, cfg.SSH.User, cfg.SSH.Key, command)
+	exitCode, err := ssh.Exec(ctx, ip, cfg.SSH.User, cfg.SSH.Key, command, cfg.EnvForward)
 	if err != nil {
 		return err
 	}

cmd/network.go

@@ -93,7 +93,7 @@ func resolveNetworkContext(cmd *cobra.Command, name string) (*networkContext, er
 
 // sshAsRoot runs a command on the container as root via SSH.
 func sshAsRoot(cmd *cobra.Command, ip string, command []string) (int, error) {
-	return ssh.Exec(cmd.Context(), ip, "root", cfg.SSH.Key, command)
+	return ssh.Exec(cmd.Context(), ip, "root", cfg.SSH.Key, command, nil)
 }
 
 func runNetworkShow(cmd *cobra.Command, args []string) error {

internal/config/config.go

@@ -11,13 +11,17 @@ import (
 )
 
 type Config struct {
-	TrueNAS    TrueNAS           `toml:"truenas"`
-	Defaults   Defaults          `toml:"defaults"`
-	SSH        SSH               `toml:"ssh"`
-	Checkpoint Checkpoint        `toml:"checkpoint"`
-	Provision  Provision         `toml:"provision"`
-	Network    Network           `toml:"network"`
-	Env        map[string]string `toml:"env"`
+	TrueNAS    TrueNAS        `toml:"truenas"`
+	Defaults   Defaults       `toml:"defaults"`
+	SSH        SSH            `toml:"ssh"`
+	Checkpoint Checkpoint     `toml:"checkpoint"`
+	Provision  Provision      `toml:"provision"`
+	Network    Network        `toml:"network"`
+	RawEnv     map[string]any `toml:"env"`
+
+	// Resolved env vars (not from TOML directly).
+	Env        map[string]string `toml:"-"` // image vars → /etc/environment
+	EnvForward map[string]string `toml:"-"` // session vars → SSH SetEnv
 }
 
 type TrueNAS struct {
@@ -109,13 +113,63 @@ func Load() (*Config, error) {
 
 	cfg.SSH.Key = expandHome(cfg.SSH.Key)
 
-	for k, v := range cfg.Env {
-		cfg.Env[k] = os.ExpandEnv(v)
+	if err := resolveEnv(cfg); err != nil {
+		return nil, err
 	}
 
 	return cfg, nil
 }
 
+// resolveEnv splits RawEnv entries into image vars (Env) and session vars (EnvForward).
+//
+// Supported forms:
+//
+//	KEY = "value"                          → image var (with $VAR expansion)
+//	KEY = { value = "x" }                 → image var (with $VAR expansion)
+//	KEY = { forward = true }              → session var (from host env, skip if unset)
+//	KEY = { value = "x", session_only = true } → session var (literal, with $VAR expansion)
+func resolveEnv(cfg *Config) error {
+	if len(cfg.RawEnv) == 0 {
+		return nil
+	}
+
+	cfg.Env = make(map[string]string)
+	cfg.EnvForward = make(map[string]string)
+
+	for k, raw := range cfg.RawEnv {
+		switch v := raw.(type) {
+		case string:
+			cfg.Env[k] = os.ExpandEnv(v)
+		case map[string]any:
+			forward, _ := v["forward"].(bool)
+			sessionOnly, _ := v["session_only"].(bool)
+			value, _ := v["value"].(string)
+
+			switch {
+			case forward:
+				if hostVal, ok := os.LookupEnv(k); ok {
+					cfg.EnvForward[k] = hostVal
+				}
+			case sessionOnly && value != "":
+				cfg.EnvForward[k] = os.ExpandEnv(value)
+			case value != "":
+				cfg.Env[k] = os.ExpandEnv(value)
+			}
+		default:
+			return fmt.Errorf("env %q: unsupported type %T", k, raw)
+		}
+	}
+
+	if len(cfg.Env) == 0 {
+		cfg.Env = nil
+	}
+	if len(cfg.EnvForward) == 0 {
+		cfg.EnvForward = nil
+	}
+
+	return nil
+}
+
 func configPath() string {
 	if dir := os.Getenv("XDG_CONFIG_HOME"); dir != "" {
 		return filepath.Join(dir, "pixels", "config.toml")

internal/config/config_test.go

@@ -288,6 +288,113 @@ LITERAL = "no-expansion-here"
 	}
 }
 
+func TestEnvForward(t *testing.T) {
+	dir := t.TempDir()
+	t.Setenv("XDG_CONFIG_HOME", dir)
+
+	cfgDir := filepath.Join(dir, "pixels")
+	if err := os.MkdirAll(cfgDir, 0o755); err != nil {
+		t.Fatal(err)
+	}
+
+	content := `
+[env]
+IMAGE_VAR = "baked-in"
+EXPLICIT_IMAGE = { value = "explicit" }
+FORWARDED = { forward = true }
+FORWARDED_UNSET = { forward = true }
+SESSION_LITERAL = { value = "session-val", session_only = true }
+EXPANDED_IMAGE = { value = "$PIXELS_TEST_EXPAND" }
+`
+	if err := os.WriteFile(filepath.Join(cfgDir, "config.toml"), []byte(content), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	t.Setenv("FORWARDED", "host-value")
+	// FORWARDED_UNSET intentionally not set.
+	t.Setenv("PIXELS_TEST_EXPAND", "expanded-val")
+
+	cfg, err := Load()
+	if err != nil {
+		t.Fatalf("Load() error: %v", err)
+	}
+
+	// Image vars (cfg.Env).
+	if cfg.Env["IMAGE_VAR"] != "baked-in" {
+		t.Errorf("Env[IMAGE_VAR] = %q, want %q", cfg.Env["IMAGE_VAR"], "baked-in")
+	}
+	if cfg.Env["EXPLICIT_IMAGE"] != "explicit" {
+		t.Errorf("Env[EXPLICIT_IMAGE] = %q, want %q", cfg.Env["EXPLICIT_IMAGE"], "explicit")
+	}
+	if cfg.Env["EXPANDED_IMAGE"] != "expanded-val" {
+		t.Errorf("Env[EXPANDED_IMAGE] = %q, want %q", cfg.Env["EXPANDED_IMAGE"], "expanded-val")
+	}
+
+	// Image vars should NOT include session entries.
+	if _, ok := cfg.Env["FORWARDED"]; ok {
+		t.Error("Env should not contain FORWARDED (it's a session var)")
+	}
+	if _, ok := cfg.Env["SESSION_LITERAL"]; ok {
+		t.Error("Env should not contain SESSION_LITERAL (it's a session var)")
+	}
+
+	// Session vars (cfg.EnvForward).
+	if cfg.EnvForward["FORWARDED"] != "host-value" {
+		t.Errorf("EnvForward[FORWARDED] = %q, want %q", cfg.EnvForward["FORWARDED"], "host-value")
+	}
+	if _, ok := cfg.EnvForward["FORWARDED_UNSET"]; ok {
+		t.Error("EnvForward should not contain FORWARDED_UNSET (not set on host)")
+	}
+	if cfg.EnvForward["SESSION_LITERAL"] != "session-val" {
+		t.Errorf("EnvForward[SESSION_LITERAL] = %q, want %q", cfg.EnvForward["SESSION_LITERAL"], "session-val")
+	}
+
+	// Session vars should NOT include image entries.
+	if _, ok := cfg.EnvForward["IMAGE_VAR"]; ok {
+		t.Error("EnvForward should not contain IMAGE_VAR (it's an image var)")
+	}
+}
+
+func TestEnvForwardNilWhenEmpty(t *testing.T) {
+	t.Setenv("XDG_CONFIG_HOME", t.TempDir())
+
+	cfg, err := Load()
+	if err != nil {
+		t.Fatalf("Load() error: %v", err)
+	}
+
+	if cfg.Env != nil {
+		t.Errorf("Env = %v, want nil (no env configured)", cfg.Env)
+	}
+	if cfg.EnvForward != nil {
+		t.Errorf("EnvForward = %v, want nil (no env configured)", cfg.EnvForward)
+	}
+}
+
+func TestEnvUnsupportedType(t *testing.T) {
+	dir := t.TempDir()
+	t.Setenv("XDG_CONFIG_HOME", dir)
+
+	cfgDir := filepath.Join(dir, "pixels")
+	if err := os.MkdirAll(cfgDir, 0o755); err != nil {
+		t.Fatal(err)
+	}
+
+	// An integer value in [env] is not a valid env entry.
+	content := `
+[env]
+BAD = 42
+`
+	if err := os.WriteFile(filepath.Join(cfgDir, "config.toml"), []byte(content), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	_, err := Load()
+	if err == nil {
+		t.Fatal("expected error for unsupported env type, got nil")
+	}
+}
+
 func TestNetworkDefaults(t *testing.T) {
 	t.Setenv("XDG_CONFIG_HOME", t.TempDir())
 	for _, key := range []string{

internal/ssh/console_unix.go

@@ -10,11 +10,12 @@ import (
 )
 
 // Console replaces the current process with an interactive SSH session.
-func Console(host, user, keyPath string) error {
+// If env is non-nil, the entries are forwarded via SSH SetEnv.
+func Console(host, user, keyPath string, env map[string]string) error {
 	sshBin, err := exec.LookPath("ssh")
 	if err != nil {
 		return fmt.Errorf("ssh binary not found: %w", err)
 	}
-	args := append([]string{"ssh"}, sshArgs(host, user, keyPath)...)
+	args := append([]string{"ssh"}, sshArgs(host, user, keyPath, env)...)
 	return syscall.Exec(sshBin, args, os.Environ())
 }

internal/ssh/console_windows.go

@@ -9,12 +9,13 @@ import (
 )
 
 // Console runs an interactive SSH session as a child process.
-func Console(host, user, keyPath string) error {
+// If env is non-nil, the entries are forwarded via SSH SetEnv.
+func Console(host, user, keyPath string, env map[string]string) error {
 	sshBin, err := exec.LookPath("ssh")
 	if err != nil {
 		return fmt.Errorf("ssh binary not found: %w", err)
 	}
-	cmd := exec.Command(sshBin, sshArgs(host, user, keyPath)...)
+	cmd := exec.Command(sshBin, sshArgs(host, user, keyPath, env)...)
 	cmd.Stdin = os.Stdin
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr

internal/ssh/ssh.go

@@ -8,6 +8,8 @@ import (
 	"net"
 	"os"
 	"os/exec"
+	"sort"
+	"strings"
 	"time"
 
 	"github.com/deevus/pixels/internal/retry"
@@ -46,8 +48,9 @@ func WaitReady(ctx context.Context, host string, timeout time.Duration, log io.W
 }
 
 // Exec runs a command on the remote host via SSH and returns its exit code.
-func Exec(ctx context.Context, host, user, keyPath string, command []string) (int, error) {
-	args := append(sshArgs(host, user, keyPath), command...)
+// If env is non-nil, the entries are forwarded via SSH SetEnv.
+func Exec(ctx context.Context, host, user, keyPath string, command []string, env map[string]string) (int, error) {
+	args := append(sshArgs(host, user, keyPath, env), command...)
 	cmd := exec.CommandContext(ctx, "ssh", args...)
 	cmd.Stdin = os.Stdin
 	cmd.Stdout = os.Stdout
@@ -65,7 +68,7 @@ func Exec(ctx context.Context, host, user, keyPath string, command []string) (in
 
 // Output runs a command on the remote host via SSH and returns its stdout.
 func Output(ctx context.Context, host, user, keyPath string, command []string) ([]byte, error) {
-	args := append(sshArgs(host, user, keyPath), command...)
+	args := append(sshArgs(host, user, keyPath, nil), command...)
 	cmd := exec.CommandContext(ctx, "ssh", args...)
 	cmd.Stderr = os.Stderr
 	return cmd.Output()
@@ -74,20 +77,20 @@ func Output(ctx context.Context, host, user, keyPath string, command []string) (
 // WaitProvisioned polls the remote host until /root/.devtools-provisioned exists.
 func WaitProvisioned(ctx context.Context, host, user, keyPath string, timeout time.Duration) error {
 	return retry.Poll(ctx, 2*time.Second, timeout, func(ctx context.Context) (bool, error) {
-		code, err := Exec(ctx, host, user, keyPath, []string{"sudo", "test", "-f", "/root/.devtools-provisioned"})
+		code, err := Exec(ctx, host, user, keyPath, []string{"sudo", "test", "-f", "/root/.devtools-provisioned"}, nil)
 		return err == nil && code == 0, nil
 	})
 }
 
 // TestAuth runs a quick SSH connection test (ssh ... true) to verify
 // key-based authentication works. Returns nil on success.
 func TestAuth(ctx context.Context, host, user, keyPath string) error {
-	args := append(sshArgs(host, user, keyPath), "true")
+	args := append(sshArgs(host, user, keyPath, nil), "true")
 	cmd := exec.CommandContext(ctx, "ssh", args...)
 	return cmd.Run()
 }
 
-func sshArgs(host, user, keyPath string) []string {
+func sshArgs(host, user, keyPath string, env map[string]string) []string {
 	args := []string{
 		"-o", "StrictHostKeyChecking=no",
 		"-o", "UserKnownHostsFile=" + os.DevNull,
@@ -97,6 +100,28 @@ func sshArgs(host, user, keyPath string) []string {
 	if keyPath != "" {
 		args = append(args, "-i", keyPath)
 	}
+
+	// Forward env vars via SSH protocol (requires AcceptEnv on server).
+	// All vars must be in a single SetEnv directive (multiple -o SetEnv
+	// flags don't stack in OpenSSH — only the first takes effect).
+	if len(env) > 0 {
+		keys := make([]string, 0, len(env))
+		for k := range env {
+			keys = append(keys, k)
+		}
+		sort.Strings(keys)
+
+		var setenv strings.Builder
+		setenv.WriteString("SetEnv=")
+		for i, k := range keys {
+			if i > 0 {
+				setenv.WriteByte(' ')
+			}
+			fmt.Fprintf(&setenv, "%s=%s", k, env[k])
+		}
+		args = append(args, "-o", setenv.String())
+	}
+
 	args = append(args, user+"@"+host)
 	return args
 }