From 23cf4e3c7693ab391e0e92fc4a52ed1fe2836fb8 Mon Sep 17 00:00:00 2001
From: Sangho Lee <sanghle@microsoft.com>
Date: Mon, 16 Mar 2026 00:43:51 +0000
Subject: [PATCH] fix heki patch offset bugs

---
 litebox_platform_lvbs/src/mshv/vsm.rs | 35 +++++++++++++++------------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/litebox_platform_lvbs/src/mshv/vsm.rs b/litebox_platform_lvbs/src/mshv/vsm.rs
index 91a8fb57a..458fb2f66 100644
--- a/litebox_platform_lvbs/src/mshv/vsm.rs
+++ b/litebox_platform_lvbs/src/mshv/vsm.rs
@@ -889,13 +889,12 @@ fn apply_vtl0_text_patch(heki_patch: HekiPatch) -> Result<(), VsmError> {
             return Err(VsmError::Vtl0CopyFailed);
         }
     } else {
-        let (patch_first, patch_second) = heki_patch.code.split_at(bytes_in_first_page);
+        let (patch_first, patch_second) =
+            heki_patch.code[..usize::from(heki_patch.size)].split_at(bytes_in_first_page);
 
         unsafe {
-            if !crate::platform_low().copy_slice_to_vtl0_phys(
-                heki_patch_pa_0 + patch_target_page_offset as u64,
-                patch_first,
-            ) || !crate::platform_low().copy_slice_to_vtl0_phys(heki_patch_pa_1, patch_second)
+            if !crate::platform_low().copy_slice_to_vtl0_phys(heki_patch_pa_0, patch_first)
+                || !crate::platform_low().copy_slice_to_vtl0_phys(heki_patch_pa_1, patch_second)
             {
                 return Err(VsmError::Vtl0CopyFailed);
             }
@@ -1683,24 +1682,30 @@ impl PatchDataMap {
 
         // the buffer looks like below:
         // [`HekiPatchInfo`, [`HekiPatch`, ...], `HekiPatchInfo`, [`HekiPatch`, ...], ...]
-        // each `HekiPatchInfo` contains the number of `HekiPatch` structures (`patch_index`) that follow it.
+        // Each `HekiPatchInfo`'s `patch_index` field specifies the number of `HekiPatch` entries that follow it.
+        // The buffer may have trailing bytes (from page-aligned VTL0 ranges) that don't form a valid record.
         let mut index: usize = 0;
-        while index <= patch_info_buf.len() - core::mem::size_of::<HekiPatchInfo>() {
-            let patch_info = HekiPatchInfo::try_from_bytes(
+        while index + core::mem::size_of::<HekiPatchInfo>() <= patch_info_buf.len() {
+            let Some(patch_info) = HekiPatchInfo::try_from_bytes(
                 &patch_info_buf[index..index + core::mem::size_of::<HekiPatchInfo>()],
-            )
-            .ok_or(PatchDataMapError::InvalidHekiPatchInfo)?;
+            ) else {
+                // Remaining bytes don't form a valid header. End of meaningful patch data.
+                break;
+            };
 
             let patch_index: usize = patch_info.patch_index.truncate();
             let total_patch_size = core::mem::size_of::<HekiPatch>()
                 .checked_mul(patch_index)
                 .ok_or(PatchDataMapError::InvalidHekiPatchInfo)?;
-            index = index
-                .checked_add(core::mem::size_of::<HekiPatchInfo>() + total_patch_size)
-                .filter(|&x| x <= patch_info_buf.len())
+            let patches_start = index
+                .checked_add(core::mem::size_of::<HekiPatchInfo>())
+                .ok_or(PatchDataMapError::InvalidHekiPatchInfo)?;
+            let patches_end = patches_start
+                .checked_add(total_patch_size)
+                .filter(|&end| end <= patch_info_buf.len())
                 .ok_or(PatchDataMapError::InvalidHekiPatchInfo)?;
 
-            for patch in patch_info_buf[index - total_patch_size..index]
+            for patch in patch_info_buf[patches_start..patches_end]
                 .chunks(core::mem::size_of::<HekiPatch>())
                 .map(HekiPatch::try_from_bytes)
             {
@@ -1745,7 +1750,7 @@ impl PatchDataMap {
                     }
                 }
             }
-            index += total_patch_size;
+            index = patches_end;
         }
 
         Ok(())