Repository: shazow/whatsabi
Pull request #120 (2024-09-18 19:37:54)

Improve devex for only resolving proxies
This is mainly when a user already has bytecode and just wants to resolve proxies on it with the minimum number of RPC requests. (A valid provider is still required because various proxies require more RPC requests to resolve.)

Example:

```ts
const address = "0x...";
const bytecode = "0x..."; // Already loaded from somewhere

const cachedCodeProvider = whatsabi.providers.WithCachedCode(provider, {
  [address]: bytecode,
});
                                                                                                                                                      
const result = whatsabi.autoload(cachedCodeProvider, address, {
  abiLoader: false, // Skip ABI loaders
  signatureLookup: false, // Skip looking up selector signatures
  followProxies: true,
})
                                                                                                                                                      
if (result.address !== address) console.log(`Resolved proxy: ${address} -> ${result.address}`); // only if followProxies is true
if (result.proxies.length > 0) console.log("Proxies detected:", result.proxies);
// Note that some proxies can only be resolved relative to a selector, like DiamondProxy. These will need to be resolved manually via result.proxies.
```



Patch:
diff --git a/src/auto.ts b/src/auto.ts
index ec42729..f10be5f 100644
--- a/src/auto.ts
+++ b/src/auto.ts
@@ -9,6 +9,7 @@ import * as errors from "./errors.js";
 import { CompatibleProvider } from "./providers.js";
 import { defaultABILoader, defaultSignatureLookup } from "./loaders.js";
 import { abiFromBytecode, disasm } from "./disasm.js";
+import { MultiABILoader } from "./loaders.js";
 
 function isAddress(address: string) {
     return address.length === 42 && address.startsWith("0x") && Number(address) >= 0;
@@ -25,7 +26,7 @@ export type AutoloadResult = {
     abi: ABI;
 
     /** Whether the `abi` is loaded from a verified source */
-    isVerified: boolean;
+    abiLoadedFrom?: ABILoader;
 
     /** List of resolveable proxies detected in the contract */
     proxies: ProxyResolver[],
@@ -121,7 +122,6 @@ export async function autoload(address: string, config: AutoloadConfig): Promise
         address,
         abi: [],
         proxies: [],
-        isVerified: false,
     };
 
     let abiLoader = config.abiLoader;
@@ -195,6 +195,23 @@ export async function autoload(address: string, config: AutoloadConfig): Promise
         // Attempt to load the ABI from a contract database, if exists
         onProgress("abiLoader", { address, facets: Object.keys(facets) });
         const loader = abiLoader;
+
+        let abiLoadedFrom;
+        let originalOnLoad;
+        if (loader instanceof MultiABILoader) {
+            // This is a workaround for avoiding to change the loadABI signature, we can remove it if we use getContract instead.
+            const onLoad = (loader: ABILoader) => {
+                abiLoadedFrom = loader;
+            }
+            originalOnLoad = loader.onLoad;
+            if (!loader.onLoad) loader.onLoad = onLoad;
+            else {
+                // Just in case someone uses this feature, let's wrap it to include both. Not ideal... Is there a better way here?
+                const original = loader.onLoad;
+                loader.onLoad = (loader: ABILoader) => { original(loader); onLoad(loader); };
+            }
+        }
+
         try {
             const addresses = Object.keys(facets);
             const promises = addresses.map(addr => loader.loadABI(addr));
@@ -204,12 +221,16 @@ export async function autoload(address: string, config: AutoloadConfig): Promise
             }));
             result.abi = pruneFacets(facets, abis);
             if (result.abi.length > 0) {
-                result.isVerified = true;
+                result.abiLoadedFrom = abiLoadedFrom;
                 return result;
             }
         } catch (error: any) {
             // TODO: Catch useful errors
             if (onError("abiLoader", error) === true) return result;
+        } finally {
+            if (loader instanceof MultiABILoader) {
+                loader.onLoad = originalOnLoad;
+            }
         }
     }
 
diff --git a/src/loaders.ts b/src/loaders.ts
index 2bfefeb..0efad37 100644
--- a/src/loaders.ts
+++ b/src/loaders.ts
@@ -35,6 +35,7 @@ export type ContractResult = {
     runs: number;
 
     ok: boolean; // False if no result is found
+    loader?: ABILoader;
 
     /**
      * getSources returns the imports -> source code mapping for the contract, if available.
@@ -77,14 +78,23 @@ const emptyContractResult: ContractResult = {
 }
 
 export interface ABILoader {
+    readonly name: string;
+
     loadABI(address: string): Promise<any[]>;
     getContract(address: string): Promise<ContractResult>
 }
 
 // Load ABIs from multiple providers until a result is found.
 export class MultiABILoader implements ABILoader {
+    readonly name: string = "MultiABILoader";
+
     loaders: ABILoader[];
 
+    /// Note: This callback is used to pull out which loader succeeded without modifying the return API.
+    /// We can remove it once we switch to using getContract for autoload.
+    /// @internal
+    onLoad?: (loader: ABILoader) => void;
+
     constructor(loaders: ABILoader[]) {
         this.loaders = loaders;
     }
@@ -93,7 +103,10 @@ export class MultiABILoader implements ABILoader {
         for (const loader of this.loaders) {
             try {
                 const r = await loader.getContract(address);
-                if (r && r.abi.length > 0) return r;
+                if (r && r.abi.length > 0) {
+                    if (this.onLoad) this.onLoad(loader);
+                    return r;
+                }
             } catch (err: any) {
                 if (err.status === 404) continue;
 
@@ -112,7 +125,10 @@ export class MultiABILoader implements ABILoader {
                 const r = await loader.loadABI(address);
 
                 // Return the first non-empty result
-                if (r.length > 0) return r;
+                if (r.length > 0) {
+                    if (this.onLoad) this.onLoad(loader);
+                    return r;
+                }
             } catch (err: any) {
                 throw new MultiABILoaderError("MultiABILoader loadABI error: " + err.message, {
                     context: { loader, address },
@@ -128,6 +144,8 @@ export class MultiABILoaderError extends errors.LoaderError { };
 
 
 export class EtherscanABILoader implements ABILoader {
+    readonly name = "EtherscanABILoader";
+
     apiKey?: string;
     baseURL: string;
 
@@ -188,6 +206,7 @@ export class EtherscanABILoader implements ABILoader {
                 },
 
                 ok: true,
+                loader: this,
             };
         } catch (err: any) {
             throw new EtherscanABILoaderError("EtherscanABILoader getContract error: " + err.message, {
@@ -233,6 +252,8 @@ function isSourcifyNotFound(error: any): boolean {
 
 // https://sourcify.dev/
 export class SourcifyABILoader implements ABILoader {
+    readonly name = "SourcifyABILoader";
+
     chainId?: number;
 
     constructor(config?: { chainId?: number }) {
@@ -281,6 +302,7 @@ export class SourcifyABILoader implements ABILoader {
                 devdoc: m.output.devdoc,
 
                 ok: true,
+                loader: this,
             };
         } catch (err: any) {
             if (isSourcifyNotFound(err)) return emptyContractResult;
diff --git a/src/providers.ts b/src/providers.ts
index 95dac40..6c1f689 100644
--- a/src/providers.ts
+++ b/src/providers.ts
@@ -65,6 +65,36 @@ export function CompatibleProvider(provider: any): Provider {
     });
 }
 
+/**
+ * Wrap an existing provider into one that will return a fixed getCode result for items defined in codeCache.
+ * The cache is treated as read-only, it will not be updated. Mainly used to avoid an extra RPC call when we already have the bytcode.
+ * @param provider - Any existing provider
+ * @param codeCache - Object containing address => code mappings
+ * @returns {Provider} - Provider that will return a fixed getCode result for items defined in codeCache.
+ * @example
+ * ```ts
+ * const address = "0x0000000000000000000000000000000000000001";
+ * const bytecode = "0x6001600101"
+ * const cachedProvider = WithCachedCode(provider, {
+ *   [address]: bytecode,
+ * });
+ * const code = await cachedProvider.getCode(address);
+ * console.log(code); // "0x6001600101"
+ * ```
+ */
+export function WithCachedCode(provider: Provider, codeCache: Record<string, string>): Provider {
+    return {
+        ...provider,
+        async getCode(address: string): Promise<string> {
+            if (codeCache[address]) {
+                return codeCache[address];
+            }
+            return await provider.getCode(address);
+        }
+    };
+}
+
+
 // RPCPRovider thesis is: let's stop trying to adapt to every RPC wrapper library's high-level functions
 // and instead have a discovery for the lowest-level RPC call function that we can use directly.
 // At least whenever possible. Higher-level functionality like getAddress is still tricky.
diff --git a/src/proxies.ts b/src/proxies.ts
index 4c09f8e..69af77c 100644
--- a/src/proxies.ts
+++ b/src/proxies.ts
@@ -1,8 +1,43 @@
+/**
+ * @module proxies
+ * @example
+ * Using WhatsABI to only resolve proxies with a known bytecode:
+ * ```ts
+ * const address = "0x...";
+ * const bytecode = "0x..."; // Already loaded from somewhere
+ * const cachedCodeProvider = whatsabi.providers.WithCachedCode(provider, {
+ *   [address]: bytecode,
+ * });
+ *
+ * const result = whatsabi.autoload(cachedCodeProvider, address, {
+ *   abiLoader: false, // Skip ABI loaders
+ *   signatureLookup: false, // Skip looking up selector signatures
+ * })
+ *
+ * if (result.address !== address) console.log(`Resolved proxy: ${address} -> ${result.address}`);
+ * if (result.proxies.length > 0) console.log("Proxies detected:", result.proxies);
+ * // Note that some proxies can only be resolved relative to a selector, like DiamondProxy. These will need to be resolved manually via result.proxies.
+ * ```
+ * @example
+ * Resolve a DiamondProxy:
+ * ```ts
+ * // Let's say we have a result with a DiamondProxy in it, from the above example
+ * const resolver = result.proxies[0] as whatsabi.proxies.DiamondProxyResolver;
+ * 
+ * // DiamondProxies have different contracts mapped relative to the selector,
+ * // so we must resolve them against a selector.
+ * const selector = "0x6e9960c3";  // function getAdmin() returns (address)
+ *
+ * const implementationAddress = await resolver.resolve(provider, address, selector);
+ * ```
+ *
+ */
 import type { StorageProvider, CallProvider } from "./providers.js";
 import { addSlotOffset, readArray, joinSlot } from "./slots.js";
 import { addressWithChecksum } from "./utils.js";
 
 export interface ProxyResolver {
+    readonly name: string;
     resolve(provider: StorageProvider|CallProvider, address: string, selector?: string): Promise<string>
     toString(): string,
 }
@@ -56,6 +91,8 @@ const EIP1967FallbackSelectors = [
 ];
 
 export class EIP1967ProxyResolver extends BaseProxyResolver implements ProxyResolver {
+    override name = "EIP1967Proxy";
+
     async resolve(provider: StorageProvider & CallProvider, address: string): Promise<string> {
         // Is there an implementation defined?
         const implAddr = addressFromPadded(await provider.getStorageAt(address, slots.EIP1967_IMPL));
@@ -95,6 +132,8 @@ const diamondSelectors = [
 
 // ERC2535 - Diamond/Facet Proxy
 export class DiamondProxyResolver extends BaseProxyResolver implements ProxyResolver {
+    override name = "DiamondProxy";
+
     async resolve(provider: StorageProvider & CallProvider, address: string, selector: string): Promise<string> {
         if (!selector) {
             throw "DiamondProxy requires a selector to resolve to a specific facet";
@@ -194,12 +233,16 @@ export class DiamondProxyResolver extends BaseProxyResolver implements ProxyReso
 }
 
 export class ZeppelinOSProxyResolver extends BaseProxyResolver implements ProxyResolver {
+    override name = "ZeppelinOSProxy";
+
     async resolve(provider: StorageProvider, address: string): Promise<string> {
         return addressFromPadded(await provider.getStorageAt(address, slots.ZEPPELINOS_IMPL));
     }
 }
 
 export class PROXIABLEProxyResolver extends BaseProxyResolver implements ProxyResolver {
+    override name = "PROXIABLEProxy";
+
     async resolve(provider: StorageProvider, address: string): Promise<string> {
         return addressFromPadded(await provider.getStorageAt(address, slots.PROXIABLE));
     }
@@ -208,19 +251,18 @@ export class PROXIABLEProxyResolver extends BaseProxyResolver implements ProxyRe
 // https://github.com/0xsequence/wallet-contracts/blob/master/contracts/Wallet.sol
 // Implementation pointer is stored in slot keyed on the deployed address.
 export class SequenceWalletProxyResolver extends BaseProxyResolver implements ProxyResolver {
+    override name = "SequenceWalletProxy";
+
     async resolve(provider: StorageProvider, address: string): Promise<string> {
         return addressFromPadded(await provider.getStorageAt(address, address.toLowerCase().slice(2)));
     }
-
-    toString(): string {
-        return "SequenceWalletProxy";
-    }
 }
 
 // FixedProxyResolver is used when we already know the resolved address
 // No additional resolving required
 // Example: EIP-1167
 export class FixedProxyResolver extends BaseProxyResolver implements ProxyResolver {
+    override name = "FixedProxy";
     readonly resolvedAddress : string;
 
     constructor(name: string, resolvedAddress: string) {
diff --git a/src/whatsabi.ts b/src/whatsabi.ts
index b4a6c5e..3747ad0 100644
--- a/src/whatsabi.ts
+++ b/src/whatsabi.ts
@@ -18,9 +18,13 @@ export type {
 } from "./auto.js";
 
 import * as loaders from "./loaders.js";
-import * as proxies from "./proxies.js";
 export { loaders };
+
+import * as proxies from "./proxies.js";
 export { proxies };
 
+import * as providers from "./providers.js";
+export { providers };
+
 import * as errors from "./errors.js";
 export { errors };


Test patch:
diff --git a/src/__tests__/auto.test.ts b/src/__tests__/auto.test.ts
index 7a65b19..c7a231f 100644
--- a/src/__tests__/auto.test.ts
+++ b/src/__tests__/auto.test.ts
@@ -70,3 +70,12 @@ online_test('autoload non-contract', async ({ provider, env }) => {
     });
     expect(abi).toStrictEqual([]);
 });
+
+online_test('autoload verified etherscan', async ({ provider, env }) => {
+    const address = "0x8f8ef111b67c04eb1641f5ff19ee54cda062f163"; // Uniswap v3 pool, verified on Etherscan and Sourcify
+    const result = await autoload(address, {
+        provider: provider,
+        ...whatsabi.loaders.defaultsWithEnv(env),
+    });
+    expect(result.abiLoadedFrom?.name).toBeTruthy()
+});
diff --git a/src/__tests__/loaders.test.ts b/src/__tests__/loaders.test.ts
index 5b75d4b..807d4d7 100644
--- a/src/__tests__/loaders.test.ts
+++ b/src/__tests__/loaders.test.ts
@@ -69,14 +69,15 @@ describe('loaders module', () => {
 
   online_test('SourcifyABILoader_getContract', async () => {
     const loader = new SourcifyABILoader();
-    const {abi, name, userdoc, devdoc} = await loader.getContract("0x7a250d5630b4cf539739df2c5dacb4c659f2488d");
-    const selectors = Object.values(selectorsFromABI(abi));
+    const result = await loader.getContract("0x7a250d5630b4cf539739df2c5dacb4c659f2488d");
+    const selectors = Object.values(selectorsFromABI(result.abi));
     const sig = "swapExactETHForTokens(uint256,address[],address,uint256)";
     expect(selectors).toContain(sig);
-    expect(name).toBe("UniswapV2Router02");
+    expect(result.name).toStrictEqual("UniswapV2Router02")
+    expect(result.loader?.name).toStrictEqual("SourcifyABILoader");
 
-    expect(userdoc).toBeDefined();
-    expect(devdoc).toBeDefined();
+    expect(result.userdoc).toBeDefined();
+    expect(result.devdoc).toBeDefined();
   })
 
   online_test('SourcifyABILoader_getContract_missing', async () => {
@@ -128,13 +129,14 @@ describe('loaders module', () => {
       new SourcifyABILoader(),
       new EtherscanABILoader({ apiKey: process.env["ETHERSCAN_API_KEY"] }),
     ]);
-    const {abi, name, userdoc, devdoc} = await loader.getContract(address);
+    const result = await loader.getContract(address);
     const sig = "getMagistrate()";
-    const selectors = Object.values(selectorsFromABI(abi));
+    const selectors = Object.values(selectorsFromABI(result.abi));
     expect(selectors).toContain(sig);
-    expect(name).toEqual("KetherSortition");
-    expect(userdoc).not.toBeDefined();
-    expect(devdoc).not.toBeDefined();
+    expect(result.name).toEqual("KetherSortition");
+    expect(result.userdoc).not.toBeDefined();
+    expect(result.devdoc).not.toBeDefined();
+    expect(result.loader?.name).toStrictEqual(EtherscanABILoader.name);
   }, SLOW_ETHERSCAN_TIMEOUT);
 
   online_test('MultiABILoader_getContract_UniswapV3Factory', async () => {
@@ -157,13 +159,14 @@ describe('loaders module', () => {
     const loader = new MultiABILoader([
       new SourcifyABILoader(),
     ]);
-    const {abi, name, userdoc, devdoc} = await loader.getContract(address);
+    const result = await loader.getContract(address);
     const sig = "owner()";
-    const selectors = Object.values(selectorsFromABI(abi));
+    const selectors = Object.values(selectorsFromABI(result.abi));
     expect(selectors).toContain(sig);
-    expect(name).toEqual("Canonical Uniswap V3 factory");
-    expect(userdoc).toBeDefined();
-    expect(devdoc).toBeDefined();
+    expect(result.name).toEqual("Canonical Uniswap V3 factory");
+    expect(result.userdoc).toBeDefined();
+    expect(result.devdoc).toBeDefined();
+    expect(result.loader?.name).toStrictEqual(SourcifyABILoader.name);
   }, SLOW_ETHERSCAN_TIMEOUT);
 
   online_test('MultiABILoader_EtherscanOnly_getContract_UniswapV3Factory', async () => {
diff --git a/src/__tests__/providers.test.ts b/src/__tests__/providers.test.ts
new file mode 100644
index 0000000..9844a91
--- /dev/null
+++ b/src/__tests__/providers.test.ts
@@ -0,0 +1,18 @@
+import { expect } from 'vitest';
+import { test } from "./env";
+
+import { whatsabi } from "../index";
+
+const fakeProvider = whatsabi.providers.CompatibleProvider({
+    request: () => { },
+});
+
+test('provider WithCachedCode helper', async () => {
+    const address = "0x0000000000000000000000000000000000000001";
+    const provider = whatsabi.providers.WithCachedCode(fakeProvider, {
+        [address]: "0xf00",
+    });
+
+    expect(await provider.getCode(address)).toStrictEqual("0xf00");
+    expect(await provider.getCode("not cached")).toBeUndefined();
+});
diff --git a/src/__tests__/proxies.test.ts b/src/__tests__/proxies.test.ts
index b23bd27..23fc5b1 100644
--- a/src/__tests__/proxies.test.ts
+++ b/src/__tests__/proxies.test.ts
@@ -19,6 +19,8 @@ describe('proxy detection', () => {
         expect(program.proxies[0]).toBeInstanceOf(proxies.FixedProxyResolver);
         const proxy = program.proxies[0] as proxies.FixedProxyResolver;
         expect(proxy.resolvedAddress).toBe("0xbebebebebebebebebebebebebebebebebebebebe");
+        expect(proxy.name).toBe("FixedProxy");
+        expect(proxy.toString()).toBe("FixedProxy");
     });
 
     test('EIP-1167 Proxy: Uniswap v1', async () => {
@@ -55,6 +57,7 @@ describe('proxy detection', () => {
 
         const program = disasm(bytecode);
         expect(program.proxies[0]).toBeInstanceOf(proxies.GnosisSafeProxyResolver);
+        expect(program.proxies[0].name).toBe("GnosisSafeProxy");
     });
 
     test('ZeppelinOS Proxy', async () => {