From 4f25a8b6a4f7761078678876e259674d736768a5 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Tue, 27 Nov 2018 16:19:27 +1100
Subject: [PATCH] tools/pydfu.py: Improve DFU reset, and auto-detect USB
 transfer size.

A DFU device must be in the idle state before it can be programmed, and
this requires either clearing the status or aborting, depending on its
current state.  Code is added to do this.  And the USB transfer size is now
automatically detected so devices with a size less than 2048 bytes work
correctly.
---
 tools/pydfu.py | 47 +++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 41 insertions(+), 6 deletions(-)

diff --git a/tools/pydfu.py b/tools/pydfu.py
index e7f4ab1780..394ecca5e1 100755
--- a/tools/pydfu.py
+++ b/tools/pydfu.py
@@ -14,6 +14,7 @@ See document UM0391 for a dscription of the DFuse file.
 from __future__ import print_function
 
 import argparse
+import collections
 import re
 import struct
 import sys
@@ -56,6 +57,9 @@ _DFU_DESCRIPTOR_TYPE                 = 0x21
 # USB device handle
 __dev = None
 
+# Configuration descriptor of the device
+__cfg_descr = None
+
 __verbose = None
 
 # USB DFU interface
@@ -76,9 +80,18 @@ else:
         return usb.util.get_string(dev, index)
 
 
+def find_dfu_cfg_descr(descr):
+    if len(descr) == 9 and descr[0] == 9 and descr[1] == _DFU_DESCRIPTOR_TYPE:
+        nt = collections.namedtuple('CfgDescr',
+            ['bLength', 'bDescriptorType', 'bmAttributes',
+            'wDetachTimeOut', 'wTransferSize', 'bcdDFUVersion'])
+        return nt(*struct.unpack('<BBBHHH', bytes(descr)))
+    return None
+
+
 def init():
     """Initializes the found DFU device so that we can program it."""
-    global __dev
+    global __dev, __cfg_descr
     devices = get_dfu_devices(idVendor=__VID, idProduct=__PID)
     if not devices:
         raise ValueError('No DFU device found')
@@ -90,8 +103,32 @@ def init():
     # Claim DFU interface
     usb.util.claim_interface(__dev, __DFU_INTERFACE)
 
-    # Clear status
-    clr_status()
+    # Find the DFU configuration descriptor, either in the device or interfaces
+    __cfg_descr = None
+    for cfg in __dev.configurations():
+        __cfg_descr = find_dfu_cfg_descr(cfg.extra_descriptors)
+        if __cfg_descr:
+            break
+        for itf in cfg.interfaces():
+            __cfg_descr = find_dfu_cfg_descr(itf.extra_descriptors)
+            if __cfg_descr:
+                break
+
+    # Get device into idle state
+    for attempt in range(4):
+        status = get_status()
+        if status == __DFU_STATE_DFU_IDLE:
+            break
+        elif (status == __DFU_STATE_DFU_DOWNLOAD_IDLE
+            or status == __DFU_STATE_DFU_UPLOAD_IDLE):
+            abort_request()
+        else:
+            clr_status()
+
+
+def abort_request():
+    """Sends an abort request."""
+    __dev.ctrl_transfer(0x21, __DFU_ABORT, 0, __DFU_INTERFACE, None, __TIMEOUT)
 
 
 def clr_status():
@@ -180,9 +217,7 @@ def write_memory(addr, buf, progress=None, progress_addr=0, progress_size=0):
         set_address(xfer_base+xfer_bytes)
 
         # Send DNLOAD with fw data
-        # the "2048" is the DFU transfer size supported by the ST DFU bootloader
-        # TODO: this number should be extracted from the USB config descriptor
-        chunk = min(2048, xfer_total-xfer_bytes)
+        chunk = min(__cfg_descr.wTransferSize, xfer_total-xfer_bytes)
         __dev.ctrl_transfer(0x21, __DFU_DNLOAD, 2, __DFU_INTERFACE,
                             buf[xfer_bytes:xfer_bytes + chunk], __TIMEOUT)