Skip to content

Commit

Permalink
♻️ Very safe transfer (#1128)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Vectorized
Vectorized authored Oct 17, 2024
1 parent 7f72ce1 commit e9c03bf
Show file tree
Hide file tree
Showing 2 changed files with 108 additions and 89 deletions.
118 changes: 50 additions & 68 deletions src/utils/SafeTransferLib.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,6 @@ pragma solidity ^0.8.4;
///
/// @dev Note:
/// - For ETH transfers, please use `forceSafeTransferETH` for DoS protection.
/// - For ERC20s, this implementation won't check that a token has code,
/// responsibility is delegated to the caller.
library SafeTransferLib {
/*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
/* CUSTOM ERRORS */
Expand Down Expand Up @@ -202,15 +200,12 @@ library SafeTransferLib {
mstore(0x40, to) // Store the `to` argument.
mstore(0x2c, shl(96, from)) // Store the `from` argument.
mstore(0x0c, 0x23b872dd000000000000000000000000) // `transferFrom(address,address,uint256)`.
// Perform the transfer, reverting upon failure.
if iszero(
and( // The arguments of `and` are evaluated from right to left.
or(eq(mload(0x00), 1), iszero(returndatasize())), // Returned 1 or nothing.
call(gas(), token, 0, 0x1c, 0x64, 0x00, 0x20)
)
) {
mstore(0x00, 0x7939f424) // `TransferFromFailed()`.
revert(0x1c, 0x04)
let success := call(gas(), token, 0, 0x1c, 0x64, 0x00, 0x20)
if iszero(and(eq(mload(0x00), 1), success)) {
if iszero(lt(or(iszero(extcodesize(token)), returndatasize()), success)) {
mstore(0x00, 0x7939f424) // `TransferFromFailed()`.
revert(0x1c, 0x04)
}
}
mstore(0x60, 0) // Restore the zero slot to zero.
mstore(0x40, m) // Restore the free memory pointer.
Expand All @@ -231,11 +226,10 @@ library SafeTransferLib {
mstore(0x40, to) // Store the `to` argument.
mstore(0x2c, shl(96, from)) // Store the `from` argument.
mstore(0x0c, 0x23b872dd000000000000000000000000) // `transferFrom(address,address,uint256)`.
success :=
and( // The arguments of `and` are evaluated from right to left.
or(eq(mload(0x00), 1), iszero(returndatasize())), // Returned 1 or nothing.
call(gas(), token, 0, 0x1c, 0x64, 0x00, 0x20)
)
success := call(gas(), token, 0, 0x1c, 0x64, 0x00, 0x20)
if iszero(and(eq(mload(0x00), 1), success)) {
success := lt(or(iszero(extcodesize(token)), returndatasize()), success)
}
mstore(0x60, 0) // Restore the zero slot to zero.
mstore(0x40, m) // Restore the free memory pointer.
}
Expand Down Expand Up @@ -268,14 +262,12 @@ library SafeTransferLib {
mstore(0x00, 0x23b872dd) // `transferFrom(address,address,uint256)`.
amount := mload(0x60) // The `amount` is already at 0x60. We'll need to return it.
// Perform the transfer, reverting upon failure.
if iszero(
and( // The arguments of `and` are evaluated from right to left.
or(eq(mload(0x00), 1), iszero(returndatasize())), // Returned 1 or nothing.
call(gas(), token, 0, 0x1c, 0x64, 0x00, 0x20)
)
) {
mstore(0x00, 0x7939f424) // `TransferFromFailed()`.
revert(0x1c, 0x04)
let success := call(gas(), token, 0, 0x1c, 0x64, 0x00, 0x20)
if iszero(and(eq(mload(0x00), 1), success)) {
if iszero(lt(or(iszero(extcodesize(token)), returndatasize()), success)) {
mstore(0x00, 0x7939f424) // `TransferFromFailed()`.
revert(0x1c, 0x04)
}
}
mstore(0x60, 0) // Restore the zero slot to zero.
mstore(0x40, m) // Restore the free memory pointer.
Expand All @@ -291,14 +283,12 @@ library SafeTransferLib {
mstore(0x34, amount) // Store the `amount` argument.
mstore(0x00, 0xa9059cbb000000000000000000000000) // `transfer(address,uint256)`.
// Perform the transfer, reverting upon failure.
if iszero(
and( // The arguments of `and` are evaluated from right to left.
or(eq(mload(0x00), 1), iszero(returndatasize())), // Returned 1 or nothing.
call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
)
) {
mstore(0x00, 0x90b8ec18) // `TransferFailed()`.
revert(0x1c, 0x04)
let success := call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
if iszero(and(eq(mload(0x00), 1), success)) {
if iszero(lt(or(iszero(extcodesize(token)), returndatasize()), success)) {
mstore(0x00, 0x90b8ec18) // `TransferFailed()`.
revert(0x1c, 0x04)
}
}
mstore(0x34, 0) // Restore the part of the free memory pointer that was overwritten.
}
Expand All @@ -325,14 +315,12 @@ library SafeTransferLib {
amount := mload(0x34) // The `amount` is already at 0x34. We'll need to return it.
mstore(0x00, 0xa9059cbb000000000000000000000000) // `transfer(address,uint256)`.
// Perform the transfer, reverting upon failure.
if iszero(
and( // The arguments of `and` are evaluated from right to left.
or(eq(mload(0x00), 1), iszero(returndatasize())), // Returned 1 or nothing.
call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
)
) {
mstore(0x00, 0x90b8ec18) // `TransferFailed()`.
revert(0x1c, 0x04)
let success := call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
if iszero(and(eq(mload(0x00), 1), success)) {
if iszero(lt(or(iszero(extcodesize(token)), returndatasize()), success)) {
mstore(0x00, 0x90b8ec18) // `TransferFailed()`.
revert(0x1c, 0x04)
}
}
mstore(0x34, 0) // Restore the part of the free memory pointer that was overwritten.
}
Expand All @@ -346,15 +334,12 @@ library SafeTransferLib {
mstore(0x14, to) // Store the `to` argument.
mstore(0x34, amount) // Store the `amount` argument.
mstore(0x00, 0x095ea7b3000000000000000000000000) // `approve(address,uint256)`.
// Perform the approval, reverting upon failure.
if iszero(
and( // The arguments of `and` are evaluated from right to left.
or(eq(mload(0x00), 1), iszero(returndatasize())), // Returned 1 or nothing.
call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
)
) {
mstore(0x00, 0x3e3f8f73) // `ApproveFailed()`.
revert(0x1c, 0x04)
let success := call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
if iszero(and(eq(mload(0x00), 1), success)) {
if iszero(lt(or(iszero(extcodesize(token)), returndatasize()), success)) {
mstore(0x00, 0x3e3f8f73) // `ApproveFailed()`.
revert(0x1c, 0x04)
}
}
mstore(0x34, 0) // Restore the part of the free memory pointer that was overwritten.
}
Expand All @@ -371,25 +356,22 @@ library SafeTransferLib {
mstore(0x34, amount) // Store the `amount` argument.
mstore(0x00, 0x095ea7b3000000000000000000000000) // `approve(address,uint256)`.
// Perform the approval, retrying upon failure.
if iszero(
and( // The arguments of `and` are evaluated from right to left.
or(eq(mload(0x00), 1), iszero(returndatasize())), // Returned 1 or nothing.
call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
)
) {
mstore(0x34, 0) // Store 0 for the `amount`.
mstore(0x00, 0x095ea7b3000000000000000000000000) // `approve(address,uint256)`.
pop(call(gas(), token, 0, 0x10, 0x44, codesize(), 0x00)) // Reset the approval.
mstore(0x34, amount) // Store back the original `amount`.
// Retry the approval, reverting upon failure.
if iszero(
and(
or(eq(mload(0x00), 1), iszero(returndatasize())), // Returned 1 or nothing.
call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
)
) {
mstore(0x00, 0x3e3f8f73) // `ApproveFailed()`.
revert(0x1c, 0x04)
let success := call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
if iszero(and(eq(mload(0x00), 1), success)) {
if iszero(lt(or(iszero(extcodesize(token)), returndatasize()), success)) {
mstore(0x34, 0) // Store 0 for the `amount`.
mstore(0x00, 0x095ea7b3000000000000000000000000) // `approve(address,uint256)`.
pop(call(gas(), token, 0, 0x10, 0x44, codesize(), 0x00)) // Reset the approval.
mstore(0x34, amount) // Store back the original `amount`.
// Retry the approval, reverting upon failure.
success := call(gas(), token, 0, 0x10, 0x44, 0x00, 0x20)
if iszero(and(eq(mload(0x00), 1), success)) {
// Check the `extcodesize` again just in case the token selfdestructs lol.
if iszero(lt(or(iszero(extcodesize(token)), returndatasize()), success)) {
mstore(0x00, 0x3e3f8f73) // `ApproveFailed()`.
revert(0x1c, 0x04)
}
}
}
}
mstore(0x34, 0) // Restore the part of the free memory pointer that was overwritten.
Expand Down
79 changes: 58 additions & 21 deletions test/SafeTransferLib.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,16 @@ contract SafeTransferLibTest is SoladyTest {
vm.etch(_PERMIT2, bytecode);
}

function testSuccessTrick(bool success, uint256 extCodeSize, uint256 returnDataSize) public {
bool expected = success && extCodeSize != 0 && returnDataSize == 0;
bool computed;
/// @solidity memory-safe-assembly
assembly {
computed := lt(or(iszero(extCodeSize), returnDataSize), success)
}
assertEq(computed, expected);
}

function testTransferWithMissingReturn() public {
verifySafeTransfer(address(missingReturn), address(0xBEEF), 1e18, _SUCCESS);
}
Expand All @@ -120,8 +130,9 @@ contract SafeTransferLibTest is SoladyTest {
verifySafeTransfer(address(returnsTooMuch), address(0xBEEF), 1e18, _SUCCESS);
}

function testTransferWithNonContract() public {
SafeTransferLib.safeTransfer(address(0xBADBEEF), address(0xBEEF), 1e18);
function testTransferWithNonContractReverts() public {
vm.expectRevert(SafeTransferLib.TransferFailed.selector);
this.safeTransfer(address(0xBADBEEF), address(0xBEEF), 1e18);
}

function testTransferFromWithMissingReturn() public {
Expand All @@ -140,8 +151,15 @@ contract SafeTransferLibTest is SoladyTest {
);
}

function testTransferFromWithNonContract() public {
SafeTransferLib.safeTransferFrom(address(0xBADBEEF), address(0xFEED), address(0xBEEF), 1e18);
function testTransferFromWithNonContractReverts() public {
vm.expectRevert(SafeTransferLib.TransferFromFailed.selector);
this.safeTransferFrom(address(0xBADBEEF), address(0xFEED), address(0xBEEF), 1e18);
}

function safeTransferFrom(address token, address from, address to, uint256 amount) public {
SafeTransferLib.safeTransferFrom(
_brutalized(token), _brutalized(from), _brutalized(to), amount
);
}

function testApproveWithMissingReturn() public {
Expand All @@ -156,12 +174,22 @@ contract SafeTransferLibTest is SoladyTest {
verifySafeApprove(address(returnsTooMuch), address(0xBEEF), 1e18, _SUCCESS);
}

function testApproveWithNonContract() public {
SafeTransferLib.safeApprove(address(0xBADBEEF), address(0xBEEF), 1e18);
function testApproveWithNonContractReverts() public {
vm.expectRevert(SafeTransferLib.ApproveFailed.selector);
this.safeApprove(address(0xBADBEEF), address(0xBEEF), 1e18);
}

function safeApprove(address token, address to, uint256 amount) public {
SafeTransferLib.safeApprove(token, to, amount);
}

function testApproveWithRetryWithNonContractReverts() public {
vm.expectRevert(SafeTransferLib.ApproveFailed.selector);
this.safeApproveWithRetry(address(0xBADBEEF), address(0xBEEF), 1e18);
}

function testApproveWithRetryWithNonContract() public {
SafeTransferLib.safeApproveWithRetry(address(0xBADBEEF), address(0xBEEF), 1e18);
function safeApproveWithRetry(address token, address to, uint256 amount) public {
SafeTransferLib.safeApproveWithRetry(token, to, amount);
}

function testTransferETH() public {
Expand Down Expand Up @@ -441,8 +469,13 @@ contract SafeTransferLibTest is SoladyTest {
verifySafeTransfer(address(returnsRawBytes), to, amount, _SUCCESS);
}

function testTransferWithNonContract(bytes32, address to, uint256 amount) public {
SafeTransferLib.safeTransfer(_brutalized(_randomHashedAddress()), _brutalized(to), amount);
function testTransferWithNonContractReverts(bytes32, address to, uint256 amount) public {
vm.expectRevert(SafeTransferLib.TransferFailed.selector);
this.safeTransfer(_randomHashedAddress(), to, amount);
}

function safeTransfer(address token, address to, uint256 amount) public {
SafeTransferLib.safeTransfer(token, to, amount);
}

function testTransferETHToContractWithoutFallbackReverts() public {
Expand Down Expand Up @@ -473,7 +506,7 @@ contract SafeTransferLibTest is SoladyTest {
verifySafeTransferFrom(address(returnsRawBytes), from, to, amount, _SUCCESS);
}

function testTransferFromWithNonContract(
function testTransferFromWithNonContractReverts(
address nonContract,
address from,
address to,
Expand All @@ -482,8 +515,8 @@ contract SafeTransferLibTest is SoladyTest {
if (uint256(uint160(nonContract)) <= 18 || nonContract.code.length > 0) {
return;
}

SafeTransferLib.safeTransferFrom(nonContract, _brutalized(from), _brutalized(to), amount);
vm.expectRevert(SafeTransferLib.TransferFromFailed.selector);
this.safeTransferFrom(nonContract, from, to, amount);
}

function testApproveWithMissingReturn(address to, uint256 amount) public {
Expand All @@ -504,22 +537,26 @@ contract SafeTransferLibTest is SoladyTest {
verifySafeApprove(address(returnsRawBytes), to, amount, _SUCCESS);
}

function testApproveWithNonContract(address nonContract, address to, uint256 amount) public {
function testApproveWithNonContractReverts(address nonContract, address to, uint256 amount)
public
{
if (uint256(uint160(nonContract)) <= 18 || nonContract.code.length > 0) {
return;
}

SafeTransferLib.safeApprove(nonContract, _brutalized(to), amount);
vm.expectRevert(SafeTransferLib.ApproveFailed.selector);
this.safeApprove(nonContract, to, amount);
}

function testApproveWithRetryWithNonContract(address nonContract, address to, uint256 amount)
public
{
function testApproveWithRetryWithNonContractReverts(
address nonContract,
address to,
uint256 amount
) public {
if (uint256(uint160(nonContract)) <= 18 || nonContract.code.length > 0) {
return;
}

SafeTransferLib.safeApproveWithRetry(nonContract, _brutalized(to), amount);
vm.expectRevert(SafeTransferLib.ApproveFailed.selector);
this.safeApproveWithRetry(nonContract, to, amount);
}

function testApproveWithRetry(address to, uint256 amount0, uint256 amount1) public {
Expand Down

0 comments on commit e9c03bf

Please sign in to comment.