diff --git a/src/lib/bundleUtils.ts b/src/lib/bundleUtils.ts index 63dcb68..8d70bc6 100644 --- a/src/lib/bundleUtils.ts +++ b/src/lib/bundleUtils.ts @@ -56,7 +56,7 @@ export const prepareUnlockTransaction = async ( ) => { const provider = getProvider(); const [baseFee, network] = await Promise.all([getBaseFee(provider, req), provider.getNetwork()]); - const unlockerWallet = WalletManager.getInstance().getWallet(ovalAddress, targetBlock, req.transactionId); + const unlockerWallet = WalletManager.getInstance().getWallet(ovalAddress, targetBlock, req.transactionId, simulate); const isSharedWallet = isOvalSharedUnlockerKey(unlockerWallet.address); // Encode the unlockLatestValue function call depending on whether the unlocker is a shared wallet or not. diff --git a/src/lib/walletManager.ts b/src/lib/walletManager.ts index 2cfb197..da76384 100644 --- a/src/lib/walletManager.ts +++ b/src/lib/walletManager.ts @@ -49,50 +49,56 @@ export class WalletManager { } // Get a wallet for a given address - public getWallet(address: string, targetBlock: number, transactionId: string): Wallet { + public getWallet(address: string, targetBlock: number, transactionId: string, skipUsageRecording = false): Wallet { if (!this.provider) { throw new Error("Provider is not initialized"); } const checkSummedAddress = getAddress(address); const wallet = this.wallets[checkSummedAddress]; if (!wallet) { - return this.getSharedWallet(address, targetBlock, transactionId); + return this.getSharedWallet(address, targetBlock, transactionId, skipUsageRecording); } return wallet.connect(this.provider); } // Get a shared wallet for a given Oval instance and target block - private getSharedWallet(ovalInstance: string, targetBlock: number, transactionId: string): Wallet { + private getSharedWallet(ovalInstance: string, targetBlock: number, transactionId: string, skipUsageRecording = false): Wallet { if (!this.provider) { throw new Error("Provider is not initialized"); } + if (!this.ovalDiscovery.isOval(ovalInstance)) { throw new Error(`Oval instance ${ovalInstance} is not found`); } - let selectedWallet: Wallet | undefined; - - // Check if a wallet has already been assigned to this Oval instance - for (const [walletPubKey, instanceUsage] of this.sharedWalletUsage.entries()) { - for (const [_, record] of instanceUsage.entries()) { - if (record.ovalInstances && record.ovalInstances.has(ovalInstance)) { - selectedWallet = this.sharedWallets.get(walletPubKey)!.connect(this.provider!); - } - } - } + let selectedWallet = this.findAssignedWallet(ovalInstance); // If no wallet has been assigned, find the least used wallet if (!selectedWallet) { selectedWallet = this.findLeastUsedWallet(transactionId); } - // Update the usage of the selected wallet - if (selectedWallet) { + if (!selectedWallet) { + throw new Error(`No available shared wallets for Oval instance ${ovalInstance} at block ${targetBlock}`); + } + + // Update the usage of the selected wallet if recording is enabled + if (!skipUsageRecording) { this.updateWalletUsage(ovalInstance, selectedWallet, targetBlock); - return selectedWallet.connect(this.provider); } - throw new Error(`No available shared wallets for Oval instance ${ovalInstance} at block ${targetBlock}`); + return selectedWallet.connect(this.provider); + } + + private findAssignedWallet(ovalInstance: string): Wallet | undefined { + for (const [walletPubKey, instanceUsage] of this.sharedWalletUsage.entries()) { + for (const record of instanceUsage.values()) { + if (record.ovalInstances?.has(ovalInstance)) { + return this.sharedWallets.get(walletPubKey)?.connect(this.provider!); + } + } + } + return undefined; } public isOvalSharedUnlocker(unlockerPublicKey: string): boolean { @@ -190,7 +196,6 @@ export class WalletManager { if (minInstances !== Infinity && minInstances !== 0) { Logger.error(transactionId, `Public key ${selectedWallet?.address} is reused in multiple Oval instances because no free wallets are available.`); } - return selectedWallet; } diff --git a/test/walletManager.ts b/test/walletManager.ts index 1d17c0b..52af9ac 100644 --- a/test/walletManager.ts +++ b/test/walletManager.ts @@ -141,6 +141,42 @@ describe('WalletManager Tests', () => { expect(errorSpy.args[0][1]).to.include(`Public key ${wallet2.address} is reused in multiple Oval instances because no free wallets are available.`); }); + it('should not record usage when skipUsageRecording is true', async () => { + const unlockerRandom = getRandomAddressAndKey(); + const sharedConfigs: OvalConfigsShared = [ + { unlockerKey: unlockerRandom.privateKey }, + ]; + const walletManager = WalletManager.getInstance(); + await walletManager.initialize(mockProvider, {}, sharedConfigs); + + const ovalInstance = getRandomAddressAndKey().address; + const targetBlock = 123; + + const wallet = walletManager.getWallet(ovalInstance, targetBlock, "transactionId", true); + + const sharedWalletUsage = walletManager['sharedWalletUsage'].get(await wallet.getAddress()); + expect(sharedWalletUsage).to.be.undefined; + }); + + it('should record usage when skipUsageRecording is false', async () => { + const unlockerRandom = getRandomAddressAndKey(); + const sharedConfigs: OvalConfigsShared = [ + { unlockerKey: unlockerRandom.privateKey }, + ]; + const walletManager = WalletManager.getInstance(); + await walletManager.initialize(mockProvider, {}, sharedConfigs); + + const ovalInstance = getRandomAddressAndKey().address; + const targetBlock = 123; + + const wallet = await walletManager.getWallet(ovalInstance, targetBlock, "transactionId"); + + const sharedWalletUsage = walletManager['sharedWalletUsage'].get(await wallet.getAddress()); + expect(sharedWalletUsage).to.not.be.undefined; + expect(sharedWalletUsage?.get(targetBlock)?.count).to.equal(1); + expect(sharedWalletUsage?.get(targetBlock)?.ovalInstances.has(ovalInstance)).to.be.true; + }); + it('should cleanup old records correctly', async () => { const unlockerRandom = getRandomAddressAndKey(); const sharedConfigs: OvalConfigsShared = [