diff --git a/src/admin/MerkleWhitelist.ts b/src/admin/MerkleWhitelist.ts new file mode 100644 index 0000000..75b5373 --- /dev/null +++ b/src/admin/MerkleWhitelist.ts @@ -0,0 +1,156 @@ +import { keccak256Bytes, hexToBytes } from '../crypto/secp256k1'; + +export class MerkleWhitelist { + private leaves: Uint8Array[]; + private tree: Uint8Array[][]; + private root: Uint8Array; + + constructor(addresses: string[]) { + if (!addresses || addresses.length === 0) { + throw new Error('Cannot build tree from empty addresses array'); + } + this.leaves = addresses.map((addr) => MerkleWhitelist.hashLeaf(addr)); + this.tree = [this.leaves]; + this.buildTree(); + this.root = this.tree[this.tree.length - 1][0]; + } + + private buildTree(): void { + let currentLevel = this.leaves; + while (currentLevel.length > 1) { + const nextLevel: Uint8Array[] = []; + for (let i = 0; i < currentLevel.length; i += 2) { + const left = currentLevel[i]; + // Duplicate the last node if the level has an odd number of elements + const right = i + 1 < currentLevel.length ? currentLevel[i + 1] : left; + nextLevel.push(MerkleWhitelist.hashPair(left, right)); + } + this.tree.push(nextLevel); + currentLevel = nextLevel; + } + } + + getRoot(): string { + return MerkleWhitelist.bytesToHex(this.root); + } + + getProof(address: string): string[] { + const leafHashHex = MerkleWhitelist.bytesToHex(MerkleWhitelist.hashLeaf(address)); + const proof: Uint8Array[] = []; + + // Find index of the leaf + let index = -1; + for (let i = 0; i < this.leaves.length; i++) { + if (MerkleWhitelist.bytesToHex(this.leaves[i]) === leafHashHex) { + index = i; + break; + } + } + + if (index === -1) { + throw new Error('Address not found in tree'); + } + + for (let level = 0; level < this.tree.length - 1; level++) { + const currentLevel = this.tree[level]; + const isRightNode = index % 2 === 1; + const siblingIndex = isRightNode ? index - 1 : index + 1; + + if (siblingIndex < currentLevel.length) { + proof.push(currentLevel[siblingIndex]); + } else { + // If there's no right sibling, the node was duplicated + proof.push(currentLevel[index]); + } + + index = Math.floor(index / 2); + } + + return proof.map(p => MerkleWhitelist.bytesToHex(p)); + } + + static verifyProof(proof: string[], address: string, root: string): boolean { + if (proof.length === 0 && this.bytesToHex(this.hashLeaf(address)) !== root) { + return false; // Empty proof logic for single leaf tree edge case + } + + let computedHash = MerkleWhitelist.hashLeaf(address); + const rootBytes = MerkleWhitelist.hexToBytesSafe(root); + + for (const sibling of proof) { + const siblingBytes = MerkleWhitelist.hexToBytesSafe(sibling); + computedHash = MerkleWhitelist.hashPair(computedHash, siblingBytes); + } + + return MerkleWhitelist.compareBytes(computedHash, rootBytes) === 0; + } + + /** + * Hashes a leaf address according to OpenZeppelin standard: + * keccak256(abi.encodePacked(address)) + */ + static hashLeaf(address: string): Uint8Array { + let cleanAddress = address.toLowerCase(); + if (cleanAddress.startsWith('0x')) { + cleanAddress = cleanAddress.slice(2); + } + // Pad to 20 bytes (40 hex chars) if short, though usually they are exactly 40 + cleanAddress = cleanAddress.padStart(40, '0'); + if (cleanAddress.length !== 40) { + throw new Error('Invalid address length'); + } + const addressBytes = hexToBytes(cleanAddress); + return keccak256Bytes(addressBytes); + } + + /** + * Hashes a pair of hashes according to OpenZeppelin standard: + * lexicographically sort the two hashes, concatenate them, then keccak256. + */ + static hashPair(a: Uint8Array, b: Uint8Array): Uint8Array { + const [left, right] = MerkleWhitelist.compareBytes(a, b) <= 0 ? [a, b] : [b, a]; + const combined = new Uint8Array(left.length + right.length); + combined.set(left, 0); + combined.set(right, left.length); + return keccak256Bytes(combined); + } + + private static compareBytes(a: Uint8Array, b: Uint8Array): number { + const len = Math.min(a.length, b.length); + for (let i = 0; i < len; i++) { + if (a[i] !== b[i]) { + return a[i] - b[i]; + } + } + return a.length - b.length; + } + + static bytesToHex(bytes: Uint8Array): string { + let hex = '0x'; + for (let i = 0; i < bytes.length; i++) { + hex += bytes[i].toString(16).padStart(2, '0'); + } + return hex; + } + + static hexToBytesSafe(hex: string): Uint8Array { + let clean = hex.startsWith('0x') ? hex.slice(2) : hex; + if (clean.length % 2 !== 0) { + clean = '0' + clean; + } + return hexToBytes(clean); + } + + static buildFromAddresses(addresses: string[]): { + root: string; + tree: MerkleWhitelist; + getProof: (address: string) => string[]; + } { + const tree = new MerkleWhitelist(addresses); + return { + root: tree.getRoot(), + tree, + getProof: (address: string) => tree.getProof(address), + }; + } +} diff --git a/src/admin/index.ts b/src/admin/index.ts index c2588ad..c9ab0de 100644 --- a/src/admin/index.ts +++ b/src/admin/index.ts @@ -1 +1,2 @@ +export * from './MerkleWhitelist'; export * from './whitelistAdmin'; diff --git a/src/admin/whitelistAdmin.ts b/src/admin/whitelistAdmin.ts index 18320fe..b692bb3 100644 --- a/src/admin/whitelistAdmin.ts +++ b/src/admin/whitelistAdmin.ts @@ -1,11 +1,11 @@ import { ethers } from 'ethers'; -import { buildWhitelistTree, MerkleTree } from '../utils/merkleTree'; +import { MerkleWhitelist } from './MerkleWhitelist'; export interface WhitelistAdmin { buildTree(addresses: string[]): { root: string; getProof: (address: string) => string[]; - tree: MerkleTree; + tree: MerkleWhitelist; }; publishRoot(root: string, version?: number): Promise; rotateWhitelist(newRoot: string, version?: number): Promise; @@ -18,7 +18,7 @@ export function createWhitelistAdmin( ): WhitelistAdmin { return { buildTree(addresses: string[]) { - const { root, tree, getProof } = buildWhitelistTree(addresses); + const { root, tree, getProof } = MerkleWhitelist.buildFromAddresses(addresses); return { root, getProof, tree }; }, @@ -66,10 +66,10 @@ export function createWhitelistAdmin( export function generateWhitelist(addresses: string[]): { root: string; - tree: MerkleTree; + tree: MerkleWhitelist; proofs: Map; } { - const { root, tree, getProof } = buildWhitelistTree(addresses); + const { root, tree, getProof } = MerkleWhitelist.buildFromAddresses(addresses); const proofs = new Map(); for (const address of addresses) { diff --git a/src/utils/index.ts b/src/utils/index.ts index 9597ff0..7cf874c 100644 --- a/src/utils/index.ts +++ b/src/utils/index.ts @@ -1,2 +1,2 @@ -export * from './merkleTree'; + export * from './constantTime'; diff --git a/src/utils/merkleTree.ts b/src/utils/merkleTree.ts deleted file mode 100644 index cfb5542..0000000 --- a/src/utils/merkleTree.ts +++ /dev/null @@ -1,116 +0,0 @@ -import { keccak256 } from 'ethers'; - -export interface MerkleProof { - path: string[]; - index: number; - leaf: string; -} - -export class MerkleTree { - private leaves: string[]; - private tree: string[][]; - private root: string; - - constructor(leaves: string[]) { - this.leaves = leaves.map((leaf) => MerkleTree.hashLeaf(leaf)); - this.tree = [this.leaves]; - this.buildTree(); - this.root = this.tree[this.tree.length - 1][0]; - } - - private buildTree(): void { - let currentLevel = this.leaves; - while (currentLevel.length > 1) { - const nextLevel: string[] = []; - for (let i = 0; i < currentLevel.length; i += 2) { - const left = currentLevel[i]; - const right = i + 1 < currentLevel.length ? currentLevel[i + 1] : left; - nextLevel.push(MerkleTree.hashPair(left, right)); - } - this.tree.push(nextLevel); - currentLevel = nextLevel; - } - } - - getRoot(): string { - return this.root; - } - - getTree(): string[][] { - return this.tree; - } - - getLeaves(): string[] { - return this.leaves; - } - - getProof(leaf: string): string[] { - const leafHash = MerkleTree.hashLeaf(leaf); - const proof: string[] = []; - let index = this.leaves.indexOf(leafHash); - - if (index === -1) { - throw new Error('Leaf not found in tree'); - } - - for (let level = 0; level < this.tree.length - 1; level++) { - const currentLevel = this.tree[level]; - const isRight = index % 2 === 1; - const siblingIndex = isRight ? index - 1 : index + 1; - - if (siblingIndex < currentLevel.length) { - proof.push(currentLevel[siblingIndex]); - } else { - proof.push(currentLevel[index]); - } - - index = Math.floor(index / 2); - } - - return proof; - } - - static verifyProof( - proof: string[], - leaf: string, - root: string - ): boolean { - let computedHash = MerkleTree.hashLeaf(leaf); - - for (const sibling of proof) { - computedHash = MerkleTree.hashPair(computedHash, sibling); - } - - return computedHash === root; - } - - static hashLeaf(leaf: string): string { - return keccak256(Buffer.from(leaf)); - } - - static hashPair(left: string, right: string): string { - const [a, b] = left <= right ? [left, right] : [right, left]; - return keccak256(Buffer.from(a + b)); - } - - static buildFromAddresses(addresses: string[]): { - root: string; - tree: MerkleTree; - getProof: (address: string) => string[]; - } { - const tree = new MerkleTree(addresses); - return { - root: tree.getRoot(), - tree, - getProof: (address: string) => tree.getProof(address), - }; - } -} - -export function buildWhitelistTree(addresses: string[]): { - root: string; - tree: MerkleTree; - getProof: (address: string) => string[]; -} { - return MerkleTree.buildFromAddresses(addresses); -} diff --git a/tests/whitelist/rotation.test.ts b/tests/whitelist/rotation.test.ts index b2bdf69..01f9adc 100644 --- a/tests/whitelist/rotation.test.ts +++ b/tests/whitelist/rotation.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, beforeEach } from 'vitest'; -import { MerkleTree } from '../../src/utils/merkleTree'; +import { MerkleWhitelist } from '../../src/admin/MerkleWhitelist'; // Mock data const testAddresses = [ @@ -9,65 +9,106 @@ const testAddresses = [ '0x4567890123456789012345678901234567890123', ]; -describe('Merkle Whitelist Rotation', () => { - let tree: MerkleTree; +// OpenZeppelin reference verifier in pure JS for cross-checking +function referenceOzVerify(proof: string[], root: string, leaf: string): boolean { + let computedHash = MerkleWhitelist.hashLeaf(leaf); + const rootBytes = MerkleWhitelist.hexToBytesSafe(root); + + for (const p of proof) { + const proofElement = MerkleWhitelist.hexToBytesSafe(p); + computedHash = MerkleWhitelist.hashPair(computedHash, proofElement); + } + + // Compare + for (let i = 0; i < computedHash.length; i++) { + if (computedHash[i] !== rootBytes[i]) return false; + } + return true; +} + +describe('Merkle Whitelist Rotation & Lifecycles', () => { + let tree: MerkleWhitelist; let root: string; let proof: string[]; beforeEach(() => { - tree = new MerkleTree(testAddresses); + tree = new MerkleWhitelist(testAddresses); root = tree.getRoot(); proof = tree.getProof(testAddresses[0]); }); it('should verify a proof against the current root', () => { - const isValid = MerkleTree.verifyProof(proof, testAddresses[0], root); + const isValid = MerkleWhitelist.verifyProof(proof, testAddresses[0], root); expect(isValid).toBe(true); + + // Check against reference implementation + expect(referenceOzVerify(proof, root, testAddresses[0])).toBe(true); }); it('should reject a proof against a different root', () => { // Create a different tree - const differentAddresses = ['0x999...', '0x888...']; - const differentTree = new MerkleTree(differentAddresses); + const differentAddresses = ['0x9999999999999999999999999999999999999999', '0x8888888888888888888888888888888888888888']; + const differentTree = new MerkleWhitelist(differentAddresses); const differentRoot = differentTree.getRoot(); - const isValid = MerkleTree.verifyProof(proof, testAddresses[0], differentRoot); + const isValid = MerkleWhitelist.verifyProof(proof, testAddresses[0], differentRoot); expect(isValid).toBe(false); }); it('should reject a proof against an old root after rotation', async () => { // Create initial tree - const initialTree = new MerkleTree(testAddresses); + const initialTree = new MerkleWhitelist(testAddresses); const initialRoot = initialTree.getRoot(); const initialProof = initialTree.getProof(testAddresses[0]); // Simulate root rotation const newAddresses = [ - '0x9876543210...', - '0x8765432109...', + '0x9876543210987654321098765432109876543210', + '0x8765432109876543210987654321098765432109', ]; - const newTree = new MerkleTree(newAddresses); + const newTree = new MerkleWhitelist(newAddresses); const newRoot = newTree.getRoot(); // The proof should be valid against the old root - expect(MerkleTree.verifyProof(initialProof, testAddresses[0], initialRoot)).toBe(true); + expect(MerkleWhitelist.verifyProof(initialProof, testAddresses[0], initialRoot)).toBe(true); // But invalid against the new root (after rotation) - expect(MerkleTree.verifyProof(initialProof, testAddresses[0], newRoot)).toBe(false); - - // This demonstrates the central security property: - // A proof valid against a previous root is correctly rejected once the root has rotated + expect(MerkleWhitelist.verifyProof(initialProof, testAddresses[0], newRoot)).toBe(false); }); - it('should handle addresses not in the tree', () => { + it('should throw for addresses not in the tree', () => { const nonWhitelistedAddress = '0x9999999999999999999999999999999999999999'; - // getProof throws for non-members rather than returning an empty/invalid proof. + // getProof throws for non-members expect(() => tree.getProof(nonWhitelistedAddress)).toThrow(); }); it('should handle empty proof arrays', () => { - const isValid = MerkleTree.verifyProof([], testAddresses[0], root); + const isValid = MerkleWhitelist.verifyProof([], testAddresses[0], root); expect(isValid).toBe(false); }); + + it('should build and verify a large tree (1000+ leaves)', () => { + const largeAddresses: string[] = []; + for (let i = 0; i < 1500; i++) { + // Generate deterministic fake addresses for performance test + let fakeHex = i.toString(16).padStart(40, '0'); + largeAddresses.push('0x' + fakeHex); + } + + const largeTree = new MerkleWhitelist(largeAddresses); + const largeRoot = largeTree.getRoot(); + + // Pick an address in the middle + const targetAddr = largeAddresses[750]; + const largeProof = largeTree.getProof(targetAddr); + + expect(largeProof.length).toBeGreaterThan(0); + + // Verify using standard static verifier + expect(MerkleWhitelist.verifyProof(largeProof, targetAddr, largeRoot)).toBe(true); + + // Verify using reference OZ verifier + expect(referenceOzVerify(largeProof, largeRoot, targetAddr)).toBe(true); + }); });