TimeLock.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.7.6;

import "forge-std/Test.sol";


// This contract is designed to act as a time vault.
// User can deposit into this contract but cannot withdraw for atleast a week.
// User can also extend the wait time beyond the 1 week waiting period.

/*
1. Alice and bob both have 1 Ether balance
2. Deploy TimeLock Contract
3. Alice and bob both deposit 1 Ether to TimeLock, they need to wait 1 week to unlock Ether
4. Bob caused an overflow on his lockTime
5, Alice can't withdraw 1 Ether, because the lock time not expired.
6. Bob can withdraw 1 Ether, because the lockTime is overflow to 0

What happened?
Attack caused the TimeLock.lockTime to overflow,
and was able to withdraw before the 1 week waiting period.
*/

contract TimeLock {
    mapping(address => uint) public balances;
    mapping(address => uint) public lockTime;

    function deposit() external payable {
        balances[msg.sender] += msg.value;
        lockTime[msg.sender] = block.timestamp + 1 weeks;
    }

    function increaseLockTime(uint _secondsToIncrease) public {
        // If you send data which is max(uint)+1 it will go back to 0 causing the lock time to be 0
        // This can allow anyone to withdraw without waiting
        lockTime[msg.sender] += _secondsToIncrease; // vulnerable
    }

    function withdraw() public {
        require(balances[msg.sender] > 0, "Insufficient funds");
        require(block.timestamp > lockTime[msg.sender], "Lock time not expired");

        uint amount = balances[msg.sender];
        balances[msg.sender] = 0;

        (bool sent, ) = msg.sender.call{value: amount}("");
        require(sent, "Failed to send Ether");
    }
}


contract TimeLockTest is Test {
    TimeLock public timeLock;
    address alice;

    function setUp() public {
        timeLock = new TimeLock();
        alice = vm.addr(1);
        vm.deal(alice, 100 ether);   
    }

    // Fuzzer to test out and find failure 
    // FOUNDRY_FUZZ_RUNS defaults to 256 and it can be set to a higher value if needed
    function testfuzz(uint time) public {
        vm.startPrank(alice);
        timeLock.deposit{value: 1 ether}();
        timeLock.increaseLockTime(time);
        uint lockTime = timeLock.lockTime(alice);
        console.log("The Block Timestamp", block.timestamp + 1 weeks);
        require(lockTime >=  block.timestamp + 1 weeks);
        vm.stopPrank();
    }

    // Succesful run to test out vulnerability
    function testWithdrawSuccess() public {
        vm.startPrank(alice);
        timeLock.deposit{value: 1 ether}();
        timeLock.increaseLockTime(type(uint).max - timeLock.lockTime(alice) + 1);
        timeLock.withdraw();
        vm.stopPrank();

    }

    // Failure test when things work normally
    function testWithdrawError() public {
        vm.startPrank(alice);
        timeLock.deposit{value: 1 ether}();
        console.log(alice.balance);
        timeLock.withdraw();
        vm.stopPrank();
    }
}