From 763de53a36243490ef46a2c702c5a1480554f286 Mon Sep 17 00:00:00 2001
From: JoscelynFarr <joscelyn@jojo.exchange>
Date: Fri, 26 Jan 2024 15:34:23 +0800
Subject: [PATCH] Fixed issue7:
 https://github.com/sherlock-audit/2023-12-jojo-exchange-update-judging/issues/7

---
 src/JOJOOperation.sol        | 4 ++++
 src/libraries/Funding.sol    | 1 +
 src/libraries/Operation.sol  | 7 +++++++
 src/libraries/Types.sol      | 1 +
 test/impl/DealerFundTest.sol | 2 +-
 test/impl/Subaccount.t.sol   | 1 +
 6 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/JOJOOperation.sol b/src/JOJOOperation.sol
index 9953d65..7ac0a1e 100644
--- a/src/JOJOOperation.sol
+++ b/src/JOJOOperation.sol
@@ -60,6 +60,10 @@ abstract contract JOJOOperation is JOJOStorage, IDealer {
         Operation.setFastWithdrawalWhitelist(state, target, isValid);
     }
 
+    function setWithdrawlWhitelist(address target, bool isValid) external onlyOwner {
+        Operation.setWithdrawalWhitelist(state, target, isValid);
+    }
+
     function disableFastWithdraw(bool disabled) external onlyOwner {
         Operation.disableFastWithdraw(state, disabled);
     }
diff --git a/src/libraries/Funding.sol b/src/libraries/Funding.sol
index fdf8a64..909fee4 100644
--- a/src/libraries/Funding.sol
+++ b/src/libraries/Funding.sol
@@ -172,6 +172,7 @@ library Funding {
 
         if (param.length != 0) {
             require(Address.isContract(to), "target is not a contract");
+            require(state.isWithdrawalWhitelist[to], "target is not in whiteList");
             (bool success,) = to.call(param);
             if (success == false) {
                 assembly {
diff --git a/src/libraries/Operation.sol b/src/libraries/Operation.sol
index 774251b..2323d05 100644
--- a/src/libraries/Operation.sol
+++ b/src/libraries/Operation.sol
@@ -25,6 +25,8 @@ library Operation {
 
     event SetFastWithdrawalWhitelist(address target, bool isValid);
 
+    event SetWithdrawalWhitelist(address target, bool isValid);
+
     event FastWithdrawDisabled(bool disabled);
 
     event SetOperator(address indexed client, address indexed operator, bool isValid);
@@ -116,6 +118,11 @@ library Operation {
         emit FastWithdrawDisabled(disabled);
     }
 
+    function setWithdrawalWhitelist(Types.State storage state, address target, bool isValid) external {
+        state.isWithdrawalWhitelist[target] = isValid;
+        emit SetFastWithdrawalWhitelist(target, isValid);
+    }
+
     function setOperator(Types.State storage state, address client, address operator, bool isValid) external {
         state.operatorRegistry[client][operator] = isValid;
         emit SetOperator(client, operator, isValid);
diff --git a/src/libraries/Types.sol b/src/libraries/Types.sol
index 13784a6..0ed68e5 100644
--- a/src/libraries/Types.sol
+++ b/src/libraries/Types.sol
@@ -55,6 +55,7 @@ library Types {
         // funding rate keeper, normally an EOA account
         address fundingRateKeeper;
         uint256 maxPositionAmount;
+        mapping(address => bool) isWithdrawalWhitelist;
     }
 
     struct Order {
diff --git a/test/impl/DealerFundTest.sol b/test/impl/DealerFundTest.sol
index 2873243..f170739 100644
--- a/test/impl/DealerFundTest.sol
+++ b/test/impl/DealerFundTest.sol
@@ -130,7 +130,7 @@ contract FundTest is Checkers {
         vm.startPrank(traders[1]);
         jojoDealer.fastWithdraw(traders[0], traders[0], 100e6, 0, false, "");
         jojoDealer.fastWithdraw(traders[0], traders[0], 0, 100e6, false, "");
-        cheats.expectRevert("Ownable: caller is not the owner");
+        cheats.expectRevert("target is not in whiteList");
         jojoDealer.fastWithdraw(
             traders[1],
             address(jojoDealer),
diff --git a/test/impl/Subaccount.t.sol b/test/impl/Subaccount.t.sol
index b6bca5c..f505d84 100644
--- a/test/impl/Subaccount.t.sol
+++ b/test/impl/Subaccount.t.sol
@@ -93,6 +93,7 @@ contract SubaccountTest is JUSDBankInitTest {
         usdc.mint(alice, 1000e6);
         jusd.mint(1000e6);
         jusd.transfer(alice, 1000e6);
+        jojoDealer.setWithdrawlWhitelist(address(jusdRepayHelper), true);
         vm.startPrank(alice);
         eth.approve(address(jusdBank), 10e18);
         address aliceSub = subaccountFactory.newSubaccount();