From 561fdfb279c45f7a9ff69a2b2e71ac272cb6aaef Mon Sep 17 00:00:00 2001
From: Scott Shawcroft <scott@tannewt.org>
Date: Tue, 17 Dec 2019 19:01:03 -0800
Subject: [PATCH] Improve USB eject by resetting on replug

---
 supervisor/shared/usb/usb.c           |  2 ++
 supervisor/shared/usb/usb_msc_flash.c | 41 ++++++++++++++++++++++-----
 supervisor/usb.h                      |  4 +++
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/supervisor/shared/usb/usb.c b/supervisor/shared/usb/usb.c
index c1d70c5b0f..a0178528d8 100644
--- a/supervisor/shared/usb/usb.c
+++ b/supervisor/shared/usb/usb.c
@@ -87,10 +87,12 @@ void usb_background(void) {
 
 // Invoked when device is mounted
 void tud_mount_cb(void) {
+    usb_msc_mount();
 }
 
 // Invoked when device is unmounted
 void tud_umount_cb(void) {
+    usb_msc_umount();
 }
 
 // Invoked when usb bus is suspended
diff --git a/supervisor/shared/usb/usb_msc_flash.c b/supervisor/shared/usb/usb_msc_flash.c
index 205cb7b500..0376cfc1df 100644
--- a/supervisor/shared/usb/usb_msc_flash.c
+++ b/supervisor/shared/usb/usb_msc_flash.c
@@ -41,6 +41,18 @@
 
 static bool ejected[1];
 
+void usb_msc_mount(void) {
+    // Reset the ejection tracking every time we're plugged into USB. This allows for us to battery
+    // power the device, eject, unplug and plug it back in to get the drive.
+    for (uint8_t i = 0; i < sizeof(ejected); i++) {
+        ejected[i] = false;
+    }
+}
+
+void usb_msc_umount(void) {
+
+}
+
 // The root FS is always at the end of the list.
 static fs_user_mount_t* get_vfs(int lun) {
     // TODO(tannewt): Return the mount which matches the lun where 0 is the end
@@ -198,19 +210,34 @@ bool tud_msc_test_unit_ready_cb(uint8_t lun) {
 // - Start = 0 : stopped power mode, if load_eject = 1 : unload disk storage
 // - Start = 1 : active mode, if load_eject = 1 : load disk storage
 bool tud_msc_start_stop_cb(uint8_t lun, uint8_t power_condition, bool start, bool load_eject) {
+    if (lun > 1) {
+        return false;
+    }
+    fs_user_mount_t* current_mount = get_vfs(lun);
+    if (current_mount == NULL) {
+        return false;
+    }
     if (load_eject) {
-        if (lun > 1) {
-            return false;
-        } else {
-            fs_user_mount_t* current_mount = get_vfs(lun);
-            if (current_mount == NULL) {
-                return false;
-            }
+        if (!start) {
+            // Eject but first flush.
             if (disk_ioctl(current_mount, CTRL_SYNC, NULL) != RES_OK) {
                 return false;
             } else {
                 ejected[lun] = true;
             }
+        } else {
+            // We can only load if it hasn't been ejected.
+            return !ejected[lun];
+        }
+    } else {
+        if (!start) {
+            // Stop the unit but don't eject.
+            if (disk_ioctl(current_mount, CTRL_SYNC, NULL) != RES_OK) {
+                return false;
+            }
+        } else {
+            // Start the unit, but only if not ejected.
+            return !ejected[lun];
         }
     }
 
diff --git a/supervisor/usb.h b/supervisor/usb.h
index c87540d408..0cc361619d 100644
--- a/supervisor/usb.h
+++ b/supervisor/usb.h
@@ -41,4 +41,8 @@ void init_usb_hardware(void);
 bool usb_enabled(void);
 void usb_init(void);
 
+// Propagate plug/unplug events to the MSC logic.
+void usb_msc_mount(void);
+void usb_msc_umount(void);
+
 #endif // MICROPY_INCLUDED_SUPERVISOR_USB_H