veFDT_mit_flash.diff

diff --git a/contracts/mocks/MockSmartWallet.sol b/contracts/mocks/MockSmartWallet.sol
index 78ce02b..7c0c6e9 100644
--- a/contracts/mocks/MockSmartWallet.sol
+++ b/contracts/mocks/MockSmartWallet.sol
@@ -42,3 +42,27 @@ contract MockSmartWallet {
         IVotingEscrow(ve).delegate(to);
     }
 }
+
+import "hardhat/console.sol";
+contract MockFlashloan {
+    event Attacked(uint votesHeldDuringAttack);
+
+    function attack(
+        IVotingEscrow ve,
+        IERC20        token,
+        address       owner, 
+        uint256       amount,
+        uint256       timestamp
+    ) external returns (uint256 votesHeld) {
+        token.transferFrom(owner, address(this), amount);
+        console.log("token balance:", token.balanceOf(address(this)));
+        token.approve(address(ve), amount);
+        ve.createLock(amount, timestamp);
+        console.log("lock created in block:", block.number, "@", block.timestamp);
+        votesHeld = ve.balanceOf(address(this));
+        console.log("votesHeld:", votesHeld);
+        emit Attacked(votesHeld);
+        ve.quitLock();
+        token.transfer(owner, token.balanceOf(address(this)));
+    }
+}
diff --git a/test/votingEscrowTest.ts b/test/votingEscrowTest.ts
index 7d3163d..d869335 100644
--- a/test/votingEscrowTest.ts
+++ b/test/votingEscrowTest.ts
@@ -115,6 +115,86 @@ describe("VotingEscrow Tests", function () {
     await restoreSnapshot(provider);
   });
 
+
+
+  // to run:
+  // rm -Rf 2022-08-fiatdao || true && git clone https://github.com/code-423n4/2022-08-fiatdao && cd 2022-08-fiatdao && npm install && npm run build && wget https://gist.githubusercontent.com/IllIllI000/67e2d48aa8b2cc33859d4ec962638c98/raw/4d697b50c7e08c33c13b263f9aa60667e4e6f162/veFDT_mit_flash.diff && git apply veFDT_mit_flash.diff && npx hardhat test --config hardhat.config.test.ts test/votingEscrowTest.ts 
+
+
+
+  describe("IllIllI Mitigation tests", async () => {
+    it("IllIllI - Quit in same block as creation, verify penalty", async () => {
+      const flashMockDeployer = await ethers.getContractFactory("MockFlashloan", admin);
+      let flashMock = await flashMockDeployer.deploy();
+      //let amountToLock = BigNumber.from("10");
+      let amountToLock = BigNumber.from("1000000000000000000000");
+      // approve the flash loan contract so it can take funds from caller
+      await fdtMock.connect(alice).approve(flashMock.address, amountToLock);
+
+      await provider.send("evm_setAutomine", [false]);
+      await provider.send("evm_setIntervalMining", [0]);
+      await provider.send("evm_mine", []);
+
+      let blockAnticipated = await getBlock() + 1;
+      let timeStartup = await getTimestamp(); 
+      let timeStart = (Math.trunc(timeStartup / WEEK) * WEEK) + WEEK + WEEK; // give us room to go backwards by adding another week
+      let timeActuallyLockedFrom = timeStart;
+      //timeStart += 284680;
+      timeStart = timeActuallyLockedFrom + 0;
+      //timeStart = timeActuallyLockedFrom + 1;
+      //timeStart = timeActuallyLockedFrom + 2;
+      //timeStart = timeActuallyLockedFrom - 1;
+      //timeStart = timeActuallyLockedFrom - 2;
+      //timeStart = timeActuallyLockedFrom - WEEK + 1;
+      //timeStart = timeActuallyLockedFrom + WEEK - 1;
+
+      await provider.send("evm_setNextBlockTimestamp", [ timeStart ]);
+
+      // alice creates a lock for some amount of time
+      let balanceStart = await fdtMock.balanceOf(alice.address);
+      let tx = await flashMock.connect(alice).attack(
+        ve.address, fdtMock.address, alice.address, amountToLock, timeStart + (2 * WEEK));
+
+      await provider.send("evm_mine", []);
+
+      let resp = await tx.wait();
+      if (!resp || !resp.events) return;
+      let amountHeld = BigNumber.from(0)
+      for (let event of resp.events) {
+        if ('Attacked' !== event.event || !event.args) {
+          continue
+        }
+        amountHeld = BigNumber.from(event.args[0])
+      }
+
+      let blockLocked = resp.blockNumber
+      expect(resp.blockNumber).to.equal(blockAnticipated);
+      let timeCreated = (await ethers.provider.getBlock(resp.blockNumber)).timestamp;
+      let timeLockUsed = (Math.trunc(timeCreated / WEEK) * WEEK);
+      console.log("timeCreated:", timeCreated, "timeLockUsed", timeLockUsed);
+
+      let timeQuit = (await ethers.provider.getBlock(resp.blockNumber)).timestamp;
+      expect(timeQuit - timeCreated).to.equal(0);
+      let blockBeforeQuit = resp.blockNumber - 1;
+      let blockAtQuit = resp.blockNumber;
+      console.log("mined:",resp.blockNumber);
+      console.log("block locked:", blockLocked, "block quit:", resp.blockNumber);
+
+      let votes = await ve.balanceOfAt(flashMock.address, blockAtQuit)
+      let penalty = await fdtMock.balanceOf(ve.address)
+      console.log("votes at start:", (await ve.balanceOfAt(flashMock.address, blockLocked)).toString());
+      console.log("votes at block before quit:", (await ve.balanceOfAt(flashMock.address, blockBeforeQuit)).toString());
+      console.log("votes at block:", votes.toString(), "penalty:", penalty.toString());
+      console.log("votes at block - penalty:", votes.sub(penalty).toString());
+      console.log("balance alice change:", balanceStart.sub(await fdtMock.balanceOf(alice.address)).toString(), "vs initial locked", amountToLock.toString());
+      console.log("votes held by contract during attack:", amountHeld.toString());
+      console.log("votes held - penalty:", amountHeld.sub(penalty).toString());
+      expect(votes).to.below(penalty);
+      expect(amountHeld).to.below(penalty);
+      expect(await ve.balanceOfAt(alice.address, blockAtQuit)).to.equal(0);
+    });
+  })
+/**
   describe("Deployment", async () => {
     it("Initialized properly", async () => {
       expect(await ve.owner()).to.equal(admin.address);
@@ -856,5 +936,5 @@ describe("VotingEscrow Tests", function () {
       await ve.connect(alice).quitLock();
       expect(await fdtMock.balanceOf(alice.address)).to.equal(aliceBalBefore);
     });
-  });
+  });/**/
 });