Merge contiguous mapped zones

[protopyte]

I faced an issue during the Ledger CTF where a call to map_space ended up with boundaries in two distinct, contiguous, previously mapped spaces. The rainbow logic ends up trying to remap already mapped regions, which the underlying unicorn doesn't like.

$ cat diff
diff --git a/rainbow/rainbow.py b/rainbow/rainbow.py
index 4f78b1a..a1ecb89 100644
--- a/rainbow/rainbow.py
+++ b/rainbow/rainbow.py
@@ -77,9 +77,13 @@ class rainbowBase:
         self.sca_values_trace = []

     # convenience function
-    def map_space(self, a, b, verbose=False):
+    def map_space(self, a, b, verbose=True):
         """ Maps area into the unicorn emulator between a and b, or nothing if it was already mapped.
         Only completes missing portions if there is overlapping with a previously-mapped segment """
+        if verbose:
+          print()
+          print(f"map_space called with {a:x} {b-a:x}")
+          print(f"current maps:", ["%08X -> %08X" % (s[0], s[1]) for s in self.mapped_regions])
         if any(map(lambda x: a >= x[0] and b <= x[1], self.mapped_regions)):
             if verbose:
                 print(f"Did not map {a:x} {b-a:x} as it is already mapped.")

$ cat bug_rainbow_map_space.py
from rainbow.generics import rainbow_arm

emu = rainbow_arm()
emu.load("picohsm-speed-dating/firmware-mcu.elf")

$ python3 bug_rainbow_map_space.py

map_space called with affffe00 220
current maps: []
Mapping : affffc00 800

map_space called with 8000000 1268
current maps: ['AFFFFC00 -> B0000400']
Mapping : 8000000 1400

map_space called with 80016d3 12
current maps: ['AFFFFC00 -> B0000400', '08000000 -> 08001400']
Mapping : 8001400 400

map_space called with 80016e8 8
current maps: ['AFFFFC00 -> B0000400', '08000000 -> 08001400', '08001400 -> 08001800']
Did not map 80016e8 8 as it is already mapped.

map_space called with 8001268 27
current maps: ['AFFFFC00 -> B0000400', '08000000 -> 08001400', '08001400 -> 08001800']
Did not map 8001268 27 as it is already mapped.

map_space called with 800128f 444
current maps: ['AFFFFC00 -> B0000400', '08000000 -> 08001400', '08001400 -> 08001800']
Mapping : 8001400 400
Traceback (most recent call last):
  File "bug_rainbow_map_space.py", line 9, in <module>
    emu.load(target)
  File "[snip]/rainbow/rainbow/rainbow.py", line 195, in load
    return load_selector(filename, self, typ, verbose=verbose)
  File "[snip]/rainbow/rainbow/loaders/__init__.py", line 33, in load_selector
    return loader(filename, rainbow_instance, verbose=verbose)
  File "[snip]/rainbow/rainbow/loaders/elfloader.py", line 36, in elfloader
    emu.map_space(
  File "[snip]/rainbow/rainbow/rainbow.py", line 124, in map_space
    ret = self.emu.mem_map(base, size)
  File "/usr/local/lib/python3.8/dist-packages/unicorn/unicorn.py", line 448, in mem_map
    raise UcError(status)
unicorn.unicorn.UcError: Invalid memory mapping (UC_ERR_MAP)

My quick fix was to merge contiguous regions at each map_space call, but you might prefer to handle this straight upon appending a newly mapped region. Obfuscated code follows:

$ cat patch
diff --git a/rainbow/rainbow.py b/rainbow/rainbow.py
index 4f78b1a..948467d 100644
--- a/rainbow/rainbow.py
+++ b/rainbow/rainbow.py
@@ -89,19 +89,30 @@ class rainbowBase:
             return

         overlap = 0
+        to_unmap = []
         for r_start, r_end in self.mapped_regions:
             # check for overlaps
-            if a < r_start and r_start < b <= r_end:
+            # Full overlap
+            if a < r_start and r_end < b:
+                # Unicorn doesn't like mapping multiple times the same area
+                # Unmap before remapping a larger area
+                self.emu.mem_unmap(r_start, r_end-r_start)
+                to_unmap += [(r_start, r_end)]
+            # Partial high overlap
+            elif a < r_start and r_start < b <= r_end:
                 overlap = 1
                 aa = a
                 bb = r_start
                 break
+            # Partial low overlap
             elif r_end > a >= r_start and b > r_end:
                 overlap = 1
                 aa = r_end
                 bb = b
                 break

+        self.mapped_regions = [r for r in self.mapped_regions if r not in to_unmap]
+
         if overlap == 0:
             aa = a
             bb = b
@@ -120,7 +131,14 @@ class rainbowBase:
         ret = self.emu.mem_map(base, size)
         if ret is not None:
             print(ret)
-        self.mapped_regions.append((base, base + size))
+
+        # Merge contiguous intervals to avoid overlaps:
+        # - discard elements that appear both as a lower and upper interval boundary
+        # - as intervals don't overlap, sorted list contains alternating lower and upper boundaries
+        regions = sorted(list(set([t[0] for t in self.mapped_regions] + [base]) ^ set([t[1] for t in self.mapped_regions] + [base + size])))
+
+        # Rebuild mapped regions list
+        self.mapped_regions = list(zip(regions[::2], regions[1::2]))

     def __setitem__(self, inp, val):
         """ Sets a register, memory address or memory range to a value. Handles writing ints or bytes. 

$ cat test_patch.py
from rainbow.generics import rainbow_arm
import unicorn as uc

tests = [
          [
            (0x1000, 0x2000), 

            # Test: same map
            (0x1000, 0x2000), 
          ],
          [
            (0x1000, 0x2000), 

            # Test: included map
            (0x1400, 0x1800), 
          ],
          [
            (0x1000, 0x2000), 

            # Test: distinct map
            (0x5000, 0x6000), 
          ],
          [
            (0x1000, 0x2000), 

            # Test: contiguous map above
            (0x2000, 0x3000), 
          ],
          [
            (0x1000, 0x2000), 

            # Test: contiguous map below
            (0x800, 0x1000), 
          ],
          [
            (0x1000, 0x2000), 
            (0x3000, 0x4000), 

            # Test: contiguous map below and above
            (0x2000, 0x3000), 
          ],
          [
            (0x1000, 0x2000), 

            # Test: overlapping below
            (0x1800, 0x3000), 
          ],
          [
            (0x1000, 0x2000), 

            # Test: exact overlapping below
            (0x1800, 0x2000), 
          ],
          [
            (0x1000, 0x2000), 

            # Test: overlapping above
            (0x800, 0x1800), 
          ],
          [
            (0x1000, 0x2000), 

            # Test: exact overlapping above
            (0x1000, 0x1800), 
          ],
          [
            (0x1000, 0x2000), 

            # Test: inclusion of previous map
            (0x800, 0x2800), 
          ],
          [
            (0x1000, 0x2000), 
            (0x3000, 0x4000), 

            # Test: multiple inclusions of previous maps
            (0x800, 0x5000), 
          ],
          [
            (0x1000, 0x2000), 
            (0x3000, 0x4000), 

            # Test: inclusion and overlap
            (0x800, 0x3800), 
          ],
        ] 

for test in tests:
  emu = rainbow_arm()

  for t in test:
    try:
      print("Mapping:", "%04X -> %04X" % t)
      emu.map_space(*t)
      print(" Mapped:", ", ".join(["%04X -> %04X" % r for r in emu.mapped_regions if r[0] < 0x10000]))
    except uc.unicorn.UcError as e:
      print(">>> CRASHED! <<<")
  print("")
$ python3 test_patch.py
Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1400 -> 1800
 Mapped: 1000 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 5000 -> 6000
 Mapped: 1000 -> 2000, 5000 -> 6000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 2000 -> 3000
 Mapped: 1000 -> 3000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 0800 -> 1000
 Mapped: 0800 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 3000 -> 4000
 Mapped: 1000 -> 2000, 3000 -> 4000
Mapping: 2000 -> 3000
 Mapped: 1000 -> 4000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1800 -> 3000
 Mapped: 1000 -> 3000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1800 -> 2000
 Mapped: 1000 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 0800 -> 1800
 Mapped: 0800 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1000 -> 1800
 Mapped: 1000 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 0800 -> 2800
 Mapped: 0800 -> 2800

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 3000 -> 4000
 Mapped: 1000 -> 2000, 3000 -> 4000
Mapping: 0800 -> 5000
 Mapped: 0800 -> 5000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 3000 -> 4000
 Mapped: 1000 -> 2000, 3000 -> 4000
Mapping: 0800 -> 3800
 Mapped: 0800 -> 4000

Current results are as follows:

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1400 -> 1800
 Mapped: 1000 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 5000 -> 6000
 Mapped: 1000 -> 2000, 5000 -> 6000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 2000 -> 3000
 Mapped: 1000 -> 2000, 2000 -> 3000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 0800 -> 1000
 Mapped: 1000 -> 2000, 0800 -> 1000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 3000 -> 4000
 Mapped: 1000 -> 2000, 3000 -> 4000
Mapping: 2000 -> 3000
 Mapped: 1000 -> 2000, 3000 -> 4000, 2000 -> 3000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1800 -> 3000
 Mapped: 1000 -> 2000, 2000 -> 3000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1800 -> 2000
 Mapped: 1000 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 0800 -> 1800
 Mapped: 1000 -> 2000, 0800 -> 1000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 1000 -> 1800
 Mapped: 1000 -> 2000

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 0800 -> 2800
>>> CRASHED! <<<

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 3000 -> 4000
 Mapped: 1000 -> 2000, 3000 -> 4000
Mapping: 0800 -> 5000
>>> CRASHED! <<<

Mapping: 1000 -> 2000
 Mapped: 1000 -> 2000
Mapping: 3000 -> 4000
 Mapped: 1000 -> 2000, 3000 -> 4000
Mapping: 0800 -> 3800
>>> CRASHED! <<<

[yhql]

Thanks for the report! I added a fix, together with (most of) your tests :)