From 7105693e3c72ad73e4d722cc3db79d065b9a0875 Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Thu, 21 Nov 2024 16:34:00 +0100 Subject: [PATCH] Change NoncesKeyed._useNonce to return a keyed value (#5312) --- contracts/utils/NoncesKeyed.sol | 31 +++++++++++++-------- test/utils/Nonces.behavior.js | 49 +++++++++++++++++++++++++++++---- 2 files changed, 63 insertions(+), 17 deletions(-) diff --git a/contracts/utils/NoncesKeyed.sol b/contracts/utils/NoncesKeyed.sol index c775700d2..7984dc70e 100644 --- a/contracts/utils/NoncesKeyed.sol +++ b/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)); } } diff --git a/test/utils/Nonces.behavior.js b/test/utils/Nonces.behavior.js index 170739664..32201f690 100644 --- a/test/utils/Nonces.behavior.js +++ b/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); }); }); });