[antlir2][image_command_alias] support env vars

Summary:
Directly support passing an env var to the command inside the image since that
is how some tools expect to be configured

Test Plan:
```
❯ buck2 test fbcode//antlir/antlir2/image_command_alias/...
Buck UI: https://www.internalfb.com/buck2/156f5365-5050-4d9c-9134-9f35ff5a79a5
Test UI: https://www.internalfb.com/intern/testinfra/testrun/10977524151792296
Network: Up: 0B  Down: 0B  (reSessionID-c47d865c-8b79-4e64-8138-a50f3f53aa10)
Jobs completed: 12. Time elapsed: 1.9s.
Tests finished: Pass 3. Fail 0. Fatal 0. Skip 0. Build failure 0
```

Reviewed By: epilatow

Differential Revision: D66710942

fbshipit-source-id: c45a3e826235ad051b29523a3859af9567ed95ee

antlir/antlir2/antlir2_isolate/isolate_cfg/src/lib.rs

@@ -414,6 +414,15 @@ impl<'a> IntoEnv<'a> for (&'a str, OsString) {
     }
 }
 
+impl<'a> IntoEnv<'a> for (String, String) {
+    fn into_env(self) -> HashMap<Cow<'a, OsStr>, Cow<'a, OsStr>> {
+        HashMap::from([(
+            Cow::Owned(OsString::from(self.0)),
+            Cow::Owned(OsString::from(self.1)),
+        )])
+    }
+}
+
 impl<'a> IntoEnv<'a> for (String, OsString) {
     fn into_env(self) -> HashMap<Cow<'a, OsStr>, Cow<'a, OsStr>> {
         HashMap::from([(Cow::Owned(OsString::from(self.0)), Cow::Owned(self.1))])

antlir/antlir2/image_command_alias/BUCK

@@ -13,5 +13,6 @@ rust_binary(
         "tracing-subscriber",
         "//antlir/antlir2/antlir2_isolate:antlir2_isolate",
         "//antlir/antlir2/antlir2_rootless:antlir2_rootless",
+        "//antlir/util/cli/json_arg:json_arg",
     ],
 )

antlir/antlir2/image_command_alias/image_command_alias.bzl

@@ -11,9 +11,14 @@ load("//antlir/buck2/bzl:ensure_single_output.bzl", "ensure_single_output")
 
 def _impl(ctx: AnalysisContext) -> list[Provider] | Promise:
     root = ensure_single_output(ctx.attrs.root)
+    if ctx.attrs.env:
+        env_json = ctx.actions.write_json("env.json", ctx.attrs.env)
+    else:
+        env_json = None
     cmd = cmd_args(
         ctx.attrs._command_alias[RunInfo],
         cmd_args(root, format = "--root={}"),
+        cmd_args(env_json, format = "--env={}") if env_json else cmd_args(),
         "--",
         ctx.attrs.exe,
         cmd_args(ctx.attrs.args),
@@ -46,6 +51,7 @@ _image_command_alias = rule(
     impl = _impl,
     attrs = {
         "args": attrs.list(attrs.arg(), default = []),
+        "env": attrs.dict(attrs.string(), attrs.arg(), default = {}),
         "exe": attrs.arg(),
         "labels": attrs.list(attrs.string(), default = []),
         "root": attrs.source(allow_directory = True),

antlir/antlir2/image_command_alias/src/main.rs

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::collections::BTreeMap;
 use std::fs;
 use std::path::Component;
 use std::path::Path;
@@ -18,11 +19,14 @@ use anyhow::ensure;
 use anyhow::Context;
 use anyhow::Result;
 use clap::Parser;
+use json_arg::JsonFile;
 
 #[derive(Debug, Parser)]
 struct Args {
     #[clap(long)]
     root: PathBuf,
+    #[clap(long)]
+    env: Option<JsonFile<BTreeMap<String, Vec<String>>>>,
     #[clap(required(true), trailing_var_arg(true), allow_hyphen_values(true))]
     command: Vec<String>,
 }
@@ -58,6 +62,16 @@ fn main() -> Result<()> {
         }
     }
 
+    if let Some(env) = args.env.map(JsonFile::into_inner) {
+        for (k, mut v) in env {
+            ensure!(
+                v.len() == 1,
+                "env var '{k}' expanded to multiple values: {v:#?}"
+            );
+            builder.setenv((k, v.remove(0)));
+        }
+    }
+
     builder
         .working_directory(cwd.as_path())
         .tmpfs(Path::new("/tmp"))

antlir/antlir2/image_command_alias/tests/BUCK

@@ -27,7 +27,8 @@ out_file=$(echo "${@: -1}")
 set -- "${@:1:$(($#-1))}"
 
 # Write remaining arguments to output file.
-echo $@ > $out_file
+echo -n $@ > $out_file
+echo "$SUFFIX" >> $out_file
 """,
         ),
         feature.rpms_install(rpms = [
@@ -49,7 +50,7 @@ image_command_alias(
 buck_genrule(
     name = "run-command-alias-echo",
     out = "out-file",
-    cmd = "$(exe :command-alias-echo) goodbye world > $OUT",
+    cmd = "$(exe :command-alias-echo) goodbye world ! > $OUT",
 )
 
 # Test running a command in an image and capturing its stdout
@@ -65,6 +66,9 @@ image_command_alias(
         "hello",
         "world",
     ],
+    env = {
+        "SUFFIX": " !",
+    },
     exe = "/bin/command-alias-write",
     layer = ":layer",
 )

antlir/antlir2/image_command_alias/tests/check-command-alias-out.sh

@@ -6,7 +6,8 @@
 
 set -ue -o pipefail
 data=$(cat "$1")
-if [[ "$data" != "hello world goodbye world" ]]; then
+if [[ "$data" != "hello world goodbye world !" ]]; then
     echo "Unexpected data in $1" >&2
+    echo "$data"
     exit 1
 fi