import { BIG_ENDIAN , BigNumber , getLogger , HashFunction , keccak256 , ONE , runInDomain , TWO , ZERO , } from '@pigi/core-utils' ; import * as assert from 'assert' ; import * as AsyncLock from 'async-lock' ; import * as domain from 'domain' ; import { MerkleTreeInclusionProof , MerkleTreeNode , MerkleUpdate , SparseMerkleTree , } from '../../types/block-production' ; import { DB } from '../../types/db/db.interface' ; const log = getLogger ( 'merkle-tree' ) ; /**
 * SparseMerkleTree implementation assuming a 256-bit hash algorithm is used.
 */ export class SparseMerkleTreeImpl implements SparseMerkleTree { public static readonly emptyBuffer : Buffer = Buffer . alloc ( 32 ) . fill ( '\x00' ) ; public static readonly siblingBuffer : Buffer = Buffer . alloc ( 1 ) . fill ( '\x00' ) ; private root : MerkleTreeNode ; private zeroHashes : Buffer [ ] ; private static readonly lockKey : string = 'lock' ; private readonly treeLock : AsyncLock = new AsyncLock ( { domainReentrant : true , } ) ; private readonly hashFunction : ( Buffer ) => Buffer ; private readonly hashBuffer : Buffer = Buffer . alloc ( 64 ) ; public static async create ( db , rootHash ? , height = 160 , hashFunction = keccak256 ) { assert ( ! rootHash || rootHash . length === 32 , 'Root hash must be 32 bytes' ) ; const tree = new SparseMerkleTreeImpl ( db , height , hashFunction ) ; await tree . init ( rootHash ) ; return tree ; } private constructor ( private db , private height = 160 , hashFunction = keccak256 ) { assert ( height > 0 , 'SMT height needs to be > 0' ) ; this . hashFunction = ( buff ) => Buffer . from ( hashFunction ( buff . toString ( 'hex' ) ) , 'hex' ) ; } private async init ( rootHash ? ) { await this . populateZeroHashesAndRoot ( rootHash ) ; } public getHeight ( ) { return this . height ; } public async getRootHash ( ) { const copy = Buffer . alloc ( this . root . hash . length ) ; this . root . hash . copy ( copy ) ; return copy ; } public async getLeaf ( leafKey , rootHash ? ) { log . debug ( `Trying to get leaf [${ leafKey . toString ( 10 ) }]` ) ; return this . treeLock . acquire ( SparseMerkleTreeImpl . lockKey , async ( ) => { if ( ! ! rootHash && ! rootHash . equals ( this . root . hash ) ) { log . debug ( `Cannot get Leaf [${ leafKey . toString ( 10 ) }] because root hash does not match.` ) ; return undefined ; } const nodesInPath = await this . getNodesInPath ( leafKey ) ; if ( ! nodesInPath || nodesInPath . length !== this . height ) { log . debug ( `Cannot get Leaf [${ leafKey . toString ( 10 ) }] because nodes in path does not equal tree height.` ) ; return undefined ; } const leaf = nodesInPath [ nodesInPath . length - 1 ] ; if ( ! leaf . key . equals ( leafKey ) ) { log . debug ( `Cannot get Leaf because leaf key does not match. Path: [${ leafKey . toString ( 10 ) }], leaf key: ${ leaf . key . toString ( 10 ) }.` ) ; return undefined ; } log . debug ( `Returning leaf value: [${ leaf . value . toString ( ) }].` ) ; return leaf . value ; } ) ; } public async verifyAndStore ( inclusionProof ) { if ( inclusionProof . siblings . length !== this . height - 1 ) { return false ; } return this . treeLock . acquire ( SparseMerkleTreeImpl . lockKey , async ( ) => { const leafHash = this . hashFunction ( inclusionProof . value ) ; if ( ! ! ( await this . getNode ( leafHash , inclusionProof . key ) ) ) { return true ; } let child = this . createNode ( leafHash , inclusionProof . value , inclusionProof . key ) ; let siblingIndex = 0 ; let parent = child ; const nodesToStore = [ child ] ; for ( let parentDepth = this . height - 2 ; parentDepth >= 0 ; parentDepth -- ) { child = parent ; const childDepth = parentDepth + 1 ; const childSiblingHash = inclusionProof . siblings [ siblingIndex ++ ] ; parent = this . calculateParentNode ( child , childSiblingHash , inclusionProof . key , parentDepth ) ; nodesToStore . push ( parent ) ; const siblingNode = await this . createProofSiblingNodeIfDoesntExist ( childSiblingHash , inclusionProof . key , childDepth ) ; if ( ! ! siblingNode ) { nodesToStore . push ( siblingNode ) ; } } if ( ! parent . hash . equals ( this . root . hash ) ) { return false ; } await Promise . all ( ( await this . getNodesInPath ( inclusionProof . key ) ) . map ( ( n ) => this . db . del ( this . getNodeID ( n ) ) ) ) ; this . root = parent ; await Promise . all ( nodesToStore . map ( ( n ) => this . db . put ( this . getNodeID ( n ) , n . value ) ) ) ; return true ; } ) ; } public async update ( leafKey , leafValue , d ? ) { return runInDomain ( d , async ( ) => { return this . treeLock . acquire ( SparseMerkleTreeImpl . lockKey , async ( ) => { let nodesToUpdate = await this . getNodesInPath ( leafKey ) ; if ( ! nodesToUpdate ) { return false ; } else if ( nodesToUpdate . length !== this . height ) { if ( ! ( await this . verifyAndStorePartiallyEmptyPath ( leafKey , nodesToUpdate . length ) ) ) { return false ; } nodesToUpdate = await this . getNodesInPath ( leafKey ) ; } const leaf = nodesToUpdate [ nodesToUpdate . length - 1 ] ; const idsToDelete = [ this . getNodeID ( leaf ) ] ; leaf . hash = this . hashFunction ( leafValue ) ; leaf . value = leafValue ; let updatedChild = leaf ; let depth = nodesToUpdate . length - 2 ; for ( ; depth >= 0 ; depth -- ) { idsToDelete . push ( this . getNodeID ( nodesToUpdate [ depth ] ) ) ; updatedChild = this . updateNode ( nodesToUpdate [ depth ] , updatedChild , leafKey , depth ) ; } await Promise . all ( [ ... nodesToUpdate . map ( ( n ) => this . db . put ( this . getNodeID ( n ) , n . value ) ) , ... idsToDelete . map ( ( id ) => this . db . del ( id ) ) , ] ) ; this . root = nodesToUpdate [ 0 ] ; return true ; } ) ; } ) ; } public async batchUpdate ( updates ) { const d = domain . create ( ) ; return runInDomain ( d , ( ) => { return this . treeLock . acquire ( SparseMerkleTreeImpl . lockKey , async ( ) => { for ( const update of updates ) { if ( ! ( await this . verifyAndStore ( { rootHash : this . root . hash , key : update . key , value : update . oldValue , siblings : update . oldValueProofSiblings , } ) ) ) { return false ; } } for ( const update of updates ) { if ( ! ( await this . update ( update . key , update . newValue , d ) ) ) { throw Error ( "Verify and Store worked but update didn't! This should never happen!" ) ; } } return true ; } ) ; } ) ; } public async getMerkleProof ( leafKey , leafValue ) { const result = await this . treeLock . acquire ( SparseMerkleTreeImpl . lockKey , async ( ) => { if ( ! this . root || ! this . root . hash ) { return undefined ; } let node = this . root ; const siblings = [ ] ; for ( let depth = 0 ; depth < this . height && ! ! node && ! ! node . value && node . value . length === 64 ; depth ++ ) { siblings . push ( this . getChildSiblingHash ( node , depth , leafKey ) ) ; node = await this . getChild ( node , depth , leafKey ) ; } if ( siblings . length !== this . height - 1 ) { log . info ( `Getting merkle proof for leaf that has not been stored. Storing siblings and leaf and creating proof.` ) ; return { rootHash : undefined , key : undefined , value : undefined , siblings : undefined , } ; } const leafHash = this . hashFunction ( leafValue ) ; if ( ! node . hash . equals ( leafHash ) ) { log . info ( `Node hash does not match leaf value hash for key [${ leafKey . toNumber ( ) }]. Passed leaf value: [${ leafValue . toString ( ) }], Passed leaf hash: ${ leafHash . toString ( 'hex' ) }, Node hash: ${ node . hash . toString ( 'hex' ) }` ) ; return undefined ; } return { rootHash : this . root . hash , key : leafKey , value : leafValue , siblings : siblings . reverse ( ) , } ; } ) ; if ( ! result || ! ! result . rootHash ) { return result ; } if ( leafValue . equals ( SparseMerkleTreeImpl . emptyBuffer ) ) { if ( await this . verifyAndStorePartiallyEmptyPath ( leafKey ) ) { return this . getMerkleProof ( leafKey , leafValue ) ; } } return undefined ; } public async verifyAndStorePartiallyEmptyPath ( leafKey , numExistingNodes ? ) { if ( numExistingNodes === undefined ) { numExistingNodes = ( await this . getNodesInPath ( leafKey ) ) . length ; } const existingChildren = Math . max ( numExistingNodes - 1 , 0 ) ; const siblings = [ ] ; let node = this . root ; for ( let i = 0 ; i < this . height - 1 ; i ++ ) { if ( i > existingChildren || ( i === existingChildren && ( ! node . value || node . value . length !== 64 ) ) ) { siblings . push ( ... this . zeroHashes . slice ( i + 1 ) ) ; break ; } siblings . push ( this . getChildSiblingHash ( node , i , leafKey ) ) ; node = await this . getChild ( node , i , leafKey ) ; } return this . verifyAndStore ( { rootHash : this . root . hash , key : leafKey , value : SparseMerkleTreeImpl . emptyBuffer , siblings : siblings . reverse ( ) , } ) ; } param parent param parentDepth param leafKey returns private getChildSiblingHash ( parent , parentDepth , leafKey ) { const isLeft = this . isLeft ( leafKey , parentDepth ) ; return isLeft ? parent . value . subarray ( 32 ) : parent . value . subarray ( 0 , 32 ) ; } param parent param parentDepth param leafKey returns private async getChild ( parent , parentDepth , leafKey ) { const childIndex = this . isLeft ( leafKey , parentDepth ) ? 0 : 32 ; const childHash = parent . value . subarray ( childIndex , childIndex + 32 ) ; return this . getNode ( childHash , this . getNodeKey ( leafKey , ++ parentDepth ) ) ; } param leafKey returns private async getNodesInPath ( leafKey ) { if ( ! this . root || ! this . root . hash ) { return [ ] ; } if ( ! this . root . value ) { return [ this . root ] ; } let node = this . root ; const nodesToUpdate = [ node ] ; let depth ; for ( depth = 0 ; depth < this . height - 1 ; depth ++ ) { const childDepth = depth + 1 ; if ( node . value . length === 64 ) { node = this . isLeft ( leafKey , depth ) ? await this . getNode ( node . value . subarray ( 0 , 32 ) , this . getNodeKey ( leafKey , childDepth ) ) : await this . getNode ( node . value . subarray ( 32 ) , this . getNodeKey ( leafKey , childDepth ) ) ; if ( node ) { nodesToUpdate . push ( node ) ; } else { break ; } } else { break ; } } return nodesToUpdate ; } param node param updatedChild param key param depth returns private updateNode ( node , updatedChild , key , depth ) { const isLeft = this . isLeft ( key , depth ) ; if ( isLeft ) { node . value . fill ( updatedChild . hash , 0 , 32 ) ; } else { node . value . fill ( updatedChild . hash , 32 ) ; } node . hash = this . hashFunction ( node . value ) ; return node ; } param nodeHash param leafKey param depth returns private async createProofSiblingNodeIfDoesntExist ( nodeHash , leafKey , depth ) { const nodeKey = this . getNodeKey ( leafKey , depth ) . xor ( ONE ) ; const node = await this . getNode ( nodeHash , nodeKey ) ; if ( ! ! node ) { return undefined ; } return this . createNode ( nodeHash , SparseMerkleTreeImpl . siblingBuffer , nodeKey ) ; } param nodeHash param nodeKey returns private async getNode ( nodeHash , nodeKey ) { const value = await this . db . get ( this . getNodeIDFromHashAndKey ( nodeHash , nodeKey ) ) ; if ( ! value ) { return undefined ; } return this . createNode ( nodeHash , value , nodeKey ) ; } param node param siblingHash param leafKey param depth returns private calculateParentNode ( node , siblingHash , leafKey , depth ) { const value = Buffer . alloc ( 64 ) ; if ( this . isLeft ( leafKey , depth ) ) { this . hashBuffer . fill ( node . hash , 0 , 32 ) . fill ( siblingHash , 32 ) . copy ( value ) ; } else { this . hashBuffer . fill ( siblingHash , 0 , 32 ) . fill ( node . hash , 32 ) . copy ( value ) ; } return this . createNode ( this . hashFunction ( value ) , value , this . getNodeKey ( leafKey , depth ) ) ; } param rootHash private async populateZeroHashesAndRoot ( rootHash ? ) { const hashes = [ this . hashFunction ( SparseMerkleTreeImpl . emptyBuffer ) , ] ; for ( let i = 1 ; i < this . height ; i ++ ) { hashes [ i ] = this . hashFunction ( this . hashBuffer . fill ( hashes [ i - 1 ] , 0 , 32 ) . fill ( hashes [ i - 1 ] , 32 ) ) ; } this . zeroHashes = hashes . reverse ( ) ; if ( ! ! rootHash ) { log . info ( `Attempting to initialize SMT with root hash ${ rootHash . toString ( 'hex' ) }` ) ; this . root = await this . getNode ( rootHash , ZERO ) ; } if ( ! this . root ) { this . root = this . createNode ( rootHash || this . zeroHashes [ 0 ] , undefined , ZERO ) ; log . info ( `Initialized Sparse Merkle Tree with root ${ ( rootHash || this . zeroHashes [ 0 ] ) . toString ( 'hex' ) }` ) ; } else { log . info ( `Initialized Sparse Merkle Tree with root ${ this . root . hash . toString ( 'hex' ) }` ) ; } } param hash param value param key returns private createNode ( hash , value , key ) { return { hash , value , key } ; } param key param depth returns private isLeft ( key , depth ) { return key . shiftLeft ( depth ) . shiftRight ( this . height - 2 ) . mod ( TWO ) . equals ( ZERO ) ; } private getNodeKey ( leafKey , depth ) { return leafKey . shiftRight ( this . height - depth - 1 ) ; } param node private getNodeID ( node ) { const id = this . getNodeIDFromHashAndKey ( node . hash , node . key ) ; log . debug ( `Node ID for key [${ node . key }] is ${ id . toString ( 'hex' ) }. (node hash: ${ node . hash . toString ( 'hex' ) }` ) ; return id ; } private getNodeIDFromHashAndKey ( nodeHash , nodeKey ) { return this . hashFunction ( this . hashBuffer . fill ( nodeHash , 0 , 32 ) . fill ( this . hashFunction ( nodeKey . toBuffer ( BIG_ENDIAN ) ) , 32 ) ) ; } }