From 670a4583546b4269ecf8c74c61d7cdfd6e0282c4 Mon Sep 17 00:00:00 2001 From: Miguel Mota Date: Thu, 17 Nov 2022 14:16:03 -0800 Subject: [PATCH] Add bigNumerify util and concatFn option. #71 --- package.json | 3 ++- src/Base.ts | 32 ++++++++++++++++++++++++++++++++ src/MerkleTree.ts | 37 +++++++++++++++++++++++-------------- test/Base.test.js | 12 +++++++++++- test/MerkleTree.test.js | 25 +++++++++++++++++++++++++ 5 files changed, 93 insertions(+), 16 deletions(-) diff --git a/package.json b/package.json index 7e85b70..1d08d3d 100644 --- a/package.json +++ b/package.json @@ -87,6 +87,7 @@ "parser": "@typescript-eslint/parser", "plugins": [ "@typescript-eslint/eslint-plugin" - ] + ], + "globals": ["BigInt"] } } diff --git a/src/Base.ts b/src/Base.ts index 946ae24..7c2ba00 100644 --- a/src/Base.ts +++ b/src/Base.ts @@ -167,6 +167,8 @@ export class Base { return Buffer.from(value.replace(/^0x/, ''), 'hex') } else if (typeof value === 'string') { return Buffer.from(value) + } else if (typeof value === 'bigint') { + return Buffer.from(value.toString(16), 'hex') } else if (typeof value === 'number') { let s = value.toString() if (s.length % 2) { @@ -181,6 +183,32 @@ export class Base { return value } + bigNumberify (value: any): BigInt { + return Base.bigNumberify(value) + } + + static bigNumberify (value: any): BigInt { + if (typeof value === 'bigint') { + return value + } + + if (typeof value === 'string') { + if (Base.isHexString(value)) { + return BigInt('0x' + value.replace('0x', '').toString()) + } + } + + if (Buffer.isBuffer(value)) { + return BigInt('0x' + value.toString('hex')) + } + + if (typeof value === 'number') { + return BigInt(value) + } + + throw new Error('cannot bigNumberify') + } + /** * isHexString * @desc Returns true if value is a hex string. @@ -281,6 +309,10 @@ export class Base { return Buffer.from(v) } + if (typeof v === 'bigint') { + return Buffer.from(value.toString(16), 'hex') + } + if (ArrayBuffer.isView(v)) { return Buffer.from(v.buffer, v.byteOffset, v.byteLength) } diff --git a/src/MerkleTree.ts b/src/MerkleTree.ts index a961992..d446f81 100644 --- a/src/MerkleTree.ts +++ b/src/MerkleTree.ts @@ -7,7 +7,7 @@ import Base from './Base' // TODO: Clean up and DRY up code // Disclaimer: The multiproof code is unaudited and may possibly contain serious issues. It's in a hacky state as is and it's begging for a rewrite! -type TValue = Buffer | string | number | null | undefined +type TValue = Buffer | BigInt | string | number | null | undefined type THashFnResult = Buffer | string type THashFn = (value: TValue) => Buffer type TLeaf = Buffer @@ -31,6 +31,8 @@ export interface Options { fillDefaultHash?: TFillDefaultHash | Buffer | string /** If set to `true`, the resulting tree will be a complete tree. Recommended for use of multiProofs. */ complete?: boolean; + + concatFn?: any } /** @@ -40,6 +42,7 @@ export interface Options { export class MerkleTree extends Base { private duplicateOdd: boolean = false private hashFn: THashFn + private concatFn: any = Buffer.concat private hashLeaves: boolean = false private isBitcoinTree: boolean = false private leaves: TLeaf[] = [] @@ -108,6 +111,10 @@ export class MerkleTree extends Base { this.duplicateOdd = !!options.duplicateOdd + if (options.concatFn) { + this.concatFn = options.concatFn + } + this.hashFn = this.bufferifyFn(hashFn) this.processLeaves(leaves) } @@ -163,14 +170,13 @@ export class MerkleTree extends Base { break } else if (i + 1 === nodes.length) { if (nodes.length % 2 === 1) { - let data = nodes[nodes.length - 1] + const data = nodes[nodes.length - 1] let hash = data // is bitcoin tree if (this.isBitcoinTree) { // Bitcoin method of duplicating the odd ending nodes - data = Buffer.concat([reverse(data), reverse(data)]) - hash = this.hashFn(data) + hash = this.hashFn(this.combineHashes([reverse(data), reverse(data)])) hash = reverse(this.hashFn(hash)) this.layers[layerIndex].push(hash) @@ -189,7 +195,6 @@ export class MerkleTree extends Base { const left = nodes[i] const right = i + 1 === nodes.length ? left : nodes[i + 1] - let data = null let combined = null if (this.isBitcoinTree) { @@ -202,8 +207,7 @@ export class MerkleTree extends Base { combined.sort(Buffer.compare) } - data = Buffer.concat(combined) - let hash = this.hashFn(data) + let hash = this.hashFn(this.combineHashes(combined)) // double hash if bitcoin tree if (this.isBitcoinTree) { @@ -1126,21 +1130,21 @@ export class MerkleTree extends Base { buffers[isLeftNode ? 'unshift' : 'push'](reverse(data)) - hash = this.hashFn(Buffer.concat(buffers)) + hash = this.hashFn(this.combineHashes(buffers)) hash = reverse(this.hashFn(hash)) } else { if (this.sortPairs) { if (Buffer.compare(hash, data) === -1) { buffers.push(hash, data) - hash = this.hashFn(Buffer.concat(buffers)) + hash = this.hashFn(this.combineHashes(buffers)) } else { buffers.push(data, hash) - hash = this.hashFn(Buffer.concat(buffers)) + hash = this.hashFn(this.combineHashes(buffers)) } } else { buffers.push(hash) buffers[isLeftNode ? 'unshift' : 'push'](data) - hash = this.hashFn(Buffer.concat(buffers)) + hash = this.hashFn(this.combineHashes(buffers)) } } } @@ -1199,7 +1203,7 @@ export class MerkleTree extends Base { pair = pair.sort(Buffer.compare) } - const hash = pair[1] ? this.hashFn(Buffer.concat(pair)) : pair[0] + const hash = pair[1] ? this.hashFn(this.combineHashes(pair)) : pair[0] tree[(index / 2) | 0] = hash indexqueue.push((index / 2) | 0) } @@ -1227,7 +1231,7 @@ export class MerkleTree extends Base { const bufA: Buffer = proofFlag[i] ? (leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++]) : proofs[proofPos++] const bufB : Buffer = leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++] const buffers = [bufA, bufB].sort(Buffer.compare) - hashes[i] = this.hashFn(Buffer.concat(buffers)) + hashes[i] = this.hashFn(this.combineHashes(buffers)) } return Buffer.compare(hashes[totalHashes - 1], root) === 0 @@ -1438,7 +1442,7 @@ export class MerkleTree extends Base { const parentNodeTreeIndex = parentIndices[i] const bufA = currentLayer[i * 2] const bufB = currentLayer[i * 2 + 1] - const hash = bufB ? this.hashFn(Buffer.concat([bufA, bufB])) : bufA + const hash = bufB ? this.hashFn(this.combineHashes([bufA, bufB])) : bufA parentLayer.push([parentNodeTreeIndex, hash]) } @@ -1447,6 +1451,11 @@ export class MerkleTree extends Base { return tree[tree.length - 1][0][1] } + + combineHashes (hashes: Buffer[]) { + hashes = hashes.filter(x => x.byteLength > 0) + return this.concatFn(hashes) + } } if (typeof window !== 'undefined') { diff --git a/test/Base.test.js b/test/Base.test.js index 2e61571..219826a 100644 --- a/test/Base.test.js +++ b/test/Base.test.js @@ -2,13 +2,14 @@ const test = require('tape') const { Base } = require('../dist/Base') test('bufferify', t => { - t.plan(4) + t.plan(5) const base = new Base() t.deepEqual(base.bufferify(''), Buffer.alloc(0)) t.deepEqual(base.bufferify('0x123'), Buffer.from('123', 'hex')) t.deepEqual(base.bufferify('123'), Buffer.from('123', 'hex')) t.deepEqual(base.bufferify(Buffer.from('123')), Buffer.from('123')) + t.deepEqual(base.bufferify(BigInt('0x123')), Buffer.from('123', 'hex')) }) test('bufferifyFn', t => { @@ -21,6 +22,15 @@ test('bufferifyFn', t => { t.deepEqual(fn('XYZ'), Buffer.from('XYZ')) }) +test('bigNumberify', t => { + t.plan(3) + + const base = new Base() + t.deepEqual(base.bigNumberify(123), BigInt(123)) + t.deepEqual(base.bigNumberify('0x123'), BigInt('0x123')) + t.deepEqual(base.bigNumberify(BigInt(123)), BigInt(123)) +}) + test('binarySearch', t => { t.plan(3) diff --git a/test/MerkleTree.test.js b/test/MerkleTree.test.js index 585374c..c1f26b0 100644 --- a/test/MerkleTree.test.js +++ b/test/MerkleTree.test.js @@ -9,6 +9,7 @@ const SHA256 = require('crypto-js/sha256') const SHA3 = require('crypto-js/sha3') const sha1 = require('sha1') const { keccak256: ethCryptoKeccak256 } = require('ethereum-cryptography/keccak') +const { poseidon } = require('circomlibjs') const { MerkleTree } = require('../') @@ -1322,3 +1323,27 @@ test.skip('1M leaves keccak256', t => { t.true(tree.verify(proof, keccak256(`${max - 1}`), tree.getRoot())) t.false(tree.verify(proof, keccak256(`${max + 1}`), tree.getRoot())) }) + +test('poseidon hash', t => { + t.plan(4) + + const _poseidon = (elements) => { + const bigInts = elements.map(MerkleTree.bigNumberify) + return Buffer.from(poseidon(bigInts).toString(16), 'hex') + } + + const result1 = _poseidon([1, 2].map(MerkleTree.bufferify)) + t.equal(result1.toString('hex'), '115cc0f5e7d690413df64c6b9662e9cf2a3617f2743245519e19607a4417189a') + + const result2 = _poseidon([1, 2, 3, 4].map(MerkleTree.bufferify)) + t.equal(result2.toString('hex'), '299c867db6c1fdd79dcefa40e4510b9837e60ebb1ce0663dbaa525df65250465') + + const leaves = [1, 2, 3, 4].map(x => _poseidon([x])) + const tree = new MerkleTree(leaves, _poseidon, { + concatFn: (hashes) => hashes + }) + t.equal(tree.getHexRoot(), '0xd24e045226875e22b37ce607ea2af7a9fbb137ee128caa0ce3663615350245') + + const proof = tree.getProof(leaves[2]) + t.true(tree.verify(proof, leaves[2], tree.getHexRoot())) +})