Add withdrawProfit zero address check. Simplify withdraw()

Author: lastperson

.env.example

@@ -26,5 +26,5 @@ ETHERSCAN_ARBITRUM_SEPOLIA=
 FORK_PROVIDER=https://eth-mainnet.public.blastapi.io
 USDC_OWNER_ADDRESS=0x7713974908Be4BEd47172370115e8b1219F4A5f0
 RPL_OWNER_ADDRESS=0xdC7b28976d6eb13082a5Be7C66f9dCFE0115738f
-UNI_OWNER_ADDRESS=0x46f34C24A7bA7a2Ac6DD76c3F09B32D41C144d08
+UNI_OWNER_ADDRESS=0x090D4613473dEE047c3f2706764f49E0821D256e
 PRIME_OWNER_ADDRESS=0x1D0065D367DA1919cD597d25F91a97B6039428C5

contracts/LiquidityPool.sol

@@ -142,19 +142,19 @@ contract LiquidityPool is ILiquidityPool, AccessControl, EIP712, Pausable {
 
     // Admin functions
 
+    /// @notice Can withdraw a maximum of totalDeposited. If anything is left, it meant to be withdrawn through
+    /// a withdrawProfit().
     function withdraw(address to, uint256 amount)
         external
         override
         onlyRole(LIQUIDITY_ADMIN_ROLE)
         whenNotPaused()
-        returns (uint256)
     {
-        uint256 withdrawn = _withdrawLogic(to, amount);
         uint256 deposited = totalDeposited;
-        require(deposited >= withdrawn, InsufficientLiquidity());
-        totalDeposited = deposited - withdrawn;
-        emit Withdraw(msg.sender, to, withdrawn);
-        return withdrawn;
+        require(deposited >= amount, InsufficientLiquidity());
+        totalDeposited = deposited - amount;
+        _withdrawLogic(to, amount);
+        emit Withdraw(msg.sender, to, amount);
     }
 
     function withdrawProfit(
@@ -240,6 +240,7 @@ contract LiquidityPool is ILiquidityPool, AccessControl, EIP712, Pausable {
         address[] calldata tokens,
         address to
     ) internal {
+        require(to != address(0), ZeroAddress());
         bool success;
         for (uint256 i = 0; i < tokens.length; i++) {
             IERC20 token = IERC20(tokens[i]);
@@ -261,9 +262,9 @@ contract LiquidityPool is ILiquidityPool, AccessControl, EIP712, Pausable {
         require(borrowToken == address(ASSETS), InvalidBorrowToken());
     }
 
-    function _withdrawLogic(address to, uint256 amount) internal virtual returns (uint256) {
+    function _withdrawLogic(address to, uint256 amount) internal virtual {
+        require(ASSETS.balanceOf(address(this)) >= amount, InsufficientLiquidity());
         ASSETS.safeTransfer(to, amount);
-        return amount;
     }
 
     function _withdrawProfitLogic(IERC20 token) internal virtual returns (uint256) {

contracts/LiquidityPoolAave.sol

@@ -166,14 +166,14 @@ contract LiquidityPoolAave is LiquidityPool {
         _checkTokenLTV(borrowToken);
     }
 
-    function _withdrawLogic(address to, uint256 amount) internal override returns (uint256) {
+    function _withdrawLogic(address to, uint256 amount) internal override {
+        require(ATOKEN.balanceOf(address(this)) >= amount, InsufficientLiquidity());
         // get USDC from AAVE
-        uint256 withdrawn = AAVE_POOL.withdraw(address(ASSETS), amount, to);
+        AAVE_POOL.withdraw(address(ASSETS), amount, to);
         // health factor after withdraw
         (,,,,,uint256 currentHealthFactor) = AAVE_POOL.getUserAccountData(address(this));
         require(currentHealthFactor >= minHealthFactor, HealthFactorTooLow());
-        emit WithdrawnFromAave(to, withdrawn);
-        return withdrawn;
+        emit WithdrawnFromAave(to, amount);
     }
 
     function _withdrawProfitLogic(IERC20 token) internal override returns (uint256) {

contracts/interfaces/ILiquidityPool.sol

@@ -14,7 +14,7 @@ interface ILiquidityPool {
 
     function depositWithPull(uint256 amount) external;
 
-    function withdraw(address to, uint256 amount) external returns (uint256 withdrawn);
+    function withdraw(address to, uint256 amount) external;
 
     function borrow(
         address borrowToken,

contracts/testing/TestLiquidityPool.sol

@@ -21,9 +21,8 @@ contract TestLiquidityPool is ILiquidityPool, AccessControl {
         emit Deposit();
     }
 
-    function withdraw(address to, uint256 amount) external override onlyRole(LIQUIDITY_ADMIN_ROLE) returns (uint256) {
+    function withdraw(address to, uint256 amount) external override onlyRole(LIQUIDITY_ADMIN_ROLE) {
         SafeERC20.safeTransfer(ASSETS, to, amount);
-        return amount;
     }
 
     function depositWithPull(uint256) external pure override {

test/LiquidityPool.ts

@@ -1,5 +1,5 @@
 import {
-  loadFixture, time
+  loadFixture, time, setBalance
 } from "@nomicfoundation/hardhat-toolbox/network-helpers";
 import {expect} from "chai";
 import hre from "hardhat";
@@ -39,6 +39,7 @@ describe("LiquidityPool", function () {
     if (!UNI_OWNER_ADDRESS) throw new Error("Env variables not configured (UNI_OWNER_ADDRESS missing)");
     const uni = await hre.ethers.getContractAt("ERC20", UNI_ADDRESS);
     const uniOwner = await hre.ethers.getImpersonatedSigner(UNI_OWNER_ADDRESS);
+    await setBalance(UNI_OWNER_ADDRESS, 10n ** 18n);
 
     const USDC_DEC = 10n ** (await usdc.decimals());
     const RPL_DEC = 10n ** (await rpl.decimals());
@@ -657,6 +658,12 @@ describe("LiquidityPool", function () {
         .to.be.revertedWithCustomError(liquidityPoolBase, "EnforcedPause");
     });
 
+    it("Should NOT withdraw profit to zero address", async function () {
+      const {liquidityPoolBase, user, usdc, withdrawProfit, pauser} = await loadFixture(deployAll);
+      await expect(liquidityPoolBase.connect(withdrawProfit).withdrawProfit([usdc.target], ZERO_ADDRESS))
+        .to.be.revertedWithCustomError(liquidityPoolBase, "ZeroAddress()");
+    });
+
     it("Should revert during withdrawing profit if no profit", async function () {
       const {liquidityPoolBase, usdc, withdrawProfit, user} = await loadFixture(deployAll);
       await expect(liquidityPoolBase.connect(withdrawProfit).withdrawProfit([usdc.target], user.address))

test/LiquidityPoolAave.ts

@@ -1,11 +1,12 @@
 import {
-  loadFixture, time
+  loadFixture, time, setBalance
 } from "@nomicfoundation/hardhat-toolbox/network-helpers";
 import {expect} from "chai";
 import hre from "hardhat";
 import {
   deploy, signBorrow
 } from "./helpers";
+import {ZERO_ADDRESS} from "../scripts/common";
 import {encodeBytes32String, AbiCoder} from "ethers";
 import {
   MockTarget, MockBorrowSwap, LiquidityPoolAave
@@ -51,6 +52,7 @@ describe("LiquidityPoolAave", function () {
     const uniOwner = await hre.ethers.getImpersonatedSigner(UNI_OWNER_ADDRESS);
     const uniData = await aavePool.getReserveData(UNI_ADDRESS);
     const uniDebtToken = await hre.ethers.getContractAt("ERC20", uniData[10]);
+    await setBalance(UNI_OWNER_ADDRESS, 10n ** 18n);
 
     // PRIME token used as not supported by aave
     const NON_SUPPORTED_TOKEN_ADDRESS = "0xb23d80f5FefcDDaa212212F028021B41DEd428CF";
@@ -627,7 +629,7 @@ describe("LiquidityPoolAave", function () {
       const usdcDebtBefore = await usdcDebtToken.balanceOf(liquidityPool.target);
       expect(usdcDebtBefore).to.be.greaterThan(amountToBorrow);
 
-      await usdc.connect(uniOwner).transfer(liquidityPool.target, amountToBorrow);
+      await usdc.connect(usdcOwner).transfer(liquidityPool.target, amountToBorrow);
 
       await expect(liquidityPool.connect(user).repay([usdc.target]))
         .to.emit(liquidityPool, "Repaid");  
@@ -892,7 +894,7 @@ describe("LiquidityPoolAave", function () {
       await expect(liquidityPool.connect(liquidityAdmin).deposit(amountCollateral))
         .to.emit(liquidityPool, "SuppliedToAave");
 
-      await expect(liquidityPool.connect(admin).setHealthFactor(4000n * 10n ** 18n / 100n))
+      await expect(liquidityPool.connect(admin).setHealthFactor(5000n * 10n ** 18n / 100n))
         .to.emit(liquidityPool, "HealthFactorSet");
 
       const amountToBorrow = 3n * UNI_DEC;
@@ -1274,6 +1276,12 @@ describe("LiquidityPoolAave", function () {
         .to.be.revertedWithCustomError(liquidityPool, "EnforcedPause");
     });
 
+    it("Should NOT withdraw profit to zero address", async function () {
+      const {liquidityPool, user, uni, withdrawProfit, pauser} = await loadFixture(deployAll);
+      await expect(liquidityPool.connect(withdrawProfit).withdrawProfit([uni.target], ZERO_ADDRESS))
+        .to.be.revertedWithCustomError(liquidityPool, "ZeroAddress()");
+    });
+
     it("Should revert during withdrawing profit if no profit", async function () {
       const {liquidityPool, uni, withdrawProfit, user} = await loadFixture(deployAll);
       await expect(liquidityPool.connect(withdrawProfit).withdrawProfit([uni.target], user.address))