Fix vmmap change_prot fragmentation and alignment bugs

src/cage/src/memory/vmmap.rs

@@ -565,40 +565,86 @@ impl VmmapOps for Vmmap {
         let new_region_end_page = page_num + npages;
         let new_region_start_page = page_num;
 
-        // Store intervals that need to be inserted after iteration
-        let mut to_insert = Vec::new();
+        // Collect information about overlapping entries that need to be modified
+        let mut entries_to_modify = Vec::new();
 
-        // Iterate over overlapping entries
         for (overlap_interval, entry) in self
             .entries
-            .overlapping_mut(ie(new_region_start_page, new_region_end_page))
+            .overlapping(ie(new_region_start_page, new_region_end_page))
         {
-            let mut ent_start = overlap_interval.start();
+            let ent_start = overlap_interval.start();
             let ent_end = overlap_interval.end();
 
-            // Case 1: Entry starts before region but extends into it
-            if ent_start < new_region_start_page && ent_end > new_region_start_page {
-                to_insert.push(ie(new_region_start_page, ent_end));
-                ent_start = new_region_start_page;
-            }
+            // Clone the entry to work with
+            let original_entry = entry.clone();
 
-            // Case 2: Entry extends beyond region end
-            if ent_start < new_region_end_page && ent_end > new_region_end_page {
-                to_insert.push(ie(ent_start, new_region_end_page));
-            } else {
-                // Case 3: Entry is fully contained - update protection
-                entry.prot = new_prot;
-            }
+            // Calculate the three potential parts:
+            // 1. Before the target region (keep old protection)
+            // 2. Inside the target region (apply new protection)
+            // 3. After the target region (keep old protection)
+
+            let overlap_start = ent_start.max(new_region_start_page);
+            let overlap_end = ent_end.min(new_region_end_page);
+
+            // Store the parts we need to create
+            entries_to_modify.push((
+                ent_start,
+                ent_end,
+                overlap_start,
+                overlap_end,
+                original_entry,
+            ));
         }
 
-        // Insert new intervals with updated protection
-        for interval in to_insert {
-            // Get and clone the entry at the start of the interval
-            let mut interval_val = self.entries.get_at_point(interval.start()).unwrap().clone();
-            // Update protection
-            interval_val.prot = new_prot;
-            // Insert the new interval
-            let _ = self.entries.insert_overwrite(interval, interval_val);
+        // Now modify the entries
+        for (ent_start, ent_end, overlap_start, overlap_end, original_entry) in entries_to_modify {
+            // Remove the original entry
+            let _ = self.entries.remove_overlapping(ie(ent_start, ent_end));
+
+            // Check if protection is actually changing
+            let prot_unchanged = original_entry.prot == new_prot;
+
+            if prot_unchanged {
+                // Protection isn't changing, keep the entry as-is (no fragmentation)
+                let _ = self
+                    .entries
+                    .insert_overwrite(ie(ent_start, ent_end), original_entry);
+            } else {
+                // Protection is changing, need to split
+
+                // Part 1: Before the target region (if exists)
+                if ent_start < overlap_start {
+                    let mut before_entry = original_entry.clone();
+                    before_entry.page_num = ent_start;
+                    before_entry.npages = overlap_start - ent_start;
+                    // Keep original protection
+                    let _ = self
+                        .entries
+                        .insert_overwrite(ie(ent_start, overlap_start), before_entry);
+                }
+
+                // Part 2: Inside the target region (apply new protection)
+                if overlap_start < overlap_end {
+                    let mut inside_entry = original_entry.clone();
+                    inside_entry.page_num = overlap_start;
+                    inside_entry.npages = overlap_end - overlap_start;
+                    inside_entry.prot = new_prot;
+                    let _ = self
+                        .entries
+                        .insert_overwrite(ie(overlap_start, overlap_end), inside_entry);
+                }
+
+                // Part 3: After the target region (if exists)
+                if overlap_end < ent_end {
+                    let mut after_entry = original_entry.clone();
+                    after_entry.page_num = overlap_end;
+                    after_entry.npages = ent_end - overlap_end;
+                    // Keep original protection
+                    let _ = self
+                        .entries
+                        .insert_overwrite(ie(overlap_end, ent_end), after_entry);
+                }
+            }
         }
     }
 
@@ -960,7 +1006,14 @@ impl VmmapOps for Vmmap {
 
             let gap_size = aligned_end_page - aligned_start_page;
             if gap_size >= rounded_num_pages {
-                return Some(ie(aligned_end_page - rounded_num_pages, aligned_end_page));
+                // Calculate the aligned end position
+                let result_end = aligned_end_page;
+                // Calculate aligned start by ensuring it's a multiple of pages_per_map
+                let result_start = result_end - rounded_num_pages;
+                // Verify both boundaries are properly aligned
+                debug_assert!(result_start % pages_per_map == 0);
+                debug_assert!(result_end % pages_per_map == 0);
+                return Some(ie(result_start, result_end));
             }
         }
 

tests/unit-tests/memory_tests/deterministic/mmap_aligned.c

@@ -0,0 +1,58 @@
+// Test: mmap with alignment requirements
+// Verifies proper page alignment for memory allocations
+#include <sys/mman.h>
+#include <stdio.h>
+#include <assert.h>
+#include <stdint.h>
+
+#define PAGESIZE 4096
+#define TEST_PAGES 10
+
+int main(void) {
+    size_t alloc_size = TEST_PAGES * PAGESIZE;
+
+    // Test 1: Basic allocation with page alignment verification
+    unsigned char *p1 = mmap(NULL, alloc_size, PROT_READ | PROT_WRITE,
+                             MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    assert(p1 != MAP_FAILED && "first mmap failed");
+
+    // All mmap allocations should be page-aligned
+    assert((uintptr_t)p1 % PAGESIZE == 0 && "first allocation not page-aligned");
+
+    // Write and verify data
+    for (int i = 0; i < TEST_PAGES; i++) {
+        p1[i * PAGESIZE] = 0xA0 + i;
+    }
+    for (int i = 0; i < TEST_PAGES; i++) {
+        assert(p1[i * PAGESIZE] == 0xA0 + i && "first allocation data verification failed");
+    }
+
+    // Test 2: Second allocation should also be page-aligned
+    unsigned char *p2 = mmap(NULL, alloc_size, PROT_READ | PROT_WRITE,
+                             MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    assert(p2 != MAP_FAILED && "second mmap failed");
+    assert((uintptr_t)p2 % PAGESIZE == 0 && "second allocation not page-aligned");
+
+    // Write and verify data for second allocation
+    for (int i = 0; i < TEST_PAGES; i++) {
+        p2[i * PAGESIZE] = 0xB0 + i;
+    }
+    for (int i = 0; i < TEST_PAGES; i++) {
+        assert(p2[i * PAGESIZE] == 0xB0 + i && "second allocation data verification failed");
+    }
+
+    // Test 3: Third allocation to verify consistency
+    unsigned char *p3 = mmap(NULL, alloc_size, PROT_READ | PROT_WRITE,
+                             MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    assert(p3 != MAP_FAILED && "third mmap failed");
+    assert((uintptr_t)p3 % PAGESIZE == 0 && "third allocation not page-aligned");
+
+    // Cleanup
+    assert(munmap(p3, alloc_size) == 0 && "munmap p3 failed");
+    assert(munmap(p2, alloc_size) == 0 && "munmap p2 failed");
+    assert(munmap(p1, alloc_size) == 0 && "munmap p1 failed");
+
+    printf("mmap_aligned test: PASS\n");
+    return 0;
+}
+

tests/unit-tests/memory_tests/deterministic/mprotect_boundary.c

@@ -0,0 +1,61 @@
+// Test: mprotect on exact page boundaries
+// Verifies correct handling of single-page and precise boundary modifications
+#include <sys/mman.h>
+#include <stdio.h>
+#include <assert.h>
+
+#define PAGESIZE 4096
+#define NUMPAGES 10
+
+int main(void) {
+    // Allocate 10 pages with READ|WRITE protection
+    size_t len = PAGESIZE * NUMPAGES;
+    unsigned char *p = mmap(NULL, len, PROT_READ | PROT_WRITE,
+                            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    assert(p != MAP_FAILED && "mmap failed");
+
+    // Write test data
+    for (int i = 0; i < NUMPAGES; i++) {
+        p[i * PAGESIZE] = 0x80 + i;
+    }
+
+    // Change protection on a single page in the middle (page 5)
+    assert(mprotect(p + (5 * PAGESIZE), PAGESIZE, PROT_READ) == 0 && "mprotect single page failed");
+
+    // Verify we can read from the protected page
+    assert(p[5 * PAGESIZE] == 0x85 && "read from single protected page failed");
+
+    // Verify we can write to pages before (0-4)
+    for (int i = 0; i < 5; i++) {
+        p[i * PAGESIZE] = 0x90 + i;
+    }
+
+    // Verify we can write to pages after (6-9)
+    for (int i = 6; i < NUMPAGES; i++) {
+        p[i * PAGESIZE] = 0x90 + i;
+    }
+
+    // Verify the writes
+    assert(p[0] == 0x90 && "boundary write verification failed");
+    assert(p[4 * PAGESIZE] == 0x94 && "boundary write verification failed");
+    assert(p[6 * PAGESIZE] == 0x96 && "boundary write verification failed");
+    assert(p[9 * PAGESIZE] == 0x99 && "boundary write verification failed");
+
+    // Change protection on first page only
+    assert(mprotect(p, PAGESIZE, PROT_READ) == 0 && "mprotect first page failed");
+
+    // Verify we can read from first page
+    assert(p[0] == 0x90 && "read from first protected page failed");
+
+    // Change protection on last page only
+    assert(mprotect(p + (9 * PAGESIZE), PAGESIZE, PROT_READ) == 0 && "mprotect last page failed");
+
+    // Verify we can read from last page
+    assert(p[9 * PAGESIZE] == 0x99 && "read from last protected page failed");
+
+    assert(munmap(p, len) == 0 && "munmap failed");
+
+    printf("mprotect_boundary test: PASS\n");
+    return 0;
+}
+

tests/unit-tests/memory_tests/deterministic/mprotect_end_region.c

@@ -0,0 +1,43 @@
+// Test: mprotect on end of memory region
+// Verifies correct splitting when changing protection at end of mapped region
+#include <sys/mman.h>
+#include <stdio.h>
+#include <assert.h>
+
+#define PAGESIZE 4096
+#define NUMPAGES 10
+
+int main(void) {
+    // Allocate 10 pages with READ|WRITE protection
+    size_t len = PAGESIZE * NUMPAGES;
+    unsigned char *p = mmap(NULL, len, PROT_READ | PROT_WRITE,
+                            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    assert(p != MAP_FAILED && "mmap failed");
+
+    // Write to all pages to ensure they work
+    for (int i = 0; i < NUMPAGES; i++) {
+        p[i * PAGESIZE] = 0xAA + i;
+    }
+
+    // Change protection on last 3 pages to READ-only
+    assert(mprotect(p + (7 * PAGESIZE), 3 * PAGESIZE, PROT_READ) == 0 && "mprotect failed");
+
+    // Verify we can still read from last 3 pages
+    assert(p[7 * PAGESIZE] == 0xAA + 7 && "read from protected region failed");
+    assert(p[9 * PAGESIZE] == 0xAA + 9 && "read from protected region failed");
+
+    // Verify we can still write to first 7 pages
+    for (int i = 0; i < 7; i++) {
+        p[i * PAGESIZE] = 0xBB + i;
+    }
+
+    // Verify writes succeeded
+    assert(p[0] == 0xBB && "write to unprotected region failed");
+    assert(p[6 * PAGESIZE] == 0xBB + 6 && "write to unprotected region failed");
+
+    assert(munmap(p, len) == 0 && "munmap failed");
+
+    printf("mprotect_end_region test: PASS\n");
+    return 0;
+}
+

tests/unit-tests/memory_tests/deterministic/mprotect_middle_region.c

@@ -0,0 +1,50 @@
+// Test: mprotect on middle of memory region
+// Verifies correct three-way splitting when changing protection in middle
+#include <sys/mman.h>
+#include <stdio.h>
+#include <assert.h>
+
+#define PAGESIZE 4096
+#define NUMPAGES 10
+
+int main(void) {
+    // Allocate 10 pages with READ|WRITE protection
+    size_t len = PAGESIZE * NUMPAGES;
+    unsigned char *p = mmap(NULL, len, PROT_READ | PROT_WRITE,
+                            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    assert(p != MAP_FAILED && "mmap failed");
+
+    // Write to all pages
+    for (int i = 0; i < NUMPAGES; i++) {
+        p[i * PAGESIZE] = 0xCC + i;
+    }
+
+    // Change protection on middle 4 pages (pages 3-6) to READ-only
+    assert(mprotect(p + (3 * PAGESIZE), 4 * PAGESIZE, PROT_READ) == 0 && "mprotect failed");
+
+    // Verify we can read from middle protected region
+    assert(p[3 * PAGESIZE] == 0xCC + 3 && "read from protected middle region failed");
+    assert(p[6 * PAGESIZE] == 0xCC + 6 && "read from protected middle region failed");
+
+    // Verify we can write to pages before protected region (0-2)
+    for (int i = 0; i < 3; i++) {
+        p[i * PAGESIZE] = 0xDD + i;
+    }
+
+    // Verify we can write to pages after protected region (7-9)
+    for (int i = 7; i < NUMPAGES; i++) {
+        p[i * PAGESIZE] = 0xDD + i;
+    }
+
+    // Verify all writes succeeded
+    assert(p[0] == 0xDD && "write to unprotected regions failed");
+    assert(p[2 * PAGESIZE] == 0xDD + 2 && "write to unprotected regions failed");
+    assert(p[7 * PAGESIZE] == 0xDD + 7 && "write to unprotected regions failed");
+    assert(p[9 * PAGESIZE] == 0xDD + 9 && "write to unprotected regions failed");
+
+    assert(munmap(p, len) == 0 && "munmap failed");
+
+    printf("mprotect_middle_region test: PASS\n");
+    return 0;
+}
+

tests/unit-tests/memory_tests/deterministic/mprotect_multiple_times.c

@@ -0,0 +1,51 @@
+// Test: multiple successive mprotect calls
+// Verifies correct state after overlapping protection changes
+#include <sys/mman.h>
+#include <stdio.h>
+#include <assert.h>
+
+#define PAGESIZE 4096
+#define NUMPAGES 10
+
+int main(void) {
+    // Allocate 10 pages with READ|WRITE protection
+    size_t len = PAGESIZE * NUMPAGES;
+    unsigned char *p = mmap(NULL, len, PROT_READ | PROT_WRITE,
+                            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    assert(p != MAP_FAILED && "mmap failed");
+
+    // Write initial data
+    for (int i = 0; i < NUMPAGES; i++) {
+        p[i * PAGESIZE] = 0x10 + i;
+    }
+
+    // First change: protect middle pages 3-6 to READ-only
+    assert(mprotect(p + (3 * PAGESIZE), 4 * PAGESIZE, PROT_READ) == 0 && "first mprotect failed");
+
+    // Second change: protect overlapping pages 5-7 to NONE
+    assert(mprotect(p + (5 * PAGESIZE), 3 * PAGESIZE, PROT_NONE) == 0 && "second mprotect failed");
+
+    // Verify we can still write to pages 0-2
+    p[0] = 0x20;
+    p[2 * PAGESIZE] = 0x22;
+
+    // Verify we can read from pages 3-4 (READ-only from first mprotect)
+    assert(p[3 * PAGESIZE] == 0x13 && "read from READ-only region failed");
+    assert(p[4 * PAGESIZE] == 0x14 && "read from READ-only region failed");
+
+    // Verify we can write to pages 8-9
+    p[8 * PAGESIZE] = 0x28;
+    p[9 * PAGESIZE] = 0x29;
+
+    // Verify final state
+    assert(p[0] == 0x20 && "final state verification failed");
+    assert(p[2 * PAGESIZE] == 0x22 && "final state verification failed");
+    assert(p[8 * PAGESIZE] == 0x28 && "final state verification failed");
+    assert(p[9 * PAGESIZE] == 0x29 && "final state verification failed");
+
+    assert(munmap(p, len) == 0 && "munmap failed");
+
+    printf("mprotect_multiple_times test: PASS\n");
+    return 0;
+}
+