allow storage.disable_usb_drive() and .enable_usb_drive() after boot.py

shared-bindings/storage/__init__.c

@@ -93,7 +93,7 @@ MP_DEFINE_CONST_FUN_OBJ_1(storage_umount_obj, storage_umount);
 //| ) -> None:
 //|     """Remounts the given path with new parameters.
 //|
-//|     This can always be done from boot.py. After boot, it can only be done when the host computer
+//|     This can always be done from ``boot.py``. After boot, it can only be done when the host computer
 //|     doesn't have write access and CircuitPython isn't currently writing to the filesystem. An
 //|     exception will be raised if this is the case. Some host OSes allow you to eject a drive which
 //|     will allow for remounting.
@@ -188,8 +188,10 @@ MP_DEFINE_CONST_FUN_OBJ_KW(storage_erase_filesystem_obj, 0, storage_erase_filesy
 
 //| def disable_usb_drive() -> None:
 //|     """Disable presenting ``CIRCUITPY`` as a USB mass storage device.
-//|     By default, the device is enabled and ``CIRCUITPY`` is visible.
-//|     Can be called in ``boot.py``, before USB is connected."""
+//|     By default, the device is enabled and ``CIRCUITPY`` is visible, if USB is available.
+//|     Must called in ``boot.py``, before USB is connected.
+//      If you want to disable the USB drive after `boot.py` has run, see `unsafe_disable_usb_drive()`.
+//|     """
 //|     ...
 //|
 //|
@@ -205,15 +207,59 @@ static mp_obj_t storage_disable_usb_drive(void) {
 }
 MP_DEFINE_CONST_FUN_OBJ_0(storage_disable_usb_drive_obj, storage_disable_usb_drive);
 
+//| def unsafe_disable_usb_drive() -> None:
+//|     """Disable presenting ``CIRCUITPY`` as a USB mass storage device.
+//|     By default, the device is enabled and ``CIRCUITPY`` is visible.
+//|     Unlike `disable_usb_drive()`, `unsafe_disable_usb_drive()` can be called
+//|     after ``code.py`` starts or from the REPL, after USB has started.
+//|
+//|     When `unsafe_disable_usb_drive()` after USB has started,
+//|     the ``CIRCUITPY`` USB drive logical unit (LUN) will report as "not ready",
+//|     causing the host to unmount it.
+//|     The drive can be made ready and available again by calling `enable_usb_drive()`.
+//|     When `disable_usb_drive` is called after ``code.py`` starts or in the REPL,
+//|     the call will delay 2.5 seconds before returning,
+//|     so that host has time to detect that the drive is not ready.
+//|     The host polls the device approximately every one or two seconds.
+//|
+//|     Note that if ``unsafe_disable_usb_drive()`` is called when the host is actively writing CIRCUITPY,
+//|     filesystem corruption could occur. Be careful to call it when the host is quiescent.
+//|
+//|     When the USB drive is disabled, CIRCUITPY becomes read/write, and can be written
+//|     from user code or the REPL. This is easier than arranging for a `remount()` in ``boot.py``.
+//|     Code editors and file uploaders can use this feature to write files via the REPL.
+//|
+//|     If `unsafe_disable_usb_drive()` is called in ``boot.py``, it is identical to calling
+//|     `disable_usb_drive()`.
+//|     """
+//|     ...
+//|
+//|
+static mp_obj_t storage_unsafe_disable_usb_drive(void) {
+    #if CIRCUITPY_USB_DEVICE && CIRCUITPY_USB_MSC
+    if (!common_hal_storage_unsafe_disable_usb_drive()) {
+    #else
+    if (true) {
+        #endif
+        mp_raise_RuntimeError(MP_ERROR_TEXT("Cannot change USB devices now"));
+    }
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_0(storage_unsafe_disable_usb_drive_obj, storage_unsafe_disable_usb_drive);
+
 //| def enable_usb_drive() -> None:
-//|     """Enabled presenting ``CIRCUITPY`` as a USB mass storage device.
+//|     """Enable presenting ``CIRCUITPY`` as a USB mass storage device.
 //|     By default, the device is enabled and ``CIRCUITPY`` is visible,
-//|     so you do not normally need to call this function.
-//|     Can be called in ``boot.py``, before USB is connected.
+//|     so you do not normally need to call this function in ``boot.py``.
+//|
+//|     If you call `enable_usb_drive()` after ``code.py`` starts or in the REPL,
+//|     you can reverse the effect of a previous `unsafe_disable_usb_drive()`.
+//|     The CIRCUITPY drive will reappear to the host, and become read-only again
+//|     if it was previously read-only.
 //|
-//|     If you enable too many devices at once, you will run out of USB endpoints.
+//|     If you enable too many USB devices at once, you will run out of USB endpoints.
 //|     The number of available endpoints varies by microcontroller.
-//|     CircuitPython will go into safe mode after running boot.py to inform you if
+//|     CircuitPython will go into safe mode after running ``boot.py`` to inform you if
 //|     not enough endpoints are available.
 //|     """
 //|     ...
@@ -234,13 +280,14 @@ MP_DEFINE_CONST_FUN_OBJ_0(storage_enable_usb_drive_obj, storage_enable_usb_drive
 static const mp_rom_map_elem_t storage_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_storage) },
 
-    { MP_ROM_QSTR(MP_QSTR_mount),             MP_ROM_PTR(&storage_mount_obj) },
-    { MP_ROM_QSTR(MP_QSTR_umount),            MP_ROM_PTR(&storage_umount_obj) },
-    { MP_ROM_QSTR(MP_QSTR_remount),           MP_ROM_PTR(&storage_remount_obj) },
-    { MP_ROM_QSTR(MP_QSTR_getmount),          MP_ROM_PTR(&storage_getmount_obj) },
-    { MP_ROM_QSTR(MP_QSTR_erase_filesystem),  MP_ROM_PTR(&storage_erase_filesystem_obj) },
-    { MP_ROM_QSTR(MP_QSTR_disable_usb_drive), MP_ROM_PTR(&storage_disable_usb_drive_obj) },
-    { MP_ROM_QSTR(MP_QSTR_enable_usb_drive),  MP_ROM_PTR(&storage_enable_usb_drive_obj) },
+    { MP_ROM_QSTR(MP_QSTR_mount),                    MP_ROM_PTR(&storage_mount_obj) },
+    { MP_ROM_QSTR(MP_QSTR_umount),                   MP_ROM_PTR(&storage_umount_obj) },
+    { MP_ROM_QSTR(MP_QSTR_remount),                  MP_ROM_PTR(&storage_remount_obj) },
+    { MP_ROM_QSTR(MP_QSTR_getmount),                 MP_ROM_PTR(&storage_getmount_obj) },
+    { MP_ROM_QSTR(MP_QSTR_erase_filesystem),         MP_ROM_PTR(&storage_erase_filesystem_obj) },
+    { MP_ROM_QSTR(MP_QSTR_disable_usb_drive),        MP_ROM_PTR(&storage_disable_usb_drive_obj) },
+    { MP_ROM_QSTR(MP_QSTR_enable_usb_drive),         MP_ROM_PTR(&storage_enable_usb_drive_obj) },
+    { MP_ROM_QSTR(MP_QSTR_unsafe_disable_usb_drive), MP_ROM_PTR(&storage_unsafe_disable_usb_drive_obj) },
 
 //| class VfsFat:
 //|     def __init__(self, block_device: BlockDevice) -> None:

shared-bindings/storage/__init__.h

@@ -19,4 +19,5 @@ mp_obj_t common_hal_storage_getmount(const char *path);
 MP_NORETURN void common_hal_storage_erase_filesystem(bool extended);
 
 bool common_hal_storage_disable_usb_drive(void);
+bool common_hal_storage_unsafe_disable_usb_drive(void);
 bool common_hal_storage_enable_usb_drive(void);

shared-module/storage/__init__.c

@@ -27,7 +27,7 @@
 #include "tusb.h"
 
 // Is the MSC device enabled?
-bool storage_usb_is_enabled;
+static volatile bool storage_usb_is_enabled;
 
 void storage_usb_set_defaults(void) {
     storage_usb_is_enabled = CIRCUITPY_USB_MSC_ENABLED_DEFAULT;
@@ -38,16 +38,29 @@ bool storage_usb_enabled(void) {
 }
 
 static bool usb_drive_set_enabled(bool enabled) {
-    // We can't change the descriptors once we're connected.
+    // We can't change the descriptors once we're connected, but we can make the LUN be ready or not ready.
+    storage_usb_is_enabled = enabled;
     if (tud_connected()) {
-        return false;
+        // The TEST UNIT READY callback in usb_msc_flash.c checks the value of storage_usb_is_enabled.
+        // If it's false, TEST UNIT READY will report "not ready"
+        // Linux and macOS send a TEST UNIT READY poll about every 1.1 seconds or faster.
+        // Windows polls every 2.1 seconds or so.
+        // So wait long enough for host to send a TEST UNIT READY and receive a reply.
+        mp_hal_delay_ms(2500);
     }
     filesystem_set_internal_writable_by_usb(enabled);
-    storage_usb_is_enabled = enabled;
     return true;
 }
 
 bool common_hal_storage_disable_usb_drive(void) {
+    if (tud_connected()) {
+        // Complain if already connected. Use `storage.unsafe_disable_usb_drive()` in that case.
+        return false;
+    }
+    return usb_drive_set_enabled(false);
+}
+
+bool common_hal_storage_unsafe_disable_usb_drive(void) {
     return usb_drive_set_enabled(false);
 }
 
@@ -59,6 +72,10 @@ bool common_hal_storage_disable_usb_drive(void) {
     return false;
 }
 
+bool common_hal_storage_unsafe_disable_usb_drive(void) {
+    return false;
+}
+
 bool common_hal_storage_enable_usb_drive(void) {
     return false;
 }

supervisor/shared/usb/usb_msc_flash.c

@@ -404,7 +404,7 @@ bool tud_msc_test_unit_ready_cb(uint8_t lun) {
         return false;
     }
 
-    if (ejected[lun] || eject_once[lun]
+    if (ejected[lun] || eject_once[lun] || (lun == 0 && !storage_usb_enabled())
         #ifdef SDCARD_LUN
         || (lun == SDCARD_LUN && !sdcard_usb_enabled())
         #endif