Harden HEKI's kernel text patch fuction

litebox_platform_lvbs/src/mshv/mem_integrity.rs

@@ -626,7 +626,10 @@ const JMP32_INSN_SIZE: u8 = 5;
 /// Refer [Linux](https://elixir.bootlin.com/linux/v6.6.85/source/arch/x86/kernel/alternative.c#L2164)
 pub fn validate_text_poke_bp_batch(patch_data: &HekiPatch, precomputed_patch: &HekiPatch) -> bool {
     // step 1
-    if patch_data.size == 1 && patch_data.code[0] == INT3_INSN_OPCODE {
+    if patch_data.size == 1
+        && patch_data.code[0] == INT3_INSN_OPCODE
+        && patch_data.pa[0] == precomputed_patch.pa[0]
+    {
         return true;
     }
 
@@ -637,7 +640,15 @@ pub fn validate_text_poke_bp_batch(patch_data: &HekiPatch, precomputed_patch: &H
             || (patch_data.pa[0] == precomputed_patch.pa[1]
                 && (precomputed_patch.pa[0] + 1).is_multiple_of(Size4KiB::SIZE)))
     {
-        1 // step 2
+        // step 2. pa[1] is dereferenced by `apply_vtl0_text_patch` only when
+        // `precomputed_patch.pa[0] + 1` is not page-aligned. In that case, it must
+        // equal `precomputed_patch.pa[1]`.
+        if !(precomputed_patch.pa[0] + 1).is_multiple_of(Size4KiB::SIZE)
+            && patch_data.pa[1] != precomputed_patch.pa[1]
+        {
+            return false;
+        }
+        1
     } else {
         return false;
     };

litebox_platform_lvbs/src/mshv/vsm.rs

@@ -1061,21 +1061,22 @@ impl Vtl0KernelInfo {
         self.system_certs.get().map(|b| &**b)
     }
 
-    // This function finds the precomputed patch data corresponding to the input patch data.
-    // We need this because each step of `mshv_vsm_patch_data`/`text_poke_bp_batch` only
-    // provides a part of the patch data and addresses (`patch[0]` or `patch[1..patch_size-1]`).
+    /// This function finds the precomputed patch data corresponding to the input patch data.
+    ///
+    /// Each step of `text_poke_bp_batch` only exposes a portion of the target's address range,
+    /// so we look up in the precomputed map by two keys derived from `patch_data.pa[0]`:
+    /// - `pa[0]` matches step 1 or 3 (target's first byte) and, for a precomputed patch that
+    ///   straddles at offset 1, step 2.
+    /// - `pa[0] - 1` matches step 2 where `patch.pa[0] == precomputed.pa[0] + 1`.
+    ///
+    /// No legitimate step requires looking up by `patch.pa[1]`.
     pub fn find_precomputed_patch(&self, patch_data: &HekiPatch) -> Option<HekiPatch> {
         self.precomputed_patches
             .get(PhysAddr::new(patch_data.pa[0]))
             .or_else(|| {
                 self.precomputed_patches
                     .get(PhysAddr::new(patch_data.pa[0].saturating_sub(1)))
             })
-            .or_else(|| {
-                self.precomputed_patches
-                    .get(PhysAddr::new(patch_data.pa[1]))
-            })
-            .or(None)
     }
 }