From c2f7b566c212aaa60f6189d8e9636a8b8c50a182 Mon Sep 17 00:00:00 2001
From: atarpara <akpatel0618@gmail.com>
Date: Sat, 29 Aug 2026 14:50:37 +0530
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9E=20Fix=20LazyShuffler.grow=20length?=
 =?UTF-8?q?=20validation?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/utils/LibPRNG.sol   | 13 +++++----
 src/utils/g/LibPRNG.sol | 13 +++++----
 test/LibPRNG.t.sol      | 59 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 10 deletions(-)

diff --git a/src/utils/LibPRNG.sol b/src/utils/LibPRNG.sol
index 4b49a69e43..8cc995a901 100644
--- a/src/utils/LibPRNG.sol
+++ b/src/utils/LibPRNG.sol
@@ -343,19 +343,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 9f218862cd..22a7846ecf 100644
--- a/src/utils/g/LibPRNG.sol
+++ b/src/utils/g/LibPRNG.sol
@@ -348,19 +348,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 f737bb640a..2a20d8ac30 100644
--- a/test/LibPRNG.t.sol
+++ b/test/LibPRNG.t.sol
@@ -620,4 +620,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);
+            }
+        }
+    }
 }