diff --git a/README.md b/README.md index a513e70..0bf9d6b 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ yarn add @d3or/slotseek ## TODO -- [ ] Add caching options to reduce the number of RPC calls and reduce the time it takes to find the same slot again +- [X] Add caching options to reduce the number of RPC calls and reduce the time it takes to find the same slot again ## Example of overriding a users balance via eth_call diff --git a/src/approval.ts b/src/approval.ts index 5d8d05a..7eb43ad 100644 --- a/src/approval.ts +++ b/src/approval.ts @@ -1,4 +1,5 @@ import { ethers } from "ethers"; +import { approvalCache } from "./cache"; /** * Generate mock approval data for a given ERC20 token @@ -114,6 +115,27 @@ export const getErc20ApprovalStorageSlot = async ( slotHash: string; isVyper: boolean; }> => { + // check the cache + const cachedValue = approvalCache.get(erc20Address.toLowerCase()); + if (cachedValue) { + if (cachedValue.isVyper) { + const { vyperSlotHash } = calculateApprovalVyperStorageSlot(ownerAddress, spenderAddress, cachedValue.slot) + return { + slot: ethers.BigNumber.from(cachedValue.slot).toHexString(), + slotHash: vyperSlotHash, + isVyper: true, + }; + + } else { + const { slotHash } = calculateApprovalSolidityStorageSlot(ownerAddress, spenderAddress, cachedValue.slot) + return { + slot: ethers.BigNumber.from(cachedValue.slot).toHexString(), + slotHash: slotHash, + isVyper: false, + } + } + } + // Get the approval for the spender, that we can use to find the slot let approval = await getErc20Approval( provider, @@ -124,50 +146,36 @@ export const getErc20ApprovalStorageSlot = async ( if (approval.gt(0)) { for (let i = 0; i < maxSlots; i++) { - // Calculate the slot hash, using the owner address and the slot index - const slot = ethers.utils.keccak256( - ethers.utils.defaultAbiCoder.encode( - ["address", "uint256"], - [ownerAddress, i] - ) - ); - // Calculate the storage slot, using the spender address and the slot hash - const storageSlot = ethers.utils.keccak256( - ethers.utils.defaultAbiCoder.encode( - ["address", "bytes32"], - [spenderAddress, slot] - ) - ); + const { storageSlot, slotHash } = calculateApprovalSolidityStorageSlot(ownerAddress, spenderAddress, i) // Get the value at the storage slot const storageValue = await provider.getStorageAt(erc20Address, storageSlot); // If the value at the storage slot is equal to the approval, return the slot as we have found the correct slot for approvals if (ethers.BigNumber.from(storageValue).eq(approval)) { + approvalCache.set(erc20Address.toLowerCase(), { + slot: i, + isVyper: false, + ts: Date.now() + }); return { slot: ethers.BigNumber.from(i).toHexString(), - slotHash: slot, + slotHash: slotHash, isVyper: false, }; } - // check via vyper storage layout, which uses keccak256(abi.encode(slot, address(this))) instead of keccak256(abi.encode(address(this), slot)) - const vyperSlotHash = ethers.utils.keccak256( - ethers.utils.defaultAbiCoder.encode( - ["uint256", "address"], - [i, ownerAddress] - ) - ); - - const vyperStorageSlot = ethers.utils.keccak256( - ethers.utils.defaultAbiCoder.encode( - ["bytes32", "address"], - [vyperSlotHash, spenderAddress] - ) - ); + const { vyperStorageSlot, vyperSlotHash } = calculateApprovalVyperStorageSlot(ownerAddress, spenderAddress, i) const vyperStorageValue = await provider.getStorageAt( erc20Address, vyperStorageSlot ); + if (ethers.BigNumber.from(vyperStorageValue).eq(approval)) { + approvalCache.set(erc20Address.toLowerCase(), { + slot: i, + isVyper: false, + ts: Date.now() + }); + return { slot: ethers.BigNumber.from(i).toHexString(), slotHash: vyperSlotHash, @@ -179,73 +187,93 @@ export const getErc20ApprovalStorageSlot = async ( throw new Error("Approval does not exist"); } - if (useFallbackSlot) { // if useFallBackSlot = true, then we are just going to assume the slot is at the slot which is most common for erc20 tokens. for approvals, this is slot #10 const fallbackSlot = 10; - // check if contract is solidity/vyper, so we know which storage slot generation method to use - // TODO: add this above check, currently not sure of how to programatically check if a contract is a vyper contract. - const isVyper = false; - if (isVyper) { - const vyperSlotHash = ethers.utils.keccak256( - ethers.utils.defaultAbiCoder.encode( - ["uint256", "address"], - [fallbackSlot, ownerAddress] - ) - ); + // check solidity, then check vyper. + // (dont have an easy way to check if a contract is solidity/vyper) + const { storageSlot, slotHash } = calculateApprovalSolidityStorageSlot(ownerAddress, spenderAddress, fallbackSlot) + // Get the value at the storage slot + const storageValue = await provider.getStorageAt(erc20Address, storageSlot); + // If the value at the storage slot is equal to the approval, return the slot as we have found the correct slot for approvals + if (ethers.BigNumber.from(storageValue).eq(approval)) { + approvalCache.set(erc20Address.toLowerCase(), { + slot: fallbackSlot, + isVyper: false, + ts: Date.now() + }); - const vyperStorageSlot = ethers.utils.keccak256( - ethers.utils.defaultAbiCoder.encode( - ["bytes32", "address"], - [vyperSlotHash, spenderAddress] - ) - ); - const vyperStorageValue = await provider.getStorageAt( - erc20Address, - vyperStorageSlot - ); - if (ethers.BigNumber.from(vyperStorageValue).eq(approval)) { - return { - slot: ethers.BigNumber.from(fallbackSlot).toHexString(), - slotHash: vyperSlotHash, - isVyper: true, - }; - } - } else { - - const slot = ethers.utils.keccak256( - ethers.utils.defaultAbiCoder.encode( - ["address", "uint256"], - [ownerAddress, fallbackSlot] - ) - ); - // Calculate the storage slot, using the spender address and the slot hash - const storageSlot = ethers.utils.keccak256( - ethers.utils.defaultAbiCoder.encode( - ["address", "bytes32"], - [spenderAddress, slot] - ) - ); - // Get the value at the storage slot - const storageValue = await provider.getStorageAt(erc20Address, storageSlot); - // If the value at the storage slot is equal to the approval, return the slot as we have found the correct slot for approvals - if (ethers.BigNumber.from(storageValue).eq(approval)) { - return { - slot: ethers.BigNumber.from(fallbackSlot).toHexString(), - slotHash: slot, - isVyper: false, - }; + return { + slot: ethers.BigNumber.from(fallbackSlot).toHexString(), + slotHash: slotHash, + isVyper: false, + }; + } - } + // check vyper + const { vyperStorageSlot, vyperSlotHash } = calculateApprovalVyperStorageSlot(ownerAddress, spenderAddress, fallbackSlot) + const vyperStorageValue = await provider.getStorageAt( + erc20Address, + vyperStorageSlot + ); + if (ethers.BigNumber.from(vyperStorageValue).eq(approval)) { + approvalCache.set(erc20Address.toLowerCase(), { + slot: fallbackSlot, + isVyper: true, + ts: Date.now() + }); + return { + slot: ethers.BigNumber.from(fallbackSlot).toHexString(), + slotHash: vyperSlotHash, + isVyper: true, + }; } - } throw new Error("Unable to find approval slot"); }; +// Generates approval solidity storage slot data +const calculateApprovalSolidityStorageSlot = (ownerAddress: string, spenderAddress: string, slotNumber: number) => { + + // Calculate the slot hash, using the owner address and the slot index + const slotHash = ethers.utils.keccak256( + ethers.utils.defaultAbiCoder.encode( + ["address", "uint256"], + [ownerAddress, slotNumber] + ) + ); + // Calculate the storage slot, using the spender address and the slot hash + const storageSlot = ethers.utils.keccak256( + ethers.utils.defaultAbiCoder.encode( + ["address", "bytes32"], + [spenderAddress, slotHash] + ) + ); + return { storageSlot, slotHash } +} + +// Generates approval vyper storage slot data +const calculateApprovalVyperStorageSlot = (ownerAddress: string, spenderAddress: string, slotNumber: number) => { + // create via vyper storage layout, which uses keccak256(abi.encode(slot, address(this))) instead of keccak256(abi.encode(address(this), slot)) + const vyperSlotHash = ethers.utils.keccak256( + ethers.utils.defaultAbiCoder.encode( + ["uint256", "address"], + [slotNumber, ownerAddress] + ) + ); + + const vyperStorageSlot = ethers.utils.keccak256( + ethers.utils.defaultAbiCoder.encode( + ["bytes32", "address"], + [vyperSlotHash, spenderAddress] + ) + ); + + return { vyperStorageSlot, vyperSlotHash } +} /** * Get the approval for a given ERC20 token * @param provider - The JsonRpcProvider instance diff --git a/src/balance.ts b/src/balance.ts index 0ad5624..3a75a8a 100644 --- a/src/balance.ts +++ b/src/balance.ts @@ -1,4 +1,5 @@ import { ethers } from "ethers"; +import { balanceCache } from "./cache"; /** * Generate mock data for a given ERC20 token balance @@ -87,12 +88,37 @@ export const getErc20BalanceStorageSlot = async ( provider: ethers.providers.JsonRpcProvider, erc20Address: string, holderAddress: string, - maxSlots = 100 + maxSlots = 30 ): Promise<{ slot: string; balance: ethers.BigNumber; isVyper: boolean; }> => { + // check the cache + const cachedValue = balanceCache.get(erc20Address.toLowerCase()); + if (cachedValue) { + if (cachedValue.isVyper) { + const { vyperSlotHash } = calculateBalanceVyperStorageSlot(holderAddress, cachedValue.slot) + const vyperBalance = await provider.getStorageAt( + erc20Address, + vyperSlotHash + ); + return { + slot: ethers.BigNumber.from(cachedValue.slot).toHexString(), + balance: ethers.BigNumber.from(vyperBalance), + isVyper: true, + }; + } else { + const { slotHash } = calculateBalanceSolidityStorageSlot(holderAddress, cachedValue.slot); + const balance = await provider.getStorageAt(erc20Address, slotHash); + return { + slot: ethers.BigNumber.from(cachedValue.slot).toHexString(), + balance: ethers.BigNumber.from(balance), + isVyper: false, + } + } + } + // Get the balance of the holder, that we can use to find the slot const userBalance = await getErc20Balance( provider, @@ -107,12 +133,16 @@ export const getErc20BalanceStorageSlot = async ( // For each slot, we compute the storage slot key [holderAddress, slot index] and get the value at that storage slot // If the value at the storage slot is equal to the balance, return the slot as we have found the correct slot for balances for (let i = 0; i < maxSlots; i++) { - const slot = ethers.utils.solidityKeccak256( - ["uint256", "uint256"], - [holderAddress, i] - ); - const balance = await provider.getStorageAt(erc20Address, slot); + const { slotHash } = calculateBalanceSolidityStorageSlot(holderAddress, i); + const balance = await provider.getStorageAt(erc20Address, slotHash); + if (ethers.BigNumber.from(balance).eq(userBalance)) { + balanceCache.set(erc20Address.toLowerCase(), { + slot: i, + isVyper: false, + ts: Date.now() + }) + return { slot: ethers.BigNumber.from(i).toHexString(), balance: ethers.BigNumber.from(balance), @@ -120,18 +150,19 @@ export const getErc20BalanceStorageSlot = async ( }; } - // check via vyper storage layout, which uses keccak256(abi.encode(slot, address(this))) instead of keccak256(abi.encode(address(this), slot)) - const vyperSlotHash = ethers.utils.keccak256( - ethers.utils.defaultAbiCoder.encode( - ["uint256", "address"], - [i, holderAddress] - ) - ); + const { vyperSlotHash } = calculateBalanceVyperStorageSlot(holderAddress, i) const vyperBalance = await provider.getStorageAt( erc20Address, vyperSlotHash ); + if (ethers.BigNumber.from(vyperBalance).eq(userBalance)) { + balanceCache.set(erc20Address.toLowerCase(), { + slot: i, + isVyper: true, + ts: Date.now() + }) + return { slot: ethers.BigNumber.from(i).toHexString(), balance: ethers.BigNumber.from(vyperBalance), @@ -142,6 +173,26 @@ export const getErc20BalanceStorageSlot = async ( throw new Error("Unable to find balance slot"); }; + +const calculateBalanceSolidityStorageSlot = (holderAddress: string, slotNumber: number) => { + const slotHash = ethers.utils.solidityKeccak256( + ["uint256", "uint256"], + [holderAddress, slotNumber] + ); + return { slotHash } +} + +const calculateBalanceVyperStorageSlot = (holderAddress: string, slotNumber: number) => { + // create hash via vyper storage layout, which uses keccak256(abi.encode(slot, address(this))) instead of keccak256(abi.encode(address(this), slot)) + const vyperSlotHash = ethers.utils.keccak256( + ethers.utils.defaultAbiCoder.encode( + ["uint256", "address"], + [slotNumber, holderAddress] + ) + ); + return { vyperSlotHash } +} + /** * Get the balance of a given address for a given ERC20 token * @param provider - The JsonRpcProvider instance diff --git a/src/cache.ts b/src/cache.ts new file mode 100644 index 0000000..3a3bdbe --- /dev/null +++ b/src/cache.ts @@ -0,0 +1,75 @@ +import { CacheMapType } from './types'; + +// check every minute +const CACHE_INTERVAL = 60 * 1000; + +export const approvalCache: CacheMapType = new Map(); +export const balanceCache: CacheMapType = new Map(); + + +const clearCacheJob = (type: 'balance' | 'approval') => { + // 1mb per map + const totalMaxSize = 1_000_000 + + const cache = type === 'balance' ? balanceCache : approvalCache; + let cacheSize = getMapSizeInBytes(cache); + + const diff = cacheSize - totalMaxSize; + if (diff < 0) return; + + + // Convert to array and sort in one pass + const sortedEntries = Array.from(cache.entries()) + .sort((a, b) => a[1].ts - b[1].ts); + + let index = 0; + while (cacheSize > totalMaxSize && index < sortedEntries.length) { + const [key, value] = sortedEntries[index]; + const entrySize = getObjectSize(key) + getObjectSize(value); + cache.delete(key); + cacheSize -= entrySize; + index++; + } +} + +const getMapSizeInBytes = (map: CacheMapType) => { + let totalSize = 0; + + for (const [key, value] of map) { + totalSize += getObjectSize(key); + totalSize += getObjectSize(value); + } + + // Add overhead for the Map structure itself + totalSize += 8 * map.size; // Assuming 8 bytes per entry for internal structure + + return totalSize; +} + +const getObjectSize = (obj: any) => { + const type = typeof obj; + switch (type) { + case 'number': + return 8; + case 'string': + return obj.length * 2; + case 'boolean': + return 4; + case 'object': + if (obj === null) { + return 0; + } + let size = 0; + for (const key in obj) { + if (obj.hasOwnProperty(key)) { + size += getObjectSize(key); + size += getObjectSize(obj[key]); + } + } + return size; + default: + return 0; + } +} + +setInterval(clearCacheJob, CACHE_INTERVAL); diff --git a/src/index.ts b/src/index.ts index 0286f55..85f95ac 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,3 +1,4 @@ +import { approvalCache, balanceCache } from './cache'; import { generateMockApprovalData, getErc20Approval, @@ -10,6 +11,8 @@ import { } from "./balance"; export { + approvalCache, + balanceCache, generateMockApprovalData, generateMockBalanceData, getErc20ApprovalStorageSlot, diff --git a/src/types.ts b/src/types.ts new file mode 100644 index 0000000..4f38c7d --- /dev/null +++ b/src/types.ts @@ -0,0 +1,11 @@ +export interface CacheData { + // Slot # + slot: number; + // if contract is vyper + isVyper: boolean; + // Timestamp added (for cleaning purposes) + ts: number; +} +export type CacheMapType = Map; + +