From 52bfe3f782e5d64c8141c37f0862474c658e86e8 Mon Sep 17 00:00:00 2001
From: jr-kenny <dalvidjr2022@gmail.com>
Date: Thu, 11 Jun 2026 12:29:18 +0100
Subject: [PATCH] fix(spammer): reject too many exponential generators instead
 of overflowing

The guard in partition_exponential that should reject more generators than
the account space can hold computes 1 << (num_generators - 1). Once that
exponent reaches the word size the shift overflows, so the guard panics in
debug builds and silently passes (with wrapped denominators) in release,
producing garbage account ranges. Short-circuit on the shift width so the
bail fires cleanly, and add a regression test.
---
 crates/spammer/src/accounts.rs | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/crates/spammer/src/accounts.rs b/crates/spammer/src/accounts.rs
index 58985de..ea5dde2 100644
--- a/crates/spammer/src/accounts.rs
+++ b/crates/spammer/src/accounts.rs
@@ -100,7 +100,13 @@ impl PartitionMode {
             (n + d / 2) / d
         }
 
-        if round_div(num_accounts, 1 << (num_generators - 1)) == 0 {
+        // The smallest bucket spans `num_accounts / 2^(num_generators - 1)`. Once the
+        // exponent reaches the word size, `2^(num_generators - 1)` already exceeds any
+        // possible `num_accounts`, so the smallest bucket is empty. Check that first:
+        // shifting by >= the word size is undefined and would otherwise panic (debug) or
+        // wrap the shift distance (release) before this guard could reject the input.
+        let shift = num_generators - 1;
+        if shift >= usize::BITS as usize || round_div(num_accounts, 1 << shift) == 0 {
             eyre::bail!("too many generators: it would result in a bucket with size 0");
         }
 
@@ -167,6 +173,24 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn partition_accounts_exponential_rejects_too_many_generators() -> Result<()> {
+        // num_generators large enough that 2^(num_generators - 1) reaches the word
+        // size. These inputs pass the num_generators>0 and num_accounts>=num_generators
+        // checks, so they reach partition_exponential, where the smallest bucket would
+        // round to 0. Expect the same graceful "too many generators" error as smaller
+        // over-subscriptions (e.g. (100, 9)) rather than a shift-overflow panic.
+        for (num_accounts, num_generators) in [(130usize, 65usize), (200, 100)] {
+            let result =
+                PartitionMode::Exponential.partition_accounts(num_accounts, num_generators);
+            assert!(
+                result.is_err(),
+                "expected graceful error for ({num_accounts}, {num_generators}), got {result:?}"
+            );
+        }
+        Ok(())
+    }
+
     #[tokio::test]
     async fn partition_accounts_exponential() -> Result<()> {
         #[rustfmt::skip]