self/circuits/tests/utils/Mgf1_sha256.test.ts

import path from 'path';
import { createHash, randomBytes } from 'node:crypto';
const wasm_tester = require('circom_tester').wasm;

describe('Mgf1_sha256 Circuit Test', function () {
  this.timeout(0); // Disable timeout
  const hashLen = 32; // SHA256 length

  const compileCircuit = async (circuitPath: string) => {
    return await wasm_tester(path.join(circuitPath), {
      include: ['node_modules'],
    });
  };

  function buffer2bitArray(b) {
    const res = [];
    for (let i = 0; i < b.length; i++) {
      for (let j = 0; j < 8; j++) {
        res.push((b[i] >> (7 - j)) & 1);
      }
    }
    return res;
  }

  function num2Bits(n, input) {
    let out = [];
    for (let i = 0; i < n; i++) {
      out[i] = (input >> i) & 1;
    }
    return out;
  }

  const bitArray2buffer = (a) => {
    const len = Math.floor((a.length - 1) / 8) + 1;
    const b = Buffer.alloc(len);

    for (let i = 0; i < a.length; i++) {
      const p = Math.floor(i / 8);
      b[p] = b[p] | (Number(a[i]) << (7 - (i % 8)));
    }
    return b;
  };

  const MGF1 = (mgfSeed: Buffer, maskLen: number) => {
    const hLen = hashLen;
    if (maskLen > 0xffffffff * hLen) {
      throw new Error('mask too long');
    }

    var T = [];
    for (var i = 0; i <= Math.ceil(maskLen / hLen) - 1; i++) {
      var C = Buffer.alloc(4);
      C.writeUInt32BE(i);
      const hash3 = createHash('sha256');
      hash3.update(Buffer.concat([mgfSeed, C]));
      T.push(hash3.digest());
    }
    return Buffer.concat(T).slice(0, maskLen);
  };

  it('Should compile', async function () {
    await compileCircuit('circuits/tests/mgf1Sha256/Mgf1Sha256_tester.circom');
  });

  it('Should generate correct MGF1 output - 4 Byte Seed', async function () {
    const seed = 12345678;
    const maskLen = 32;
    const seedLen = 4; // 4 bytes - set in the circuit

    const bitArray = num2Bits(seedLen * 8, seed);
    const mgfSeed = bitArray2buffer(bitArray);

    const expected = MGF1(mgfSeed, maskLen);

    const circuit = await compileCircuit('circuits/tests/mgf1Sha256/Mgf1Sha256_tester.circom');
    const expected_mask_output = buffer2bitArray(expected);
    const inputs = {
      seed: seed,
      expected_mask_output,
    };

    const witness = await circuit.calculateWitness(inputs);
    await circuit.checkConstraints(witness);
  });

  it('Should generate correct MGF1 output - 32 Byte Seed', async function () {
    const randBytes = randomBytes(32);

    const maskLen = 32;
    const seedLen = 32; // set in circuit
    const expected = MGF1(randBytes, maskLen);

    const circuit = await compileCircuit(
      'circuits/tests/mgf1Sha256/Mgf1Sha256_32Bytes_tester.circom'
    );
    const expected_mask_output = buffer2bitArray(expected);
    const inputs = {
      seed: buffer2bitArray(randBytes),
      expected_mask_output,
    };

    const witness = await circuit.calculateWitness(inputs);
    await circuit.checkConstraints(witness);
  });

  it('Should generate correct MGF1 output - seedLen value > than actual seed length', async function () {
    const seed = 1234;
    const maskLen = 32;
    const seedLen = 4; //set in circuit

    const bitArray = num2Bits(seedLen * 8, seed);
    const mgfSeed = bitArray2buffer(bitArray);

    const expected = MGF1(mgfSeed, maskLen);

    const circuit = await compileCircuit('circuits/tests/mgf1Sha256/Mgf1Sha256_tester.circom');
    const expected_mask_output = buffer2bitArray(expected);
    const inputs = {
      seed: seed,
      expected_mask_output,
    };

    const witness = await circuit.calculateWitness(inputs);
    await circuit.checkConstraints(witness);
  });

  it('Should generate correct MGF1 output - maskLen == 1', async function () {
    const seed = 12345678;
    const maskLen = 1;
    const seedLen = 4;

    const bitArray = num2Bits(seedLen * 8, seed);
    const mgfSeed = bitArray2buffer(bitArray);

    const expected = MGF1(mgfSeed, maskLen);

    const circuit = await compileCircuit(
      'circuits/tests/mgf1Sha256/Mgf1Sha256_1ByteMask_tester.circom'
    );
    const expected_mask_output = buffer2bitArray(expected);
    const inputs = {
      seed: seed,
      expected_mask_output,
    };

    const witness = await circuit.calculateWitness(inputs);
    await circuit.checkConstraints(witness);
  });
});