[runner] Rework env variables exporting (#2535)

Replace `~/.ssh/environment` (processed by OpenSSH) with a shell script
exporting env variables. The shell script is sourced in `/etc/profile`,
which is only executed by login shell.

Advantages:

* multi-line variables support (e.g, `DSTACK_NODES_IPS`)
* no hard limit of 1000 variables
* shared by all users
* vars are not redefined by `pam_env` (`/etc/environment`)
* it's unlikely that `/tmp/dstack_profile` and `/etc/profile`
  are located on a volume, unlike `~/.ssh`; no need for cleanup

Tested supported shells:

* dash (Debian and derivatives, BusyBox distros, e.g., Alpine)
* bash as sh (Fedora and downstream distros)
* bash (almost every distro)

Not supported shells (do not use `/etc/profile`):

* zsh
* fish

Fixes: #2371

Author: un-def

runner/internal/executor/executor.go

@@ -310,17 +310,18 @@ func (ex *RunExecutor) execJob(ctx context.Context, jobLogFile io.Writer) error
 	// `env` interpolation feature is postponed to some future release
 	envMap.Update(ex.jobSpec.Env, false)
 
+	const profilePath = "/etc/profile"
+	const dstackProfilePath = "/tmp/dstack_profile"
+	if err := writeDstackProfile(envMap, dstackProfilePath); err != nil {
+		log.Warning(ctx, "failed to write dstack_profile", "path", dstackProfilePath, "err", err)
+	} else if err := includeDstackProfile(profilePath, dstackProfilePath); err != nil {
+		log.Warning(ctx, "failed to include dstack_profile", "path", profilePath, "err", err)
+	}
+
 	// As of 2024-11-29, ex.homeDir is always set to /root
 	rootSSHDir, err := prepareSSHDir(-1, -1, ex.homeDir)
 	if err != nil {
 		log.Warning(ctx, "failed to prepare ssh dir", "home", ex.homeDir, "err", err)
-	} else {
-		rootSSHEnvPath := filepath.Join(rootSSHDir, "environment")
-		restoreRootSSHEnv := backupFile(ctx, rootSSHEnvPath)
-		defer restoreRootSSHEnv(ctx)
-		if err := writeSSHEnvironment(envMap, -1, -1, rootSSHEnvPath); err != nil {
-			log.Warning(ctx, "failed to write SSH environment", "path", ex.homeDir, "err", err)
-		}
 	}
 	userSSHDir := ""
 	uid := -1
@@ -337,12 +338,6 @@ func (ex *RunExecutor) execJob(ctx context.Context, jobLogFile io.Writer) error
 			if err != nil {
 				log.Warning(ctx, "failed to prepare ssh dir", "home", homeDir, "err", err)
 			} else {
-				userSSHEnvPath := filepath.Join(userSSHDir, "environment")
-				restoreUserSSHEnv := backupFile(ctx, userSSHEnvPath)
-				defer restoreUserSSHEnv(ctx)
-				if err := writeSSHEnvironment(envMap, uid, gid, userSSHEnvPath); err != nil {
-					log.Warning(ctx, "failed to write SSH environment", "path", homeDir, "err", err)
-				}
 				rootSSHKeysPath := filepath.Join(rootSSHDir, "authorized_keys")
 				userSSHKeysPath := filepath.Join(userSSHDir, "authorized_keys")
 				restoreUserSSHKeys := backupFile(ctx, userSSHKeysPath)
@@ -676,53 +671,40 @@ func prepareSSHDir(uid int, gid int, homeDir string) (string, error) {
 	return sshDir, nil
 }
 
-func writeSSHEnvironment(env map[string]string, uid int, gid int, envPath string) error {
-	info, err := os.Stat(envPath)
-	if err == nil {
-		if info.IsDir() {
-			return fmt.Errorf("is a directory: %s", envPath)
-		}
-		if err = os.Chmod(envPath, 0o600); err != nil {
-			return err
-		}
-	} else if !errors.Is(err, os.ErrNotExist) {
-		return err
-	}
-
-	envFile, err := os.OpenFile(envPath, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0o600)
+func writeDstackProfile(env map[string]string, path string) error {
+	file, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
 	if err != nil {
 		return err
 	}
-	defer envFile.Close()
+	defer file.Close()
 	for key, value := range env {
 		switch key {
-		case "USER", "HOME", "SHELL", "PWD", "_":
-			continue
-		}
-		// sshd doesn't support multiline variable values in .ssh/environment:
-		// VAR1=line1
-		// line2
-		// line3
-		// VAR2=singleline
-		// leads to:
-		// Bad line 2 in /root/.ssh/environment
-		// Bad line 3 in /root/.ssh/environment
-		// Assuming a single trailing newline is not a big deal
-		value := strings.TrimSuffix(value, "\n")
-		// If there is any non-trailing newline, or more than one
-		// trailing newline, skip the variable
-		if strings.Contains(value, "\n") {
+		case "HOSTNAME", "USER", "HOME", "SHELL", "SHLVL", "PWD", "_":
 			continue
 		}
-		line := fmt.Sprintf("%s=%s\n", key, value)
-		if _, err = envFile.WriteString(line); err != nil {
+		line := fmt.Sprintf("export %s='%s'\n", key, strings.ReplaceAll(value, `'`, `'"'"'`))
+		if _, err = file.WriteString(line); err != nil {
 			return err
 		}
 	}
-	if err = os.Chown(envPath, uid, gid); err != nil {
+	if err = os.Chmod(path, 0o644); err != nil {
 		return err
 	}
+	return nil
+}
 
+func includeDstackProfile(profilePath string, dstackProfilePath string) error {
+	file, err := os.OpenFile(profilePath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+	if _, err = file.WriteString(fmt.Sprintf("\n. '%s'\n", dstackProfilePath)); err != nil {
+		return err
+	}
+	if err = os.Chmod(profilePath, 0o644); err != nil {
+		return err
+	}
 	return nil
 }
 

runner/internal/executor/executor_test.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"testing"
 	"time"
@@ -207,3 +208,22 @@ func makeCodeTar(t *testing.T, path string) {
 	}
 	require.NoError(t, tw.Close())
 }
+
+func TestWriteDstackProfile(t *testing.T) {
+	testCases := []string{
+		"",
+		"string 'with 'single' quotes",
+		"multi\nline\tstring",
+	}
+	tmp := t.TempDir()
+	path := tmp + "/dstack_profile"
+	script := fmt.Sprintf(`. '%s'; printf '%%s' "$VAR"`, path)
+	for _, value := range testCases {
+		env := map[string]string{"VAR": value}
+		writeDstackProfile(env, path)
+		cmd := exec.Command("/bin/sh", "-c", script)
+		out, err := cmd.Output()
+		assert.NoError(t, err)
+		assert.Equal(t, value, string(out))
+	}
+}

runner/internal/shim/docker.go

@@ -896,7 +896,6 @@ func getSSHShellCommands(openSSHPort int, publicSSHKey string) []string {
 		"chmod 700 ~/.ssh",
 		fmt.Sprintf("echo '%s' > ~/.ssh/authorized_keys", publicSSHKey),
 		"chmod 600 ~/.ssh/authorized_keys",
-		`if [ -f ~/.profile ]; then sed -ie '1s@^@export PATH="'"$PATH"':$PATH"\n\n@' ~/.profile; fi`,
 		// regenerate host keys
 		"rm -rf /etc/ssh/ssh_host_*",
 		"ssh-keygen -A > /dev/null",
@@ -914,7 +913,6 @@ func getSSHShellCommands(openSSHPort int, publicSSHKey string) []string {
 				" -o PidFile=none"+
 				" -o PasswordAuthentication=no"+
 				" -o AllowTcpForwarding=yes"+
-				" -o PermitUserEnvironment=yes"+
 				" -o ClientAliveInterval=30"+
 				" -o ClientAliveCountMax=4",
 			openSSHPort,

src/dstack/_internal/core/backends/base/compute.py

@@ -555,9 +555,7 @@ def get_gateway_user_data(authorized_key: str) -> str:
     )
 
 
-def get_docker_commands(
-    authorized_keys: List[str], fix_path_in_dot_profile: bool = True
-) -> List[str]:
+def get_docker_commands(authorized_keys: list[str]) -> list[str]:
     authorized_keys_content = "\n".join(authorized_keys).strip()
     commands = [
         # save and unset ld.so variables
@@ -581,9 +579,6 @@ def get_docker_commands(
         "chmod 700 ~/.ssh",
         f"echo '{authorized_keys_content}' > ~/.ssh/authorized_keys",
         "chmod 600 ~/.ssh/authorized_keys",
-        r"""if [ -f ~/.profile ]; then sed -ie '1s@^@export PATH="'"$PATH"':$PATH"\n\n@' ~/.profile; fi"""
-        if fix_path_in_dot_profile
-        else ":",
         # regenerate host keys
         "rm -rf /etc/ssh/ssh_host_*",
         "ssh-keygen -A > /dev/null",
@@ -601,7 +596,6 @@ def get_docker_commands(
             " -o PidFile=none"
             " -o PasswordAuthentication=no"
             " -o AllowTcpForwarding=yes"
-            " -o PermitUserEnvironment=yes"
             " -o ClientAliveInterval=30"
             " -o ClientAliveCountMax=4"
         ),

src/dstack/_internal/core/backends/runpod/compute.py

@@ -260,7 +260,7 @@ def _clean_stale_container_registry_auths(self) -> None:
 
 
 def _get_docker_args(authorized_keys: List[str]) -> str:
-    commands = get_docker_commands(authorized_keys, False)
+    commands = get_docker_commands(authorized_keys)
     command = " && ".join(commands)
     docker_args = {"cmd": [command], "entrypoint": ["/bin/sh", "-c"]}
     docker_args_escaped = json.dumps(json.dumps(docker_args)).strip('"')