From 367ea4a0a8a59ade42f20d1ef638f84006e5666c Mon Sep 17 00:00:00 2001 From: Arseniy Klempner Date: Mon, 15 Dec 2025 16:26:55 -0800 Subject: [PATCH] fix: consolidate util functions --- packages/rln/src/proof.spec.ts | 26 +++----- packages/rln/src/utils/bytes.ts | 24 ++++--- packages/rln/src/utils/merkle.ts | 109 +++---------------------------- packages/rln/src/zerokit.ts | 7 +- 4 files changed, 35 insertions(+), 131 deletions(-) diff --git a/packages/rln/src/proof.spec.ts b/packages/rln/src/proof.spec.ts index 5d15798283..61b084e617 100644 --- a/packages/rln/src/proof.spec.ts +++ b/packages/rln/src/proof.spec.ts @@ -5,7 +5,7 @@ import { RLNInstance } from "./rln.js"; import { BytesUtils } from "./utils/index.js"; import { calculateRateCommitment, - extractPathDirectionsFromProof, + getPathDirectionsFromIndex, MERKLE_TREE_DEPTH, reconstructMerkleRoot } from "./utils/merkle.js"; @@ -15,24 +15,23 @@ describe("RLN Proof Integration Tests", function () { this.timeout(30000); it("validate stored merkle proof data", function () { - // Convert stored merkle proof strings to bigints const merkleProof = TEST_KEYSTORE_DATA.merkleProof.map((p) => BigInt(p)); expect(merkleProof).to.be.an("array"); - expect(merkleProof).to.have.lengthOf(MERKLE_TREE_DEPTH); // RLN uses fixed depth merkle tree + expect(merkleProof).to.have.lengthOf(MERKLE_TREE_DEPTH); - merkleProof.forEach((element, i) => { + for (let i = 0; i < merkleProof.length; i++) { + const element = merkleProof[i]; expect(element).to.be.a( "bigint", `Proof element ${i} should be a bigint` ); expect(element).to.not.equal(0n, `Proof element ${i} should not be zero`); - }); + } }); it("should generate a valid RLN proof", async function () { const rlnInstance = await RLNInstance.create(); - // Load credential from test keystore const keystore = Keystore.fromString(TEST_KEYSTORE_DATA.keystoreJson); if (!keystore) { throw new Error("Failed to load test keystore"); @@ -53,14 +52,7 @@ describe("RLN Proof Integration Tests", function () { const rateCommitment = calculateRateCommitment(idCommitment, rateLimit); - const proofElementIndexes = extractPathDirectionsFromProof( - merkleProof, - rateCommitment, - merkleRoot - ); - if (!proofElementIndexes) { - throw new Error("Failed to extract proof element indexes"); - } + const proofElementIndexes = getPathDirectionsFromIndex(membershipIndex); expect(proofElementIndexes).to.have.lengthOf(MERKLE_TREE_DEPTH); @@ -82,7 +74,9 @@ describe("RLN Proof Integration Tests", function () { Number(membershipIndex), new Date(), credential.identity.IDSecretHash, - merkleProof.map((proof) => BytesUtils.fromBigInt(proof, 32, "little")), + merkleProof.map((element) => + BytesUtils.bytes32FromBigInt(element, "little") + ), proofElementIndexes.map((index) => BytesUtils.writeUIntLE(new Uint8Array(1), index, 0, 1) ), @@ -94,7 +88,7 @@ describe("RLN Proof Integration Tests", function () { BytesUtils.writeUIntLE(new Uint8Array(8), testMessage.length, 0, 8), testMessage, proof, - [BytesUtils.fromBigInt(merkleRoot, 32, "little")] + [BytesUtils.bytes32FromBigInt(merkleRoot, "little")] ); expect(isValid).to.be.true; }); diff --git a/packages/rln/src/utils/bytes.ts b/packages/rln/src/utils/bytes.ts index 296ced4dd8..61d4528fae 100644 --- a/packages/rln/src/utils/bytes.ts +++ b/packages/rln/src/utils/bytes.ts @@ -50,30 +50,34 @@ export class BytesUtils { } /** - * Convert a BigInt to a Uint8Array with configurable output endianness - * @param value - The BigInt to convert - * @param byteLength - The desired byte length of the output (optional, auto-calculated if not provided) + * Convert a BigInt to a bytes32 (32-byte Uint8Array) + * @param value - The BigInt to convert (must fit in 32 bytes) * @param outputEndianness - Endianness of the output bytes ('big' or 'little') - * @returns Uint8Array representation of the BigInt + * @returns 32-byte Uint8Array representation of the BigInt */ - public static fromBigInt( + public static bytes32FromBigInt( value: bigint, - byteLength: number, outputEndianness: "big" | "little" = "little" ): Uint8Array { if (value < 0n) { throw new Error("Cannot convert negative BigInt to bytes"); } - if (value === 0n) { - return new Uint8Array(byteLength); + if (value >> 256n !== 0n) { + throw new Error( + `BigInt value is too large to fit in 32 bytes (max bit length: 256)` + ); } - const result = new Uint8Array(byteLength); + if (value === 0n) { + return new Uint8Array(32); + } + + const result = new Uint8Array(32); let workingValue = value; // Extract bytes in big-endian order - for (let i = byteLength - 1; i >= 0; i--) { + for (let i = 31; i >= 0; i--) { result[i] = Number(workingValue & 0xffn); workingValue = workingValue >> 8n; } diff --git a/packages/rln/src/utils/merkle.ts b/packages/rln/src/utils/merkle.ts index 2080c74439..fda9af680b 100644 --- a/packages/rln/src/utils/merkle.ts +++ b/packages/rln/src/utils/merkle.ts @@ -26,64 +26,23 @@ export function reconstructMerkleRoot( ); } - let currentValue = leafValue; + let currentValue = BytesUtils.bytes32FromBigInt(leafValue); - // Process each level of the tree (0 to MERKLE_TREE_DEPTH-1) for (let level = 0; level < MERKLE_TREE_DEPTH; level++) { - // Check if bit `level` is set in the leaf index const bit = (leafIndex >> BigInt(level)) & 1n; - // Convert bigints to Uint8Array for hashing - const currentBytes = bigIntToBytes32(currentValue); - const proofBytes = bigIntToBytes32(proof[level]); - - let hashResult: Uint8Array; + const proofBytes = BytesUtils.bytes32FromBigInt(proof[level]); if (bit === 0n) { // Current node is a left child: hash(current, proof[level]) - hashResult = poseidonHash(currentBytes, proofBytes); + currentValue = poseidonHash(currentValue, proofBytes); } else { // Current node is a right child: hash(proof[level], current) - hashResult = poseidonHash(proofBytes, currentBytes); - } - - // Convert hash result back to bigint for next iteration - currentValue = BytesUtils.toBigInt(hashResult, "little"); - } - - return currentValue; -} - -/** - * Extracts index information from a Merkle proof by attempting to reconstruct - * the root with different possible indices and comparing against the expected root - * - * @param proof - Array of MERKLE_TREE_DEPTH bigint elements representing the Merkle proof - * @param leafValue - The value of the leaf (typically the rate commitment) - * @param expectedRoot - The expected root to match against - * @param maxIndex - Maximum index to try (default: 2^MERKLE_TREE_DEPTH - 1) - * @returns The index that produces the expected root, or null if not found - */ -function extractIndexFromProof( - proof: readonly bigint[], - leafValue: bigint, - expectedRoot: bigint, - maxIndex: bigint = (1n << BigInt(MERKLE_TREE_DEPTH)) - 1n -): bigint | null { - // Try different indices to see which one produces the expected root - for (let index = 0n; index <= maxIndex; index++) { - try { - const reconstructedRoot = reconstructMerkleRoot(proof, index, leafValue); - if (reconstructedRoot === expectedRoot) { - return index; - } - } catch (error) { - // Continue trying other indices if reconstruction fails - continue; + currentValue = poseidonHash(proofBytes, currentValue); } } - return null; + return BytesUtils.toBigInt(currentValue, "little"); } /** @@ -98,65 +57,13 @@ export function calculateRateCommitment( idCommitment: bigint, rateLimit: bigint ): bigint { - const idBytes = bigIntToBytes32(idCommitment); - const rateLimitBytes = bigIntToBytes32(rateLimit); + const idBytes = BytesUtils.bytes32FromBigInt(idCommitment); + const rateLimitBytes = BytesUtils.bytes32FromBigInt(rateLimit); const hashResult = poseidonHash(idBytes, rateLimitBytes); return BytesUtils.toBigInt(hashResult, "little"); } -/** - * Converts a bigint to a 32-byte Uint8Array in little-endian format - * - * @param value - The bigint value to convert - * @returns 32-byte Uint8Array representation - */ -function bigIntToBytes32(value: bigint): Uint8Array { - const bytes = new Uint8Array(32); - let temp = value; - - for (let i = 0; i < 32; i++) { - bytes[i] = Number(temp & 0xffn); - temp >>= 8n; - } - - return bytes; -} - -/** - * Extracts the path direction bits from a Merkle proof by finding the leaf index - * that produces the expected root, then converting that index to path directions - * - * @param proof - Array of MERKLE_TREE_DEPTH bigint elements representing the Merkle proof - * @param leafValue - The value of the leaf (typically the rate commitment) - * @param expectedRoot - The expected root to match against - * @param maxIndex - Maximum index to try (default: 2^MERKLE_TREE_DEPTH - 1) - * @returns Array of MERKLE_TREE_DEPTH numbers (0 or 1) representing path directions, or null if no valid path found - * - 0 means the node is a left child (hash order: current, sibling) - * - 1 means the node is a right child (hash order: sibling, current) - */ -export function extractPathDirectionsFromProof( - proof: readonly bigint[], - leafValue: bigint, - expectedRoot: bigint, - maxIndex: bigint = (1n << BigInt(MERKLE_TREE_DEPTH)) - 1n -): number[] | null { - // First, find the leaf index that produces the expected root - const leafIndex = extractIndexFromProof( - proof, - leafValue, - expectedRoot, - maxIndex - ); - - if (leafIndex === null) { - return null; - } - - // Convert the leaf index to path directions - return getPathDirectionsFromIndex(leafIndex); -} - /** * Converts a leaf index to an array of path direction bits * @@ -165,7 +72,7 @@ export function extractPathDirectionsFromProof( * - 0 means the node is a left child (hash order: current, sibling) * - 1 means the node is a right child (hash order: sibling, current) */ -function getPathDirectionsFromIndex(leafIndex: bigint): number[] { +export function getPathDirectionsFromIndex(leafIndex: bigint): number[] { const pathDirections: number[] = []; // For each level (0 to MERKLE_TREE_DEPTH-1), extract the bit that determines left/right diff --git a/packages/rln/src/zerokit.ts b/packages/rln/src/zerokit.ts index ca743fec86..2e138a5153 100644 --- a/packages/rln/src/zerokit.ts +++ b/packages/rln/src/zerokit.ts @@ -41,7 +41,7 @@ export class Zerokit { idSecretHash: Uint8Array, pathElements: Uint8Array[], identityPathIndex: Uint8Array[], - x: Uint8Array, + msg: Uint8Array, epoch: Uint8Array, rateLimit: number, messageId: number // number of message sent by the user in this epoch @@ -69,6 +69,7 @@ export class Zerokit { // We assume that each identity path index is already in little-endian format identityPathIndexBytes.set(identityPathIndex[i], 8 + i * 1); } + const x = sha256(msg); return BytesUtils.concatenate( idSecretHash, BytesUtils.writeUIntLE(new Uint8Array(32), rateLimit, 0, 32), @@ -108,13 +109,11 @@ export class Zerokit { ); } - const x = sha256(msg); - const serializedWitness = await this.serializeWitness( idSecretHash, pathElements, identityPathIndex, - x, + msg, epoch, rateLimit, messageId