Skip to content

Commit

Permalink
feat(*): support customising prompt for KMD password and Mnemonic phrase
Browse files Browse the repository at this point in the history
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
  • Loading branch information
PatrickDinh committed Dec 3, 2024
1 parent 049d831 commit 8bbca2d
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 18 deletions.
27 changes: 26 additions & 1 deletion packages/use-wallet/src/__tests__/wallets/kmd.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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)
})
})
})
26 changes: 26 additions & 0 deletions packages/use-wallet/src/__tests__/wallets/mnemonic.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
})
})
})
19 changes: 10 additions & 9 deletions packages/use-wallet/src/wallets/kmd.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,11 @@ interface KmdConstructor {
baseServer?: string
port?: string | number
headers?: Record<string, string>
promptForPassword: () => Promise<string>
}

export type KmdOptions = Partial<Pick<KmdConstructor, 'token'>> &
Omit<KmdConstructor, 'token'> & {
export type KmdOptions = Partial<Pick<KmdConstructor, 'token' | 'promptForPassword'>> &
Omit<KmdConstructor, 'token' | 'promptForPassword'> & {
wallet?: string
}

Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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())

Expand Down Expand Up @@ -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,
Expand All @@ -301,11 +302,11 @@ export class KmdWallet extends BaseWallet {
this.logger.debug('Token released successfully')
}

private getPassword(): string {
private async getPassword(): Promise<string> {
if (this.password !== null) {
return this.password
}
const password = prompt('KMD password') || ''
const password = await this.options.promptForPassword()
this.password = password
return password
}
Expand Down
21 changes: 13 additions & 8 deletions packages/use-wallet/src/wallets/mnemonic.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string | null>
}

export type MnemonicOptions = Partial<Pick<MnemonicConstructor, 'promptForMnemonic'>> &
Omit<MnemonicConstructor, 'promptForMnemonic'>


export const LOCAL_STORAGE_MNEMONIC_KEY = `${LOCAL_STORAGE_KEY}_mnemonic`

const ICON = `data:image/svg+xml;base64,${btoa(`
Expand All @@ -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<State>

Expand All @@ -36,8 +41,8 @@ export class MnemonicWallet extends BaseWallet {
}: WalletConstructor<WalletId.MNEMONIC>) {
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

Expand Down Expand Up @@ -80,10 +85,10 @@ export class MnemonicWallet extends BaseWallet {
}
}

private initializeAccount(): algosdk.Account {
private async initializeAccount(): Promise<algosdk.Account> {
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')
Expand All @@ -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`,
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 8bbca2d

Please sign in to comment.