diff --git a/src/utils/LibPRNG.sol b/src/utils/LibPRNG.sol index 4b49a69e4..f35699bed 100644 --- a/src/utils/LibPRNG.sol +++ b/src/utils/LibPRNG.sol @@ -13,7 +13,8 @@ library LibPRNG { /// @dev The initial length must be greater than zero and less than `2**32 - 1`. error InvalidInitialLazyShufflerLength(); - /// @dev The new length must not be less than the current length. + /// @dev The new length must not be less than the current length, + /// and must be less than `2**32 - 1`. error InvalidNewLazyShufflerLength(); /// @dev The lazy shuffler has not been initialized. @@ -343,19 +344,22 @@ library LibPRNG { /// @dev Increases the length of `$`. /// Reverts if `$` has not been initialized. + /// Reverts if `n` is less than the current length, or if `n >= 2**32 - 1`. + /// Reverts if `n` crosses the entry width boundary at a length of 65535. function grow(LazyShuffler storage $, uint256 n) internal { /// @solidity memory-safe-assembly assembly { let state := sload($.slot) // The packed value at `$`. - // If the new length is smaller than the old length, revert. - if lt(n, shr(224, state)) { - mstore(0x00, 0xbed37c6e) // `InvalidNewLazyShufflerLength()`. - revert(0x1c, 0x04) - } if iszero(state) { mstore(0x00, 0x1ead2566) // `LazyShufflerNotInitialized()`. revert(0x1c, 0x04) } + let o := shr(224, state) // The old length. + let limit := or(0xfffe, mul(0xffff0000, gt(o, 0xfffe))) + if or(lt(n, o), gt(n, limit)) { + mstore(0x00, 0xbed37c6e) // `InvalidNewLazyShufflerLength()`. + revert(0x1c, 0x04) + } sstore($.slot, or(shl(224, n), shr(32, shl(32, state)))) } } diff --git a/src/utils/g/LibPRNG.sol b/src/utils/g/LibPRNG.sol index 9f218862c..5407ba721 100644 --- a/src/utils/g/LibPRNG.sol +++ b/src/utils/g/LibPRNG.sol @@ -36,7 +36,8 @@ library LibPRNG { /// @dev The initial length must be greater than zero and less than `2**32 - 1`. error InvalidInitialLazyShufflerLength(); - /// @dev The new length must not be less than the current length. + /// @dev The new length must not be less than the current length, + /// and must be less than `2**32 - 1`. error InvalidNewLazyShufflerLength(); /// @dev The lazy shuffler has not been initialized. @@ -348,19 +349,22 @@ library LibPRNG { /// @dev Increases the length of `$`. /// Reverts if `$` has not been initialized. + /// Reverts if `n` is less than the current length, or if `n >= 2**32 - 1`. + /// Reverts if `n` crosses the entry width boundary at a length of 65535. function grow(LazyShuffler storage $, uint256 n) internal { /// @solidity memory-safe-assembly assembly { let state := sload($.slot) // The packed value at `$`. - // If the new length is smaller than the old length, revert. - if lt(n, shr(224, state)) { - mstore(0x00, 0xbed37c6e) // `InvalidNewLazyShufflerLength()`. - revert(0x1c, 0x04) - } if iszero(state) { mstore(0x00, 0x1ead2566) // `LazyShufflerNotInitialized()`. revert(0x1c, 0x04) } + let o := shr(224, state) // The old length. + let limit := or(0xfffe, mul(0xffff0000, gt(o, 0xfffe))) + if or(lt(n, o), gt(n, limit)) { + mstore(0x00, 0xbed37c6e) // `InvalidNewLazyShufflerLength()`. + revert(0x1c, 0x04) + } sstore($.slot, or(shl(224, n), shr(32, shl(32, state)))) } } diff --git a/test/LibPRNG.t.sol b/test/LibPRNG.t.sol index f737bb640..b71f7ebab 100644 --- a/test/LibPRNG.t.sol +++ b/test/LibPRNG.t.sol @@ -571,11 +571,14 @@ contract LibPRNGTest is SoladyTest { function testLazyShufflerRevertsOnGrowWithInvalidLength(uint256 n, uint256 nGrow) public { n = _bound(n, 1, 2 ** 32 - 2); this.lazyShufflerInitialize(n); - nGrow = _bound(n, 0, 2 ** 32 - 2); - if (nGrow < n) { + nGrow = _bound(nGrow, 0, 2 ** 32 - 2); + uint256 limit = n > 65534 ? 2 ** 32 - 2 : 65534; + bool reverts = nGrow < n || nGrow > limit; + if (reverts) { vm.expectRevert(LibPRNG.InvalidNewLazyShufflerLength.selector); } - this.lazyShufflerGrow(n); + this.lazyShufflerGrow(nGrow); + assertEq(_lazyShuffler0.length(), reverts ? n : nGrow); } function testLazyShufflerRevertsOnDoubleInit() public { @@ -620,4 +623,63 @@ contract LibPRNGTest is SoladyTest { function lazyShuffler1Get(uint256 i) public view returns (uint256) { return _lazyShuffler1.get(i); } + + function testLazyShufflerRevertsOnGrowAcrossWidthBoundary() public { + _lazyShuffler0.initialize(2); + _lazyShuffler0.next(0); + vm.expectRevert(LibPRNG.InvalidNewLazyShufflerLength.selector); + this.lazyShufflerGrow(65535); + } + + function testLazyShufflerRevertsOnGrowAcrossWidthBoundaryUndrawn() public { + _lazyShuffler0.initialize(2); + vm.expectRevert(LibPRNG.InvalidNewLazyShufflerLength.selector); + this.lazyShufflerGrow(65535); + } + + // `grow` had no upper bound check, so the length silently truncated to zero. + function testLazyShufflerRevertsOnGrowOutOfRange(uint256 n) public { + _lazyShuffler0.initialize(10); + n = _bound(n, 2 ** 32 - 1, type(uint256).max); + vm.expectRevert(LibPRNG.InvalidNewLazyShufflerLength.selector); + this.lazyShufflerGrow(n); + assertEq(_lazyShuffler0.length(), 10); + } + + function testLazyShufflerGrowWithinSameWidth() public { + _lazyShuffler0.initialize(2); + uint256 first = _lazyShuffler0.next(0); + _lazyShuffler0.grow(1000); + assertEq(_lazyShuffler0.get(0), first); + assertLt(_lazyShuffler0.get(0), 1000); + } + + function testLazyShufflerGrowWithinWideWidth() public { + _lazyShuffler0.initialize(70000); + uint256 first = _lazyShuffler0.next(0); + _lazyShuffler0.grow(200000); + assertEq(_lazyShuffler0.get(0), first); + assertLt(_lazyShuffler0.get(0), 200000); + } + + // A 32-bit shuffler still draws each value at most once across a grow. + function testLazyShufflerWideProducesNoDuplicatesAcrossGrow() public { + _lazyShuffler0.initialize(65535); + uint256[] memory seen = new uint256[](8); + unchecked { + for (uint256 i; i != 4; ++i) { + seen[i] = _lazyShuffler0.next(_random()); + } + _lazyShuffler0.grow(65600); + for (uint256 i = 4; i != 8; ++i) { + seen[i] = _lazyShuffler0.next(_random()); + } + LibSort.sort(seen); + LibSort.uniquifySorted(seen); + assertEq(seen.length, 8); + for (uint256 i; i != 8; ++i) { + assertLt(seen[i], 65600); + } + } + } }