Add Azure Linux 4.0 PXE boot support for LiveOS images

test/vmtests/vmtests/imagecustomizer/test_pxe.py

@@ -0,0 +1,274 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import logging
+import platform
+import random
+import string
+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,
+    )
+
+    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_tar_path,
+        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: