// AereProofAggregatorV2 — regression tests for INNER VKEYS COUNTED BUT NEVER ENFORCED.
//
// FINDING (2026-07-17 sweep, code-read): AereProofAggregator 0x6a260238…F84C5 pins the
// aggregation vkey and the guest is genuinely recursive — but the contract only COUNTED how
// many folded inner vkeys were registered (into a `recognized` field) and recorded the
// aggregation regardless. An unregistered inner program was not an error, just a lower number.
//
// THE ATTACK: an attacker writes an SP1 guest with NO constraints —
//     fn main() { sp1_zkvm::io::commit_slice(&attacker_chosen_bytes); }
// — has it commit the exact AereZKScreen public tuple, proves it (trivial: it asserts nothing),
// folds it with the REAL aggregation guest (which honestly verifies it, because it IS a valid
// SP1 proof of that vacuous program), and submits. V1 records an aggregation that `exists`,
// carries the pinned aggVKey, and is indistinguishable from a genuine one except for
// `recognized < count` — a field nothing ever enforced.
//
// Each attack test is PAIRED with a control on the REAL V1 contract showing the same call
// SUCCEEDS on V1, i.e. the bug was real and V2 closes it.
//
// THE GUEST IS UNCHANGED. The composite the guest commits ALREADY binds every inner vkey
// digest, and V1 already recomputed it correctly — so the contract always KNEW which programs
// were folded. The fix is purely what the contract DOES with that knowledge. This is verified
// below: the composite recomputed by these tests reproduces the REAL live on-chain aggregation
// digest 0x8e246396…4e36 byte-for-byte.
//
// The SP1 verifier is MOCKED (MockSp1Verifier: reverts unless the exact (vkey, publicValues,
// proof) triple was pre-registered — mirroring the real gateway). That is the CORRECT model:
// the attacker's aggregation proof is CRYPTOGRAPHICALLY VALID (the recursion honestly verified
// a real proof of a real — if vacuous — program), so the gateway accepts it. The bug is not a
// broken proof; it is a proof about an unvetted program.
const { expect } = require("chai");
const { ethers } = require("hardhat");

const PROOF = "0x4388a21c1234"; // opaque; the mock keys acceptance on the triple

/* ---------------- REAL live values (chain 2800, read 2026-07-17) ------------------ */
// The pinned aggregation program (registered on the live V1).
const AGG_VKEY = "0x003a70854f258b13c6612b0272421f2d7d97f90eee63e124864c80be61a2672c";
// The three REAL registered inner programs: KoalaBear digest (what the recursion binds)
// + BN254 on-chain vkey (what is published) + label.
const REAL_INNERS = [
  {
    label: "zkscreen",
    digest: "0x3510880a6bb8ae42705a4cd76f3aa69422a8d836122361ba12a913635de3e5b7",
    onChain: "0x006a211015aee2b90b82d266bef3aa6944551b06c488d86e895489b1dde3e5b7",
    pvHash: "0xe4d363f0108963cb592f1c592d855a46b6d9579cccef2d4689b93c2441c78d88",
  },
  {
    label: "over18",
    digest: "0x2d54a563446f19740eab67524f3622a67463b716375c230a30d2b0fe633343c5",
    onChain: "0x005aa94ac711bc65d0755b3a94f3622a6e8c76e2cdd708c29869587f633343c5",
    pvHash: "0x44d77dcf2f60c490ffad39fb2cbe75648ae3c402ce6cbe1ac29f6525af3c6ce6",
  },
  {
    label: "zkml-mnist",
    digest: "0x30a21e711cd157425f9b5e2203ea0a525cb86cbc625b720f3313ce3727c4bceb",
    onChain: "0x0061443ce273455d0afcdaf1103ea0a52b970d979896dc83d989e71ba7c4bceb",
    pvHash: "0xa04ed64252a5388bee0a757a0c2f83da7594cfaee6c7c25e6b873f46f43aad7e",
  },
];
// The composite the LIVE contract actually recorded for those three (aggregationCount==2).
const REAL_COMPOSITE = "0x8e246396a2f40c837501d06c965163db644284ffc877d7cd0d38681b612b4e36";

const FOUNDATION = "0x0243A4f47D44b40b65D33f20329dE20D00c6f3C3"; // live V1 owner

// The attacker's unconstrained guest: a real SP1 program, so it has a real vkey digest —
// just one that was never vetted or registered.
const ATTACKER_DIGEST = "0x" + "de".repeat(32);

// composite = sha256( u32_be(n) || (digest_i || pvHash_i)... ) — the guest's formula.
function composite(digests, pvHashes) {
  let packed = ethers.solidityPacked(["uint32"], [digests.length]);
  for (let i = 0; i < digests.length; i++) packed = ethers.concat([packed, digests[i], pvHashes[i]]);
  return ethers.sha256(packed);
}
function marker(vkey, pv, proof) {
  return ethers.keccak256(ethers.AbiCoder.defaultAbiCoder().encode(["bytes32", "bytes", "bytes"], [vkey, pv, proof]));
}

describe("AereProofAggregatorV2 — inner-vkey enforcement fix (2026-07-17)", function () {
  let mock, v1, v2, owner, attacker;

  const digests = REAL_INNERS.map((p) => p.digest);
  const pvHashes = REAL_INNERS.map((p) => p.pvHash);

  async function accept(vkey, pv, proof = PROOF) {
    await mock.setValid(marker(vkey, pv, proof), true);
  }

  async function deployV2(seal = false) {
    return await (await ethers.getContractFactory("AereProofAggregatorV2")).deploy(
      await mock.getAddress(),
      seal ? ethers.ZeroAddress : owner.address,
      [AGG_VKEY],
      ["recursive-aggregation agg-program (real, live)"],
      digests,
      REAL_INNERS.map((p) => p.onChain),
      REAL_INNERS.map((p) => p.label),
      seal
    );
  }

  beforeEach(async function () {
    [owner, attacker] = await ethers.getSigners();
    mock = await (await ethers.getContractFactory("MockSp1Verifier")).deploy();
    await mock.waitForDeployment();

    // The REAL V1 contract, registered EXACTLY as the live one is — the control.
    v1 = await (await ethers.getContractFactory("AereProofAggregator")).deploy(await mock.getAddress());
    await v1.waitForDeployment();
    await v1.registerAggProgram(AGG_VKEY, "recursive-aggregation agg-program (real, live)");
    for (const p of REAL_INNERS) await v1.registerInnerProgram(p.digest, p.onChain, p.label);

    v2 = await deployV2(false);
    await v2.waitForDeployment();
  });

  /* ---------------- the guest is unchanged: composite reproduces live data --------- */

  describe("the recursion guest is UNCHANGED (the composite already bound the inner vkeys)", function () {
    it("recomputes the REAL live on-chain composite 0x8e246396…4e36 byte-for-byte", async function () {
      expect(composite(digests, pvHashes)).to.equal(REAL_COMPOSITE);
    });

    it("V2's on-chain _composite matches the guest formula on the REAL live inner set", async function () {
      await accept(AGG_VKEY, REAL_COMPOSITE);
      // verifyAggregation returns the composite it recomputed on-chain
      expect(await v2.verifyAggregation(AGG_VKEY, digests, pvHashes, REAL_COMPOSITE, PROOF)).to.equal(REAL_COMPOSITE);
    });

    it("V1 recomputed it correctly too — V1 always KNEW which programs were folded", async function () {
      await accept(AGG_VKEY, REAL_COMPOSITE);
      expect(await v1.verifyAggregation(AGG_VKEY, digests, pvHashes, REAL_COMPOSITE, PROOF)).to.equal(REAL_COMPOSITE);
    });
  });

  /* ================================================================================
     THE ATTACK: fold a proof of an UNCONSTRAINED, unregistered guest.
     ================================================================================ */

  describe("ATTACK — unconstrained attacker guest folded into an aggregation", function () {
    // The attacker's vacuous guest commits the EXACT zkscreen public tuple (so its pvHash is
    // the real zkscreen one), but the program that produced it is the attacker's own.
    const atkDigests = () => [ATTACKER_DIGEST];
    const atkPvHashes = () => [REAL_INNERS[0].pvHash]; // the real AereZKScreen public-values hash
    const atkComposite = () => composite(atkDigests(), atkPvHashes());

    it("PAIRED CONTROL: V1 RECORDS the attacker's aggregation as if legitimate", async function () {
      const c = atkComposite();
      await accept(AGG_VKEY, c);
      await expect(v1.connect(attacker).recordAggregation(AGG_VKEY, atkDigests(), atkPvHashes(), c, PROOF)).to.emit(
        v1,
        "AggregationVerified"
      );

      const rec = await v1.getAggregation(c);
      expect(rec[0]).to.equal(true); // exists  <-- a forged aggregation is ON THE BOOKS
      expect(rec[1]).to.equal(AGG_VKEY); // carries the pinned, legitimate agg vkey
      expect(rec[2]).to.equal(1); // count
      expect(rec[3]).to.equal(0); // recognized == 0: the ONLY tell, and nothing enforced it
      expect(await v1.aggregationCount()).to.equal(1n);
    });

    it("V2 REVERTS with UnregisteredInnerProgram — the attack is closed", async function () {
      const c = atkComposite();
      await accept(AGG_VKEY, c);
      await expect(v2.connect(attacker).recordAggregation(AGG_VKEY, atkDigests(), atkPvHashes(), c, PROOF))
        .to.be.revertedWithCustomError(v2, "UnregisteredInnerProgram")
        .withArgs(ATTACKER_DIGEST, 0);
      expect((await v2.getAggregation(c))[0]).to.equal(false); // nothing recorded
      expect(await v2.aggregationCount()).to.equal(0n);
    });

    it("V2's verifyAggregation view REVERTS too (a view used as an oracle must bind)", async function () {
      const c = atkComposite();
      await accept(AGG_VKEY, c);
      await expect(v2.verifyAggregation(AGG_VKEY, atkDigests(), atkPvHashes(), c, PROOF))
        .to.be.revertedWithCustomError(v2, "UnregisteredInnerProgram")
        .withArgs(ATTACKER_DIGEST, 0);
    });

    it("PAIRED CONTROL: V1's verifyAggregation view ACCEPTS the attacker's fold", async function () {
      const c = atkComposite();
      await accept(AGG_VKEY, c);
      expect(await v1.verifyAggregation(AGG_VKEY, atkDigests(), atkPvHashes(), c, PROOF)).to.equal(c);
    });
  });

  describe("ATTACK — attacker guest SMUGGLED IN alongside real registered programs", function () {
    // Harder to spot: 3 real programs + 1 attacker program. V1 reports recognized=3, count=4.
    const mixDigests = () => [...digests, ATTACKER_DIGEST];
    const mixPvHashes = () => [...pvHashes, "0x" + "77".repeat(32)];
    const mixComposite = () => composite(mixDigests(), mixPvHashes());

    it("PAIRED CONTROL: V1 records it (recognized=3 of count=4) — looks almost legitimate", async function () {
      const c = mixComposite();
      await accept(AGG_VKEY, c);
      await v1.connect(attacker).recordAggregation(AGG_VKEY, mixDigests(), mixPvHashes(), c, PROOF);
      const rec = await v1.getAggregation(c);
      expect(rec[0]).to.equal(true);
      expect(rec[2]).to.equal(4); // count
      expect(rec[3]).to.equal(3); // recognized — 3 real programs carry the forged 4th
    });

    it("V2 REVERTS at the attacker's index (3), rejecting the whole aggregation", async function () {
      const c = mixComposite();
      await accept(AGG_VKEY, c);
      await expect(v2.connect(attacker).recordAggregation(AGG_VKEY, mixDigests(), mixPvHashes(), c, PROOF))
        .to.be.revertedWithCustomError(v2, "UnregisteredInnerProgram")
        .withArgs(ATTACKER_DIGEST, 3);
      expect(await v2.aggregationCount()).to.equal(0n);
    });
  });

  describe("ATTACK — zero inner vkey digest (default value of an unregistered slot)", function () {
    it("V2 REVERTS: a zero digest can never be allowed, so the gate cannot be no-op'd", async function () {
      const d = [ethers.ZeroHash];
      const h = [REAL_INNERS[0].pvHash];
      const c = composite(d, h);
      await accept(AGG_VKEY, c);
      await expect(v2.recordAggregation(AGG_VKEY, d, h, c, PROOF))
        .to.be.revertedWithCustomError(v2, "UnregisteredInnerProgram")
        .withArgs(ethers.ZeroHash, 0);
    });

    it("V2 refuses to REGISTER a zero digest (would turn the gate into a no-op)", async function () {
      await expect(v2.registerInnerProgram(ethers.ZeroHash, "0x" + "11".repeat(32), "evil"))
        .to.be.revertedWithCustomError(v2, "ZeroVKey");
    });

    it("PAIRED CONTROL: V1 HAPPILY registers a zero digest", async function () {
      await expect(v1.registerInnerProgram(ethers.ZeroHash, "0x" + "11".repeat(32), "evil")).to.not.be.reverted;
      expect((await v1.innerPrograms(ethers.ZeroHash))[0]).to.equal(true);
    });
  });

  /* ---------------- V2 does not over-reject: the honest path still works ----------- */

  describe("V2 does not over-reject — real aggregations still record", function () {
    it("records the REAL 3-program aggregation, recognized == count (the V2 invariant)", async function () {
      await accept(AGG_VKEY, REAL_COMPOSITE);
      await expect(v2.recordAggregation(AGG_VKEY, digests, pvHashes, REAL_COMPOSITE, PROOF)).to.emit(
        v2,
        "AggregationVerified"
      );
      const rec = await v2.getAggregation(REAL_COMPOSITE);
      expect(rec[0]).to.equal(true);
      expect(rec[1]).to.equal(AGG_VKEY);
      expect(rec[2]).to.equal(3); // count
      expect(rec[3]).to.equal(3); // recognized == count, ALWAYS, by construction
      expect(rec[5]).to.equal(owner.address);
      expect(await v2.aggregationCount()).to.equal(1n);
    });

    it("seeds the allowed set at deploy — functional with NO post-deploy signature", async function () {
      for (const p of REAL_INNERS) {
        const r = await v2.innerPrograms(p.digest);
        expect(r[0]).to.equal(true);
        expect(r[1]).to.equal(p.onChain);
        expect(r[2]).to.equal(p.label);
      }
      expect(await v2.isAggProgram(AGG_VKEY)).to.equal(true);
      expect(await v2.owner()).to.equal(owner.address);
      expect(await v2.SP1_VERIFIER()).to.equal(await mock.getAddress());
    });

    it("areInnerProgramsAllowed() previews the gate without spending gas", async function () {
      expect(await v2.areInnerProgramsAllowed(digests)).to.equal(true);
      expect(await v2.areInnerProgramsAllowed([...digests, ATTACKER_DIGEST])).to.equal(false);
    });

    it("a newly-registered program becomes foldable (the governed extension path)", async function () {
      const NEWD = "0x" + "5a".repeat(32);
      const d = [NEWD];
      const h = ["0x" + "99".repeat(32)];
      const c = composite(d, h);
      await accept(AGG_VKEY, c);
      await expect(v2.recordAggregation(AGG_VKEY, d, h, c, PROOF)).to.be.revertedWithCustomError(
        v2,
        "UnregisteredInnerProgram"
      );
      await v2.registerInnerProgram(NEWD, "0x" + "6b".repeat(32), "new-program");
      await expect(v2.recordAggregation(AGG_VKEY, d, h, c, PROOF)).to.emit(v2, "AggregationVerified");
    });
  });

  /* ---------------- inherited protections still hold ------------------------------ */

  describe("inherited protections (unchanged from V1)", function () {
    it("an unregistered AGGREGATION program still reverts", async function () {
      const bad = "0x" + "01".repeat(32);
      await expect(v2.recordAggregation(bad, digests, pvHashes, REAL_COMPOSITE, PROOF))
        .to.be.revertedWithCustomError(v2, "UnknownAggProgram")
        .withArgs(bad);
    });

    it("a composite that does not match the committed public values reverts", async function () {
      const wrong = "0x" + "ee".repeat(32);
      await accept(AGG_VKEY, wrong);
      await expect(v2.recordAggregation(AGG_VKEY, digests, pvHashes, wrong, PROOF)).to.be.revertedWithCustomError(
        v2,
        "CompositeMismatch"
      );
    });

    it("an invalid Groth16 proof reverts InvalidProof", async function () {
      await expect(v2.recordAggregation(AGG_VKEY, digests, pvHashes, REAL_COMPOSITE, PROOF))
        .to.be.revertedWithCustomError(v2, "InvalidProof");
    });

    it("dropping a folded child changes the composite => rejected", async function () {
      await accept(AGG_VKEY, REAL_COMPOSITE);
      await expect(
        v2.recordAggregation(AGG_VKEY, digests.slice(0, 2), pvHashes.slice(0, 2), REAL_COMPOSITE, PROOF)
      ).to.be.revertedWithCustomError(v2, "CompositeMismatch");
    });

    it("empty / mismatched-length / wrong-pv-length all revert", async function () {
      await expect(v2.recordAggregation(AGG_VKEY, [], [], REAL_COMPOSITE, PROOF)).to.be.revertedWithCustomError(
        v2,
        "EmptyAggregation"
      );
      await expect(
        v2.recordAggregation(AGG_VKEY, digests, pvHashes.slice(0, 2), REAL_COMPOSITE, PROOF)
      ).to.be.revertedWithCustomError(v2, "LengthMismatch");
      await expect(v2.recordAggregation(AGG_VKEY, digests, pvHashes, "0x1234", PROOF))
        .to.be.revertedWithCustomError(v2, "BadPublicValuesLength")
        .withArgs(2);
    });
  });

  /* ---------------- registry governance + the seal latch -------------------------- */

  describe("registry governance and the one-way seal", function () {
    it("only the owner may register", async function () {
      await expect(
        v2.connect(attacker).registerInnerProgram("0x" + "5a".repeat(32), "0x" + "6b".repeat(32), "evil")
      ).to.be.revertedWith("Ownable: caller is not the owner");
      await expect(v2.connect(attacker).registerAggProgram("0x" + "5a".repeat(32), "evil")).to.be.revertedWith(
        "Ownable: caller is not the owner"
      );
    });

    it("sealRegistry() permanently freezes the set and renounces ownership", async function () {
      await expect(v2.sealRegistry()).to.emit(v2, "RegistrySealed");
      expect(await v2.registrySealed()).to.equal(true);
      expect(await v2.owner()).to.equal(ethers.ZeroAddress);
      // no one can extend the set any more — not even the old owner
      await expect(
        v2.registerInnerProgram("0x" + "5a".repeat(32), "0x" + "6b".repeat(32), "post-seal")
      ).to.be.revertedWith("Ownable: caller is not the owner");
      await expect(v2.sealRegistry()).to.be.revertedWith("Ownable: caller is not the owner");
    });

    it("sealed aggregations still record (sealing freezes the SET, not the contract)", async function () {
      await v2.sealRegistry();
      await accept(AGG_VKEY, REAL_COMPOSITE);
      await expect(v2.recordAggregation(AGG_VKEY, digests, pvHashes, REAL_COMPOSITE, PROOF)).to.emit(
        v2,
        "AggregationVerified"
      );
    });

    it("can be deployed SEALED at construction: seeded set, zero admin, ever", async function () {
      const sealed = await deployV2(true);
      await sealed.waitForDeployment();
      expect(await sealed.registrySealed()).to.equal(true);
      expect(await sealed.owner()).to.equal(ethers.ZeroAddress);
      expect(await sealed.isAggProgram(AGG_VKEY)).to.equal(true);
      expect((await sealed.innerPrograms(REAL_INNERS[0].digest))[0]).to.equal(true);
      // and it still rejects the attack
      const c = composite([ATTACKER_DIGEST], [REAL_INNERS[0].pvHash]);
      await accept(AGG_VKEY, c);
      await expect(
        sealed.recordAggregation(AGG_VKEY, [ATTACKER_DIGEST], [REAL_INNERS[0].pvHash], c, PROOF)
      ).to.be.revertedWithCustomError(sealed, "UnregisteredInnerProgram");
    });
  });

  /* ---------------- constructor guards -------------------------------------------- */

  describe("constructor guards", function () {
    let F;
    beforeEach(async function () {
      F = await ethers.getContractFactory("AereProofAggregatorV2");
    });

    it("rejects a code-less SP1 verifier (V1 accepted one)", async function () {
      await expect(
        F.deploy("0x000000000000000000000000000000000000dEaD", owner.address, [], [], [], [], [], false)
      ).to.be.revertedWithCustomError(F, "VerifierHasNoCode");
    });

    it("rejects the zero verifier / zero owner when not sealed", async function () {
      await expect(F.deploy(ethers.ZeroAddress, owner.address, [], [], [], [], [], false)).to.be.revertedWithCustomError(
        F,
        "ZeroAddress"
      );
      await expect(
        F.deploy(await mock.getAddress(), ethers.ZeroAddress, [], [], [], [], [], false)
      ).to.be.revertedWithCustomError(F, "ZeroAddress");
    });

    it("rejects mismatched seed array lengths", async function () {
      const m = await mock.getAddress();
      await expect(F.deploy(m, owner.address, [AGG_VKEY], [], [], [], [], false)).to.be.revertedWithCustomError(
        F,
        "SeedLengthMismatch"
      );
      await expect(
        F.deploy(m, owner.address, [], [], digests, [], [], false)
      ).to.be.revertedWithCustomError(F, "SeedLengthMismatch");
    });

    it("rejects seeding a zero inner digest or zero agg vkey", async function () {
      const m = await mock.getAddress();
      await expect(
        F.deploy(m, owner.address, [], [], [ethers.ZeroHash], ["0x" + "11".repeat(32)], ["evil"], false)
      ).to.be.revertedWithCustomError(F, "ZeroVKey");
      await expect(
        F.deploy(m, owner.address, [ethers.ZeroHash], ["evil"], [], [], [], false)
      ).to.be.revertedWithCustomError(F, "ZeroVKey");
    });
  });
});