From 9e292b85023981cfeb0913506d4bbbb7ee53059e Mon Sep 17 00:00:00 2001
From: Asger Hautop Drewsen <asgerdrewsen@gmail.com>
Date: Thu, 4 Jun 2015 18:55:10 +0200
Subject: [PATCH] Fix some hanging issues with the repl

---
 src/frida/application.py | 10 +++++++---
 src/frida/repl.py        | 41 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/src/frida/application.py b/src/frida/application.py
index 39390d7..9a631a2 100755
--- a/src/frida/application.py
+++ b/src/frida/application.py
@@ -49,7 +49,7 @@ def await_enter(reactor):
         pass
 
 class ConsoleApplication(object):
-    def __init__(self, run_until_return=await_enter):
+    def __init__(self, run_until_return=await_enter, on_stop=None):
         colorama.init()
 
         parser = OptionParser(usage=self._usage(), version=frida.__version__)
@@ -87,7 +87,7 @@ class ConsoleApplication(object):
         self._schedule_on_session_detached = lambda: self._reactor.schedule(self._on_session_detached)
         self._started = False
         self._resumed = False
-        self._reactor = Reactor(run_until_return)
+        self._reactor = Reactor(run_until_return, on_stop)
         self._exit_status = None
         self._status_updated = False
 
@@ -250,9 +250,10 @@ def expand_target(target):
 
 
 class Reactor(object):
-    def __init__(self, run_until_return):
+    def __init__(self, run_until_return, on_stop):
         self._running = False
         self._run_until_return = run_until_return
+        self._on_stop = on_stop
         self._pending = collections.deque([])
         self._lock = threading.Lock()
         self._cond = threading.Condition(self._lock)
@@ -290,6 +291,9 @@ class Reactor(object):
                     self._cond.wait(timeout)
                 running = self._running
 
+        if self._on_stop:
+            self._on_stop()
+
     def stop(self):
         with self._lock:
             self._running = False
diff --git a/src/frida/repl.py b/src/frida/repl.py
index 8cfeb1d..33a30cd 100644
--- a/src/frida/repl.py
+++ b/src/frida/repl.py
@@ -8,9 +8,10 @@ def main():
     import sys
     import threading
     import os
-    from prompt_toolkit.shortcuts import get_input
+    from prompt_toolkit.shortcuts import create_default_application, create_default_output, create_eventloop
     from prompt_toolkit.history import FileHistory
     from prompt_toolkit.completion import Completion, Completer
+    from prompt_toolkit.interface import CommandLineInterface
     from pygments.lexers import JavascriptLexer
     from pygments.token import Token
 
@@ -24,8 +25,9 @@ def main():
             self._completor_locals = []
             self._history = FileHistory(os.path.join(os.path.expanduser('~'), '.frida_history'))
             self._completer = FridaCompleter(self)
+            self._cli = None
 
-            super(REPLApplication, self).__init__(self._process_input)
+            super(REPLApplication, self).__init__(self._process_input, self._on_stop)
 
         def _add_options(self, parser):
             parser.add_option("-l", "--load", help="load SCRIPT", metavar="SCRIPT",
@@ -54,6 +56,15 @@ def main():
                 sys.stdout.write("\033[A")
             self._ready.set()
 
+        def _on_stop(self):
+            def set_return():
+                raise EOFError()
+
+            try:
+                self._cli.eventloop.call_from_executor(set_return)
+            except Exception:
+                pass
+
         def _stop(self):
             self._unload_script()
 
@@ -78,15 +89,37 @@ def main():
 
         def _process_input(self, reactor):
             self._print_startup_message()
-            self._ready.wait()
+            while self._ready.wait(0.5) != True:
+                if not reactor._running:
+                    return
 
             while True:
                 expression = ""
                 line = ""
                 while len(expression) == 0 or line.endswith("\\"):
+                    if not reactor._running:
+                        return
                     try:
                         prompt = "[%s]" % self._prompt_string + "-> " if len(expression) == 0 else "... "
-                        line = get_input(prompt, history=self._history, lexer=JavascriptLexer, completer=self._completer)
+
+                        # We create the prompt manually instead of using get_input,
+                        # so we can use the cli in the _on_stop method
+                        eventloop = create_eventloop()
+
+                        self._cli = CommandLineInterface(
+                            application=create_default_application(prompt, history=self._history, completer=self._completer, lexer=JavascriptLexer),
+                            eventloop=eventloop,
+                            output=create_default_output())
+
+                        try:
+                            line = None
+
+                            document = self._cli.run()
+
+                            if document:
+                                line = document.text
+                        finally:
+                            eventloop.close()
                     except EOFError:
                         # An extra newline after EOF to exit the REPL cleanly
                         print("\nThank you for using Frida!")