Memory manager refactor

[Calastrophe]

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.

[mrexodia]

Been a bit busy recently, but will review when I have some more time.

[mrexodia]

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:

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).

[mrexodia]

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