Trie/StateManager: create partial state tries / state managers from proofs (#3186)

* trie: add methods to create and update tries from proofs

* stateManager: add fromProof support

* Add readme updates for new fromProof constructors

* Apply suggested changes

* Add missing line

* statemanager/trie: further update readme examples

---------

Co-authored-by: acolytec3 <[email protected]>

Author: jochem-brouwer

packages/statemanager/README.md

@@ -25,11 +25,13 @@ Note: this library was part of the [@ethereumjs/vm](../vm/) package up till VM `
 
 The `StateManager` provides high-level access and manipulation methods to and for the Ethereum state, thinking in terms of accounts or contract code rather then the storage operations of the underlying data structure (e.g. a [Trie](../trie/)).
 
-The library includes a TypeScript interface `StateManager` to ensure a unified interface (e.g. when passed to the VM) as well as a concrete Trie-based implementation `DefaultStateManager` as well as an `RPCStateManager` implementation that sources state and history data from an external JSON-RPC provider.
+The library includes a TypeScript interface `StateManager` to ensure a unified interface (e.g. when passed to the VM), a concrete Trie-based `DefaultStateManager` implementation, as well as an `RPCStateManager` implementation that sources state and history data from an external JSON-RPC provider.
 
 It also includes a checkpoint/revert/commit mechanism to either persist or revert state changes and provides a sophisticated caching mechanism under the hood to reduce the need for direct state accesses.
 
-### `DefaultStateManager` Example
+### `DefaultStateManager`
+
+#### Usage example
 
 ```typescript
 import { Account, Address } from '@ethereumjs/util'
@@ -45,7 +47,7 @@ await stateManager.commit()
 await stateManager.flush()
 ```
 
-### Account, Storage and Code Caches
+#### Account, Storage and Code Caches
 
 Starting with the v2 release and complemented by the v2.1 release the StateManager comes with a significantly more elaborate caching mechanism for account, storage and code caches.
 
@@ -55,6 +57,25 @@ Caches now "survive" a flush operation and especially long-lived usage scenarios
 
 Have a loot at the extended `CacheOptions` on how to use and leverage the new cache system.
 
+#### Instantiating from a proof
+
+The `DefaultStateManager` has a static constructor `fromProof` that accepts one or more [EIP-1186](https://eips.ethereum.org/EIPS/eip-1186) [proofs](./src/stateManager.ts) and will instantiate a `DefaultStateManager` with a partial trie containing the state provided by the proof(s). See below example:
+
+```typescript
+// setup `stateManager` with some existing address
+const proof = await stateManager.getProof(address)
+const proofWithStorage = await stateManger.getProof(contractAddress, [storageKey1, storageKey2])
+
+const partialStateManager = await DefaultStateManager.fromProof(proof)
+// To add more proof data, use `addProofData`
+await partialStateManager.addProofData(proofWithStorage)
+const accountFromNewSM = await partialStateManager.getAccount(address)
+const accountFromOldSM = await stateManager.getAccount(address)
+console.log(accountFromNewSM, accountFromOldSM) // should match
+const slot1FromNewSM = await stateManager.getContractStorage(contractAddress, storageKey1)
+const slot2FromNewSM = await stateManager.getContractStorage(contractAddress, storageKey1) // should also match
+```
+
 ### `RPCStateManager`
 
 First, a simple example of usage:
@@ -74,9 +95,9 @@ The `RPCStateManager` can be be used with any JSON-RPC provider that supports th
 
 **Note:** Usage of this StateManager can cause a heavy load regarding state request API calls, so be careful (or at least: aware) if used in combination with a JSON-RPC provider connecting to a third-party API service like Infura!
 
-### Points on usage:
+#### Points on `RPCStateManager` usage
 
-#### Instantiating the EVM
+##### Instantiating the EVM
 
 In order to have an EVM instance that supports the BLOCKHASH opcode (which requires access to block history), you must instantiate both the `RPCStateManager` and the `RpcBlockChain` and use that when initalizing your EVM instance as below:
 
@@ -92,24 +113,24 @@ const evm = new EVM({ blockchain, stateManager: state })
 
 Note: Failing to provide the `RPCBlockChain` instance when instantiating the EVM means that the `BLOCKHASH` opcode will fail to work correctly during EVM execution.
 
-#### Provider selection
+##### Provider selection
 
 - The provider you select must support the `eth_getProof`, `eth_getCode`, and `eth_getStorageAt` RPC methods.
 - Not all providers support retrieving state from all block heights so refer to your provider's documentation. Trying to use a block height not supported by your provider (e.g. any block older than the last 256 for CloudFlare) will result in RPC errors when using the state manager.
 
-#### Block Tag selection
+##### Block Tag selection
 
 - You have to pass a block number or `earliest` in the constructor that specifies the block height you want to pull state from.
 - The `latest`/`pending` values supported by the Ethereum JSON-RPC are not supported as longer running scripts run the risk of state values changing as blocks are mined while your script is running.
 - If using a very recent block as your block tag, be aware that reorgs could occur and potentially alter the state you are interacting with.
 - If you want to rerun transactions from block X or run block X, you need to specify the block tag as X-1 in the state manager constructor to ensure you are pulling the state values at the point in time the transactions or block was run.
 
-#### Potential gotchas
+##### Potential gotchas
 
 - The RPC State Manager cannot compute valid state roots when running blocks as it does not have access to the entire Ethereum state trie so can not compute correct state roots, either for the account trie or for storage tries.
 - If you are replaying mainnet transactions and an account or account storage is touched by multiple transactions in a block, you must replay those transactions in order (with regard to their position in that block) or calculated gas will likely be different than actual gas consumed.
 
-#### Further reference
+##### Further reference
 
 Refer to [this test script](./test/rpcStateManager.spec.ts) for complete examples of running transactions and blocks in the `vm` with data sourced from a provider.
 

packages/statemanager/src/stateManager.ts

@@ -705,6 +705,7 @@ export class DefaultStateManager implements EVMStateManagerInterface {
    * @param storageSlots storage slots to get proof of
    */
   async getProof(address: Address, storageSlots: Uint8Array[] = []): Promise<Proof> {
+    await this.flush()
     const account = await this.getAccount(address)
     if (!account) {
       // throw new Error(`getProof() can only be called for an existing account`)
@@ -748,6 +749,93 @@ export class DefaultStateManager implements EVMStateManagerInterface {
     return returnValue
   }
 
+  /**
+   * Create a StateManager and initialize this with proof(s) gotten previously from getProof
+   * This generates a (partial) StateManager where one can retrieve all items from the proof
+   * @param proof Either a proof retrieved from `getProof`, or an array of those proofs
+   * @param safe Wether or not to verify that the roots of the proof items match the reported roots
+   * @param verifyRoot verify that all proof root nodes match statemanager's stateroot - should be
+   * set to `false` when constructing a state manager where the underlying trie has proof nodes from different state roots
+   * @returns A new DefaultStateManager with elements from the given proof included in its backing state trie
+   */
+  static async fromProof(
+    proof: Proof | Proof[],
+    safe: boolean = false,
+    opts: DefaultStateManagerOpts = {}
+  ): Promise<DefaultStateManager> {
+    if (Array.isArray(proof)) {
+      if (proof.length === 0) {
+        return new DefaultStateManager(opts)
+      } else {
+        const trie =
+          opts.trie ??
+          (await Trie.createTrieFromProof(
+            proof[0].accountProof.map((e) => hexToBytes(e)),
+            { useKeyHashing: true }
+          ))
+        const sm = new DefaultStateManager({ ...opts, trie })
+        const address = Address.fromString(proof[0].address)
+        await sm.addStorageProof(proof[0].storageProof, proof[0].storageHash, address, safe)
+        for (let i = 1; i < proof.length; i++) {
+          const proofItem = proof[i]
+          await sm.addProofData(proofItem, true)
+        }
+        await sm.flush() // TODO verify if this is necessary
+        return sm
+      }
+    } else {
+      return DefaultStateManager.fromProof([proof])
+    }
+  }
+
+  /**
+   * Adds a storage proof to the state manager
+   * @param storageProof The storage proof
+   * @param storageHash The root hash of the storage trie
+   * @param address The address
+   * @param safe Whether or not to verify if the reported roots match the current storage root
+   */
+  private async addStorageProof(
+    storageProof: StorageProof[],
+    storageHash: string,
+    address: Address,
+    safe: boolean = false
+  ) {
+    const trie = this._getStorageTrie(address)
+    trie.root(hexToBytes(storageHash))
+    for (let i = 0; i < storageProof.length; i++) {
+      await trie.updateTrieFromProof(
+        storageProof[i].proof.map((e) => hexToBytes(e)),
+        safe
+      )
+    }
+  }
+
+  /**
+   * Add proof(s) into an already existing trie
+   * @param proof The proof(s) retrieved from `getProof`
+   * @param verifyRoot verify that all proof root nodes match statemanager's stateroot - should be
+   * set to `false` when constructing a state manager where the underlying trie has proof nodes from different state roots
+   */
+  async addProofData(proof: Proof | Proof[], safe: boolean = false) {
+    if (Array.isArray(proof)) {
+      for (let i = 0; i < proof.length; i++) {
+        await this._trie.updateTrieFromProof(
+          proof[i].accountProof.map((e) => hexToBytes(e)),
+          safe
+        )
+        await this.addStorageProof(
+          proof[i].storageProof,
+          proof[i].storageHash,
+          Address.fromString(proof[i].address),
+          safe
+        )
+      }
+    } else {
+      await this.addProofData([proof], safe)
+    }
+  }
+
   /**
    * Verify an EIP-1186 proof. Throws if proof is invalid, otherwise returns true.
    * @param proof the proof to prove

packages/statemanager/test/proofStateManager.spec.ts

@@ -4,6 +4,7 @@ import {
   Address,
   bytesToHex,
   bytesToUnprefixedHex,
+  equalsBytes,
   hexToBytes,
   randomBytes,
   zeros,
@@ -28,6 +29,20 @@ describe('ProofStateManager', () => {
     assert.equal(proof.nonce, '0x0', 'Nonce is in quantity-encoded RPC representation')
   })
 
+  it(`should correctly return the right storage root / account root`, async () => {
+    const address = Address.zero()
+    const key = zeros(32)
+    const stateManager = new DefaultStateManager()
+
+    await stateManager.putAccount(address, new Account(BigInt(100), BigInt(200)))
+    const storageRoot = (await stateManager.getAccount(address))!.storageRoot
+
+    await stateManager.putContractStorage(address, key, new Uint8Array([10]))
+
+    const proof = await stateManager.getProof(address, [key])
+    assert.ok(!equalsBytes(hexToBytes(proof.storageHash), storageRoot))
+  })
+
   it(`should return quantity-encoded RPC representation for existing accounts`, async () => {
     const address = Address.zero()
     const key = zeros(32)
@@ -132,7 +147,6 @@ describe('ProofStateManager', () => {
       await trie._db.put(key, bufferData)
     }
     trie.root(stateRoot!)
-    await stateManager.putAccount(address, new Account())
     const proof = await stateManager.getProof(address)
     assert.deepEqual((ropsten_nonexistentAccount as any).default, proof)
     assert.ok(await stateManager.verifyProof(ropsten_nonexistentAccount))