From 441a17fb91c5e237c99073e02685b040c9ef09d1 Mon Sep 17 00:00:00 2001 From: Tom Dyas Date: Fri, 3 Jan 2025 19:17:43 +0000 Subject: [PATCH] substitute `{chroot}` marker in environment variables for remote execution (#21810) As described in #17519, remote execution environments did not substitute the `{chroot}` marker with the path to the execution sandbox as is done with local and Docker environments. #21803 added this support for program arguments. This PR completes the feature and implements that support for environment variables. Closes #17519. --- docs/notes/2.25.x.md | 2 +- .../remote/src/remote_tests.rs | 6 +- src/rust/engine/process_execution/src/lib.rs | 55 +++++++++++++++---- .../engine/process_execution/src/tests.rs | 55 ++++++++++++++++++- 4 files changed, 102 insertions(+), 16 deletions(-) diff --git a/docs/notes/2.25.x.md b/docs/notes/2.25.x.md index f893dcda712..75995b77cd7 100644 --- a/docs/notes/2.25.x.md +++ b/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 diff --git a/src/rust/engine/process_execution/remote/src/remote_tests.rs b/src/rust/engine/process_execution/remote/src/remote_tests.rs index 2ef7f27496a..916028b0a61 100644 --- a/src/rust/engine/process_execution/remote/src/remote_tests.rs +++ b/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, diff --git a/src/rust/engine/process_execution/src/lib.rs b/src/rust/engine/process_execution/src/lib.rs index 7e74770b807..19bee39114c 100644 --- a/src/rust/engine/process_execution/src/lib.rs +++ b/src/rust/engine/process_execution/src/lib.rs @@ -1023,12 +1023,12 @@ fn quote_path(path: &Path) -> Result { /// 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, append_only_caches_base_path: Option<&str>, working_directory: Option<&str>, sandbox_root_token: Option<&str>, + env_vars_to_substitute: Vec<&str>, ) -> Result, String> { let caches_fragment = match append_only_caches_base_path { Some(base_path) if !caches.is_empty() => { @@ -1088,7 +1088,7 @@ 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", @@ -1096,6 +1096,18 @@ fn maybe_make_wrapper_script( ), 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 { 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::>(); + 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::>(); + + 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), }); } diff --git a/src/rust/engine/process_execution/src/tests.rs b/src/rust/engine/process_execution/src/tests.rs index 6cbfba66d4c..ab82d61bf14 100644 --- a/src/rust/engine/process_execution/src/tests.rs +++ b/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"); +}