MMIO unmap fix

qiling/hw/hw.py

@@ -182,6 +182,6 @@ def restore(self, saved_state):
 
         # a dirty hack to rehydrate non-pickleable hwman
         # a proper fix would require a deeper refactoring to how peripherals are created and managed
-        for ph in self.ql.mem.mmio_cbs.values():
+        for *_, ph in self.ql.mem.map_info:
             if isinstance(ph, QlPripheralHandler):
                 setattr(ph, '_hwman', self)

qiling/os/memory.py

@@ -13,12 +13,6 @@
 from qiling import Qiling
 from qiling.exception import *
 
-# tuple: range start, range end, permissions mask, range label, is mmio?
-MapInfoEntry = Tuple[int, int, int, str, bool]
-
-MmioReadCallback  = Callable[[Qiling, int, int], int]
-MmioWriteCallback = Callable[[Qiling, int, int, int], None]
-
 
 class QlMmioHandler(Protocol):
     """A simple MMIO handler boilerplate that can be used to implement memory mapped devices.
@@ -36,6 +30,13 @@ def write(self, ql: Qiling, offset: int, size: int, value: int) -> None:
         ...
 
 
+# tuple: range start, range end, permissions mask, range label, mmio hander object (if mmio range)
+MapInfoEntry = Tuple[int, int, int, str, Optional[QlMmioHandler]]
+
+MmioReadCallback  = Callable[[Qiling, int, int], int]
+MmioWriteCallback = Callable[[Qiling, int, int, int], None]
+
+
 class QlMemoryManager:
     """
     some ideas and code from:
@@ -45,7 +46,6 @@ class QlMemoryManager:
     def __init__(self, ql: Qiling, pagesize: int = 0x1000):
         self.ql = ql
         self.map_info: List[MapInfoEntry] = []
-        self.mmio_cbs: Dict[Tuple[int, int], QlMmioHandler] = {}
 
         bit_stuff = {
             64: (1 << 64) - 1,
@@ -121,18 +121,18 @@ def string(self, addr: int, value=None, encoding='utf-8') -> Optional[str]:
 
         self.__write_string(addr, value, encoding)
 
-    def add_mapinfo(self, mem_s: int, mem_e: int, mem_p: int, mem_info: str, is_mmio: bool = False):
+    def add_mapinfo(self, mem_s: int, mem_e: int, mem_p: int, mem_info: str, mmio_ctx: Optional[QlMmioHandler] = None):
         """Add a new memory range to map.
 
         Args:
             mem_s: memory range start
             mem_e: memory range end
             mem_p: permissions mask
             mem_info: map entry label
-            is_mmio: memory range is mmio
+            mmio_ctx: mmio handler object; if specified the range will be treated as mmio
         """
 
-        bisect.insort(self.map_info, (mem_s, mem_e, mem_p, mem_info, is_mmio))
+        bisect.insort(self.map_info, (mem_s, mem_e, mem_p, mem_info, mmio_ctx))
 
     def del_mapinfo(self, mem_s: int, mem_e: int):
         """Subtract a memory range from map.
@@ -146,13 +146,13 @@ def del_mapinfo(self, mem_s: int, mem_e: int):
 
         def __split_overlaps():
             for idx in overlap_ranges:
-                lbound, ubound, perms, label, is_mmio = self.map_info[idx]
+                lbound, ubound, perms, label, mmio_ctx = self.map_info[idx]
 
                 if lbound < mem_s:
-                    yield (lbound, mem_s, perms, label, is_mmio)
+                    yield (lbound, mem_s, perms, label, mmio_ctx)
 
                 if mem_e < ubound:
-                    yield (mem_e, ubound, perms, label, is_mmio)
+                    yield (mem_e, ubound, perms, label, mmio_ctx)
 
         # indices of first and last overlapping ranges. since map info is always
         # sorted, we know that all overlapping rages are consecutive, so i1 > i0
@@ -170,27 +170,30 @@ def __split_overlaps():
         for entry in new_entries:
             bisect.insort(self.map_info, entry)
 
-    def change_mapinfo(self, mem_s: int, mem_e: int, mem_p: Optional[int] = None, mem_info: Optional[str] = None):
-        tmp_map_info: Optional[MapInfoEntry] = None
-        info_idx: int = -1
-
-        for idx, map_info in enumerate(self.map_info):
-            if mem_s >= map_info[0] and mem_e <= map_info[1]:
-                tmp_map_info = map_info
-                info_idx = idx
-                break
+    def change_mapinfo(self, mem_s: int, mem_e: int, *, new_perms: Optional[int] = None, new_info: Optional[str] = None) -> None:
+        if new_perms is None and new_info is None:
+            # nothing to do
+            return
 
-        if tmp_map_info is None:
+        try:
+            # locate the map info entry to change
+            entry = next(entry for entry in self.map_info if mem_s >= entry[0] and mem_e <= entry[1])
+        except StopIteration:
             self.ql.log.error(f'Cannot change mapinfo at {mem_s:#08x}-{mem_e:#08x}')
             return
 
-        if mem_p is not None:
-            self.del_mapinfo(mem_s, mem_e)
-            self.add_mapinfo(mem_s, mem_e, mem_p, mem_info if mem_info else tmp_map_info[3])
-            return
+        _, _, perms, info, mmio_ctx = entry
+
+        # caller wants to change perms?
+        if new_perms is not None:
+            perms = new_perms
 
-        if mem_info is not None:
-            self.map_info[info_idx] = (tmp_map_info[0], tmp_map_info[1], tmp_map_info[2], mem_info, tmp_map_info[4])
+        # caller wants to change info?
+        if new_info is not None:
+            info = new_info
+
+        self.del_mapinfo(mem_s, mem_e)
+        self.add_mapinfo(mem_s, mem_e, perms, info, mmio_ctx)
 
     def get_mapinfo(self) -> Sequence[Tuple[int, int, str, str, str]]:
         """Get memory map info.
@@ -209,18 +212,18 @@ def __perms_mapping(ps: int) -> str:
 
             return ''.join(val if idx & ps else '-' for idx, val in perms_d.items())
 
-        def __process(lbound: int, ubound: int, perms: int, label: str, is_mmio: bool) -> Tuple[int, int, str, str, str]:
-            perms_str = __perms_mapping(perms)
+        def __process(entry: MapInfoEntry) -> Tuple[int, int, str, str, str]:
+            lbound, ubound, perms, label, mmio_ctx = entry
 
             if hasattr(self.ql, 'loader'):
                 image = self.ql.loader.find_containing_image(lbound)
-                container = image.path if image and not is_mmio else ''
+                container = image.path if image and mmio_ctx is None else ''
             else:
                 container = ''
 
-            return (lbound, ubound, perms_str, label, container)
+            return (lbound, ubound, __perms_mapping(perms), label, container)
 
-        return tuple(__process(*entry) for entry in self.map_info)
+        return tuple(__process(entry) for entry in self.map_info)
 
     def get_formatted_mapinfo(self) -> Sequence[str]:
         """Get memory map info in a nicely formatted table.
@@ -311,12 +314,13 @@ def save(self):
             "mmio" : []
         }
 
-        for lbound, ubound, perm, label, is_mmio in self.map_info:
-            if is_mmio:
-                mem_dict['mmio'].append((lbound, ubound, perm, label, self.mmio_cbs[(lbound, ubound)]))
+        for lbound, ubound, perm, label, mmio_ctx in self.map_info:
+            if mmio_ctx is None:
+                key, data = 'ram', bytes(self.read(lbound, ubound - lbound))
             else:
-                data = self.read(lbound, ubound - lbound)
-                mem_dict['ram'].append((lbound, ubound, perm, label, bytes(data)))
+                key, data = 'mmio', mmio_ctx
+
+            mem_dict[key].append((lbound, ubound, perm, label, data))
 
         return mem_dict
 
@@ -429,7 +433,7 @@ def search(self, needle: Union[bytes, Pattern[bytes]], begin: Optional[int] = No
         assert begin < end, 'search arguments do not make sense'
 
         # narrow the search down to relevant ranges; mmio ranges are excluded due to potential read side effects
-        ranges = [(max(begin, lbound), min(ubound, end)) for lbound, ubound, _, _, is_mmio in self.map_info if not (end < lbound or ubound < begin or is_mmio)]
+        ranges = [(max(begin, lbound), min(ubound, end)) for lbound, ubound, _, _, mmio_ctx in self.map_info if not (end < lbound or ubound < begin or mmio_ctx is not None)]
         results = []
 
         # if needle is a bytes sequence use it verbatim, not as a pattern
@@ -455,9 +459,6 @@ def unmap(self, addr: int, size: int) -> None:
         self.del_mapinfo(addr, addr + size)
         self.ql.uc.mem_unmap(addr, size)
 
-        if (addr, addr + size) in self.mmio_cbs:
-            del self.mmio_cbs[(addr, addr+size)]
-
     def unmap_between(self, mem_s: int, mem_e: int) -> None:
         """Reclaim any allocated memory region within the specified range.
 
@@ -616,7 +617,7 @@ def protect(self, addr: int, size: int, perms):
         aligned_size = self.align_up((addr & (self.pagesize - 1)) + size)
 
         self.ql.uc.mem_protect(aligned_address, aligned_size, perms)
-        self.change_mapinfo(aligned_address, aligned_address + aligned_size, perms)
+        self.change_mapinfo(aligned_address, aligned_address + aligned_size, new_perms=perms)
 
     def map(self, addr: int, size: int, perms: int = UC_PROT_ALL, info: Optional[str] = None):
         """Map a new memory range.
@@ -638,7 +639,7 @@ def map(self, addr: int, size: int, perms: int = UC_PROT_ALL, info: Optional[str
             raise QlMemoryMappedError('Requested memory is unavailable')
 
         self.ql.uc.mem_map(addr, size, perms)
-        self.add_mapinfo(addr, addr + size, perms, info or '[mapped]', is_mmio=False)
+        self.add_mapinfo(addr, addr + size, perms, info or '[mapped]', None)
 
     def map_mmio(self, addr: int, size: int, handler: QlMmioHandler, info: str = '[mmio]'):
         # TODO: mmio memory overlap with ram? Is that possible?
@@ -664,9 +665,7 @@ def __mmio_write(uc, offset: int, size: int, value: int, user_data: MmioWriteCal
             cb(self.ql, offset, size, value)
 
         self.ql.uc.mmio_map(addr, size, __mmio_read, handler.read, __mmio_write, handler.write)
-        self.add_mapinfo(addr, addr + size, prot, info, is_mmio=True)
-
-        self.mmio_cbs[(addr, addr + size)] = handler
+        self.add_mapinfo(addr, addr + size, prot, info, handler)
 
 
 class Chunk: