Test _callbackResult is not unecessarily deleted

alcueca · Jul 30, 2023 · 5b3a9e2 · 5b3a9e2
1 parent 6f72724
commit 5b3a9e2
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 2 deletions.
diff --git a/src/aave/AaveWrapper.sol b/src/aave/AaveWrapper.sol
@@ -90,7 +90,10 @@ contract AaveWrapper is IERC3156PPFlashLender, IFlashLoanSimpleReceiver {
         });
 
         result = _callbackResult;
-        delete _callbackResult;
+        // Avoid storage write if not needed
+        if (result.length > 0) {
+            delete _callbackResult;
+        }
         return result;
     }
 

diff --git a/src/test/FlashBorrower.sol b/src/test/FlashBorrower.sol
@@ -74,6 +74,21 @@ contract FlashBorrower {
         return abi.encode(ERC3156PP_CALLBACK_SUCCESS);
     }
 
+    function onFlashLoanVoid(address initiator, address paymentReceiver, IERC20 asset, uint256 amount, uint256 fee, bytes calldata) external returns(bytes memory) {
+        require(msg.sender == address(lender), "FlashBorrower: Untrusted lender");
+        require(initiator == address(this), "FlashBorrower: External loan initiator");
+
+        flashInitiator = initiator;
+        flashAsset = asset;
+        flashAmount = amount;
+        flashFee = fee;
+        loanReceiver.retrieve(asset);
+        flashBalance = IERC20(asset).balanceOf(address(this));
+        asset.transfer(paymentReceiver, amount + fee);
+
+        return "";
+    }
+
     function flashBorrow(IERC20 asset, uint256 amount) public returns(bytes memory) {
         return lender.flashLoan(address(loanReceiver), asset, amount, "", this.onFlashLoan);
     }
@@ -85,4 +100,8 @@ contract FlashBorrower {
     function flashBorrowAndReenter(IERC20 asset, uint256 amount) public returns(bytes memory) {
         return lender.flashLoan(address(loanReceiver), asset, amount, "", this.onReenter);
     }
+
+    function flashBorrowVoid(IERC20 asset, uint256 amount) public returns(bytes memory) {
+        return lender.flashLoan(address(loanReceiver), asset, amount, "", this.onFlashLoanVoid);
+    }
 }
diff --git a/test/AaveWrapper.t.sol b/test/AaveWrapper.t.sol
@@ -67,6 +67,22 @@ contract AaveWrapperTest is PRBTest, StdCheats {
         assertEq(vm.load(address(wrapper), bytes32(uint256(0))), "");
     }
 
+    function test_flashLoan_void() external {
+        console2.log("test_flashLoan");
+        uint256 loan = 1e18;
+        uint256 fee = wrapper.flashFee(dai, loan);
+        dai.transfer(address(borrower), fee);
+
+        vm.record();
+        bytes memory result = borrower.flashBorrowVoid(dai, loan);
+
+        // Test the return values
+        assertEq(result, "", "Void result");
+
+        (, bytes32[] memory writeSlots) = vm.accesses(address(wrapper));
+        assertEq(writeSlots.length, 0, "writeSlots");
+    }
+
     function test_executeOperation() public {
         bytes memory data = abi.encode(address(0), address(this), this._voidCallback.encodeFunction(), "");
 
@@ -85,7 +101,18 @@ contract AaveWrapperTest is PRBTest, StdCheats {
         assertEq(writeSlots.length, 0, "writeSlots");
     }
 
-    function _voidCallback(address, address, IERC20, uint256, uint256, bytes memory) external pure returns (bytes memory) {
+    function _voidCallback(
+        address,
+        address,
+        IERC20,
+        uint256,
+        uint256,
+        bytes memory
+    )
+        external
+        pure
+        returns (bytes memory)
+    {
         return "";
     }
 }