fix(dstack-util): escape backslash in the systemd env file

dstack/dstack-util/src/parse_env_file.rs

@@ -7,29 +7,23 @@ use serde::Deserialize;
 use std::collections::{BTreeMap, BTreeSet};
 use tracing::warn;
 
+/// Escape a value for systemd's `EnvironmentFile=` parser. A newline cannot be
+/// represented there, so it is written as a literal `\n`.
 fn escape_value(v: &str) -> String {
-    let mut needs_quotes = false;
+    let needs_quotes = v.contains(|c: char| c.is_whitespace() || r#"|&;<>()$`\"'"#.contains(c));
     let mut escaped = String::with_capacity(v.len());
-
-    // Check if we need quotes (spaces or special chars)
-    if v.chars().any(|c| " \t|&;<>()$`\\\"'\n".contains(c)) {
-        needs_quotes = true;
-    }
-
-    // Escape special characters
     for c in v.chars() {
         match c {
-            '\n' => escaped.push_str("\\n"),
-            '"' => escaped.push_str("\\\""),
-            '$' => escaped.push_str("\\$"),
-            '`' => escaped.push_str("\\`"),
+            '\n' => escaped.push_str(r"\n"),
+            '\\' | '"' | '$' | '`' => {
+                escaped.push('\\');
+                escaped.push(c);
+            }
             _ => escaped.push(c),
         }
     }
-
-    // Wrap in quotes if needed
     if needs_quotes {
-        format!("\"{}\"", escaped)
+        format!(r#""{escaped}""#)
     } else {
         escaped
     }
@@ -96,7 +90,12 @@ pub fn convert_env_to_str(parsed_env: &BTreeMap<String, String>) -> String {
     #[allow(clippy::format_collect)]
     parsed_env
         .iter()
-        .map(|(key, value)| format!("{}={}\n", key, escape_value(value)))
+        .map(|(key, value)| {
+            if value.contains('\n') {
+                warn!("env var {key} contains a newline, delivering it as a literal \\n");
+            }
+            format!("{}={}\n", key, escape_value(value))
+        })
         .collect()
 }
 
@@ -107,10 +106,70 @@ mod tests {
     #[test]
     fn test_escape_value() {
         assert_eq!(escape_value("simple"), "simple");
-        assert_eq!(escape_value("hello world"), "\"hello world\"");
-        assert_eq!(escape_value("say \"hello\""), "\"say \\\"hello\\\"\"");
-        assert_eq!(escape_value("line1\nline2"), "\"line1\\nline2\"");
-        assert_eq!(escape_value("price=$100"), "\"price=\\$100\"");
-        assert_eq!(escape_value("command=`date`"), "\"command=\\`date\\`\"");
+        assert_eq!(escape_value("hello world"), r#""hello world""#);
+        assert_eq!(escape_value(r#"say "hello""#), r#""say \"hello\"""#);
+        assert_eq!(escape_value("line1\nline2"), r#""line1\nline2""#);
+        assert_eq!(escape_value("price=$100"), r#""price=\$100""#);
+        assert_eq!(escape_value("command=`date`"), r#""command=\`date\`""#);
+        assert_eq!(escape_value(r"trail\"), r#""trail\\""#);
+        assert_eq!(escape_value("cr\r"), "\"cr\r\"");
+    }
+
+    /// Round-trip through systemd's own parser; skipped without a user manager.
+    #[test]
+    fn systemd_round_trip() {
+        use std::process::Command;
+
+        let user_manager = Command::new("systemctl")
+            .args(["--user", "show-environment"])
+            .output()
+            .is_ok_and(|o| o.status.success());
+        if !user_manager {
+            eprintln!("skipping: no systemd user manager");
+            return;
+        }
+        let cases = [
+            "",
+            "a b",
+            "tab\there",
+            r#"say "hi""#,
+            "$HOME",
+            "`date`",
+            r"back\slash",
+            r"trail\",
+            r"\",
+            r#"a\"b"#,
+            r"a\$b",
+            r"a\`b",
+            r"literal\n",
+            "cr\rlf",
+            "it's",
+            "a;b|c&d",
+            "#x",
+            "ünï",
+        ];
+        let env: BTreeMap<_, _> = cases
+            .iter()
+            .enumerate()
+            .map(|(i, v)| (format!("DSTACK_ENV_CASE_{i}"), v.to_string()))
+            .collect();
+        let dir = tempfile::tempdir().expect("failed to create temp dir");
+        let file = dir.path().join("env");
+        std::fs::write(&file, convert_env_to_str(&env)).expect("failed to write env file");
+        let output = Command::new("systemd-run")
+            .args(["--user", "--wait", "--collect", "--quiet", "--pipe"])
+            .arg(format!("--property=EnvironmentFile={}", file.display()))
+            .args(["env", "-0"])
+            .output()
+            .expect("failed to run systemd-run");
+        assert!(output.status.success(), "systemd-run failed");
+        let stdout = String::from_utf8(output.stdout).expect("env output is not utf-8");
+        let delivered: BTreeMap<_, _> = stdout
+            .split('\0')
+            .filter_map(|kv| kv.split_once('='))
+            .collect();
+        for (key, value) in &env {
+            assert_eq!(delivered.get(key.as_str()), Some(&value.as_str()), "{key}");
+        }
     }
 }