Helper: Improve connection switching when toggling interfaces

helper/helper/device.py

@@ -202,7 +202,12 @@ class DevicesNode(RpcNode):
     def __call__(self, *args, **kwargs):
         with self._get_state:
             try:
-                return super().__call__(*args, **kwargs)
+                response = super().__call__(*args, **kwargs)
+                if "device_closed" in response.flags:
+                    self._list_state = 0
+                    self._device_mapping = {}
+                    response.flags.remove("device_closed")
+                return response
             except ConnectionException as e:
                 if self._failing_connection == e.body:
                     self._retries += 1
@@ -270,6 +275,13 @@ class AbstractDeviceNode(RpcNode):
     def __call__(self, *args, **kwargs):
         try:
             response = super().__call__(*args, **kwargs)
+
+            # The command resulted in the device closing
+            if "device_closed" in response.flags:
+                self.close()
+                return response
+
+            # The command resulted in device_info modification
             if "device_info" in response.flags:
                 old_info = self._info
                 # Refresh data
@@ -296,7 +308,9 @@ class AbstractDeviceNode(RpcNode):
             raise NoSuchNodeException(name)
 
     def get_data(self):
+        if self._data:
             return self._data
+        raise ChildResetException("Unable to read device data")
 
     def _refresh_data(self):
         ...
@@ -326,7 +340,12 @@ class UsbDeviceNode(AbstractDeviceNode):
         if self._child and not self._child.closed:
             # Make sure to close any open session
             self._child._close_child()
+            try:
                 return self._read_data(self._child._connection)
+            except Exception:
+                logger.warning(
+                    f"Unable to use {self._child._connection}", exc_info=True
+                )
 
         # No child, open new connection
         for conn_type in (SmartCardConnection, OtpConnection, FidoConnection):
@@ -336,7 +355,9 @@ class UsbDeviceNode(AbstractDeviceNode):
                         return self._read_data(conn)
                 except Exception:
                     logger.warning(f"Unable to connect via {conn_type}", exc_info=True)
-        raise ValueError("No supported connections")
+        # Failed to refresh, close
+        self.close()
+        return None
 
     @child(condition=lambda self: self._supports_connection(SmartCardConnection))
     def ccid(self):

helper/helper/fido.py

@@ -203,7 +203,7 @@ class Ctap2Node(RpcNode):
             raise
         self._info = self.ctap.get_info()
         self._token = None
-        return RpcResponse(dict(), ["device_info"])
+        return RpcResponse(dict(), ["device_info", "device_closed"])
 
     @action(condition=lambda self: self._info.options["clientPin"])
     def unlock(self, params, event, signal):

helper/helper/management.py

@@ -51,10 +51,6 @@ class ManagementNode(RpcNode):
     def _await_reboot(self, serial, usb_enabled):
         ifaces = CAPABILITY(usb_enabled or 0).usb_interfaces
 
-        # Prefer to use the "same" connection type as before
-        if self._connection_type.usb_interface in ifaces:
-            connection_types = [self._connection_type]
-        else:
         types: list[Type[Connection]] = [
             SmartCardConnection,
             OtpConnection,
@@ -62,9 +58,13 @@ class ManagementNode(RpcNode):
             FidoConnection,  # type: ignore
         ]
         connection_types = [t for t in types if t.usb_interface in ifaces]
+        # Prefer to use the "same" connection type as before
+        if self._connection_type in connection_types:
+            connection_types.remove(self._connection_type)
+            connection_types.insert(0, self._connection_type)
 
         self.session.close()
-        logger.debug("Waiting for device to re-appear...")
+        logger.debug(f"Waiting for device to re-appear over {connection_types}...")
         for _ in range(10):
             sleep(0.2)  # Always sleep initially
             for dev, info in list_all_devices(connection_types):
@@ -87,10 +87,12 @@ class ManagementNode(RpcNode):
         )
         serial = self.session.read_device_info().serial
         self.session.write_device_config(config, reboot, cur_lock_code, new_lock_code)
+        flags = ["device_info"]
         if reboot:
             enabled = config.enabled_capabilities.get(TRANSPORT.USB)
+            flags.append("device_closed")
             self._await_reboot(serial, enabled)
-        return RpcResponse(dict(), ["device_info"])
+        return RpcResponse(dict(), flags)
 
     @action
     def set_mode(self, params, event, signal):