Frida windows check shadow fix (#2159)

* Fix check_shadow and add additional tests

* add some additional documentation

Author: mineo333

libafl_frida/src/alloc.rs

@@ -434,67 +434,107 @@ impl Allocator {
         (shadow_mapping_start, (end - start) / 8 + 1)
     }
 
+    #[inline]
+    #[must_use]
+    fn check_shadow_aligned(&mut self, address: *const c_void, size: usize) -> bool {
+        assert_eq!((address as usize) & 7, 0, "check_shadow_aligned used when address is not aligned. Use check_shadow");
+        assert_eq!(size & 7, 0, "check_shadow_aligned used when size is not aligned. Use check_shadow");
+
+        if size == 0 {
+            return true;
+        }
+
+        let shadow_addr = map_to_shadow!(self, (address as usize));
+        let shadow_size = size>>3;
+        let buf =
+            unsafe { std::slice::from_raw_parts_mut(shadow_addr as *mut u8, shadow_size) };
+        let (prefix, aligned, suffix) = unsafe { buf.align_to::<u128>() };
+        if !prefix.iter().all(|&x| x == 0xff)
+            || !suffix.iter().all(|&x| x == 0xff)
+            || !aligned
+                .iter()
+                .all(|&x| x == 0xffffffffffffffffffffffffffffffffu128)
+        {
+            return false;
+        }
+
+        return true;
+    }
     /// Checks whether the given address up till size is valid unpoisoned shadow memory.
     /// TODO: check edge cases
     #[inline]
     #[must_use]
     pub fn check_shadow(&mut self, address: *const c_void, size: usize) -> bool {
+        //the algorithm for check_shadow is as follows:
+        //1. we first check if its managed. if is not then exit
+        //2. we check if it is aligned. this should be 99% of accesses. If it is do an aligned check and leave
+        //3. if it is not split the check into 3 parts: the pre-aligment bytes, the aligned portion, and the post alignment posts
+        //3. The prealignment bytes are the unaligned bytes (if any) located in the qword preceding the aligned portion. Perform a specialied check to ensure that the bytes from [start, align(start, 8)) are valid. In this case align(start,8) aligns start to the next 8 byte boundary.
+        //4. The aligned check is where the address and the size is 8 byte aligned. Use check_shadow_aligned to check it
+        //5. The post-alignment is the same as pre-alignment except it is the qword following the aligned portion. Use a specialized check to ensure that [end & ~7, end) is valid. 
+
         if size == 0 || !self.is_managed(address as *mut c_void) {
             return true;
         }
-        let address = address as usize;
-        let shadow_size = size / 8;
-
-        let shadow_addr = map_to_shadow!(self, address);
-
-        // self.map_shadow_for_region(address, address + size, false);
 
-        log::info!(
-            "check_shadow: {:x}, {:x}, {:x}, {:x}",
-            address,
-            shadow_size,
-            shadow_addr,
+        log::trace!(
+            "check_shadow: {:x}, {:x}",
+            address as usize,
             size
         );
 
-        let offset = address & 7;
-        // if we are not aligned to 8 bytes, we need to check the high bits of the shadow
-        if offset != 0 {
-            let val = (unsafe { (shadow_addr as *const u16).read() }) >> offset;
-            let mask = ((1 << (size % 9)) - 1) as u16;
-            if val & mask != mask {
+        //fast path. most buffers are likely 8 byte aligned in size and address
+        if (address as usize) & 7 == 0 && size & 7 == 0 {
+            return self.check_shadow_aligned(address, size); 
+        }
+
+        //slow path. check everything
+        let start_address = address as usize;
+        let end_address = start_address + size;
+        
+        //8 byte align the start/end so we can use check_shadow_aligned for the majority of it
+        //in the case of subqword accesses (i.e,, the entire access is located within 1 qword), aligned_start > aligned_end naturally
+        let aligned_start = (start_address + 7) & !7;
+        let aligned_end = end_address & !7;
+
+        let start_offset = start_address & 7;
+        let end_offset = end_address & 7;
+
+
+        //if the start is unaligned
+        if start_address != aligned_start {
+            let start_shadow = map_to_shadow!(self, start_address);
+            //the purpose of the min is to account for sub-qword accesses. If the access is subqword then size will be less than the distance to the end of the qword
+            
+        
+            let start_mask: u8 = 0xff << std::cmp::min(size, 8-start_offset);
+
+            if unsafe { (start_shadow as *const u8).read() } & start_mask != start_mask {
                 return false;
             }
         }
 
-        if size >= 8 {
-            let buf =
-                unsafe { std::slice::from_raw_parts_mut(shadow_addr as *mut u8, shadow_size) };
-            let (prefix, aligned, suffix) = unsafe { buf.align_to::<u128>() };
-            if prefix.iter().all(|&x| x == 0xff)
-                && suffix.iter().all(|&x| x == 0xff)
-                && aligned
-                    .iter()
-                    .all(|&x| x == 0xffffffffffffffffffffffffffffffffu128)
-            {
-                if size % 8 != 0 {
-                    let val = unsafe { ((shadow_addr + shadow_size) as *mut u8).read() };
-                    let mask = (((1 << (size % 8)) - 1) as u8).rotate_left(8 - (size % 8) as u32);
-                    if val & mask != mask {
-                        return false;
-                    }
-                }
-                return true;
+        //if this is not true then it must be a subqword access as the start will be larger than the end
+        if aligned_start <= aligned_end {
+            if !self.check_shadow_aligned(aligned_start as *const c_void, aligned_end-aligned_start) {
+                return false;
             }
-        }
-        if size % 8 != 0 {
-            let val = unsafe { ((shadow_addr + shadow_size) as *mut u8).read() };
-            let mask = (((1 << (size % 8)) - 1) as u8).rotate_left(8 - (size % 8) as u32);
-            if val & mask == mask {
-                return true;
+
+            if end_address != aligned_end {
+                let end_shadow = map_to_shadow!(self, end_address);
+            
+                let end_mask = 0xff << (8 - end_offset); //we want to check from the beginning of the qword to the offset
+                if unsafe { (end_shadow as *const u8).read() } & end_mask != end_mask {
+                    return false;
+                }
             }
+
+
         }
-        return false;
+        // self.map_shadow_for_region(address, address + size, false);
+
+        
+        return true;
     }
     /// Maps the address to a shadow address
     #[inline]
@@ -570,10 +610,9 @@ impl Allocator {
                 let start = range.memory_range().base_address().0 as usize;
                 let end = start + range.memory_range().size();
                 log::trace!(
-                    "Start: {:#x}, end: {:#x}, prot: {:?}",
+                    "Start: {:#x}, end: {:#x}",
                     start,
                     end,
-                    range.protection()
                 );
                 occupied_ranges.push((start, end));
                 let base: usize = 2;
@@ -743,6 +782,17 @@ fn check_shadow() {
     assert!(allocator.check_shadow(unsafe { allocation.offset(1) }, 8) == false);
     assert!(allocator.check_shadow(unsafe { allocation.offset(2) }, 8) == false);
     assert!(allocator.check_shadow(unsafe { allocation.offset(3) }, 8) == false);
+    let allocation = unsafe { allocator.alloc(0xc, 0) };
+    assert!(allocator.check_shadow(unsafe { allocation.offset(4)}, 8) == true);
+    //subqword access
+    assert!(allocator.check_shadow(unsafe { allocation.offset(3)}, 2) == true);
+    //unaligned access
+    assert!(allocator.check_shadow(unsafe { allocation.offset(3)}, 8) == true);
+    let allocation = unsafe { allocator.alloc(0x20, 0) };
+    //access with unaligned parts at the beginning and end
+    assert!(allocator.check_shadow(unsafe { allocation.offset(10)}, 21) == true);
+    //invalid, unaligned access
+    assert!(allocator.check_shadow(unsafe { allocation.offset(10)}, 29) == false);
     let allocation = unsafe { allocator.alloc(4, 0) };
     assert!(allocator.check_shadow(allocation, 1) == true);
     assert!(allocator.check_shadow(allocation, 2) == true);