Skip to content

Commit c351b05

Browse files
committed
fix: harden hardware wallet removal
1 parent 92a01a7 commit c351b05

7 files changed

Lines changed: 81 additions & 19 deletions

File tree

app/src/main/java/to/bitkit/models/HwWallet.kt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ import androidx.compose.runtime.Immutable
44
import androidx.compose.runtime.Stable
55
import com.synonym.bitkitcore.Activity
66
import kotlinx.collections.immutable.ImmutableList
7+
import kotlinx.collections.immutable.ImmutableSet
8+
import kotlinx.collections.immutable.persistentSetOf
79
import kotlinx.serialization.Serializable
810

911
/** A paired hardware wallet tracked as a watch-only balance. */
@@ -16,6 +18,7 @@ data class HwWallet(
1618
val isConnected: Boolean,
1719
val balanceSats: ULong,
1820
val activities: ImmutableList<Activity>,
21+
val deviceIds: ImmutableSet<String> = persistentSetOf(id),
1922
)
2023

2124
/** Serializable per-device balance snapshot carried by [BalanceState]. */

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import com.synonym.bitkitcore.WatcherEvent
99
import kotlinx.collections.immutable.ImmutableList
1010
import kotlinx.collections.immutable.persistentListOf
1111
import kotlinx.collections.immutable.toImmutableList
12+
import kotlinx.collections.immutable.toImmutableSet
1213
import kotlinx.coroutines.CoroutineDispatcher
1314
import kotlinx.coroutines.CoroutineScope
1415
import kotlinx.coroutines.SupervisorJob
@@ -124,11 +125,9 @@ class HwWalletRepo @Inject constructor(
124125
activeWatchers.toList()
125126
.filter { it.toDeviceId() in ids }
126127
.forEach { stopActiveWatcher(it) }
127-
// forgetDevice can report a benign cleanup failure (e.g. disconnecting the live
128-
// transport) after it has already dropped the device from the store, so treat the
129-
// removal as successful as long as the device is actually gone afterwards.
130-
ids.forEach { trezorRepo.forgetDevice(it) }
128+
val failures = ids.mapNotNull { trezorRepo.forgetDevice(it).exceptionOrNull() }
131129
val remaining = hwWalletStore.loadKnownDevices().map { it.id }.toSet()
130+
failures.firstOrNull()?.let { throw it }
132131
check(ids.none { it in remaining }) { "Hardware wallet '$deviceId' still present after removal" }
133132
}
134133
}
@@ -159,6 +158,7 @@ class HwWalletRepo @Inject constructor(
159158
activities = deviceWatchers
160159
.toMergedActivities()
161160
.toImmutableList(),
161+
deviceIds = ids.toImmutableSet(),
162162
)
163163
}
164164
.toImmutableList()

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -655,8 +655,11 @@ class TrezorRepo @Inject constructor(
655655
val updated = _state.value.knownDevices.filter { it.id != deviceId }
656656
saveKnownDevices(updated)
657657
_state.update { it.copy(knownDevices = updated.toImmutableList()) }
658-
disconnectResult.getOrThrow()
659658
clearCredentialsResult.getOrThrow()
659+
disconnectResult.onFailure {
660+
TrezorDebugLog.log("FORGET", "Ignored disconnect failure: ${it.message}")
661+
Logger.warn("Ignored disconnect failure while forgetting device '$deviceId'", it, context = TAG)
662+
}
660663
TrezorDebugLog.log("FORGET", "Device forgotten successfully")
661664
Logger.info("Forgot device: '$deviceId'", context = TAG)
662665
}.onFailure { e ->

app/src/main/java/to/bitkit/ui/screens/wallets/HardwareWalletScreen.kt

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ fun HardwareWalletScreen(
6969
) {
7070
val wallets by viewModel.wallets.collectAsStateWithLifecycle()
7171
val uiState by viewModel.uiState.collectAsStateWithLifecycle()
72-
val wallet = remember(wallets, deviceId) { wallets.find { it.id == deviceId } }
72+
val wallet = remember(wallets, deviceId) { wallets.find { deviceId in it.deviceIds } }
7373

7474
// Leave the screen once the device is gone, whether removed here or forgotten elsewhere.
7575
LaunchedEffect(wallet) {
@@ -195,17 +195,23 @@ private fun HardwareWalletContent(
195195
showFooter = false,
196196
hardwareIds = hardwareIds,
197197
footerContent = {
198-
TertiaryButton(
199-
text = stringResource(R.string.hardware__remove_button, wallet.name),
198+
RemoveHardwareWalletButton(
199+
walletName = wallet.name,
200200
onClick = onRemoveClick,
201-
modifier = Modifier
202-
.wrapContentWidth()
203-
.padding(top = 8.dp)
204-
.testTag("RemoveHardwareWallet")
205201
)
206202
},
207203
)
208204
}
205+
if (showEmptyState) {
206+
item { VerticalSpacer(32.dp) }
207+
item {
208+
RemoveHardwareWalletButton(
209+
walletName = wallet.name,
210+
onClick = onRemoveClick,
211+
)
212+
}
213+
item { VerticalSpacer(120.dp) }
214+
}
209215
}
210216

211217
if (showRemoveDialog) {
@@ -221,6 +227,22 @@ private fun HardwareWalletContent(
221227
}
222228
}
223229

230+
@Composable
231+
private fun RemoveHardwareWalletButton(
232+
walletName: String,
233+
onClick: () -> Unit,
234+
modifier: Modifier = Modifier,
235+
) {
236+
TertiaryButton(
237+
text = stringResource(R.string.hardware__remove_button, walletName),
238+
onClick = onClick,
239+
modifier = modifier
240+
.wrapContentWidth()
241+
.padding(top = 8.dp)
242+
.testTag("RemoveHardwareWallet")
243+
)
244+
}
245+
224246
private fun previewWallet(
225247
balanceSats: ULong = 10_562_411uL,
226248
activities: ImmutableList<Activity> = previewOnchainActivityItems(),

app/src/main/java/to/bitkit/ui/screens/wallets/activity/components/ActivityListGrouped.kt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,10 @@ fun LazyListScope.activityListGroupedItems(
255255
)
256256
}
257257
}
258+
footerContent?.let { content ->
259+
item { content() }
260+
item { VerticalSpacer(120.dp) }
261+
}
258262
}
259263
}
260264

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

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ class HwWalletRepoTest : BaseUnitTest() {
8484

8585
val wallet = sut.wallets.value.single()
8686
assertEquals("dev1", wallet.id)
87+
assertEquals(setOf("dev1"), wallet.deviceIds)
8788
assertEquals("Trezor", wallet.name)
8889
assertEquals(0uL, wallet.balanceSats)
8990
assertEquals(0uL, sut.totalSats.value)
@@ -506,6 +507,7 @@ class HwWalletRepoTest : BaseUnitTest() {
506507
assertEquals(421_900uL, wallet.balanceSats)
507508
assertEquals(421_900uL, sut.totalSats.value)
508509
assertEquals(1, wallet.activities.size)
510+
assertEquals(setOf("ble1", "usb1"), wallet.deviceIds)
509511
assertEquals(TransportType.USB, wallet.transportType)
510512
}
511513

@@ -523,6 +525,7 @@ class HwWalletRepoTest : BaseUnitTest() {
523525

524526
val wallet = sut.wallets.value.single()
525527
assertEquals("usb1", wallet.id)
528+
assertEquals(setOf("ble1", "usb1"), wallet.deviceIds)
526529
assertEquals(TransportType.USB, wallet.transportType)
527530
assertEquals(true, wallet.isConnected)
528531
}
@@ -599,14 +602,14 @@ class HwWalletRepoTest : BaseUnitTest() {
599602
}
600603

601604
@Test
602-
fun `removeDevice succeeds when the device is gone despite a cleanup failure`() = test {
605+
fun `removeDevice fails when forget reports credential cleanup failure despite the device being gone`() = test {
603606
whenever(hwWalletStore.loadKnownDevices()).thenReturn(listOf(device), emptyList())
604-
whenever { trezorRepo.forgetDevice(any()) }.thenReturn(Result.failure(AppError("disconnect failed")))
607+
whenever { trezorRepo.forgetDevice(any()) }.thenReturn(Result.failure(AppError("clear failed")))
605608
val sut = createRepo()
606609

607610
val result = sut.removeDevice("dev1")
608611

609-
assertEquals(true, result.isSuccess)
612+
assertEquals(true, result.isFailure)
610613
verify(trezorRepo).forgetDevice("dev1")
611614
}
612615

@@ -631,8 +634,8 @@ class HwWalletRepoTest : BaseUnitTest() {
631634

632635
@Test
633636
fun `removeDevice fails when the device is still present afterwards`() = test {
634-
whenever(hwWalletStore.loadKnownDevices()).thenReturn(listOf(device))
635-
whenever { trezorRepo.forgetDevice(any()) }.thenReturn(Result.failure(AppError("forget failed")))
637+
whenever(hwWalletStore.loadKnownDevices()).thenReturn(listOf(device), listOf(device))
638+
whenever { trezorRepo.forgetDevice(any()) }.thenReturn(Result.success(Unit))
636639
val sut = createRepo()
637640

638641
val result = sut.removeDevice("dev1")

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

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -898,7 +898,7 @@ class TrezorRepoTest : BaseUnitTest() {
898898
// region forgetDevice
899899

900900
@Test
901-
fun `forgetDevice should remove known device when service cleanup fails`() = test {
901+
fun `forgetDevice should remove known device when disconnect cleanup fails`() = test {
902902
val knownDevice = mockKnownDevice()
903903
val features = mockFeatures()
904904
val device = mockDeviceInfo()
@@ -911,6 +911,32 @@ class TrezorRepoTest : BaseUnitTest() {
911911
sut.scan()
912912
sut.connect(DEVICE_ID)
913913
whenever(trezorService.disconnect()).thenThrow(RuntimeException("disconnect failed"))
914+
915+
val result = sut.forgetDevice(DEVICE_ID)
916+
917+
assertTrue(result.isSuccess)
918+
assertTrue(sut.state.value.knownDevices.isEmpty())
919+
assertNull(sut.state.value.connectedDevice)
920+
assertNull(sut.state.value.connectedDeviceId)
921+
assertNull(sut.state.value.error)
922+
verify(trezorTransport).clearDeviceCredential(DEVICE_ID)
923+
verify(trezorService).clearCredentials(DEVICE_ID)
924+
verify(hwWalletStore).saveKnownDevices(emptyList())
925+
}
926+
927+
@Test
928+
fun `forgetDevice should fail when credential cleanup fails`() = test {
929+
val knownDevice = mockKnownDevice()
930+
val features = mockFeatures()
931+
val device = mockDeviceInfo()
932+
whenever(hwWalletStore.loadKnownDevices()).thenReturn(listOf(knownDevice))
933+
whenever(trezorService.connect(eq(DEVICE_ID), any())).thenReturn(features)
934+
whenever(trezorService.scan()).thenReturn(listOf(device))
935+
sut = createSut()
936+
937+
sut.initialize()
938+
sut.scan()
939+
sut.connect(DEVICE_ID)
914940
whenever(trezorService.clearCredentials(DEVICE_ID)).thenThrow(RuntimeException("clear failed"))
915941

916942
val result = sut.forgetDevice(DEVICE_ID)
@@ -919,7 +945,8 @@ class TrezorRepoTest : BaseUnitTest() {
919945
assertTrue(sut.state.value.knownDevices.isEmpty())
920946
assertNull(sut.state.value.connectedDevice)
921947
assertNull(sut.state.value.connectedDeviceId)
922-
assertEquals("disconnect failed", sut.state.value.error)
948+
assertEquals("clear failed", result.exceptionOrNull()?.message)
949+
assertEquals("clear failed", sut.state.value.error)
923950
verify(trezorTransport).clearDeviceCredential(DEVICE_ID)
924951
verify(trezorService).clearCredentials(DEVICE_ID)
925952
verify(hwWalletStore).saveKnownDevices(emptyList())

0 commit comments

Comments
 (0)