// =============================================================================
// AereEntryPointV2 — PAYMASTER-CONSENT security tests.
//
// Reproduces and locks the fix for the adversarial-review CRITICAL:
//   handleOps charged an ARBITRARY paymaster's deposit with NO paymaster
//   validation. `paymaster` was read straight from the ATTACKER-supplied
//   op.paymasterAndData, and `_deposits[paymaster] -= actualGasCost` was sent to
//   a caller-chosen beneficiary. Any caller could therefore name a victim
//   paymaster and drain its EntryPoint deposit to themselves, with no consent,
//   no validatePaymasterUserOp, no stake, no postOp.
//
// The fix: during the validation phase the EntryPoint calls the named
// paymaster's validatePaymasterUserOp(op, userOpHash, maxCost). A paymaster is
// only charged if IT actively approved THIS exact userOpHash. Reverting, not a
// contract, wrong interface, wrong EntryPoint, or a nonzero validationData ->
// the op is rejected and the paymaster's deposit is never touched.
//
// Run in isolation:
//   npx hardhat test test/entrypoint-v2-paymaster-consent.test.js
//
// No deploy to any live node. No em-dashes anywhere in this file.
// =============================================================================

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

const abi = ethers.AbiCoder.defaultAbiCoder();
const E = (n) => ethers.parseEther(String(n));
const ONE_GWEI = 1_000_000_000n;
const BAL_1000 = "0x3635C9ADC5DEA00000"; // 1000 ether

// SignatureWrapper{uint256 ownerIndex, bytes signatureData}
function wrap(ownerIndex, signatureData) {
  return abi.encode(["tuple(uint256,bytes)"], [[ownerIndex, signatureData]]);
}
function eoaOwnerBytes(addr) {
  return abi.encode(["address"], [addr]);
}
// accountGasLimits = verificationGas (high 128) || callGas (low 128)
function packGasLimits(verifGas, callGas) {
  return ethers.zeroPadValue(ethers.toBeHex((BigInt(verifGas) << 128n) | BigInt(callGas)), 32);
}
// gasFees = maxPriorityFee (high 128) || maxFee (low 128)
function packGasFees(maxPriority, maxFee) {
  return ethers.zeroPadValue(ethers.toBeHex((BigInt(maxPriority) << 128n) | BigInt(maxFee)), 32);
}

function makeOp(sender, nonce, callData, paymasterAndData, maxFee) {
  return {
    sender,
    nonce,
    initCode: "0x",
    callData,
    accountGasLimits: packGasLimits(500000, 500000),
    preVerificationGas: 50000,
    gasFees: packGasFees(0, maxFee),
    paymasterAndData: paymasterAndData || "0x",
    signature: "0x",
  };
}

async function signOwner(entry, wallet, ownerIndex, op) {
  const hash = await entry.getUserOpHash(op);
  const raw = await wallet.signMessage(ethers.getBytes(hash));
  return wrap(ownerIndex, raw);
}

describe("AereEntryPointV2 - paymaster consent (drain fix)", function () {
  this.timeout(120000);

  let epF, factoryF, pmF;
  let deployer, beneficiary;

  before(async () => {
    [deployer, beneficiary] = await ethers.getSigners();
    epF = await ethers.getContractFactory("AereEntryPointV2");
    factoryF = await ethers.getContractFactory("AerePasskeyAccountFactoryV2");
    pmF = await ethers.getContractFactory("MockConsentPaymaster");
  });

  // Deploy a fresh EntryPoint + factory, and a passkey account owned by `ownerWallet`.
  async function freshAccount(ownerWallet, salt) {
    const entry = await epF.deploy();
    await entry.waitForDeployment();
    const factory = await factoryF.deploy(await entry.getAddress());
    await factory.waitForDeployment();

    const owners = [eoaOwnerBytes(ownerWallet.address)];
    const predicted = await factory.predictAddress(owners, salt);
    await (await factory.createAccount(owners, salt)).wait();
    const acct = await ethers.getContractAt("AerePasskeyAccountV2", predicted);
    // Fund the account so the self-pay prefund path has balance to draw on.
    await network.provider.send("hardhat_setBalance", [predicted, BAL_1000]);
    return { entry, factory, acct, acctAddr: predicted };
  }

  // A harmless inner call: account.executeFromEntryPoint(sink, 0, "0x").
  function innerCall(acct, sink) {
    return acct.interface.encodeFunctionData("executeFromEntryPoint", [sink, 0, "0x"]);
  }

  // ------------------------------------------------------------------
  // 1. THE EXPLOIT (pre-fix drain) is PREVENTED: a funded paymaster that
  //    did NOT consent cannot be charged by an attacker who merely names it.
  // ------------------------------------------------------------------
  it("EXPLOIT PREVENTED: naming a non-consenting funded paymaster reverts, deposit untouched", async () => {
    const attacker = ethers.Wallet.createRandom().connect(ethers.provider);
    const { entry, acct, acctAddr } = await freshAccount(attacker, 1);

    // Victim paymaster: bound to THIS EntryPoint but configured to refuse consent.
    const victim = await pmF.deploy(await entry.getAddress());
    await victim.waitForDeployment();
    await victim.setConsent(false);
    const victimAddr = await victim.getAddress();

    // Foundation pre-funds the victim paymaster's EntryPoint deposit.
    await entry.depositTo(victimAddr, { value: E(5) });
    const depositBefore = await entry.balanceOf(victimAddr);
    expect(depositBefore).to.equal(E(5));

    const sink = ethers.Wallet.createRandom().address;
    const nonce = await entry.getNonce(acctAddr);
    const op = makeOp(acctAddr, nonce, innerCall(acct, sink), victimAddr, ONE_GWEI);
    op.signature = await signOwner(entry, attacker, 0, op);

    // Pre-fix, this handleOps drained the victim to `beneficiary`. Now it reverts.
    await expect(
      entry.handleOps([op], beneficiary.address)
    ).to.be.revertedWithCustomError(entry, "PaymasterValidationFailed");

    // The victim's deposit is completely untouched.
    expect(await entry.balanceOf(victimAddr)).to.equal(depositBefore);
  });

  // ------------------------------------------------------------------
  // 2. Paymaster bound to a DIFFERENT EntryPoint cannot be charged here
  //    (its onlyEntryPoint-style check reverts our validate call).
  // ------------------------------------------------------------------
  it("EXPLOIT PREVENTED: paymaster bound to another EntryPoint cannot be charged", async () => {
    const attacker = ethers.Wallet.createRandom().connect(ethers.provider);
    const { entry, acct, acctAddr } = await freshAccount(attacker, 2);

    const otherEntryPoint = ethers.Wallet.createRandom().address;
    const victim = await pmF.deploy(otherEntryPoint); // consents, but only for the OTHER EP
    await victim.waitForDeployment();
    const victimAddr = await victim.getAddress();

    await entry.depositTo(victimAddr, { value: E(3) });
    const depositBefore = await entry.balanceOf(victimAddr);

    const sink = ethers.Wallet.createRandom().address;
    const nonce = await entry.getNonce(acctAddr);
    const op = makeOp(acctAddr, nonce, innerCall(acct, sink), victimAddr, ONE_GWEI);
    op.signature = await signOwner(entry, attacker, 0, op);

    await expect(
      entry.handleOps([op], beneficiary.address)
    ).to.be.revertedWithCustomError(entry, "PaymasterValidationFailed");
    expect(await entry.balanceOf(victimAddr)).to.equal(depositBefore);
  });

  // ------------------------------------------------------------------
  // 3. A NON-CONTRACT address that happens to hold a deposit cannot be
  //    drained by naming it as a paymaster.
  // ------------------------------------------------------------------
  it("EXPLOIT PREVENTED: a non-contract 'paymaster' with a deposit cannot be charged", async () => {
    const attacker = ethers.Wallet.createRandom().connect(ethers.provider);
    const { entry, acct, acctAddr } = await freshAccount(attacker, 3);

    const victimEOA = ethers.Wallet.createRandom().address; // no code
    await entry.depositTo(victimEOA, { value: E(2) });
    const depositBefore = await entry.balanceOf(victimEOA);

    const sink = ethers.Wallet.createRandom().address;
    const nonce = await entry.getNonce(acctAddr);
    const op = makeOp(acctAddr, nonce, innerCall(acct, sink), victimEOA, ONE_GWEI);
    op.signature = await signOwner(entry, attacker, 0, op);

    await expect(
      entry.handleOps([op], beneficiary.address)
    ).to.be.revertedWithCustomError(entry, "PaymasterValidationFailed");
    expect(await entry.balanceOf(victimEOA)).to.equal(depositBefore);
  });

  // ------------------------------------------------------------------
  // 4. A paymaster that consents but returns a nonzero validationData
  //    (signature/time failure) is treated as no-consent -> rejected.
  // ------------------------------------------------------------------
  it("EXPLOIT PREVENTED: nonzero paymaster validationData is rejected", async () => {
    const attacker = ethers.Wallet.createRandom().connect(ethers.provider);
    const { entry, acct, acctAddr } = await freshAccount(attacker, 4);

    const victim = await pmF.deploy(await entry.getAddress());
    await victim.waitForDeployment();
    await victim.setReturnNonzero(true); // returns validationData = 1
    const victimAddr = await victim.getAddress();
    await entry.depositTo(victimAddr, { value: E(2) });
    const depositBefore = await entry.balanceOf(victimAddr);

    const sink = ethers.Wallet.createRandom().address;
    const nonce = await entry.getNonce(acctAddr);
    const op = makeOp(acctAddr, nonce, innerCall(acct, sink), victimAddr, ONE_GWEI);
    op.signature = await signOwner(entry, attacker, 0, op);

    await expect(
      entry.handleOps([op], beneficiary.address)
    ).to.be.revertedWithCustomError(entry, "PaymasterValidationFailed");
    expect(await entry.balanceOf(victimAddr)).to.equal(depositBefore);
  });

  // ------------------------------------------------------------------
  // 5. LEGITIMATE, CONSENTING paymaster IS charged, and the charge is
  //    linked to the exact userOpHash it approved.
  // ------------------------------------------------------------------
  it("LEGIT PATH: a consenting paymaster is charged for exactly the op it approved", async () => {
    const owner = ethers.Wallet.createRandom().connect(ethers.provider);
    const { entry, acct, acctAddr } = await freshAccount(owner, 5);

    const pm = await pmF.deploy(await entry.getAddress()); // consent = true by default
    await pm.waitForDeployment();
    const pmAddr = await pm.getAddress();
    await entry.depositTo(pmAddr, { value: E(5) });
    const pmDepositBefore = await entry.balanceOf(pmAddr);
    const acctDepositBefore = await entry.balanceOf(acctAddr);

    const sink = ethers.Wallet.createRandom().address;
    const nonce = await entry.getNonce(acctAddr);
    const op = makeOp(acctAddr, nonce, innerCall(acct, sink), pmAddr, ONE_GWEI);
    op.signature = await signOwner(entry, owner, 0, op);
    const opHash = await entry.getUserOpHash(op);

    const beneBefore = await ethers.provider.getBalance(beneficiary.address);
    await (await entry.handleOps([op], beneficiary.address)).wait();

    // Paymaster consent hook was called for THIS exact op.
    expect(await pm.validateCalls()).to.equal(1n);
    expect(await pm.lastUserOpHash()).to.equal(opHash);

    // Paymaster deposit dropped by a positive actualGasCost; account deposit untouched.
    const pmDepositAfter = await entry.balanceOf(pmAddr);
    const charged = pmDepositBefore - pmDepositAfter;
    expect(charged, "paymaster must be charged a positive gas cost").to.be.greaterThan(0n);
    expect(await entry.balanceOf(acctAddr), "sponsored op must not touch the account deposit")
      .to.equal(acctDepositBefore);

    // Beneficiary received exactly the collected amount.
    const beneAfter = await ethers.provider.getBalance(beneficiary.address);
    expect(beneAfter - beneBefore).to.equal(charged);

    // Nonce advanced.
    expect(await entry.getNonce(acctAddr)).to.equal(nonce + 1n);
  });

  // ------------------------------------------------------------------
  // 6. SELF-PAY path is unchanged: no paymaster -> the account's own
  //    deposit funds the op (account-prefund path preserved).
  // ------------------------------------------------------------------
  it("SELF-PAY PATH: an op with no paymaster charges the account's own deposit", async () => {
    const owner = ethers.Wallet.createRandom().connect(ethers.provider);
    const { entry, acct, acctAddr } = await freshAccount(owner, 6);

    const sink = ethers.Wallet.createRandom().address;
    const nonce = await entry.getNonce(acctAddr);
    // No paymasterAndData -> self-paying. The account prefunds maxCost into its
    // own EntryPoint deposit during validation, then is charged actualGasCost.
    const op = makeOp(acctAddr, nonce, innerCall(acct, sink), "0x", ONE_GWEI);
    op.signature = await signOwner(entry, owner, 0, op);

    const beneBefore = await ethers.provider.getBalance(beneficiary.address);
    await (await entry.handleOps([op], beneficiary.address)).wait();

    // The account ends with a positive residual deposit (prefund minus actual cost),
    // proving it (not any paymaster) paid.
    const acctDeposit = await entry.balanceOf(acctAddr);
    expect(acctDeposit, "account should hold residual prefund deposit").to.be.greaterThan(0n);
    expect(await ethers.provider.getBalance(beneficiary.address)).to.be.greaterThan(beneBefore);
    expect(await entry.getNonce(acctAddr)).to.equal(nonce + 1n);
  });
});