From 2ec02a17049dc5e95d7c6f3506b531435c7c5899 Mon Sep 17 00:00:00 2001
From: mio <mio@lazym.io>
Date: Sat, 5 Mar 2022 00:16:19 +0100
Subject: [PATCH 01/11] Use multitask unicorn for mcu

---
 qiling/arch/cortex_m.py        |  38 +++--
 qiling/extensions/multitask.py | 272 +++++++++++++++++++++++++++++++++
 2 files changed, 300 insertions(+), 10 deletions(-)
 create mode 100644 qiling/extensions/multitask.py

diff --git a/qiling/arch/cortex_m.py b/qiling/arch/cortex_m.py
index a9e0a64ee..8bd9fbe14 100644
--- a/qiling/arch/cortex_m.py
+++ b/qiling/arch/cortex_m.py
@@ -6,7 +6,7 @@
 from functools import cached_property
 from contextlib import ContextDecorator
 
-from unicorn import Uc, UC_ARCH_ARM, UC_MODE_ARM, UC_MODE_MCLASS, UC_MODE_THUMB
+from unicorn import Uc, UcError, UC_ARCH_ARM, UC_MODE_ARM, UC_MODE_MCLASS, UC_MODE_THUMB, UC_ERR_OK
 from capstone import Cs, CS_ARCH_ARM, CS_MODE_ARM, CS_MODE_MCLASS, CS_MODE_THUMB
 from keystone import Ks, KS_ARCH_ARM, KS_MODE_ARM, KS_MODE_THUMB
 
@@ -18,6 +18,8 @@
 from qiling.const import QL_ARCH, QL_ENDIAN, QL_VERBOSE
 from qiling.exception import QlErrorNotImplemented
 
+from qiling.extensions.multitask import MultiTaskUnicorn, UnicornTask
+
 class QlInterruptContext(ContextDecorator):
     def __init__(self, ql: Qiling):
         self.ql = ql
@@ -59,6 +61,25 @@ def __exit__(self, *exc):
         if self.ql.verbose >= QL_VERBOSE.DISASM:
             self.ql.log.info('Exit from interrupt')
 
+
+# This class exits to emulate clock interrupt.
+class QlArchCORTEX_MThread(UnicornTask):
+
+    def __init__(self, ql: "Qiling", begin: int, end: int, task_id=None):
+        super().__init__(ql.uc, begin, end, task_id)
+        self.ql = ql
+    
+    def on_start(self):
+        # Don't save anything.
+        return None
+    
+    def on_interrupted(self, ucerr: int):
+        # And don't restore anything.
+        if ucerr != UC_ERR_OK:
+            raise UcError(ucerr)
+
+        self.ql.hw.step()
+
 class QlArchCORTEX_M(QlArchARM):
     type = QL_ARCH.ARM
     bits = 32
@@ -94,10 +115,6 @@ def is_thumb(self) -> bool:
     def endian(self) -> QL_ENDIAN:
         return QL_ENDIAN.EL
 
-    def step(self):
-        self.ql.emu_start(self.effective_pc, 0, count=1)
-        self.ql.hw.step()
-
     def stop(self):
         self.ql.emu_stop()
         self.runable = False
@@ -108,12 +125,13 @@ def run(self, count=-1, end=None):
         if type(end) is int:
             end |= 1        
         
-        while self.runable and count != 0:
-            if self.effective_pc == end:
-                break
+        if end is None:
+            end = 0
 
-            self.step()
-            count -= 1    
+        self.mtuc = MultiTaskUnicorn(self.uc)
+        utk = QlArchCORTEX_MThread(self.ql, self.effective_pc, end)
+        self.mtuc.task_create(utk)
+        self.mtuc.start()
 
     def is_handler_mode(self):
         return self.regs.ipsr > 1
diff --git a/qiling/extensions/multitask.py b/qiling/extensions/multitask.py
new file mode 100644
index 000000000..30f2706ed
--- /dev/null
+++ b/qiling/extensions/multitask.py
@@ -0,0 +1,272 @@
+# Ziqiao Kong (mio@lazym.io)
+
+from typing import Dict
+from unicorn import *
+from unicorn.x86_const import UC_X86_REG_EIP, UC_X86_REG_RIP
+from unicorn.arm64_const import UC_ARM64_REG_PC
+from unicorn.arm_const import UC_ARM_REG_PC, UC_ARM_REG_CPSR
+from unicorn.mips_const import UC_MIPS_REG_PC
+from unicorn.m68k_const import UC_M68K_REG_PC
+from unicorn.riscv_const import UC_RISCV_REG_PC
+from unicorn.ppc_const import UC_PPC_REG_PC
+from unicorn.sparc_const import UC_SPARC_REG_PC
+
+import gevent
+import gevent.threadpool
+import gevent.lock
+import threading
+
+# This class is named UnicornTask be design since it's not a
+# real thread. The expected usage is to inherit this class
+# and overwrite specific methods.
+class UnicornTask:
+
+    def __init__(self, uc: Uc, begin: int, end: int, task_id = None):
+        self._uc = uc
+        self._begin = begin
+        self._end = end
+        self._stop_request = False
+        self._ctx = None
+        self._task_id = None
+
+    @property
+    def pc(self):
+        arch = self._uc.ctl_get_arch()
+        mode = self._uc.ctl_get_mode()
+
+        # This extension is designed to be independent of Qiling, so let's
+        # do this manually...
+        if arch == UC_ARCH_X86:
+            if (mode & UC_MODE_32) != 0:
+                return self._uc.reg_read(UC_X86_REG_EIP)
+            elif (mode & UC_MODE_64) != 0:
+                return self._uc.reg_read(UC_X86_REG_RIP)
+        elif arch == UC_ARCH_MIPS:
+            return self._uc.reg_read(UC_MIPS_REG_PC)
+        elif arch == UC_ARCH_ARM:
+            pc = self._uc.reg_read(UC_ARM_REG_PC)
+            if (self._uc.reg_read(UC_ARM_REG_CPSR) & (1 << 5)):
+                return pc | 1
+            else:
+                return pc
+        elif arch == UC_ARCH_ARM64:
+            return self._uc.reg_read(UC_ARM64_REG_PC)
+        elif arch == UC_ARCH_PPC:
+            return self._uc.reg_read(UC_PPC_REG_PC)
+        elif arch == UC_ARCH_M68K:
+            return self._uc.reg_read(UC_M68K_REG_PC)
+        elif arch == UC_ARCH_SPARC:
+            return self._uc.reg_read(UC_SPARC_REG_PC)
+        elif arch == UC_ARCH_RISCV:
+            return self._uc.reg_read(UC_RISCV_REG_PC)
+        
+        # Really?
+        return 0
+
+    def save(self):
+        """ This method is used to save the task context.
+            Overwrite this method to implement specifc logic.
+        """
+        return self._uc.context_save()
+    
+    def restore(self, context):
+        """ This method is used to restore the task context.
+            Overwrite this method to implement specific logic.
+        """
+        self._uc.context_restore(context)
+        self._begin = self.pc
+
+    def _run(self):
+        # This method is not intended to be overwritten!
+        try:
+            self._uc.emu_start(self._begin, self._end, 0, 0)
+        except UcError as err:
+            return err.errno
+        
+        return UC_ERR_OK
+
+    def on_start(self):
+        """ This callback is triggered when a task gets scheduled.
+        """
+        if self._ctx:
+            self.restore(self._ctx)
+
+    def on_interrupted(self, ucerr: int):
+        """ This callback is triggered when a task gets interrupted, which
+            is useful to emulate a clock interrupt.
+        """
+        self._ctx = self.save()
+
+    def on_exit(self):
+        """ This callback is triggered when a task is about to exit.
+        """
+        pass
+
+# This is the core scheduler of a multi task unicorn.
+# To implement a non-block syscall:
+#    1. Record the syscall in the hook, but **do nothing**
+#    2. Stop emulation.
+#    3. Handle the syscall in the on_interruped callback with
+#       proper **gevent functions** like gevent.sleep instead
+#       of time.sleep.
+#    4. In this case, gevent would schedule another task to
+#       take emulation if the task is blocked.
+#
+# Bear in mind that only one task can be picked to emulate at
+# the same time.
+class MultiTaskUnicorn:
+
+    def __init__(self, uc: Uc, interval: float = 0.5):
+        # This takes over the ownershtip of uc instance
+        self._uc = uc
+        self._interval = interval
+        self._tasks = {} # type: Dict[int, UnicornTask]
+        self._task_id_counter = 2000
+        self._to_stop = False
+        self._cur_utk_id = None
+        self._running = False
+        self._run_lock = threading.RLock()
+
+    @property
+    def current_thread(self):
+        return self._tasks[self._cur_utk_id]
+
+    @property
+    def running(self):
+        return self._running
+    
+    def _next_task_id(self):
+        while self._task_id_counter in self._tasks:
+            self._task_id_counter += 1
+        
+        return self._task_id_counter
+
+    def _utk_run(self, utk: UnicornTask):
+        # Only one thread can start emulation at the same time, but
+        # some other greenlets may do other async work.
+        with self._run_lock:
+            self._running = True
+            result = utk._run()
+            self._running = False
+        return result
+
+    def _task_main(self, utk_id: int):
+
+        utk = self._tasks[utk_id]
+
+        while True:
+            # utk may be stopped before running once, check it.
+            if utk._stop_request:
+                break
+            
+            self._cur_utk_id = utk_id
+
+            utk.on_start()
+
+            with gevent.Timeout(self._interval, False):
+                try:
+                    pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
+                    task = pool.spawn(self._utk_run, utk) # Run unicorn in a separate thread.
+                    task.wait()
+                finally:
+                    if not task.done():
+                        # Interrupted by timeout, in this case we call uc_emu_stop.
+                        self._uc.emu_stop()
+                    
+                    # Wait until we get the result.
+                    ucerr = task.get()
+
+            if utk._stop_request:
+                utk.on_exit()
+                break
+            else:
+                utk.on_interrupted(ucerr)
+
+            if self._to_stop:
+                break
+            
+            # on_interrupted callback may have asked us to stop.
+            if utk._stop_request:
+                break
+
+            # Give up control at once.
+            gevent.sleep(0)
+
+        del self._tasks[utk_id]
+
+    def save(self):
+        return { k: v.save() for k, v in self._tasks.items() }
+
+    def restore(self, threads_context: dict):
+        for task_id, context in threads_context:
+            if task_id in self._tasks:
+                self._tasks[task_id].restore(context)
+
+    def task_create(self, utk: UnicornTask):
+        """ Create a unicorn task. utk should be a initialized UnicornTask object.
+            If the task_id is not set, we generate one.
+
+            utk: The task to add.
+        """
+        if not isinstance(utk, UnicornTask):
+            raise TypeError("Expect a UnicornTask or derived class")
+        if utk._task_id is None:
+            utk._task_id = self._next_task_id()
+        self._tasks[utk._task_id] = utk
+        return utk._task_id
+
+    def task_exit(self, utk_id):
+        """ Stop a task.
+            
+            utk_id: The id returned from task_create.
+        """
+        if utk_id not in self._tasks:
+            return
+        
+        if utk_id == self._cur_utk_id and self._running:
+            self._uc.emu_stop()
+        
+        self._tasks[utk_id]._stop_request = True
+
+    def emu_once(self, begin: int, end: int, timeout: int, count: int):
+        """ Emulate an area of code just once. This is equivalent to uc_emu_start but is gevent-aware.
+            NOTE: Calling this method may cause current greenlet to be switched out.
+
+            begin, end, timeout, count: refer to uc_emu_start
+        """
+        def _once(begin: int, end: int, timeout: int, count: int):
+            with self._run_lock:
+                try:
+                    self._uc.emu_start(begin, end, timeout, count)
+                except UcError as err:
+                    return err.errno
+            
+            return UC_ERR_OK
+
+        pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
+        task = pool.spawn(_once, begin, end, timeout, count)
+        return task.get()
+
+    def stop(self):
+        """ This will stop all running tasks.
+        """
+        self._to_stop = True
+        if self._running:
+            self._uc.emu_stop()
+
+    def start(self):
+        """ This will start emulation until all tasks get done.
+        """
+        workset = {} # type: Dict[int, gevent.Greenlet]
+        self._to_stop = False
+
+        while len(self._tasks) != 0:
+            
+            new_workset = { k: v for k, v in workset.items() if not v.dead}
+
+            for utk_id in self._tasks:
+                new_workset[utk_id] = gevent.spawn(self._task_main, utk_id)
+
+            workset = new_workset
+
+            gevent.joinall(list(workset.values()), raise_error=True)
\ No newline at end of file

From fd9e6b1520c1261e2f718a36473fc69c06a0dd30 Mon Sep 17 00:00:00 2001
From: mio <mio@lazym.io>
Date: Sat, 5 Mar 2022 21:37:27 +0100
Subject: [PATCH 02/11] Fix task restore

---
 qiling/arch/cortex_m.py        | 7 +++++--
 qiling/extensions/multitask.py | 5 +++++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/qiling/arch/cortex_m.py b/qiling/arch/cortex_m.py
index 8bd9fbe14..2c142910d 100644
--- a/qiling/arch/cortex_m.py
+++ b/qiling/arch/cortex_m.py
@@ -74,6 +74,8 @@ def on_start(self):
         return None
     
     def on_interrupted(self, ucerr: int):
+        self._begin = self.pc
+
         # And don't restore anything.
         if ucerr != UC_ERR_OK:
             raise UcError(ucerr)
@@ -128,7 +130,8 @@ def run(self, count=-1, end=None):
         if end is None:
             end = 0
 
-        self.mtuc = MultiTaskUnicorn(self.uc)
+        self.mtuc = MultiTaskUnicorn(self.uc, 0.1)
+        self.ql.hook_code(lambda x, y, z: 0)
         utk = QlArchCORTEX_MThread(self.ql, self.effective_pc, end)
         self.mtuc.task_create(utk)
         self.mtuc.start()
@@ -200,4 +203,4 @@ def hard_interrupt_handler(self, ql, intno):
             self.regs.write('pc', entry)
             self.regs.write('lr', exc_return) 
 
-            self.ql.emu_start(self.effective_pc, 0, count=0xffffff)
+            self.mtuc.emu_once(self.effective_pc, 0, 0, 0)
diff --git a/qiling/extensions/multitask.py b/qiling/extensions/multitask.py
index 30f2706ed..62daaf072 100644
--- a/qiling/extensions/multitask.py
+++ b/qiling/extensions/multitask.py
@@ -247,6 +247,11 @@ def _once(begin: int, end: int, timeout: int, count: int):
         task = pool.spawn(_once, begin, end, timeout, count)
         return task.get()
 
+    def emu_stop(self):
+        # Interrupt the emulation
+        if self._running:
+            self._uc.emu_stop()
+
     def stop(self):
         """ This will stop all running tasks.
         """

From 9f3a05ab3ac4aba37bdbf866726b92d626aff644 Mon Sep 17 00:00:00 2001
From: mio <mio@lazym.io>
Date: Sat, 5 Mar 2022 22:10:11 +0100
Subject: [PATCH 03/11] Use a more proper interval

---
 qiling/arch/cortex_m.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/qiling/arch/cortex_m.py b/qiling/arch/cortex_m.py
index 2c142910d..f159fd9ef 100644
--- a/qiling/arch/cortex_m.py
+++ b/qiling/arch/cortex_m.py
@@ -130,7 +130,7 @@ def run(self, count=-1, end=None):
         if end is None:
             end = 0
 
-        self.mtuc = MultiTaskUnicorn(self.uc, 0.1)
+        self.mtuc = MultiTaskUnicorn(self.uc, 0.01)
         self.ql.hook_code(lambda x, y, z: 0)
         utk = QlArchCORTEX_MThread(self.ql, self.effective_pc, end)
         self.mtuc.task_create(utk)

From ab56097658ee87158a61a82a86e24ecef11e5148 Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sat, 5 Mar 2022 22:33:38 +0100
Subject: [PATCH 04/11] Remove the dumb hook

This requires unicorn commit 7cd475d45b398c9954b8685b038ed5b37a668c60
---
 qiling/arch/cortex_m.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/qiling/arch/cortex_m.py b/qiling/arch/cortex_m.py
index f159fd9ef..bfb1ff719 100644
--- a/qiling/arch/cortex_m.py
+++ b/qiling/arch/cortex_m.py
@@ -131,7 +131,6 @@ def run(self, count=-1, end=None):
             end = 0
 
         self.mtuc = MultiTaskUnicorn(self.uc, 0.01)
-        self.ql.hook_code(lambda x, y, z: 0)
         utk = QlArchCORTEX_MThread(self.ql, self.effective_pc, end)
         self.mtuc.task_create(utk)
         self.mtuc.start()

From 02dcb939892e97fc8957ae9e9d14f7d4d2e4adac Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sat, 5 Mar 2022 23:04:12 +0100
Subject: [PATCH 05/11] Change how MultitaskUnicorn works

---
 qiling/arch/cortex_m.py        |   9 ++-
 qiling/extensions/multitask.py | 113 +++++++++++++++++----------------
 2 files changed, 63 insertions(+), 59 deletions(-)

diff --git a/qiling/arch/cortex_m.py b/qiling/arch/cortex_m.py
index bfb1ff719..b115c687f 100644
--- a/qiling/arch/cortex_m.py
+++ b/qiling/arch/cortex_m.py
@@ -91,7 +91,7 @@ def __init__(self, ql: Qiling):
 
     @cached_property
     def uc(self):
-        return Uc(UC_ARCH_ARM, UC_MODE_ARM + UC_MODE_MCLASS + UC_MODE_THUMB)
+        return MultiTaskUnicorn(UC_ARCH_ARM, UC_MODE_ARM + UC_MODE_MCLASS + UC_MODE_THUMB, 0.01)
 
     @cached_property
     def regs(self) -> QlRegisterManager:
@@ -130,10 +130,9 @@ def run(self, count=-1, end=None):
         if end is None:
             end = 0
 
-        self.mtuc = MultiTaskUnicorn(self.uc, 0.01)
         utk = QlArchCORTEX_MThread(self.ql, self.effective_pc, end)
-        self.mtuc.task_create(utk)
-        self.mtuc.start()
+        self.uc.task_create(utk)
+        self.uc.tasks_start()
 
     def is_handler_mode(self):
         return self.regs.ipsr > 1
@@ -202,4 +201,4 @@ def hard_interrupt_handler(self, ql, intno):
             self.regs.write('pc', entry)
             self.regs.write('lr', exc_return) 
 
-            self.mtuc.emu_once(self.effective_pc, 0, 0, 0)
+            self.uc.emu_start(self.effective_pc, 0, 0, 0)
diff --git a/qiling/extensions/multitask.py b/qiling/extensions/multitask.py
index 62daaf072..4cacd2463 100644
--- a/qiling/extensions/multitask.py
+++ b/qiling/extensions/multitask.py
@@ -76,15 +76,6 @@ def restore(self, context):
         self._uc.context_restore(context)
         self._begin = self.pc
 
-    def _run(self):
-        # This method is not intended to be overwritten!
-        try:
-            self._uc.emu_start(self._begin, self._end, 0, 0)
-        except UcError as err:
-            return err.errno
-        
-        return UC_ERR_OK
-
     def on_start(self):
         """ This callback is triggered when a task gets scheduled.
         """
@@ -114,11 +105,10 @@ def on_exit(self):
 #
 # Bear in mind that only one task can be picked to emulate at
 # the same time.
-class MultiTaskUnicorn:
+class MultiTaskUnicorn(Uc):
 
-    def __init__(self, uc: Uc, interval: float = 0.5):
-        # This takes over the ownershtip of uc instance
-        self._uc = uc
+    def __init__(self, arch, mode, interval=0.1):
+        super().__init__(arch, mode)
         self._interval = interval
         self._tasks = {} # type: Dict[int, UnicornTask]
         self._task_id_counter = 2000
@@ -126,6 +116,7 @@ def __init__(self, uc: Uc, interval: float = 0.5):
         self._cur_utk_id = None
         self._running = False
         self._run_lock = threading.RLock()
+        self._multitask_enabled = False
 
     @property
     def current_thread(self):
@@ -134,21 +125,26 @@ def current_thread(self):
     @property
     def running(self):
         return self._running
-    
+
     def _next_task_id(self):
         while self._task_id_counter in self._tasks:
             self._task_id_counter += 1
         
         return self._task_id_counter
 
-    def _utk_run(self, utk: UnicornTask):
-        # Only one thread can start emulation at the same time, but
-        # some other greenlets may do other async work.
+    def _emu_start_locked(self, begin: int, end: int, timeout: int, count: int):
         with self._run_lock:
-            self._running = True
-            result = utk._run()
-            self._running = False
-        return result
+            try:
+                self._running = True
+                super().emu_start(begin, end, timeout, count)
+                self._running = False
+            except UcError as err:
+                return err.errno
+        
+        return UC_ERR_OK
+
+    def _emu_start_utk_locked(self, utk: UnicornTask):
+        return self._emu_start_locked(utk._begin, utk._end, 0, 0)
 
     def _task_main(self, utk_id: int):
 
@@ -166,12 +162,12 @@ def _task_main(self, utk_id: int):
             with gevent.Timeout(self._interval, False):
                 try:
                     pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
-                    task = pool.spawn(self._utk_run, utk) # Run unicorn in a separate thread.
+                    task = pool.spawn(self._emu_start_utk_locked, utk) # Run unicorn in a separate thread.
                     task.wait()
                 finally:
                     if not task.done():
                         # Interrupted by timeout, in this case we call uc_emu_stop.
-                        self._uc.emu_stop()
+                        self.emu_stop()
                     
                     # Wait until we get the result.
                     ucerr = task.get()
@@ -194,11 +190,11 @@ def _task_main(self, utk_id: int):
 
         del self._tasks[utk_id]
 
-    def save(self):
+    def tasks_save(self):
         return { k: v.save() for k, v in self._tasks.items() }
 
-    def restore(self, threads_context: dict):
-        for task_id, context in threads_context:
+    def tasks_restore(self, tasks_context: dict):
+        for task_id, context in tasks_context:
             if task_id in self._tasks:
                 self._tasks[task_id].restore(context)
 
@@ -210,6 +206,7 @@ def task_create(self, utk: UnicornTask):
         """
         if not isinstance(utk, UnicornTask):
             raise TypeError("Expect a UnicornTask or derived class")
+        self._multitask_enabled = True
         if utk._task_id is None:
             utk._task_id = self._next_task_id()
         self._tasks[utk._task_id] = utk
@@ -224,54 +221,62 @@ def task_exit(self, utk_id):
             return
         
         if utk_id == self._cur_utk_id and self._running:
-            self._uc.emu_stop()
+            self.emu_stop()
         
         self._tasks[utk_id]._stop_request = True
 
-    def emu_once(self, begin: int, end: int, timeout: int, count: int):
+    def emu_start(self, begin: int, end: int, timeout: int, count: int):
         """ Emulate an area of code just once. This is equivalent to uc_emu_start but is gevent-aware.
             NOTE: Calling this method may cause current greenlet to be switched out.
 
             begin, end, timeout, count: refer to uc_emu_start
         """
-        def _once(begin: int, end: int, timeout: int, count: int):
-            with self._run_lock:
-                try:
-                    self._uc.emu_start(begin, end, timeout, count)
-                except UcError as err:
-                    return err.errno
-            
-            return UC_ERR_OK
 
-        pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
-        task = pool.spawn(_once, begin, end, timeout, count)
-        return task.get()
+        if self._multitask_enabled:
+            pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
+            task = pool.spawn(self._emu_start_locked, begin, end, timeout, count)
+            ucerr = task.get()
+
+            if ucerr != UC_ERR_OK:
+                raise UcError(ucerr)
+            
+            return ucerr
+        else:
+            return super().emu_start(begin, end, timeout, count)
 
     def emu_stop(self):
-        # Interrupt the emulation
-        if self._running:
-            self._uc.emu_stop()
+        """ Interrupt the emulation. This method mimic the standard Uc object.
+        """
+        if self._multitask_enabled:
+            if self._running:
+                super().emu_stop()
+        else:
+            super().emu_stop()
 
-    def stop(self):
+    def tasks_stop(self):
         """ This will stop all running tasks.
         """
-        self._to_stop = True
-        if self._running:
-            self._uc.emu_stop()
+        if self._multitask_enabled:
+            self._to_stop = True
+            if self._running:
+                super().emu_stop()
 
-    def start(self):
+    def tasks_start(self):
         """ This will start emulation until all tasks get done.
         """
         workset = {} # type: Dict[int, gevent.Greenlet]
         self._to_stop = False
 
-        while len(self._tasks) != 0:
-            
-            new_workset = { k: v for k, v in workset.items() if not v.dead}
+        if self._multitask_enabled:
+            while len(self._tasks) != 0:
+                
+                new_workset = { k: v for k, v in workset.items() if not v.dead}
 
-            for utk_id in self._tasks:
-                new_workset[utk_id] = gevent.spawn(self._task_main, utk_id)
+                for utk_id in self._tasks:
+                    new_workset[utk_id] = gevent.spawn(self._task_main, utk_id)
 
-            workset = new_workset
+                workset = new_workset
 
-            gevent.joinall(list(workset.values()), raise_error=True)
\ No newline at end of file
+                gevent.joinall(list(workset.values()), raise_error=True)
+            
+            self._multitask_enabled = False
\ No newline at end of file

From 7c727c6f9e270876057a5f0d4c3a42f4c1b36f86 Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sat, 5 Mar 2022 23:10:45 +0100
Subject: [PATCH 06/11] Update comments

---
 qiling/extensions/multitask.py | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/qiling/extensions/multitask.py b/qiling/extensions/multitask.py
index 4cacd2463..fc5f28340 100644
--- a/qiling/extensions/multitask.py
+++ b/qiling/extensions/multitask.py
@@ -19,12 +19,14 @@
 # This class is named UnicornTask be design since it's not a
 # real thread. The expected usage is to inherit this class
 # and overwrite specific methods.
+#
+# This class is a friend class of MultiTaskUnicorn
 class UnicornTask:
 
     def __init__(self, uc: Uc, begin: int, end: int, task_id = None):
         self._uc = uc
         self._begin = begin
-        self._end = end
+        self._end = end 
         self._stop_request = False
         self._ctx = None
         self._task_id = None
@@ -93,8 +95,11 @@ def on_exit(self):
         """
         pass
 
-# This is the core scheduler of a multi task unicorn.
-# To implement a non-block syscall:
+# This mimic a Unicorn object by maintaining the same interface.
+# If no task is registered, the behavior is exactly the same as
+# a normal unicorn.
+#
+# Note: To implement a non-block syscall:
 #    1. Record the syscall in the hook, but **do nothing**
 #    2. Stop emulation.
 #    3. Handle the syscall in the on_interruped callback with
@@ -191,9 +196,13 @@ def _task_main(self, utk_id: int):
         del self._tasks[utk_id]
 
     def tasks_save(self):
+        """ Save all tasks' contexts.
+        """
         return { k: v.save() for k, v in self._tasks.items() }
 
     def tasks_restore(self, tasks_context: dict):
+        """ Restore the contexts of all tasks.
+        """
         for task_id, context in tasks_context:
             if task_id in self._tasks:
                 self._tasks[task_id].restore(context)
@@ -226,10 +235,13 @@ def task_exit(self, utk_id):
         self._tasks[utk_id]._stop_request = True
 
     def emu_start(self, begin: int, end: int, timeout: int, count: int):
-        """ Emulate an area of code just once. This is equivalent to uc_emu_start but is gevent-aware.
+        """ Emulate an area of code just once. This overwrites the original emu_start interface and
+            provides extra cares when multitask is enabled. If no task is registerd, this call bahaves
+            like the original emu_start.
+
             NOTE: Calling this method may cause current greenlet to be switched out.
 
-            begin, end, timeout, count: refer to uc_emu_start
+            begin, end, timeout, count: refer to Uc.emu_start
         """
 
         if self._multitask_enabled:
@@ -245,7 +257,7 @@ def emu_start(self, begin: int, end: int, timeout: int, count: int):
             return super().emu_start(begin, end, timeout, count)
 
     def emu_stop(self):
-        """ Interrupt the emulation. This method mimic the standard Uc object.
+        """ Stop the emulation. If no task is registerd, this call bahaves like the original emu_stop.
         """
         if self._multitask_enabled:
             if self._running:
@@ -254,7 +266,7 @@ def emu_stop(self):
             super().emu_stop()
 
     def tasks_stop(self):
-        """ This will stop all running tasks.
+        """ This will stop all running tasks. If no task is registered, this call does nothing.
         """
         if self._multitask_enabled:
             self._to_stop = True

From 1c64b41244c717940f8ae6484d1f6d48a5986b68 Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sat, 5 Mar 2022 23:45:09 +0100
Subject: [PATCH 07/11] Support count and timeout primitive

---
 qiling/arch/cortex_m.py        |  4 +--
 qiling/extensions/multitask.py | 58 ++++++++++++++++++++++++++--------
 2 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/qiling/arch/cortex_m.py b/qiling/arch/cortex_m.py
index b115c687f..1b9a10bd7 100644
--- a/qiling/arch/cortex_m.py
+++ b/qiling/arch/cortex_m.py
@@ -91,7 +91,7 @@ def __init__(self, ql: Qiling):
 
     @cached_property
     def uc(self):
-        return MultiTaskUnicorn(UC_ARCH_ARM, UC_MODE_ARM + UC_MODE_MCLASS + UC_MODE_THUMB, 0.01)
+        return MultiTaskUnicorn(UC_ARCH_ARM, UC_MODE_ARM + UC_MODE_MCLASS + UC_MODE_THUMB, 10)
 
     @cached_property
     def regs(self) -> QlRegisterManager:
@@ -132,7 +132,7 @@ def run(self, count=-1, end=None):
 
         utk = QlArchCORTEX_MThread(self.ql, self.effective_pc, end)
         self.uc.task_create(utk)
-        self.uc.tasks_start()
+        self.uc.tasks_start(count=count)
 
     def is_handler_mode(self):
         return self.regs.ipsr > 1
diff --git a/qiling/extensions/multitask.py b/qiling/extensions/multitask.py
index fc5f28340..45470e20d 100644
--- a/qiling/extensions/multitask.py
+++ b/qiling/extensions/multitask.py
@@ -1,6 +1,6 @@
 # Ziqiao Kong (mio@lazym.io)
 
-from typing import Dict
+from typing import Dict, List
 from unicorn import *
 from unicorn.x86_const import UC_X86_REG_EIP, UC_X86_REG_RIP
 from unicorn.arm64_const import UC_ARM64_REG_PC
@@ -112,7 +112,12 @@ def on_exit(self):
 # the same time.
 class MultiTaskUnicorn(Uc):
 
-    def __init__(self, arch, mode, interval=0.1):
+    def __init__(self, arch, mode, interval: int = 100):
+        """ Create a MultiTaskUnicorn object.
+
+            Interval: Sceduling interval in **ms**. The longger interval, the better
+            performance but less interrupts.
+        """
         super().__init__(arch, mode)
         self._interval = interval
         self._tasks = {} # type: Dict[int, UnicornTask]
@@ -122,6 +127,7 @@ def __init__(self, arch, mode, interval=0.1):
         self._running = False
         self._run_lock = threading.RLock()
         self._multitask_enabled = False
+        self._count = 0
 
     @property
     def current_thread(self):
@@ -151,20 +157,27 @@ def _emu_start_locked(self, begin: int, end: int, timeout: int, count: int):
     def _emu_start_utk_locked(self, utk: UnicornTask):
         return self._emu_start_locked(utk._begin, utk._end, 0, 0)
 
+    def _tmieout_main(self, timeout: int):
+
+        gevent.sleep(timeout / 1000)
+
+        self.tasks_stop()
+
     def _task_main(self, utk_id: int):
 
         utk = self._tasks[utk_id]
+        use_count = (self._count > 0)
 
         while True:
             # utk may be stopped before running once, check it.
             if utk._stop_request:
-                break
+                return # If we have to stop due to a tasks_stop, we preserve all threads so that we may resume.
             
             self._cur_utk_id = utk_id
 
             utk.on_start()
 
-            with gevent.Timeout(self._interval, False):
+            with gevent.Timeout(self._interval / 1000, False):
                 try:
                     pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
                     task = pool.spawn(self._emu_start_utk_locked, utk) # Run unicorn in a separate thread.
@@ -177,6 +190,13 @@ def _task_main(self, utk_id: int):
                     # Wait until we get the result.
                     ucerr = task.get()
 
+            if use_count:
+                self._count -= 1
+            
+                if self._count <= 0:
+                    self.tasks_stop()
+                    return
+
             if utk._stop_request:
                 utk.on_exit()
                 break
@@ -184,7 +204,7 @@ def _task_main(self, utk_id: int):
                 utk.on_interrupted(ucerr)
 
             if self._to_stop:
-                break
+                return
             
             # on_interrupted callback may have asked us to stop.
             if utk._stop_request:
@@ -273,22 +293,34 @@ def tasks_stop(self):
             if self._running:
                 super().emu_stop()
 
-    def tasks_start(self):
+    def tasks_start(self, count: int = 0, timeout: int = 0):
         """ This will start emulation until all tasks get done.
+
+            count: Stop after sceduling *count* times. <=0 disables this check.
+            timeout: Stop after *timeout* ms. <=0 disables this check.
         """
-        workset = {} # type: Dict[int, gevent.Greenlet]
+        workset = [] # type: List[gevent.Greenlet]
         self._to_stop = False
+        self._count = count
+
+        if self._count <= 0:
+            self._count = 0
 
         if self._multitask_enabled:
-            while len(self._tasks) != 0:
+
+            if timeout > 0:
+                workset.append(gevent.spawn(self._tmieout_main, timeout))
+    
+            while len(self._tasks) != 0 and not self._to_stop:
                 
-                new_workset = { k: v for k, v in workset.items() if not v.dead}
+                new_workset = [ v for v in workset if not v.dead]
 
                 for utk_id in self._tasks:
-                    new_workset[utk_id] = gevent.spawn(self._task_main, utk_id)
+                    new_workset.append(gevent.spawn(self._task_main, utk_id))
 
                 workset = new_workset
 
-                gevent.joinall(list(workset.values()), raise_error=True)
-            
-            self._multitask_enabled = False
\ No newline at end of file
+                gevent.joinall(workset, raise_error=True)
+
+            if len(self._tasks) == 0:
+                self._multitask_enabled = False
\ No newline at end of file

From 7d985422e66c73a9675b6d3e4fa0f9143fafc32e Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sun, 6 Mar 2022 00:33:20 +0100
Subject: [PATCH 08/11] Fix task stop

---
 qiling/extensions/multitask.py | 47 ++++++++++++++++++++++------------
 1 file changed, 30 insertions(+), 17 deletions(-)

diff --git a/qiling/extensions/multitask.py b/qiling/extensions/multitask.py
index 45470e20d..8e79cd788 100644
--- a/qiling/extensions/multitask.py
+++ b/qiling/extensions/multitask.py
@@ -30,41 +30,51 @@ def __init__(self, uc: Uc, begin: int, end: int, task_id = None):
         self._stop_request = False
         self._ctx = None
         self._task_id = None
+        self._arch = self._uc._arch
+        self._mode = self._uc._mode
 
     @property
     def pc(self):
-        arch = self._uc.ctl_get_arch()
-        mode = self._uc.ctl_get_mode()
+        """ Get current PC of the thread. This property should only be accessed when
+            the task is running.
+        """
+        raw_pc = self._raw_pc()
+        if (self._uc.reg_read(UC_ARM_REG_CPSR) & (1 << 5)):
+            return raw_pc | 1
+        else:
+            return raw_pc
 
+    def _raw_pc(self):
         # This extension is designed to be independent of Qiling, so let's
         # do this manually...
-        if arch == UC_ARCH_X86:
-            if (mode & UC_MODE_32) != 0:
+        if self._arch == UC_ARCH_X86:
+            if (self._mode & UC_MODE_32) != 0:
                 return self._uc.reg_read(UC_X86_REG_EIP)
-            elif (mode & UC_MODE_64) != 0:
+            elif (self._mode & UC_MODE_64) != 0:
                 return self._uc.reg_read(UC_X86_REG_RIP)
-        elif arch == UC_ARCH_MIPS:
+        elif self._arch == UC_ARCH_MIPS:
             return self._uc.reg_read(UC_MIPS_REG_PC)
-        elif arch == UC_ARCH_ARM:
-            pc = self._uc.reg_read(UC_ARM_REG_PC)
-            if (self._uc.reg_read(UC_ARM_REG_CPSR) & (1 << 5)):
-                return pc | 1
-            else:
-                return pc
-        elif arch == UC_ARCH_ARM64:
+        elif self._arch == UC_ARCH_ARM:
+            return self._uc.reg_read(UC_ARM_REG_PC)
+            
+        elif self._arch == UC_ARCH_ARM64:
             return self._uc.reg_read(UC_ARM64_REG_PC)
-        elif arch == UC_ARCH_PPC:
+        elif self._arch == UC_ARCH_PPC:
             return self._uc.reg_read(UC_PPC_REG_PC)
-        elif arch == UC_ARCH_M68K:
+        elif self._arch == UC_ARCH_M68K:
             return self._uc.reg_read(UC_M68K_REG_PC)
-        elif arch == UC_ARCH_SPARC:
+        elif self._arch == UC_ARCH_SPARC:
             return self._uc.reg_read(UC_SPARC_REG_PC)
-        elif arch == UC_ARCH_RISCV:
+        elif self._arch == UC_ARCH_RISCV:
             return self._uc.reg_read(UC_RISCV_REG_PC)
         
         # Really?
         return 0
 
+    def _reach_end(self):
+        # We may stop due to the scheduler asks us to, so check it manually.
+        return self._raw_pc() == self._end
+
     def save(self):
         """ This method is used to save the task context.
             Overwrite this method to implement specifc logic.
@@ -190,6 +200,9 @@ def _task_main(self, utk_id: int):
                     # Wait until we get the result.
                     ucerr = task.get()
 
+            if utk._reach_end():
+                utk._stop_request = True
+
             if use_count:
                 self._count -= 1
             

From 687a84d4377c72e1c121cd5a21cbd6e5ac849913 Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sun, 6 Mar 2022 15:20:11 +0100
Subject: [PATCH 09/11] No need to |1 for the *end* of thumb mode

---
 qiling/arch/cortex_m.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/qiling/arch/cortex_m.py b/qiling/arch/cortex_m.py
index 1b9a10bd7..fbb356c95 100644
--- a/qiling/arch/cortex_m.py
+++ b/qiling/arch/cortex_m.py
@@ -124,9 +124,6 @@ def stop(self):
     def run(self, count=-1, end=None):
         self.runable = True
 
-        if type(end) is int:
-            end |= 1        
-        
         if end is None:
             end = 0
 

From 2bf2a90c3dd56676bf17ff86421eccb8d6783aa6 Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Mon, 7 Mar 2022 00:44:08 +0100
Subject: [PATCH 10/11] Mimic unicorn interface better

---
 qiling/extensions/multitask.py | 117 +++++++++++++++++++++------------
 1 file changed, 75 insertions(+), 42 deletions(-)

diff --git a/qiling/extensions/multitask.py b/qiling/extensions/multitask.py
index 8e79cd788..a9da37bbf 100644
--- a/qiling/extensions/multitask.py
+++ b/qiling/extensions/multitask.py
@@ -1,4 +1,4 @@
-# Ziqiao Kong (mio@lazym.io)
+# Lazymio (mio@lazym.io)
 
 from typing import Dict, List
 from unicorn import *
@@ -56,7 +56,7 @@ def _raw_pc(self):
             return self._uc.reg_read(UC_MIPS_REG_PC)
         elif self._arch == UC_ARCH_ARM:
             return self._uc.reg_read(UC_ARM_REG_PC)
-            
+
         elif self._arch == UC_ARCH_ARM64:
             return self._uc.reg_read(UC_ARM64_REG_PC)
         elif self._arch == UC_ARCH_PPC:
@@ -67,12 +67,13 @@ def _raw_pc(self):
             return self._uc.reg_read(UC_SPARC_REG_PC)
         elif self._arch == UC_ARCH_RISCV:
             return self._uc.reg_read(UC_RISCV_REG_PC)
-        
+
         # Really?
         return 0
 
     def _reach_end(self):
         # We may stop due to the scheduler asks us to, so check it manually.
+        #print(f"{hex(self._raw_pc())} {hex(self._end)}")
         return self._raw_pc() == self._end
 
     def save(self):
@@ -80,7 +81,7 @@ def save(self):
             Overwrite this method to implement specifc logic.
         """
         return self._uc.context_save()
-    
+
     def restore(self, context):
         """ This method is used to restore the task context.
             Overwrite this method to implement specific logic.
@@ -105,6 +106,20 @@ def on_exit(self):
         """
         pass
 
+# This manages nested uc_emu_start calls and is designed as a friend
+# class of MultiTaskUnicorn.
+class NestedCounter:
+
+    def __init__(self, mtuc: "MultiTaskUnicorn"):
+        self._mtuc = mtuc
+    
+    def __enter__(self, *args, **kwargs):
+        self._mtuc._nested_started += 1
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self._mtuc._nested_started -= 1
+
 # This mimic a Unicorn object by maintaining the same interface.
 # If no task is registered, the behavior is exactly the same as
 # a normal unicorn.
@@ -124,7 +139,6 @@ class MultiTaskUnicorn(Uc):
 
     def __init__(self, arch, mode, interval: int = 100):
         """ Create a MultiTaskUnicorn object.
-
             Interval: Sceduling interval in **ms**. The longger interval, the better
             performance but less interrupts.
         """
@@ -138,6 +152,7 @@ def __init__(self, arch, mode, interval: int = 100):
         self._run_lock = threading.RLock()
         self._multitask_enabled = False
         self._count = 0
+        self._nested_started = 0
 
     @property
     def current_thread(self):
@@ -150,7 +165,7 @@ def running(self):
     def _next_task_id(self):
         while self._task_id_counter in self._tasks:
             self._task_id_counter += 1
-        
+
         return self._task_id_counter
 
     def _emu_start_locked(self, begin: int, end: int, timeout: int, count: int):
@@ -161,13 +176,13 @@ def _emu_start_locked(self, begin: int, end: int, timeout: int, count: int):
                 self._running = False
             except UcError as err:
                 return err.errno
-        
+
         return UC_ERR_OK
 
     def _emu_start_utk_locked(self, utk: UnicornTask):
         return self._emu_start_locked(utk._begin, utk._end, 0, 0)
 
-    def _tmieout_main(self, timeout: int):
+    def _timeout_main(self, timeout: int):
 
         gevent.sleep(timeout / 1000)
 
@@ -182,7 +197,7 @@ def _task_main(self, utk_id: int):
             # utk may be stopped before running once, check it.
             if utk._stop_request:
                 return # If we have to stop due to a tasks_stop, we preserve all threads so that we may resume.
-            
+
             self._cur_utk_id = utk_id
 
             utk.on_start()
@@ -195,8 +210,8 @@ def _task_main(self, utk_id: int):
                 finally:
                     if not task.done():
                         # Interrupted by timeout, in this case we call uc_emu_stop.
-                        self.emu_stop()
-                    
+                        super().emu_stop()
+
                     # Wait until we get the result.
                     ucerr = task.get()
 
@@ -205,7 +220,7 @@ def _task_main(self, utk_id: int):
 
             if use_count:
                 self._count -= 1
-            
+
                 if self._count <= 0:
                     self.tasks_stop()
                     return
@@ -218,9 +233,10 @@ def _task_main(self, utk_id: int):
 
             if self._to_stop:
                 return
-            
+
             # on_interrupted callback may have asked us to stop.
             if utk._stop_request:
+                utk.on_exit()
                 break
 
             # Give up control at once.
@@ -243,7 +259,6 @@ def tasks_restore(self, tasks_context: dict):
     def task_create(self, utk: UnicornTask):
         """ Create a unicorn task. utk should be a initialized UnicornTask object.
             If the task_id is not set, we generate one.
-
             utk: The task to add.
         """
         if not isinstance(utk, UnicornTask):
@@ -261,31 +276,41 @@ def task_exit(self, utk_id):
         """
         if utk_id not in self._tasks:
             return
-        
+
         if utk_id == self._cur_utk_id and self._running:
             self.emu_stop()
-        
+
         self._tasks[utk_id]._stop_request = True
 
     def emu_start(self, begin: int, end: int, timeout: int, count: int):
         """ Emulate an area of code just once. This overwrites the original emu_start interface and
             provides extra cares when multitask is enabled. If no task is registerd, this call bahaves
             like the original emu_start.
-
             NOTE: Calling this method may cause current greenlet to be switched out.
-
             begin, end, timeout, count: refer to Uc.emu_start
         """
-
         if self._multitask_enabled:
-            pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
-            task = pool.spawn(self._emu_start_locked, begin, end, timeout, count)
-            ucerr = task.get()
+            if self._nested_started > 0:
+                with NestedCounter(self):
+                    pool = gevent.get_hub().threadpool # type: gevent.threadpool.ThreadPool
+                    task = pool.spawn(self._emu_start_locked, begin, end, timeout, count)
+                    ucerr = task.get()
 
-            if ucerr != UC_ERR_OK:
-                raise UcError(ucerr)
-            
-            return ucerr
+                    if ucerr != UC_ERR_OK:
+                        raise UcError(ucerr)
+
+                    return ucerr
+            else:
+
+                # Assume users resume on the last thread (and that should be the case)
+                if self._cur_utk_id in self._tasks:
+                    self._tasks[self._cur_utk_id]._begin = begin
+                    self._tasks[self._cur_utk_id]._end = end
+                else:
+                    print(f"Warning: Can't found last thread we scheduled")
+
+                # This translation is not accurate, though.
+                self.tasks_start(count, timeout)
         else:
             return super().emu_start(begin, end, timeout, count)
 
@@ -295,6 +320,9 @@ def emu_stop(self):
         if self._multitask_enabled:
             if self._running:
                 super().emu_stop()
+            # Stop the world as original uc_emu_stop does
+            if self._nested_started == 1:
+                self.tasks_stop()
         else:
             super().emu_stop()
 
@@ -308,7 +336,6 @@ def tasks_stop(self):
 
     def tasks_start(self, count: int = 0, timeout: int = 0):
         """ This will start emulation until all tasks get done.
-
             count: Stop after sceduling *count* times. <=0 disables this check.
             timeout: Stop after *timeout* ms. <=0 disables this check.
         """
@@ -316,24 +343,30 @@ def tasks_start(self, count: int = 0, timeout: int = 0):
         self._to_stop = False
         self._count = count
 
-        if self._count <= 0:
-            self._count = 0
+        if self._nested_started != 0:
+            print("Warning: tasks_start is called inside an uc_emu_start!")
+            return
 
-        if self._multitask_enabled:
+        with NestedCounter(self):
 
-            if timeout > 0:
-                workset.append(gevent.spawn(self._tmieout_main, timeout))
-    
-            while len(self._tasks) != 0 and not self._to_stop:
-                
-                new_workset = [ v for v in workset if not v.dead]
+            if self._count <= 0:
+                self._count = 0
+
+            if self._multitask_enabled:
+
+                if timeout > 0:
+                    workset.append(gevent.spawn(self._timeout_main, timeout))
+
+                while len(self._tasks) != 0 and not self._to_stop:
+
+                    new_workset = [ v for v in workset if not v.dead]
 
-                for utk_id in self._tasks:
-                    new_workset.append(gevent.spawn(self._task_main, utk_id))
+                    for utk_id in self._tasks:
+                        new_workset.append(gevent.spawn(self._task_main, utk_id))
 
-                workset = new_workset
+                    workset = new_workset
 
-                gevent.joinall(workset, raise_error=True)
+                    gevent.joinall(workset, raise_error=True)
 
-            if len(self._tasks) == 0:
-                self._multitask_enabled = False
\ No newline at end of file
+                if len(self._tasks) == 0:
+                    self._multitask_enabled = False
\ No newline at end of file

From 4b4a48b592ed437f7848a482198cabeb64c40b45 Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sun, 24 Apr 2022 00:51:22 +0200
Subject: [PATCH 11/11] Dirty switch by ql.multithread

---
 qiling/arch/cortex_m.py | 39 ++++++++++++++++++++++++++++++++-------
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/qiling/arch/cortex_m.py b/qiling/arch/cortex_m.py
index fbb356c95..ad6a9e0a0 100644
--- a/qiling/arch/cortex_m.py
+++ b/qiling/arch/cortex_m.py
@@ -91,7 +91,14 @@ def __init__(self, ql: Qiling):
 
     @cached_property
     def uc(self):
-        return MultiTaskUnicorn(UC_ARCH_ARM, UC_MODE_ARM + UC_MODE_MCLASS + UC_MODE_THUMB, 10)
+        # XXX:
+        #   The `multithread` argument is reused here to decide if multitask is enabled, though, the exact meaning
+        #   here is not accurate enough. This switch could be safely removed once current MCU code is aware of the
+        #   multitask unicorn.
+        if self.ql.multithread:
+            return MultiTaskUnicorn(UC_ARCH_ARM, UC_MODE_ARM + UC_MODE_MCLASS + UC_MODE_THUMB, 10)
+        else:
+            return Uc(UC_ARCH_ARM, UC_MODE_ARM + UC_MODE_MCLASS + UC_MODE_THUMB)
 
     @cached_property
     def regs(self) -> QlRegisterManager:
@@ -117,6 +124,10 @@ def is_thumb(self) -> bool:
     def endian(self) -> QL_ENDIAN:
         return QL_ENDIAN.EL
 
+    def step(self):
+        self.ql.emu_start(self.effective_pc, 0, count=1)
+        self.ql.hw.step()
+
     def stop(self):
         self.ql.emu_stop()
         self.runable = False
@@ -124,12 +135,23 @@ def stop(self):
     def run(self, count=-1, end=None):
         self.runable = True
 
-        if end is None:
-            end = 0
+        if not self.ql.multithread:
+            if type(end) is int:
+                end |= 1        
 
-        utk = QlArchCORTEX_MThread(self.ql, self.effective_pc, end)
-        self.uc.task_create(utk)
-        self.uc.tasks_start(count=count)
+            while self.runable and count != 0:
+                if self.effective_pc == end:
+                    break
+
+                self.step()
+                count -= 1
+        else:
+            if end is None:
+                end = 0
+
+            utk = QlArchCORTEX_MThread(self.ql, self.effective_pc, end)
+            self.uc.task_create(utk)
+            self.uc.tasks_start(count=count)
 
     def is_handler_mode(self):
         return self.regs.ipsr > 1
@@ -198,4 +220,7 @@ def hard_interrupt_handler(self, ql, intno):
             self.regs.write('pc', entry)
             self.regs.write('lr', exc_return) 
 
-            self.uc.emu_start(self.effective_pc, 0, 0, 0)
+            if not self.ql.multithread:
+                self.ql.emu_start(self.effective_pc, 0, count=0xffffff)
+            else:
+                self.uc.emu_start(self.effective_pc, 0, 0, 0)