From 973be920b165e720b1f76f184ecd1f8cedf2bf60 Mon Sep 17 00:00:00 2001
From: smit-gardhariya <gardhariya.smit@gmail.com>
Date: Wed, 20 Nov 2024 13:58:54 +0530
Subject: [PATCH] Add support for mariner-3.0 in dom0 kernel installer

dom0 image based on mariner 3.0 has different
parameter that need to be changed while installing kernel
under different file. Add support for it under kernel
installer transformer.

Signed-off-by: Smit Gardhariya <sgardhariya@microsoft.com>
---
 lisa/transformers/dom0_kernel_installer.py | 49 ++++++++++++++++++----
 1 file changed, 40 insertions(+), 9 deletions(-)

diff --git a/lisa/transformers/dom0_kernel_installer.py b/lisa/transformers/dom0_kernel_installer.py
index 19a63c0b48..0b8d00d632 100644
--- a/lisa/transformers/dom0_kernel_installer.py
+++ b/lisa/transformers/dom0_kernel_installer.py
@@ -80,6 +80,9 @@ def install(self) -> str:
         uname = node.tools[Uname]
         current_kernel = uname.get_linux_information().kernel_version_raw
 
+        node_info = node.get_information()
+        distro = node_info.get("distro_version", "")
+
         # Kernel absolute path: /home/user/vmlinuz-5.15.57.1+
         # Naming convention : vmlinuz-<version>
         new_kernel = os.path.basename(kernel_image_path).split("-")[1].strip()
@@ -124,10 +127,19 @@ def install(self) -> str:
                 node.get_pure_path(f"/boot/initrd.img-{new_kernel}"),
             )
         else:
+            # Mariner 2.0 initrd
+            target = f"/boot/initrd.img-{current_kernel}"
+            link = f"/boot/initrd.img-{new_kernel}"
+
+            if distro == "Microsoft Azure Linux 3.0":
+                # Mariner 3.0 initrd
+                target = f"/boot/initramfs-{current_kernel}.img"
+                link = f"/boot/initramfs-{new_kernel}.img"
+
             ln = node.tools[Ln]
             ln.create_link(
-                target=f"/boot/initrd.img-{current_kernel}",
-                link=f"/boot/initrd.img-{new_kernel}",
+                target=target,
+                link=link,
             )
 
         if kernel_config_path:
@@ -148,6 +160,7 @@ def install(self) -> str:
             node,
             current_kernel,
             new_kernel,
+            distro,
         )
 
         return new_kernel
@@ -185,10 +198,13 @@ def install(self) -> str:
         uname = node.tools[Uname]
         current_kernel = uname.get_linux_information().kernel_version_raw
 
+        node_info = node.get_information()
+        distro = node_info.get("distro_version", "")
         _update_mariner_config(
             node,
             current_kernel,
             new_kernel,
+            distro,
         )
 
         return new_kernel
@@ -211,22 +227,37 @@ def _update_mariner_config(
     node: Node,
     current_kernel: str,
     new_kernel: str,
+    distro: str,
 ) -> None:
-    mariner_config: str = "/boot/mariner-mshv.cfg"
     sed = node.tools[Sed]
 
-    # Modify the /boot/mariner-mshv.cfg to point new kernel binary
+    # Param for Dom0 2.0 kernel installation
+    mariner_config = "/boot/mariner-mshv.cfg"
+    vmlinuz_regexp = f"vmlinuz-{current_kernel}"
+    vmlinuz_replacement = f"vmlinuz-{new_kernel}"
+    initrd_regexp = f"mariner_initrd_mshv=initrd.img-{current_kernel}"
+    initrd_replacement = f"mariner_initrd_mshv=initrd.img-{new_kernel}"
+
+    if distro == "Microsoft Azure Linux 3.0":
+        # Change param for Dom0 3.0 kernel installation
+        mariner_config = "/boot/grub2/grub.cfg"
+        initrd_regexp = f"initramfs-{current_kernel}.img"
+        initrd_replacement = f"initramfs-{new_kernel}.img"
+
+    # Modify file to point new kernel binary
     sed.substitute(
-        regexp=f"mariner_linux_mshv=vmlinuz-{current_kernel}",
-        replacement=f"mariner_linux_mshv=vmlinuz-{new_kernel}",
+        regexp=vmlinuz_regexp,
+        replacement=vmlinuz_replacement,
         file=mariner_config,
         sudo=True,
     )
 
-    # Modify the /boot/mariner-mshv.cfg to point new initrd binary
+    # Modify file to point new initrd binary
     sed.substitute(
-        regexp=f"mariner_initrd_mshv=initrd.img-{current_kernel}",
-        replacement=f"mariner_initrd_mshv=initrd.img-{new_kernel}",
+        regexp=initrd_regexp,
+        replacement=initrd_replacement,
         file=mariner_config,
         sudo=True,
     )
+    res = node.execute(f"cat {mariner_config}", sudo=True)
+    print(f"MARINER CONFIG: {res.stdout}")