From fecfd5269632e86f86c6b31a9bc2ec22876f8934 Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Sat, 5 Feb 2022 23:29:44 +1100
Subject: [PATCH] tools/mpremote: Fix special handling of ctrl-D when host FS
 is mounted.

Changes are:
- decision to remount local filesystem on remote device is made only if
  "MPY: soft reboot" is seen in the output after sending a ctrl-D
- a nice message is printed to the user when the remount occurs
- soft reset during raw REPL is now handled correctly

Fixes issue #7731.

Signed-off-by: Damien George <damien@micropython.org>
---
 tools/mpremote/mpremote/main.py            |  4 +-
 tools/mpremote/mpremote/pyboardextended.py | 65 ++++++++++++++++------
 2 files changed, 49 insertions(+), 20 deletions(-)

diff --git a/tools/mpremote/mpremote/main.py b/tools/mpremote/mpremote/main.py
index 492d2d26d4..1073bec7fa 100644
--- a/tools/mpremote/mpremote/main.py
+++ b/tools/mpremote/mpremote/main.py
@@ -304,8 +304,8 @@ def do_repl_main_loop(pyb, console_in, console_out_write, *, code_to_inject, fil
             if c == b"\x1d":  # ctrl-], quit
                 break
             elif c == b"\x04":  # ctrl-D
-                # do a soft reset and reload the filesystem hook
-                pyb.soft_reset_with_mount(console_out_write)
+                # special handling needed for ctrl-D if filesystem is mounted
+                pyb.write_ctrl_d(console_out_write)
             elif c == b"\x0a" and code_to_inject is not None:  # ctrl-j, inject code
                 pyb.serial.write(code_to_inject)
             elif c == b"\x0b" and file_to_inject is not None:  # ctrl-k, inject script
diff --git a/tools/mpremote/mpremote/pyboardextended.py b/tools/mpremote/mpremote/pyboardextended.py
index 2ec1fbbee9..6817e7bbf4 100644
--- a/tools/mpremote/mpremote/pyboardextended.py
+++ b/tools/mpremote/mpremote/pyboardextended.py
@@ -621,37 +621,66 @@ class PyboardExtended(Pyboard):
         self.cmd = PyboardCommand(self.serial, fout, path)
         self.serial = SerialIntercept(self.serial, self.cmd)
 
-    def soft_reset_with_mount(self, out_callback):
+    def write_ctrl_d(self, out_callback):
         self.serial.write(b"\x04")
         if not self.mounted:
             return
 
-        # Clear flag while board reboots, it will be re-set once mounted.
-        self.mounted = False
+        # Read response from the device until it is quiet (with a timeout).
+        INITIAL_TIMEOUT = 0.5
+        QUIET_TIMEOUT = 0.2
+        FULL_TIMEOUT = 5
+        t_start = t_last_activity = time.monotonic()
+        data_all = b""
+        while True:
+            t = time.monotonic()
+            n = self.serial.inWaiting()
+            if n > 0:
+                data = self.serial.read(n)
+                out_callback(data)
+                data_all += data
+                t_last_activity = t
+            else:
+                if len(data_all) == 0:
+                    if t - t_start > INITIAL_TIMEOUT:
+                        return
+                else:
+                    if t - t_start > FULL_TIMEOUT:
+                        return
+                    if t - t_last_activity > QUIET_TIMEOUT:
+                        break
 
-        # Wait for a response to the soft-reset command.
-        for i in range(10):
-            if self.serial.inWaiting():
-                break
-            time.sleep(0.05)
+        # Check if a soft reset occurred.
+        if data_all.find(b"MPY: soft reboot") == -1:
+            return
+        if data_all.endswith(b">>> "):
+            in_friendly_repl = True
+        elif data_all.endswith(b">"):
+            in_friendly_repl = False
         else:
-            # Device didn't respond so it wasn't in a state to do a soft reset.
             return
 
-        out_callback(self.serial.read(1))
+        # Clear state while board remounts, it will be re-set once mounted.
+        self.mounted = False
         self.serial = self.serial.orig_serial
-        n = self.serial.inWaiting()
-        while n > 0:
-            buf = self.serial.read(n)
-            out_callback(buf)
-            time.sleep(0.2)
-            n = self.serial.inWaiting()
+
+        # Provide a message about the remount.
+        out_callback(bytes(f"\r\nRemount local directory {self.cmd.root} at /remote\r\n", "utf8"))
+
+        # Enter raw REPL and re-mount the remote filesystem.
         self.serial.write(b"\x01")
         self.exec_(fs_hook_code)
         self.exec_("__mount()")
         self.mounted = True
-        self.exit_raw_repl()
-        self.read_until(4, b">>> ")
+
+        # Exit raw REPL if needed, and wait for the friendly REPL prompt.
+        if in_friendly_repl:
+            self.exit_raw_repl()
+            prompt = b">>> "
+        else:
+            prompt = b">"
+        self.read_until(len(prompt), prompt)
+        out_callback(prompt)
         self.serial = SerialIntercept(self.serial, self.cmd)
 
     def umount_local(self):