Skip to content

Commit

Permalink
fix: example in tests
Browse files Browse the repository at this point in the history
  • Loading branch information
prestwich committed Oct 4, 2023
1 parent 81cd99c commit 8658d63
Showing 1 changed file with 42 additions and 0 deletions.
42 changes: 42 additions & 0 deletions src/test/ExcessivelySafeCall.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ contract ContractTest is DSTest {
using ExcessivelySafeCall for address;

address target;
Intermediary intermediary;
CallTarget t;

function returnSize() internal pure returns (uint256 _bytes) {
Expand All @@ -19,6 +20,7 @@ contract ContractTest is DSTest {
function setUp() public {
t = new CallTarget();
target = address(t);
intermediary = new Intermediary();
}

function testCall() public {
Expand Down Expand Up @@ -177,9 +179,38 @@ contract ContractTest is DSTest {
assertEq(_ret.length, 32, "revert didn't truncate");
}

function test_drain() public {
intermediary.getDrained{gas: 1_000_000}(target, 10_000);
}

function test_drain_safe() public {
intermediary.getSafeDrained{gas: 1_000_000}(target, 10_000);
}
}


contract Intermediary {
using ExcessivelySafeCall for address;
function getDrained(address target, uint256 drainTo) public {
bool _success;
bytes memory _ret;
(_success, _ret) = target.call(
abi.encodeWithSelector(CallTarget.drainTo.selector, drainTo)
);
}

function getSafeDrained(address target, uint256 drainTo) public {
bool _success;
bytes memory _ret;
(_success, _ret) = target.excessivelySafeCall(
gasleft(),
0,
0,
abi.encodeWithSelector(CallTarget.drainTo.selector, drainTo)
);
}
}

contract CallTarget {
uint256 public called;
constructor () {}
Expand Down Expand Up @@ -219,4 +250,15 @@ contract CallTarget {
function badRev() external pure {
revBytes(1_000_000);
}

function drainTo(uint256 gas) public view {
uint256 len;
while (gasleft() > gas) {
bytes memory a = new bytes(0);
len += 32 + a.length;
}
assembly {
return (0, len)
}
}
}

0 comments on commit 8658d63

Please sign in to comment.