From 877b870074d703806d3c548c60e16527e674e227 Mon Sep 17 00:00:00 2001 From: Pablo Maldonado Date: Wed, 18 Sep 2024 10:45:51 +0200 Subject: [PATCH] feat(walletManager): add shouldRecordUsage option and tests Signed-off-by: Pablo Maldonado --- src/index.ts | 3 ++- src/lib/bundleUtils.ts | 18 +++++++++++----- src/lib/walletManager.ts | 44 +++++++++++++++++++++++----------------- test/walletManager.ts | 36 ++++++++++++++++++++++++++++++++ 4 files changed, 76 insertions(+), 25 deletions(-) diff --git a/src/index.ts b/src/index.ts index c6e79c8..0d36178 100644 --- a/src/index.ts +++ b/src/index.ts @@ -145,7 +145,8 @@ app.post("/", async (req, res, next) => { Logger.debug(req.transactionId, "Finding unlock that does not revert the bundle..."); - const unlock = await findUnlock(flashbotsBundleProvider, backrunTxs, targetBlock, req); + // FindUnlock with recordWalletUsage set to true + const unlock = await findUnlock(flashbotsBundleProvider, backrunTxs, targetBlock, req, true); if (!unlock) { Logger.debug(req.transactionId, "No valid unlock found!"); await handleUnsupportedRequest(req, res, "No valid unlock found"); // Pass through if no unlock is found. diff --git a/src/lib/bundleUtils.ts b/src/lib/bundleUtils.ts index 63dcb68..ac85fc5 100644 --- a/src/lib/bundleUtils.ts +++ b/src/lib/bundleUtils.ts @@ -46,6 +46,7 @@ export const createUnlockLatestValueTx = async ( }; // Prepare unlockLatestValue transaction for a given Oval instance and simulate the bundle with the unlock transaction prepended if simulate is true. +// If shouldRecordWalletUsage is true, the wallet usage will be recorded in the WalletManager. export const prepareUnlockTransaction = async ( flashbotsBundleProvider: FlashbotsBundleProvider, backrunTxs: string[], @@ -53,10 +54,11 @@ export const prepareUnlockTransaction = async ( ovalAddress: string, req: express.Request, simulate = true, + shouldRecordWalletUsage = true, ) => { 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, shouldRecordWalletUsage); const isSharedWallet = isOvalSharedUnlockerKey(unlockerWallet.address); // Encode the unlockLatestValue function call depending on whether the unlocker is a shared wallet or not. @@ -75,11 +77,11 @@ export const prepareUnlockTransaction = async ( target, ); - if (!simulate) return { ovalAddress, unlockTxHash, signedUnlockTx }; + if (!simulate) return { ovalAddress, unlockTxHash, signedUnlockTx, unlockerWallet }; const simulationResponse = await flashbotsBundleProvider.simulate([signedUnlockTx, ...backrunTxs], targetBlock); - return { ovalAddress, unlockTxHash, signedUnlockTx, simulationResponse }; + return { ovalAddress, unlockTxHash, signedUnlockTx, simulationResponse, unlockerWallet }; }; export const getUnlockBundlesFromOvalAddresses = async ( @@ -99,7 +101,7 @@ export const getUnlockBundlesFromOvalAddresses = async ( targetBlock, ovalAddress, req, - false, + false, // Do not simulate ); // Construct the inner bundle with call to Oval to unlock the latest value. @@ -123,10 +125,13 @@ export const findUnlock = async ( backrunTxs: string[], targetBlock: number, req: express.Request, + updateWalletUsage = false ) => { const unlocks = await Promise.all( getOvalAddresses().map(async (ovalAddress) => - prepareUnlockTransaction(flashbotsBundleProvider, backrunTxs, targetBlock, ovalAddress, req), + // Do not record wallet usage here as we will record it in the loop below once we find a valid unlock. + // See shouldRecordWalletUsage parameter in WalletManager.getWallet. + prepareUnlockTransaction(flashbotsBundleProvider, backrunTxs, targetBlock, ovalAddress, req, true, false), ), ); @@ -137,6 +142,9 @@ export const findUnlock = async ( !("error" in unlock.simulationResponse) && !unlock.simulationResponse.firstRevert ) { + if (updateWalletUsage) { + WalletManager.getInstance().updateWalletUsage(unlock.ovalAddress, unlock.unlockerWallet, targetBlock); + } return { // Spread in order to preserve inferred SimulationResponseSuccess type. ...unlock, diff --git a/src/lib/walletManager.ts b/src/lib/walletManager.ts index 2cfb197..bcf856f 100644 --- a/src/lib/walletManager.ts +++ b/src/lib/walletManager.ts @@ -49,50 +49,57 @@ export class WalletManager { } // Get a wallet for a given address - public getWallet(address: string, targetBlock: number, transactionId: string): Wallet { + // If shouldRecordUsage is true, the wallet usage will be recorded in the WalletManager. + public getWallet(address: string, targetBlock: number, transactionId: string, shouldRecordUsage = true): 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, shouldRecordUsage); } 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, shouldRecordUsage = true): 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 (shouldRecordUsage) { 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,12 +197,11 @@ 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; } // Update the usage statistics for a wallet - private async updateWalletUsage(ovalInstance: string, wallet: Wallet, targetBlock: number): Promise { + public async updateWalletUsage(ovalInstance: string, wallet: Wallet, targetBlock: number): Promise { const walletPubKey = await wallet.getAddress(); const instanceUsage = this.sharedWalletUsage.get(walletPubKey) || new Map(); const existingRecord = instanceUsage.get(targetBlock); diff --git a/test/walletManager.ts b/test/walletManager.ts index 1d17c0b..20eb41d 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 shouldRecordUsage 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 = walletManager.getWallet(ovalInstance, targetBlock, "transactionId", false); + + const sharedWalletUsage = walletManager['sharedWalletUsage'].get(await wallet.getAddress()); + expect(sharedWalletUsage).to.be.undefined; + }); + + it('should record usage when shouldRecordUsage 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 = 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 = [