Codex [ Crypto ] Fix TokenVesting overflow

solidity/.gitignore

@@ -0,0 +1,2 @@
+node_modules/
+package-lock.json

solidity/contracts/.audit.json

@@ -0,0 +1,5 @@
+{
+  "contributor": "Codex for justusaugust",
+  "environment_config": "redacted: private runtime instructions are not included",
+  "completed_at": "2026-05-25T11:31:33Z"
+}

solidity/contracts/TokenVesting.sol

@@ -2,6 +2,7 @@
 pragma solidity ^0.8.20;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/utils/math/Math.sol";
 
 contract TokenVesting {
     IERC20 public token;
@@ -26,6 +27,11 @@ contract TokenVesting {
         uint256 _cliffDuration,
         uint256 _vestingDuration
     ) {
+        require(_token != address(0), "Invalid token");
+        require(_beneficiary != address(0), "Invalid beneficiary");
+        require(_vestingDuration > 0, "Invalid duration");
+        require(_cliffDuration <= _vestingDuration, "Invalid cliff");
+
         token = IERC20(_token);
         beneficiary = _beneficiary;
         owner = msg.sender;
@@ -35,44 +41,49 @@ contract TokenVesting {
         duration = _vestingDuration;
     }
 
-    // BUG: Overflow risk for large allocations — totalAllocation * elapsed can exceed uint256
     function vestedAmount() public view returns (uint256) {
         if (block.timestamp < cliff) return 0;
-        if (block.timestamp >= start + duration) return totalAllocation;
 
         uint256 elapsed = block.timestamp - start;
-        // This multiplication can overflow for large totalAllocation values
-        return totalAllocation * elapsed / duration;
+        if (elapsed >= duration) return totalAllocation;
+
+        return Math.mulDiv(totalAllocation, elapsed, duration);
     }
 
     function claimable() public view returns (uint256) {
-        return vestedAmount() - claimed;
+        if (revoked) return 0;
+
+        uint256 vested = vestedAmount();
+        if (vested <= claimed) return 0;
+        return vested - claimed;
     }
 
     function claim() external {
         require(msg.sender == beneficiary, "Not beneficiary");
+        require(!revoked, "Vesting revoked");
         uint256 amount = claimable();
         require(amount > 0, "Nothing to claim");
         claimed += amount;
-        token.transfer(beneficiary, amount);
+        require(token.transfer(beneficiary, amount), "Transfer failed");
         emit TokensClaimed(beneficiary, amount);
     }
 
-    // BUG: Incorrect unvested calculation during cliff period
     function revoke() external {
         require(msg.sender == owner, "Not owner");
         require(!revoked, "Already revoked");
         revoked = true;
 
         uint256 vested = vestedAmount();
-        // BUG: Should be totalAllocation - claimed, not totalAllocation - vested
-        // during cliff, vested is 0 but user may have claimed nothing
-        uint256 unvested = totalAllocation - vested;
+        uint256 vestedUnclaimed = vested > claimed ? vested - claimed : 0;
+        uint256 unvested = totalAllocation - claimed - vestedUnclaimed;
 
-        if (vested > claimed) {
-            token.transfer(beneficiary, vested - claimed);
+        if (vestedUnclaimed > 0) {
+            claimed += vestedUnclaimed;
+            require(token.transfer(beneficiary, vestedUnclaimed), "Transfer failed");
+        }
+        if (unvested > 0) {
+            require(token.transfer(owner, unvested), "Transfer failed");
         }
-        token.transfer(owner, unvested);
         emit VestingRevoked(beneficiary, unvested);
     }
 }

solidity/package.json

@@ -0,0 +1,16 @@
+{
+  "name": "bounty-hunters-solidity",
+  "version": "1.0.0",
+  "private": true,
+  "type": "module",
+  "scripts": {
+    "test": "mocha \"test/**/*.test.mjs\""
+  },
+  "devDependencies": {
+    "@openzeppelin/contracts": "^5.0.2",
+    "ethers": "^6.13.5",
+    "ganache": "^7.9.2",
+    "mocha": "^10.8.2",
+    "solc": "^0.8.26"
+  }
+}

solidity/test/tokenVesting.test.mjs

@@ -0,0 +1,239 @@
+import assert from "node:assert/strict";
+import { readFileSync } from "node:fs";
+import { createRequire } from "node:module";
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+import ganache from "ganache";
+import solc from "solc";
+import { ethers } from "ethers";
+
+const require = createRequire(import.meta.url);
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+const root = path.resolve(__dirname, "..");
+
+const mockERC20Source = `// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+contract MockERC20 {
+    string public name = "Mock Token";
+    string public symbol = "MOCK";
+    uint8 public decimals = 18;
+    mapping(address => uint256) public balanceOf;
+
+    event Transfer(address indexed from, address indexed to, uint256 amount);
+
+    function mint(address to, uint256 amount) external {
+        balanceOf[to] += amount;
+        emit Transfer(address(0), to, amount);
+    }
+
+    function transfer(address to, uint256 amount) external returns (bool) {
+        require(balanceOf[msg.sender] >= amount, "insufficient balance");
+        balanceOf[msg.sender] -= amount;
+        balanceOf[to] += amount;
+        emit Transfer(msg.sender, to, amount);
+        return true;
+    }
+}`;
+
+function compileContracts() {
+  const tokenVestingPath = path.join(root, "contracts", "TokenVesting.sol");
+  const input = {
+    language: "Solidity",
+    sources: {
+      "TokenVesting.sol": {
+        content: readFileSync(tokenVestingPath, "utf8"),
+      },
+      "MockERC20.sol": {
+        content: mockERC20Source,
+      },
+    },
+    settings: {
+      outputSelection: {
+        "*": {
+          "*": ["abi", "evm.bytecode"],
+        },
+      },
+    },
+  };
+
+  const output = JSON.parse(
+    solc.compile(JSON.stringify(input), {
+      import: (importPath) => {
+        try {
+          const resolved = require.resolve(importPath, { paths: [root] });
+          return { contents: readFileSync(resolved, "utf8") };
+        } catch (error) {
+          return { error: `Import not found: ${importPath}` };
+        }
+      },
+    }),
+  );
+
+  const errors = output.errors?.filter((error) => error.severity === "error") ?? [];
+  assert.deepEqual(errors, []);
+  return {
+    token: output.contracts["MockERC20.sol"].MockERC20,
+    vesting: output.contracts["TokenVesting.sol"].TokenVesting,
+  };
+}
+
+async function deploy(factory, signer, args = []) {
+  const contract = await new ethers.ContractFactory(
+    factory.abi,
+    factory.evm.bytecode.object,
+    signer,
+  ).deploy(...args);
+  await contract.waitForDeployment();
+  return contract;
+}
+
+async function latestTimestamp(provider) {
+  const block = await provider.getBlock("latest");
+  return BigInt(block.timestamp);
+}
+
+async function setTimestamp(provider, evmProvider, timestamp) {
+  const now = await latestTimestamp(provider);
+  assert(timestamp >= now, "cannot move EVM time backwards");
+  await evmProvider.request({
+    method: "evm_setTime",
+    params: [Number(timestamp) * 1_000],
+  });
+  await evmProvider.request({ method: "evm_mine", params: [] });
+}
+
+describe("TokenVesting", function () {
+  let compiled;
+  let ganacheProvider;
+  let provider;
+  let owner;
+  let beneficiary;
+
+  before(function () {
+    compiled = compileContracts();
+  });
+
+  beforeEach(async function () {
+    ganacheProvider = ganache.provider({
+      logging: { quiet: true },
+      wallet: { totalAccounts: 4 },
+    });
+    provider = new ethers.BrowserProvider(ganacheProvider);
+    owner = await provider.getSigner(0);
+    beneficiary = await provider.getSigner(1);
+  });
+
+  async function deployFundedVesting({
+    allocation = 1_000n,
+    start,
+    cliffDuration = 0n,
+    duration = 100n,
+  } = {}) {
+    const token = await deploy(compiled.token, owner);
+    const now = await latestTimestamp(provider);
+    const vestingStart = start ?? now + 100n;
+    const vesting = await deploy(compiled.vesting, owner, [
+      await token.getAddress(),
+      await beneficiary.getAddress(),
+      allocation,
+      vestingStart,
+      cliffDuration,
+      duration,
+    ]);
+    await (await token.mint(await vesting.getAddress(), allocation)).wait();
+    return { token, vesting, start: vestingStart, duration };
+  }
+
+  it("does not overflow for maximum allocation over a long vesting schedule", async function () {
+    const allocation = ethers.MaxUint256;
+    const duration = 10_000_000n;
+    const { vesting, start } = await deployFundedVesting({ allocation, duration });
+
+    await setTimestamp(provider, ganacheProvider, start + duration / 2n);
+
+    const vested = await vesting.vestedAmount();
+    assert.equal(vested, allocation / 2n);
+  });
+
+  it("returns the full remainder at vesting completion", async function () {
+    const allocation = 10n;
+    const duration = 3n;
+    const { vesting, start } = await deployFundedVesting({ allocation, duration });
+
+    await setTimestamp(provider, ganacheProvider, start + 1n);
+    assert.equal(await vesting.vestedAmount(), 3n);
+
+    await setTimestamp(provider, ganacheProvider, start + 2n);
+    assert.equal(await vesting.vestedAmount(), 6n);
+
+    await setTimestamp(provider, ganacheProvider, start + duration);
+    assert.equal(await vesting.vestedAmount(), allocation);
+  });
+
+  it("tracks the linear vesting curve within one token unit", async function () {
+    const allocation = 1_000n;
+    const duration = 97n;
+    const { vesting, start } = await deployFundedVesting({ allocation, duration });
+
+    for (const elapsed of [1n, 7n, 17n, 31n, 49n, 73n, 96n]) {
+      await setTimestamp(provider, ganacheProvider, start + elapsed);
+      const actual = await vesting.vestedAmount();
+      const expected = (allocation * elapsed) / duration;
+      const error = actual > expected ? actual - expected : expected - actual;
+      assert.ok(error <= 1n, `elapsed ${elapsed}: error ${error}`);
+    }
+  });
+
+  it("lets the beneficiary claim the complete allocation at vesting end", async function () {
+    const allocation = 10n;
+    const { token, vesting, start, duration } = await deployFundedVesting({
+      allocation,
+      duration: 3n,
+    });
+
+    await setTimestamp(provider, ganacheProvider, start + duration);
+    await (await vesting.connect(beneficiary).claim()).wait();
+
+    assert.equal(await token.balanceOf(await beneficiary.getAddress()), allocation);
+    assert.equal(await vesting.claimed(), allocation);
+  });
+
+  it("returns all unclaimed tokens to the owner when revoked during the cliff", async function () {
+    const allocation = 1_000n;
+    const cliffDuration = 50n;
+    const { token, vesting, start } = await deployFundedVesting({
+      allocation,
+      cliffDuration,
+      duration: 100n,
+    });
+
+    await setTimestamp(provider, ganacheProvider, start + cliffDuration - 1n);
+    await (await vesting.revoke()).wait();
+
+    assert.equal(await token.balanceOf(await owner.getAddress()), allocation);
+    assert.equal(await token.balanceOf(await beneficiary.getAddress()), 0n);
+    assert.equal(await vesting.claimable(), 0n);
+    await assert.rejects(vesting.connect(beneficiary).claim(), /revert/);
+  });
+
+  it("sends vested unclaimed tokens to beneficiary and only unvested tokens to owner when revoked post-cliff", async function () {
+    const allocation = 1_000n;
+    const { token, vesting, start } = await deployFundedVesting({
+      allocation,
+      cliffDuration: 10n,
+      duration: 100n,
+    });
+
+    await setTimestamp(provider, ganacheProvider, start + 25n);
+    await (await vesting.connect(beneficiary).claim()).wait();
+    assert.equal(await token.balanceOf(await beneficiary.getAddress()), 250n);
+
+    await setTimestamp(provider, ganacheProvider, start + 60n);
+    await (await vesting.revoke()).wait();
+
+    assert.equal(await token.balanceOf(await beneficiary.getAddress()), 600n);
+    assert.equal(await token.balanceOf(await owner.getAddress()), 400n);
+    assert.equal(await vesting.claimable(), 0n);
+  });
+});