Add PXE boot support and VM tests for LiveOS images

Add the grub handling needed to PXE (network) boot Azure Linux LiveOS
images, plus VM tests that exercise it.

Product (toolkit/tools/pkg/imagecustomizerlib):
- Expand BLS Type #1 boot entries into explicit grub menuentries for the
  PXE grub config, since grub cannot enumerate boot/loader/entries over
  TFTP. This is done per distro through a new
  DistroHandler.UpdateLiveOSGrubCfgForPxe hook and a shared BLS renderer.

Tests (test/vmtests):
- test_pxe.py: bootstrap and full-OS PXE boot tests for AZL3 and AZL4,
  driven by a transient libvirt NAT network (dnsmasq DHCP and TFTP) and
  an HTTP server over the extracted PXE artifacts (pxe_server.py).
- The tests are skipped on Azure Linux hosts, whose edk2-ovmf lacks the
  UEFI network stack (EDK II NetworkPkg) and so cannot netboot a VM. They
  run on Ubuntu hosts, whose OVMF includes it.

Author: vinceaperri

test/vmtests/vmtests/imagecustomizer/test_pxe.py

@@ -0,0 +1,280 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import logging
+import platform
+import random
+import string
+import tarfile
+from pathlib import Path
+from typing import List, Tuple
+
+import libvirt  # type: ignore
+import pytest
+from docker import DockerClient
+
+from ..conftest import TEST_CONFIGS_DIR
+from ..utils.closeable import Closeable
+from ..utils.host_utils import get_host_distro
+from ..utils.imagecustomizer import (
+    add_preview_features_to_config,
+    add_pxe_bootstrap_base_url_to_config,
+    add_ssh_to_config,
+    run_image_customizer,
+)
+from ..utils.libvirt_utils import VmSpec, create_libvirt_domain_xml
+from ..utils.libvirt_vm import LibvirtVm
+from ..utils.pxe_server import PXE_HTTP_PORT, PXE_NETWORK_GATEWAY_IP, PxeEnvironment
+from ..utils.user_utils import get_username
+from .test_min_change import run_basic_checks
+
+# The full-OS image is downloaded into RAM during PXE bootstrap, so the VM needs more memory than a disk/ISO boot.
+PXE_VM_MEMORY_MIB = 8192
+PXE_VM_CORE_COUNT = 4
+
+# PXE boot adds a firmware netboot phase plus an over-the-network bootstrap-image download before the OS requests its
+# DHCP lease, so it needs additional time to boot.
+PXE_BOOT_IP_WAIT_TIME_EXTRA_SECONDS = 600
+
+
+def run_pxe_test(
+    docker_client: DockerClient,
+    image_customizer_container_url: str,
+    input_image: Path,
+    input_image_azl_release: int,
+    initramfs_type: str,
+    config_path: Path,
+    ssh_key: Tuple[str, Path],
+    test_temp_dir: Path,
+    test_instance_name: str,
+    logs_dir: Path,
+    libvirt_conn: libvirt.virConnect,
+    close_list: List[Closeable],
+) -> None:
+
+    ssh_public_key, ssh_private_key_path = ssh_key
+
+    if platform.machine() == "x86_64":
+        boot_loader_file = "bootx64.efi"
+    else:
+        boot_loader_file = "bootaa64.efi"
+
+    username = get_username()
+
+    modified_config_path = add_ssh_to_config(config_path, username, ssh_public_key, close_list)
+
+    if initramfs_type == "bootstrap":
+        bootstrap_base_url = f"http://{PXE_NETWORK_GATEWAY_IP}:{PXE_HTTP_PORT}"
+        modified_config_path = add_pxe_bootstrap_base_url_to_config(
+            modified_config_path, bootstrap_base_url, close_list
+        )
+
+    pxe_tar_path = test_temp_dir.joinpath("pxe-artifacts.tar.gz")
+    run_image_customizer(
+        docker_client,
+        image_customizer_container_url,
+        "customize",
+        modified_config_path,
+        "pxe-tar",
+        pxe_tar_path,
+        image_file=input_image,
+    )
+
+    pxe_artifacts_dir = test_temp_dir.joinpath("pxe-artifacts")
+    pxe_artifacts_dir.mkdir()
+    with tarfile.open(pxe_tar_path, "r:gz") as tar:
+        tar.extractall(pxe_artifacts_dir)
+
+    customized_name = (
+        "pxe_"
+        + initramfs_type.replace("-", "_")
+        + "_"
+        + get_host_distro()
+        + "_efi_azl"
+        + str(input_image_azl_release)
+        + "_to_efi"
+    )
+    customized_log_path = str(logs_dir) + "/" + customized_name
+    http_log_file_path = Path(customized_log_path + ".http.log")
+    vm_console_log_file_path = customized_log_path + ".console.log"
+
+    suffix = "".join(random.choice(string.ascii_lowercase) for _ in range(5))
+    network_name = test_instance_name + "-pxe"
+    bridge_name = "pxebr" + suffix
+
+    pxe_env = PxeEnvironment(
+        libvirt_conn,
+        network_name,
+        bridge_name,
+        pxe_artifacts_dir,
+        boot_loader_file,
+        http_log_file_path,
+    )
+    close_list.append(pxe_env)
+
+    # Create the VM: no disk, boots from the PXE network.
+    vm_spec = VmSpec(
+        test_instance_name,
+        PXE_VM_MEMORY_MIB,
+        PXE_VM_CORE_COUNT,
+        None,
+        "efi",
+        secure_boot=False,
+        pxe_boot=True,
+        network_name=pxe_env.network_name,
+    )
+    domain_xml = create_libvirt_domain_xml(libvirt_conn, vm_spec)
+    logging.debug(f"\n\ndomain_xml = {domain_xml}\n\n")
+
+    vm = LibvirtVm(test_instance_name, domain_xml, vm_console_log_file_path, libvirt_conn)
+    close_list.append(vm)
+
+    # Start the VM.
+    vm.start()
+
+    # Connect to the VM and run the basic boot validation.
+    with vm.create_ssh_client(
+        ssh_private_key_path,
+        test_temp_dir,
+        username,
+        ip_wait_time_extra=PXE_BOOT_IP_WAIT_TIME_EXTRA_SECONDS,
+    ) as ssh_client:
+        run_basic_checks(ssh_client, input_image_azl_release, test_temp_dir)
+
+
+@pytest.mark.skipif(
+    get_host_distro() == "azurelinux",
+    reason="PXE requires a network-enabled host UEFI firmware (EDK II NetworkPkg), which only Ubuntu hosts provide",
+)
+def test_pxe_bootstrap_efi_azl3(
+    docker_client: DockerClient,
+    image_customizer_container_url: str,
+    core_efi_azl3: Path,
+    ssh_key: Tuple[str, Path],
+    test_temp_dir: Path,
+    test_instance_name: str,
+    logs_dir: Path,
+    libvirt_conn: libvirt.virConnect,
+    close_list: List[Closeable],
+) -> None:
+    azl_release = 3
+    config_path = TEST_CONFIGS_DIR.joinpath("pxe-bootstrap-vm-azl3.yaml")
+
+    run_pxe_test(
+        docker_client,
+        image_customizer_container_url,
+        core_efi_azl3,
+        azl_release,
+        "bootstrap",
+        config_path,
+        ssh_key,
+        test_temp_dir,
+        test_instance_name,
+        logs_dir,
+        libvirt_conn,
+        close_list,
+    )
+
+
+@pytest.mark.skipif(
+    get_host_distro() == "azurelinux",
+    reason="PXE requires a network-enabled host UEFI firmware (EDK II NetworkPkg), which only Ubuntu hosts provide",
+)
+def test_pxe_bootstrap_efi_azl4(
+    docker_client: DockerClient,
+    image_customizer_container_url: str,
+    core_efi_azl4: Path,
+    ssh_key: Tuple[str, Path],
+    test_temp_dir: Path,
+    test_instance_name: str,
+    logs_dir: Path,
+    libvirt_conn: libvirt.virConnect,
+    close_list: List[Closeable],
+) -> None:
+    azl_release = 4
+    config_path = TEST_CONFIGS_DIR.joinpath("pxe-bootstrap-vm-azl4.yaml")
+    config_path = add_preview_features_to_config(config_path, "preview-distro-version", close_list)
+
+    run_pxe_test(
+        docker_client,
+        image_customizer_container_url,
+        core_efi_azl4,
+        azl_release,
+        "bootstrap",
+        config_path,
+        ssh_key,
+        test_temp_dir,
+        test_instance_name,
+        logs_dir,
+        libvirt_conn,
+        close_list,
+    )
+
+
+@pytest.mark.skipif(
+    get_host_distro() == "azurelinux",
+    reason="PXE requires a network-enabled host UEFI firmware (EDK II NetworkPkg), which only Ubuntu hosts provide",
+)
+def test_pxe_full_os_efi_azl3(
+    docker_client: DockerClient,
+    image_customizer_container_url: str,
+    core_efi_azl3: Path,
+    ssh_key: Tuple[str, Path],
+    test_temp_dir: Path,
+    test_instance_name: str,
+    logs_dir: Path,
+    libvirt_conn: libvirt.virConnect,
+    close_list: List[Closeable],
+) -> None:
+    azl_release = 3
+    config_path = TEST_CONFIGS_DIR.joinpath("pxe-full-os-vm-azl3.yaml")
+
+    run_pxe_test(
+        docker_client,
+        image_customizer_container_url,
+        core_efi_azl3,
+        azl_release,
+        "full-os",
+        config_path,
+        ssh_key,
+        test_temp_dir,
+        test_instance_name,
+        logs_dir,
+        libvirt_conn,
+        close_list,
+    )
+
+
+@pytest.mark.skipif(
+    get_host_distro() == "azurelinux",
+    reason="PXE requires a network-enabled host UEFI firmware (EDK II NetworkPkg), which only Ubuntu hosts provide",
+)
+def test_pxe_full_os_efi_azl4(
+    docker_client: DockerClient,
+    image_customizer_container_url: str,
+    core_efi_azl4: Path,
+    ssh_key: Tuple[str, Path],
+    test_temp_dir: Path,
+    test_instance_name: str,
+    logs_dir: Path,
+    libvirt_conn: libvirt.virConnect,
+    close_list: List[Closeable],
+) -> None:
+    azl_release = 4
+    config_path = TEST_CONFIGS_DIR.joinpath("pxe-full-os-vm-azl4.yaml")
+    config_path = add_preview_features_to_config(config_path, "preview-distro-version", close_list)
+
+    run_pxe_test(
+        docker_client,
+        image_customizer_container_url,
+        core_efi_azl4,
+        azl_release,
+        "full-os",
+        config_path,
+        ssh_key,
+        test_temp_dir,
+        test_instance_name,
+        logs_dir,
+        libvirt_conn,
+        close_list,
+    )

test/vmtests/vmtests/utils/host_utils.py

@@ -4,13 +4,14 @@
 
 def get_host_distro() -> str:
     file_path = "/etc/os-release"
-    name_value = ""
+    id_value = ""
     with open(file_path, "r") as file:
         for line in file:
             if line.startswith("ID="):
-                name_value = line.strip().split("=", 1)[1]  # Get the value part
+                id_value = line.strip().split("=", 1)[1]
                 break
-    if name_value == "":
+
+    if id_value == "":
         raise Exception("ID field not found in os-release file")
 
-    return name_value
+    return id_value

test/vmtests/vmtests/utils/imagecustomizer.py

@@ -209,6 +209,38 @@ def add_preview_features_to_config(config_path: Path, preview_feature: str, clos
     return path
 
 
+def add_pxe_bootstrap_base_url_to_config(
+    config_path: Path, bootstrap_base_url: str, close_list: List[Closeable]
+) -> Path:
+    """Modify an image customizer config file to set the PXE bootstrap base URL.
+
+    This URL is baked into the generated PXE artifacts (grub.cfg's root=live:<url>/image.iso), so it must match where
+    the bootstrap image is actually served at boot time.
+
+    Args:
+        config_path: Path to the base config file
+        bootstrap_base_url: Base URL where the bootstrap image is served at boot time
+        close_list: List of resources to be cleaned up
+
+    Returns:
+        Path to the modified config file
+    """
+    config_str = config_path.read_text()
+    config = yaml.safe_load(config_str)
+
+    pxe = dict_get_or_set(config, "pxe", {})
+    pxe["bootstrapBaseUrl"] = bootstrap_base_url
+
+    # Write out new config file to a temporary file.
+    fd, modified_config_path = tempfile.mkstemp(prefix=config_path.name + "~", suffix=".tmp", dir=config_path.parent)
+    with fdopen(fd, mode="w") as file:
+        yaml.safe_dump(config, file)
+
+    path = Path(modified_config_path)
+    close_list.append(RemoveFileOnClose(path))
+    return path
+
+
 def dict_get_or_set(dictionary: Dict[Any, Any], value_name: str, default: Any = None) -> Any:
     value = dictionary.get(value_name)
     if value is None: