Allow signature scanning to cross page boundaries

src/signature.rs

@@ -1,6 +1,6 @@
 //! Support for finding patterns in a process's memory.
 
-use core::mem::{self, MaybeUninit};
+use core::{mem, slice};
 
 use bytemuck::AnyBitPattern;
 
@@ -194,23 +194,86 @@
         process: &Process,
         (addr, len): (impl Into<Address>, u64),
     ) -> Option<Address> {
+        const MEM_SIZE: usize = 0x1000;
         let mut addr: Address = Into::into(addr);
-        // TODO: Handle the case where a signature may be cut in half by a page
-        // boundary.
         let overall_end = addr.value() + len;
-        let mut buf = [MaybeUninit::uninit(); 4 << 10];
+
+        // The sigscan essentially works by reading one memory page (0x1000 bytes)
+        // at a time and looking for the signature in each page. We create a buffer
+        // sligthly larger than 0x1000 bytes in order to accomodate the size of
+        // the memory page + the signature - 1. The very first bytes of the
+        // buffer are intended to be used as the tail of the previous memory page.
+        // This allows to scan across the memory page boundaries.
+
+        // We should use N - 1 but we resort to MEM_SIZE - 1 to avoid using [feature(generic_const_exprs)]
+        #[repr(packed)]
+        struct Buffer<const N: usize> {
+            _head: [u8; N], 
+            _buffer: [u8; MEM_SIZE - 1] 
+        }
+
+        let mut global_buffer =
+            // SAFETY: Using mem::zeroed is safe in this instance as the Buffer struct
+            // only contains u8, for which zeroed data represent a valid bit pattern.
+            unsafe {
+                mem::zeroed::<Buffer<N>>()
+        };
+
+        let buf = 
+            // SAFETY: The buffer is guaranteed to be initialized due
+            // to using mem::zeroed() so we can safely return an u8 slice of it.
+            unsafe {
+                slice::from_raw_parts_mut(&mut global_buffer as *mut _ as *mut u8, size_of::<Buffer<N>>())
+        };
+
+        let first =
+            // SAFETY: The buffer is guaranteed to be initialized due
+            // to using mem::zeroed() so we can safely return an u8 slice of it.
+            unsafe {
+                slice::from_raw_parts_mut(buf.as_mut_ptr(), N - 1)
+        };
+
+        let last =
+            // SAFETY: The buffer is guaranteed to be initialized due
+            // to using mem::zeroed() so we can safely return an u8 slice of it.
+            unsafe {
+                slice::from_raw_parts_mut(buf.as_mut_ptr().add(MEM_SIZE), N - 1)
+        };
+
+        let mut last_page_success = false;
         while addr.value() < overall_end {
             // We round up to the 4 KiB address boundary as that's a single
             // page, which is safe to read either fully or not at all. We do
-            // this to do a single read rather than many small ones as the
-            // syscall overhead is a quite high.
-            let end = (addr.value() & !((4 << 10) - 1)) + (4 << 10).min(overall_end);
-            let len = end - addr.value();
-            let current_read_buf = &mut buf[..len as usize];
-            if let Ok(current_read_buf) = process.read_into_uninit_buf(addr, current_read_buf) {
-                if let Some(pos) = self.scan(current_read_buf) {
-                    return Some(addr.add(pos as u64));
+            // this to reduce the number of syscalls as much as possible, as the
+            // syscall overhead is quite high.
+            let end = ((addr.value() & !((4 << 10) - 1)) + (4 << 10)).min(overall_end);
+            let len = end.saturating_sub(addr.value()) as usize;
+
+            // If we read the previous memory page successfully, then we can copy the last
+            // elements to the start of the buffer. Otherwise, we just fill it with zeroes.
+            if last_page_success {
+                first.copy_from_slice(last);
+            }
+
+            if process.read_into_buf(addr, &mut buf[N - 1..][..len]).is_ok() {
+                // If the previous memory page has been read successfully, then we have copied
+                // the last N - 1 bytes to the start of the buf array, and we need to check
+                // starting from those in order to correctly identify possible signatures crossing
+                // memory page boundaries
+                if last_page_success {
+                    if let Some(pos) = self.scan(&mut buf[..len + N - 1]) {
+                        return Some(addr.add(pos as u64).add_signed(-(N as i64 - 1)));
+                    }
+                // If the prevbious memory page wasn't read successfully, the first N - 1
+                // bytes are excluded from the signature verification 
+                } else {
+                    if let Some(pos) = self.scan(&mut buf[N - 1..][..len]) {
+                        return Some(addr.add(pos as u64));
+                    }
                 }
+
+                last_page_success = true;
+
             };
             addr = Address::new(end);
         }