From 53e2945680e8383a6147b862e2a2f8c0a150772b Mon Sep 17 00:00:00 2001
From: Smit Gardhariya <113590758+smit-gardhariya@users.noreply.github.com>
Date: Fri, 22 Nov 2024 10:33:48 +0530
Subject: [PATCH] dom0_kernel_installer: Add support for mariner-3.0 (#3522)

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 | 46 +++++++++++++++++-----
 1 file changed, 37 insertions(+), 9 deletions(-)

diff --git a/lisa/transformers/dom0_kernel_installer.py b/lisa/transformers/dom0_kernel_installer.py
index 19a63c0b48..011d1f368d 100644
--- a/lisa/transformers/dom0_kernel_installer.py
+++ b/lisa/transformers/dom0_kernel_installer.py
@@ -9,6 +9,7 @@
 
 from lisa import schema
 from lisa.node import Node
+from lisa.operating_system import CBLMariner
 from lisa.tools import Cp, Echo, Ln, Ls, Sed, Tar, Uname
 from lisa.util import field_metadata
 
@@ -80,6 +81,8 @@ def install(self) -> str:
         uname = node.tools[Uname]
         current_kernel = uname.get_linux_information().kernel_version_raw
 
+        mariner_version = int(node.os.information.version.major)
+
         # 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 3.0 initrd
+            target = f"/boot/initramfs-{current_kernel}.img"
+            link = f"/boot/initramfs-{new_kernel}.img"
+
+            if isinstance(node.os, CBLMariner) and mariner_version == 2:
+                # Mariner 2.0 initrd
+                target = f"/boot/initrd.img-{current_kernel}"
+                link = f"/boot/initrd.img-{new_kernel}"
+
             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,
+            mariner_version,
         )
 
         return new_kernel
@@ -185,10 +198,12 @@ def install(self) -> str:
         uname = node.tools[Uname]
         current_kernel = uname.get_linux_information().kernel_version_raw
 
+        mariner_version = int(node.os.information.version.major)
         _update_mariner_config(
             node,
             current_kernel,
             new_kernel,
+            mariner_version,
         )
 
         return new_kernel
@@ -211,22 +226,35 @@ def _update_mariner_config(
     node: Node,
     current_kernel: str,
     new_kernel: str,
+    mariner_version: int,
 ) -> 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 3.0 kernel installation
+    mariner_config = "/boot/grub2/grub.cfg"
+    vmlinuz_regexp = f"vmlinuz-{current_kernel}"
+    vmlinuz_replacement = f"vmlinuz-{new_kernel}"
+    initrd_regexp = f"initramfs-{current_kernel}.img"
+    initrd_replacement = f"initramfs-{new_kernel}.img"
+
+    if isinstance(node.os, CBLMariner) and mariner_version == 2:
+        # Change param for Dom0 2.0 kernel installation
+        mariner_config = "/boot/mariner-mshv.cfg"
+        initrd_regexp = f"mariner_initrd_mshv=initrd.img-{current_kernel}"
+        initrd_replacement = f"mariner_initrd_mshv=initrd.img-{new_kernel}"
+
+    # 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,
     )