merkletreets/merkletree.ts

merkletree for freeserver

/**
 * 默克尔树
 * Copyright(C) 2018 liumurong
 */

import * as crypto from 'crypto';

/**
 * 选项参数
 */
export interface MerkleOptions {
	/**
	 * 叶子节点的个数
	 */
	leafcount: number;
	hashTree?: MerkleTreeOption;
}

/**
 * 树的选项参数
 */
export interface MerkleTreeOption {
	leaves: string[];
	_depth: number;
	rows: string[][];
}

/**
 * 树节点
 */
export interface MerkleNode {
    hash: string,
    left: MerkleNode | null,
    right: MerkleNode | null,
    depth: number,
    index: number
}

/**
 * MerkleTree
 */
export class MerkleTree {

    /**
     * 叶节点数组
     */
    private leaves: string[];
    /**
     * 树高度
     */
    private _depth: number;
    /**
     * 树节点
     */
    private rows: string[][];

    /**
     * MerkleTree
     * @param opts 选项
     */
    constructor(opts: MerkleOptions = { leafcount: 0 }) {
        this.leaves = opts.hashTree ? opts.hashTree.leaves : new Array(opts.leafcount).fill(this.hashFunc("0x"));  // 默认填充0x字符
        this._depth = opts.hashTree ? opts.hashTree._depth : 0;
        this.rows = opts.hashTree ? opts.hashTree.rows : [];
    }
    /**
     * 添加节点
     * @param data 叶子节点数据
     * @param hashed 是否是Hash值
     */
    addLeaf(data: Buffer | string, hashed: boolean = false) {
        if (hashed === true) {
            this.leaves.push(data.toString());
            return;
        }
        this.leaves.push(this.hashFunc(data))
    }
    /**
     * 更新叶子节点
     * @param index 叶子节点索引
     * @param data 数据
     * @param hashed 数据是否直接是Hash值
     */
    updateLeaf(index: number, data: Buffer | string, hashed: boolean = false) {
        let leafLength = this.leaves.length;
        if (leafLength <= index) {  // 索引超范围
            return;
        }
        this.leaves[index] = hashed === true ? data.toString() : this.hashFunc(data);
    }

    /**
     * 重新计算整棵树
     * 当更新过叶子节点或者添加新节点之后需要重新更新整棵树
     * 使用leafcount初始化的树也需要调用此方法
     */
    update() {
        let depth = this.depth(true);
        this.rows = new Array(depth);
        this.rows[depth] = this.leaves;
        for (let i = depth - 1; i >= 0; i--) {
            this.rows[i] = this.getNodes(this.rows[i + 1])
        }
    }
    /**
     * 据叶子节点获取上一级节点
     * @param leaves 叶子节点
     */
    private getNodes(leaves: string[]) {
        let remainder = leaves.length % 2;
        let nodes: string[] = [];
        for (let i = 0; i < leaves.length - 1; i = i + 2) {
            nodes[i / 2] = this.hashFunc(leaves[i] + leaves[i + 1]);
        }
        if (remainder === 1) {
            // 将最后一个节点提升等级，并未复制一个新的
            nodes[((leaves.length - remainder) / 2)] = leaves[leaves.length - 1];
        }
        return nodes;
    }

    /**
     * 获取Root节点
     */
    root(): MerkleNode {
        return <MerkleNode>this.getTreeNode(0, 0);
    }

    /**
     * 获取树节点
     * @param depth 深度
     * @param index 索引(本层)
     */
    getTreeNode(depth: number, index: number): MerkleNode | null {
        if (depth < 0 || depth > this.depth()) {
            return null;
        }
        let levelNodes = this.getLevelNodes(depth);
        if (index < 0 || index >= levelNodes.length) {
            return null;
        }
        let nextLevelNodes = this.getLevelNodes(depth + 1);
        let nextLevelIndex = index * 2;
        return {
            hash: levelNodes[index],
            depth: depth,
            index: index,
            left: {
                hash: nextLevelNodes[nextLevelIndex],
                depth: depth + 1,
                index: nextLevelIndex,
                left: null,
                right: null
            },
            right: nextLevelIndex + 1 >= nextLevelNodes.length ? null : {
                hash: nextLevelNodes[nextLevelIndex + 1],
                depth: depth + 1,
                index: nextLevelIndex + 1,
                left: null,
                right: null
            }
        }
    }

    /**
     * 树高
     * @param force 强制刷新树高
     */
    depth(force: boolean = false): number {
        if (!this._depth || force === true) {
            let pow = 0;
            while (Math.pow(2, pow) < this.leaves.length) {
                pow++;
            }
            this._depth = pow;
        }
        return this._depth;
    }
    /**
     * 获取此级别的所有节点
     * @param level 级别
     */
    getLevelNodes(level: number): string[] {
        if (level >= this.rows.length) {
            return [];
        }
        return this.rows[level];
    }

    /**
    * 获取叶子节点的hash值
    * @param index 叶子节点的索引
    */
    getLeafHash(index: number) {
        return this.leaves[index];
    }

    /**
     * 计算数据的Hash值，采用sha256算法
     * @param data 数据
     */
    hashFunc(data: Buffer | string): string {
        let hash = crypto.createHash("sha256");
        if (Buffer.isBuffer(data)) {
            return hash.update(data).digest("hex");
        }
        return hash.update(Buffer.from(data)).digest("hex");
    }
}