Improve USB eject by resetting on replug

3 files changed, 40 insertions, 7 deletions

diff --git a/supervisor/shared/usb/usb.c b/supervisor/shared/usb/usb.c
index c1d70c5b0..a0178528d 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 205cb7b50..0376cfc1d 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 c87540d40..0cc361619 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