fix: harden hw pairing flow

Author: ovitrif

app/src/main/java/to/bitkit/repositories/HwWalletRepo.kt

@@ -127,7 +127,10 @@ class HwWalletRepo @Inject constructor(
     suspend fun hasKnownDevice(deviceId: String): Boolean = trezorRepo.hasKnownDevice(deviceId)
 
     /** Connects and pairs a discovered device, persisting it as a watch-only known device. */
-    suspend fun connect(deviceId: String): Result<TrezorFeatures> = trezorRepo.connect(deviceId)
+    suspend fun connect(deviceId: String): Result<TrezorFeatures> {
+        trezorRepo.resetWalletSelection()
+        return trezorRepo.connect(deviceId)
+    }
 
     /**
      * Persists the Bitkit-side funds label for a paired device. Applied to every entry sharing the

app/src/main/java/to/bitkit/repositories/TrezorRepo.kt

@@ -828,9 +828,9 @@ class TrezorRepo @Inject constructor(
     }
 
     private suspend fun addOrUpdateKnownDevice(deviceInfo: TrezorDeviceInfo, features: TrezorFeatures) {
-        val existing = _state.value.knownDevices
-        val existingIds = existing.map { it.id }.toSet()
-        val knownDevices = existing + hwWalletStore.loadKnownDevices().filter { it.id !in existingIds }
+        val stored = hwWalletStore.loadKnownDevices()
+        val storedIds = stored.map { it.id }.toSet()
+        val knownDevices = stored + _state.value.knownDevices.filter { it.id !in storedIds }
         val previous = knownDevices.find { it.id == deviceInfo.id }
         val known = KnownDevice(
             id = deviceInfo.id,

app/src/main/java/to/bitkit/ui/sheets/hardware/HardwareSheet.kt

@@ -160,7 +160,10 @@ fun HardwareSheet(
                     isConnecting = uiState.isConnecting,
                     errorMessage = uiState.errorMessage,
                     onConnect = { viewModel.onConnectClick(route.deviceId) },
-                    onCancel = appViewModel::hideSheet,
+                    onCancel = {
+                        viewModel.cancelConnect()
+                        appViewModel.hideSheet()
+                    },
                 )
             }
             composableWithDefaultTransitions<HardwareRoute.Paired> {
@@ -180,6 +183,7 @@ fun HardwareSheet(
     }
 
     BackHandler {
+        viewModel.cancelConnect()
         appViewModel.hideSheet()
     }
 

app/src/main/java/to/bitkit/ui/sheets/hardware/HwConnectViewModel.kt

@@ -47,6 +47,7 @@ class HwConnectViewModel @Inject constructor(
     val effects = _effects.asSharedFlow()
 
     private var searchJob: Job? = null
+    private var connectJob: Job? = null
     private var labelInitialized = false
     private var includeBluetoothInScan = true
     private var scanUsbBeforeConnect = false
@@ -85,26 +86,64 @@ class HwConnectViewModel @Inject constructor(
     fun onConnectClick(deviceIdOverride: String? = null) {
         val state = _uiState.value
         val deviceId = deviceIdOverride ?: state.foundDeviceId ?: return
+        if (connectJob?.isActive == true) return
         val shouldScanUsbBeforeConnect = scanUsbBeforeConnect
         searchJob?.cancel()
         _uiState.update { it.copy(isConnecting = true, errorMessage = null) }
-        viewModelScope.launch {
+        connectJob = viewModelScope.launch {
+            var resolvedDeviceId = deviceId
+            var resolvedDeviceModel = state.deviceModel
             if (shouldScanUsbBeforeConnect) {
                 hwWalletRepo.scan(includeBluetooth = false)
-            }
-            hwWalletRepo.connect(deviceId)
-                .onSuccess { onConnected(deviceId, it) }
-                .onFailure {
-                    _uiState.update { state ->
-                        state.copy(
-                            isConnecting = false,
-                            errorMessage = context.getString(R.string.hardware__connect_error),
-                        )
+                    .onSuccess { devices ->
+                        devices.firstOrNull { it.id == deviceId || it.path == deviceId }?.let { device ->
+                            resolvedDeviceId = device.id
+                            resolvedDeviceModel = resolveHwWalletName(label = null, model = device.model)
+                            _uiState.update {
+                                it.copy(
+                                    foundDeviceId = resolvedDeviceId,
+                                    deviceModel = resolvedDeviceModel,
+                                )
+                            }
+                        }
                     }
-                }
+                    .onFailure {
+                        onConnectFailed(resolvedDeviceId, resolvedDeviceModel)
+                        return@launch
+                    }
+            }
+            hwWalletRepo.connect(resolvedDeviceId)
+                .onSuccess { onConnected(resolvedDeviceId, it) }
+                .onFailure { onConnectFailed(resolvedDeviceId, resolvedDeviceModel) }
+            connectJob = null
         }
     }
 
+    private fun onConnectFailed(deviceId: String, deviceModel: String) {
+        _uiState.update {
+            it.copy(
+                isConnecting = false,
+                foundDeviceId = deviceId,
+                deviceModel = deviceModel,
+                errorMessage = context.getString(R.string.hardware__connect_error),
+            )
+        }
+        setEffect(
+            HwConnectEffect.NavigateToFound(
+                deviceId = deviceId,
+                deviceModel = deviceModel,
+            )
+        )
+        connectJob = null
+    }
+
+    fun cancelConnect() {
+        connectJob?.cancel()
+        connectJob = null
+        hwWalletRepo.cancelPairingCode()
+        _uiState.update { it.copy(isConnecting = false) }
+    }
+
     fun onLabelChange(value: String) = _uiState.update { it.copy(labelInput = value.take(DEVICE_LABEL_MAX_LENGTH)) }
 
     fun onFinishClick() {
@@ -122,6 +161,9 @@ class HwConnectViewModel @Inject constructor(
     fun resetState() {
         searchJob?.cancel()
         searchJob = null
+        connectJob?.cancel()
+        connectJob = null
+        hwWalletRepo.cancelPairingCode()
         labelInitialized = false
         includeBluetoothInScan = true
         scanUsbBeforeConnect = false

app/src/test/java/to/bitkit/repositories/HwWalletRepoTest.kt

@@ -773,6 +773,7 @@ class HwWalletRepoTest : BaseUnitTest() {
 
         sut.connect("dev1")
 
+        verify(trezorRepo).resetWalletSelection()
         verify(trezorRepo).connect("dev1")
     }
 

app/src/test/java/to/bitkit/repositories/TrezorRepoTest.kt

@@ -146,6 +146,7 @@ class TrezorRepoTest : BaseUnitTest() {
         model: String? = DEVICE_MODEL,
         transportType: TransportType = TransportType.USB,
         xpubs: Map<String, String> = emptyMap(),
+        customLabel: String? = null,
     ) = KnownDevice(
         id = id,
         name = name,
@@ -155,6 +156,7 @@ class TrezorRepoTest : BaseUnitTest() {
         model = model,
         lastConnectedAt = 123L,
         xpubs = xpubs,
+        customLabel = customLabel,
     )
 
     // region initialize
@@ -584,6 +586,26 @@ class TrezorRepoTest : BaseUnitTest() {
         )
     }
 
+    @Test
+    fun `connect preserves stored custom label over stale state label`() = test {
+        val features = mockFeatures()
+        val device = mockDeviceInfo()
+        whenever(hwWalletStore.loadKnownDevices())
+            .thenReturn(listOf(mockKnownDevice()))
+            .thenReturn(listOf(mockKnownDevice(customLabel = "Cold Storage")))
+        whenever(trezorService.connect(eq(DEVICE_ID), any())).thenReturn(features)
+        whenever(trezorService.scan()).thenReturn(listOf(device))
+        sut = createSut()
+
+        sut.scan()
+        val result = sut.connect(DEVICE_ID)
+
+        assertTrue(result.isSuccess)
+        val captor = argumentCaptor<List<KnownDevice>>()
+        verify(hwWalletStore).saveKnownDevices(captor.capture())
+        assertEquals("Cold Storage", captor.lastValue.single().customLabel)
+    }
+
     @Test
     fun `connect should retry once for retryable THP errors`() = test {
         val features = mockFeatures()

app/src/test/java/to/bitkit/ui/sheets/hardware/HwConnectViewModelTest.kt

@@ -118,6 +118,46 @@ class HwConnectViewModelTest : BaseUnitTest() {
         verify(hwWalletRepo).connect("usb1")
     }
 
+    @Test
+    fun `onConnectClick uses scanned device id for usb route path`() = test {
+        val path = "/dev/bus/usb/001/002"
+        val connectedFeatures = features(model = "Safe 5")
+        val usbDevice = deviceInfo(
+            id = "core-usb-id",
+            model = "Safe 5",
+            transportType = TrezorTransportType.USB,
+            path = path,
+        )
+        whenever(hwWalletRepo.scan(includeBluetooth = false)).thenReturn(Result.success(listOf(usbDevice)))
+        whenever(hwWalletRepo.connect("core-usb-id")).thenReturn(Result.success(connectedFeatures))
+        sut.onFoundRoute(deviceId = path, deviceModel = "Trezor")
+
+        sut.effects.test {
+            sut.onConnectClick()
+            assertEquals(HwConnectEffect.NavigateToPaired, awaitItem())
+            cancelAndIgnoreRemainingEvents()
+        }
+
+        verify(hwWalletRepo).connect("core-usb-id")
+        assertEquals("core-usb-id", sut.uiState.value.foundDeviceId)
+    }
+
+    @Test
+    fun `onConnectClick returns to found when connect fails from pair code`() = test {
+        whenever(hwWalletRepo.scan(includeBluetooth = false)).thenReturn(Result.success(emptyList()))
+        whenever(hwWalletRepo.connect("usb1")).thenReturn(Result.failure(AppError("connect failed")))
+        sut.onFoundRoute(deviceId = "usb1", deviceModel = "Trezor Safe 5")
+
+        sut.effects.test {
+            sut.onConnectClick()
+            assertEquals(HwConnectEffect.NavigateToFound("usb1", "Trezor Safe 5"), awaitItem())
+            cancelAndIgnoreRemainingEvents()
+        }
+
+        assertFalse(sut.uiState.value.isConnecting)
+        assertEquals(CONNECT_ERROR, sut.uiState.value.errorMessage)
+    }
+
     @Test
     fun `onConnectClick connects the found device and advances to paired`() = test {
         givenDeviceFound()
@@ -203,11 +243,16 @@ class HwConnectViewModelTest : BaseUnitTest() {
         sut.onIntroContinue()
     }
 
-    private fun deviceInfo(id: String, model: String?) = TrezorDeviceInfo(
+    private fun deviceInfo(
+        id: String,
+        model: String?,
+        transportType: TrezorTransportType = TrezorTransportType.BLUETOOTH,
+        path: String = "ble:$id",
+    ) = TrezorDeviceInfo(
         id = id,
-        transportType = TrezorTransportType.BLUETOOTH,
+        transportType = transportType,
         name = null,
-        path = "ble:$id",
+        path = path,
         label = null,
         model = model,
         isBootloader = false,