Skip to content

Commit

Permalink
extend merkle_tree implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
tungng98 committed Feb 16, 2023
1 parent bf8eec4 commit deab806
Show file tree
Hide file tree
Showing 4 changed files with 80 additions and 16 deletions.
14 changes: 9 additions & 5 deletions core/src/core/hash.service.ts
Original file line number Diff line number Diff line change
@@ -1,13 +1,17 @@
import { SHA256 } from 'crypto-js'
import create from 'keccak'
import { SHA256 } from 'crypto-js';
import create from 'keccak';

export class HashService {

static keckka256(input: string | Buffer): Buffer {
return create('keccak256').update(input).digest()
return create('keccak256').update(input).digest();
}

static sha256(message: string): Buffer {
return Buffer.from(SHA256(message).toString(), 'hex')
static sha256(input: string | Buffer): Buffer {
if (typeof input == 'string') {
return Buffer.from(SHA256(input).toString(), 'hex');
}
const words = CryptoJS.enc.Hex.parse(input.toString('hex'));
return Buffer.from(SHA256(words).toString(), 'hex');
}
}
2 changes: 1 addition & 1 deletion core/src/core/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,6 @@ export { Ed25519SignService } from './ed25519_sign.service';
export { HashService } from './hash.service';
export { IdlParserService } from './idl_parser.service';
export { InstructionLog, ProgramLogCategory, SignatureTuple, TransactionLog } from './interfaces';
export { MerkleNode, MerkleTree } from './merkle_tree';
export { MerkleNode, MerkleTreeKeccak, MerkleTreeSha256 } from './merkle_tree';
export { SolanaService } from './solana.service';
export { DEFAULT_PUBKEY, getProgramReturn, sendTransaction } from './solana_web3.service';
78 changes: 69 additions & 9 deletions core/src/core/merkle_tree.ts
Original file line number Diff line number Diff line change
@@ -1,20 +1,68 @@
import { HashService } from "./hash.service";
import { HashService } from './hash.service';

const LEVEL_ARRAY = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M']
const SIZE_ARRAY = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
const LEVEL_ARRAY = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M'];
const SIZE_ARRAY = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096];

export interface MerkleNode {
row: string;
index: number;
hash: Buffer;
}

export class MerkleTree {
export class MerkleTreeKeccak {
private _tree: MerkleTree;

constructor(hashes: Buffer[]) {
this._tree = new MerkleTree(hashes, HashService.keckka256);
}

height(): number {
return this._tree.height();
}

nodes(): MerkleNode[][] {
return this._tree.nodes();
}

root(): MerkleNode {
return this._tree.root();
}

proofs(index: number): MerkleNode[] {
return this._tree.proofs(index);
}
}

export class MerkleTreeSha256 {
private _tree: MerkleTree;

constructor(hashes: Buffer[]) {
this._tree = new MerkleTree(hashes, HashService.sha256);
}

height(): number {
return this._tree.height();
}

nodes(): MerkleNode[][] {
return this._tree.nodes();
}

root(): MerkleNode {
return this._tree.root();
}

proofs(index: number): MerkleNode[] {
return this._tree.proofs(index);
}
}

class MerkleTree {
private _height: number;
private _nodes: MerkleNode[][];
private _root: MerkleNode;

constructor(hashes: Buffer[]) {
constructor(hashes: Buffer[], hashFn: (input: Buffer) => Buffer) {
// detect tree height
for (let i = 0; i < SIZE_ARRAY.length; i++) {
if (SIZE_ARRAY[i] >= hashes.length) {
Expand Down Expand Up @@ -42,10 +90,9 @@ export class MerkleTree {
for (let j = 0; j < subNodes.length; j += 2) {
const hash0: Buffer = subNodes[j].hash;
const hash1: Buffer = subNodes[j + 1].hash;
const newHash: Buffer =
hash0 <= hash1
? HashService.keckka256(Buffer.concat([hash0, hash1]))
: HashService.keckka256(Buffer.concat([hash1, hash0]));
const newHash: Buffer = hash0.compare(hash1) <= 0
? hashFn(Buffer.concat([hash0, hash1]))
: hashFn(Buffer.concat([hash1, hash0]));
newNodes.push(<MerkleNode>{
row: LEVEL_ARRAY[i],
index: j / 2,
Expand All @@ -69,4 +116,17 @@ export class MerkleTree {
root(): MerkleNode {
return this._root;
}

proofs(index: number): MerkleNode[] {
const nodes = this._nodes;
const proofs: MerkleNode[] = [];
let currentIndex = index;
for (let i = 0; i < nodes.length - 1; i++) {
const proof = currentIndex % 2 == 0 ? nodes[i][currentIndex + 1] : nodes[i][currentIndex - 1];
currentIndex = (currentIndex - (currentIndex % 2)) / 2;
proofs.push(proof);
}

return proofs;
}
}
2 changes: 1 addition & 1 deletion core/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ export { Ed25519SignService } from './core/ed25519_sign.service';
export { HashService } from './core/hash.service';
export { IdlParserService } from './core/idl_parser.service';
export { InstructionLog, ProgramLogCategory, SignatureTuple, TransactionLog } from './core/interfaces';
export { MerkleNode, MerkleTree } from './core/merkle_tree';
export { MerkleNode, MerkleTreeKeccak, MerkleTreeSha256 } from './core/merkle_tree';
export { SolanaService } from './core/solana.service';
export { DEFAULT_PUBKEY, getProgramReturn, sendRawTransaction, sendRawTransaction2, sendTransaction, sendTransaction2 } from './core/solana_web3.service';
export { DurableNonceService } from './durable_nonce.service';
Expand Down

0 comments on commit deab806

Please sign in to comment.