Sitelet https://github.com/Vectorized/solady/commit/c2f7b566c212aaa60f6189d8e9636a8b8c50a182
Skip to content

Commit c2f7b56

Browse files
atarparaatarpara
committed
🐞 Fix LazyShuffler.grow length validation
1 parent 4225d5f commit c2f7b56

3 files changed

Lines changed: 75 additions & 10 deletions

File tree

‎src/utils/LibPRNG.sol‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -343,19 +343,22 @@ library LibPRNG {
343343

344344
/// @dev Increases the length of `$`.
345345
/// Reverts if `$` has not been initialized.
346+
/// Reverts if `n` is less than the current length, or if `n >= 2**32 - 1`.
347+
/// Reverts if `n` crosses the entry width boundary at a length of 65535.
346348
function grow(LazyShuffler storage $, uint256 n) internal {
347349
/// @solidity memory-safe-assembly
348350
assembly {
349351
let state := sload($.slot) // The packed value at `$`.
350-
// If the new length is smaller than the old length, revert.
351-
if lt(n, shr(224, state)) {
352-
mstore(0x00, 0xbed37c6e) // `InvalidNewLazyShufflerLength()`.
353-
revert(0x1c, 0x04)
354-
}
355352
if iszero(state) {
356353
mstore(0x00, 0x1ead2566) // `LazyShufflerNotInitialized()`.
357354
revert(0x1c, 0x04)
358355
}
356+
let o := shr(224, state) // The old length.
357+
let limit := or(0xfffe, mul(0xffff0000, gt(o, 0xfffe)))
358+
if or(lt(n, o), gt(n, limit)) {
359+
mstore(0x00, 0xbed37c6e) // `InvalidNewLazyShufflerLength()`.
360+
revert(0x1c, 0x04)
361+
}
359362
sstore($.slot, or(shl(224, n), shr(32, shl(32, state))))
360363
}
361364
}

‎src/utils/g/LibPRNG.sol‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -348,19 +348,22 @@ library LibPRNG {
348348

349349
/// @dev Increases the length of `$`.
350350
/// Reverts if `$` has not been initialized.
351+
/// Reverts if `n` is less than the current length, or if `n >= 2**32 - 1`.
352+
/// Reverts if `n` crosses the entry width boundary at a length of 65535.
351353
function grow(LazyShuffler storage $, uint256 n) internal {
352354
/// @solidity memory-safe-assembly
353355
assembly {
354356
let state := sload($.slot) // The packed value at `$`.
355-
// If the new length is smaller than the old length, revert.
356-
if lt(n, shr(224, state)) {
357-
mstore(0x00, 0xbed37c6e) // `InvalidNewLazyShufflerLength()`.
358-
revert(0x1c, 0x04)
359-
}
360357
if iszero(state) {
361358
mstore(0x00, 0x1ead2566) // `LazyShufflerNotInitialized()`.
362359
revert(0x1c, 0x04)
363360
}
361+
let o := shr(224, state) // The old length.
362+
let limit := or(0xfffe, mul(0xffff0000, gt(o, 0xfffe)))
363+
if or(lt(n, o), gt(n, limit)) {
364+
mstore(0x00, 0xbed37c6e) // `InvalidNewLazyShufflerLength()`.
365+
revert(0x1c, 0x04)
366+
}
364367
sstore($.slot, or(shl(224, n), shr(32, shl(32, state))))
365368
}
366369
}

‎test/LibPRNG.t.sol‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -620,4 +620,63 @@ contract LibPRNGTest is SoladyTest {
620620
function lazyShuffler1Get(uint256 i) public view returns (uint256) {
621621
return _lazyShuffler1.get(i);
622622
}
623+
624+
function testLazyShufflerRevertsOnGrowAcrossWidthBoundary() public {
625+
_lazyShuffler0.initialize(2);
626+
_lazyShuffler0.next(0);
627+
vm.expectRevert(LibPRNG.InvalidNewLazyShufflerLength.selector);
628+
this.lazyShufflerGrow(65535);
629+
}
630+
631+
function testLazyShufflerRevertsOnGrowAcrossWidthBoundaryUndrawn() public {
632+
_lazyShuffler0.initialize(2);
633+
vm.expectRevert(LibPRNG.InvalidNewLazyShufflerLength.selector);
634+
this.lazyShufflerGrow(65535);
635+
}
636+
637+
// `grow` had no upper bound check, so the length silently truncated to zero.
638+
function testLazyShufflerRevertsOnGrowOutOfRange(uint256 n) public {
639+
_lazyShuffler0.initialize(10);
640+
n = _bound(n, 2 ** 32 - 1, type(uint256).max);
641+
vm.expectRevert(LibPRNG.InvalidNewLazyShufflerLength.selector);
642+
this.lazyShufflerGrow(n);
643+
assertEq(_lazyShuffler0.length(), 10);
644+
}
645+
646+
function testLazyShufflerGrowWithinSameWidth() public {
647+
_lazyShuffler0.initialize(2);
648+
uint256 first = _lazyShuffler0.next(0);
649+
_lazyShuffler0.grow(1000);
650+
assertEq(_lazyShuffler0.get(0), first);
651+
assertLt(_lazyShuffler0.get(0), 1000);
652+
}
653+
654+
function testLazyShufflerGrowWithinWideWidth() public {
655+
_lazyShuffler0.initialize(70000);
656+
uint256 first = _lazyShuffler0.next(0);
657+
_lazyShuffler0.grow(200000);
658+
assertEq(_lazyShuffler0.get(0), first);
659+
assertLt(_lazyShuffler0.get(0), 200000);
660+
}
661+
662+
// A 32-bit shuffler still draws each value at most once across a grow.
663+
function testLazyShufflerWideProducesNoDuplicatesAcrossGrow() public {
664+
_lazyShuffler0.initialize(65535);
665+
uint256[] memory seen = new uint256[](8);
666+
unchecked {
667+
for (uint256 i; i != 4; ++i) {
668+
seen[i] = _lazyShuffler0.next(_random());
669+
}
670+
_lazyShuffler0.grow(65600);
671+
for (uint256 i = 4; i != 8; ++i) {
672+
seen[i] = _lazyShuffler0.next(_random());
673+
}
674+
LibSort.sort(seen);
675+
LibSort.uniquifySorted(seen);
676+
assertEq(seen.length, 8);
677+
for (uint256 i; i != 8; ++i) {
678+
assertLt(seen[i], 65600);
679+
}
680+
}
681+
}
623682
}

0 commit comments

Comments
 (0)