merkle.ts

import Decimal from "decimal.js";

import { webcrypto } from "crypto"
const { subtle } = webcrypto;

// Required to provide precise computation of currency balances
Decimal.set({
    precision: 1000,
    toExpNeg: -1000,
    toExpPos: 1000
})

// Core hash function taking a string as input and returning an hexstring hash value
async function hash(input: string) {
    const ec = new TextEncoder();
    const digest = await subtle.digest('SHA-256', ec.encode(input));
    return Buffer.from(digest).toString('hex');
}

// Concatenation with ',' separator of strings
// Non-async so that we can inline such calls
function concat(...args: (string)[]) {
    return Array.prototype.slice.call(args).join(',');
}

// Core Key derivation function
// Input parameters key, ctx and output are strings
async function kdf(key: string, ctx: string) {
    const ec = new TextEncoder();
    const keyData = await subtle.importKey('raw', ec.encode(key), {
        name: 'HMAC',
        hash: 'SHA-256',
    }, true, ['sign', 'verify']);

    const digest = await subtle.sign({
        name: 'HMAC'
    }, keyData, ec.encode(ctx));

    return Buffer.from(digest).toString('hex');
}

// Derivation of Audit nonce
// userCred, auditId and output are strings
async function auditNonceDerive(userCred: string, auditId: string) {
    return await kdf(userCred, concat('SB PoL - Audit nonce derivation', auditId));
}

// Derivation of Merkle tree nonce
// auditNonce, currency and output are strings
async function merkleTreeNonceDerive(auditNonce: string, currency: string) {
    return await kdf(auditNonce, concat('SB PoL - Currency Merkle tree nonce derivation', currency));
}

// Derivation of leaf nonce
// mtNonce is a string, leafIndex is an integer (js number), output is a string
async function leafNonceDerive(mtNonce: string, leafIndex: number) {
    return await kdf(mtNonce, concat('SB PoL - Leaf nonce derivation', leafIndex.toString()));
}

// Leaf hash computation
// leafNonce and userId are string and liability is Decimal
async function leafHashDerive(leafNonce: string, userId: string, liability: Decimal) {
    return await hash(concat('leaf', leafNonce, userId, liability.toString()));
}

// Derivation of parent hash based on children
// leftBal: Decimal, rightBal: Decimal, leftHash: str, rightHash: str, output: str
async function innerHashDerive(leftBal: Decimal, rightBal: Decimal, leftHash: string, rightHash: string) {
    return await hash(concat('inner', leftBal.toString(), rightBal.toString(), leftHash, rightHash))
}

// Direct derivation of leaf hash from all user-related parameters
// All inputs are strings except leafIndex is an integer (js number) and liability is Decimal
export async function fullLeafHashDerive(userCred: string, auditId: string, currency: string, leafIndex: number, userId: string, liability: Decimal) {
    const auditNonce = await auditNonceDerive(userCred, auditId);
    const mtNonce = await merkleTreeNonceDerive(auditNonce, currency);
    const leafNonce = await leafNonceDerive(mtNonce, leafIndex);

    return await leafHashDerive(leafNonce, userId, liability);
}

interface NodeLeaf {
    liability: Decimal
    digest: string
}

interface NodeLeafWithCurrency extends NodeLeaf {
    currency: string
}

interface MaybeEmptyNodeLeaf {
    liability: Decimal | null
    digest: string | null
}

// Verification of the audit commitment hash derived from the Merkle tree root nodes (audit partition commitments)
// Return boolean
export async function audCommHashCheck(partIdToData: Map<string, NodeLeafWithCurrency>, auditHash: string) {
    let input = 'SB PoL - Audit commitment hash';
    const auditIdToDataEntries = Array.from(partIdToData, ([auditId, value]) => ({ auditId, value }));
    auditIdToDataEntries.sort((a, b) => {
        if (a.value.currency.toUpperCase() < b.value.currency.toUpperCase()) { 
            return - 1
        } else {
            return 1
        }
    })

    auditIdToDataEntries.forEach(
        (node) => { input = concat(input, node.auditId.toString(), node.value.liability.toString(), node.value.digest); }
    )

    const computedHash = await hash(input);

    if (auditHash !== computedHash) {
        console.log('The audit commitment hash does not match.');
        console.log('Audit Hash: ' + auditHash);
        console.log('Computed Hash: ' + computedHash);
        return false;
    }

    return true;
}

// Build the Merkle tree and output an array of 'NodeLeaf' (liability: , digest: )
// Input: array of elements of type 'NodeLeaf'
export async function buildMerkleTree(leaves: NodeLeaf[]) {
    const n = leaves.length;
    const tree: MaybeEmptyNodeLeaf[] = leaves.map(a => ({ ...a }));

    for (let i = 0; i < n - 1; i++) {
        tree.unshift(
            {
                liability: null,
                digest: null
            });
    }

    for (let i = n - 2; i >= 0; i--) {
        const left = tree[2 * i + 1];
        const right = tree[2 * i + 2];
        if (left.liability !== null && right.liability !== null && left.digest !== null && right.digest !== null) {
            tree[i].digest = await innerHashDerive(left.liability, right.liability, left.digest, right.digest);
            tree[i].liability = left.liability.plus(right.liability);
        }
    }

    return tree;
}

// leafNode: NodeLeaf object ({liability: Decimal , digest: hexString})
// leafIndex: position of the leaf (zero-based and left-to-right numbering)
// witnesses: array of node objects to reconstruct path up to the root. Root is included here.
// Warning: No validation that the root hash is correct here, i.e., with respect to the audit commitment hash.
// Output: boolean
export async function validateProof(leafNode: NodeLeaf, leafIndex: number, witnesses: NodeLeaf[]) {
    const d = witnesses.length - 1;
    if (d < 1) {
        console.log("The number of witnesses is too small.");
        return false;
    }

    let nodeLia = leafNode.liability;
    if (!nodeLia.isPositive()) {
        console.log("The user liability is non-positive.");
        return false;
    }

    let nodeHash = leafNode.digest;
    let treeIdx = leafIndex + Math.pow(2, d) - 1;

    for (let i = 0; i < d; i++) {
        const witLia = witnesses[i].liability;
        const witHash = witnesses[i].digest;

        if (!witLia.isPositive()) {
            console.log("The witness " + i + "is non-positive.");
            return false;
        }

        if (treeIdx % 2 == 1) {
            nodeHash = await innerHashDerive(nodeLia, witLia, nodeHash, witHash);
        } else {
            nodeHash = await innerHashDerive(witLia, nodeLia, witHash, nodeHash);
        }

        nodeLia = nodeLia.plus(witLia);
        treeIdx = Math.floor((treeIdx - 1) / 2);
    }

    if (!nodeLia.equals(witnesses[d].liability)) {
        console.log("Merkle Root tree liability does not match.")
        return false;
    }

    if (nodeHash !== witnesses[d].digest) {
        console.log("Merkle Root tree hash does not match.")
        console.log("Computed hash: " + nodeHash);
        console.log("Provided hash: " + witnesses[d].digest);
        return false;
    }

    return true;
}

export interface AuditRootInterface {
    audit: {
        commitment: {
            digest: string
        },
        id: string
        partitions: [{
            commitment: NodeLeaf
            currency: string
            id: string
        },]
        time: string
    }
}

export interface UserLiabilitiesInterface {
    user_id: string,
    audit_id: string,
    liabilities: {
        audit_partition_id: string
        liability: Decimal
        proof: {
            leaf_index: number
            nonce: string
            witnesses: NodeLeaf[]
        }
    }[]
}

// Inputs: user credential as string
//         JSON object containing all liabilities of a user
//         JSON object containing all audit-level data such as audit commitment hash and Merkle tree root nodes.
// Outputs: array with 2 elements. One boolean telling whether the liabilities are valid.
//          A second one which is a hashmap between currencies and user liability.
export async function validateUserLiabilities(userCred: string, auditRootData: AuditRootInterface, userLiabilities: UserLiabilitiesInterface) {

    // Check that the auditId is the same in both JSON files.
    if (userLiabilities['audit_id'] !== auditRootData['audit']['id']) {
        console.log("Both auditIds do not match.");
        return [false, null];
    }

    const auditHash = auditRootData['audit']['commitment']['digest'];
    const auditPartitions = auditRootData['audit']['partitions'];

    // Populate a Map between an audit partition id (identifier for a Merkle tree of a given currency)
    // and an object containing the Merkle tree relevant data: (currency, liability, digest).
    const partIdToData = new Map();
    auditPartitions.forEach(
        part => partIdToData.set(part['id'], {
            currency: part['currency'],
            liability: part['commitment']['liability'],
            digest: part['commitment']['digest']
        })
    );

    // Validation of all Merkle tree commitments against the audit commitment hash.
    const auditComOk = await audCommHashCheck(partIdToData, auditHash);

    if (!auditComOk) {
        return [false, null];
    }

    const userOutput = new Map();
    const auditId = userLiabilities['audit_id'];
    const userId = userLiabilities['user_id'];

    // Process each user liability: validate and accumulate into userOutput Map.
    for (const liab of userLiabilities['liabilities']) {
        const partitionData = partIdToData.get(liab['audit_partition_id']);
        const currency = partitionData.currency;

        // Construct the leaf node.
        const leafHash = await fullLeafHashDerive(
            userCred,
            auditId,
            currency,
            liab['proof']['leaf_index'],
            userId,
            liab['liability']
        );

        const leafNode = {
            liability: new Decimal(liab['liability']),
            digest: leafHash
        };

        const witnesses = liab['proof']['witnesses'].map(
            val => ({
                liability: new Decimal(val['liability']),
                digest: val['digest']
            }));

        // Append the root of Merkle tree (partition) to the list of witnesses.
        witnesses.push(
            {
                liability: new Decimal(partitionData.liability),
                digest: partitionData.digest
            }
        )

        // Check inclusion proof for this leaf node
        const res = await validateProof(leafNode, liab['proof']['leaf_index'], witnesses);

        if (!res) {
            return [false, null];
        }

        // Accumulate liability
        if (userOutput.has(currency)) {
            userOutput.set(currency, userOutput.get(currency).plus(leafNode.liability));
        } else {
            userOutput.set(currency, leafNode.liability);
        }
    }

    return [true, userOutput];
}