antlir/vm: add sidecar on-host service support

Summary:
This is worse than zeroxoneb's proposed systemd container runtime environment,
but until that is ready this can serve as a stopgap measure for tests that
require some supporting service to run on the host before the VM is booted (for
example services like DHCP)

Reviewed By: zeroxoneb

Differential Revision: D28098482

fbshipit-source-id: 0f187d4be997a55dd1ad4953246437f92feafe19

antlir/bzl/vm/types.bzl

@@ -108,15 +108,20 @@ _vm_runtime_t = shape.shape(
 
     # Details of the emulator being used to run the VM
     emulator = shape.field(_vm_emulator_t),
+
+    # Shell commands to run before booting the VM
+    sidecar_services = shape.list(str),
 )
 
 def _new_vm_runtime(
         connection = None,
-        emulator = None):
+        emulator = None,
+        sidecar_services = None):
     return shape.new(
         _vm_runtime_t,
         connection = connection or _new_vm_connection(),
         emulator = emulator or _new_vm_emulator(),
+        sidecar_services = sidecar_services or [],
     )
 
 _vm_runtime_api = struct(

antlir/vm/tap.py

@@ -75,8 +75,17 @@ def __post_init__(self):
                 text=True,
                 stdin=subprocess.DEVNULL,
             )
+            subprocess.run(
+                self.netns.nsenter_as_root(
+                    "ip", "addr", "add", self.host_ipv6_ll, "dev", TAPDEV
+                ),
+                check=True,
+                capture_output=True,
+                text=True,
+                stdin=subprocess.DEVNULL,
+            )
         except subprocess.CalledProcessError as e:
-            raise TapError(f"Failed to create tap device: {e.stderr}")
+            raise TapError(f"Failed to setup tap device: {e.stderr}")
 
     def _ensure_dev_net_tun(self) -> None:
         # See class docblock, this should eventually be handled by the
@@ -114,6 +123,10 @@ def guest_mac(self) -> str:
     def guest_ipv6_ll(self) -> str:
         return f"fe80::200:0ff:fe00:1%{TAPDEV}"
 
+    @property
+    def host_ipv6_ll(self) -> str:
+        return "fe80::200:0ff:fe00:2"
+
     @property
     def qemu_args(self) -> Iterable[str]:
         return (

antlir/vm/tests/BUCK

@@ -1,6 +1,6 @@
 load("//antlir/bzl:image.bzl", "image")
 load("//antlir/bzl:image_unittest_helpers.bzl", helpers = "image_unittest_helpers")
-load("//antlir/bzl:oss_shim.bzl", "export_file", "python_unittest")
+load("//antlir/bzl:oss_shim.bzl", "export_file", "python_unittest", "rust_binary", "third_party")
 load("//antlir/bzl:shape.bzl", "shape")
 load("//antlir/bzl/vm:defs.bzl", "vm")
 
@@ -151,10 +151,45 @@ vm.cpp_unittest(
     env = test_env_vars,
 )
 
+vm.rust_unittest(
+    name = "rust",
+    srcs = ["rust_test.rs"],
+    env = test_env_vars,
+    crate_root = "rust_test.rs",
+)
+
 vm.python_unittest(
     name = "test-with-kernel-devel",
     srcs = ["test_kernel_devel.py"],
     vm_opts = vm.types.opts.new(
         devel = True,
     ),
 )
+
+rust_binary(
+    name = "sidecar",
+    srcs = ["sidecar.rs"],
+    deps = [
+        third_party.library(
+            "tokio",
+            platform = "rust",
+        ),
+    ],
+)
+
+vm.rust_unittest(
+    name = "with-sidecar",
+    srcs = ["test_with_sidecar.rs"],
+    vm_opts = vm.types.opts.new(
+        runtime = vm.types.runtime.new(
+            sidecar_services = ["$(exe :sidecar)"],
+        ),
+    ),
+    crate_root = "test_with_sidecar.rs",
+    deps = [
+        third_party.library(
+            "tokio",
+            platform = "rust",
+        ),
+    ],
+)

antlir/vm/tests/sidecar.rs

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::TcpListener;
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let listener = TcpListener::bind("[::]:8080").await?;
+
+    loop {
+        let (mut socket, _) = listener.accept().await?;
+
+        tokio::spawn(async move {
+            let mut buf = [0; 1024];
+
+            // In a loop, read data from the socket and write the data back.
+            loop {
+                let n = match socket.read(&mut buf).await {
+                    // socket closed
+                    Ok(n) if n == 0 => return,
+                    Ok(n) => n,
+                    Err(e) => {
+                        eprintln!("failed to read from socket; err = {:?}", e);
+                        return;
+                    }
+                };
+
+                // Write the data back
+                if let Err(e) = socket.write_all(&buf[0..n]).await {
+                    eprintln!("failed to write to socket; err = {:?}", e);
+                    return;
+                }
+            }
+        });
+    }
+}

antlir/vm/tests/test_with_sidecar.rs

@@ -0,0 +1,27 @@
+/*
+ * Copyright (c) Facebook, Inc. and its 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::net::SocketAddrV6;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::TcpStream;
+
+#[tokio::test]
+async fn sidecar() {
+    let sock = SocketAddrV6::new("fe80::200:0ff:fe00:2".parse().unwrap(), 8080, 0, 2);
+    let mut stream = TcpStream::connect(sock).await.unwrap();
+
+    // Write some data.
+    stream.write_all(b"hello world!").await.unwrap();
+    let mut buf = [0; 1024];
+    stream.read(&mut buf).await.unwrap();
+    assert_eq!(
+        std::str::from_utf8(&buf)
+            .unwrap()
+            .trim_end_matches(char::from(0)),
+        "hello world!"
+    );
+}

antlir/vm/vm.py

@@ -358,6 +358,22 @@ async def __vm_with_stack(
         f"{' '.join(ns.nsenter_as_user('/bin/bash'))}"
     )
 
+    logger.debug(
+        f"Starting sidecars {opts.runtime.sidecar_services} before QEMU"
+    )
+    sidecar_procs = await asyncio.gather(
+        # Execing in the shell is not the safest thing, but it makes it easy to
+        # pass extra arguments from the TARGETS files that define tests which
+        # require sidecars as well as easily handling the `$(exe)` expansion of
+        # `python_binary` rules into `python3 -Es $par`
+        *(
+            asyncio.create_subprocess_exec(
+                *ns.nsenter_as_user("/bin/sh", "-c", sidecar)
+            )
+            for sidecar in opts.runtime.sidecar_services
+        )
+    )
+
     tapdev = VmTap(netns=ns, uid=os.getuid(), gid=os.getgid())
     args = [
         "-no-reboot",
@@ -511,7 +527,6 @@ async def __vm_with_stack(
         logger.error(f"VM failed: {e}")
         raise RuntimeError(f"VM failed: {e}")
     finally:
-
         # Future: unless we are running in `--shell=console` mode, the VM
         # hasn't been told to shutdown yet.  So this is the 'default'
         # behavior for termination, but really this should be a last resort
@@ -550,6 +565,8 @@ async def __vm_with_stack(
 
         logger.debug(f"VM exited with: {proc.returncode}")
 
+        subprocess.run(["sudo", "kill", "-KILL", "--", *[str(proc.pid) for proc in sidecar_procs]])
+
 
 @asynccontextmanager
 async def vm(*args, **kwargs) -> AsyncContextManager[GuestSSHConnection]: