CP-12136: Fix avalanche_getAccountPubKey for ledger (#3650)

Author: ruijialin-avalabs

packages/core-mobile/app/new/features/ledger/hooks/useLedgerWallet.ts

@@ -9,7 +9,6 @@ import {
   LedgerDerivationPathType,
   LedgerKeys,
   LedgerTransportState,
-  PublicKeyInfo,
   WalletCreationOptions,
   WalletUpdateOptions,
   WalletUpdateSolanaOptions
@@ -23,10 +22,8 @@ import Logger from 'utils/Logger'
 import { uuid } from 'utils/uuid'
 import { CoreAccountType } from '@avalabs/types'
 import BiometricsSDK from 'utils/BiometricsSDK'
-import { Curve } from 'utils/publicKeys'
 import { LedgerWalletSecretSchema } from '../utils'
 import { useLedgerWalletMap } from '../store'
-import { DerivationPathKey, getLedgerDerivationPath } from '../consts'
 
 export interface UseLedgerWalletReturn {
   // Connection state
@@ -135,19 +132,9 @@ export function useLedgerWallet(): UseLedgerWalletReturn {
 
         const newWalletId = uuid()
 
-        // Use addresses for display and xpubs for wallet functionality
-        const { addresses, xpubs } = avalancheKeys
-
+        const { addresses, xpubs, publicKeys } = avalancheKeys
         const formattedAddresses = getFormattedAddresses(addresses)
 
-        // Create the public keys array
-        const publicKeysToStore = getPublicKeysForAccount(
-          formattedAddresses,
-          solanaKeys,
-          0 // For wallet creation, we are only adding the first account (index 0)
-        )
-        // Store the Ledger wallet with the specified derivation path type
-        // For BIP44, store xpub in per-account format for future account additions
         await dispatch(
           storeWallet({
             walletId: newWalletId,
@@ -166,7 +153,12 @@ export function useLedgerWallet(): UseLedgerWalletReturn {
                   }
                 }
               }),
-              publicKeys: publicKeysToStore
+              publicKeys: {
+                0: [
+                  ...publicKeys,
+                  ...(solanaKeys?.length > 0 ? [solanaKeys[0]] : [])
+                ].filter(Boolean)
+              }
             }),
             type:
               derivationPathType === LedgerDerivationPathType.BIP44
@@ -235,18 +227,9 @@ export function useLedgerWallet(): UseLedgerWalletReturn {
           throw new Error('Missing Avalanche keys for account creation')
         }
 
-        // Use addresses for display and xpubs for wallet functionality
-        const { addresses, xpubs } = avalancheKeys
-
+        const { addresses, xpubs, publicKeys: newPublicKeys } = avalancheKeys
         const formattedAddresses = getFormattedAddresses(addresses)
 
-        // Create the public keys array
-        const publicKeysToUpdate = getPublicKeysForAccount(
-          formattedAddresses,
-          solanaKeys,
-          accountIndexToUse
-        )
-
         const walletSecretResult = await BiometricsSDK.loadWalletSecret(
           walletId
         )
@@ -292,7 +275,13 @@ export function useLedgerWallet(): UseLedgerWalletReturn {
                   }
                 }
               }),
-              publicKeys: [...publicKeys, ...publicKeysToUpdate] // Append new account public keys to existing array
+              publicKeys: {
+                ...publicKeys,
+                [accountIndexToUse]: [
+                  ...newPublicKeys,
+                  ...(solanaKeys.length > 0 ? [solanaKeys[0]] : [])
+                ].filter(Boolean)
+              }
             })
           })
         ).unwrap()
@@ -372,6 +361,7 @@ export function useLedgerWallet(): UseLedgerWalletReturn {
         }
 
         const { publicKeys, ...baseWalletSecret } = parsedWalletSecret
+        const accountIndex = account.index
 
         // Update the Ledger wallet extended public keys for new account
         await dispatch(
@@ -381,18 +371,13 @@ export function useLedgerWallet(): UseLedgerWalletReturn {
             type: walletType,
             walletSecret: JSON.stringify({
               ...baseWalletSecret,
-              publicKeys: [
+              publicKeys: {
                 ...publicKeys,
-                ...(solanaKeys.length > 0 && solanaKeys[0]?.key
-                  ? [
-                      {
-                        key: solanaKeys[0].key, // Solana addresses don't use 0x prefix
-                        derivationPath: solanaKeys[0].derivationPath, // Use the same path from getSolanaKeys
-                        curve: Curve.ED25519
-                      }
-                    ]
-                  : [])
-              ] // Append new account public keys to existing array
+                [accountIndex]: [
+                  ...(publicKeys[accountIndex] ?? []),
+                  ...(solanaKeys.length > 0 ? [solanaKeys[0]] : [])
+                ].filter(Boolean)
+              }
             })
           })
         ).unwrap()
@@ -454,53 +439,3 @@ const getFormattedAddresses = (address: {
       : address.coreEth
   }
 }
-
-const getPublicKeysForAccount = (
-  address: {
-    evm: string
-    avm: string
-    pvm: string
-    btc: string
-    coreEth: string
-  },
-  solanaKeys: PublicKeyInfo[],
-  accountIndex = 0
-): PublicKeyInfo[] => {
-  return [
-    // Use formatted addresses
-    {
-      key: address.evm, // Use formatted address
-      derivationPath: getLedgerDerivationPath(
-        DerivationPathKey.EVM,
-        accountIndex
-      ),
-      curve: Curve.SECP256K1
-    },
-    {
-      key: address.avm,
-      derivationPath: getLedgerDerivationPath(
-        DerivationPathKey.AVALANCHE,
-        accountIndex
-      ),
-      curve: Curve.SECP256K1
-    },
-    {
-      key: address.pvm,
-      derivationPath: getLedgerDerivationPath(
-        DerivationPathKey.AVALANCHE,
-        accountIndex
-      ),
-      curve: Curve.SECP256K1
-    },
-    // Only include Solana key if it exists
-    ...(solanaKeys.length > 0 && solanaKeys[0]?.key
-      ? [
-          {
-            key: solanaKeys[0].key, // Solana addresses don't use 0x prefix
-            derivationPath: solanaKeys[0].derivationPath, // Use the same path from getSolanaKeys
-            curve: Curve.ED25519
-          }
-        ]
-      : [])
-  ]
-}

packages/core-mobile/app/new/features/ledger/screens/AppConnectionAddAccountScreen.tsx

@@ -117,6 +117,7 @@ export const AppConnectionAddAccountScreen = (): JSX.Element => {
 
   return (
     <AppConnectionScreen
+      selectedDerivationPath={derivationPathType}
       completeStepTitle={`Your Account\nis being set up`}
       handleComplete={handleComplete}
       deviceId={device?.id}

packages/core-mobile/app/new/features/ledger/screens/AppConnectionOnboardingScreen.tsx

@@ -4,7 +4,10 @@ import { useSelector } from 'react-redux'
 import { selectIsDeveloperMode } from 'store/settings/advanced'
 import { Alert } from 'react-native'
 import LedgerService from 'services/ledger/LedgerService'
-import { LedgerKeysByNetwork } from 'services/ledger/types'
+import {
+  LedgerDerivationPathType,
+  LedgerKeysByNetwork
+} from 'services/ledger/types'
 import { useRouter } from 'expo-router'
 import { useLedgerWallet } from '../hooks/useLedgerWallet'
 import { useLedgerSetupContext } from '../contexts/LedgerSetupContext'
@@ -109,6 +112,9 @@ export const AppConnectionOnboardingScreen = (): JSX.Element => {
 
   return (
     <AppConnectionScreen
+      selectedDerivationPath={
+        selectedDerivationPath ?? LedgerDerivationPathType.BIP44
+      }
       completeStepTitle={`Your Ledger wallet\nis being set up`}
       handleComplete={handleComplete}
       deviceId={connectedDeviceId}

packages/core-mobile/app/new/features/ledger/screens/AppConnectionScreen.tsx

@@ -18,12 +18,16 @@ import React, {
 import { Alert, Platform, View } from 'react-native'
 import { useSelector } from 'react-redux'
 import LedgerService from 'services/ledger/LedgerService'
-import { LedgerKeysByNetwork } from 'services/ledger/types'
+import {
+  LedgerDerivationPathType,
+  LedgerKeysByNetwork
+} from 'services/ledger/types'
 import { selectIsSolanaSupportBlocked } from 'store/posthog'
 import { selectIsDeveloperMode } from 'store/settings/advanced'
 import Logger from 'utils/Logger'
 
 export default function AppConnectionScreen({
+  selectedDerivationPath = LedgerDerivationPathType.BIP44,
   completeStepTitle,
   isUpdatingWallet,
   handleComplete,
@@ -32,6 +36,7 @@ export default function AppConnectionScreen({
   disconnectDevice,
   accountIndex
 }: {
+  selectedDerivationPath?: LedgerDerivationPathType
   completeStepTitle: string
   isUpdatingWallet: boolean
   deviceId?: string | null
@@ -149,11 +154,13 @@ export default function AppConnectionScreen({
       // Get keys from service
       const avalancheKeys = await LedgerService.getAvalancheKeys(
         accountIndex,
-        isDeveloperMode
+        isDeveloperMode,
+        selectedDerivationPath
       )
       const oppositeAvalancheKeys = await LedgerService.getAvalancheKeys(
         accountIndex,
-        !isDeveloperMode
+        !isDeveloperMode,
+        selectedDerivationPath
       )
 
       // Update local state
@@ -189,7 +196,13 @@ export default function AppConnectionScreen({
         [{ text: 'OK' }]
       )
     }
-  }, [accountIndex, deviceId, isDeveloperMode, isSolanaSupportBlocked])
+  }, [
+    accountIndex,
+    deviceId,
+    isDeveloperMode,
+    isSolanaSupportBlocked,
+    selectedDerivationPath
+  ])
 
   const handleConnectSolana = useCallback(async () => {
     try {

packages/core-mobile/app/new/features/ledger/utils/index.ts

@@ -4,6 +4,7 @@ import { LedgerAppType, LedgerDerivationPathType } from 'services/ledger/types'
 import { OnDelegationProgress } from 'contexts/DelegationContext'
 import { z } from 'zod'
 import { RpcMethod } from '@avalabs/vm-module-types'
+import { Curve } from 'utils/publicKeys'
 import { ledgerParamsStore, StakingProgressParams } from '../store'
 
 export const showLedgerReviewTransaction = ({
@@ -53,22 +54,42 @@ export const getLedgerAppName = (network?: Network): LedgerAppType => {
     : LedgerAppType.UNKNOWN
 }
 
+const BtcWalletPolicySchema = z.object({
+  hmacHex: z.string(),
+  masterFingerprint: z.string(),
+  xpub: z.string(),
+  name: z.string()
+})
+
+const PublicKeyInfoSchema = z.object({
+  key: z.string(),
+  derivationPath: z.string(),
+  curve: z.enum([Curve.SECP256K1, Curve.ED25519]),
+  btcWalletPolicy: BtcWalletPolicySchema.optional()
+})
+
 export const LedgerWalletSecretSchema = z.looseObject({
   deviceId: z.string(),
   deviceName: z.string(),
-  derivationPathSpec: z.nativeEnum(LedgerDerivationPathType),
-  extendedPublicKeys: z.record(
-    z.string(),
-    z.object({
-      evm: z.string().optional(),
-      avalanche: z.string().optional()
-    })
-  ),
-  publicKeys: z.array(
-    z.object({
-      key: z.string(),
-      derivationPath: z.string(),
-      curve: z.string()
-    })
-  )
+  derivationPathSpec: z.enum([
+    LedgerDerivationPathType.BIP44,
+    LedgerDerivationPathType.LedgerLive
+  ]),
+  extendedPublicKeys: z
+    .record(
+      z.string(),
+      z.object({
+        evm: z.string().optional(),
+        avalanche: z.string().optional()
+      })
+    )
+    .optional(),
+  publicKeys: z
+    .record(z.string(), z.array(PublicKeyInfoSchema))
+    .transform(
+      record =>
+        Object.fromEntries(
+          Object.entries(record).map(([k, v]) => [Number(k), v])
+        ) as Record<number, z.infer<typeof PublicKeyInfoSchema>[]>
+    )
 })