feat: add "assumeNoBlacklisted" stdCheat (#384)

foundry-rs · May 26, 2023 · 56db9b5 · 56db9b5
1 parent af5aeeb
commit 56db9b5
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 0 deletions.
diff --git a/src/StdCheats.sol b/src/StdCheats.sol
@@ -5,6 +5,7 @@ pragma experimental ABIEncoderV2;
 
 import {StdStorage, stdStorage} from "./StdStorage.sol";
 import {Vm} from "./Vm.sol";
+import {console2} from "./console2.sol";
 
 abstract contract StdCheatsSafe {
     Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
@@ -193,6 +194,27 @@ abstract contract StdCheatsSafe {
         uint256 key;
     }
 
+    // Checks that `addr` is not blacklisted by token contracts that have a blacklist.
+    function assumeNoBlacklisted(address token, address addr) internal virtual {
+        // Nothing to check if `token` is not a contract.
+        uint256 tokenCodeSize;
+        assembly {
+            tokenCodeSize := extcodesize(token)
+        }
+        require(tokenCodeSize > 0, "StdCheats assumeNoBlacklisted(address,address): Token address is not a contract.");
+
+        bool success;
+        bytes memory returnData;
+
+        // 4-byte selector for `isBlacklisted(address)`, used by USDC.
+        (success, returnData) = token.staticcall(abi.encodeWithSelector(0xfe575a87, addr));
+        vm.assume(!success || abi.decode(returnData, (bool)) == false);
+
+        // 4-byte selector for `isBlackListed(address)`, used by USDT.
+        (success, returnData) = token.staticcall(abi.encodeWithSelector(0xe47d6060, addr));
+        vm.assume(!success || abi.decode(returnData, (bool)) == false);
+    }
+
     function assumeNoPrecompiles(address addr) internal virtual {
         // Assembly required since `block.chainid` was introduced in 0.8.0.
         uint256 chainId;

diff --git a/test/StdCheats.t.sol b/test/StdCheats.t.sol
@@ -368,6 +368,60 @@ contract StdCheatsTest is Test {
     }
 }
 
+contract StdCheatsMock is StdCheats {
+    // We deploy a mock version so we can properly test expected reverts.
+    function assumeNoBlacklisted_(address token, address addr) external {
+        return assumeNoBlacklisted(token, addr);
+    }
+}
+
+contract StdCheatsForkTest is Test {
+    address internal constant SHIB = 0x95aD61b0a150d79219dCF64E1E6Cc01f0B64C4cE;
+    address internal constant USDC = 0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48;
+    address internal constant USDC_BLACKLISTED_USER = 0x1E34A77868E19A6647b1f2F47B51ed72dEDE95DD;
+    address internal constant USDT = 0xdAC17F958D2ee523a2206206994597C13D831ec7;
+    address internal constant USDT_BLACKLISTED_USER = 0x8f8a8F4B54a2aAC7799d7bc81368aC27b852822A;
+
+    // We deploy a mock version so we can properly test the revert.
+    StdCheatsMock private stdCheats = new StdCheatsMock();
+
+    function setUp() public {
+        // All tests of the `assumeNoBlacklisted` method are fork tests using live contracts.
+        vm.createSelectFork({urlOrAlias: "mainnet", blockNumber: 16_428_900});
+    }
+
+    function testCannotAssumeNoBlacklisted_EOA() external {
+        address eoa = vm.addr({privateKey: 1});
+        vm.expectRevert("StdCheats assumeNoBlacklisted(address,address): Token address is not a contract.");
+        assumeNoBlacklisted(eoa, address(0));
+    }
+
+    function testAssumeNoBlacklisted_TokenWithoutBlacklist(address addr) external {
+        assumeNoBlacklisted(SHIB, addr);
+        assertTrue(true);
+    }
+
+    function testAssumeNoBlacklisted_USDC() external {
+        vm.expectRevert();
+        stdCheats.assumeNoBlacklisted_(USDC, USDC_BLACKLISTED_USER);
+    }
+
+    function testAssumeNoBlacklisted_USDC(address addr) external {
+        assumeNoBlacklisted(USDC, addr);
+        assertFalse(USDCLike(USDC).isBlacklisted(addr));
+    }
+
+    function testAssumeNoBlacklisted_USDT() external {
+        vm.expectRevert();
+        stdCheats.assumeNoBlacklisted_(USDT, USDT_BLACKLISTED_USER);
+    }
+
+    function testAssumeNoBlacklisted_USDT(address addr) external {
+        assumeNoBlacklisted(USDT, addr);
+        assertFalse(USDTLike(USDT).isBlackListed(addr));
+    }
+}
+
 contract Bar {
     constructor() payable {
         /// `DEAL` STDCHEAT
@@ -438,6 +492,14 @@ contract BarERC721 {
     mapping(address => uint256) private _balances;
 }
 
+interface USDCLike {
+    function isBlacklisted(address) external view returns (bool);
+}
+
+interface USDTLike {
+    function isBlackListed(address) external view returns (bool);
+}
+
 contract RevertingContract {
     constructor() {
         revert();