feat: custom overflow errors

morpho-org · Feb 16, 2024 · 6335a15 · 6335a15
1 parent ef372d5
commit 6335a15
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 2 deletions.
diff --git a/src/PublicAllocator.sol b/src/PublicAllocator.sol
@@ -103,7 +103,9 @@ contract PublicAllocator is IPublicAllocatorStaticTyping {
             uint128 withdrawnAssets = withdrawals[i].amount;
             totalWithdrawn += withdrawnAssets;
             flowCap[id].maxIn += withdrawnAssets;
+            if (flowCap[id].maxOut < withdrawnAssets) revert ErrorsLib.MaxOutflowExceeded(id);
             flowCap[id].maxOut -= withdrawnAssets;
+            if (assets < withdrawnAssets) revert ErrorsLib.NotEnoughSupply(id);
             allocations[i].assets = assets - withdrawnAssets;
             allocations[i].marketParams = withdrawals[i].marketParams;
 
@@ -112,6 +114,7 @@ contract PublicAllocator is IPublicAllocatorStaticTyping {
 
         allocations[withdrawals.length].marketParams = supplyMarketParams;
         allocations[withdrawals.length].assets = type(uint256).max;
+        if (flowCap[supplyMarketId].maxIn < totalWithdrawn) revert ErrorsLib.MaxInflowExceeded(supplyMarketId);
         flowCap[supplyMarketId].maxIn -= totalWithdrawn;
         flowCap[supplyMarketId].maxOut += totalWithdrawn;
 

diff --git a/src/libraries/ErrorsLib.sol b/src/libraries/ErrorsLib.sol
@@ -31,4 +31,13 @@ library ErrorsLib {
 
     /// @notice Thrown when the PublicAllocatorFactory is called with a vault not made by the MetaMorphoFactory.
     error NotMetaMorpho();
+
+    /// @notice Thrown when attempting to withdraw more than the available supply of a market.
+    error NotEnoughSupply(Id id);
+
+    /// @notice Thrown when attempting to withdraw more than the max outflow of a market.
+    error MaxOutflowExceeded(Id id);
+
+    /// @notice Thrown when attempting to supply more than the max inflow of a market.
+    error MaxInflowExceeded(Id id);
 }
diff --git a/test/PublicAllocatorTest.sol b/test/PublicAllocatorTest.sol
@@ -77,7 +77,7 @@ contract PublicAllocatorTest is IntegrationTest {
         vm.prank(OWNER);
         publicAllocator.setFlowCaps(flowCaps);
         withdrawals.push(Withdrawal(idleParams, flow));
-        vm.expectRevert(stdError.arithmeticError);
+        vm.expectRevert(abi.encodeWithSelector(ErrorsLib.MaxOutflowExceeded.selector,idleParams.id()));
         publicAllocator.reallocateTo(withdrawals, allMarkets[0]);
     }
 
@@ -88,7 +88,7 @@ contract PublicAllocatorTest is IntegrationTest {
         vm.prank(OWNER);
         publicAllocator.setFlowCaps(flowCaps);
         withdrawals.push(Withdrawal(idleParams, flow));
-        vm.expectRevert(stdError.arithmeticError);
+        vm.expectRevert(abi.encodeWithSelector(ErrorsLib.MaxInflowExceeded.selector,allMarkets[0].id()));
         publicAllocator.reallocateTo(withdrawals, allMarkets[0]);
     }
 
@@ -484,4 +484,49 @@ contract PublicAllocatorTest is IntegrationTest {
         vm.prank(OWNER);
         publicAllocator.setFlowCaps(flowCaps);
     }
+
+    function testNotEnoughSupply() public {
+        uint128 flow = 1e18;
+        // Set flow limits with withdraw market's maxIn to max
+        flowCaps.push(FlowConfig(idleParams.id(), FlowCap(MAX_SETTABLE_FLOW_CAP, MAX_SETTABLE_FLOW_CAP)));
+        flowCaps.push(FlowConfig(allMarkets[0].id(), FlowCap(MAX_SETTABLE_FLOW_CAP, MAX_SETTABLE_FLOW_CAP)));
+        vm.prank(OWNER);
+        publicAllocator.setFlowCaps(flowCaps);
+
+        withdrawals.push(Withdrawal(idleParams, flow));
+        publicAllocator.reallocateTo(withdrawals, allMarkets[0]);
+
+        delete withdrawals;
+
+        withdrawals.push(Withdrawal(allMarkets[0],flow+1));
+        vm.expectRevert(abi.encodeWithSelector(ErrorsLib.NotEnoughSupply.selector,allMarkets[0].id()));
+        publicAllocator.reallocateTo(withdrawals,idleParams);
+    }
+
+    function testMaxOutflowExceeded() public {
+        uint128 cap = 1e18;
+        // Set flow limits with withdraw market's maxIn to max
+        flowCaps.push(FlowConfig(idleParams.id(), FlowCap(MAX_SETTABLE_FLOW_CAP, cap)));
+        flowCaps.push(FlowConfig(allMarkets[0].id(), FlowCap(MAX_SETTABLE_FLOW_CAP, MAX_SETTABLE_FLOW_CAP)));
+        vm.prank(OWNER);
+        publicAllocator.setFlowCaps(flowCaps);
+
+        withdrawals.push(Withdrawal(idleParams, cap+1));
+        vm.expectRevert(abi.encodeWithSelector(ErrorsLib.MaxOutflowExceeded.selector,idleParams.id()));
+        publicAllocator.reallocateTo(withdrawals, allMarkets[0]);
+    }
+
+
+    function testMaxInflowExceeded() public {
+        uint128 cap = 1e18;
+        // Set flow limits with withdraw market's maxIn to max
+        flowCaps.push(FlowConfig(idleParams.id(), FlowCap(MAX_SETTABLE_FLOW_CAP, MAX_SETTABLE_FLOW_CAP)));
+        flowCaps.push(FlowConfig(allMarkets[0].id(), FlowCap(cap, MAX_SETTABLE_FLOW_CAP)));
+        vm.prank(OWNER);
+        publicAllocator.setFlowCaps(flowCaps);
+
+        withdrawals.push(Withdrawal(idleParams, cap+1));
+        vm.expectRevert(abi.encodeWithSelector(ErrorsLib.MaxInflowExceeded.selector,allMarkets[0].id()));
+        publicAllocator.reallocateTo(withdrawals, allMarkets[0]);
+    }
 }