const { expect } = require("chai");
const { ethers } = require("hardhat");

// FINDING 3 (medium): EVM-validity canonical binding.
//
// V1 (AereEVMValidity / AereEVMValidityBatch) records an SP1 postStateRoot as a
// block WITHOUT binding it to the real chain: no blockHash/parentHash in the
// public values, no blockhash() check, no inter-record continuity. A genuine SP1
// proof of a FABRICATED execution can therefore be recorded as "canonical", and
// first-write-wins dedup lets an attacker squat slot N.
//
// This suite (1) reproduces the attack on V1, then (2) proves V2 closes it:
//   - a fabricated proof with a wrong/absent committed blockHash is REJECTED,
//   - a real in-window proof with the correct committed blockHash is accepted,
//   - batch V2 enforces inter-record continuity.
//
// The SP1 verifier is mocked (MockSp1Verifier: verifyProof reverts unless the
// exact (vkey, publicValues, proof) triple was pre-marked valid), and blockhash
// binding is exercised against real hardhat block hashes.

const coder = ethers.AbiCoder.defaultAbiCoder();
const VKEY = "0x" + "ab".repeat(32);
const PROOF = "0x4388a21c" + "00".repeat(64); // dummy selector-prefixed bytes
const ZERO32 = "0x" + "00".repeat(32);

function marker(vkey, pv, proof) {
  return ethers.keccak256(
    coder.encode(["bytes32", "bytes", "bytes"], [vkey, pv, proof])
  );
}

// V1 single-block public values: 7 words.
function encodeV1(o) {
  return coder.encode(
    ["uint256", "uint256", "bytes32", "bytes32", "bytes32", "uint256", "uint256"],
    [o.chainId, o.blockNumber, o.prevStateRoot, o.postStateRoot, o.txRoot, o.gasUsed, o.numTxns]
  );
}

// V2 single-block public values: 9 words (adds parentHash + blockHash).
function encodeV2(o) {
  return coder.encode(
    ["uint256", "uint256", "bytes32", "bytes32", "bytes32", "bytes32", "bytes32", "uint256", "uint256"],
    [o.chainId, o.blockNumber, o.parentHash, o.blockHash, o.prevStateRoot, o.postStateRoot, o.txRoot, o.gasUsed, o.numTxns]
  );
}

// V2 batch public values: 10 words.
function encodeBatchV2(o) {
  return coder.encode(
    ["uint256", "uint256", "uint256", "uint256", "bytes32", "bytes32", "bytes32", "bytes32", "uint256", "uint256"],
    [o.chainId, o.firstBlock, o.lastBlock, o.numBlocks, o.firstParentHash, o.lastBlockHash, o.prevStateRoot, o.postStateRoot, o.totalGas, o.numTxns]
  );
}

const FAKE_POST = "0x" + "fa".repeat(32);
const PREV_ROOT = "0x" + "11".repeat(32);
const TX_ROOT = "0x" + "22".repeat(32);

describe("FINDING 3 - EVM-validity canonical binding", () => {
  let mock, chainId;

  beforeEach(async () => {
    const M = await ethers.getContractFactory("MockSp1Verifier");
    mock = await M.deploy();
    await mock.waitForDeployment();
    chainId = (await ethers.provider.getNetwork()).chainId; // 31337 on hardhat
  });

  // ------------------------------------------------------------------ V1 hole
  describe("V1 attack surface (documents the vulnerability)", () => {
    let v1;
    beforeEach(async () => {
      const F = await ethers.getContractFactory("AereEVMValidity");
      v1 = await F.deploy(await mock.getAddress(), VKEY);
      await v1.waitForDeployment();
    });

    it("records a FABRICATED execution as canonical (no blockhash binding)", async () => {
      // A genuine-for-vkey proof whose committed postStateRoot is fabricated and
      // whose blockNumber never existed as claimed. V1 has no way to tell.
      const pv = encodeV1({
        chainId, blockNumber: 999_000, prevStateRoot: PREV_ROOT,
        postStateRoot: FAKE_POST, txRoot: TX_ROOT, gasUsed: 21000, numTxns: 1,
      });
      await mock.setValid(marker(VKEY, pv, PROOF), true);

      await v1.recordBlock(pv, PROOF); // ATTACK SUCCEEDS: no revert
      expect(await v1.provenCount()).to.equal(1n);
      const rec = await v1.getRecord(0);
      expect(rec[2]).to.equal(FAKE_POST); // fabricated root anchored as canonical
    });

    it("lets an attacker SQUAT slot N via first-write-wins", async () => {
      const pvFake = encodeV1({
        chainId, blockNumber: 999_000, prevStateRoot: PREV_ROOT,
        postStateRoot: FAKE_POST, txRoot: TX_ROOT, gasUsed: 21000, numTxns: 1,
      });
      await mock.setValid(marker(VKEY, pvFake, PROOF), true);
      await v1.recordBlock(pvFake, PROOF);

      // A later, DIFFERENT (even honest) proof for the same height is locked out.
      const pvReal = encodeV1({
        chainId, blockNumber: 999_000, prevStateRoot: PREV_ROOT,
        postStateRoot: "0x" + "cc".repeat(32), txRoot: TX_ROOT, gasUsed: 42000, numTxns: 2,
      });
      await mock.setValid(marker(VKEY, pvReal, PROOF), true);
      await expect(v1.recordBlock(pvReal, PROOF)).to.be.revertedWithCustomError(
        v1, "AlreadyProven"
      );
    });
  });

  // ------------------------------------------------------------------ V2 fix
  describe("V2 single-block canonical binding", () => {
    let v2;
    beforeEach(async () => {
      const F = await ethers.getContractFactory("AereEVMValidityV2");
      v2 = await F.deploy(await mock.getAddress(), VKEY);
      await v2.waitForDeployment();
    });

    it("REJECTS a fabricated in-window block whose committed blockHash is wrong", async () => {
      const bn = await ethers.provider.getBlockNumber();
      const parent = await ethers.provider.getBlock(bn - 1);
      const pv = encodeV2({
        chainId, blockNumber: bn,
        parentHash: parent.hash,
        blockHash: "0x" + "de".repeat(32), // WRONG: not the real hash of block bn
        prevStateRoot: PREV_ROOT, postStateRoot: FAKE_POST, txRoot: TX_ROOT,
        gasUsed: 21000, numTxns: 1,
      });
      await mock.setValid(marker(VKEY, pv, PROOF), true);
      await expect(v2.recordBlock(pv, PROOF)).to.be.revertedWithCustomError(
        v2, "BlockHashMismatch"
      );
    });

    it("REJECTS a block whose hash is unavailable (out of window / future)", async () => {
      const bn = await ethers.provider.getBlockNumber();
      const pv = encodeV2({
        chainId, blockNumber: bn + 5000, // future: blockhash() == 0
        parentHash: ZERO32, blockHash: "0x" + "de".repeat(32),
        prevStateRoot: PREV_ROOT, postStateRoot: FAKE_POST, txRoot: TX_ROOT,
        gasUsed: 21000, numTxns: 1,
      });
      await mock.setValid(marker(VKEY, pv, PROOF), true);
      await expect(v2.recordBlock(pv, PROOF)).to.be.revertedWithCustomError(
        v2, "BlockHashUnavailable"
      );
    });

    it("ACCEPTS a real in-window proof with the correct committed blockHash", async () => {
      const bn = await ethers.provider.getBlockNumber();
      const blk = await ethers.provider.getBlock(bn);
      const parent = await ethers.provider.getBlock(bn - 1);
      const pv = encodeV2({
        chainId, blockNumber: bn,
        parentHash: parent.hash, blockHash: blk.hash,
        prevStateRoot: PREV_ROOT, postStateRoot: "0x" + "77".repeat(32), txRoot: TX_ROOT,
        gasUsed: 21000, numTxns: 1,
      });
      await mock.setValid(marker(VKEY, pv, PROOF), true); // mines a block
      await v2.recordBlock(pv, PROOF); // bn still in the 256 window here
      expect(await v2.provenCount()).to.equal(1n);
      const rec = await v2.getRecord(0);
      expect(rec[0]).to.equal(BigInt(bn));
      expect(rec[2]).to.equal(blk.hash); // blockHash stored (getRecord field #3)
    });

    it("REJECTS a wrong committed parentHash even when blockHash is right", async () => {
      const bn = await ethers.provider.getBlockNumber();
      const blk = await ethers.provider.getBlock(bn);
      const pv = encodeV2({
        chainId, blockNumber: bn,
        parentHash: "0x" + "ee".repeat(32), // WRONG parent
        blockHash: blk.hash,
        prevStateRoot: PREV_ROOT, postStateRoot: "0x" + "77".repeat(32), txRoot: TX_ROOT,
        gasUsed: 21000, numTxns: 1,
      });
      await mock.setValid(marker(VKEY, pv, PROOF), true);
      await expect(v2.recordBlock(pv, PROOF)).to.be.revertedWithCustomError(
        v2, "ParentHashMismatch"
      );
    });
  });

  describe("V2 batch continuity + tip binding", () => {
    let b2;
    beforeEach(async () => {
      const F = await ethers.getContractFactory("AereEVMValidityBatchV2");
      b2 = await F.deploy(await mock.getAddress(), VKEY);
      await b2.waitForDeployment();
    });

    // Bind a batch's tip to a real recent block after mining a margin of blocks,
    // so logical firstBlock..lastBlock ranges never go negative.
    async function freshTip() {
      await ethers.provider.send("hardhat_mine", ["0x14"]); // mine 20 blocks
      const l = await ethers.provider.getBlockNumber();
      const h = (await ethers.provider.getBlock(l)).hash;
      return { l, h };
    }

    it("REJECTS a batch whose tip blockHash is wrong", async () => {
      const { l } = await freshTip();
      const pv = encodeBatchV2({
        chainId, firstBlock: l - 2, lastBlock: l, numBlocks: 3,
        firstParentHash: ZERO32, lastBlockHash: "0x" + "de".repeat(32), // wrong
        prevStateRoot: PREV_ROOT, postStateRoot: "0x" + "a1".repeat(32),
        totalGas: 63000, numTxns: 3,
      });
      await mock.setValid(marker(VKEY, pv, PROOF), true);
      await expect(b2.recordBatch(pv, PROOF)).to.be.revertedWithCustomError(
        b2, "BlockHashMismatch"
      );
    });

    it("ACCEPTS a first batch, then ENFORCES continuity on the next", async () => {
      const R1 = "0x" + "a1".repeat(32);
      const R2 = "0x" + "b2".repeat(32);

      // First batch [la-2 .. la] anchors the chain (no continuity needed yet).
      const a = await freshTip();
      const la = a.l;
      const pvA = encodeBatchV2({
        chainId, firstBlock: la - 2, lastBlock: la, numBlocks: 3,
        firstParentHash: ZERO32, lastBlockHash: a.h,
        prevStateRoot: PREV_ROOT, postStateRoot: R1, totalGas: 63000, numTxns: 3,
      });
      await mock.setValid(marker(VKEY, pvA, PROOF), true);
      await b2.recordBatch(pvA, PROOF);
      expect(await b2.provenCount()).to.equal(1n);

      // Discontinuous firstBlock (gap) is rejected. lb is far above la+5.
      const g = await freshTip();
      const pvGap = encodeBatchV2({
        chainId, firstBlock: la + 5, lastBlock: g.l, numBlocks: g.l - (la + 5) + 1,
        firstParentHash: ZERO32, lastBlockHash: g.h,
        prevStateRoot: R1, postStateRoot: R2, totalGas: 42000, numTxns: 2,
      });
      await mock.setValid(marker(VKEY, pvGap, PROOF), true);
      await expect(b2.recordBatch(pvGap, PROOF)).to.be.revertedWithCustomError(
        b2, "Discontinuous"
      );

      // Correct firstBlock but wrong prevStateRoot is rejected.
      const br = await freshTip();
      const pvBadRoot = encodeBatchV2({
        chainId, firstBlock: la + 1, lastBlock: br.l, numBlocks: br.l - (la + 1) + 1,
        firstParentHash: ZERO32, lastBlockHash: br.h,
        prevStateRoot: "0x" + "99".repeat(32), postStateRoot: R2, totalGas: 42000, numTxns: 2,
      });
      await mock.setValid(marker(VKEY, pvBadRoot, PROOF), true);
      await expect(b2.recordBatch(pvBadRoot, PROOF)).to.be.revertedWithCustomError(
        b2, "Discontinuous"
      );

      // Contiguous + matching prevStateRoot is accepted.
      const gd = await freshTip();
      const pvGood = encodeBatchV2({
        chainId, firstBlock: la + 1, lastBlock: gd.l, numBlocks: gd.l - (la + 1) + 1,
        firstParentHash: ZERO32, lastBlockHash: gd.h,
        prevStateRoot: R1, postStateRoot: R2, totalGas: 42000, numTxns: 2,
      });
      await mock.setValid(marker(VKEY, pvGood, PROOF), true);
      await b2.recordBatch(pvGood, PROOF);
      expect(await b2.provenCount()).to.equal(2n);
    });
  });
});