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

// Save as test/AccessControlExample.t.sol in a Foundry project.
// Run: forge test --match-contract AccessControlExampleTest -vv
interface Vm {
    function deal(address account, uint256 balance) external;
    function prank(address sender) external;
    function expectRevert(bytes4 selector) external;
}

contract VulnerableTreasury {
    constructor() payable {}

    function withdraw(address payable recipient, uint256 amount) external {
        (bool ok,) = recipient.call{value: amount}("");
        require(ok, "transfer failed");
    }
}

contract FixedTreasury {
    error Unauthorized();
    address public immutable owner;

    constructor() payable { owner = msg.sender; }

    function withdraw(address payable recipient, uint256 amount) external {
        if (msg.sender != owner) revert Unauthorized();
        (bool ok,) = recipient.call{value: amount}("");
        require(ok, "transfer failed");
    }
}

contract AccessControlExampleTest {
    Vm constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
    address payable constant attacker = payable(address(0xBEEF));

    function testDemonstratesUnauthorizedWithdrawal() public {
        vm.deal(address(this), 1 ether);
        vm.deal(attacker, 0);
        VulnerableTreasury treasury = new VulnerableTreasury{value: 1 ether}();
        vm.prank(attacker);
        treasury.withdraw(attacker, 1 ether);
        require(attacker.balance == 1 ether, "unauthorized withdrawal reproduced");
        require(address(treasury).balance == 0, "treasury drained");
    }

    function testFixedBlocksUnauthorizedWithdrawal() public {
        vm.deal(address(this), 1 ether);
        FixedTreasury treasury = new FixedTreasury{value: 1 ether}();
        vm.expectRevert(FixedTreasury.Unauthorized.selector);
        vm.prank(attacker);
        treasury.withdraw(attacker, 1 ether);
        require(address(treasury).balance == 1 ether, "treasury retains funds");
    }

    function testFixedAllowsOwnerWithdrawal() public {
        vm.deal(address(this), 1 ether);
        vm.deal(attacker, 0);
        FixedTreasury treasury = new FixedTreasury{value: 1 ether}();
        treasury.withdraw(attacker, 1 ether);
        require(attacker.balance == 1 ether, "authorized withdrawal works");
    }
}