4 files changed, 30 insertions, 40 deletions

diff --git a/embassy-usb-dfu/src/application.rs b/embassy-usb-dfu/src/application.rs
index 2646d100d..4b7b72073 100644
--- a/embassy-usb-dfu/src/application.rs
+++ b/embassy-usb-dfu/src/application.rs
@@ -2,7 +2,7 @@ use embassy_boot::BlockingFirmwareState;
 use embassy_time::{Duration, Instant};
 use embassy_usb::control::{InResponse, OutResponse, Recipient, RequestType};
 use embassy_usb::driver::Driver;
-use embassy_usb::{msos, Builder, Handler};
+use embassy_usb::{Builder, FunctionBuilder, Handler};
 use embedded_storage::nor_flash::NorFlash;
 
 use crate::consts::{
@@ -130,22 +130,13 @@ pub fn usb_dfu<'d, D: Driver<'d>, MARK: DfuMarker, RST: Reset>(
     builder: &mut Builder<'d, D>,
     handler: &'d mut Control<MARK, RST>,
     timeout: Duration,
-    winusb_guids: Option<&'d [&str]>,
+    func_modifier: impl Fn(&mut FunctionBuilder<'_, 'd, D>),
 ) {
     let mut func = builder.function(0x00, 0x00, 0x00);
-    if let Some(winusb_guids) = winusb_guids {
-        // We add MSOS headers so that the device automatically gets assigned the WinUSB driver on Windows.
-        // Otherwise users need to do this manually using a tool like Zadig.
-        //
-        // Adding them here on the function level appears to only be needed for compositive devices.
-        // In addition to being on the function level, they should also be added to the device level.
-        //
-        func.msos_feature(msos::CompatibleIdFeatureDescriptor::new("WINUSB", ""));
-        func.msos_feature(msos::RegistryPropertyFeatureDescriptor::new(
-            "DeviceInterfaceGUIDs",
-            msos::PropertyData::RegMultiSz(winusb_guids),
-        ));
-    }
+
+    // Here we give users the opportunity to add their own function level MSOS headers for instance.
+    // This is useful when DFU functionality is part of a composite USB device.
+    func_modifier(&mut func);
 
     let mut iface = func.interface();
     let mut alt = iface.alt_setting(USB_CLASS_APPN_SPEC, APPN_SPEC_SUBCLASS_DFU, DFU_PROTOCOL_RT, None);
diff --git a/embassy-usb-dfu/src/dfu.rs b/embassy-usb-dfu/src/dfu.rs
index 43a35637d..0f39d906b 100644
--- a/embassy-usb-dfu/src/dfu.rs
+++ b/embassy-usb-dfu/src/dfu.rs
@@ -1,7 +1,7 @@
 use embassy_boot::{AlignedBuffer, BlockingFirmwareUpdater, FirmwareUpdaterError};
 use embassy_usb::control::{InResponse, OutResponse, Recipient, RequestType};
 use embassy_usb::driver::Driver;
-use embassy_usb::{msos, Builder, Handler};
+use embassy_usb::{Builder, FunctionBuilder, Handler};
 use embedded_storage::nor_flash::{NorFlash, NorFlashErrorKind};
 
 use crate::consts::{
@@ -186,22 +186,13 @@ impl<'d, DFU: NorFlash, STATE: NorFlash, RST: Reset, const BLOCK_SIZE: usize> Ha
 pub fn usb_dfu<'d, D: Driver<'d>, DFU: NorFlash, STATE: NorFlash, RST: Reset, const BLOCK_SIZE: usize>(
     builder: &mut Builder<'d, D>,
     handler: &'d mut Control<'d, DFU, STATE, RST, BLOCK_SIZE>,
-    winusb_guids: Option<&'d [&str]>,
+    func_modifier: impl Fn(&mut FunctionBuilder<'_, 'd, D>),
 ) {
     let mut func = builder.function(USB_CLASS_APPN_SPEC, APPN_SPEC_SUBCLASS_DFU, DFU_PROTOCOL_DFU);
-    if let Some(winusb_guids) = winusb_guids {
-        // We add MSOS headers so that the device automatically gets assigned the WinUSB driver on Windows.
-        // Otherwise users need to do this manually using a tool like Zadig.
-        //
-        // Adding them here on the function level appears to only be needed for compositive devices.
-        // In addition to being on the function level, they should also be added to the device level.
-        //
-        func.msos_feature(msos::CompatibleIdFeatureDescriptor::new("WINUSB", ""));
-        func.msos_feature(msos::RegistryPropertyFeatureDescriptor::new(
-            "DeviceInterfaceGUIDs",
-            msos::PropertyData::RegMultiSz(winusb_guids),
-        ));
-    }
+
+    // Here we give users the opportunity to add their own function level MSOS headers for instance.
+    // This is useful when DFU functionality is part of a composite USB device.
+    func_modifier(&mut func);
 
     let mut iface = func.interface();
     let mut alt = iface.alt_setting(USB_CLASS_APPN_SPEC, APPN_SPEC_SUBCLASS_DFU, DFU_PROTOCOL_DFU, None);
diff --git a/examples/boot/application/stm32wb-dfu/src/main.rs b/examples/boot/application/stm32wb-dfu/src/main.rs
index 6236dfe52..4d6556597 100644
--- a/examples/boot/application/stm32wb-dfu/src/main.rs
+++ b/examples/boot/application/stm32wb-dfu/src/main.rs
@@ -70,11 +70,15 @@ async fn main(_spawner: Spawner) {
         msos::PropertyData::RegMultiSz(DEVICE_INTERFACE_GUIDS),
     ));
 
-    // For non-composite devices:
-    usb_dfu(&mut builder, &mut state, Duration::from_millis(2500), None);
-
-    // Or for composite devices:
-    // usb_dfu(&mut builder, &mut state, Duration::from_millis(2500), Some(DEVICE_INTERFACE_GUIDS));
+    usb_dfu(&mut builder, &mut state, Duration::from_millis(2500), |func| {
+        // You likely don't have to add these function level headers if your USB device is not composite
+        // (i.e. if your device does not expose another interface in addition to DFU)
+        func.msos_feature(msos::CompatibleIdFeatureDescriptor::new("WINUSB", ""));
+        func.msos_feature(msos::RegistryPropertyFeatureDescriptor::new(
+            "DeviceInterfaceGUIDs",
+            msos::PropertyData::RegMultiSz(DEVICE_INTERFACE_GUIDS),
+        ));
+    });
 
     let mut dev = builder.build();
     dev.run().await
diff --git a/examples/boot/bootloader/stm32wb-dfu/src/main.rs b/examples/boot/bootloader/stm32wb-dfu/src/main.rs
index 8cfd4daa7..fea6f4a0d 100644
--- a/examples/boot/bootloader/stm32wb-dfu/src/main.rs
+++ b/examples/boot/bootloader/stm32wb-dfu/src/main.rs
@@ -78,11 +78,15 @@ fn main() -> ! {
             msos::PropertyData::RegMultiSz(DEVICE_INTERFACE_GUIDS),
         ));
 
-        // For non-composite devices:
-        usb_dfu::<_, _, _, _, 4096>(&mut builder, &mut state, None);
-
-        // Or for composite devices:
-        // usb_dfu::<_, _, _, _, 4096>(&mut builder, &mut state, Some(DEVICE_INTERFACE_GUIDS));
+        usb_dfu::<_, _, _, _, 4096>(&mut builder, &mut state, |func| {
+            // You likely don't have to add these function level headers if your USB device is not composite
+            // (i.e. if your device does not expose another interface in addition to DFU)
+            func.msos_feature(msos::CompatibleIdFeatureDescriptor::new("WINUSB", ""));
+            func.msos_feature(msos::RegistryPropertyFeatureDescriptor::new(
+                "DeviceInterfaceGUIDs",
+                msos::PropertyData::RegMultiSz(DEVICE_INTERFACE_GUIDS),
+            ));
+        });
 
         let mut dev = builder.build();
         embassy_futures::block_on(dev.run());