src/test/basetest.sol

1// SPDX-License-Identifier: UNLICENSED
2pragma solidity ^0.8.15;
3 
4import "./tokenhelper.sol";
5import "forge-std/Test.sol";
6 
7// Base contract for DeFi exploit POCs. Inherit from this instead of Test directly.
8// Provides before/after balance logging via the `balanceLog` modifier.
9//
10// Single-asset mode (default):
11//   Set `fundingToken` to the token you profit in (address(0) = native ETH/chain coin).
12//   The `balanceLog` modifier logs that one token before and after testExploit().
13//
14// Multi-asset mode:
15//   Set `multiAssetLog = true` and populate `fundingTokens` with every token you want tracked.
16//   The same `balanceLog` modifier will log all of them. No override needed.
17//   Optionally set `attacker` to log a different address (e.g. a separate profit contract).
18//   If `attacker` is left as address(0), it resolves to address(this).
19contract BaseTestWithBalanceLog is Test {
20    // Single-asset mode: the token to log profit in. address(0) = native coin.
21    address fundingToken = address(0);
22    // Multi-asset mode: full list of tokens to track (ERC-20 or address(0) for native).
23    address[] fundingTokens;
24    // Set to true to enable multi-asset logging via fundingTokens[].
25    bool multiAssetLog = false;
26    // Address whose balances are logged. Defaults to address(this) when left as address(0).
27    address attacker;
28 
29    struct ChainInfo {
30        string name;
31        string symbol;
32    }
33 
34    mapping(uint256 => ChainInfo) private chainIdToInfo;
35 
36    constructor() {
37        chainIdToInfo[1] = ChainInfo("MAINNET", "ETH");
38        chainIdToInfo[238] = ChainInfo("BLAST", "ETH");
39        chainIdToInfo[10] = ChainInfo("OPTIMISM", "ETH");
40        chainIdToInfo[250] = ChainInfo("FANTOM", "FTM");
41        chainIdToInfo[42_161] = ChainInfo("ARBITRUM", "ETH");
42        chainIdToInfo[56] = ChainInfo("BSC", "BNB");
43        chainIdToInfo[1285] = ChainInfo("MOONRIVER", "MOVR");
44        chainIdToInfo[100] = ChainInfo("GNOSIS", "XDAI");
45        chainIdToInfo[43_114] = ChainInfo("AVALANCHE", "AVAX");
46        chainIdToInfo[137] = ChainInfo("POLYGON", "MATIC");
47        chainIdToInfo[42_220] = ChainInfo("CELO", "CELO");
48        chainIdToInfo[8453] = ChainInfo("BASE", "ETH");
49        chainIdToInfo[1329] = ChainInfo("SEI", "SEI");
50    }
51 
52    function getChainSymbol(
53        uint256 chainId
54    ) internal view returns (string memory symbol) {
55        symbol = chainIdToInfo[chainId].symbol;
56        if (bytes(symbol).length == 0) symbol = "ETH";
57    }
58 
59    function _getTokenData(
60        address token,
61        address account
62    ) internal returns (string memory symbol, uint256 balance, uint8 decimals) {
63        if (token == address(0)) {
64            symbol = getChainSymbol(block.chainid);
65            balance = account.balance;
66            decimals = 18;
67        } else {
68            symbol = TokenHelper.getTokenSymbol(token);
69            balance = TokenHelper.getTokenBalance(token, account);
70            decimals = TokenHelper.getTokenDecimals(token);
71        }
72    }
73 
74    function _logTokenBalance(
75        address token,
76        address account,
77        string memory label
78    ) private {
79        (string memory symbol, uint256 balance, uint8 decimals) = _getTokenData(token, account);
80        emit log_named_decimal_uint(string(abi.encodePacked(label, " ", symbol, " Balance")), balance, decimals);
81    }
82 
83    function _attacker() private view returns (address) {
84        return attacker == address(0) ? address(this) : attacker;
85    }
86 
87    modifier balanceLog() virtual {
88        if (multiAssetLog) {
89            if (fundingToken == address(0)) vm.deal(_attacker(), 0);
90            _logMultiAssetBalances("Before exploit");
91        } else {
92            if (fundingToken == address(0)) vm.deal(_attacker(), 0);
93            _logTokenBalance(fundingToken, _attacker(), "Attacker Before exploit");
94        }
95        _;
96        if (multiAssetLog) {
97            _logMultiAssetBalances("After exploit");
98        } else {
99            _logTokenBalance(fundingToken, _attacker(), "Attacker After exploit");
100        }
101    }
102 
103    modifier balanceLog2(address target) virtual {
104        if (fundingToken == address(0)) vm.deal(target, 0);
105        _logTokenBalance(fundingToken, target, string(abi.encodePacked("Attacker Before exploit")));
106        _;
107        _logTokenBalance(fundingToken, target, string(abi.encodePacked("Attacker After exploit")));
108    }
109 
110    function logTokenBalance(
111        address token,
112        address account,
113        string memory label
114    ) internal {
115        _logTokenBalance(token, account, label);
116    }
117 
118    function logMultipleTokenBalances(
119        address[] memory tokens,
120        address account,
121        string memory label
122    ) internal {
123        emit log_string(string(abi.encodePacked("=== ", label, " ===")));
124        for (uint256 i = 0; i < tokens.length; i++) {
125            _logTokenBalance(tokens[i], account, "");
126        }
127    }
128 
129    function _addFundingToken(address token) internal {
130        fundingTokens.push(token);
131    }
132 
133    function _addFundingTokens(address[] memory tokens) internal {
134        for (uint256 i = 0; i < tokens.length; i++) {
135            fundingTokens.push(tokens[i]);
136        }
137    }
138 
139    function _logMultiAssetBalances(string memory label) internal {
140        logMultipleTokenBalances(fundingTokens, _attacker(), label);
141    }
142}