From f4d5fb38b34c9371657bed5c01560daaacd9287a Mon Sep 17 00:00:00 2001 From: d3or Date: Fri, 12 Jul 2024 11:51:54 -0400 Subject: [PATCH 1/3] feat: add cache & cache clearing job --- src/approval.ts | 194 +++++++++++++++++++++++++++--------------------- src/balance.ts | 77 +++++++++++++++---- src/cache.ts | 75 +++++++++++++++++++ src/index.ts | 3 + src/types.ts | 11 +++ 5 files changed, 264 insertions(+), 96 deletions(-) create mode 100644 src/cache.ts create mode 100644 src/types.ts 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..27a59c1 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 } = calculateApprovalVyperStorageSlot(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 } = calculateApprovalSolidityStorageSlot(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 } = calculateApprovalSolidityStorageSlot(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 } = calculateApprovalVyperStorageSlot(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 calculateApprovalSolidityStorageSlot = (holderAddress: string, slotNumber: number) => { + const slotHash = ethers.utils.solidityKeccak256( + ["uint256", "uint256"], + [holderAddress, slotNumber] + ); + return { slotHash } +} + +const calculateApprovalVyperStorageSlot = (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; + + From e30549823eeb0ec8568e2e4329693f4b7dd354da Mon Sep 17 00:00:00 2001 From: d3or Date: Fri, 12 Jul 2024 11:54:14 -0400 Subject: [PATCH 2/3] feat: rename method --- src/balance.ts | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/balance.ts b/src/balance.ts index 27a59c1..3a75a8a 100644 --- a/src/balance.ts +++ b/src/balance.ts @@ -98,7 +98,7 @@ export const getErc20BalanceStorageSlot = async ( const cachedValue = balanceCache.get(erc20Address.toLowerCase()); if (cachedValue) { if (cachedValue.isVyper) { - const { vyperSlotHash } = calculateApprovalVyperStorageSlot(holderAddress, cachedValue.slot) + const { vyperSlotHash } = calculateBalanceVyperStorageSlot(holderAddress, cachedValue.slot) const vyperBalance = await provider.getStorageAt( erc20Address, vyperSlotHash @@ -109,7 +109,7 @@ export const getErc20BalanceStorageSlot = async ( isVyper: true, }; } else { - const { slotHash } = calculateApprovalSolidityStorageSlot(holderAddress, cachedValue.slot); + const { slotHash } = calculateBalanceSolidityStorageSlot(holderAddress, cachedValue.slot); const balance = await provider.getStorageAt(erc20Address, slotHash); return { slot: ethers.BigNumber.from(cachedValue.slot).toHexString(), @@ -133,7 +133,7 @@ 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 { slotHash } = calculateApprovalSolidityStorageSlot(holderAddress, i); + const { slotHash } = calculateBalanceSolidityStorageSlot(holderAddress, i); const balance = await provider.getStorageAt(erc20Address, slotHash); if (ethers.BigNumber.from(balance).eq(userBalance)) { @@ -150,7 +150,7 @@ export const getErc20BalanceStorageSlot = async ( }; } - const { vyperSlotHash } = calculateApprovalVyperStorageSlot(holderAddress, i) + const { vyperSlotHash } = calculateBalanceVyperStorageSlot(holderAddress, i) const vyperBalance = await provider.getStorageAt( erc20Address, vyperSlotHash @@ -174,7 +174,7 @@ export const getErc20BalanceStorageSlot = async ( }; -const calculateApprovalSolidityStorageSlot = (holderAddress: string, slotNumber: number) => { +const calculateBalanceSolidityStorageSlot = (holderAddress: string, slotNumber: number) => { const slotHash = ethers.utils.solidityKeccak256( ["uint256", "uint256"], [holderAddress, slotNumber] @@ -182,7 +182,7 @@ const calculateApprovalSolidityStorageSlot = (holderAddress: string, slotNumber: return { slotHash } } -const calculateApprovalVyperStorageSlot = (holderAddress: string, slotNumber: number) => { +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( From dc6941a0dca6aecace5ec2fa48a3197935722f76 Mon Sep 17 00:00:00 2001 From: d3or Date: Fri, 12 Jul 2024 11:57:00 -0400 Subject: [PATCH 3/3] Feat: update readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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