supervisor: Add supervisor.set_usb_identification() function

s-ol · s-ol · commit bee9957efd9f · 2022-05-04T09:54:24.000+02:00
diff --git a/shared-bindings/supervisor/__init__.c b/shared-bindings/supervisor/__init__.c
@@ -38,6 +38,7 @@
 #include "supervisor/shared/traceback.h"
 #include "supervisor/shared/translate.h"
 #include "supervisor/shared/workflow.h"
+#include "supervisor/usb.h"
 
 #include "shared-bindings/microcontroller/__init__.h"
 #include "shared-bindings/supervisor/__init__.h"
@@ -311,6 +312,62 @@ STATIC mp_obj_t supervisor_reset_terminal(mp_obj_t x_pixels, mp_obj_t y_pixels)
 }
 MP_DEFINE_CONST_FUN_OBJ_2(supervisor_reset_terminal_obj, supervisor_reset_terminal);
 
+//| def set_usb_identification(manufacturer: Optional[str] = None, product: Optional[str] = None, vid: int = -1, pid: int = -1) -> None:
+//|     """Override identification constants in the USB Device Descriptor.
+//|
+//|     If passed, `manufacturer` and `product` must be ASCII strings (or buffers) of at most 127
+//|     characters. Any omitted arguments will be left at their default values.
+//|
+//|     This method must be called in boot.py to have any effect."""
+//|     ...
+//|
+STATIC mp_obj_t supervisor_set_usb_identification(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
+    static const mp_arg_t allowed_args[] = {
+        { MP_QSTR_manufacturer, MP_ARG_OBJ, {.u_rom_obj = mp_const_none} },
+        { MP_QSTR_product, MP_ARG_OBJ, {.u_rom_obj = mp_const_none} },
+        { MP_QSTR_vid, MP_ARG_INT, {.u_int = -1} },
+        { MP_QSTR_pid, MP_ARG_INT, {.u_int = -1} },
+    };
+    struct {
+        mp_arg_val_t manufacturer;
+        mp_arg_val_t product;
+        mp_arg_val_t vid;
+        mp_arg_val_t pid;
+    } args;
+    mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, (mp_arg_val_t *)&args);
+
+    if (!usb_identification_allocation) {
+        usb_identification_allocation = allocate_memory(sizeof(usb_identification_t), false, true);
+    }
+    usb_identification_t *identification = (usb_identification_t *)usb_identification_allocation->ptr;
+
+    mp_arg_validate_int_range(args.vid.u_int, -1, (1 << 16) - 1, MP_QSTR_vid);
+    mp_arg_validate_int_range(args.pid.u_int, -1, (1 << 16) - 1, MP_QSTR_pid);
+
+    identification->vid = args.vid.u_int > -1 ? args.vid.u_int : USB_VID;
+    identification->pid = args.pid.u_int > -1 ? args.pid.u_int : USB_PID;
+
+    mp_buffer_info_t info;
+    if (args.manufacturer.u_obj != mp_const_none) {
+        mp_get_buffer_raise(args.manufacturer.u_obj, &info, MP_BUFFER_READ);
+        mp_arg_validate_length_range(info.len, 0, 127, MP_QSTR_manufacturer);
+        memcpy(identification->manufacturer_name, info.buf, info.len);
+    } else {
+        memcpy(identification->manufacturer_name, USB_MANUFACTURER, sizeof(USB_MANUFACTURER));
+    }
+
+    if (args.product.u_obj != mp_const_none) {
+        mp_get_buffer_raise(args.product.u_obj, &info, MP_BUFFER_READ);
+        mp_arg_validate_length_range(info.len, 0, 127, MP_QSTR_product);
+        memcpy(identification->product_name, info.buf, info.len);
+    } else {
+        memcpy(identification->product_name, USB_MANUFACTURER, sizeof(USB_MANUFACTURER));
+    }
+
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_KW(supervisor_set_usb_identification_obj, 0, supervisor_set_usb_identification);
+
 STATIC const mp_rom_map_elem_t supervisor_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_supervisor) },
     { MP_ROM_QSTR(MP_QSTR_enable_autoreload),  MP_ROM_PTR(&supervisor_enable_autoreload_obj) },
@@ -325,6 +382,7 @@ STATIC const mp_rom_map_elem_t supervisor_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR_get_previous_traceback),  MP_ROM_PTR(&supervisor_get_previous_traceback_obj) },
     { MP_ROM_QSTR(MP_QSTR_disable_ble_workflow),  MP_ROM_PTR(&supervisor_disable_ble_workflow_obj) },
     { MP_ROM_QSTR(MP_QSTR_reset_terminal),  MP_ROM_PTR(&supervisor_reset_terminal_obj) },
+    { MP_ROM_QSTR(MP_QSTR_set_usb_identification),  MP_ROM_PTR(&supervisor_set_usb_identification_obj) },
 };
 
 STATIC MP_DEFINE_CONST_DICT(supervisor_module_globals, supervisor_module_globals_table);
diff --git a/supervisor/shared/usb/usb.c b/supervisor/shared/usb/usb.c
@@ -118,9 +118,11 @@ void usb_set_defaults(void) {
     #endif
 };
 
+supervisor_allocation *usb_identification_allocation;
+
 // Some dynamic USB data must be saved after boot.py. How much is needed?
 size_t usb_boot_py_data_size(void) {
-    size_t size = 0;
+    size_t size = sizeof(usb_identification_t);
 
     #if CIRCUITPY_USB_HID
     size += usb_hid_report_descriptor_length();
@@ -131,19 +133,41 @@ size_t usb_boot_py_data_size(void) {
 
 // Fill in the data to save.
 void usb_get_boot_py_data(uint8_t *temp_storage, size_t temp_storage_size) {
+    if (usb_identification_allocation) {
+        memcpy(temp_storage, usb_identification_allocation->ptr, sizeof(usb_identification_t));
+        free_memory(usb_identification_allocation);
+    } else {
+        usb_identification_t defaults = {
+            .vid = USB_VID,
+            .pid = USB_PID,
+            .manufacturer_name = USB_MANUFACTURER,
+            .product_name = USB_PRODUCT,
+        };
+        memcpy(temp_storage, &defaults, sizeof(defaults));
+    }
+
+    temp_storage += sizeof(usb_identification_t);
+    temp_storage_size -= sizeof(usb_identification_t);
+
     #if CIRCUITPY_USB_HID
     usb_hid_build_report_descriptor(temp_storage, temp_storage_size);
     #endif
 }
 
 // After VM is gone, save data into non-heap storage (storage_allocations).
 void usb_return_boot_py_data(uint8_t *temp_storage, size_t temp_storage_size) {
+    usb_identification_t identification;
+    memcpy(&identification, temp_storage, sizeof(usb_identification_t));
+
+    temp_storage += sizeof(usb_identification_t);
+    temp_storage_size -= sizeof(usb_identification_t);
+
     #if CIRCUITPY_USB_HID
     usb_hid_save_report_descriptor(temp_storage, temp_storage_size);
     #endif
 
     // Now we can also build the rest of the descriptors and place them in storage_allocations.
-    usb_build_descriptors();
+    usb_build_descriptors(&identification);
 }
 
 // Call this when ready to run code.py or a REPL, and a VM has been started.
diff --git a/supervisor/shared/usb/usb_desc.c b/supervisor/shared/usb/usb_desc.c
@@ -68,9 +68,6 @@ static supervisor_allocation *device_descriptor_allocation;
 static supervisor_allocation *configuration_descriptor_allocation;
 static supervisor_allocation *string_descriptors_allocation;
 
-static const char manufacturer_name[] = USB_MANUFACTURER;
-static const char product_name[] = USB_PRODUCT;
-
 // Serial number string is UID length * 2 (2 nibbles per byte) + 1 byte for null termination.
 static char serial_number_hex_string[COMMON_HAL_MCU_PROCESSOR_UID_LENGTH * 2 + 1];
 
@@ -113,23 +110,23 @@ static const uint8_t configuration_descriptor_template[] = {
     0x32,        // 8 bMaxPower 100mA
 };
 
-static void usb_build_device_descriptor(uint16_t vid, uint16_t pid) {
+static void usb_build_device_descriptor(const usb_identification_t *identification) {
     device_descriptor_allocation =
         allocate_memory(align32_size(sizeof(device_descriptor_template)),
             /*high_address*/ false, /*movable*/ false);
     uint8_t *device_descriptor = (uint8_t *)device_descriptor_allocation->ptr;
     memcpy(device_descriptor, device_descriptor_template, sizeof(device_descriptor_template));
 
-    device_descriptor[DEVICE_VID_LO_INDEX] = vid & 0xFF;
-    device_descriptor[DEVICE_VID_HI_INDEX] = vid >> 8;
-    device_descriptor[DEVICE_PID_LO_INDEX] = pid & 0xFF;
-    device_descriptor[DEVICE_PID_HI_INDEX] = pid >> 8;
+    device_descriptor[DEVICE_VID_LO_INDEX] = identification->vid & 0xFF;
+    device_descriptor[DEVICE_VID_HI_INDEX] = identification->vid >> 8;
+    device_descriptor[DEVICE_PID_LO_INDEX] = identification->pid & 0xFF;
+    device_descriptor[DEVICE_PID_HI_INDEX] = identification->pid >> 8;
 
-    usb_add_interface_string(current_interface_string, manufacturer_name);
+    usb_add_interface_string(current_interface_string, identification->manufacturer_name);
     device_descriptor[DEVICE_MANUFACTURER_STRING_INDEX] = current_interface_string;
     current_interface_string++;
 
-    usb_add_interface_string(current_interface_string, product_name);
+    usb_add_interface_string(current_interface_string, identification->product_name);
     device_descriptor[DEVICE_PRODUCT_STRING_INDEX] = current_interface_string;
     current_interface_string++;
 
@@ -319,7 +316,7 @@ static void usb_build_interface_string_table(void) {
 
 // After boot.py runs, the USB devices to be used have been chosen, and the descriptors can be set up.
 // This is called after the VM is finished, because it uses storage_allocations.
-void usb_build_descriptors(void) {
+void usb_build_descriptors(const usb_identification_t *identification) {
     uint8_t raw_id[COMMON_HAL_MCU_PROCESSOR_UID_LENGTH];
     common_hal_mcu_processor_get_uid(raw_id);
 
@@ -336,7 +333,7 @@ void usb_build_descriptors(void) {
     current_interface_string = 1;
     collected_interface_strings_length = 0;
 
-    usb_build_device_descriptor(USB_VID, USB_PID);
+    usb_build_device_descriptor(identification);
     usb_build_configuration_descriptor();
     usb_build_interface_string_table();
 }
diff --git a/supervisor/usb.h b/supervisor/usb.h
@@ -31,6 +31,8 @@
 #include <stddef.h>
 #include <stdint.h>
 
+#include "supervisor/memory.h"
+
 // Ports must call this as frequently as they can in order to keep the USB
 // connection alive and responsive.  Normally this is called from background
 // tasks after the USB IRQ handler is executed, but in specific circumstances
@@ -58,10 +60,19 @@ typedef struct {
     size_t num_out_endpoints;
 } descriptor_counts_t;
 
+typedef struct {
+    uint16_t vid;
+    uint16_t pid;
+    char manufacturer_name[128];
+    char product_name[128];
+} usb_identification_t;
+
+extern supervisor_allocation *usb_identification_allocation;
+
 // Shared implementation.
 bool usb_enabled(void);
 void usb_add_interface_string(uint8_t interface_string_index, const char str[]);
-void usb_build_descriptors(void);
+void usb_build_descriptors(const usb_identification_t *identification);
 void usb_disconnect(void);
 void usb_init(void);
 void usb_set_defaults(void);
