[opentitantool] Improvement to HyperDebug firmware upgrade

Firmware upgrading of HyperDebug requires entering and leaving "DFU
mode", each of which involves resetting connection with the USB device.

Previous code waited one second for the device to reset after upgrading
concluded, before attempting to restablish USB connection with the newly
flashed firmware.  This turns out to be insufficient on some machines.

This patch will wait up to five seconds, checking twice a second if the
device is available, and terminating early if it is.

Change-Id: Ib59d7f5ebaf234a5c2508253236f6bc3ee96c83f
Signed-off-by: Jes B. Klinke <[email protected]>

sw/host/opentitanlib/src/transport/hyperdebug/dfu.rs

@@ -4,7 +4,7 @@
 
 // Firmware update protocol for HyperDebug
 
-use anyhow::{anyhow, bail, Context, Result};
+use anyhow::{anyhow, bail, Result};
 use once_cell::sync::Lazy;
 use regex::Regex;
 use serde_annotate::Annotate;
@@ -23,8 +23,12 @@ const PID_DFU_BOOTLOADER: u16 = 0xdf11;
 /// of the `opentitantool` invocation (and presenting itself with STMs VID:DID, rather than
 /// Google's).
 pub struct HyperdebugDfu {
+    // Handle to USB device which may or may not already be in DFU mode
     usb_backend: RefCell<UsbBackend>,
     current_firmware_version: Option<String>,
+    // expected USB VID:PID of the HyperDebug device when not in DFU mode
+    usb_vid: u16,
+    usb_pid: u16,
 }
 
 impl HyperdebugDfu {
@@ -41,18 +45,18 @@ impl HyperdebugDfu {
         // will put the desired firmware on the HyperDebug, both in the case of previous
         // interrupted update, as well as the ordinary case of outdated or current HyperDebug
         // firmware already running.
-        if let Ok(usb_backend) = UsbBackend::new(
-            usb_vid.unwrap_or(VID_ST_MICROELECTRONICS),
-            usb_pid.unwrap_or(PID_DFU_BOOTLOADER),
-            usb_serial,
-        ) {
+        if let Ok(usb_backend) =
+            UsbBackend::new(VID_ST_MICROELECTRONICS, PID_DFU_BOOTLOADER, usb_serial)
+        {
             // HyperDebug device is already in DFU mode, we cannot query firmware version through
             // USB strings.  (And the fact that it was left in DFU mode, probably as a result of a
             // previous incomplete update attempt, should mean that we would not want to trust the
             // version, even if we could extract it through the DFU firmware.)
             return Ok(Self {
                 usb_backend: RefCell::new(usb_backend),
                 current_firmware_version: None,
+                usb_vid: usb_vid.unwrap_or(super::VID_GOOGLE),
+                usb_pid: usb_pid.unwrap_or(super::PID_HYPERDEBUG),
             });
         }
 
@@ -76,6 +80,8 @@ impl HyperdebugDfu {
         Ok(Self {
             usb_backend: RefCell::new(usb_backend),
             current_firmware_version,
+            usb_vid: usb_vid.unwrap_or(super::VID_GOOGLE),
+            usb_pid: usb_pid.unwrap_or(super::PID_HYPERDEBUG),
         })
     }
 }
@@ -94,6 +100,8 @@ impl Transport for HyperdebugDfu {
                 &update_firmware_action.firmware,
                 update_firmware_action.progress.as_ref(),
                 update_firmware_action.force,
+                self.usb_vid,
+                self.usb_pid,
             )
         } else {
             bail!(TransportError::UnsupportedOperation)
@@ -167,6 +175,8 @@ pub fn update_firmware(
     firmware: &Option<Vec<u8>>,
     progress: &dyn ProgressIndicator,
     force: bool,
+    usb_vid: u16,
+    usb_pid: u16,
 ) -> Result<Option<Box<dyn Annotate>>> {
     let firmware: &[u8] = if let Some(vec) = firmware.as_ref() {
         validate_firmware_image(vec)?;
@@ -196,6 +206,7 @@ pub fn update_firmware(
     if wait_for_idle(usb_device, dfu_desc.dfu_interface)? != DFU_STATE_APP_IDLE {
         // Device is already running DFU bootloader, proceed to firmware transfer.
         do_update_firmware(usb_device, dfu_desc, firmware, progress)?;
+        restablish_connection(usb_vid, usb_pid, usb_device.get_serial_number())?;
         return Ok(None);
     }
 
@@ -233,19 +244,24 @@ pub fn update_firmware(
     let dfu_desc = scan_usb_descriptor(&dfu_device)?;
     dfu_device.claim_interface(dfu_desc.dfu_interface)?;
     do_update_firmware(&dfu_device, dfu_desc, firmware, progress)?;
+    restablish_connection(usb_vid, usb_pid, usb_device.get_serial_number())?;
+    Ok(None)
+}
 
+fn restablish_connection(usb_vid: u16, usb_pid: u16, serial_number: &str) -> Result<()> {
     // At this point, the new firmware has been completely transferred, and the USB device is
-    // resetting and booting the new firmware.  Wait a second, then verify that device can now be
-    // found on the USB bus with the original DID:VID.
-    std::thread::sleep(std::time::Duration::from_millis(1000));
+    // resetting and booting the new firmware.  Wait up to five seconds, repeatedly testing if the
+    // device can be found on the USB bus with the original DID:VID.
     log::info!("Connecting to newly flashed firmware...");
-    let _new_device = UsbBackend::new(
-        usb_device.get_vendor_id(),
-        usb_device.get_product_id(),
-        Some(usb_device.get_serial_number()),
-    )
-    .context("Unable to establish connection after flashing.  Possibly bad image.")?;
-    Ok(None)
+    for _ in 0..10 {
+        std::thread::sleep(std::time::Duration::from_millis(500));
+        if UsbBackend::new(usb_vid, usb_pid, Some(serial_number)).is_ok() {
+            return Ok(());
+        }
+    }
+    bail!(TransportError::FirmwareProgramFailed(
+        "Unable to establish connection after flashing.  Possibly bad image.".to_string()
+    ));
 }
 
 fn do_update_firmware(

sw/host/opentitanlib/src/transport/hyperdebug/mod.rs

@@ -808,12 +808,16 @@ impl<T: Flavor> Transport for Hyperdebug<T> {
 
     fn dispatch(&self, action: &dyn Any) -> Result<Option<Box<dyn Annotate>>> {
         if let Some(update_firmware_action) = action.downcast_ref::<UpdateFirmware>() {
+            let usb_vid = self.inner.usb_device.borrow().get_vendor_id();
+            let usb_pid = self.inner.usb_device.borrow().get_product_id();
             dfu::update_firmware(
                 &mut self.inner.usb_device.borrow_mut(),
                 self.current_firmware_version.as_deref(),
                 &update_firmware_action.firmware,
                 update_firmware_action.progress.as_ref(),
                 update_firmware_action.force,
+                usb_vid,
+                usb_pid,
             )
         } else if let Some(jtag_set_pins) = action.downcast_ref::<SetJtagPins>() {
             match (