Turboproof EE in JS

ee.ts

import * as assert from 'assert'
import BN = require('bn.js')
import { keccak256 } from 'ethereumjs-util'
import { encode, decode } from 'rlp'
import { verifyMultiproof, decodeInstructions } from './multiproof'

interface TestSuite {
  preStateRoot: Buffer
  blockData: Buffer
}

// TODO
interface BlockData {}

interface Account {
  nonce: BN
  balance: BN
  stateRoot: Buffer
  codeHash: Buffer
}

interface Tx {
  toIdx: number
  value: BN
  nonce: BN
  fromIdx: number
}

export function decodeBlockData(blockData: Buffer): any {
  const d = decode(blockData)
  return {
    txes: d[0],
    addrs: d[1],
    multiproof: {
      // @ts-ignore
      hashes: d[2][0],
      // @ts-ignore
      keyvals: d[2][1],
      // @ts-ignore
      instructions: decodeInstructions(d[2][2]),
    },
  }
}

function decodeTx(raw: Buffer[]): Tx {
  // TODO: How many addresses possible? is u8 enough for indices?
  return {
    toIdx: bufToU8(raw[0]),
    value: new BN(raw[1]),
    nonce: new BN(raw[2]),
    fromIdx: bufToU8(raw[3]),
  }
}

function bufToU8(b: Buffer): number {
  // RLP decoding of 0 is empty buffer
  if (b.length === 0) {
    return 0
  }
  return b.readUInt8(0)
}

function decodeAccount(raw: Buffer): Account {
  const d = decode(raw)
  // @ts-ignore
  return { nonce: new BN(d[0]), balance: new BN(d[1]), stateRoot: d[2], codeHash: d[3] }
}

function encodeAccount(acc: Account): Buffer {
  return encode([
    acc.nonce.toArrayLike(Buffer, 'be'),
    acc.balance.toArrayLike(Buffer, 'be'),
    acc.stateRoot,
    acc.codeHash,
  ])
}

export function main(testSuite: TestSuite): boolean {
  const preStateRoot = testSuite.preStateRoot
  const blockData = decodeBlockData(testSuite.blockData)

  const addrs = blockData.addrs
  const proof = blockData.multiproof
  const updatedLeaves = new Array(blockData.multiproof.keyvals.length).fill(undefined)

  for (const rawTx of blockData.txes) {
    const tx = decodeTx(rawTx)
    const fromIdx = tx.fromIdx
    const toIdx = tx.toIdx

    const encodedUnsignedTx = encode([addrs[toIdx], tx.value, tx.nonce])
    const txHash = keccak256(encodedUnsignedTx)
    // TODO: assert signature recovered address matches the given address
    // const recoveredAddr = ecrecover(txHash, tx.sig)
    // assert(addrs[fromIdx].equals(recoveredAddr))

    // If `from` or `to` has been modified, use the updated version
    let fromLeaf = proof.keyvals[fromIdx]
    if (updatedLeaves[fromIdx] !== undefined) {
      fromLeaf = updatedLeaves[fromIdx]
    }
    let toLeaf = proof.keyvals[toIdx]
    if (updatedLeaves[toIdx] !== undefined) {
      toLeaf = updatedLeaves[toIdx]
    }

    const decodedFrom = decode(fromLeaf)
    const decodedTo = decode(toLeaf)
    // @ts-ignore
    const fromAcc = decodeAccount(decodedFrom[1])
    // @ts-ignore
    const toAcc = decodeAccount(decodedTo[1])

    assert(fromAcc.balance.gte(tx.value))
    assert(fromAcc.nonce.eq(tx.nonce))

    fromAcc.balance.isub(tx.value)
    toAcc.balance.iadd(tx.value)
    fromAcc.nonce.iaddn(1)

    updatedLeaves[fromIdx] = encode([decodedFrom[0], encodeAccount(fromAcc)])
    updatedLeaves[toIdx] = encode([decodedTo[0], encodeAccount(toAcc)])
  }

  const keys = addrs.map((a: Buffer) => keccak256(a))

  return verifyMultiproof(preStateRoot, blockData.multiproof, keys)
}

multiproof.ts

import * as assert from 'assert'
import { decode, encode } from 'rlp'
import { keccak256 } from 'ethereumjs-util'
import { Trie } from './baseTrie'
import { BranchNode, ExtensionNode, LeafNode, EmbeddedNode, decodeRawNode } from './trieNode'
import { stringToNibbles, nibblesToBuffer, matchingNibbleLength } from './util/nibbles'
import { addHexPrefix, removeHexPrefix } from './util/hex'
const promisify = require('util.promisify')

export enum Opcode {
  Branch = 0,
  Hasher = 1,
  Leaf = 2,
  Extension = 3,
  Add = 4,
}

export enum NodeType {
  Branch = 0,
  Leaf = 1,
  Extension = 2,
  Hash = 3,
}

export interface Instruction {
  kind: Opcode
  value: number | number[]
}

export interface Multiproof {
  hashes: Buffer[]
  keyvals: Buffer[]
  instructions: Instruction[]
}

export function verifyMultiproof(root: Buffer, proof: Multiproof, keys: Buffer[]): boolean {
  const stack: any[] = []

  const leaves = proof.keyvals.map((l: Buffer) => decode(l))
  assert(leaves.length === keys.length)
  let leafIdx = 0
  let hashIdx = 0
  const paths = new Array(leaves.length).fill(undefined)

  for (const instr of proof.instructions) {
    if (instr.kind === Opcode.Hasher) {
      const h = proof.hashes[hashIdx++]
      if (!h) {
        throw new Error('Not enough hashes in multiproof')
      }
      stack.push([NodeType.Hash, [h, instr.value as number], []])
    } else if (instr.kind === Opcode.Leaf) {
      const l = leaves[leafIdx++]
      if (!l) {
        throw new Error('Expected leaf in multiproof')
      }
      // TODO: Nibble from prefix `digit`
      // @ts-ignore
      //stack.push([NodeType.Leaf, [l[0].slice(l[0].length - instr.value), l[1]]])
      // Disregard leaf operand
      stack.push([NodeType.Leaf, [l[0], l[1]], [leafIdx - 1]])
      // @ts-ignore
      paths[leafIdx - 1] = removeHexPrefix(stringToNibbles(l[0]))
    } else if (instr.kind === Opcode.Branch) {
      const n = stack.pop()
      if (!n) {
        throw new Error('Stack underflow')
      }
      const children = new Array(16).fill(null)
      children[instr.value as number] = n
      stack.push([NodeType.Branch, children, n[2].slice()])
      for (let i = 0; i < n[2].length; i++) {
        paths[n[2][i]] = [instr.value as number, ...paths[n[2][i]]]
      }
    } else if (instr.kind === Opcode.Extension) {
      const n = stack.pop()
      if (!n) {
        throw new Error('Stack underflow')
      }
      stack.push([NodeType.Extension, [instr.value, n], n[2].slice()])
      for (let i = 0; i < n[2].length; i++) {
        paths[n[2][i]] = [...(instr.value as number[]), ...paths[n[2][i]]]
      }
    } else if (instr.kind === Opcode.Add) {
      const n1 = stack.pop()
      const n2 = stack.pop()
      if (!n1 || !n2) {
        throw new Error('Stack underflow')
      }
      assert(n2[0] === NodeType.Branch, 'expected branch node on stack')
      assert(instr.value < 17)
      n2[1][instr.value as number] = n1
      n2[2] = Array.from(new Set([...n1[2], ...n2[2]]))
      stack.push(n2)
      for (let i = 0; i < n1[2].length; i++) {
        paths[n1[2][i]] = [instr.value as number, ...paths[n1[2][i]]]
      }
    } else {
      throw new Error('Invalid opcode')
    }
  }

  const r = stack.pop()
  if (!r) {
    throw new Error('Expected root node on top of stack')
  }
  let h = hashTrie(r)
  // Special case, if trie contains only one leaf
  // and that leaf has length < 32
  if (h.length < 32) {
    h = keccak256(encode(h))
  }

  // Assuming sorted keys
  for (let i = 0; i < paths.length; i++) {
    const addr = nibblesToBuffer(paths[i])
    assert(addr.equals(keys[i]))
  }
  return h.equals(root)
}

function hashTrie(node: any): Buffer {
  const typ = node[0]
  node = node[1]
  if (typ === NodeType.Branch) {
    const res = new Array(17).fill(Buffer.alloc(0))
    for (let i = 0; i < 16; i++) {
      if (node[i] === null) {
        continue
      }
      res[i] = hashTrie(node[i])
    }
    const e = encode(res)
    if (e.length > 32) {
      return keccak256(e)
    } else {
      return e
    }
  } else if (typ === NodeType.Leaf) {
    const e = encode(node)
    if (e.length > 32) {
      return keccak256(e)
    } else {
      return node
    }
  } else if (typ === NodeType.Hash) {
    // TODO: What if it's an embedded node with length === 32?
    // Maybe try decoding and if it fails assume it's a hash
    if (node[0].length < 32) {
      // Embedded node, decode to get correct serialization for parent node
      return decode(node[0])
    }
    return node[0]
  } else if (typ === NodeType.Extension) {
    const hashedNode = hashTrie(node[1])
    node = [nibblesToBuffer(addHexPrefix(node[0], false)), hashedNode]
    const e = encode(node)
    if (e.length > 32) {
      return keccak256(e)
    } else {
      return e
    }
  } else {
    throw new Error('Invalid node')
  }
}

export async function makeMultiproof(trie: Trie, keys: Buffer[]): Promise<Multiproof> {
  if (keys.length === 0) {
    return {
      hashes: [trie.root],
      keyvals: [],
      instructions: [{ kind: Opcode.Hasher, value: 0 }],
    }
  }

  const keysNibbles = []
  for (const k of keys) {
    keysNibbles.push(stringToNibbles(k))
  }

  return _makeMultiproof(trie, trie.root, keysNibbles)
}

async function _makeMultiproof(
  trie: Trie,
  rootHash: EmbeddedNode,
  keys: number[][],
): Promise<Multiproof> {
  let proof: Multiproof = {
    hashes: [],
    keyvals: [],
    instructions: [],
  }

  let root
  if (Buffer.isBuffer(rootHash)) {
    root = await promisify(trie._lookupNode.bind(trie))(rootHash)
  } else if (Array.isArray(rootHash)) {
    // Embedded node
    root = decodeRawNode(rootHash)
  } else {
    throw new Error('Unexpected root')
  }

  if (root instanceof BranchNode) {
    // Truncate first nibble of keys
    const table = new Array(16).fill(undefined)
    // Group target keys based by their first nibbles.
    // Also implicitly sorts the keys.
    for (const k of keys) {
      const idx = k[0]
      if (!table[idx]) table[idx] = []
      table[idx].push(k.slice(1))
    }

    let addBranchOp = true
    for (let i = 0; i < 16; i++) {
      if (table[i] === undefined) {
        // Empty subtree, hash it and add a HASHER op
        const child = root.getBranch(i)
        if (child) {
          proof.instructions.push({ kind: Opcode.Hasher, value: 0 })
          // TODO: Make sure child is a hash
          // what to do if embedded?
          if (Buffer.isBuffer(child)) {
            proof.hashes.push(child)
          } else if (Array.isArray(child)) {
            proof.hashes.push(encode(child))
          } else {
            throw new Error('Invalid branch child')
          }
          if (addBranchOp) {
            proof.instructions.push({ kind: Opcode.Branch, value: i })
            addBranchOp = false
          } else {
            proof.instructions.push({ kind: Opcode.Add, value: i })
          }
        }
      } else {
        const child = root.getBranch(i) as Buffer
        if (!child) {
          throw new Error('Key not in trie')
        }
        const p = await _makeMultiproof(trie, child, table[i])
        proof.hashes.push(...p.hashes)
        proof.keyvals.push(...p.keyvals)
        proof.instructions.push(...p.instructions)

        if (addBranchOp) {
          proof.instructions.push({ kind: Opcode.Branch, value: i })
          addBranchOp = false
        } else {
          proof.instructions.push({ kind: Opcode.Add, value: i })
        }
      }
    }
  } else if (root instanceof ExtensionNode) {
    const extkey = root.key
    // Make sure all keys follow the extension node
    // and truncate them.
    for (let i = 0; i < keys.length; i++) {
      const k = keys[i]
      if (matchingNibbleLength(k, extkey) !== extkey.length) {
        // TODO: Maybe allow proving non-existent keys
        throw new Error('Key not in trie')
      }
      keys[i] = k.slice(extkey.length)
    }
    const p = await _makeMultiproof(trie, root.value, keys)
    proof.hashes.push(...p.hashes)
    proof.keyvals.push(...p.keyvals)
    proof.instructions.push(...p.instructions)
    proof.instructions.push({ kind: Opcode.Extension, value: extkey })
  } else if (root instanceof LeafNode) {
    if (keys.length !== 1) {
      throw new Error('Expected 1 remaining key')
    }
    if (matchingNibbleLength(keys[0], root.key) !== root.key.length) {
      throw new Error("Leaf key doesn't match target key")
    }
    // TODO: Check key matches leaf's key
    proof = {
      hashes: [],
      keyvals: [root.serialize()],
      instructions: [{ kind: Opcode.Leaf, value: root.key.length }],
    }
  } else {
    throw new Error('Unexpected node type')
  }

  return proof
}

export function decodeMultiproof(raw: Buffer): Multiproof {
  const dec = decode(raw)
  assert(dec.length === 3)

  return {
    // @ts-ignore
    hashes: dec[0],
    // @ts-ignore
    keyvals: dec[1],
    // @ts-ignore
    instructions: decodeInstructions(dec[2]),
  }
}

export function encodeMultiproof(proof: Multiproof): Buffer {
  return encode(rawMultiproof(proof))
}

export function rawMultiproof(proof: Multiproof): any {
  return [proof.hashes, proof.keyvals, proof.instructions.map(i => [i.kind, i.value])]
}

export function decodeInstructions(instructions: Buffer[][]) {
  const res = []
  for (const op of instructions) {
    switch (bufToU8(op[0])) {
      case Opcode.Branch:
        res.push({ kind: Opcode.Branch, value: bufToU8(op[1]) })
        break
      case Opcode.Hasher:
        res.push({ kind: Opcode.Hasher, value: bufToU8(op[1]) })
        break
      case Opcode.Leaf:
        res.push({ kind: Opcode.Leaf, value: bufToU8(op[1]) })
        break
      case Opcode.Extension:
        // @ts-ignore
        res.push({ kind: Opcode.Extension, value: op[1].map(v => bufToU8(v)) })
        break
      case Opcode.Add:
        res.push({ kind: Opcode.Add, value: bufToU8(op[1]) })
        break
    }
  }
  return res
}

function bufToU8(b: Buffer): number {
  // RLP decoding of 0 is empty buffer
  if (b.length === 0) {
    return 0
  }
  return b.readUInt8(0)
}

relayer.js

const assert = require('assert')
const { promisify } = require('util')
const BN = require('bn.js')
const Trie = require('../../dist').SecureTrie
const { verifyMultiproof, makeMultiproof, rawMultiproof } = require('../../dist/multiproof')
const Account = require('ethereumjs-account').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')

async function main () {
  const testSuite = await generateTestSuite()
  writeScoutConfig(testSuite)
}

async function generateTestSuite () {
  const testSuite = {}
  const trie = new Trie()
  // Generate random accounts
  let accounts = await generateAccounts(trie, 5000)

  testSuite.preStateRoot = trie.root

  // Generate txes
  let [txes, addrs, multiproof, simulationData] = await generateTxes(trie, accounts, 70)

  // Serialize witnesses and tx data
  const blockData = encode([txes, addrs, rawMultiproof(multiproof)])
  console.log(`block data length: ${blockData.length}`)
  testSuite.blockData = blockData

  // Apply txes on top of trie to compute post state root
  for (tx of simulationData) {
    await transfer(trie, tx)
  }

  testSuite.postStateRoot = trie.root
  return testSuite
}

async function generateTxes (trie, accounts, count=50) {
  let txes = []
  let simulationData = []
  const root = trie.root
  const toProve = {}
  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)
    simulationData.push({ from, to, value, nonce })

    const fromAccount = await getAccount(trie, from)
    const fromKey = from.toString('hex')
    if (!toProve[fromKey]) {
      toProve[fromKey] = []
    }
    toProve[fromKey].push({ txId: i, fieldIdx: 3 })

    const toAccount = await getAccount(trie, to)
    const toKey = to.toString('hex')
    if (!toProve[toKey]) {
      toProve[toKey] = []
    }
    toProve[toKey].push({ txId: i, fieldIdx: 0 })

    const txRlp = encode([to, stripZeros(value.toBuffer('be', 32)), stripZeros(nonce.toBuffer('be', 32))])
    const txHash = keccak256(txRlp)
    const txSig = ecsign(txHash, accounts[i].privateKey)
    assert(txSig.r.byteLength === 32)
    assert(txSig.s.byteLength === 32)
    assert(txSig.v < 256)
    const v = new Uint8Array(1)
    v[0] = txSig.v
    const sigBytes = Buffer.concat([txSig.r, txSig.s, v], 65)

    txes.push([to, stripZeros(value.toBuffer('be', 32)), stripZeros(nonce.toBuffer('be', 32)), from])
    //txes.push([to, stripZeros(value.toBuffer('be', 32)), stripZeros(nonce.toBuffer('be', 32)), sigBytes])
    //txes.push([to, stripZeros(value.toBuffer('be', 32)), stripZeros(nonce.toBuffer('be', 32)), [stripZeros(txSig.r), stripZeros(txSig.s), txSig.v]])
    //txes.push([to, stripZeros(value.toBuffer('be', 32)), stripZeros(nonce.toBuffer('be', 32)), from])
  }
  // Make sure keys are unique and sort them
  const unsortedAddrs = Object.keys(toProve).map((s) => Buffer.from(s, 'hex'))
  const keys = unsortedAddrs.map((a) => keccak256(a))
  keys.sort(Buffer.compare)
  // Sort addresses based on their hashes.
  // Naive algorithm
  const sortedAddrs = new Array(keys.length).fill(undefined)
  for (const a of unsortedAddrs) {
    let idx = -1
    const h = keccak256(a)
    for (let i = 0; i < keys.length; i++) {
      const k = keys[i]
      if (h.equals(k)) {
        idx = i
      }
    }
    assert(idx >= 0)
    sortedAddrs[idx] = a
  }

  const proof = await makeMultiproof(trie, keys)
  // Verify proof is valid
  assert(await verifyMultiproof(root, proof, keys))

  // Modify txes and replace from and to addresses
  // with their index in the keys array
  for (let i = 0; i < sortedAddrs.length; i++) {
    const addr = sortedAddrs[i]
    const addrData = toProve[addr.toString('hex')]
    for (const instance of addrData) {
      txes[instance.txId][instance.fieldIdx] = i
    }
  }

  return [txes, sortedAddrs, proof, simulationData]
}

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

  const fromAcc = await getAccount(trie, from)
  const toAcc = await getAccount(trie, 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 putAccount(trie, from, fromAcc)
  await putAccount(trie, to, toAcc)
}

async function generateAccounts (trie, 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 putAccount(trie, address, account)
  }
  return accounts
}

async function putAccount (trie, address, account) {
  await promisify(trie.put.bind(trie))(address, account.serialize())
}

async function getAccount (trie, address) {
  let raw = await promisify(trie.get.bind(trie))(address)
  if (!raw) {
    return new Account()
  } else {
    return new Account(raw)
  }
}

function writeScoutConfig (data) {
  const testSuite = {
    'beacon_state': {
      'execution_scripts': [
        'target/wasm32-unknown-unknown/release/turboproof.wasm'
      ],
    },
    'shard_pre_state': {
      'exec_env_states': [
        data.preStateRoot.toString('hex')
      ]
    },
    'shard_blocks': [
      { env: 0, data: data.blockData.toString('hex') }
    ],
    'shard_post_state': {
      'exec_env_states': [
        data.postStateRoot.toString('hex')
      ]
    }
  }

  const serializedTestSuite = yaml.safeDump(testSuite)
  fs.writeFileSync('turboproof.yaml', serializedTestSuite)
}

module.exports = {
  main,
  generateTestSuite
}

test.js

const tape = require('tape')
const { main, decodeBlockData } = require('../dist/ee')
const { generateTestSuite } = require('../src/relayer/lib')

tape('turboproof ee', async t => {
  const testSuite = await generateTestSuite()
  // Only verifies multiproof, doesn't update trie
  const isValid = main({ preStateRoot: testSuite.preStateRoot, blockData: testSuite.blockData })
  t.true(isValid)
  t.end()
})