[antlir2][vm] read env values instead of passing them in

Summary:
Currently, vmtest pass in `--setenv` with values from
`ExternalRunnerTestInfo.env`. This mostly works except when the env values
contain macros like `$(location)`. Based on testing, those macros will not be
evaluated at the time our provider generates the full command for spawning the
test. Thus we get them passed in verbatim into container and VM, which obviously
won't work. Instead, defer reading of the env values until our VM process spawns,
at which point the macros must have been evaluated.

This does mean that [shell] sub target will not be able to reproduce the same
environment as the real test. `buck run` simply won't set any of those test
envs. However, we don't really have a choice given `buck run` will not get the
evaluated value anyway.

More context:
https://fb.workplace.com/groups/askbuck/posts/25235126456109298
https://fb.workplace.com/groups/antlirusers/posts/1780377695727348
https://www.internalfb.com/diff/D49933136

Test Plan:
Enhance unit test and vmtests to cover this case. CI signal should
cover all these tests.

```name=Verify the failing test
$ buck2 test 'fbcode//mode/dbgo-cov' '-c' 'fbcode.cxx_coverage_only=strobelight' fbcode//strobelight/server/vmtest:pylatency-bpf-verification-vm-test-5.19.0-0_fbk4_10711_g9cbb08e4dfe0 -- --exact 'strobelight/server/vmtest:pylatency-bpf-verification-vm-test-5.19.0-0_fbk4_10711_g9cbb08e4dfe0 - PyLatencyVerificationTest.Load' --collect-coverage '--code-coverage-session=test_session' --force-tpx
File changed: fbsource//default.profraw
File changed: fbcode//strobelight/server/vmtest/strobelight_test_hhvm_usdt.py
File changed: fbcode//strobelight/server/test/USDTDummyLoopRRM.cpp
12 additional file change events
Buck UI: https://www.internalfb.com/buck2/28bf261e-14d0-4f8b-8c07-3e6fb2186a53
Test UI: https://www.internalfb.com/intern/testinfra/testrun/4503599830155826
Network: Up: 174B  Down: 11MiB  (reSessionID-c122a95b-c397-4508-aa43-ce8ccb5cd48c)
Jobs completed: 149. Time elapsed: 52.2s.
Cache hits: 100%. Commands: 1 (cached: 1, remote: 0, local: 0)
Tests finished: Pass 1. Fail 0. Fatal 0. Skip 0. Build failure 0
```

Reviewed By: epilatow

Differential Revision: D49981391

fbshipit-source-id: db7304bad30c6180fd1a0ffc5c411ef9bfa1a971

antlir/antlir2/antlir2_vm/bzl/test.bzl

@@ -22,10 +22,7 @@ def _impl(ctx: AnalysisContext) -> list[Provider]:
         cmd_args(ctx.attrs.vm_host[VMHostInfo].image[LayerInfo].subvol_symlink, format = "--image={}"),
         cmd_args(ctx.attrs.vm_host[VMHostInfo].machine_spec, format = "--machine-spec={}"),
         cmd_args(ctx.attrs.vm_host[VMHostInfo].runtime_spec, format = "--runtime-spec={}"),
-        cmd_args(
-            ["{}={}".format(k, v) for k, v in ctx.attrs.test[ExternalRunnerTestInfo].env.items()],
-            format = "--setenv={}",
-        ),
+        cmd_args([k for k in ctx.attrs.test[ExternalRunnerTestInfo].env], format = "--passenv={}"),
     )
     test_args = cmd_args(
         ctx.attrs.test[ExternalRunnerTestInfo].test_type,

antlir/antlir2/antlir2_vm/src/isolation.rs

@@ -8,7 +8,6 @@
 use std::collections::BTreeMap;
 use std::collections::HashSet;
 use std::env;
-use std::ffi::OsString;
 use std::path::PathBuf;
 use std::process::Command;
 
@@ -87,21 +86,6 @@ impl Platform {
     }
 }
 
-/// If these env exist, always pass them through.
-const PASSTHROUGH_ENVS: &[&str] = &["RUST_LOG", "ANTLIR_BUCK"];
-
-/// Generate default passthrough env vars
-pub(crate) fn default_passthrough_envs() -> Vec<KvPair> {
-    PASSTHROUGH_ENVS
-        .iter()
-        .filter(|x| env::var(*x).is_ok())
-        .map(|x| KvPair {
-            key: x.to_string(),
-            value: OsString::from(env::var(x).expect("must exist")),
-        })
-        .collect()
-}
-
 /// Return IsolatedContext ready for executing a command inside isolation
 /// # Arguments
 /// * `image` - container image that would be used to run the VM
@@ -136,21 +120,3 @@ pub(crate) fn is_isolated() -> Result<bool> {
     debug!("systemd-detect-virt returned: {}", virt);
     Ok(virt == "systemd-nspawn")
 }
-
-#[cfg(test)]
-mod test {
-    use super::*;
-
-    #[test]
-    fn test_default_envs() {
-        assert_eq!(default_passthrough_envs(), Vec::new());
-        env::set_var("RUST_LOG", "hello");
-        assert_eq!(
-            default_passthrough_envs(),
-            vec![KvPair {
-                key: "RUST_LOG".into(),
-                value: "hello".into(),
-            }],
-        );
-    }
-}

antlir/antlir2/antlir2_vm/src/main.rs

@@ -27,15 +27,13 @@ use anyhow::Context;
 use clap::Args;
 use clap::Parser;
 use clap::Subcommand;
-use image_test_lib::KvPair;
 use image_test_lib::Test;
 use json_arg::JsonFile;
 use tempfile::tempdir;
 use tracing::debug;
 use tracing_subscriber::filter::LevelFilter;
 use tracing_subscriber::prelude::*;
 
-use crate::isolation::default_passthrough_envs;
 use crate::isolation::is_isolated;
 use crate::isolation::isolated;
 use crate::isolation::Platform;
@@ -44,6 +42,7 @@ use crate::types::MachineOpts;
 use crate::types::RuntimeOpts;
 use crate::types::VMArgs;
 use crate::utils::console_output_path_for_tpx;
+use crate::utils::env_names_to_kvpairs;
 use crate::utils::log_command;
 use crate::vm::VM;
 
@@ -84,10 +83,9 @@ struct IsolateCmdArgs {
     /// Path to container image.
     #[arg(long)]
     image: PathBuf,
-    /// Set these env vars in the container. If VM executes a command, these
-    /// env vars will also be prepended to the command.
+    /// Pass through these environment variables into the container and VM, if they exist.
     #[arg(long)]
-    setenv: Vec<KvPair>,
+    passenv: Vec<String>,
     /// Args for run command
     #[clap(flatten)]
     run_cmd_args: RunCmdArgs,
@@ -111,9 +109,8 @@ fn run(args: &RunCmdArgs) -> Result<()> {
 /// Enter isolated container and then respawn itself inside it with `run`
 /// command and its parameters.
 fn respawn(args: &IsolateCmdArgs) -> Result<()> {
-    let mut envs = default_passthrough_envs();
-    envs.extend(args.setenv.clone());
     let mut vm_args = args.run_cmd_args.vm_args.clone();
+    let envs = env_names_to_kvpairs(args.passenv.clone());
     vm_args.command_envs = envs.clone();
 
     // Let's always capture console output unless it's console mode
@@ -142,21 +139,6 @@ fn respawn(args: &IsolateCmdArgs) -> Result<()> {
     Ok(())
 }
 
-/// Merge all sources of our envs into final list of env vars we should use
-/// everywhere for tests. Dedup is handled by functions that use the result.
-fn get_test_envs(from_cli: &[KvPair]) -> Vec<KvPair> {
-    // This handles common envs like RUST_LOG
-    let mut envs = default_passthrough_envs();
-    envs.extend_from_slice(from_cli);
-    // forward test runner env vars to the inner test
-    for (key, val) in std::env::vars() {
-        if key.starts_with("TEST_PILOT") {
-            envs.push((key, OsString::from(val)).into());
-        }
-    }
-    envs
-}
-
 /// Validated `VMArgs` and other necessary metadata for tests.
 struct ValidatedVMArgs {
     /// VMArgs that will be passed into the VM with modified fields
@@ -167,7 +149,7 @@ struct ValidatedVMArgs {
 
 /// Further validate `VMArgs` parsed by clap and generate a new `VMArgs` with
 /// content specific to test execution.
-fn get_test_vm_args(orig_args: &VMArgs, cli_envs: &[KvPair]) -> Result<ValidatedVMArgs> {
+fn get_test_vm_args(orig_args: &VMArgs, cli_envs: Vec<String>) -> Result<ValidatedVMArgs> {
     if orig_args.timeout_secs.is_none() {
         return Err(anyhow!("Test command must specify --timeout-secs."));
     }
@@ -177,7 +159,15 @@ fn get_test_vm_args(orig_args: &VMArgs, cli_envs: &[KvPair]) -> Result<Validated
             This will be parsed from env and test command parameters instead."
         ));
     }
-    let envs = get_test_envs(cli_envs);
+
+    // Forward test runner env vars to the inner test
+    let mut env_names = cli_envs;
+    for (key, _) in std::env::vars() {
+        if key.starts_with("TEST_PILOT") {
+            env_names.push(key);
+        }
+    }
+    let envs = env_names_to_kvpairs(env_names);
 
     #[derive(Debug, Parser)]
     struct TestArgsParser {
@@ -255,7 +245,7 @@ fn vm_test_command(args: &IsolateCmdArgs, validated_args: &ValidatedVMArgs) -> R
 /// inputs instead of allowing caller to pass them in. Some inputs are parsed
 /// from the test command.
 fn test(args: &IsolateCmdArgs) -> Result<()> {
-    let validated_args = get_test_vm_args(&args.run_cmd_args.vm_args, &args.setenv)?;
+    let validated_args = get_test_vm_args(&args.run_cmd_args.vm_args, args.passenv.clone())?;
     let mut command = if validated_args.is_list {
         list_test_command(args, &validated_args)
     } else {
@@ -292,24 +282,11 @@ fn main() -> Result<()> {
 
 #[cfg(test)]
 mod test {
+    use image_test_lib::KvPair;
+
     use super::*;
     use crate::types::VMModeArgs;
 
-    #[test]
-    fn test_get_test_envs() {
-        env::set_var("RUST_LOG", "hello");
-        env::set_var("TEST_PILOT_A", "A");
-        let from_cli = vec![KvPair::from(("foo", "bar"))];
-        assert_eq!(
-            get_test_envs(&from_cli),
-            vec![
-                KvPair::from(("RUST_LOG", "hello")),
-                KvPair::from(("foo", "bar")),
-                KvPair::from(("TEST_PILOT_A", "A")),
-            ],
-        )
-    }
-
     #[test]
     fn test_get_test_vm_args() {
         let valid = VMArgs {
@@ -320,25 +297,35 @@ mod test {
             },
             ..Default::default()
         };
-        let empty_env = Vec::<KvPair>::new();
         let mut expected = valid.clone();
         expected.mode.command = Some(vec![OsString::from("whatever")]);
-        let parsed = get_test_vm_args(&valid, &empty_env).expect("Parsing should succeed");
-        assert_eq!(parsed.inner, expected);
+        let parsed = get_test_vm_args(&valid, vec![]).expect("Parsing should succeed");
+        assert_eq!(parsed.inner.mode, expected.mode);
         assert!(!parsed.is_list);
 
         let mut timeout = valid.clone();
         timeout.timeout_secs = None;
-        assert!(get_test_vm_args(&timeout, &empty_env).is_err());
+        assert!(get_test_vm_args(&timeout, vec![]).is_err());
 
         let mut output_dirs = valid.clone();
         output_dirs.output_dirs = vec![PathBuf::from("/some")];
-        assert!(get_test_vm_args(&output_dirs, &empty_env).is_err());
+        assert!(get_test_vm_args(&output_dirs, vec![]).is_err());
 
-        let mut command = valid;
+        let mut command = valid.clone();
         command.mode.command = None;
-        assert!(get_test_vm_args(&command, &empty_env).is_err());
+        assert!(get_test_vm_args(&command, vec![]).is_err());
         command.mode.command = Some(vec![OsString::from("invalid")]);
-        assert!(get_test_vm_args(&command, &empty_env).is_err());
+        assert!(get_test_vm_args(&command, vec![]).is_err());
+
+        let env_var_test = valid;
+        std::env::set_var("TEST_PILOT_A", "A");
+        let parsed = get_test_vm_args(&env_var_test, vec![]).expect("Parsing should succeed");
+        assert!(
+            parsed
+                .inner
+                .command_envs
+                .iter()
+                .any(|x| *x == KvPair::from(("TEST_PILOT_A", "A")))
+        );
     }
 }

antlir/antlir2/antlir2_vm/src/utils.rs

@@ -5,12 +5,14 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::collections::HashSet;
 use std::fs;
 use std::io::ErrorKind;
 use std::path::Path;
 use std::path::PathBuf;
 use std::process::Command;
 
+use image_test_lib::KvPair;
 use tracing::debug;
 use tracing::error;
 
@@ -85,6 +87,24 @@ pub(crate) fn console_output_path_for_tpx() -> Result<Option<PathBuf>, std::io::
     }
 }
 
+/// Convert a list of env names into KvPair with its values
+pub(crate) fn env_names_to_kvpairs(env_names: Vec<String>) -> Vec<KvPair> {
+    let mut names: HashSet<_> = env_names.into_iter().collect();
+    // If these env exist, always pass them through.
+    ["RUST_LOG", "RUST_BACKTRACE", "ANTLIR_BUCK"]
+        .iter()
+        .for_each(|name| {
+            names.insert(name.to_string());
+        });
+    names
+        .iter()
+        .filter_map(|name| match std::env::var(name) {
+            Ok(value) => Some(KvPair::from((name, value))),
+            _ => None,
+        })
+        .collect()
+}
+
 #[cfg(test)]
 /// Helper function for converting qemu args to a single string for asserting in tests.
 /// This is usually only needed for string only functions like `contains`.
@@ -97,6 +117,10 @@ pub(crate) fn qemu_args_to_string(args: &[std::ffi::OsString]) -> String {
 
 #[cfg(test)]
 mod test {
+    use std::collections::HashMap;
+    use std::env;
+    use std::ffi::OsString;
+
     use super::*;
 
     #[test]
@@ -110,4 +134,103 @@ mod test {
             format!("Program: `hello`. Args: `{:?}`", vec!["world"]),
         );
     }
+
+    struct EnvTest {
+        envs: Vec<(&'static str, &'static str)>,
+        passenv: Vec<&'static str>,
+        result: HashMap<String, OsString>,
+    }
+
+    const ALLOWED_ENV_NAMES: &[&str] = &[
+        "RUST_LOG",
+        "RUST_BACKTRACE",
+        "ANTLIR_BUCK",
+        "TEST_PILOT_A",
+        "OTHER",
+    ];
+
+    impl EnvTest {
+        fn new(
+            envs: Vec<(&'static str, &'static str)>,
+            passenv: Vec<&'static str>,
+            result: Vec<(&'static str, &'static str)>,
+        ) -> Self {
+            // We have to clear all envs across tests, so we must know the ones to clear.
+            envs.iter().for_each(|(name, _)| {
+                assert!(
+                    ALLOWED_ENV_NAMES.contains(name),
+                    "{name} not allowed for testing"
+                )
+            });
+            Self {
+                envs,
+                passenv,
+                result: result
+                    .into_iter()
+                    .map(|(k, v)| (k.into(), v.into()))
+                    .collect(),
+            }
+        }
+
+        fn test(&self) {
+            ALLOWED_ENV_NAMES.iter().for_each(|name| {
+                env::remove_var(name);
+            });
+            self.envs.iter().for_each(|(name, val)| {
+                env::set_var(name, val);
+            });
+            let result: HashMap<_, _> =
+                env_names_to_kvpairs(self.passenv.iter().map(|s| s.to_string()).collect())
+                    .drain(..)
+                    .map(|pair| (pair.key, pair.value))
+                    .collect();
+            assert_eq!(result, self.result);
+        }
+    }
+
+    #[test]
+    fn test_env_names_to_kvpairs() {
+        [
+            EnvTest::new(vec![], vec![], vec![]),
+            // Reads nothing
+            EnvTest::new(
+                vec![],
+                vec!["RUST_LOG", "ANTLIR_BUCK", "TEST_PILOT_A", "OTHER"],
+                vec![],
+            ),
+            // Always pass through
+            EnvTest::new(
+                vec![
+                    ("RUST_LOG", "info"),
+                    ("ANTLIR_BUCK", "1"),
+                    ("TEST_PILOT_A", "A"),
+                    ("OTHER", "other"),
+                ],
+                vec![],
+                vec![("RUST_LOG", "info"), ("ANTLIR_BUCK", "1")],
+            ),
+            // Selection
+            EnvTest::new(
+                vec![("TEST_PILOT_A", "A"), ("OTHER", "other")],
+                vec!["TEST_PILOT_A"],
+                vec![("TEST_PILOT_A", "A")],
+            ),
+            // Mixed
+            EnvTest::new(
+                vec![
+                    ("RUST_LOG", "info"),
+                    ("TEST_PILOT_A", "A"),
+                    ("OTHER", "other"),
+                ],
+                vec!["TEST_PILOT_A"],
+                vec![("TEST_PILOT_A", "A"), ("RUST_LOG", "info")],
+            ),
+        ]
+        .iter()
+        .enumerate()
+        .for_each(|(i, test)| {
+            println!("Running test #{i}");
+            test.test();
+        });
+    }
 }