From 88edccb893295d34d59e30a40533bcca2b2da716 Mon Sep 17 00:00:00 2001
From: vsevolod <5249361+Vsevo1od@users.noreply.github.com>
Date: Wed, 26 Jan 2022 21:48:30 +0300
Subject: [PATCH] Add nonce to ControlWalletProxy

---
 contracts/periphery/ControlWalletProxy.sol | 13 ++++++++++-
 test/ControlWalletProxy.test.ts            | 27 +++++++++++++++++++---
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/contracts/periphery/ControlWalletProxy.sol b/contracts/periphery/ControlWalletProxy.sol
index cc51e965..e1f282a5 100644
--- 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
index 47c7234c..3f2040ac 100644
--- 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', () => {