feat(services): recover orphaned Android Keystore keys via on-chain matching

Author: justjoolz

app/src/main/java/com/flowfoundation/wallet/manager/account/AccountManager.kt

@@ -129,6 +129,21 @@ object AccountManager {
                 }
 
                 logd(TAG, "AccountManager initialization completed successfully")
+
+                // Attempt to recover orphaned Android Keystore keys that were
+                // generated during key rotation but never persisted to the
+                // account cache or KeyStorageManager. If any are found, they
+                // will be written to KeyStorageManager so that
+                // buildLocalKeyAccounts() can surface them in the account switcher.
+                try {
+                    val recovered = KeyStoreMigrationManager.recoverOrphanedKeystoreKeys()
+                    if (recovered > 0) {
+                        logd(TAG, "Recovered $recovered orphaned Keystore key(s)")
+                    }
+                } catch (e: Exception) {
+                    loge(TAG, "Orphaned key recovery failed (non-fatal): ${e.message}")
+                }
+
                 // Update Accounts info with WalletDataManager
                 WalletDataManager.updateWalletData()
 

app/src/main/java/com/flowfoundation/wallet/manager/account/KeyStoreMigrationManager.kt

@@ -1,7 +1,26 @@
 package com.flowfoundation.wallet.manager.account
 
+import com.flow.wallet.KeyManager
+import com.flow.wallet.toFormatString
+import com.flowfoundation.wallet.cache.AccountCacheManager
+import com.flowfoundation.wallet.cache.UserPrefixCacheManager
+import com.flowfoundation.wallet.manager.account.Account
+import com.flowfoundation.wallet.manager.account.Accounts
+import com.flowfoundation.wallet.manager.account.UserPrefix
+import com.flowfoundation.wallet.manager.account.UserPrefixes
+import com.flowfoundation.wallet.manager.app.chainNetWorkString
+import com.flowfoundation.wallet.manager.flow.FlowCadenceApi
+import com.flowfoundation.wallet.manager.walletdata.FlowWallet
+import com.flowfoundation.wallet.network.model.UserInfoData
+import com.flowfoundation.wallet.utils.Env
 import com.flowfoundation.wallet.utils.logd
+import com.flowfoundation.wallet.utils.loge
 import java.security.KeyStore
+import java.security.cert.Certificate
+import java.security.interfaces.ECPublicKey
+import java.security.spec.ECPoint
+import kotlinx.coroutines.runBlocking
+import androidx.core.content.edit
 
 /**
  * Manages migration of private keys from the old Android Keystore system
@@ -10,9 +29,27 @@ import java.security.KeyStore
  * This is critical for users upgrading from versions that used direct Android Keystore access
  * to versions that use the Flow-Wallet-Kit storage abstraction.
  *
+ * ## Orphaned Key Recovery
+ *
+ * When the app performs a key rotation, it:
+ * 1. Generates a new key pair in the Android Keystore (alias = `{KEYSTORE_ALIAS_PREFIX}{prefix}`)
+ * 2. Submits the rotation transaction on-chain (adding the new public key)
+ * 3. Persists the new `prefix` in AccountManager's cache
+ *
+ * If step 3 fails (crash, timeout, serialization error), the key physically exists in the
+ * Android Keystore hardware but the app has no record of the alias/prefix.
+ *
+ * [recoverOrphanedKeystoreKeys] bridges this gap by:
+ * - Enumerating all aliases in the Android Keystore
+ * - Extracting the EC public key from each entry's certificate
+ * - Matching it against the on-chain public keys for all known Flow addresses
+ * - Writing matched prefixes back into [KeyStorageManager] so that
+ *   `AccountManager.buildLocalKeyAccounts()` can surface them in the account switcher
  */
 object KeyStoreMigrationManager {
     private const val TAG = "KeyStoreMigration"
+    private const val RECOVERY_PREFS = "keystore_recovery_state"
+    private const val KEY_LAST_SCAN_ALIASES_HASH = "last_scan_aliases_hash"
 
     /**
      * Diagnostic function to list all keys in Android Keystore
@@ -38,4 +75,305 @@ object KeyStoreMigrationManager {
         }
     }
 
+    /**
+     * Scans the Android Keystore for keys that match on-chain public keys but have no
+     * corresponding entry in [KeyStorageManager]. When a match is found, the prefix is
+     * saved so that `buildLocalKeyAccounts()` can surface the account in the switcher.
+     *
+     * This function is idempotent and safe to call on every cold start. It short-circuits
+     * if the set of Keystore aliases has not changed since the last successful scan.
+     *
+     * @param knownAddresses Flow addresses to check against on-chain keys. If empty,
+     *   the function will collect addresses from the current AccountManager cache.
+     * @return The number of newly recovered key mappings.
+     */
+    fun recoverOrphanedKeystoreKeys(knownAddresses: Set<String> = emptySet()): Int {
+        logd(TAG, "=== recoverOrphanedKeystoreKeys START ===")
+
+        val keyStore: KeyStore
+        val walletAliases: List<String>
+
+        try {
+            keyStore = KeyStore.getInstance("AndroidKeyStore")
+            keyStore.load(null)
+
+            // Collect only aliases that follow the Flow Wallet naming convention
+            val aliasPrefix = KeyManager.KEYSTORE_ALIAS_PREFIX
+            walletAliases = mutableListOf<String>()
+            val enumeration = keyStore.aliases()
+            while (enumeration.hasMoreElements()) {
+                val alias = enumeration.nextElement()
+                if (alias.startsWith(aliasPrefix)) {
+                    walletAliases.add(alias)
+                }
+            }
+
+            logd(TAG, "Found ${walletAliases.size} wallet-related aliases in Android Keystore")
+
+            if (walletAliases.isEmpty()) {
+                logd(TAG, "No wallet aliases found, nothing to recover")
+                return 0
+            }
+        } catch (e: Exception) {
+            loge(TAG, "Failed to enumerate Android Keystore: ${e.message}")
+            return 0
+        }
+
+        // Short-circuit if the alias set hasn't changed since the last scan
+        val aliasesHash = walletAliases.sorted().hashCode()
+        val prefs = Env.getApp().getSharedPreferences(RECOVERY_PREFS, 0)
+        val lastHash = prefs.getInt(KEY_LAST_SCAN_ALIASES_HASH, 0)
+        if (aliasesHash == lastHash && lastHash != 0) {
+            logd(TAG, "Alias set unchanged since last scan, skipping")
+            return 0
+        }
+
+        // Collect all known Flow addresses to check against
+        val addresses = if (knownAddresses.isNotEmpty()) {
+            knownAddresses
+        } else {
+            collectKnownAddresses()
+        }
+
+        if (addresses.isEmpty()) {
+            logd(TAG, "No known Flow addresses to match against")
+            return 0
+        }
+
+        logd(TAG, "Checking ${walletAliases.size} aliases against ${addresses.size} addresses")
+
+        // Build a map of public key hex -> on-chain address for all known accounts
+        val onChainKeyMap = buildOnChainKeyMap(addresses)
+
+        if (onChainKeyMap.isEmpty()) {
+            logd(TAG, "No on-chain keys fetched, cannot match")
+            // Don't cache the hash - we want to retry when network is available
+            return 0
+        }
+
+        // Extract public keys from each Keystore alias and match against on-chain keys
+        var recoveredCount = 0
+        for (alias in walletAliases) {
+            val prefix = alias.removePrefix(KeyManager.KEYSTORE_ALIAS_PREFIX)
+
+            // Skip if this prefix is already known to KeyStorageManager
+            // We need to check all UIDs since we don't know which UID this prefix belongs to
+            if (isAliasAlreadyMapped(prefix)) {
+                logd(TAG, "Alias already mapped: $alias -> prefix $prefix")
+                continue
+            }
+
+            try {
+                val publicKeyHex = extractPublicKeyHex(keyStore, alias)
+                if (publicKeyHex == null) {
+                    logd(TAG, "Could not extract public key from alias: $alias")
+                    continue
+                }
+
+                logd(TAG, "Extracted public key from $alias: ${publicKeyHex.take(20)}...")
+
+                // Check if this public key matches any on-chain key
+                val matchedAddress = onChainKeyMap[publicKeyHex.lowercase()]
+                if (matchedAddress != null) {
+                    logd(TAG, "✓ MATCH FOUND: alias=$alias matches on-chain key for address=$matchedAddress")
+
+                    // Generate a synthetic UID for this orphaned key.
+                    // We use a deterministic ID based on the prefix so repeated runs
+                    // produce the same UID and don't create duplicates.
+                    val syntheticUid = "recovered_${prefix.take(16)}"
+
+                    // Create UserPrefix mapping
+                    val currentPrefixes = UserPrefixCacheManager.read()?.toMutableList() ?: mutableListOf()
+                    currentPrefixes.removeAll { it.userId == syntheticUid }
+                    currentPrefixes.add(UserPrefix(syntheticUid, prefix))
+                    UserPrefixCacheManager.cache(UserPrefixes().apply { addAll(currentPrefixes) })
+
+                    // Create a mock Account so it appears in the Switch Account list
+                    val newAccount = Account(
+                        userInfo = UserInfoData(
+                            username = "Recovered Account",
+                            avatar = "https://source.boringavatars.com/marble/120/$matchedAddress",
+                            created = System.currentTimeMillis().toString(),
+                            nickname = "Recovered Key"
+                        ),
+                        isActive = false,
+                        prefix = prefix,
+                        walletNodes = listOf(FlowWallet(address = matchedAddress, name = "Flow Wallet", emojiId = 1, chainIdString = chainNetWorkString()))
+                    )
+                    
+                    val currentAccounts = AccountCacheManager.read()?.toMutableList() ?: mutableListOf()
+                    currentAccounts.removeAll { it.prefix == prefix }
+                    currentAccounts.add(newAccount)
+                    AccountCacheManager.cache(Accounts().apply { addAll(currentAccounts) })
+
+                    // NOTE: Injecting directly into CacheManager as a temporary fallback.
+                    // Ideally, this should be handled by a dedicated Keystore abstraction layer in the future.
+
+                    logd(TAG, "Saved recovered mapping: uid=$syntheticUid, prefix=$prefix, address=$matchedAddress")
+                    recoveredCount++
+                } else {
+                    logd(TAG, "No on-chain match for alias: $alias")
+                }
+            } catch (e: Exception) {
+                loge(TAG, "Error processing alias $alias: ${e.message}")
+            }
+        }
+
+        // Cache the aliases hash so we don't re-scan unnecessarily
+        prefs.edit { putInt(KEY_LAST_SCAN_ALIASES_HASH, aliasesHash) }
+
+        logd(TAG, "=== recoverOrphanedKeystoreKeys DONE: recovered $recoveredCount keys ===")
+        return recoveredCount
+    }
+
+    /**
+     * Forces a re-scan on the next call to [recoverOrphanedKeystoreKeys].
+     * Call this whenever a key rotation occurs or an account is removed.
+     */
+    fun invalidateRecoveryCache() {
+        val prefs = Env.getApp().getSharedPreferences(RECOVERY_PREFS, 0)
+        prefs.edit { remove(KEY_LAST_SCAN_ALIASES_HASH) }
+        logd(TAG, "Recovery cache invalidated")
+    }
+
+    // ──────────────────────────────────────────────────────────────────────
+    // Private helpers
+    // ──────────────────────────────────────────────────────────────────────
+
+    /**
+     * Collects all Flow addresses known to AccountManager (from the cache and walletNodes).
+     */
+    private fun collectKnownAddresses(): Set<String> {
+        val addresses = mutableSetOf<String>()
+        try {
+            val accounts = AccountManager.list()
+            for (account in accounts) {
+                // Get address from walletNodes
+                account.firstFlowWalletAddress()?.let { addresses.add(it) }
+
+                // Get address from wallet blockchain data
+                account.wallet?.wallets?.forEach { wallet ->
+                    wallet.blockchain?.forEach { chain ->
+                        if (chain.address.isNotBlank()) {
+                            val addr = if (chain.address.startsWith("0x")) chain.address else "0x${chain.address}"
+                            addresses.add(addr)
+                        }
+                    }
+                }
+            }
+        } catch (e: Exception) {
+            loge(TAG, "Error collecting known addresses: ${e.message}")
+        }
+        logd(TAG, "Collected ${addresses.size} known Flow addresses: $addresses")
+        return addresses
+    }
+
+    /**
+     * For each Flow address, fetches the on-chain account keys and builds a map of
+     * `normalizedPublicKeyHex -> flowAddress` for all non-revoked keys.
+     */
+    private fun buildOnChainKeyMap(addresses: Set<String>): Map<String, String> {
+        val keyMap = mutableMapOf<String, String>()
+
+        for (address in addresses) {
+            try {
+                val account = runBlocking { FlowCadenceApi.getAccount(address) }
+                val keys = account.keys ?: continue
+
+                for (key in keys) {
+                    if (key.revoked) continue
+
+                    // Normalize the public key: strip 0x prefix, strip 04 uncompressed prefix, lowercase
+                    val rawKey = key.publicKey.removePrefix("0x").lowercase()
+                    val stripped = if (rawKey.startsWith("04") && rawKey.length == 130) {
+                        rawKey.substring(2)
+                    } else {
+                        rawKey
+                    }
+
+                    keyMap[rawKey] = address
+                    keyMap[stripped] = address
+
+                    logd(TAG, "On-chain key for $address: index=${key.index}, " +
+                            "revoked=${key.revoked}, pubKey=${rawKey.take(20)}...")
+                }
+            } catch (e: Exception) {
+                loge(TAG, "Failed to fetch on-chain keys for $address: ${e.message}")
+            }
+        }
+
+        logd(TAG, "Built on-chain key map with ${keyMap.size} entries")
+        return keyMap
+    }
+
+    /**
+     * Extracts the EC public key from an Android Keystore alias as a hex string.
+     *
+     * The public key is extracted from the self-signed certificate associated with the
+     * Keystore entry. This works for both hardware-backed and software-backed keys
+     * because the certificate always contains the public key, even when the private key
+     * is non-extractable.
+     *
+     * The returned hex string is the raw uncompressed point (x || y) without the 04 prefix,
+     * matching Flow's on-chain key format.
+     */
+    private fun extractPublicKeyHex(keyStore: KeyStore, alias: String): String? {
+        return try {
+            val cert: Certificate = keyStore.getCertificate(alias) ?: return null
+            val publicKey = cert.publicKey
+
+            if (publicKey !is ECPublicKey) {
+                logd(TAG, "Key $alias is not an EC key: ${publicKey.javaClass.simpleName}")
+                return null
+            }
+
+            val ecPoint: ECPoint = publicKey.w
+            val xBytes = ecPoint.affineX.toByteArray()
+            val yBytes = ecPoint.affineY.toByteArray()
+
+            // Normalize to exactly 32 bytes each (strip leading zero, or left-pad)
+            val x = normalizeCoordinate(xBytes)
+            val y = normalizeCoordinate(yBytes)
+
+            // Return as raw 64-byte hex (no 04 prefix), matching Flow's format
+            (x + y).joinToString("") { "%02x".format(it) }
+        } catch (e: Exception) {
+            loge(TAG, "Failed to extract public key from $alias: ${e.message}")
+            null
+        }
+    }
+
+    /**
+     * Normalizes a BigInteger byte array to exactly 32 bytes:
+     * - Strips leading zero byte (from positive BigInteger encoding)
+     * - Left-pads with zeros if shorter than 32 bytes
+     */
+    private fun normalizeCoordinate(bytes: ByteArray): ByteArray {
+        return when {
+            bytes.size == 32 -> bytes
+            bytes.size == 33 && bytes[0] == 0.toByte() -> bytes.copyOfRange(1, 33)
+            bytes.size > 33 -> {
+                // Take the last 32 bytes
+                bytes.copyOfRange(bytes.size - 32, bytes.size)
+            }
+            bytes.size < 32 -> {
+                val padded = ByteArray(32)
+                System.arraycopy(bytes, 0, padded, 32 - bytes.size, bytes.size)
+                padded
+            }
+            else -> bytes
+        }
+    }
+
+    /**
+     * Checks whether a given prefix is already tracked by any UID in [UserPrefixCacheManager].
+     */
+    private fun isAliasAlreadyMapped(prefix: String): Boolean {
+        return try {
+            val userPrefixes = UserPrefixCacheManager.read() ?: emptyList()
+            userPrefixes.any { it.prefix == prefix }
+        } catch (e: Exception) {
+            false
+        }
+    }
 }