From 8bbca2d8093829569b81d3b61ba1f94584f67285 Mon Sep 17 00:00:00 2001 From: Hoang Dinh Date: Tue, 3 Dec 2024 16:20:55 +1000 Subject: [PATCH] feat(*): support customising prompt for KMD password and Mnemonic phrase There are scenarios where customsing window.prompt are needed: - window.prompt isn't supported, for example, on macOS WebKit - a better UI to display the message To implement this change, I have: - add options to customise the prompt to KMD and Mnemonic wallets - when the option isn't set, fallback to the default window.prompt --- .../src/__tests__/wallets/kmd.test.ts | 27 ++++++++++++++++++- .../src/__tests__/wallets/mnemonic.test.ts | 26 ++++++++++++++++++ packages/use-wallet/src/wallets/kmd.ts | 19 ++++++------- packages/use-wallet/src/wallets/mnemonic.ts | 21 +++++++++------ 4 files changed, 75 insertions(+), 18 deletions(-) diff --git a/packages/use-wallet/src/__tests__/wallets/kmd.test.ts b/packages/use-wallet/src/__tests__/wallets/kmd.test.ts index b73b4578..bb6564d8 100644 --- a/packages/use-wallet/src/__tests__/wallets/kmd.test.ts +++ b/packages/use-wallet/src/__tests__/wallets/kmd.test.ts @@ -428,7 +428,7 @@ describe('KmdWallet', () => { it('should handle null from cancelled prompt', async () => { // Mock prompt to return null (user cancelled) - global.prompt = vi.fn().mockReturnValue(null) + global.prompt = vi.fn().mockReturnValue(null) mockKmd.listKeys.mockResolvedValueOnce({ addresses: [account1.address] }) await wallet.connect() @@ -437,4 +437,29 @@ describe('KmdWallet', () => { expect(mockKmd.initWalletHandle).toHaveBeenCalledWith(mockWallet.id, '') }) }) + + describe('custom prompt for password', () => { + const customPassword = 'customPassword' + + beforeEach(() => { + wallet = new KmdWallet({ + id: WalletId.KMD, + metadata: {}, + getAlgodClient: {} as any, + store, + subscribe: vi.fn(), + options: { + promptForPassword: () => Promise.resolve(customPassword) + } + }) + }) + + it('should return password from custom prompt', async () => { + mockKmd.listKeys.mockResolvedValueOnce({ addresses: [account1.address] }) + await wallet.connect() + + expect(global.prompt).toHaveBeenCalledTimes(0) + expect(mockKmd.initWalletHandle).toHaveBeenCalledWith(mockWallet.id, customPassword) + }) + }) }) diff --git a/packages/use-wallet/src/__tests__/wallets/mnemonic.test.ts b/packages/use-wallet/src/__tests__/wallets/mnemonic.test.ts index b53c5a30..6c096605 100644 --- a/packages/use-wallet/src/__tests__/wallets/mnemonic.test.ts +++ b/packages/use-wallet/src/__tests__/wallets/mnemonic.test.ts @@ -361,4 +361,30 @@ describe('MnemonicWallet', () => { }) }) }) + + describe('custom prompt for mnemonic', () => { + const MOCK_ACCOUNT_MNEMONIC = 'just aim reveal time update elegant column reunion lazy ritual room unusual notice camera forward couple quantum gym laundry absurd drill pyramid tip able outdoor' + + beforeEach(() => { + wallet = new MnemonicWallet({ + id: WalletId.MNEMONIC, + options: { promptForMnemonic: () => Promise.resolve(MOCK_ACCOUNT_MNEMONIC), persistToStorage: true }, + metadata: {}, + getAlgodClient: {} as any, + store, + subscribe: vi.fn() + }) + }) + + it('should save mnemonic into storage', async () => { + const storageSetItemSpy = vi.spyOn(StorageAdapter, 'setItem') + // Simulate no mnemonic in storage + vi.mocked(StorageAdapter.getItem).mockImplementation(() => null) + + await wallet.connect() + + expect(global.prompt).toHaveBeenCalledTimes(0) + expect(storageSetItemSpy).toHaveBeenCalledWith(LOCAL_STORAGE_MNEMONIC_KEY, MOCK_ACCOUNT_MNEMONIC) + }) + }) }) diff --git a/packages/use-wallet/src/wallets/kmd.ts b/packages/use-wallet/src/wallets/kmd.ts index 4f29b43e..4a20adc7 100644 --- a/packages/use-wallet/src/wallets/kmd.ts +++ b/packages/use-wallet/src/wallets/kmd.ts @@ -10,10 +10,11 @@ interface KmdConstructor { baseServer?: string port?: string | number headers?: Record + promptForPassword: () => Promise } -export type KmdOptions = Partial> & - Omit & { +export type KmdOptions = Partial> & + Omit & { wallet?: string } @@ -78,12 +79,12 @@ export class KmdWallet extends BaseWallet { token = 'a'.repeat(64), baseServer = 'http://127.0.0.1', port = 4002, - wallet = 'unencrypted-default-wallet' + wallet = 'unencrypted-default-wallet', + promptForPassword = () => Promise.resolve(prompt('KMD password') || '') } = options || {} - this.options = { token, baseServer, port } + this.options = { token, baseServer, port, promptForPassword } this.walletName = wallet - this.store = store } @@ -238,7 +239,7 @@ export class KmdWallet extends BaseWallet { // Get token and password const token = await this.fetchToken() - const password = this.getPassword() + const password = await this.getPassword() const client = this.client || (await this.initializeClient()) @@ -284,7 +285,7 @@ export class KmdWallet extends BaseWallet { const client = this.client || (await this.initializeClient()) const walletId = this.walletId || (await this.fetchWalletId()) - const password = this.getPassword() + const password = await this.getPassword() const { wallet_handle_token }: InitWalletHandleResponse = await client.initWalletHandle( walletId, @@ -301,11 +302,11 @@ export class KmdWallet extends BaseWallet { this.logger.debug('Token released successfully') } - private getPassword(): string { + private async getPassword(): Promise { if (this.password !== null) { return this.password } - const password = prompt('KMD password') || '' + const password = await this.options.promptForPassword() this.password = password return password } diff --git a/packages/use-wallet/src/wallets/mnemonic.ts b/packages/use-wallet/src/wallets/mnemonic.ts index a729d1cf..0d62b7fc 100644 --- a/packages/use-wallet/src/wallets/mnemonic.ts +++ b/packages/use-wallet/src/wallets/mnemonic.ts @@ -7,10 +7,15 @@ import { BaseWallet } from 'src/wallets/base' import type { Store } from '@tanstack/store' import type { WalletAccount, WalletConstructor, WalletId } from 'src/wallets/types' -export type MnemonicOptions = { +interface MnemonicConstructor { persistToStorage?: boolean + promptForMnemonic: () => Promise } +export type MnemonicOptions = Partial> & + Omit + + export const LOCAL_STORAGE_MNEMONIC_KEY = `${LOCAL_STORAGE_KEY}_mnemonic` const ICON = `data:image/svg+xml;base64,${btoa(` @@ -22,7 +27,7 @@ const ICON = `data:image/svg+xml;base64,${btoa(` export class MnemonicWallet extends BaseWallet { private account: algosdk.Account | null = null - private options: MnemonicOptions + private options: MnemonicConstructor protected store: Store @@ -36,8 +41,8 @@ export class MnemonicWallet extends BaseWallet { }: WalletConstructor) { super({ id, metadata, getAlgodClient, store, subscribe }) - const { persistToStorage = false } = options || {} - this.options = { persistToStorage } + const { persistToStorage = false, promptForMnemonic = () => Promise.resolve(prompt('Enter 25-word mnemonic passphrase:')) } = options || {} + this.options = { persistToStorage, promptForMnemonic } this.store = store @@ -80,10 +85,10 @@ export class MnemonicWallet extends BaseWallet { } } - private initializeAccount(): algosdk.Account { + private async initializeAccount(): Promise { let mnemonic = this.loadMnemonicFromStorage() if (!mnemonic) { - mnemonic = prompt('Enter 25-word mnemonic passphrase:') + mnemonic = await this.options.promptForMnemonic() if (!mnemonic) { this.account = null this.logger.error('No mnemonic provided') @@ -106,7 +111,7 @@ export class MnemonicWallet extends BaseWallet { this.checkMainnet() this.logger.info('Connecting...') - const account = this.initializeAccount() + const account = await this.initializeAccount() const walletAccount = { name: `${this.metadata.name} Account`, @@ -153,7 +158,7 @@ export class MnemonicWallet extends BaseWallet { // If persisting to storage is enabled, then resume session if (this.options.persistToStorage) { try { - this.initializeAccount() + await this.initializeAccount() this.logger.info('Session resumed successfully') } catch (error: any) { this.logger.error('Error resuming session:', error.message)