[antlir2][image_test] test execution subcommand

Summary:
Create a subcommand that actually does the test execution.
This does not change any net behavior of `image_test`, but it will soon make it
possible for this to be invoked by a user in an interactive shell.

Test Plan:
```
❯ buck2 test fbcode//antlir/antlir2/testing/tests:test-rust-centos9 fbcode//antlir/antlir2/testing/tests:test-rust-nobody-centos9 fbcode//antlir/antlir2/testing/tests:test-rust-boot-centos9 fbcode//antlir/antlir2/testing/tests:test-rust-boot-nobody-centos9
Buck UI: https://www.internalfb.com/buck2/7b40707e-700d-43a5-a43b-6aa790f29c1b
Test UI: https://www.internalfb.com/intern/testinfra/testrun/9288674289047734
Network: Up: 6.2KiB  Down: 2.0KiB  (reSessionID-07bd93dd-40c3-4a2d-afad-25ac6a24c372)
Jobs completed: 60. Time elapsed: 7.4s.
Cache hits: 0%. Commands: 6 (cached: 0, remote: 0, local: 6)
Tests finished: Pass 20. Fail 0. Fatal 0. Skip 0. Build failure 0
```

Reviewed By: naveedgol

Differential Revision: D65825538

fbshipit-source-id: dff6f2307a6740a80a963277289de33fc568d522

antlir/antlir2/testing/image_test/BUCK

@@ -21,8 +21,11 @@ rust_binary(
     visibility = ["PUBLIC"],
     deps = [
         "anyhow",
+        "bon",
         "clap",
+        "nix",
         "serde",
+        "serde_json",
         "tempfile",
         "textwrap",
         "tracing",

antlir/antlir2/testing/image_test/antlir2_image_test.service

@@ -14,6 +14,6 @@ FailureAction=exit-force
 # happen). {Failure,Success}Action are still respected when the test process
 # exits either way.
 Type=simple
-Environment=USER=%u
+ExecStart=/__antlir2_image_test__/image-test exec
 StandardOutput=truncate:/antlir2/test_stdout
 StandardError=truncate:/antlir2/test_stderr

antlir/antlir2/testing/image_test/src/exec.rs

@@ -0,0 +1,76 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::collections::BTreeMap;
+use std::ffi::OsString;
+use std::os::unix::process::CommandExt;
+use std::path::PathBuf;
+use std::process::Command;
+
+use anyhow::Context;
+use anyhow::Result;
+use bon::Builder;
+use clap::Parser;
+use json_arg::JsonFile;
+use nix::unistd::User;
+use serde::Deserialize;
+use serde::Serialize;
+
+#[derive(Debug, Clone, Builder, Serialize, Deserialize)]
+/// Specification of how to execute the test.
+/// This specification is just how to invoke the inner test binary, the
+/// containerization should already have been set up by 'spawn'.
+pub(crate) struct Spec {
+    /// The test command
+    cmd: Vec<OsString>,
+    /// CWD of the test
+    working_directory: PathBuf,
+    /// Run the test as this user
+    user: String,
+    /// Set these env vars in the test environment
+    #[serde(default)]
+    env: BTreeMap<String, String>,
+}
+
+#[derive(Debug, Parser)]
+/// Execute the inner test
+pub(crate) struct Args {
+    #[clap(default_value = "/__antlir2_image_test__/exec_spec.json")]
+    spec: JsonFile<Spec>,
+}
+
+impl Args {
+    pub(crate) fn run(self) -> Result<()> {
+        let spec = self.spec.into_inner();
+        std::env::set_current_dir(&spec.working_directory)
+            .with_context(|| format!("while changing to '{}'", spec.working_directory.display()))?;
+        let mut env = spec.env;
+        env.insert("USER".into(), spec.user.clone());
+        env.insert(
+            "PWD".into(),
+            spec.working_directory
+                .to_str()
+                .with_context(|| {
+                    format!("pwd '{}' was not utf8", spec.working_directory.display())
+                })?
+                .into(),
+        );
+
+        let user = User::from_name(&spec.user)
+            .context("failed to lookup user")?
+            .with_context(|| format!("no such user '{}'", spec.user))?;
+
+        let mut cmd = spec.cmd.into_iter();
+        let err = Command::new(cmd.next().context("test command was empty")?)
+            .args(cmd)
+            .envs(env)
+            .uid(user.uid.into())
+            .gid(user.gid.into())
+            .exec();
+        Err(err.into())
+    }
+}

antlir/antlir2/testing/image_test/src/main.rs

@@ -8,16 +8,17 @@
 use anyhow::Result;
 use clap::Parser;
 
+mod exec;
 mod runtime;
 mod shell_help;
 mod spawn;
 
-pub(crate) use runtime::RuntimeSpec;
-
 #[derive(Parser, Debug)]
 enum Args {
     /// Spawn a container to run the test
     Spawn(spawn::Args),
+    /// Execute the test from inside the container
+    Exec(exec::Args),
     ShellHelp(shell_help::Args),
 }
 
@@ -28,6 +29,7 @@ fn main() -> Result<()> {
 
     match args {
         Args::Spawn(a) => a.run(),
+        Args::Exec(a) => a.run(),
         Args::ShellHelp(a) => a.run(),
     }
 }

antlir/antlir2/testing/image_test/src/runtime.rs

@@ -13,7 +13,7 @@ use serde::Deserialize;
 
 #[derive(Debug, Clone, Deserialize)]
 /// Specification of the test runtime (the rootfs layer, environment, etc)
-pub(crate) struct RuntimeSpec {
+pub(crate) struct Spec {
     /// Path to layer to run the test in
     pub(crate) layer: PathBuf,
     /// Run the test as this user

antlir/antlir2/testing/image_test/src/spawn.rs

@@ -12,7 +12,6 @@ use std::fs::Permissions;
 use std::io::Write;
 use std::os::fd::AsRawFd;
 use std::os::fd::FromRawFd;
-use std::os::unix::ffi::OsStrExt;
 use std::os::unix::fs::PermissionsExt;
 use std::os::unix::process::CommandExt;
 use std::path::Path;
@@ -32,7 +31,8 @@ use tempfile::NamedTempFile;
 use tracing::debug;
 use tracing::trace;
 
-use crate::RuntimeSpec;
+use crate::exec;
+use crate::runtime;
 
 fn make_log_files(_base: &str) -> Result<(NamedTempFile, NamedTempFile)> {
     Ok((NamedTempFile::new()?, NamedTempFile::new()?))
@@ -42,7 +42,7 @@ fn make_log_files(_base: &str) -> Result<(NamedTempFile, NamedTempFile)> {
 #[derive(Parser, Debug)]
 pub(crate) struct Args {
     #[clap(long)]
-    spec: JsonFile<RuntimeSpec>,
+    spec: JsonFile<runtime::Spec>,
     #[clap(subcommand)]
     test: Test,
 }
@@ -189,53 +189,6 @@ impl Args {
                     writeln!(test_unit_dropin, "Wants={unit}")?;
                 }
 
-                writeln!(test_unit_dropin, "[Service]")?;
-
-                writeln!(test_unit_dropin, "User={}", spec.user)?;
-                write!(test_unit_dropin, "WorkingDirectory=")?;
-                let cwd = std::env::current_dir().context("while getting cwd")?;
-                test_unit_dropin.write_all(cwd.as_os_str().as_bytes())?;
-                test_unit_dropin.write_all(b"\n")?;
-
-                write!(test_unit_dropin, "Environment=PWD=")?;
-                test_unit_dropin.write_all(cwd.as_os_str().as_bytes())?;
-                test_unit_dropin.write_all(b"\n")?;
-
-                write!(test_unit_dropin, "ExecStart=")?;
-                let mut iter = self.test.into_inner_cmd().into_iter().peekable();
-                if let Some(exe) = iter.next() {
-                    let realpath = std::fs::canonicalize(&exe)
-                        .with_context(|| format!("while getting absolute path of {exe:?}"))?;
-                    test_unit_dropin.write_all(realpath.as_os_str().as_bytes())?;
-                    if iter.peek().is_some() {
-                        test_unit_dropin.write_all(b" ")?;
-                    }
-                }
-                while let Some(arg) = iter.next() {
-                    test_unit_dropin.write_all(arg.as_os_str().as_bytes())?;
-                    if iter.peek().is_some() {
-                        test_unit_dropin.write_all(b" ")?;
-                    }
-                }
-                test_unit_dropin.write_all(b"\n")?;
-
-                for (key, val) in &setenv {
-                    writeln!(
-                        test_unit_dropin,
-                        "Environment=\"{key}={}\"",
-                        val.replace('"', "\\\"")
-                    )?;
-                }
-                // forward test runner env vars to the inner test
-                for (key, val) in std::env::vars() {
-                    if key.starts_with("TEST_PILOT") {
-                        writeln!(
-                            test_unit_dropin,
-                            "Environment=\"{key}={}\"",
-                            val.replace('"', "\\\"")
-                        )?;
-                    }
-                }
                 // wire the test output to the parent process's std{out,err}
                 ctx.outputs(HashMap::from([
                     (Path::new("/antlir2/test_stdout"), test_stdout.path()),
@@ -246,6 +199,29 @@ impl Args {
                     test_unit_dropin.path(),
                 ));
 
+                let mut exec_env = setenv.clone();
+                // forward test runner env vars to the inner test
+                for (key, val) in std::env::vars() {
+                    if key.starts_with("TEST_PILOT") {
+                        exec_env.insert(key, val);
+                    }
+                }
+
+                let exec_spec = exec::Spec::builder()
+                    .cmd(self.test.into_inner_cmd())
+                    .user(spec.user)
+                    .working_directory(std::env::current_dir().context("while getting cwd")?)
+                    .env(exec_env)
+                    .build();
+                let exec_spec_file = tempfile::NamedTempFile::new()
+                    .context("while creating temp file for exec spec")?;
+                serde_json::to_writer_pretty(&exec_spec_file, &exec_spec)
+                    .context("while serializing exec spec to file")?;
+                ctx.inputs((
+                    Path::new("/__antlir2_image_test__/exec_spec.json"),
+                    exec_spec_file.path(),
+                ));
+
                 // Register the test container with systemd-machined so manual debugging
                 // is a easier.
                 ctx.register(true);

antlir/antlir2/testing/tests/BUCK

@@ -46,6 +46,7 @@ test_variants(
     lang = "rust",
     test_rule = image_rust_test,
     deps = [
+        "nix",
         "rustix",
         "serde_json",
         "whoami",

antlir/antlir2/testing/tests/test.rs

@@ -7,13 +7,20 @@
 
 use std::path::Path;
 
+use nix::unistd::getuid;
+use nix::unistd::User;
 use rustix::fs::statfs;
 
 #[test]
 fn user() {
     let expected = std::env::var("TEST_USER").expect("TEST_USER not set");
     let actual = whoami::username();
     assert_eq!(expected, actual);
+    let expected_uid = User::from_name(&expected)
+        .expect("failed to lookup user")
+        .expect("no such user")
+        .uid;
+    assert_eq!(getuid(), expected_uid);
 }
 
 #[test]