From 3ead274a4485457e9c2660db5c31671be6117894 Mon Sep 17 00:00:00 2001 From: Vinnie Magro Date: Tue, 3 Dec 2024 13:42:53 -0800 Subject: [PATCH] [antlir2][image_command_alias] support env vars MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../antlir2/antlir2_isolate/isolate_cfg/src/lib.rs | 9 +++++++++ antlir/antlir2/image_command_alias/BUCK | 1 + .../image_command_alias/image_command_alias.bzl | 6 ++++++ antlir/antlir2/image_command_alias/src/main.rs | 14 ++++++++++++++ antlir/antlir2/image_command_alias/tests/BUCK | 8 ++++++-- .../tests/check-command-alias-out.sh | 3 ++- 6 files changed, 38 insertions(+), 3 deletions(-) diff --git a/antlir/antlir2/antlir2_isolate/isolate_cfg/src/lib.rs b/antlir/antlir2/antlir2_isolate/isolate_cfg/src/lib.rs index d783aff47f..96d94a10cf 100644 --- a/antlir/antlir2/antlir2_isolate/isolate_cfg/src/lib.rs +++ b/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>> { + 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>> { HashMap::from([(Cow::Owned(OsString::from(self.0)), Cow::Owned(self.1))]) diff --git a/antlir/antlir2/image_command_alias/BUCK b/antlir/antlir2/image_command_alias/BUCK index 0b280e04dc..c13ec7f5d0 100644 --- a/antlir/antlir2/image_command_alias/BUCK +++ b/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", ], ) diff --git a/antlir/antlir2/image_command_alias/image_command_alias.bzl b/antlir/antlir2/image_command_alias/image_command_alias.bzl index 5f01413bf6..79d1eddbbc 100644 --- a/antlir/antlir2/image_command_alias/image_command_alias.bzl +++ b/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), diff --git a/antlir/antlir2/image_command_alias/src/main.rs b/antlir/antlir2/image_command_alias/src/main.rs index 82e279d09c..48cd2ce361 100644 --- a/antlir/antlir2/image_command_alias/src/main.rs +++ b/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>>>, #[clap(required(true), trailing_var_arg(true), allow_hyphen_values(true))] command: Vec, } @@ -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")) diff --git a/antlir/antlir2/image_command_alias/tests/BUCK b/antlir/antlir2/image_command_alias/tests/BUCK index aacd1e53dd..e6dca5ad0a 100644 --- a/antlir/antlir2/image_command_alias/tests/BUCK +++ b/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", ) diff --git a/antlir/antlir2/image_command_alias/tests/check-command-alias-out.sh b/antlir/antlir2/image_command_alias/tests/check-command-alias-out.sh index f130c4b0fe..d29d62f381 100755 --- a/antlir/antlir2/image_command_alias/tests/check-command-alias-out.sh +++ b/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