substitute {chroot} marker in environment variables for remote execution

docs/notes/2.25.x.md

@@ -29,7 +29,7 @@ Pants now sends a `user-agent` header with every request to a remote store or a
 even when other headers are configured. If necessary, the user may override the user agent by specifying
 one in `remote_store_headers` or `remote_execution_headers`.
 
-Pants now supports the `{chroot}` replacement marker in program arguments for remote execution. (With local and Docker execution, the `{chroot}` marker is replaced with the absolute path of the sandbox directory if it appears in program arguments or environment variables. Pants will now replace `{chroot}` for program arguments with remote execution, but still does not yet support doing so in environment variables.  This requires `/bin/bash` to be available on the remote execution server.)
+Pants now supports the `{chroot}` replacement marker in remote execution contexts. (With local and Docker execution, the `{chroot}` marker is replaced with the absolute path of the sandbox directory if it appears in program arguments or environment variables. Pants will do the same as well in remote execution contexts. This requires `/bin/bash` to be available on the remote execution server.)
 
 ### New Options System
 

src/rust/engine/process_execution/remote/src/remote_tests.rs

@@ -724,7 +724,7 @@ async fn make_execute_request_with_append_only_caches() {
         input_root_digest: Some(
             (Digest::new(
                 Fingerprint::from_hex_string(
-                    "d6f65a7f56380de341fac43fcb4e2006fd0e716579325c4640ec11b7c040bb9c",
+                    "40ff253312fe7bb50acad02e81e3a16b0109eb934c37bb61c3735792b8bead5f",
                 )
                 .unwrap(),
                 178,
@@ -739,7 +739,7 @@ async fn make_execute_request_with_append_only_caches() {
         action_digest: Some(
             (&Digest::new(
                 Fingerprint::from_hex_string(
-                    "d516c6b3b9781777486025ebc87cb4810bc7088dc45092a70a592d8ca74f48d8",
+                    "798ac2aaf68b36d571fcb90719e8dafdf5929081c491f178ac7f4663b591183e",
                 )
                 .unwrap(),
                 146,
@@ -752,7 +752,7 @@ async fn make_execute_request_with_append_only_caches() {
 
     let want_input_root_digest = DirectoryDigest::from_persisted_digest(Digest::new(
         Fingerprint::from_hex_string(
-            "d6f65a7f56380de341fac43fcb4e2006fd0e716579325c4640ec11b7c040bb9c",
+            "40ff253312fe7bb50acad02e81e3a16b0109eb934c37bb61c3735792b8bead5f",
         )
         .unwrap(),
         178,

src/rust/engine/process_execution/src/lib.rs

@@ -1023,12 +1023,12 @@ fn quote_path(path: &Path) -> Result<String, String> {
 
 /// Check supplied parameters from an execution request and create a wrapper script if a wrapper script
 /// is needed.
-// TODO: Support `{chroot}` replacement in environment variables and not just program arguments.
 fn maybe_make_wrapper_script(
     caches: &BTreeMap<CacheName, RelativePath>,
     append_only_caches_base_path: Option<&str>,
     working_directory: Option<&str>,
     sandbox_root_token: Option<&str>,
+    env_vars_to_substitute: Vec<&str>,
 ) -> Result<Option<String>, String> {
     let caches_fragment = match append_only_caches_base_path {
         Some(base_path) if !caches.is_empty() => {
@@ -1088,14 +1088,26 @@ fn maybe_make_wrapper_script(
 
     // Generate code to replace `{chroot}` markers if any are present in the command.
     let sandbox_root_fragment = if let Some(token) = sandbox_root_token {
-        let fragment = format!(
+        let mut fragment = format!(
             concat!(
                 "sandbox_root=\"$(/bin/pwd)\"\n",
                 "args=(\"${{@//{0}/$sandbox_root}}\")\n",
                 "set -- \"${{args[@]}}\"\n",
             ),
             token
         );
+        if !env_vars_to_substitute.is_empty() {
+            let env_vars_str = shlex::try_join(env_vars_to_substitute)
+                .map_err(|e| format!("Failed to shell quote environment names: {e}"))?;
+            writeln!(&mut fragment, "for env_var in {env_vars_str}; do")
+                .map_err(|err| format!("write! failed: {err}"))?;
+            writeln!(
+                &mut fragment,
+                "  eval \"export $env_var=\\${{$env_var//{token}/$sandbox_root}}\""
+            )
+            .map_err(|err| format!("write! failed: {err}"))?;
+            writeln!(&mut fragment, "done").map_err(|err| format!("write! failed: {err}"))?;
+        }
         Some(fragment)
     } else {
         None
@@ -1115,15 +1127,15 @@ fn maybe_make_wrapper_script(
         writeln!(&mut script, "#!{shebang}").map_err(|err| format!("write! failed: {err:?}"))?;
 
         if let Some(sandbox_root_fragment) = sandbox_root_fragment {
-            writeln!(&mut script, "{sandbox_root_fragment}")
+            write!(&mut script, "{sandbox_root_fragment}")
                 .map_err(|err| format!("write! failed: {err:?}"))?;
         }
         if let Some(caches_fragment) = caches_fragment {
-            writeln!(&mut script, "{caches_fragment}")
+            write!(&mut script, "{caches_fragment}")
                 .map_err(|err| format!("write! failed: {err:?}"))?;
         }
         if let Some(chdir_fragment) = chdir_fragment {
-            writeln!(&mut script, "{chdir_fragment}")
+            write!(&mut script, "{chdir_fragment}")
                 .map_err(|err| format!("write! failed: {err:?}"))?;
         }
 
@@ -1144,20 +1156,41 @@ pub async fn make_execute_request(
 ) -> Result<EntireExecuteRequest, String> {
     const WRAPPER_SCRIPT: &str = "./__pants_wrapper__";
     const SANDBOX_ROOT_TOKEN: &str = "__PANTS_SANDBOX_ROOT__";
+    const CHROOT_MARKER: &str = "{chroot}";
 
     // Implement append-only caches by running a wrapper script before the actual program
     // to be invoked in the remote environment.
     let wrapper_script_content_opt = {
-        let sandbox_root_token = req
-            .argv
+        let args_have_chroot_marker = req.argv.iter().any(|arg| arg.contains(CHROOT_MARKER));
+
+        let mut env_vars_with_chroot_marker = req
+            .env
+            .iter()
+            .filter_map(|(key, value)| {
+                if value.contains(CHROOT_MARKER) {
+                    Some(key.to_owned())
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>();
+        env_vars_with_chroot_marker.sort();
+
+        let env_var_with_chroot_marker_refs = env_vars_with_chroot_marker
             .iter()
-            .any(|arg| arg.contains("{chroot}"))
-            .then_some(SANDBOX_ROOT_TOKEN);
+            .map(|s| s.as_str())
+            .collect::<Vec<_>>();
+
+        let sandbox_root_token = (args_have_chroot_marker
+            || !env_vars_with_chroot_marker.is_empty())
+        .then_some(SANDBOX_ROOT_TOKEN);
+
         maybe_make_wrapper_script(
             &req.append_only_caches,
             append_only_caches_base_path,
             req.working_directory.as_ref().and_then(|p| p.to_str()),
             sandbox_root_token,
+            env_var_with_chroot_marker_refs,
         )?
     };
     let wrapper_script_digest_opt = if let Some(script) = wrapper_script_content_opt {
@@ -1181,7 +1214,7 @@ pub async fn make_execute_request(
             args.extend(
                 req.argv
                     .iter()
-                    .map(|orig_arg| orig_arg.replace("{chroot}", SANDBOX_ROOT_TOKEN)),
+                    .map(|orig_arg| orig_arg.replace(CHROOT_MARKER, SANDBOX_ROOT_TOKEN)),
             );
             args
         }
@@ -1207,7 +1240,7 @@ pub async fn make_execute_request(
             .environment_variables
             .push(remexec::command::EnvironmentVariable {
                 name: name.to_string(),
-                value: value.to_string(),
+                value: value.replace(CHROOT_MARKER, SANDBOX_ROOT_TOKEN),
             });
     }
 

src/rust/engine/process_execution/src/tests.rs

@@ -182,6 +182,7 @@ async fn wrapper_script_supports_append_only_caches() {
         dummy_caches_base_path.path().to_str(),
         Some(SUBDIR_NAME),
         None,
+        vec![],
     )
     .unwrap()
     .unwrap();
@@ -242,7 +243,7 @@ async fn wrapper_script_supports_append_only_caches() {
 async fn wrapper_script_supports_sandbox_root_replacements_in_args() {
     let caches = BTreeMap::new();
 
-    let script_content = maybe_make_wrapper_script(&caches, None, None, Some("__ROOT__"))
+    let script_content = maybe_make_wrapper_script(&caches, None, None, Some("__ROOT__"), vec![])
         .unwrap()
         .unwrap();
 
@@ -281,3 +282,55 @@ async fn wrapper_script_supports_sandbox_root_replacements_in_args() {
         .unwrap();
     assert_eq!(content, "xyzzy\n");
 }
+
+#[tokio::test]
+async fn wrapper_script_supports_sandbox_root_replacements_in_environmenbt() {
+    let caches = BTreeMap::new();
+
+    let script_content = maybe_make_wrapper_script(
+        &caches,
+        None,
+        None,
+        Some("__ROOT__"),
+        vec!["TEST_FILE_PATH"],
+    )
+    .unwrap()
+    .unwrap();
+
+    let dummy_sandbox_path = TempDir::new().unwrap();
+    let script_path = dummy_sandbox_path.path().join("wrapper");
+    tokio::fs::write(&script_path, script_content.as_bytes())
+        .await
+        .unwrap();
+    tokio::fs::set_permissions(&script_path, Permissions::from_mode(0o755))
+        .await
+        .unwrap();
+
+    let mut cmd = tokio::process::Command::new("./wrapper");
+    cmd.args(&[
+        "/bin/sh",
+        "-c",
+        "echo xyzzy > $TEST_FILE_PATH && echo $TEST_FILE_PATH",
+    ]);
+    cmd.env("TEST_FILE_PATH", "__ROOT__/foo.txt");
+    cmd.current_dir(dummy_sandbox_path.path());
+    cmd.stdin(Stdio::null());
+    cmd.stdout(Stdio::piped());
+    cmd.stderr(Stdio::piped());
+
+    let child = cmd.spawn().unwrap();
+    let output = child.wait_with_output().await.unwrap();
+    if output.status.code() != Some(0) {
+        let stdout = String::from_utf8_lossy(&output.stdout);
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        println!("stdout:{}\n\nstderr: {}", &stdout, &stderr);
+        panic!("Wrapper script failed to run: {}", output.status);
+    }
+
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    assert!(!stdout.contains("__ROOT__"));
+    let content = tokio::fs::read_to_string(Path::new(stdout.trim()))
+        .await
+        .unwrap();
+    assert_eq!(content, "xyzzy\n");
+}