Repository: mrexodia/dumpulator (language: python)

Issue:
Memory manager refactor
I rewrote the memory manager to use the paging, although an issue that arose was that the checking function is a bit limited at the moment.

I have had my hand at trying to improve it - it will detect if something was allocating in already allocated space, or "enveloping" allocated space, but if it were to have one 'foot' in at any side of the region - it won't catch the allocation error.

This one is way faster ( mainly because the checker function had to be limited ), hopefully you guys can point me in the right direction if there is something glaringly wrong. Thanks.


Discussion / hints:
Been a bit busy recently, but will review when I have some more time.
Okay I checked the code and it's still a bit difficult to reason about. Maybe to recap, Windows appears to manage memory in the following way:

- You _reserve_ a chunk of memory, let's call it _parent region_
- You can _(de)commit_/_protect_ ranges of pages inside the _parent region_
- Any operation going outside of the parent region will fail
- All regions are in multiple of `PAGE_SIZE`

Here is a start of my attempt:

```python
from enum import Enum
from typing import Any, List, Dict
import bisect
import copy

PAGE_SIZE = 0x1000

class RegionProtect(Enum):
    PAGE_EXECUTE = 0x10
    PAGE_EXECUTE_READ = 0x20
    PAGE_EXECUTE_READWRITE = 0x40
    PAGE_EXECUTE_WRITECOPY = 0x80
    PAGE_NOACCESS = 0x1
    PAGE_READONLY = 0x2
    PAGE_READWRITE = 0x4
    PAGE_WRITECOPY = 0x8
    PAGE_GUARD = 0x100
    PAGE_NOCACHE = 0x200
    PAGE_WRITECOMBINE = 0x400

class RegionType(Enum):
    MEM_IMAGE = 0x1000000
    MEM_MAPPED = 0x40000
    MEM_PRIVATE = 0x20000

class MemoryRegion:
    def __init__(self, start: int, size: int, protect: RegionProtect = RegionProtect.PAGE_NOACCESS, type: RegionType = RegionType.MEM_PRIVATE):
        assert start & 0xFFF == 0
        assert size & 0xFFF == 0 # TODO: support pseudo page of one address
        self.start = start
        self.size = size
        self.protect = protect
        self.type = type

    @property
    def end(self):
        return self.start + self.size

    def __lt__(self, other: Any):
        if isinstance(other, int):
            return self.start < other
        elif isinstance(other, MemoryRegion):
            return self.start < other.start
        raise TypeError()

    def __contains__(self, other: object) -> bool:
        if isinstance(other, int):
            return other >= self.start and other < self.end
        elif isinstance(other, MemoryRegion):
            if other.size == 0:
                return other.start >= self.start and other.end < self.end
            else:
                return other.start >= self.start and other.end <= self.end
        raise TypeError()

    def overlaps(self, other):
        if isinstance(other, MemoryRegion):
            if self.start <= other.start:
                return other.start < self.end
            else:
                return self.start < other.end
        raise TypeError()

    def __str__(self):
        return f"{hex(self.start)}[{hex(self.size)}]"

    def __repr__(self) -> str:
        return f"MemoryRegion({hex(self.start)}, {hex(self.size)}, {self.protect}, {self.type})"

    def pages(self):
        for page in range(self.start, self.end, PAGE_SIZE):
            page_region = copy.deepcopy(self)
            page_region.start = page
            page_region.size = PAGE_SIZE
            yield page_region

class MemoryManager:
    def __init__(self):
        self._regions: List[MemoryRegion] = []
        self._committed: Dict[int, MemoryRegion] = {}
    
    def find_parent(self, region: MemoryRegion):
        index = bisect.bisect_right(self._regions, region)
        if index == 0:
            return None
        else:
            closest = self._regions[index - 1]
            if region in closest:
                return closest
            else:
                return None

    def reserve(self, region: MemoryRegion):
        assert region.size > 0
        def check_overlaps(index):
            if index >= 0 and index < len(self._regions):
                value = self._regions[index]
                if region.overlaps(value):
                    raise KeyError(f"Requested region {region} overlaps with {value}")

        index = bisect.bisect_right(self._regions, region)
        if index == 0:
            check_overlaps(index)
        else:
            check_overlaps(index - 1)
            check_overlaps(index)
        self._regions.insert(index, region)

    def commit(self, start: int, size: int):
        region = MemoryRegion(start, size)
        parent_region = self.find_parent(region)
        if parent_region is None:
            raise KeyError(f"Could not find parent for {region}")
        
        for page in region.pages():
            self._committed[page.start] = page

    def protect(self, start: int, size: int, protect: RegionProtect):
        region = MemoryRegion(start, size)
        parent_region = self.find_parent(region)
        if parent_region is None:
            raise KeyError(f"Could not find parent for {region}")

        for page in region.pages():
            page.protect = protect
            # TODO: what happens if the page wasn't committed before?
            self._committed[page.start] = page

    def decommit(self, start: int, size: int):
        region = MemoryRegion(start, size)
        parent_region = self.find_parent(region)
        if parent_region is None:
            raise KeyError(f"Could not find parent for {region}")
        
        for page in region.pages():
            # TODO: what happens if the page wasn't committed before?
            del self._committed[page.start]

    def free(self, start: int, size: int):
        region = MemoryRegion(start, size)
        parent_region = self.find_parent(region)
        if parent_region is None:
            raise KeyError(f"Could not find parent for {region}")

        for page in parent_region.pages():
            if page.start in self._committed:
                del self._committed[page.start]
        self._regions.remove(parent_region)

    def query(self, start: int):
        parent_region = self.find_parent(MemoryRegion(start, 0))
        if parent_region is None:
            # TODO: handle free pages
            raise NotImplementedError()
        
        # Reference: https://learn.microsoft.com/en-us/windows/win32/api/memoryapi/nf-memoryapi-virtualquery#remarks
        final_region: MemoryRegion = None
        for page in parent_region.pages():
            if page.start < start:
                continue
            elif final_region is None:
                # TODO: what happens if the initial page is not commited?
                final_region = self._committed[page.start]
            else:
                commited_page = self._committed.get(page.start, None)
                if commited_page is None:
                    # TODO: what happens if you query a reserved page?
                    break
                if commited_page.type == final_region.type and commited_page.protect == final_region.protect:
                    final_region.size += PAGE_SIZE
                else:
                    break
        return final_region, parent_region

mm = MemoryManager()
pr = MemoryRegion(0x10000, 0x5000)
mm.reserve(pr)
mm.commit(0x11000, 0x4000)
mm.protect(0x11000, 0x3000, RegionProtect.PAGE_EXECUTE_READ)
print(mm.query(0x11000))
```

The way it's stored is `self._regions` contains the (reserved) parent regions and `self._committed` goes `page_address -> MemoryRegion` for each commited page. The only thing that gets complicated is the `query` function, but that's just a matter of handling a few edge cases and as you can see the rest of the functions are pretty trivial.

I think this solves the issue you were having with the overlapping regions. If you reserve `(0x10000, 0x3000)` it will not find a parent region if you try to find `(0x1000, 0x40000)`.

Hm, I think it will be easier if I continue this. Will keep you posted.

Fix patch:
diff --git a/src/dumpulator/details.py b/src/dumpulator/details.py
index 5d38a96..7953796 100644
--- a/src/dumpulator/details.py
+++ b/src/dumpulator/details.py
@@ -1,24 +1,26 @@
 import struct
 from collections import namedtuple
 
-from minidump.minidumpfile import AllocationProtect
 from unicorn import *
 from unicorn.x86_const import *
 
-def map_unicorn_perms(protect: AllocationProtect):
+from dumpulator.memory import MemoryProtect
+
+def map_unicorn_perms(protect: MemoryProtect):
     if isinstance(protect, int):
-        protect = AllocationProtect(protect)
+        protect = MemoryProtect(protect)
+    assert isinstance(protect, MemoryProtect)
     mapping = {
-        AllocationProtect.PAGE_EXECUTE: UC_PROT_EXEC | UC_PROT_READ,
-        AllocationProtect.PAGE_EXECUTE_READ: UC_PROT_EXEC | UC_PROT_READ,
-        AllocationProtect.PAGE_EXECUTE_READWRITE: UC_PROT_ALL,
-        AllocationProtect.PAGE_EXECUTE_WRITECOPY: UC_PROT_ALL,
-        AllocationProtect.PAGE_NOACCESS: UC_PROT_NONE,
-        AllocationProtect.PAGE_READONLY: UC_PROT_READ,
-        AllocationProtect.PAGE_READWRITE: UC_PROT_READ | UC_PROT_WRITE,
-        AllocationProtect.PAGE_WRITECOPY: UC_PROT_READ | UC_PROT_WRITE,
+        MemoryProtect.PAGE_EXECUTE: UC_PROT_EXEC | UC_PROT_READ,
+        MemoryProtect.PAGE_EXECUTE_READ: UC_PROT_EXEC | UC_PROT_READ,
+        MemoryProtect.PAGE_EXECUTE_READWRITE: UC_PROT_ALL,
+        MemoryProtect.PAGE_EXECUTE_WRITECOPY: UC_PROT_ALL,
+        MemoryProtect.PAGE_NOACCESS: UC_PROT_NONE,
+        MemoryProtect.PAGE_READONLY: UC_PROT_READ,
+        MemoryProtect.PAGE_READWRITE: UC_PROT_READ | UC_PROT_WRITE,
+        MemoryProtect.PAGE_WRITECOPY: UC_PROT_READ | UC_PROT_WRITE,
     }
-    return mapping.get(protect, UC_PROT_NONE)
+    return mapping[protect]
 
 
 class Registers:
@@ -561,197 +563,3 @@ def __setitem__(self, index, value):
     "Reserved"
 ]
 assert len(interrupt_names) == 32
-
-
-
-class MemoryManager():
-        def __init__(self, uc: Uc):
-            self._memory_map = []
-            self._uc = uc
-
-        def mem_protect(self, addr, size, perms):
-            for memory_region in self._memory_map: # Query parents
-                if addr in memory_region and addr+size in memory_region: # Verify that the memory you are altering is inside a parent memory chunk
-                    perms = map_unicorn_perms(perms)
-                    memory_region.children.confliction_check(addr,size) # Identify if there are any conflictions
-                    new_addr, new_size = memory_region.children.mem_protect_correct(addr, size, perms) # Using the conflictions, correct the memory layout
-                    memory_region.children.add_child(ChildMemoryRegion(new_addr,new_size,perms, None, memory_region.start))
-                    self._uc.mem_protect(addr, size, perms)
-                else:
-                    print("MEMORY_MANAGER | Attempted to change write permissions to memory that isn't mapped...")
-            return
-
-        def mem_free(self, addr): # MAJOR : NEED TO MANAGE FREE'ING PARENT MEMORY AND SHRINKING EFFECTED CHILD MEMORY OR DELETING CHILDREN ENTIRELY
-            for memory_region in self._memory_map:
-                if memory_region.start == addr:
-                    self._memory_map.remove(memory_region) # NEED TO ADD SELF._UC.MEM_UNMAP(ADDRESS,SIZE) ASSERT SIZE >= 0x1000
-                    return
-            return
-
-        def mem_allocate(self, addr:int, size:int, perms:int=UC_PROT_ALL, comment=None): # WORKS
-            assert size >= 0x1000
-            if self.check_memory_avail(addr, size) == True:
-                self._memory_map.append(ParentMemoryRegion(addr,size,perms,comment))
-                self._uc.mem_map(addr, size, perms)
-                return
-            else:
-                print("MEMORY_MANAGER | There was an attempt to illegally overmap memory from", addr, "to", addr+size)
-                return
-
-        def mem_find(self, addr): # WORKS
-            for memory_region in self._memory_map: 
-                if addr in memory_region:
-                    return memory_region
-            self.mem_error("MEMORY_MANAGER | Unable to find parent memory for given address...")
-            return 0
-
-        def check_memory_avail(self, addr, size):
-            for memory_region in self._memory_map:
-                if memory_region.query_parent(addr, size) == False:
-                    return True
-                else:
-                    return False
-            return True
-
-        def write(self, addr, data):
-            self._uc.mem_write(addr, data)
-
-        def read(self, addr, size):
-            return self._uc.mem_read(addr, size)
-
-        def readable_contents(self):
-            if len(self._memory_map):
-                for mem_region in self._memory_map:
-                    print("MEMORY_MANAGER | There is a parent block which starts at", mem_region.start, "and ends at", mem_region.end)
-            else:
-                print("MEMORY_MANAGER | There is no memory allocated...")
-
-class MemoryRegionChildL:
-    def __init__(self):
-        self._child_mem_list = []
-        self._conflicting_regions = []
-    
-    def sort(self):
-        length = len(self._child_mem_list)
-        for i in range(length-1):
-            for j in range(0, length-i-1):
-                if self._child_mem_list[j] > self._child_mem_list[j+1]:
-                    swapped = True
-                    self._child_mem_list[j], self._child_mem_list[j+1] = self._child_mem_list[j+1], self._child_mem_list[j]
-            if not swapped:
-                return
-
-    def confliction_check(self, addr, size): # Updates the conflicting memory regions list
-        self._conflicting_regions = []
-        assert len(self._conflicting_regions) == 0 # Make sure that confliction regions is cleared
-        for gen_address in range(addr, addr+size, 0x100):
-            for child_mem in self._child_mem_list:
-                if gen_address in child_mem:
-                    self._conflicting_regions.append(child_mem)
-        return
-
-    def mem_protect_correct(self, addr, size, perms):
-        new_addr = addr
-        new_size = size
-        for mem_region in self._conflicting_regions: #Clear the area for the new write permissions
-            if new_addr not in mem_region and new_addr+new_size not in mem_region: # No need to correct to split this area, it will be overwritten anyway.
-                self.remove_child(mem_region)
-            elif new_addr in mem_region and new_addr+new_size not in mem_region and mem_region.perms != perms: # The start of the region is in the mem_region, but not the end.
-                mod_region = self.find_child(mem_region)
-                mod_region.end = new_addr
-                mod_region.size = mem_region.end - mem_region.start
-            elif new_addr not in mem_region and new_addr+new_size in mem_region and mem_region.perms != perms:
-                mod_region = self.find_child(mem_region)
-                mod_region.start = new_addr+new_size
-                mod_region.size = mem_region.end - mem_region.start
-            elif new_addr in mem_region and new_addr+new_size in mem_region and mem_region.perms != perms:
-                new_region = ChildMemoryRegion(addr+size, mem_region.end-addr+size, mem_region.perms, mem_region.comment, mem_region.start)
-                mod_region = self.find_child(mem_region)
-                mod_region.end = new_addr
-                mod_region.size = mem_region.end - mem_region.start
-                self._child_mem_list.append(new_region)
-            elif new_addr in mem_region and new_addr+new_size in mem_region and mem_region.perms == perms:
-                return False # The attempted write permissions is within a single child memory block, with the same permissions, its safe to fail the change protection.
-            elif new_addr in mem_region and new_addr+new_size not in mem_region and mem_region.perms == perms:
-                new_addr = mem_region.start
-                new_size += addr - mem_region.start
-                self.remove_child(mem_region)
-            elif new_addr not in mem_region and new_addr+new_size in mem_region and mem_region.perms == perms:
-                new_size += mem_region.end - new_addr+new_size
-                self.remove_child(mem_region)
-            else:
-                print("MEMORY_REGION_CHILD | Unhandled logic in mem_protect_correct")
-                return False
-        return new_addr, new_size
-
-    def add_child(self, memory_region):
-        self._child_mem_list.append(memory_region)
-        return
-    
-    def find_child(self, mem_region):
-        for i, mem in enumerate(self._child_mem_list):
-            if mem_region.start == mem.start:
-                return self._child_mem_list[i]
-
-    def remove_child(self, mem_region):
-        for i, mem in enumerate(self._child_mem_list):
-            if mem_region.start == mem.start:
-                del self._child_mem_list[i]
-                break
-
-    def print_contents(self):
-        for i in self._child_mem_list:
-            print("A child memory block starts at", i.start, "ends at", i.start+i.size, "with permissions of", i.perms, "and a comment of", i.comment)
-
-class ChildMemoryRegion:
-    def __init__(self, start_addr, size, perms, comment=None, alloc_base=None):
-        self.region = comment
-        self.start = start_addr
-        self.end = self.start+size
-        self.alloc_base = alloc_base
-        self.size = size
-        self.perms = perms
-        self.comment = comment
-
-    def __gt__(self, other):
-        if self.start > other.start:
-            return True
-        return False
-
-    def __contains__(self, addr):
-        if self.start <= addr and addr <= self.end:
-            # print(self.start, "<=",addr, " & ",addr, "<=", self.end)
-            return True
-        else:
-            return False
-
-class ParentMemoryRegion:
-    def __init__(self, start_addr, size, perms, comment=None, alloc_base=None):
-        self.region = comment
-        self.start = start_addr
-        self.end = self.start+size
-        self.alloc_base = alloc_base
-        self.size = size
-        self.perms = perms
-        self.comment = comment
-        self.children = MemoryRegionChildL()
-
-    @property
-    def is_parent(self):
-        return self.alloc_base is None
-
-    def __contains__(self, addr):
-        if self.start <= addr and addr <= self.end:
-            return True
-        else:
-            return False
-
-    def query_parent(self, addr, size):
-        for gen_address in range(addr, addr+size, 0x100): # Could change the range step
-            if gen_address in self:
-                return True
-        return False
-
-    def print_contents(self):
-        print("MEMORY_REGION | This block starts at", self.start, "ends at", self.end, "with permissions of", self.perms,)
-        return
\ No newline at end of file
diff --git a/src/dumpulator/dumpulator.py b/src/dumpulator/dumpulator.py
index 2eec5f8..0a51d71 100644
--- a/src/dumpulator/dumpulator.py
+++ b/src/dumpulator/dumpulator.py
@@ -15,6 +15,7 @@
 from .handles import HandleManager, FileObject
 from .native import *
 from .details import *
+from .memory import *
 from capstone import *
 from capstone.x86_const import *
 
@@ -55,6 +56,20 @@ def __init__(self):
     def __str__(self):
         return f"{self.type}, ({hex(self.tb_start)}, {hex(self.tb_size)}, {self.tb_icount})"
 
+class UnicornPageManager(PageManager):
+    def __init__(self, uc: Uc):
+        self._uc = uc
+
+    def commit(self, addr: int, size: int, protect: MemoryProtect) -> None:
+        perms = map_unicorn_perms(protect)
+        self._uc.mem_map(addr, size, perms)
+
+    def decommit(self, addr: int, size: int) -> None:
+        self._uc.mem_unmap(addr, size)
+
+    def protect(self, addr: int, size: int, protect: MemoryProtect) -> None:
+        self._uc.mem_protect(addr, size, map_unicorn_perms(protect))
+
 class Dumpulator(Architecture):
     def __init__(self, minidump_file, *, trace=False, quiet=False, thread_id=None):
         self._quiet = quiet
@@ -87,6 +102,7 @@ def __init__(self, minidump_file, *, trace=False, quiet=False, thread_id=None):
 
         self.regs = Registers(self._uc, self._x64)
         self.args = Arguments(self._uc, self.regs, self._x64)
+        self.memory = MemoryManager(UnicornPageManager(self._uc))
         self._allocate_base = None
         self._allocate_size = 1024 * 1024 * 10  # NOTE: 10 megs
         self._allocate_ptr = None
@@ -147,16 +163,38 @@ def _setup_emulator(self, thread):
         self._uc.mem_write(KERNEL_CAVE, bytes(kernel_code))
 
         info: minidump.MinidumpMemoryInfo
+        regions: List[List[minidump.MinidumpMemoryInfo]] = []
         for info in self._minidump.memory_info.infos:
+            if len(regions) == 0 or info.AllocationBase != regions[-1][0].AllocationBase or info.State == minidump.MemoryState.MEM_FREE:
+                regions.append([])
+            regions[-1].append(info)
+        # NOTE: The HYPERVISOR_SHARED_DATA does not respect the allocation granularity
+        self.memory._granularity = PAGE_SIZE
+        for i in range(len(regions)):
+            region = regions[i]
+            reserve_addr = None
+            reserve_size = 0
+            assert len(region) >= 1
+            for j in range(len(region)):
+                info = region[j]
+                if reserve_addr is None:
+                    reserve_addr = info.BaseAddress
+                reserve_size += info.RegionSize
+            info = region[0]
             emu_addr = info.BaseAddress & self.addr_mask
-            if info.State == minidump.MemoryState.MEM_COMMIT:
-                self.info(f"committed: 0x{emu_addr:x}, size: 0x{info.RegionSize:x}, protect: {info.Protect}")
-                self._uc.mem_map(emu_addr, info.RegionSize, map_unicorn_perms(info.Protect))
-            elif info.State == minidump.MemoryState.MEM_FREE and emu_addr > 0x10000 and info.RegionSize >= self._allocate_size:
-                self._allocate_base = emu_addr
-            elif info.State == minidump.MemoryState.MEM_RESERVE:
-                self.info(f"reserved: {hex(emu_addr)}, size: {hex(info.RegionSize)}")
-                self._uc.mem_map(emu_addr, info.RegionSize, UC_PROT_NONE)
+            if info.State == minidump.MemoryState.MEM_FREE:
+                continue
+            reserve_protect = MemoryProtect(info.AllocationProtect)
+            reserve_type = MemoryType(info.Type.value)
+            print(f" reserved: {hex(reserve_addr)}, size: {hex(reserve_size)}, protect: {reserve_protect}, type: {reserve_type}")
+            self.memory.reserve(reserve_addr, reserve_size, reserve_protect, reserve_type)
+            for info in region:
+                emu_addr = info.BaseAddress & self.addr_mask
+                if info.State == minidump.MemoryState.MEM_COMMIT:
+                    protect = reserve_protect if info.Protect is None else MemoryProtect(info.Protect.value)
+                    print(f"committed: {hex(emu_addr)}, size: {hex(info.RegionSize)}, protect: {protect}")
+                    self.memory.commit(info.BaseAddress, info.RegionSize, protect)
+        self.memory._granularity = 0x10000
 
         memory = self._minidump.get_reader().get_buffered_reader()
         seg: minidump.MinidumpMemorySegment
@@ -378,10 +416,6 @@ def read(self, addr, size):
     def write(self, addr, data):
         self._uc.mem_write(addr, data)
 
-    def protect(self, addr, size, protect):
-        perms = map_unicorn_perms(protect)
-        self._uc.mem_protect(addr, size, perms)
-
     def call(self, addr, args: List[int] = [], regs: dict = {}, count=0):
         # allow passing custom registers
         for name, value in regs.items():
@@ -402,17 +436,26 @@ def call(self, addr, args: List[int] = [], regs: dict = {}, count=0):
 
     def allocate(self, size, page_align=False):
         if not self._allocate_ptr:
-            self._uc.mem_map(self._allocate_base, self._allocate_size)
+            self._allocate_base = self.memory.find_free(self._allocate_size)
+            self.memory.reserve(
+                start=self._allocate_base,
+                size=self._allocate_size,
+                protect=MemoryProtect.PAGE_EXECUTE_READWRITE,
+                type=MemoryType.MEM_PRIVATE,
+                comment="allocated region"
+            )
             self._allocate_ptr = self._allocate_base
 
         if page_align:
             self._allocate_ptr = round_to_pages(self._allocate_ptr)
             size = round_to_pages(size)
 
-        ptr = self._allocate_ptr + size
-        if ptr > self._allocate_base + self._allocate_size:
+        if self._allocate_ptr + size > self._allocate_base + self._allocate_size:
             raise Exception("not enough room to allocate!")
-        self._allocate_ptr = ptr
+
+        ptr = self._allocate_ptr
+        self._allocate_ptr += size
+        self.memory.commit(self.memory.page_align(ptr), self.memory.page_align(size))
         return ptr
 
     def handle_exception(self):
@@ -445,7 +488,7 @@ def handle_exception(self):
       alignment        @ rsp + 588 : 8
       MACHINE_FRAME    @ rsp + 590 : 28                       | alignas(16) from RSP in exception / xstate
       alignment        @ rsp + 5b8 : 8
-      xstate           @ rsp + 5c0 : CONTEXT_EX.Xstate.Length | alignas(64) from RSP in exception 
+      xstate           @ rsp + 5c0 : CONTEXT_EX.Xstate.Length | alignas(64) from RSP in exception
             """
             allocation_size = 0x720
             context_flags = 0x10005F
@@ -629,17 +672,17 @@ def map_module(self, file_data: bytes, file_name: str = "", requested_base: int
         section_alignment = pe.OPTIONAL_HEADER.SectionAlignment
         assert section_alignment == 0x1000
         if requested_base == 0:
-            image_base = self.allocate(image_size, True)
+            image_base = self.memory.find_free(image_size)
         else:
             image_base = requested_base
-            self._uc.mem_map(image_base, image_size)
+        self.memory.reserve(image_base, image_size, MemoryProtect.PAGE_EXECUTE_WRITECOPY)
 
         # TODO: map the header properly
         header = pe.header
         header_size = pe.sections[0].VirtualAddress_adj
         print(f"Mapping header {hex(image_base)}[{hex(header_size)}]")
+        self.memory.commit(image_base, header_size, MemoryProtect.PAGE_READONLY)
         self.write(image_base, header)
-        self.protect(image_base, header_size, PAGE_READONLY)
 
         for section in pe.sections:
             name = section.Name.rstrip(b"\0")
@@ -653,14 +696,14 @@ def map_module(self, file_data: bytes, file_name: str = "", requested_base: int
             assert flags & IMAGE_SCN_MEM_READ != 0
             execute = flags & IMAGE_SCN_MEM_EXECUTE
             write = flags & IMAGE_SCN_MEM_WRITE
-            protect = PAGE_READONLY
+            protect = MemoryProtect.PAGE_READONLY
             if write:
-                protect = PAGE_READWRITE
+                protect = MemoryProtect.PAGE_READWRITE
             if execute:
-                protect <<= 4
-            print(f"Mapping section '{name.decode()}' {hex(rva)}[{hex(rva)}] -> {hex(va)}")
+                protect = MemoryProtect(protect.value << 4)
+            print(f"Mapping section '{name.decode()}' {hex(rva)}[{hex(rva)}] -> {hex(va)} as {protect}")
+            self.memory.commit(va, size, protect)
             self.write(va, data)
-            self.protect(va, size, protect)
 
         # TODO: implement relocations
         reloc_dir = pe.OPTIONAL_HEADER.DATA_DIRECTORY[5]
diff --git a/src/dumpulator/memory.py b/src/dumpulator/memory.py
new file mode 100644
index 0000000..0feca48
--- /dev/null
+++ b/src/dumpulator/memory.py
@@ -0,0 +1,303 @@
+from enum import Enum, IntFlag, Flag
+from typing import Any, List, Dict, Union
+import bisect
+
+PAGE_SIZE = 0x1000
+
+class MemoryProtect(Flag):
+    UNDEFINED = 0x0
+    PAGE_EXECUTE = 0x10
+    PAGE_EXECUTE_READ = 0x20
+    PAGE_EXECUTE_READWRITE = 0x40
+    PAGE_EXECUTE_WRITECOPY = 0x80
+    PAGE_NOACCESS = 0x1
+    PAGE_READONLY = 0x2
+    PAGE_READWRITE = 0x4
+    PAGE_WRITECOPY = 0x8
+    # Only these can be combined
+    PAGE_GUARD = 0x100
+    PAGE_NOCACHE = 0x200
+    PAGE_WRITECOMBINE = 0x400
+
+class MemoryType(Enum):
+    UNDEFINED = 0
+    MEM_IMAGE = 0x1000000
+    MEM_MAPPED = 0x40000
+    MEM_PRIVATE = 0x20000
+
+class MemoryState(Enum):
+    MEM_COMMIT = 0x1000
+    MEM_RESERVE = 0x2000
+    MEM_FREE = 0x10000
+
+class MemoryBasicInformation:
+    def __init__(self, base: int, allocation_base: int, allocation_protect: MemoryProtect):
+        self.base = base
+        self.allocation_base = allocation_base
+        self.allocation_protect = allocation_protect
+        self.region_size: int = PAGE_SIZE
+        self.state: MemoryState = None
+        self.protect: MemoryProtect = None
+        self.type: MemoryType = None
+
+    def __str__(self):
+        return f"MemoryBasicInformation(base: {hex(self.base)}, allocation_base: {hex(self.allocation_base)}, region_size: {hex(self.region_size)}, state: {self.state}, protect: {self.protect}, type: {self.type})"
+
+class MemoryRegion:
+    def __init__(self, start: int, size: int, protect: MemoryProtect = MemoryProtect.PAGE_NOACCESS, type: MemoryType = MemoryType.MEM_PRIVATE, comment: str = ""):
+        assert start & 0xFFF == 0
+        assert size & 0xFFF == 0
+        self.start = start
+        self.size = size
+        self.protect = protect
+        self.type = type
+        self.comment = comment
+
+    @property
+    def end(self):
+        return self.start + self.size
+
+    def __lt__(self, other: Any):
+        if isinstance(other, int):
+            return self.start < other
+        elif isinstance(other, MemoryRegion):
+            return self.start < other.start
+        raise TypeError()
+
+    def __contains__(self, other: object) -> bool:
+        if isinstance(other, int):
+            return other >= self.start and other < self.end
+        elif isinstance(other, MemoryRegion):
+            if other.size == 0:
+                return other.start >= self.start and other.end < self.end
+            else:
+                return other.start >= self.start and other.end <= self.end
+        raise TypeError()
+
+    def overlaps(self, other):
+        if isinstance(other, MemoryRegion):
+            if self.start <= other.start:
+                return other.start < self.end
+            else:
+                return self.start < other.end
+        raise TypeError()
+
+    def __str__(self):
+        result = f"{hex(self.start)}[{hex(self.size)}]"
+        if len(self.comment) > 0:
+            result += f" ({self.comment})"
+        return result
+
+    def __repr__(self) -> str:
+        return f"MemoryRegion({hex(self.start)}, {hex(self.size)}, {self.protect}, {self.type}, {self.comment}"
+
+    def pages(self):
+        for page in range(self.start, self.end, PAGE_SIZE):
+            yield page
+
+class PageManager:
+    def commit(self, addr: int, size: int, protect: MemoryProtect) -> None:
+        raise NotImplementedError()
+
+    def decommit(self, addr: int, size: int) -> None:
+        raise NotImplementedError()
+
+    def protect(self, addr: int, size: int, protect: MemoryProtect) -> None:
+        raise NotImplementedError()
+
+class MemoryManager:
+    def __init__(self, page_manager: PageManager, minimum = 0x10000, maximum = 0x7fffffff0000, granularity = 0x10000):
+        self._page_manager = page_manager
+        self._minimum = minimum
+        self._maximum = maximum
+        self._granularity = granularity
+        self._regions: List[MemoryRegion] = []
+        self._committed: Dict[int, MemoryRegion] = {}
+
+    def find_parent(self, region: Union[MemoryRegion, int]):
+        if isinstance(region, int):
+            region = MemoryRegion(self.page_align(region), 0)
+        index = bisect.bisect_right(self._regions, region)
+        if index == 0:
+            return None
+        else:
+            closest = self._regions[index - 1]
+            if region in closest:
+                return closest
+            else:
+                return None
+
+    def page_align(self, addr: int):
+        mask = PAGE_SIZE - 1
+        return (addr + mask) & ~mask
+
+    def allocation_align(self, addr: int):
+        mask = self._granularity - 1
+        return (addr + mask) & ~mask
+
+    def find_free(self, size: int):
+        assert size > 0 and self.page_align(size) == size
+        base = self._minimum
+        while base < self._maximum:
+            info = self.query(base)
+            assert info.base == base
+            if info.state == MemoryState.MEM_FREE and info.region_size >= size and self.allocation_align(base) == base:
+                return info.base
+            base += info.region_size
+        return None
+
+    def reserve(self, start: int, size: int, protect: MemoryProtect, type: MemoryType = MemoryType.MEM_PRIVATE, comment: str = ""):
+        assert isinstance(protect, MemoryProtect)
+        assert isinstance(type, MemoryType)
+        assert size > 0 and self.page_align(size) == size
+        assert self.allocation_align(start) == start
+        region = MemoryRegion(start, size, protect, type, comment)
+        if region.start < self._minimum or region.end > self._maximum:
+            raise KeyError(f"Requested region {region} is out of bounds")
+
+        def check_overlaps(index):
+            if index >= 0 and index < len(self._regions):
+                value = self._regions[index]
+                if region.overlaps(value):
+                    raise KeyError(f"Requested region {region} overlaps with {value}")
+
+        index = bisect.bisect_right(self._regions, region)
+        if index == 0:
+            check_overlaps(index)
+        else:
+            check_overlaps(index - 1)
+            check_overlaps(index)
+        self._regions.insert(index, region)
+
+    def release(self, start: int):
+        assert self.allocation_align(start) == start
+
+        parent_region = self.find_parent(start)
+        if parent_region is None:
+            raise KeyError(f"Could not find parent for {hex(start)}")
+        if parent_region.start != start:
+            raise KeyError(f"You can only release the whole parent region")
+
+        if all([page in self._committed for page in parent_region.pages()]):
+            self._page_manager.decommit(parent_region.start, parent_region.size)
+            for page in parent_region.pages():
+                del self._committed[page]
+        else:
+            for page in parent_region.pages():
+                if page in self._committed:
+                    self._page_manager.decommit(page, PAGE_SIZE)
+                    del self._committed[page]
+        self._regions.remove(parent_region)
+
+    def commit(self, start: int, size: int, protect: MemoryProtect = MemoryProtect.UNDEFINED):
+        assert isinstance(protect, MemoryProtect)
+        assert size > 0 and self.page_align(size) == size
+        assert self.page_align(start) == start
+        region = MemoryRegion(start, size)
+        parent_region = self.find_parent(region)
+        if parent_region is None:
+            raise KeyError(f"Could not find parent for {region}")
+
+        if protect == MemoryProtect.UNDEFINED:
+            protect = parent_region.protect
+
+        if all([page not in self._committed for page in region.pages()]):
+            self._page_manager.commit(region.start, region.size, protect)
+            for page in region.pages():
+                self._committed[page] = MemoryRegion(page, PAGE_SIZE, protect, parent_region.type)
+        else:
+            for page in region.pages():
+                if page in self._committed:
+                    self._page_manager.protect(page, PAGE_SIZE, protect)
+                    self._committed[page].protect = protect
+                else:
+                    self._page_manager.commit(page, PAGE_SIZE, protect)
+                    self._committed[page] = MemoryRegion(page, PAGE_SIZE, protect, parent_region.type)
+
+    def decommit(self, start: int, size: int):
+        assert size > 0 and self.page_align(size) == size
+        assert self.page_align(start) == start
+        region = MemoryRegion(start, size)
+        parent_region = self.find_parent(region)
+        if parent_region is None:
+            raise KeyError(f"Could not find parent for {region}")
+
+        if all([page in self._committed for page in region.pages()]):
+            self._page_manager.decommit(region.start, region.size)
+            for page in self._committed:
+                del self._committed[page]
+        else:
+            for page in region.pages():
+                if page in self._committed:
+                    self._page_manager.decommit(page, PAGE_SIZE)
+                    del self._committed[page]
+
+    def protect(self, start: int, size: int, protect: MemoryProtect):
+        assert isinstance(protect, MemoryProtect)
+        assert size > 0 and self.page_align(size) == size
+        assert self.page_align(start) == start
+        region = MemoryRegion(start, size)
+        parent_region = self.find_parent(region)
+        if parent_region is None:
+            raise KeyError(f"Could not find parent for {region}")
+
+        # Make sure all pages in the region are commited
+        for page in region.pages():
+            if not self._committed[page]:
+                raise KeyError(f"Could not protect uncommited page {hex(page)}")
+
+        # Change the protection
+        old_protect = self._committed[region.start].protect
+        self._page_manager.protect(region.start, region.size, protect)
+        for page in region.pages():
+            self._committed[page].protect = protect
+
+        return old_protect
+
+    def query(self, start: int):
+        assert self.page_align(start) == start
+
+        region = MemoryRegion(start, 0)
+        parent_region = self.find_parent(region)
+        if parent_region is None:
+            index = bisect.bisect_right(self._regions, region)
+            next_start = self._maximum
+            if index < len(self._regions):
+                next_start = self._regions[index].start
+            result = MemoryBasicInformation(start, 0, MemoryProtect.UNDEFINED)
+            result.region_size = next_start - start
+            assert result.base + result.region_size == next_start
+            result.state = MemoryState.MEM_FREE
+            result.protect = MemoryProtect.UNDEFINED
+            result.type = MemoryType.UNDEFINED
+            return result
+
+        # Reference: https://learn.microsoft.com/en-us/windows/win32/api/memoryapi/nf-memoryapi-virtualquery#remarks
+        result: MemoryBasicInformation = None
+        for page in parent_region.pages():
+            if page < start:
+                continue
+            elif result is None:
+                result = MemoryBasicInformation(page, parent_region.start, parent_region.protect)
+                if page in self._committed:
+                    result.state = MemoryState.MEM_COMMIT
+                    commited_page = self._committed[page]
+                    result.protect = commited_page.protect
+                    result.type = commited_page.type
+                    assert commited_page.type == parent_region.type
+                else:
+                    result.state = MemoryState.MEM_RESERVE
+                    result.protect = MemoryProtect.UNDEFINED
+                    result.type = parent_region.type
+            else:
+                commited_page = self._committed.get(page, None)
+                if result.state == MemoryState.MEM_RESERVE:
+                    result.region_size += PAGE_SIZE
+                elif result.state == MemoryState.MEM_COMMIT:
+                    if commited_page is not None and commited_page.type == result.type and commited_page.protect == result.protect:
+                        result.region_size += PAGE_SIZE
+                    else:
+                        break
+                else:
+                    assert False  # unreachable
+        return result
diff --git a/src/dumpulator/ntsyscalls.py b/src/dumpulator/ntsyscalls.py
index d7bab6b..db87ff4 100644
--- a/src/dumpulator/ntsyscalls.py
+++ b/src/dumpulator/ntsyscalls.py
@@ -1,9 +1,11 @@
 import ctypes
 import struct
 import unicorn
+
 from .dumpulator import Dumpulator, syscall_functions, ExceptionInfo, ExceptionType
 from .native import *
 from .handles import *
+from .memory import *
 from pathlib import Path
 
 
@@ -303,33 +305,29 @@ def ZwAllocateVirtualMemory(dp: Dumpulator,
                             ):
     assert ZeroBits == 0
     assert ProcessHandle == dp.NtCurrentProcess()
-    assert Protect == PAGE_READWRITE
     base = dp.read_ptr(BaseAddress.ptr)
     assert base & 0xFFF == 0
     size = round_to_pages(dp.read_ptr(RegionSize.ptr))
     assert size != 0
+    protect = MemoryProtect(Protect)
     if AllocationType == MEM_COMMIT:
         if base == 0:
             base = dp.allocate(size, True)
         print(f"commit({hex(base)}[{hex(size)}])")
-        dp._uc.mem_protect(base, size, unicorn.UC_PROT_READ | unicorn.UC_PROT_WRITE)
+        dp.memory.commit(base, size, protect)
     elif AllocationType == MEM_RESERVE:
         if base == 0:
-            base = dp.allocate(size, True)
-            dp._uc.mem_protect(base, size, unicorn.UC_PROT_NONE)
+            base = dp.memory.find_free(size)
             BaseAddress.write_ptr(base)
-        else:
-            dp._uc.mem_map(base, size, unicorn.UC_PROT_NONE)
         print(f"reserve({hex(base)}[{hex(size)}])")
+        dp.memory.reserve(base, size, protect)
     elif AllocationType == MEM_COMMIT | MEM_RESERVE:
         if base == 0:
-            base = dp.allocate(size, True)
-            dp._uc.mem_protect(base, size, unicorn.UC_PROT_READ | unicorn.UC_PROT_WRITE)
+            base = dp.memory.find_free(size)
             BaseAddress.write_ptr(base)
-        else:
-            dp._uc.mem_map(base, size, unicorn.UC_PROT_READ | unicorn.UC_PROT_WRITE)
         print(f"reserve+commit({hex(base)}[{hex(size)}])")
-        return STATUS_SUCCESS
+        dp.memory.reserve(base, size, protect)
+        dp.memory.commit(base, size)
     else:
         assert False
     return STATUS_SUCCESS
@@ -2129,7 +2127,7 @@ def ZwMapViewOfSection(dp: Dumpulator,
     section = dp.handles.get(SectionHandle, SectionObject)
     data = section.file.read()
     image_base, image_size, pe = dp.map_module(data, section.file.path, requested_base)
-    
+
     # Handle out parameters
     BaseAddress.write_ptr(image_base)
     ViewSize.write_ptr(image_size)
@@ -2613,10 +2611,11 @@ def ZwProtectVirtualMemory(dp: Dumpulator,
                            ):
     base = BaseAddress.read_ptr() & 0xFFFFFFFFFFFFF000
     size = round_to_pages(RegionSize.read_ptr())
+    protect = MemoryProtect(NewProtect)
 
-    print(f"protect {base:x}[{size:x}] = {NewProtect:x}")
-    dp.protect(base, size, NewProtect)
-    # TODO: OldProtect is not implemented
+    print(f"protect {hex(base)}[{hex(size)}] = {protect}")
+    old_protect = dp.memory.protect(base, size, protect)
+    OldProtect.write_ulong(old_protect.value)
     return STATUS_SUCCESS
 
 @syscall
@@ -3225,47 +3224,33 @@ def ZwQueryVirtualMemory(dp: Dumpulator,
                          ):
     assert ProcessHandle == dp.NtCurrentProcess()
     if MemoryInformationClass == MEMORY_INFORMATION_CLASS.MemoryBasicInformation:
-        info = MEMORY_BASIC_INFORMATION(dp)
         assert MemoryInformationLength == ctypes.sizeof(info)
-        # TODO: hardcoded for the example
-        if BaseAddress.ptr in [0x140001050, 0x140001e00]:
-            info.BaseAddress = 0x140001000
-            info.AllocationBase = 0x140000000
-            info.RegionSize = 0xD000
-        elif BaseAddress.ptr in [0x401aef, 0x401080]:
-            info.BaseAddress = 0x401000
-            info.AllocationBase = 0x400000
-            info.RegionSize = 0xC000
-        else:
-            raise NotImplementedError()
-        info.AllocationProtect = PAGE_EXECUTE_WRITECOPY
-        info.ParitionId = 0
-        info.State = MEM_COMMIT
-        info.Type = MEM_IMAGE
-        MemoryInformation.write(bytes(info))
+        info = dp.memory.query(BaseAddress.ptr)
+        mbi = MEMORY_BASIC_INFORMATION(dp)
+        mbi.BaseAddress = info.base
+        mbi.AllocationBase = info.allocation_base
+        mbi.AllocationProtect = info.allocation_protect
+        mbi.RegionSize = info.region_size
+        mbi.State = info.state.value
+        mbi.Protect = info.protect.value
+        mbi.Type = info.type.value
+        MemoryInformation.write(bytes(mbi))
         if ReturnLength.ptr:
-            ReturnLength.write_ulong(ctypes.sizeof(info))
+            ReturnLength.write_ulong(ctypes.sizeof(mbi))
         return STATUS_SUCCESS
     elif MemoryInformationClass == MEMORY_INFORMATION_CLASS.MemoryRegionInformation:
-        info = MEMORY_REGION_INFORMATION(dp)
-        assert MemoryInformationLength >= ctypes.sizeof(info)
-        # TODO: hardcoded for the example
-        if BaseAddress.ptr in [0x140001050, 0x140001e00]:
-            info.AllocationBase = 0x140000000
-            info.RegionSize = 0x1c000  # ImageSize
-        elif BaseAddress.ptr in [0x401aef, 0x401080]:
-            info.AllocationBase = 0x400000
-            info.RegionSize = 0x16000  # ImageSize
-            pass
-        else:
-            raise NotImplementedError()
-        info.AllocationProtect = PAGE_EXECUTE_WRITECOPY
-        info.Flags = REGION_MAPPED_IMAGE
-        info.CommitSize = info.RegionSize
-        MemoryInformation.write(bytes(info))
-        extra_size = MemoryInformationLength - ctypes.sizeof(info)
+        parent_region = dp.memory.find_parent(BaseAddress.ptr)
+        mri = MEMORY_REGION_INFORMATION(dp)
+        mri.AllocationBase = parent_region.start
+        mri.AllocationProtect = parent_region.protect
+        mri.Flags = REGION_MAPPED_IMAGE if parent_region.type == MemoryType.MEM_IMAGE else REGION_PRIVATE
+        mri.RegionSize = parent_region.size
+        mri.CommitSize = parent_region.size  # TODO
+        assert MemoryInformationLength >= ctypes.sizeof(mri)
+        MemoryInformation.write(bytes(mri))
+        extra_size = MemoryInformationLength - ctypes.sizeof(mri)
         if extra_size > 0:
-            dp.write(MemoryInformation.ptr + ctypes.sizeof(info), b"\x69" * extra_size)
+            dp.write(MemoryInformation.ptr + ctypes.sizeof(mri), b"\x69" * extra_size)
         if ReturnLength.ptr:
             ReturnLength.write_ulong(MemoryInformationLength)
         return STATUS_SUCCESS


Test patch:
diff --git a/tests/memory_test.py b/tests/memory_test.py
new file mode 100644
index 0000000..7576e08
--- /dev/null
+++ b/tests/memory_test.py
@@ -0,0 +1,106 @@
+import unittest
+from typing import Dict
+
+from dumpulator.memory import *
+
+class MockPageManager(PageManager):
+    def __init__(self):
+        self.pages: Dict[int, MemoryProtect] = {}
+
+    def commit(self, addr: int, size: int, protect: MemoryProtect) -> None:
+        print(f"commit({hex(addr)}, {hex(size)}, {protect})")
+        for page in range(addr, addr + size, PAGE_SIZE):
+            self.pages[page] = protect
+
+    def decommit(self, addr: int, size: int) -> None:
+        print(f"decommit({hex(addr)}, {hex(size)})")
+        for page in range(addr, addr + size, PAGE_SIZE):
+            del self.pages[page]
+
+    def protect(self, addr: int, size: int, protect: MemoryProtect) -> None:
+        print(f"protect({hex(addr)}, {hex(size)}, {protect})")
+        for page in range(addr, addr + size, PAGE_SIZE):
+            assert page in self.pages
+            self.pages[page] = protect
+
+class TestMemoryManager(unittest.TestCase):
+    def setUp(self) -> None:
+        self.pm = MockPageManager()
+        self.mm = MemoryManager(self.pm)
+
+    def test_reserve(self):
+        self.mm.reserve(0x20000, 0x10000, MemoryProtect.PAGE_READWRITE)
+        assert self.pm.pages == {}
+        assert self.mm.query(0x20000).region_size == 0x10000
+        assert self.mm.query(0x22000).region_size == 0x10000 - 0x2000
+        self.assertRaises(KeyError, lambda: self.mm.reserve(0x10000, 0x100000, MemoryProtect.PAGE_READWRITE))
+        self.assertRaises(KeyError, lambda: self.mm.reserve(0x0, 0x2000, MemoryProtect.PAGE_NOACCESS))
+        self.assertRaises(KeyError, lambda: self.mm.reserve(self.mm._maximum, 0x1000, MemoryProtect.PAGE_READWRITE))
+        self.mm.reserve(0x30000, 0x2000, MemoryProtect.PAGE_READONLY)
+        assert self.pm.pages == {}
+        assert self.mm.query(0x30000).region_size == 0x2000
+
+    def test_commit(self):
+        self.mm.reserve(0x20000, 0x30000, MemoryProtect.PAGE_NOACCESS)
+        self.mm.commit(0x20000, 0x2000, MemoryProtect.PAGE_READWRITE)
+        assert self.pm.pages == {
+            0x20000: MemoryProtect.PAGE_READWRITE,
+            0x21000: MemoryProtect.PAGE_READWRITE,
+        }
+        info = self.mm.query(0x20000)
+        assert info.base == 0x20000
+        assert info.allocation_base == 0x20000
+        assert info.allocation_protect == MemoryProtect.PAGE_NOACCESS
+        assert info.region_size == 0x2000
+        assert info.state == MemoryState.MEM_COMMIT
+        assert info.protect == MemoryProtect.PAGE_READWRITE
+        assert info.type == MemoryType.MEM_PRIVATE
+        info = self.mm.query(0x23000)
+        assert info.base == 0x23000
+        assert info.allocation_base == 0x20000
+        assert info.allocation_protect == MemoryProtect.PAGE_NOACCESS
+        assert info.region_size == 0x30000 - 0x3000
+        assert info.state == MemoryState.MEM_RESERVE
+        assert info.protect == MemoryProtect.UNDEFINED
+        assert info.type == MemoryType.MEM_PRIVATE
+
+    def test_protect(self):
+        # Protect parent region that doesn't exist
+        self.assertRaises(KeyError, lambda: self.mm.protect(0x20000, 0x5000, MemoryProtect.PAGE_NOACCESS))
+        # Reserve the region
+        self.mm.reserve(0x20000, 0x5000, MemoryProtect.PAGE_READWRITE)
+        # Protect uncommited pages
+        self.assertRaises(KeyError, lambda: self.mm.protect(0x20000, 0x2000, MemoryProtect.PAGE_EXECUTE_READ))
+        # Commit part of the region (use parent protection)
+        self.mm.commit(0x20000, 0x2000, MemoryProtect.UNDEFINED)
+        assert self.pm.pages == {
+            0x20000: MemoryProtect.PAGE_READWRITE,
+            0x21000: MemoryProtect.PAGE_READWRITE,
+        }
+        info = self.mm.query(0x20000)
+        assert info.protect == MemoryProtect.PAGE_READWRITE
+        assert info.region_size == 0x2000
+        # Change the protection
+        old_protect = self.mm.protect(0x20000, 0x2000, MemoryProtect.PAGE_EXECUTE_READ)
+        assert old_protect == info.protect
+        assert self.pm.pages == {
+            0x20000: MemoryProtect.PAGE_EXECUTE_READ,
+            0x21000: MemoryProtect.PAGE_EXECUTE_READ,
+        }
+        info = self.mm.query(0x20000)
+        assert info.protect == MemoryProtect.PAGE_EXECUTE_READ
+        # Commit the whole region (overrides the protection)
+        self.mm.commit(0x20000, 0x5000, MemoryProtect.PAGE_EXECUTE_READWRITE)
+        assert self.pm.pages == {
+            0x20000: MemoryProtect.PAGE_EXECUTE_READWRITE,
+            0x21000: MemoryProtect.PAGE_EXECUTE_READWRITE,
+            0x22000: MemoryProtect.PAGE_EXECUTE_READWRITE,
+            0x23000: MemoryProtect.PAGE_EXECUTE_READWRITE,
+            0x24000: MemoryProtect.PAGE_EXECUTE_READWRITE,
+        }
+        info = self.mm.query(0x20000)
+        assert info.region_size == 0x5000
+        assert info.protect == MemoryProtect.PAGE_EXECUTE_READWRITE
+
+if __name__ == "__main__":
+    unittest.main()