From c6d27245c1f81d760f73b1c8cada4ccb2befc3ee Mon Sep 17 00:00:00 2001
From: kongwen686 <chaoqiang.tian01@gmail.com>
Date: Sun, 17 May 2026 01:15:44 +0800
Subject: [PATCH] Validate batch distribution sizes

---
 .../src/__tests__/batchDistribution.test.ts   | 94 +++++++++++++++++++
 packages/sdk/src/utils/batchDistribution.ts   | 16 ++++
 2 files changed, 110 insertions(+)
 create mode 100644 packages/sdk/src/__tests__/batchDistribution.test.ts

diff --git a/packages/sdk/src/__tests__/batchDistribution.test.ts b/packages/sdk/src/__tests__/batchDistribution.test.ts
new file mode 100644
index 0000000..60e617c
--- /dev/null
+++ b/packages/sdk/src/__tests__/batchDistribution.test.ts
@@ -0,0 +1,94 @@
+import { describe, expect, it, vi, beforeEach } from "vitest";
+import type { DistributorClient } from "../DistributorClient";
+import {
+  createBatches,
+  prepareBatchEqualDistribution,
+  prepareBatchWeightedDistribution,
+} from "../utils/batchDistribution";
+
+const SENDER = "GAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAWHF";
+const TOKEN = "CAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD2KM";
+const RECIPIENT_A = "GBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB";
+const RECIPIENT_B = "GCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC";
+const RECIPIENT_C = "GDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDD";
+
+function createMockClient(): DistributorClient {
+  return {
+    distributeEqual: vi.fn().mockResolvedValue({ signAndSend: vi.fn() }),
+    distributeWeighted: vi.fn().mockResolvedValue({ signAndSend: vi.fn() }),
+  } as unknown as DistributorClient;
+}
+
+describe("batchDistribution", () => {
+  let client: DistributorClient;
+
+  beforeEach(() => {
+    client = createMockClient();
+  });
+
+  describe("createBatches", () => {
+    it("splits arrays into fixed-size chunks", () => {
+      expect(createBatches([1, 2, 3], 2)).toEqual([[1, 2], [3]]);
+    });
+
+    it("rejects zero batch size", () => {
+      expect(() => createBatches([1, 2, 3], 0)).toThrow(
+        "batchSize must be a positive integer"
+      );
+    });
+
+    it("rejects negative batch size", () => {
+      expect(() => createBatches([1, 2, 3], -1)).toThrow(
+        "batchSize must be a positive integer"
+      );
+    });
+
+    it("rejects fractional batch size", () => {
+      expect(() => createBatches([1, 2, 3], 1.5)).toThrow(
+        "batchSize must be a positive integer"
+      );
+    });
+  });
+
+  describe("prepareBatchEqualDistribution", () => {
+    it.each([0, -1, 1.5])(
+      "rejects maxRecipientsPerBatch=%s before RPC calls",
+      async (maxRecipientsPerBatch) => {
+        await expect(
+          prepareBatchEqualDistribution(client, {
+            sender: SENDER,
+            token: TOKEN,
+            total_amount: 300n,
+            recipients: [RECIPIENT_A, RECIPIENT_B],
+            config: { maxRecipientsPerBatch },
+          })
+        ).rejects.toThrow(
+          "config.maxRecipientsPerBatch must be a positive integer"
+        );
+
+        expect(client.distributeEqual).not.toHaveBeenCalled();
+      }
+    );
+  });
+
+  describe("prepareBatchWeightedDistribution", () => {
+    it.each([0, -1, 1.5])(
+      "rejects maxRecipientsPerBatch=%s before RPC calls",
+      async (maxRecipientsPerBatch) => {
+        await expect(
+          prepareBatchWeightedDistribution(client, {
+            sender: SENDER,
+            token: TOKEN,
+            recipients: [RECIPIENT_A, RECIPIENT_B, RECIPIENT_C],
+            amounts: [100n, 200n, 300n],
+            config: { maxRecipientsPerBatch },
+          })
+        ).rejects.toThrow(
+          "config.maxRecipientsPerBatch must be a positive integer"
+        );
+
+        expect(client.distributeWeighted).not.toHaveBeenCalled();
+      }
+    );
+  });
+});
diff --git a/packages/sdk/src/utils/batchDistribution.ts b/packages/sdk/src/utils/batchDistribution.ts
index 3805d77..efb0a88 100644
--- a/packages/sdk/src/utils/batchDistribution.ts
+++ b/packages/sdk/src/utils/batchDistribution.ts
@@ -153,6 +153,12 @@ export interface BatchDistributionResult {
   amountBatches?: bigint[][];
 }
 
+function assertPositiveInteger(value: number, name: string): void {
+  if (!Number.isInteger(value) || value <= 0) {
+    throw new Error(`${name} must be a positive integer (got ${value})`);
+  }
+}
+
 /**
  * Splits recipients into batches and creates assembled transactions for equal distribution.
  * 
@@ -197,6 +203,10 @@ export async function prepareBatchEqualDistribution(
 ): Promise<BatchDistributionResult> {
   const { sender, token, total_amount, recipients, config = {} } = params;
   const maxRecipientsPerBatch = config.maxRecipientsPerBatch ?? 100;
+  assertPositiveInteger(
+    maxRecipientsPerBatch,
+    'config.maxRecipientsPerBatch'
+  );
 
   if (recipients.length === 0) {
     throw new Error('Recipients array cannot be empty');
@@ -280,6 +290,10 @@ export async function prepareBatchWeightedDistribution(
 ): Promise<BatchDistributionResult> {
   const { sender, token, recipients, amounts, config = {} } = params;
   const maxRecipientsPerBatch = config.maxRecipientsPerBatch ?? 100;
+  assertPositiveInteger(
+    maxRecipientsPerBatch,
+    'config.maxRecipientsPerBatch'
+  );
 
   if (recipients.length === 0) {
     throw new Error('Recipients array cannot be empty');
@@ -341,6 +355,8 @@ export async function prepareBatchWeightedDistribution(
  * \`\`\`
  */
 export function createBatches<T>(array: T[], batchSize: number): T[][] {
+  assertPositiveInteger(batchSize, 'batchSize');
+
   const batches: T[][] = [];
   
   for (let i = 0; i < array.length; i += batchSize) {