From e39c9047d12a8f25300a1fdf7cc02e0bc7b0547e Mon Sep 17 00:00:00 2001 From: Sam Glenn Date: Sun, 20 Jul 2025 21:39:27 -0500 Subject: [PATCH] chore(workshop): trap academy workshop 2 - response function demonstrations --- .../sudden-balance-drop-trap/bun.lockb | Bin 1688 -> 2030 bytes .../sudden-balance-drop-trap/drosera.toml | 8 +- .../sudden-balance-drop-trap/package.json | 3 +- .../sudden-balance-drop-trap/remappings.txt | 3 +- .../src/SuddenBalanceDropResponse.sol | 89 +++- .../src/SuddenBalanceDropTrap.sol | 18 +- .../test/SuddenBalanceDropTrap.t.sol | 500 +++++++++++++----- .../test/mocks/MockVault.sol | 114 ++++ 8 files changed, 565 insertions(+), 170 deletions(-) create mode 100644 defi-automation/sudden-balance-drop-trap/test/mocks/MockVault.sol diff --git a/defi-automation/sudden-balance-drop-trap/bun.lockb b/defi-automation/sudden-balance-drop-trap/bun.lockb index 3fc7c28e3741d77424f1f6c5c6a658389d96d058..76881625c8acb0cdee2f14ce0c0ad37e27765465 100755 GIT binary patch delta 590 zcmbQi`;LEtp61jY{dejs&5QmE$gzAs?bglK?tMF~{Vh**=YNl^1;J5q>qPEXb&?@B_%@LNdL~P|rfokb!}Tk%2)5 zC_Ztcj43lv4G4n(It>#;7iXS4ld;&5{s$$%9k z$-rO(5>gTaxf&!4^ejUGRB9jVWG@yqkj&(K7E={iG{AzJ2^Q829H3AJ0uDBy?IM#8 zvFLDFK!w8CC;wnkfN~~Fuud||C@Co@w$j%xN=?r!E-9+i%PYvuD%Q&{O4o6!uKz(B7gy=w9*)^`B68Idpm delta 374 zcmaFIKZAFIp5{)zvj@5!Br&r;DZdl@>De`H|6>#XZ)0|zS+D-Zs%Oe24mJibV4fH* zFI3J9;Xp_u28M=-6C@^X2$<}^$RRliD9R5M>IBl9KpLWmfnnj~L`HRmLqIMUP&5Q8 zTV|+dp=Zdza0bYi0U2~qW)dTdA`?^>j0SN*Y#5(u@=wO%$ui6?lQ%F$Oip0Vnf!n` ziJS2N)L~bcC)csq$%75~_a6daMli5|g&834u6_Kx_dOb^r^DfrLSpfustc zQcGASyRoWGE?_m`dIuGH$1?c@tAmyVD@Z*MKtqYkrqset!BC+%vnn+|O$W$I%*jm8 P%TME)Y{52Z68k#<$=^(O diff --git a/defi-automation/sudden-balance-drop-trap/drosera.toml b/defi-automation/sudden-balance-drop-trap/drosera.toml index 1aa4497..00f9a4a 100644 --- a/defi-automation/sudden-balance-drop-trap/drosera.toml +++ b/defi-automation/sudden-balance-drop-trap/drosera.toml @@ -7,11 +7,11 @@ drosera_address = "" [traps.sudden_balance_drop] path = "out/SuddenBalanceDropTrap.sol/SuddenBalanceDropTrap.json" response_contract = "0x0000000000000000000000000000000000000000" # placeholder - would be response contract address -response_function = "handleBalanceDrop(address,uint256,uint256)" # function to handle sudden balance drops +response_function = "handleBalanceDropWithEvent(address,uint256,uint256)" # function to handle sudden balance drops cooldown_period_blocks = 50 -min_number_of_operators = 2 +min_number_of_operators = 1 max_number_of_operators = 5 -block_sample_size = 1 -private_trap = false +block_sample_size = 2 +private_trap = true whitelist = [] address = "0x0000000000000000000000000000000000000000" # placeholder - would be deployed trap address \ No newline at end of file diff --git a/defi-automation/sudden-balance-drop-trap/package.json b/defi-automation/sudden-balance-drop-trap/package.json index a22d00d..6de9dab 100644 --- a/defi-automation/sudden-balance-drop-trap/package.json +++ b/defi-automation/sudden-balance-drop-trap/package.json @@ -5,6 +5,7 @@ "forge-std": "github:foundry-rs/forge-std#v1.8.1" }, "dependencies": { - "contracts": "https://github.com/drosera-network/contracts" + "contracts": "https://github.com/drosera-network/contracts", + "solmate": "^6.8.0" } } \ No newline at end of file diff --git a/defi-automation/sudden-balance-drop-trap/remappings.txt b/defi-automation/sudden-balance-drop-trap/remappings.txt index 50e7c4e..92dd484 100644 --- a/defi-automation/sudden-balance-drop-trap/remappings.txt +++ b/defi-automation/sudden-balance-drop-trap/remappings.txt @@ -1,2 +1,3 @@ forge-std/=node_modules/forge-std/src/ -drosera-contracts/=node_modules/contracts/src/ \ No newline at end of file +drosera-contracts/=node_modules/contracts/src/ +solmate/=node_modules/solmate/ \ No newline at end of file diff --git a/defi-automation/sudden-balance-drop-trap/src/SuddenBalanceDropResponse.sol b/defi-automation/sudden-balance-drop-trap/src/SuddenBalanceDropResponse.sol index 74d2f4e..0e3c05c 100644 --- a/defi-automation/sudden-balance-drop-trap/src/SuddenBalanceDropResponse.sol +++ b/defi-automation/sudden-balance-drop-trap/src/SuddenBalanceDropResponse.sol @@ -1,6 +1,12 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.19; +interface ISecuredVault { + function pause() external; + function activateCircuitBreaker() external; + function emergencyTransferToOwner() external; +} + /** * @title SuddenBalanceDropResponse * @notice Response contract for handling sudden balance drop alerts @@ -12,6 +18,9 @@ contract SuddenBalanceDropResponse { uint256 newBalance, uint256 timestamp ); + event VaultPaused(address indexed vault); + event CircuitBreakerActivated(address indexed vault); + event EmergencyTransferTriggered(address indexed vault); // State to track responses mapping(address => bool) public balanceDropHandled; @@ -29,20 +38,92 @@ contract SuddenBalanceDropResponse { } /** - * @notice Handle a sudden balance drop + * @notice Handle a sudden balance drop - Event only + * @param vault The vault address that experienced the drop + * @param oldBalance The previous balance + * @param newBalance The current balance + * @dev This only emits an event for monitoring purposes + */ + function handleBalanceDropWithEvent( + address vault, + uint256 oldBalance, + uint256 newBalance + ) external onlyTrapConfig { + // Mark the vault as handled + balanceDropHandled[vault] = true; + emit BalanceDropHandled(vault, oldBalance, newBalance, block.timestamp); + } + + /** + * @notice Handle a sudden balance drop - Pause vault + * @param vault The vault address that experienced the drop + * @param oldBalance The previous balance + * @param newBalance The current balance + * @dev This emits an event and attempts to pause the vault + */ + function handleBalanceDropWithPause( + address vault, + uint256 oldBalance, + uint256 newBalance + ) external onlyTrapConfig { + // Mark the vault as handled + balanceDropHandled[vault] = true; + emit BalanceDropHandled(vault, oldBalance, newBalance, block.timestamp); + + // Attempt to pause the vault if it supports pausing + try ISecuredVault(vault).pause() { + emit VaultPaused(vault); + } catch { + // Vault doesn't support pausing or pause failed + } + } + + /** + * @notice Activate circuit breaker for a vault (alternative to pausing) + * @param vault The vault address that experienced the drop + * @param oldBalance The previous balance + * @param newBalance The current balance + * @dev This provides a time-based freeze instead of indefinite pause + */ + function handleBalanceDropWithCircuitBreaker( + address vault, + uint256 oldBalance, + uint256 newBalance + ) external onlyTrapConfig { + // Mark the vault as handled + balanceDropHandled[vault] = true; + emit BalanceDropHandled(vault, oldBalance, newBalance, block.timestamp); + + // Activate the vault's circuit breaker + try ISecuredVault(vault).activateCircuitBreaker() { + emit CircuitBreakerActivated(vault); + } catch { + // Vault doesn't support circuit breaker or activation failed + } + } + + /** + * @notice Transfer all vault funds to owner (emergency protection) * @param vault The vault address that experienced the drop * @param oldBalance The previous balance * @param newBalance The current balance - * @dev This function matches the response_function signature in drosera.toml + * @dev This transfers all remaining funds to vault owner */ - function handleBalanceDrop( + function handleBalanceDropWithTransfer( address vault, uint256 oldBalance, uint256 newBalance ) external onlyTrapConfig { - // NOTE: This is a simplified example of a response but can be extended to real world use cases. + // Mark the vault as handled balanceDropHandled[vault] = true; emit BalanceDropHandled(vault, oldBalance, newBalance, block.timestamp); + + // Trigger emergency transfer to owner + try ISecuredVault(vault).emergencyTransferToOwner() { + emit EmergencyTransferTriggered(vault); + } catch { + // Vault doesn't support emergency transfer or transfer failed + } } // View function for testing diff --git a/defi-automation/sudden-balance-drop-trap/src/SuddenBalanceDropTrap.sol b/defi-automation/sudden-balance-drop-trap/src/SuddenBalanceDropTrap.sol index 4bdfe72..50b0180 100644 --- a/defi-automation/sudden-balance-drop-trap/src/SuddenBalanceDropTrap.sol +++ b/defi-automation/sudden-balance-drop-trap/src/SuddenBalanceDropTrap.sol @@ -43,22 +43,8 @@ contract SuddenBalanceDropTrap is ITrap { // Monitor real treasury/vault addresses with major stablecoins monitoredVaults.push( VaultInfo({ - vault: 0x742d35cC6634C0532925a3B8d80a6B24C5D06e41, // Example treasury - token: 0xa0B86a33e6441fD9Eec086d4E61ef0b5D31a5e7D // USDC - }) - ); - - monitoredVaults.push( - VaultInfo({ - vault: 0x8Ba1f109551Bd432803012645Aac136C40872A5F, // Example vault - token: 0xdAC17F958D2ee523a2206206994597C13D831ec7 // USDT - }) - ); - - monitoredVaults.push( - VaultInfo({ - vault: 0x742d35cC6634C0532925a3B8d80a6B24C5D06e41, // Example treasury - token: 0x6B175474E89094C44Da98b954EedeAC495271d0F // DAI + vault: 0x742d35cC6634C0532925a3B8d80a6B24C5D06e41, + token: 0xa0B86a33e6441fD9Eec086d4E61ef0b5D31a5e7D }) ); } diff --git a/defi-automation/sudden-balance-drop-trap/test/SuddenBalanceDropTrap.t.sol b/defi-automation/sudden-balance-drop-trap/test/SuddenBalanceDropTrap.t.sol index 3653a1d..ffe5436 100644 --- a/defi-automation/sudden-balance-drop-trap/test/SuddenBalanceDropTrap.t.sol +++ b/defi-automation/sudden-balance-drop-trap/test/SuddenBalanceDropTrap.t.sol @@ -2,185 +2,397 @@ pragma solidity ^0.8.19; import {Test} from "forge-std/Test.sol"; +import {ERC20} from "solmate/src/tokens/ERC20.sol"; import {SuddenBalanceDropTrap} from "../src/SuddenBalanceDropTrap.sol"; import {SuddenBalanceDropResponse} from "../src/SuddenBalanceDropResponse.sol"; +import {MockVault} from "./mocks/MockVault.sol"; + +// Simple mock USDC token for testing +contract MockUSDC is ERC20 { + constructor() ERC20("USD Coin", "USDC", 6) {} + + function mint(address to, uint256 amount) external { + _mint(to, amount); + } +} contract SuddenBalanceDropTrapTest is Test { SuddenBalanceDropTrap public trap; SuddenBalanceDropResponse public responseContract; - - address constant TEST_VAULT = 0x1111111111111111111111111111111111111111; - address constant TEST_TOKEN = 0x2222222222222222222222222222222222222222; - address constant MOCK_TRAP_CONFIG = 0x7E1b5cA35bd6BcAe8Ff33C0dDf79EffCFf0Ad19e; - + MockUSDC public mockUSDC; + MockVault public mockVault; + + address constant MOCK_TRAP_CONFIG = + 0x7E1b5cA35bd6BcAe8Ff33C0dDf79EffCFf0Ad19e; + function setUp() public { - trap = new SuddenBalanceDropTrap(); + // Deploy contracts first + mockUSDC = new MockUSDC(); responseContract = new SuddenBalanceDropResponse(MOCK_TRAP_CONFIG); + + // Deploy mock vault with the target USDC address (will be etched) + mockVault = new MockVault( + 0xa0B86a33e6441fD9Eec086d4E61ef0b5D31a5e7D, + address(responseContract) + ); + + // Etch the contracts to the hardcoded addresses in the trap + vm.etch( + 0xa0B86a33e6441fD9Eec086d4E61ef0b5D31a5e7D, + address(mockUSDC).code + ); + vm.etch( + 0x742d35cC6634C0532925a3B8d80a6B24C5D06e41, + address(mockVault).code + ); + + // Manually set the storage for the etched vault + // Slot 0: paused (bool) and guardian (address) are packed + // paused takes 1 byte at position 0, guardian takes 20 bytes after that + bytes32 packedSlot0 = bytes32( + uint256(uint160(address(responseContract))) << 8 + ); + vm.store( + 0x742d35cC6634C0532925a3B8d80a6B24C5D06e41, + bytes32(uint256(0)), + packedSlot0 + ); + + // Slot 1: owner address + vm.store( + 0x742d35cC6634C0532925a3B8d80a6B24C5D06e41, + bytes32(uint256(1)), + bytes32(uint256(uint160(address(this)))) // Set test contract as owner + ); + + // Now deploy the trap which will use the etched addresses + trap = new SuddenBalanceDropTrap(); } - - function test_InitialState() public view { - SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap.getMonitoredVaults(); - uint256 threshold = trap.getDropThreshold(); - - assertEq(vaults.length, 3); - assertEq(threshold, 1000); // 10% in BPS - } - - function test_CollectData() public { - // Mock balances for all monitored vaults - SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap.getMonitoredVaults(); - for (uint256 i = 0; i < vaults.length; i++) { - vm.mockCall( - vaults[i].token, - abi.encodeWithSignature("balanceOf(address)", vaults[i].vault), - abi.encode(100000e6) // 100k tokens - ); - } - - bytes memory data = trap.collect(); - assertTrue(data.length > 0); - - SuddenBalanceDropTrap.VaultBalance[] memory balances = - abi.decode(data, (SuddenBalanceDropTrap.VaultBalance[])); - assertEq(balances.length, 3); - } - - function test_NoBalanceDrops() public { - // Mock stable balances for all vaults - SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap.getMonitoredVaults(); - for (uint256 i = 0; i < vaults.length; i++) { - vm.mockCall( - vaults[i].token, - abi.encodeWithSignature("balanceOf(address)", vaults[i].vault), - abi.encode(100000e6) // 100k tokens (stable) - ); - } - - bytes memory data1 = trap.collect(); - bytes memory data2 = trap.collect(); - - bytes[] memory dataArray = new bytes[](2); - dataArray[0] = data2; - dataArray[1] = data1; - - (bool shouldTrigger,) = trap.shouldRespond(dataArray); - assertFalse(shouldTrigger); - } - - function test_BalanceDrop() public { - // Mock stable balances for original vaults first - SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap.getMonitoredVaults(); - for (uint256 i = 0; i < vaults.length; i++) { - vm.mockCall( - vaults[i].token, - abi.encodeWithSignature("balanceOf(address)", vaults[i].vault), - abi.encode(100000e6) // 100k tokens (stable) - ); - } - - // Add test vault and mock initial high balance - trap.addMonitoredVault(TEST_VAULT, TEST_TOKEN); - - vm.mockCall( - TEST_TOKEN, - abi.encodeWithSignature("balanceOf(address)", TEST_VAULT), - abi.encode(100000e6) // 100k tokens initially - ); - + + // FILE: drosera.toml + // [traps.sudden_balance_drop] + // path = "out/SuddenBalanceDropTrap.sol/SuddenBalanceDropTrap.json" + // response_contract = "0xresponsecontractaddress0000000000000000" + // response_function = "handleBalanceDropWithEvent(address,uint256,uint256)" + // cooldown_period_blocks = 50 + // min_number_of_operators = 1 + // max_number_of_operators = 5 + // block_sample_size = 2 + // + // forge test --contracts ./test/SuddenBalanceDropTrap.t.sol --match-test test_EventOnlyResponse -vvvv + function test_EventOnlyResponse() public { + // Get the monitored vault info + SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap + .getMonitoredVaults(); + address vaultAddr = vaults[0].vault; + address tokenAddr = vaults[0].token; + + // Mint initial balance + MockUSDC(tokenAddr).mint(vaultAddr, 100000e6); + + // Collect data before drop bytes memory data1 = trap.collect(); - - // Mock balance drop to 80k (20% drop) - vm.mockCall( - TEST_TOKEN, - abi.encodeWithSignature("balanceOf(address)", TEST_VAULT), - abi.encode(80000e6) // 80k tokens (20% drop) - ); - + + // Simulate balance drop + vm.roll(block.number + 1); + vm.prank(vaultAddr); + MockUSDC(tokenAddr).transfer(address(0xdead), 20000e6); + + // Collect data after drop bytes memory data2 = trap.collect(); - + bytes[] memory dataArray = new bytes[](2); dataArray[0] = data2; dataArray[1] = data1; - - (bool shouldTrigger, bytes memory responseData) = trap.shouldRespond(dataArray); - + + (bool shouldTrigger, bytes memory responseData) = trap.shouldRespond( + dataArray + ); assertTrue(shouldTrigger); - assertTrue(responseData.length > 0); - } - - function test_AddVault() public { - trap.addMonitoredVault(TEST_VAULT, TEST_TOKEN); - - SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap.getMonitoredVaults(); - assertEq(vaults.length, 4); - assertEq(vaults[3].vault, TEST_VAULT); - assertEq(vaults[3].token, TEST_TOKEN); + + // Expect the BalanceDropHandled event to be emitted + vm.expectEmit(true, false, false, true); + emit SuddenBalanceDropResponse.BalanceDropHandled( + vaultAddr, + 100000e6, + 80000e6, + block.timestamp + ); + + // Call the event-only response + vm.prank(MOCK_TRAP_CONFIG); + (bool success, ) = address(responseContract).call( + abi.encodePacked( + bytes4( + keccak256( + "handleBalanceDropWithEvent(address,uint256,uint256)" + ) + ), + responseData + ) + ); + assertTrue(success); + + // Verify balance drop was handled + assertTrue(responseContract.wasBalanceDropHandled(vaultAddr)); + + // Verify vault is NOT paused (event-only response) + assertFalse( + MockVault(vaultAddr).paused(), + "Vault should not be paused" + ); + + // Verify withdrawals still work normally + address user = address(0x1234); + vm.prank(user); + MockVault(vaultAddr).withdraw(10000e6); + assertEq(MockUSDC(tokenAddr).balanceOf(user), 10000e6); + assertEq(MockUSDC(tokenAddr).balanceOf(vaultAddr), 70000e6); } - // Test response contract integration using actual shouldRespond data - function test_ResponseContractWithTrapData() public { - // Mock stable balances for original vaults first - SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap.getMonitoredVaults(); - for (uint256 i = 0; i < vaults.length; i++) { - vm.mockCall( - vaults[i].token, - abi.encodeWithSignature("balanceOf(address)", vaults[i].vault), - abi.encode(100000e6) // 100k tokens (stable) - ); - } - - // Add test vault and mock initial high balance - trap.addMonitoredVault(TEST_VAULT, TEST_TOKEN); - - vm.mockCall( - TEST_TOKEN, - abi.encodeWithSignature("balanceOf(address)", TEST_VAULT), - abi.encode(100000e6) // 100k tokens initially - ); - + // FILE: drosera.toml + // [traps.sudden_balance_drop] + // path = "out/SuddenBalanceDropTrap.sol/SuddenBalanceDropTrap.json" + // response_contract = "0xresponsecontractaddress0000000000000000" + // response_function = "handleBalanceDropWithPause(address,uint256,uint256)" + // cooldown_period_blocks = 50 + // min_number_of_operators = 1 + // max_number_of_operators = 5 + // block_sample_size = 2 + // + // forge test --contracts ./test/SuddenBalanceDropTrap.t.sol --match-test test_PauseResponse -vvvv + function test_PauseResponse() public { + // Get the monitored vault info from the trap (uses etched addresses) + SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap + .getMonitoredVaults(); + address vaultAddr = vaults[0].vault; + address tokenAddr = vaults[0].token; + + // Mint initial balance of 100k USDC to the etched vault + MockUSDC(tokenAddr).mint(vaultAddr, 100000e6); + + // Get first data collection with 100k balance bytes memory data1 = trap.collect(); - - // Mock balance drop to 80k (20% drop) - vm.mockCall( - TEST_TOKEN, - abi.encodeWithSignature("balanceOf(address)", TEST_VAULT), - abi.encode(80000e6) // 80k tokens (20% drop) - ); - + + // Simulate balance drop - transfer 20k to reduce vault balance to 80k (20% drop) + vm.roll(block.number + 1); + vm.prank(vaultAddr); + MockUSDC(tokenAddr).transfer(address(0xdead), 20000e6); + bytes memory data2 = trap.collect(); - + bytes[] memory dataArray = new bytes[](2); dataArray[0] = data2; dataArray[1] = data1; - - (bool shouldTrigger, bytes memory responseData) = trap.shouldRespond(dataArray); + + (bool shouldTrigger, bytes memory responseData) = trap.shouldRespond( + dataArray + ); assertTrue(shouldTrigger, "Trap should trigger for balance drop"); - - // Verify we can decode the response data (just to check format) - (address vault, uint256 oldBalance, uint256 newBalance) = - abi.decode(responseData, (address, uint256, uint256)); - assertEq(vault, TEST_VAULT, "Vault should be test vault"); + + // Verify we can decode the response data + (address vault, uint256 oldBalance, uint256 newBalance) = abi.decode( + responseData, + (address, uint256, uint256) + ); + assertEq(vault, vaultAddr, "Vault should be the etched vault"); assertEq(oldBalance, 100000e6, "Old balance should be 100k"); assertEq(newBalance, 80000e6, "New balance should be 80k"); - + + // Verify vault is not paused before response + assertFalse( + MockVault(vaultAddr).paused(), + "Vault should not be paused initially" + ); + // This is how the Drosera operator would call the response contract: // It combines the function selector with the response data from shouldRespond vm.prank(MOCK_TRAP_CONFIG); - (bool success,) = address(responseContract).call( + (bool success, ) = address(responseContract).call( abi.encodePacked( - bytes4(keccak256("handleBalanceDrop(address,uint256,uint256)")), + bytes4( + keccak256( + "handleBalanceDropWithPause(address,uint256,uint256)" + ) + ), responseData ) ); assertTrue(success, "Response contract call should succeed"); // Verify balance drop was handled - assertTrue(responseContract.wasBalanceDropHandled(TEST_VAULT)); + assertTrue(responseContract.wasBalanceDropHandled(vaultAddr)); + + // Verify vault was paused by the response contract + assertTrue( + MockVault(vaultAddr).paused(), + "Vault should be paused after balance drop" + ); + + // Demonstrate that the paused vault blocks withdrawals + address user = address(0x1234); + vm.prank(user); + vm.expectRevert("Vault is paused"); + MockVault(vaultAddr).withdraw(10000e6); + + // Verify vault balance remains unchanged after failed withdrawal + assertEq(MockUSDC(tokenAddr).balanceOf(vaultAddr), 80000e6); } - function test_ResponseContractAccessControl() public { - // Test that only TrapConfig can call response functions - vm.expectRevert("Only TrapConfig can call this"); - responseContract.handleBalanceDrop(TEST_VAULT, 100000e6, 80000e6); + // FILE: drosera.toml + // [traps.sudden_balance_drop] + // path = "out/SuddenBalanceDropTrap.sol/SuddenBalanceDropTrap.json" + // response_contract = "0xresponsecontractaddress0000000000000000" + // response_function = "handleBalanceDropWithCircuitBreaker(address,uint256,uint256)" + // cooldown_period_blocks = 50 + // min_number_of_operators = 1 + // max_number_of_operators = 5 + // block_sample_size = 2 + // + // forge test --contracts ./test/SuddenBalanceDropTrap.t.sol --match-test test_CircuitBreakerResponse -vvvv + function test_CircuitBreakerResponse() public { + // Get the monitored vault info from the trap + SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap + .getMonitoredVaults(); + address vaultAddr = vaults[0].vault; + address tokenAddr = vaults[0].token; + + // Mint initial balance + MockUSDC(tokenAddr).mint(vaultAddr, 100000e6); + + // Get first data collection + bytes memory data1 = trap.collect(); + + // Simulate balance drop + vm.roll(block.number + 1); + vm.prank(vaultAddr); + MockUSDC(tokenAddr).transfer(address(0xdead), 20000e6); + + bytes memory data2 = trap.collect(); + + bytes[] memory dataArray = new bytes[](2); + dataArray[0] = data2; + dataArray[1] = data1; + + (bool shouldTrigger, bytes memory responseData) = trap.shouldRespond( + dataArray + ); + assertTrue(shouldTrigger); + + // Call the circuit breaker response instead of pause + vm.prank(MOCK_TRAP_CONFIG); + (bool success, ) = address(responseContract).call( + abi.encodePacked( + bytes4( + keccak256( + "handleBalanceDropWithCircuitBreaker(address,uint256,uint256)" + ) + ), + responseData + ) + ); + assertTrue(success); + + // Verify circuit breaker is active + assertTrue( + MockVault(vaultAddr).isCircuitBreakerActive(), + "Circuit breaker should be active" + ); + + // Try to withdraw - should fail due to circuit breaker + address user = address(0x5678); + vm.prank(user); + vm.expectRevert("Vault is frozen by circuit breaker"); + MockVault(vaultAddr).withdraw(10000e6); + + // Fast forward 32 blocks + vm.roll(block.number + 32); + + // Now withdrawal should work + vm.prank(user); + MockVault(vaultAddr).withdraw(10000e6); + assertEq(MockUSDC(tokenAddr).balanceOf(user), 10000e6); + assertEq(MockUSDC(tokenAddr).balanceOf(vaultAddr), 70000e6); + } + + // FILE: drosera.toml + // [traps.sudden_balance_drop] + // path = "out/SuddenBalanceDropTrap.sol/SuddenBalanceDropTrap.json" + // response_contract = "0xresponsecontractaddress0000000000000000" + // response_function = "handleBalanceDropWithTransfer(address,uint256,uint256)" + // cooldown_period_blocks = 50 + // min_number_of_operators = 1 + // max_number_of_operators = 5 + // block_sample_size = 2 + // + // forge test --contracts ./test/SuddenBalanceDropTrap.t.sol --match-test test_EmergencyTransferResponse -vvvv + function test_EmergencyTransferResponse() public { + // Get the monitored vault info + SuddenBalanceDropTrap.VaultInfo[] memory vaults = trap + .getMonitoredVaults(); + address vaultAddr = vaults[0].vault; + address tokenAddr = vaults[0].token; + + // Mint initial balance + MockUSDC(tokenAddr).mint(vaultAddr, 100000e6); + + // Get owner address (this test contract) + address vaultOwner = address(this); + + // Verify initial balances + assertEq(MockUSDC(tokenAddr).balanceOf(vaultAddr), 100000e6); + assertEq(MockUSDC(tokenAddr).balanceOf(vaultOwner), 0); + + // Collect data before drop + bytes memory data1 = trap.collect(); + + // Simulate balance drop + vm.roll(block.number + 1); + vm.prank(vaultAddr); + MockUSDC(tokenAddr).transfer(address(0xdead), 20000e6); + + // Collect data after drop + bytes memory data2 = trap.collect(); + + bytes[] memory dataArray = new bytes[](2); + dataArray[0] = data2; + dataArray[1] = data1; + + (bool shouldTrigger, bytes memory responseData) = trap.shouldRespond( + dataArray + ); + assertTrue(shouldTrigger); + + // Call the emergency transfer response + vm.prank(MOCK_TRAP_CONFIG); + (bool success, ) = address(responseContract).call( + abi.encodePacked( + bytes4( + keccak256( + "handleBalanceDropWithTransfer(address,uint256,uint256)" + ) + ), + responseData + ) + ); + assertTrue(success); + + // Verify all remaining funds were transferred to owner + assertEq( + MockUSDC(tokenAddr).balanceOf(vaultAddr), + 0, + "Vault should be empty" + ); + assertEq( + MockUSDC(tokenAddr).balanceOf(vaultOwner), + 80000e6, + "Owner should have all remaining funds" + ); + + // Verify the vault can no longer be used for withdrawals (no funds) + address user = address(0x9999); + vm.prank(user); + vm.expectRevert("Insufficient balance"); + MockVault(vaultAddr).withdraw(1000e6); } -} \ No newline at end of file +} diff --git a/defi-automation/sudden-balance-drop-trap/test/mocks/MockVault.sol b/defi-automation/sudden-balance-drop-trap/test/mocks/MockVault.sol new file mode 100644 index 0000000..d695285 --- /dev/null +++ b/defi-automation/sudden-balance-drop-trap/test/mocks/MockVault.sol @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.19; + +interface IERC20 { + function balanceOf(address account) external view returns (uint256); + function symbol() external view returns (string memory); + function transfer(address to, uint256 amount) external returns (bool); +} + +/** + * @title MockVault + * @notice Mock vault contract that holds USDC and can be paused or frozen + * @dev Used for testing the response contract's pause and circuit breaker functionality + */ +contract MockVault { + bool public paused; + address public guardian; + address public owner; + IERC20 public immutable usdcToken; + + // Circuit breaker state + uint256 public withdrawalFreezeUntil; + uint256 public constant CIRCUIT_BREAKER_BLOCKS = 32; + + event Paused(address indexed by); + event Unpaused(address indexed by); + event CircuitBreakerActivated(uint256 freezeUntilBlock); + event EmergencyTransfer(address indexed to, uint256 amount); + + modifier whenNotPaused() { + require(!paused, "Vault is paused"); + require(block.number >= withdrawalFreezeUntil, "Vault is frozen by circuit breaker"); + _; + } + + modifier onlyGuardian() { + require(msg.sender == guardian, "Only guardian can call"); + _; + } + + constructor(address _usdcToken, address _guardian) { + usdcToken = IERC20(_usdcToken); + guardian = _guardian; + owner = msg.sender; // Set deployer as owner + } + + /** + * @notice Pause the vault - callable by guardian (response contract) + */ + function pause() external onlyGuardian { + require(!paused, "Already paused"); + paused = true; + emit Paused(msg.sender); + } + + /** + * @notice Unpause the vault - callable by guardian + */ + function unpause() external onlyGuardian { + require(paused, "Not paused"); + paused = false; + emit Unpaused(msg.sender); + } + + /** + * @notice Activate circuit breaker to freeze withdrawals for 32 blocks + * @dev Alternative to pausing - provides automatic recovery after time period + */ + function activateCircuitBreaker() external onlyGuardian { + withdrawalFreezeUntil = block.number + CIRCUIT_BREAKER_BLOCKS; + emit CircuitBreakerActivated(withdrawalFreezeUntil); + } + + /** + * @notice Check if circuit breaker is currently active + */ + function isCircuitBreakerActive() external view returns (bool) { + return block.number < withdrawalFreezeUntil; + } + + /** + * @notice Withdraw USDC from the vault - blocked when paused + * @param amount The amount of USDC to withdraw + */ + function withdraw(uint256 amount) external whenNotPaused { + require(amount > 0, "Invalid amount"); + require(usdcToken.balanceOf(address(this)) >= amount, "Insufficient balance"); + + // Transfer tokens to the caller + bool success = usdcToken.transfer(msg.sender, amount); + require(success, "Transfer failed"); + } + + /** + * @notice Emergency transfer all funds to owner - callable by guardian + * @dev Used as an emergency protection mechanism + */ + function emergencyTransferToOwner() external onlyGuardian { + uint256 balance = usdcToken.balanceOf(address(this)); + require(balance > 0, "No balance to transfer"); + + bool success = usdcToken.transfer(owner, balance); + require(success, "Transfer failed"); + + emit EmergencyTransfer(owner, balance); + } + + /** + * @notice Get current USDC balance + */ + function getBalance() external view returns (uint256) { + return usdcToken.balanceOf(address(this)); + } +}