Change NoncesKeyed._useNonce to return a keyed value (#5312)

contracts/utils/NoncesKeyed.sol

@@ -13,7 +13,7 @@ abstract contract NoncesKeyed is Nonces {
 
     /// @dev Returns the next unused nonce for an address and key. Result contains the key prefix.
     function nonces(address owner, uint192 key) public view virtual returns (uint256) {
-        return key == 0 ? nonces(owner) : ((uint256(key) << 64) | _nonces[owner][key]);
+        return key == 0 ? nonces(owner) : _pack(key, _nonces[owner][key]);
     }
 
     /**
@@ -27,7 +27,7 @@ abstract contract NoncesKeyed is Nonces {
         // decremented or reset. This guarantees that the nonce never overflows.
         unchecked {
             // It is important to do x++ and not ++x here.
-            return key == 0 ? _useNonce(owner) : _nonces[owner][key]++;
+            return key == 0 ? _useNonce(owner) : _pack(key, _nonces[owner][key]++);
         }
     }
 
@@ -39,7 +39,13 @@ abstract contract NoncesKeyed is Nonces {
      * - use the last 24 bytes for the nonce
      */
     function _useCheckedNonce(address owner, uint256 keyNonce) internal virtual override {
-        _useCheckedNonce(owner, uint192(keyNonce >> 64), uint64(keyNonce));
+        (uint192 key, ) = _unpack(keyNonce);
+        if (key == 0) {
+            super._useCheckedNonce(owner, keyNonce);
+        } else {
+            uint256 current = _useNonce(owner, key);
+            if (keyNonce != current) revert InvalidAccountNonce(owner, current);
+        }
     }
 
     /**
@@ -48,13 +54,16 @@ abstract contract NoncesKeyed is Nonces {
      * This version takes the key and the nonce as two different parameters.
      */
     function _useCheckedNonce(address owner, uint192 key, uint64 nonce) internal virtual {
-        if (key == 0) {
-            super._useCheckedNonce(owner, nonce);
-        } else {
-            uint256 current = _useNonce(owner, key);
-            if (nonce != current) {
-                revert InvalidAccountNonce(owner, current);
+        _useCheckedNonce(owner, _pack(key, nonce));
     }
+
+    /// @dev Pack key and nonce into a keyNonce
+    function _pack(uint192 key, uint64 nonce) private pure returns (uint256) {
+        return (uint256(key) << 64) | nonce;
     }
+
+    /// @dev Unpack a keyNonce into its key and nonce components
+    function _unpack(uint256 keyNonce) private pure returns (uint192 key, uint64 nonce) {
+        return (uint192(keyNonce >> 64), uint64(keyNonce));
     }
 }

test/utils/Nonces.behavior.js

@@ -86,7 +86,7 @@ function shouldBehaveLikeNoncesKeyed() {
 
         await expect(this.mock.$_useNonce(sender, ethers.Typed.uint192(0n)))
           .to.emit(this.mock, 'return$_useNonce_address_uint192')
-          .withArgs(1n);
+          .withArgs(keyOffset(0n) + 1n);
 
         expect(this.mock.nonces(sender, ethers.Typed.uint192(0n))).to.eventually.equal(keyOffset(0n) + 2n);
         expect(this.mock.nonces(sender, ethers.Typed.uint192(17n))).to.eventually.equal(keyOffset(17n) + 0n);
@@ -98,18 +98,18 @@ function shouldBehaveLikeNoncesKeyed() {
 
         await expect(this.mock.$_useNonce(sender, ethers.Typed.uint192(17n)))
           .to.emit(this.mock, 'return$_useNonce_address_uint192')
-          .withArgs(0n);
+          .withArgs(keyOffset(17n) + 0n);
 
         await expect(this.mock.$_useNonce(sender, ethers.Typed.uint192(17n)))
           .to.emit(this.mock, 'return$_useNonce_address_uint192')
-          .withArgs(1n);
+          .withArgs(keyOffset(17n) + 1n);
 
         expect(this.mock.nonces(sender, ethers.Typed.uint192(0n))).to.eventually.equal(keyOffset(0n) + 0n);
         expect(this.mock.nonces(sender, ethers.Typed.uint192(17n))).to.eventually.equal(keyOffset(17n) + 2n);
       });
     });
 
-    describe('_useCheckedNonce', function () {
+    describe('_useCheckedNonce(address, uint256)', function () {
       it('default variant uses key 0', async function () {
         const currentNonce = await this.mock.nonces(sender, ethers.Typed.uint192(0n));
 
@@ -135,12 +135,49 @@ function shouldBehaveLikeNoncesKeyed() {
         // reuse same nonce
         await expect(this.mock.$_useCheckedNonce(sender, currentNonce))
           .to.be.revertedWithCustomError(this.mock, 'InvalidAccountNonce')
-          .withArgs(sender, 1);
+          .withArgs(sender, currentNonce + 1n);
 
         // use "future" nonce too early
         await expect(this.mock.$_useCheckedNonce(sender, currentNonce + 10n))
           .to.be.revertedWithCustomError(this.mock, 'InvalidAccountNonce')
-          .withArgs(sender, 1);
+          .withArgs(sender, currentNonce + 1n);
+      });
+    });
+
+    describe('_useCheckedNonce(address, uint192, uint64)', function () {
+      const MASK = 0xffffffffffffffffn;
+
+      it('default variant uses key 0', async function () {
+        const currentNonce = await this.mock.nonces(sender, ethers.Typed.uint192(0n));
+
+        await this.mock.$_useCheckedNonce(sender, ethers.Typed.uint192(0n), currentNonce);
+
+        expect(this.mock.nonces(sender, ethers.Typed.uint192(0n))).to.eventually.equal(currentNonce + 1n);
+      });
+
+      it('use nonce at another key', async function () {
+        const currentNonce = await this.mock.nonces(sender, ethers.Typed.uint192(17n));
+
+        await this.mock.$_useCheckedNonce(sender, ethers.Typed.uint192(17n), currentNonce & MASK);
+
+        expect(this.mock.nonces(sender, ethers.Typed.uint192(17n))).to.eventually.equal(currentNonce + 1n);
+      });
+
+      it('reverts when nonce is not the expected', async function () {
+        const currentNonce = await this.mock.nonces(sender, ethers.Typed.uint192(42n));
+
+        // use and increment
+        await this.mock.$_useCheckedNonce(sender, ethers.Typed.uint192(42n), currentNonce & MASK);
+
+        // reuse same nonce
+        await expect(this.mock.$_useCheckedNonce(sender, ethers.Typed.uint192(42n), currentNonce & MASK))
+          .to.be.revertedWithCustomError(this.mock, 'InvalidAccountNonce')
+          .withArgs(sender, currentNonce + 1n);
+
+        // use "future" nonce too early
+        await expect(this.mock.$_useCheckedNonce(sender, ethers.Typed.uint192(42n), (currentNonce & MASK) + 10n))
+          .to.be.revertedWithCustomError(this.mock, 'InvalidAccountNonce')
+          .withArgs(sender, currentNonce + 1n);
       });
     });
   });