Refreshing stale MPT proofs

README.md

Background

In the context of stateless execution, a node does not need to store the full state of the blockchain and expects each transaction to come with all the necessary parts of the state and proofs for the validity of those parts.

The responsibility of attaching the proofs falls on the transaction sender or a third party who's storing the necessary parts of the trie. If a transaction doesn't make it immediately into the next block, the blockchain advances forward, modifying the state (and the state root). This causes all the proofs that were generated against a previous state root to become stale.

Simulation

This script simulates the following scenario:

• Proofs for two txes (T1, T2) are generated against state root A.
• Then, verifier executes T1, resulting in state root B.
• Now the proofs in T2 are stale.

We want to refresh the proofs of the second tx, only having A, B, T1, T2 without having to store the whole state. This is done by refreshTx, which simply replays T1 against A, and re-generate proofs for T2 based on this intermediate state.

It could be however that after T1, T2 is no longer valid (e.g. T1 and T2 have the same sender, and T1 spends all balance).

index.js

const assert = require('assert')
const { promisify } = require('util')
const BN = require('bn.js')
const Trie = require('merkle-patricia-tree/secure')
const Account = require('ethereumjs-account').default
const StateManager = require('ethereumjs-vm/dist/state/stateManager').default
const PStateManager = require('ethereumjs-vm/dist/state/promisified').default
const { keccak256, ecsign, stripZeros } = require('ethereumjs-util')
const { encode } = require('rlp')
const Wallet = require('ethereumjs-wallet')
const yaml = require('js-yaml')
const fs = require('fs')

const prove = promisify(Trie.prove)
const verifyProof = promisify(Trie.verifyProof)

/**
 * This script tries to simulate the following scenario:
 *
 * - Proofs for two txes (T1, T2) are generated against state root A.
 * - Then, verifier executes T1, resulting in state root B.
 * - Now the proofs in T2 are stale.
 *
 * We want to refresh the proofs of the second tx, only having A, B, T1, T2
 * and not the whole state.
 */
async function main () {
  const rawState = new StateManager()
  const state = new PStateManager(rawState)

  // Generate random accounts
  let accounts = await generateAccounts(state, 5000)

  let root = await state.getStateRoot()
  const preStateRoot = root

  // Generate txes against initial state root
  let txes = await generateTxes(state, accounts, 2)

  // Apply txes on top of trie to compute expected post state root
  for (tx of txes) {
    await transfer(state, tx)
  }
  root = await state.getStateRoot()
  const postStateRoot = root

  // Verify proofs for T1 and execute it statelessly, returning B
  const intermediateStateRoot = await verifyAndUpdate(preStateRoot, txes[0])

  // Now T2 proofs are stale, refresh them
  const t2 = await refreshTx(txes[1], preStateRoot, intermediateStateRoot, txes[0])

  // Run verifier with refereshed T2 against root B
  const resRoot = await verifyAndUpdate(intermediateStateRoot, t2)

  // Computed state root should match expected post state root
  assert(postStateRoot.equals(resRoot))
}

async function refreshTx (tx, oldRoot, newRoot, previousTx) {
  // Verify proofs against old root
  await verifyProof(oldRoot, keccak256(tx.from), tx.fromWitness)
  await verifyProof(oldRoot, keccak256(tx.to), tx.toWitness)

  // If roots are equal, proof is already valid
  if (oldRoot.equals(newRoot)) {
    return tx
  }

  // Batch insert proof nodes for both txes into trie's underlying db, i.e
  // create a partial trie from the proof nodes for both txes.
  const trie = await trieFromProofNodes(oldRoot, previousTx.fromWitness.concat(previousTx.toWitness, tx.fromWitness, tx.toWitness))

  // Replay previous tx against old root to get intermediate state
  const state = new PStateManager(new StateManager({ trie }))
  await transfer(state, previousTx)
  const intermediateStateRoot = await state.getStateRoot()
  assert(newRoot.equals(intermediateStateRoot))

  // Generate fresh proofs against intermediate state
  const fromWitness = await prove(state._wrapped._trie, keccak256(tx.from))
  const toWitness = await prove(state._wrapped._trie, keccak256(tx.to))
  tx.fromWitness = fromWitness
  tx.toWitness = toWitness

  return tx
}

async function verifyAndUpdate (preStateRoot, tx) {
  // Verify proofs
  await verifyProof(preStateRoot, keccak256(tx.from), tx.fromWitness)
  await verifyProof(preStateRoot, keccak256(tx.to), tx.toWitness)

  const trie = await trieFromProofNodes(preStateRoot, tx.fromWitness.concat(tx.toWitness))
  const state = new PStateManager(new StateManager({ trie }))

  await transfer(state, tx)

  return state.getStateRoot()
}

async function transfer (state, tx) {
  let { from, to, value, nonce } = tx
  assert(value.gten(0))

  const fromAcc = await state.getAccount(from)
  const toAcc = await state.getAccount(to)

  assert(new BN(fromAcc.balance).gte(value))
  assert(new BN(fromAcc.nonce).eq(nonce))

  const newFromBalance = new BN(fromAcc.balance).sub(value)
  fromAcc.balance = newFromBalance.toBuffer()
  fromAcc.nonce = nonce.addn(1).toBuffer()
  const newToBalance = new BN(toAcc.balance).add(value)
  toAcc.balance = newToBalance.toBuffer()

  await state.putAccount(from, fromAcc)
  await state.putAccount(to, toAcc)
}

async function trieFromProofNodes (root, proofNodes) {
  const trie = new Trie(null, root)
  let opStack = proofNodes.map((nodeValue) => {
    return { type: 'put', key: keccak256(nodeValue), value: nodeValue }
  })

  const batchP = promisify(trie._batchNodes.bind(trie))
  await batchP(opStack)

  return trie
}

async function generateTxes (state, accounts, count=50) {
  let txes = []
  const root = await state.getStateRoot()

  for (let i = 0; i < count; i++) {
    const from = accounts[i].address
    const to = accounts[i + 1].address
    const value = new BN('00000000000000000000000000000000000000000000000000000000000000ff', 16)
    const nonce = new BN('0000000000000000000000000000000000000000000000000000000000000000', 16)

    const fromAccount = await state.getAccount(from)
    const fromWitness = await prove(state._wrapped._trie, keccak256(from))
    let val = await verifyProof(root, keccak256(from), fromWitness)
    assert(val.equals(fromAccount.serialize()), "valid from witness")

    const toAccount = await state.getAccount(to)
    const toWitness = await prove(state._wrapped._trie, keccak256(to))
    val = await verifyProof(root, keccak256(to), toWitness)
    assert(val.equals(toAccount.serialize()), "valid to witness")

    txes.push({
      from, to, value, nonce,
      fromWitness, toWitness
    })
  }

  return txes
}

async function generateAccounts (state, count=500) {
  let accounts = []
  for (let i = 0; i < count; i++) {
    let wallet = Wallet.generate()
    let address = wallet.getAddress()
    let privateKey = wallet.getPrivateKey()
    let account = new Account()
    account.balance = new BN('ffffff', 16).toBuffer()
    accounts.push({
      address,
      privateKey,
      account
    })
    await state.putAccount(address, account)
  }
  return accounts
}

main().then(() => {}).catch((e) => console.log(e))