Add nonce to ControlWalletProxy

debridge-finance · Jan 26, 2022 · 88edccb · 88edccb
1 parent ccebe4e
commit 88edccb
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 4 deletions.
diff --git a/contracts/periphery/ControlWalletProxy.sol b/contracts/periphery/ControlWalletProxy.sol
@@ -18,6 +18,8 @@ contract ControlWalletProxy is Initializable {
     // chainIdFrom => list of addresses that can control this contract
     mapping(uint256 => mapping(bytes => bool)) public controlParams;
 
+    uint256 public nonce;
+
     /* ========== ERRORS ========== */
 
     error CallProxyBadRole();
@@ -28,6 +30,8 @@ contract ControlWalletProxy is Initializable {
     error RemovingMissingAddress();
     error RemovingLastAddress();
 
+    error WrongNonce();
+
     /* ========== EVENTS ========== */
 
     // emitted when controlling address updated
@@ -67,8 +71,13 @@ contract ControlWalletProxy is Initializable {
         address token,
         uint256 amount,
         address destination,
-        bytes memory data
+        bytes memory data,
+        uint256 _nonce
     ) external payable onlyCallProxyFromControllingAddress returns (bool _result) {
+        if (_nonce != nonce){
+            revert WrongNonce();
+        }
+
         if (token != address(0)) {
             IERC20Upgradeable(token).safeApprove(destination, 0);
             IERC20Upgradeable(token).safeApprove(destination, amount);
@@ -83,6 +92,8 @@ contract ControlWalletProxy is Initializable {
         if (!_result) {
             revert ExternalCallFailed();
         }
+
+        nonce++;
     }
 
     function addControllingAddress(

diff --git a/test/ControlWalletProxy.test.ts b/test/ControlWalletProxy.test.ts
@@ -48,7 +48,7 @@ beforeEach(async () => {
 });
 
 
-describe('initializer', async () => {
+describe('initializer', () => {
     test('can be called only once', async () => {
         await expect(initializeControlWalletProxy()).to.be.revertedWith("Initializable: contract is already initialized");
     })
@@ -151,7 +151,7 @@ const testOnlyCallProxyFromControllingAddressModifier = (
 }
 
 describe('`call` method', () => {
-    let mockArgsForCall: [string, BigNumberish, string, BytesLike];
+    let mockArgsForCall: [string, BigNumberish, string, BytesLike, BigNumberish];
     let callWithMockArgsCallData: string;
 
     beforeEach(async () => {
@@ -165,7 +165,7 @@ describe('`call` method', () => {
             .encodeFunctionData('someFunction');
 
         mockArgsForCall = [
-            erc20Mock.address, 0, mockAcceptAnyCall.address, callDataForMockAcceptAnyCall
+            erc20Mock.address, 0, mockAcceptAnyCall.address, callDataForMockAcceptAnyCall, 0
         ];
         callWithMockArgsCallData = controlWalletProxy.interface.encodeFunctionData('call', mockArgsForCall);
     })
@@ -180,6 +180,27 @@ describe('`call` method', () => {
             .encodeFunctionData('unknownFunction');
         await expectCallThroughCallProxyFail(callDataThatWillBeRejectedByWallet);
     })
+
+    test('reverts on wrong nonce', async () => {
+        await expectCallThroughCallProxySuccess(callWithMockArgsCallData);
+        await expectCallThroughCallProxyFail(callWithMockArgsCallData);
+        // replace last element (nonce) with new nonce
+        mockArgsForCall[4] = 1;
+        await expectCallThroughCallProxySuccess(
+            controlWalletProxy.interface.encodeFunctionData('call', mockArgsForCall)
+        );
+
+        // nonce is bigger, should be 2
+        mockArgsForCall[4] = 3;
+        await expectCallThroughCallProxyFail(
+            controlWalletProxy.interface.encodeFunctionData('call', mockArgsForCall)
+        );
+
+        mockArgsForCall[4] = 2;
+        await expectCallThroughCallProxySuccess(
+            controlWalletProxy.interface.encodeFunctionData('call', mockArgsForCall)
+        );
+    });
 })
 
 describe('add/remove calling address', () => {