updated flow and unit test

axelarnetwork · Nov 5, 2024 · 2d32939 · 2d32939
1 parent b875a83
commit 2d32939
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 15 deletions.
diff --git a/contracts/interfaces/IFlowLimit.sol b/contracts/interfaces/IFlowLimit.sol
@@ -8,7 +8,8 @@ pragma solidity ^0.8.0;
  */
 interface IFlowLimit {
     error FlowLimitExceeded(uint256 limit, uint256 flowAmount, address tokenManager);
-    error InvalidFlowLimit(uint256 flowLimit, bytes32 tokenId);
+    error FlowAdditionOverflow(uint256 flowToAdd, uint256 flowAmount, address tokenManager);
+    error FlowLimitOverflow(uint256 flowLimit, uint256 flowToCompare, address tokenManager);
 
     event FlowLimitSet(bytes32 indexed tokenId, address operator, uint256 flowLimit_);
 

diff --git a/contracts/utils/FlowLimit.sol b/contracts/utils/FlowLimit.sol
@@ -35,8 +35,6 @@ contract FlowLimit is IFlowLimit {
      * @param tokenId The id of the token to set the flow limit for.
      */
     function _setFlowLimit(uint256 flowLimit_, bytes32 tokenId) internal {
-        if (flowLimit_ == type(uint256).max) revert InvalidFlowLimit(flowLimit_, tokenId);
-
         assembly {
             sstore(FLOW_LIMIT_SLOT, flowLimit_)
         }
@@ -104,6 +102,10 @@ contract FlowLimit is IFlowLimit {
             flowToCompare := sload(slotToCompare)
         }
 
+        // Overflow checks for flowToAdd + flowAmount and flowToCompare + flowLimit_
+        if (flowToAdd > type(uint256).max - flowAmount) revert FlowAdditionOverflow(flowAmount, flowToAdd, address(this));
+        if (flowToCompare > type(uint256).max - flowLimit_) revert FlowLimitOverflow(flowLimit_, flowToCompare, address(this));
+
         if (flowToAdd + flowAmount > flowToCompare + flowLimit_)
             revert FlowLimitExceeded((flowToCompare + flowLimit_), flowToAdd + flowAmount, address(this));
         if (flowAmount > flowLimit_) revert FlowLimitExceeded(flowLimit_, flowAmount, address(this));

diff --git a/test/InterchainTokenService.js b/test/InterchainTokenService.js
@@ -2714,22 +2714,13 @@ describe('Interchain Token Service', () => {
         let tokenManager, tokenId;
         const sendAmount = 1234;
         const flowLimit = (sendAmount * 3) / 2;
-        const mintAmount = flowLimit * 3;
+        const mintAmount = MaxUint256;
 
         before(async () => {
             [, tokenManager, tokenId] = await deployFunctions.mintBurn(service, 'Test Token Lock Unlock', 'TT', 12, mintAmount);
             await tokenManager.setFlowLimit(flowLimit).then((tx) => tx.wait);
         });
 
-        it('Should be reverted setting invalid flow limit', async () => {
-            const invalidFlowLimit = MaxUint256;
-
-            await expectRevert((gasOptions) => tokenManager.setFlowLimit(invalidFlowLimit, gasOptions), tokenManager, 'InvalidFlowLimit', [
-                invalidFlowLimit,
-                tokenId,
-            ]);
-        });
-
         it('Should be able to send token only if it does not trigger the mint limit', async () => {
             await service.interchainTransfer(tokenId, destinationChain, destinationAddress, sendAmount, '0x', 0).then((tx) => tx.wait);
 
@@ -2755,10 +2746,10 @@ describe('Interchain Token Service', () => {
             expect(flowIn).to.eq(0);
             expect(flowOut).to.eq(sendAmount);
 
-            async function receiveToken(sendAmount) {
+            async function receiveToken(amount) {
                 const payload = defaultAbiCoder.encode(
                     ['uint256', 'bytes32', 'bytes', 'bytes', 'uint256', 'bytes'],
-                    [MESSAGE_TYPE_INTERCHAIN_TRANSFER, tokenId, hexlify(wallet.address), wallet.address, sendAmount, '0x'],
+                    [MESSAGE_TYPE_INTERCHAIN_TRANSFER, tokenId, hexlify(wallet.address), wallet.address, amount, '0x'],
                 );
                 const commandId = await approveContractCall(gateway, destinationChain, service.address, service.address, payload);
 
@@ -2785,6 +2776,68 @@ describe('Interchain Token Service', () => {
             ]);
         });
 
+        it('Should revert if the flow limit overflows', async () => {
+            const tokenFlowLimit = await service.flowLimit(tokenId);
+            expect(tokenFlowLimit).to.eq(flowLimit);
+
+            let flowIn = await service.flowInAmount(tokenId);
+            let flowOut = await service.flowOutAmount(tokenId);
+
+            expect(flowIn).to.eq(sendAmount);
+            expect(flowOut).to.eq(sendAmount);
+
+            let newFlowLimit = MaxUint256;
+            let newSendAmount = 1;
+
+            await tokenManager.setFlowLimit(newFlowLimit).then((tx) => tx.wait);
+
+            async function receiveToken(amount) {
+                const payload = defaultAbiCoder.encode(
+                    ['uint256', 'bytes32', 'bytes', 'bytes', 'uint256', 'bytes'],
+                    [MESSAGE_TYPE_INTERCHAIN_TRANSFER, tokenId, hexlify(wallet.address), wallet.address, amount, '0x'],
+                );
+                const commandId = await approveContractCall(gateway, destinationChain, service.address, service.address, payload);
+
+                return service.execute(commandId, destinationChain, service.address, payload);
+            }
+
+            const errorSignatureHash = id('FlowLimitOverflow(uint256,uint256,address)');
+            const selector = errorSignatureHash.substring(0, 10);
+            const errorData = defaultAbiCoder.encode(['uint256', 'uint256', 'address'], [newFlowLimit, flowIn, tokenManager.address]);
+
+            await expectRevert((gasOptions) => receiveToken(newSendAmount, gasOptions), service, 'GiveTokenFailed', [
+                selector + errorData.substring(2),
+            ]);
+        });
+
+        it('Should revert if the flow addition overflows', async () => {
+            let newFlowLimit = MaxUint256;
+            let newSendAmount = MaxUint256;
+
+            const tokenFlowLimit = await service.flowLimit(tokenId);
+            expect(tokenFlowLimit).to.eq(MaxUint256);
+
+            let flowIn = await service.flowInAmount(tokenId);
+            let flowOut = await service.flowOutAmount(tokenId);
+
+            expect(flowIn).to.eq(sendAmount);
+            expect(flowOut).to.eq(sendAmount);
+
+            await tokenManager.setFlowLimit(newFlowLimit).then((tx) => tx.wait);
+
+            const errorSignatureHash = id('FlowAdditionOverflow(uint256,uint256,address)');
+            const selector = errorSignatureHash.substring(0, 10);
+            const errorData = defaultAbiCoder.encode(['uint256', 'uint256', 'address'], [newSendAmount, flowOut, tokenManager.address]);
+
+            await expectRevert(
+                (gasOptions) =>
+                    service.interchainTransfer(tokenId, destinationChain, destinationAddress, newSendAmount, '0x', 0, gasOptions),
+                service,
+                'TakeTokenFailed',
+                [selector + errorData.substring(2)],
+            );
+        });
+
         it('Should be able to set flow limits for each token manager', async () => {
             const tokenIds = [];
             const tokenManagers = [];