├── .env.template ├── .github └── workflows │ └── test.yml ├── .gitignore ├── .gitmodules ├── README.md ├── foundry.toml ├── script ├── CounterExploitUUPSProxy.s.sol └── DeployCanonicalCounterExpoit.s.sol ├── src └── CounterExploit.sol └── test ├── CounterExploit.t.sol └── mock ├── ERC1967Proxy.mock.sol ├── ERC20.mock.sol └── Vault.mock.sol /.env.template: -------------------------------------------------------------------------------- 1 | PRIVATE_KEY= 2 | PROXY_ADDRESS= -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: test 2 | 3 | on: workflow_dispatch 4 | 5 | env: 6 | FOUNDRY_PROFILE: ci 7 | 8 | jobs: 9 | check: 10 | strategy: 11 | fail-fast: true 12 | 13 | name: Foundry project 14 | runs-on: ubuntu-latest 15 | steps: 16 | - uses: actions/checkout@v3 17 | with: 18 | submodules: recursive 19 | 20 | - name: Install Foundry 21 | uses: foundry-rs/foundry-toolchain@v1 22 | with: 23 | version: nightly 24 | 25 | - name: Run Forge build 26 | run: | 27 | forge --version 28 | forge build --sizes 29 | id: build 30 | 31 | - name: Run Forge tests 32 | run: | 33 | forge test -vvv 34 | id: test 35 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiler files 2 | cache/ 3 | out/ 4 | 5 | # Ignores development broadcast logs 6 | !/broadcast 7 | /broadcast/*/31337/ 8 | /broadcast/**/dry-run/ 9 | 10 | # Docs 11 | docs/ 12 | 13 | # Dotenv file 14 | .env 15 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "lib/forge-std"] 2 | path = lib/forge-std 3 | url = https://github.com/foundry-rs/forge-std 4 | branch = v1.4.0 5 | [submodule "lib/solmate"] 6 | path = lib/solmate 7 | url = https://github.com/transmissions11/solmate 8 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # The Counter Exploit Toolkit 2 | 3 | > NOTICE: FOR RESEARCH PURPOSES ONLY. 4 | 5 | ## Overiew 6 | 7 | This repository contains the contract `CounterExploit`, which contains all of the logic needed for 8 | an upgradeable proxy-based counter exploit. 9 | 10 | It provides arbitrary storage write access, arbitrary token withdrawal access, ether withdrawal 11 | access, an ether deposit honeypot, and, provided the attacker has sufficient balance and allowance 12 | for the proxy contract, it has the ability to take tokens directly from the attacker. 13 | 14 | Each function includes a batchable interface for executing multiple operations in the same 15 | transaction. 16 | 17 | It may be upgraded to, then upgraded from once the counter exploit is complete. 18 | 19 | ## Who Is This For? 20 | 21 | This contract is for two kinds of actors. First is the institutional actors that intend to follow 22 | the demands or requests of law enforcement to counter exploit the attacker. Second is for users of 23 | protocols with upgradeable proxies that may not be aware of the implications of such technology. 24 | 25 | ### For Institutional Actors 26 | 27 | This contract will have a canonical deployment that may be delecatecalled. All you need to do is 28 | send a transaction from the current proxy's authorized address that upgrades the current 29 | implementation address to the `CounterExploit` address and call `initialize()` 30 | _IN THE SAME TRANSACTION_. It is important to note that not atomically upgradingn and initializing 31 | will open a front-running opportunity that will give the attacker full control over your contract. 32 | This is the most crucial step, do not get this wrong. 33 | 34 | Once upgraded and initialized, you may step on your technology accelerationist gas pedal and see 35 | upradeable proxies through to their end; controlled by coercive, powerful actors for any purpose 36 | they see fit. 37 | 38 | See the API section below for more about specifics. 39 | 40 | ### For Users of Protocols with Upgradeable Proxies 41 | 42 | It is important to know exactly what you, a user, are getting into when interacting with protocols 43 | that use upgradeable proxies. Any contract that has an authorized address to upgrade the proxy has 44 | the ability to do this, regardless of whether the controller is a EOA, multisig, DAO, or an 45 | attacker through some other potential exploit in the protocol. 46 | 47 | Definitions: 48 | 49 | **Arbitrary storage write access**: Any storage slot within the contract may be overwritten. This 50 | includes, but is not limited to recordings of deposits, internal permissions, and authorized 51 | addresses. 52 | 53 | **Arbitrary token withdrawal access**: Any tokens that are in the contract may be withdrawn from 54 | the contract itself. This includes up to the full amount but may be more precise by targeting 55 | attackers in particular. Note that this will likely need to be used in parallel with storage writes. 56 | For example, if an attacker's funds need to be taken, the protocol must both take the tokens from 57 | the contract itself and overwrite the attacker's deposit slot to ensure they may not withdraw any 58 | other tokens from the protocol. 59 | 60 | **Ether withdrawal access**: Similar to arbitrary token withdrawal access, any ether that the 61 | contract holds may be withdrawn. 62 | 63 | **Steal token access**: This is likely the most important to understand. Protocols often require 64 | some kind of ERC20 token approval to act on the user's behalf, particularly with things like 65 | deposits and swaps. It is a common user experience pattern to ask the user to give "infinite" 66 | approval to the contract to prevent the user from having to run redundant transactions. However, 67 | combining this with upgradeable proxies means that not only can tokens be withdrawn from the proxy 68 | itself, but they may also be taken **from the user**, provided they have sufficient allowance 69 | granted to the contract. This may be useful to recover an attacker's stolen funds, but this is not 70 | enforceable to limit it to just the attacker. If you want to see and potentially revoke your token 71 | allowances, you may do so from here https://revoke.cash/ 72 | 73 | ## API 74 | 75 | The following methods are implemented on the `CounterExploit` contract. 76 | 77 | ### Write Access 78 | 79 | ```solidity 80 | /// @notice Writes a value to storage. 81 | /// @dev Reverts if caller is not admin. 82 | /// @param slot The storage slot to write to. 83 | /// @param value The value to write. 84 | function write(bytes32 slot, bytes32 value); 85 | 86 | /// @notice Batch writes values to storage. 87 | /// @dev Reverts if caller is not admin. 88 | /// @param slots The storage slots to write to. 89 | /// @param values The values to write. 90 | function writeBatch(bytes32[] calldata slots, bytes32[] calldata values); 91 | ``` 92 | 93 | ### Token Withdrawal Access 94 | 95 | ```solidity 96 | /// @notice Transfers tokens from this address to a receiver. 97 | /// @dev Reverts if caller is not admin. 98 | /// @param receiver The address to transfer tokens to. 99 | /// @param token The token to transfer. 100 | /// @param amount The amount of tokens to transfer. 101 | function takeToken(address receiver, ERC20 token, uint256 amount); 102 | 103 | /// @notice Batch transfers tokens from this address to a receiver. 104 | /// @dev Reverts if caller is not admin. 105 | /// @param receiver The address to transfer tokens to. 106 | /// @param tokens The tokens to transfer. 107 | /// @param amounts The amounts of tokens to transfer. 108 | function takeTokenBatch(address receiver, ERC20[] calldata tokens, uint256[] calldata amounts); 109 | ``` 110 | 111 | 112 | ### Ether Withdrawal Access 113 | 114 | ```solidity 115 | /// @notice Transfers ether to a receiver. 116 | /// @dev Reverts if caller is not admin. 117 | /// @param receiver The address to transfer ether to. 118 | /// @param amount The amount of ether to transfer. 119 | /// @return success True if the transfer succeeded (used to remove solc warning). 120 | function takeEther(address receiver, uint256 amount) 121 | ``` 122 | 123 | ### Steal Token Access 124 | 125 | ```solidity 126 | /// @notice Transfers tokens from the attacker to a receiver. 127 | /// @dev Reverts if caller is not admin. Gets the attacker's balance and allowance and transfers 128 | /// the lesser of the two. 129 | /// @param attacker The address to transfer tokens from. 130 | /// @param receiver The address to transfer tokens to. 131 | /// @param token The token to transfer. 132 | /// @param amount The amount of tokens to transfer. 133 | function stealToken( 134 | address attacker, 135 | address receiver, 136 | ERC20 token, 137 | uint256 amount 138 | ); 139 | 140 | /// @notice Batch transfers tokens from the attacker to a receiver. 141 | /// @dev Reverts if caller is not admin. 142 | /// @param receiver The address to transfer tokens to. 143 | /// @param tokens The tokens to transfer. 144 | /// @param amounts The amounts of tokens to transfer. 145 | function stealTokenBatch( 146 | address attacker, 147 | address receiver, 148 | ERC20[] calldata tokens, 149 | uint256[] calldata amounts 150 | ); 151 | ``` 152 | 153 | ### Ether Receive Honeypot 154 | 155 | ```solidity 156 | /// @notice Receives any amount of Ether without taking action. 157 | receive() external payable {} 158 | ``` 159 | -------------------------------------------------------------------------------- /foundry.toml: -------------------------------------------------------------------------------- 1 | [profile.default] 2 | src = 'src' 3 | out = 'out' 4 | libs = ['lib'] 5 | 6 | # See more config options https://github.com/foundry-rs/foundry/tree/master/config -------------------------------------------------------------------------------- /script/CounterExploitUUPSProxy.s.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity 0.8.18; 3 | 4 | import "lib/forge-std/src/Script.sol"; 5 | import "src/CounterExploit.sol"; 6 | 7 | interface UUPSProxy { 8 | function upgradeToAndCall(address imlementation, bytes calldata data) external; 9 | } 10 | 11 | contract CounterExploitERC1967Script is Script { 12 | address canonicalCounterExploit; 13 | 14 | function run() public { 15 | UUPSProxy proxy = UUPSProxy(vm.envAddress("PROXY_ADDRESS")); 16 | bytes memory initializer = abi.encodeCall(CounterExploit.init, ()); 17 | 18 | vm.startBroadcast(vm.envUint("PRIVATE_KEY")); 19 | proxy.upgradeToAndCall(canonicalCounterExploit, initializer); 20 | vm.stopBroadcast(); 21 | } 22 | } 23 | -------------------------------------------------------------------------------- /script/DeployCanonicalCounterExpoit.s.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity 0.8.18; 3 | 4 | import "lib/forge-std/src/Script.sol"; 5 | import "src/CounterExploit.sol"; 6 | 7 | contract DeployCanonicalCounterExploitScript is Script { 8 | function run() public { 9 | bytes32 slot = bytes32(uint256(keccak256("counter-exploit-toolkit.admin")) - 1); 10 | bytes32 admin = 0x0000000000000000000000000000000000000000000000000000000000000001; 11 | 12 | vm.startBroadcast(vm.envUint("PRIVATE_KEY")); 13 | CounterExploit ce = new CounterExploit(); 14 | ce.initialize(); 15 | ce.write(slot, admin); 16 | vm.stopBroadcast(); 17 | } 18 | } 19 | -------------------------------------------------------------------------------- /src/CounterExploit.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity 0.8.18; 3 | 4 | import "lib/solmate/src/tokens/ERC20.sol"; 5 | 6 | error Initialized(); 7 | error Unauthorized(); 8 | error LengthMismatch(); 9 | error InsufficientBalance(); 10 | error InsufficientAllowance(); 11 | 12 | /// @title Counter Exploit Contract. 13 | /// @author jtriley.eth 14 | /// @notice FOR RESEARCH PURPOSES ONLY! 15 | /// @dev This contract can be upgraded to by an upgradeable proxy, which enables arbitrary storage 16 | /// writes, token transfers, and ether transfers. 17 | contract CounterExploit { 18 | modifier onlyAdmin() { 19 | if (msg.sender != admin()) revert Unauthorized(); 20 | _; 21 | } 22 | 23 | /// @notice Returns the admin storage slot. 24 | /// @dev Computed as: `uint256(keccak256("counter-exploit-toolkit.admin")) - 1`. 25 | bytes32 public constant adminSlot = 26 | 0x1b681a0f4a91c141cc6ef34e55895f703d1ef56884d08c0fb5c6bc79a29fec42; 27 | 28 | /// @notice Initializer, callable once, sets the admin address. 29 | function initialize() public { 30 | if (admin() != address(0)) revert Initialized(); 31 | assembly { 32 | sstore(adminSlot, caller()) 33 | } 34 | } 35 | 36 | /// @notice Returns the admin address. 37 | function admin() public view returns (address _admin) { 38 | assembly { 39 | _admin := sload(adminSlot) 40 | } 41 | } 42 | 43 | /// @notice Writes a value to storage. 44 | /// @dev Reverts if caller is not admin. 45 | /// @param slot The storage slot to write to. 46 | /// @param value The value to write. 47 | function write(bytes32 slot, bytes32 value) public onlyAdmin { 48 | assembly { 49 | sstore(slot, value) 50 | } 51 | } 52 | 53 | /// @notice Batch writes values to storage. 54 | /// @dev Reverts if caller is not admin. 55 | /// @param slots The storage slots to write to. 56 | /// @param values The values to write. 57 | function writeBatch(bytes32[] calldata slots, bytes32[] calldata values) 58 | public 59 | onlyAdmin 60 | { 61 | unchecked { 62 | uint256 len = slots.length; 63 | if (len != values.length) revert LengthMismatch(); 64 | for (uint256 i; i < slots.length; ++i) { 65 | write(slots[i], values[i]); 66 | } 67 | } 68 | } 69 | 70 | /// @notice Transfers tokens from this address to a receiver. 71 | /// @dev Reverts if caller is not admin. 72 | /// @param receiver The address to transfer tokens to. 73 | /// @param token The token to transfer. 74 | /// @param amount The amount of tokens to transfer. 75 | function takeToken( 76 | address receiver, 77 | ERC20 token, 78 | uint256 amount 79 | ) public onlyAdmin { 80 | token.transfer(receiver, amount); 81 | } 82 | 83 | /// @notice Batch transfers tokens from this address to a receiver. 84 | /// @dev Reverts if caller is not admin. 85 | /// @param receiver The address to transfer tokens to. 86 | /// @param tokens The tokens to transfer. 87 | /// @param amounts The amounts of tokens to transfer. 88 | function takeTokenBatch( 89 | address receiver, 90 | ERC20[] calldata tokens, 91 | uint256[] calldata amounts 92 | ) public onlyAdmin { 93 | unchecked { 94 | uint256 len = tokens.length; 95 | if (len != amounts.length) revert LengthMismatch(); 96 | for (uint256 i; i < len; ++i) { 97 | takeToken(receiver, tokens[i], amounts[i]); 98 | } 99 | } 100 | } 101 | 102 | /// @notice Transfers ether to a receiver. 103 | /// @dev Reverts if caller is not admin. 104 | /// @param receiver The address to transfer ether to. 105 | /// @param amount The amount of ether to transfer. 106 | /// @return success True if the transfer succeeded (used to remove solc warning). 107 | function takeEther(address receiver, uint256 amount) public onlyAdmin returns (bool success) { 108 | (success, ) = payable(receiver).call{value: amount}(""); 109 | } 110 | 111 | /// @notice Transfers tokens from the attacker to a receiver. 112 | /// @dev Reverts if caller is not admin. Gets the attacker's balance and allowance and transfers 113 | /// the lesser of the two. 114 | /// @param attacker The address to trasnfer tokens from. 115 | /// @param receiver The address to transfer tokens to. 116 | /// @param token The token to transfer. 117 | /// @param amount The amount of tokens to transfer. 118 | function stealToken( 119 | address attacker, 120 | address receiver, 121 | ERC20 token, 122 | uint256 amount 123 | ) public onlyAdmin { 124 | if (amount > token.balanceOf(attacker)) revert InsufficientBalance(); 125 | if (amount > token.allowance(attacker, address(this))) revert InsufficientAllowance(); 126 | token.transferFrom(attacker, receiver, amount); 127 | } 128 | 129 | /// @notice Batch transfers tokens from the attacker to a receiver. 130 | /// @dev Reverts if caller is not admin. 131 | /// @param receiver The address to transfer tokens to. 132 | /// @param tokens The tokens to transfer. 133 | /// @param amounts The amounts of tokens to transfer. 134 | function stealTokenBatch( 135 | address attacker, 136 | address receiver, 137 | ERC20[] calldata tokens, 138 | uint256[] calldata amounts 139 | ) public onlyAdmin { 140 | unchecked { 141 | uint256 len = tokens.length; 142 | if (len != amounts.length) revert LengthMismatch(); 143 | for (uint256 i; i < len; ++i) { 144 | stealToken(attacker, receiver, tokens[i], amounts[i]); 145 | } 146 | } 147 | } 148 | 149 | /// @notice Receives any amount of Ether without taking action. 150 | receive() external payable {} 151 | } 152 | -------------------------------------------------------------------------------- /test/CounterExploit.t.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity 0.8.18; 3 | 4 | import "lib/forge-std/src/Test.sol"; 5 | import "src/CounterExploit.sol"; 6 | 7 | import "test/mock/ERC1967Proxy.mock.sol"; 8 | import "test/mock/Vault.mock.sol"; 9 | import "test/mock/ERC20.mock.sol"; 10 | 11 | contract CounterExploitTest is Test { 12 | VaultMock vault; 13 | ERC20Mock token; 14 | ERC1967ProxyMock proxy; 15 | address counterExploit; 16 | 17 | address admin = vm.addr(1); 18 | address attacker = vm.addr(2); 19 | address temporaryHolder = vm.addr(3); 20 | 21 | bytes32 oneBytes32 = 0x0000000000000000000000000000000000000000000000000000000000000001; 22 | bytes32 twoBytes32 = 0x0000000000000000000000000000000000000000000000000000000000000002; 23 | 24 | function setUp() public { 25 | vault = new VaultMock(); 26 | token = new ERC20Mock(); 27 | proxy = new ERC1967ProxyMock(address(vault)); 28 | counterExploit = address(new CounterExploit()); 29 | 30 | token.mint(attacker, 1 ether); 31 | 32 | vm.prank(attacker); 33 | token.approve(payable(address(proxy)), type(uint256).max); 34 | } 35 | 36 | // --- SUCCESS CASES --- 37 | 38 | function testCanUpgradeTo() public { 39 | proxy.upgradeTo(counterExploit); 40 | assertEq(proxy.implementation(), counterExploit); 41 | } 42 | 43 | function testCanUpgradeToAndCall() public { 44 | vm.prank(admin); 45 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 46 | assertEq(CounterExploit(payable(address(proxy))).admin(), admin); 47 | } 48 | 49 | function testCanWrite() public { 50 | vm.startPrank(admin); 51 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 52 | 53 | bytes32 slot = oneBytes32; 54 | bytes32 value1 = oneBytes32; 55 | bytes32 value2 = twoBytes32; 56 | 57 | CounterExploit(payable(address(proxy))).write(slot, value1); 58 | assertEq(vm.load(payable(address(proxy)), slot), value1); 59 | 60 | CounterExploit(payable(address(proxy))).write(slot, value2); 61 | assertEq(vm.load(payable(address(proxy)), slot), value2); 62 | 63 | vm.stopPrank(); 64 | } 65 | 66 | function testCanWriteBatch() public { 67 | vm.startPrank(admin); 68 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 69 | 70 | bytes32[] memory slots = new bytes32[](2); 71 | bytes32[] memory values = new bytes32[](2); 72 | slots[0] = oneBytes32; 73 | slots[1] = twoBytes32; 74 | values[0] = oneBytes32; 75 | values[1] = twoBytes32; 76 | 77 | CounterExploit(payable(address(proxy))).writeBatch(slots, values); 78 | 79 | vm.stopPrank(); 80 | 81 | assertEq(vm.load(payable(address(proxy)), slots[0]), values[0]); 82 | assertEq(vm.load(payable(address(proxy)), slots[1]), values[1]); 83 | } 84 | 85 | function canTakeToken() public { 86 | vm.prank(attacker); 87 | VaultMock(payable(address(proxy))).deposit(token, 1 ether); 88 | 89 | vm.startPrank(admin); 90 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 91 | 92 | CounterExploit(payable(address(proxy))).takeToken(temporaryHolder, token, 1 ether); 93 | vm.stopPrank(); 94 | 95 | assertEq(token.balanceOf(temporaryHolder), 1 ether); 96 | } 97 | 98 | function canTakeTokenBatch() public { 99 | vm.prank(attacker); 100 | VaultMock(payable(address(proxy))).deposit(token, 1 ether); 101 | 102 | vm.startPrank(admin); 103 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 104 | 105 | ERC20[] memory tokens = new ERC20[](2); 106 | uint256[] memory amounts = new uint256[](2); 107 | tokens[0] = token; 108 | tokens[1] = token; 109 | amounts[0] = 0.5 ether; 110 | amounts[1] = 0.5 ether; 111 | 112 | CounterExploit(payable(address(proxy))).takeTokenBatch(temporaryHolder, tokens, amounts); 113 | vm.stopPrank(); 114 | 115 | assertEq(token.balanceOf(temporaryHolder), 1 ether); 116 | } 117 | 118 | function canTakeEther() public { 119 | vm.deal(attacker, 1 ether); 120 | vm.prank(attacker); 121 | VaultMock(payable(address(proxy))).depositEther{value: 1 ether}(); 122 | 123 | vm.startPrank(admin); 124 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 125 | 126 | bool success = CounterExploit(payable(address(proxy))).takeEther(temporaryHolder, 1 ether); 127 | vm.stopPrank(); 128 | 129 | assertTrue(success); 130 | assertEq(temporaryHolder.balance, 1 ether); 131 | } 132 | 133 | function canStealToken() public { 134 | vm.startPrank(admin); 135 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 136 | 137 | CounterExploit(payable(address(proxy))).stealToken(attacker, temporaryHolder, token, 0.5 ether); 138 | vm.stopPrank(); 139 | 140 | assertEq(token.balanceOf(attacker), 0.5 ether); 141 | assertEq(token.balanceOf(temporaryHolder), 0.5 ether); 142 | } 143 | 144 | function canStealTokenBatch() public { 145 | vm.startPrank(admin); 146 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 147 | 148 | ERC20[] memory tokens = new ERC20[](2); 149 | uint256[] memory amounts = new uint256[](2); 150 | tokens[0] = token; 151 | tokens[1] = token; 152 | amounts[0] = 0.5 ether; 153 | amounts[1] = 0.5 ether; 154 | 155 | CounterExploit(payable(address(proxy))).stealTokenBatch(attacker, temporaryHolder, tokens, amounts); 156 | vm.stopPrank(); 157 | 158 | assertEq(token.balanceOf(attacker), 0.5 ether); 159 | assertEq(token.balanceOf(temporaryHolder), 0.5 ether); 160 | } 161 | 162 | function canReceive() public { 163 | vm.deal(attacker, 1 ether); 164 | 165 | vm.prank(admin); 166 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 167 | 168 | vm.prank(attacker); 169 | (bool success, ) = payable(address(proxy)).call{value: 1 ether}(""); 170 | 171 | assertTrue(success); 172 | assertEq(attacker.balance, 0); 173 | assertEq(payable(address(proxy)).balance, 1 ether); 174 | } 175 | 176 | // --- FAILURE CASES --- 177 | 178 | function testCantInitTwice() public { 179 | vm.prank(admin); 180 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 181 | 182 | vm.expectRevert(); 183 | CounterExploit(payable(address(proxy))).initialize(); 184 | } 185 | 186 | function testCantWriteWithoutAdmin() public { 187 | vm.prank(admin); 188 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 189 | 190 | vm.expectRevert(); 191 | vm.prank(attacker); 192 | CounterExploit(payable(address(proxy))).write(0x00, oneBytes32); 193 | } 194 | 195 | function testCantTakeTokenWithoutAdmin() public { 196 | vm.prank(admin); 197 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 198 | 199 | vm.expectRevert(); 200 | vm.prank(attacker); 201 | CounterExploit(payable(address(proxy))).takeToken(attacker, token, 1 ether); 202 | } 203 | 204 | function testCantTakeEtherWithoutAdmin() public { 205 | vm.prank(admin); 206 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 207 | 208 | vm.expectRevert(); 209 | vm.prank(attacker); 210 | CounterExploit(payable(address(proxy))).takeToken(attacker, token, 1 ether); 211 | } 212 | 213 | function testCantStealTokenWithoutAdmin() public { 214 | vm.prank(admin); 215 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 216 | 217 | vm.expectRevert(); 218 | vm.prank(attacker); 219 | CounterExploit(payable(address(proxy))).stealToken(admin, attacker, token, 1 ether); 220 | } 221 | 222 | function testCantStealTokenWithInsufficientBalance() public { 223 | vm.prank(attacker); 224 | token.transfer(address(0x01), 1 ether); 225 | 226 | vm.startPrank(admin); 227 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 228 | 229 | vm.expectRevert(); 230 | CounterExploit(payable(address(proxy))).stealToken(attacker, temporaryHolder, token, 1 ether); 231 | } 232 | 233 | function testCantStealTokenWithInsufficientAllowance() public { 234 | vm.prank(attacker); 235 | token.approve(payable(address(proxy)), 0); 236 | 237 | vm.startPrank(admin); 238 | proxy.upgradeToAndCall(counterExploit, abi.encodeCall(CounterExploit.initialize, ())); 239 | 240 | vm.expectRevert(); 241 | CounterExploit(payable(address(proxy))).stealToken(attacker, temporaryHolder, token, 1 ether); 242 | } 243 | } 244 | -------------------------------------------------------------------------------- /test/mock/ERC1967Proxy.mock.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity 0.8.18; 3 | 4 | error DelegatecallFailed(); 5 | 6 | contract ERC1967ProxyMock { 7 | bytes32 private constant IMPL_SLOT = 8 | 0x360894a13ba1a3210667c828492db98dca3e2076cc3733d6b56f0d58356e9bf3; 9 | 10 | constructor(address _implementation) { 11 | assembly { sstore(IMPL_SLOT, _implementation) } 12 | } 13 | 14 | function implementation() public view returns (address _implementation) { 15 | assembly { _implementation := sload(IMPL_SLOT) } 16 | } 17 | 18 | function upgradeTo(address _implementation) public { 19 | assembly { sstore(IMPL_SLOT, _implementation) } 20 | } 21 | 22 | function upgradeToAndCall(address _implementation, bytes calldata _data) public { 23 | upgradeTo(_implementation); 24 | (bool success, ) = _implementation.delegatecall(_data); 25 | if (!success) revert DelegatecallFailed(); 26 | } 27 | 28 | fallback() external { 29 | assembly { 30 | let _impl := sload(IMPL_SLOT) 31 | calldatacopy(0x00, 0x00, calldatasize()) 32 | let result := delegatecall(gas(), _impl, 0x00, calldatasize(), 0x00, 0x00) 33 | returndatacopy(0x00, 0x00, returndatasize()) 34 | switch result 35 | case 0x00 { revert(0x00, returndatasize()) } 36 | default { return(0x00, returndatasize()) } 37 | } 38 | } 39 | } 40 | -------------------------------------------------------------------------------- /test/mock/ERC20.mock.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity 0.8.18; 3 | 4 | import "lib/solmate/src/tokens/ERC20.sol"; 5 | 6 | contract ERC20Mock is ERC20("test token", "tt", 18) { 7 | function mint(address receiver, uint256 amount) public { 8 | _mint(receiver, amount); 9 | } 10 | } 11 | -------------------------------------------------------------------------------- /test/mock/Vault.mock.sol: -------------------------------------------------------------------------------- 1 | // SPDX-License-Identifier: MIT 2 | pragma solidity 0.8.18; 3 | 4 | import "lib/solmate/src/tokens/ERC20.sol"; 5 | 6 | contract VaultMock { 7 | mapping(address => uint256) public etherDeposit; 8 | mapping(address => mapping(ERC20 => uint256)) public tokenDeposit; 9 | 10 | function deposit(ERC20 token, uint256 amount) public { 11 | token.transferFrom(msg.sender, address(this), amount); 12 | tokenDeposit[msg.sender][token] += amount; 13 | } 14 | 15 | function depositEther() public payable { 16 | etherDeposit[msg.sender] += msg.value; 17 | } 18 | } 19 | --------------------------------------------------------------------------------