From a589c6d49265d24cf202e8f68c6b2ad4fafd398c Mon Sep 17 00:00:00 2001
From: didinele <didinele.dev@gmail.com>
Date: Mon, 30 Dec 2024 20:49:21 +0200
Subject: [PATCH] fix(SimpleIdentifyThrottler): don't sleep negative amounts

---
 .../util/SimpleIdentifyThrottler.test.ts      | 47 ++++++++++++++-----
 .../src/throttling/SimpleIdentifyThrottler.ts |  2 +-
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/packages/ws/__tests__/util/SimpleIdentifyThrottler.test.ts b/packages/ws/__tests__/util/SimpleIdentifyThrottler.test.ts
index 536625a04..188ffdb1e 100644
--- a/packages/ws/__tests__/util/SimpleIdentifyThrottler.test.ts
+++ b/packages/ws/__tests__/util/SimpleIdentifyThrottler.test.ts
@@ -1,33 +1,58 @@
-// @ts-nocheck
-import { setTimeout as sleep } from 'node:timers/promises';
-import { expect, test, vi } from 'vitest';
+import * as timers from 'node:timers/promises';
+import { expect, test, vi, beforeEach, afterEach } from 'vitest';
 import { SimpleIdentifyThrottler } from '../../src/index.js';
 
 vi.mock('node:timers/promises', () => ({
 	setTimeout: vi.fn(),
 }));
 
-const throttler = new SimpleIdentifyThrottler(2);
+let throttler: SimpleIdentifyThrottler;
+const controller = new AbortController();
 
 vi.useFakeTimers();
 
-const NOW = vi.fn().mockReturnValue(Date.now());
+const TIME = Date.now();
+const NOW = vi.fn().mockReturnValue(TIME);
 global.Date.now = NOW;
 
+const sleep = vi.spyOn(timers, 'setTimeout');
+
+beforeEach(() => {
+	throttler = new SimpleIdentifyThrottler(2);
+});
+
+afterEach(() => {
+	sleep.mockClear();
+});
+
 test('basic case', async () => {
 	// Those shouldn't wait since they're in different keys
-
-	await throttler.waitForIdentify(0);
+	await throttler.waitForIdentify(0, controller.signal);
 	expect(sleep).not.toHaveBeenCalled();
 
-	await throttler.waitForIdentify(1);
+	await throttler.waitForIdentify(1, controller.signal);
 	expect(sleep).not.toHaveBeenCalled();
 
 	// Those should wait
-
-	await throttler.waitForIdentify(2);
+	await throttler.waitForIdentify(2, controller.signal);
 	expect(sleep).toHaveBeenCalledTimes(1);
 
-	await throttler.waitForIdentify(3);
+	await throttler.waitForIdentify(3, controller.signal);
 	expect(sleep).toHaveBeenCalledTimes(2);
 });
+
+test('does not call sleep with a negative time', async () => {
+	await throttler.waitForIdentify(0, controller.signal);
+	expect(sleep).not.toHaveBeenCalled();
+
+	await throttler.waitForIdentify(0, controller.signal);
+	expect(sleep).toHaveBeenCalledTimes(1);
+
+	// By overshooting, we're simulating a bug that existed prior to this test, where-in by enough time
+	// passing before the shard tried to identify for a subsequent time, the passed value would end up being negative
+	// (and this was unchecked). Node simply treats that as 1ms, so it wasn't particularly harmful, but they
+	// recently introduced a warning for it, so we want to avoid that.
+	NOW.mockReturnValueOnce(TIME + 10_000);
+	await throttler.waitForIdentify(0, controller.signal);
+	expect(sleep).toHaveBeenCalledTimes(1);
+});
diff --git a/packages/ws/src/throttling/SimpleIdentifyThrottler.ts b/packages/ws/src/throttling/SimpleIdentifyThrottler.ts
index bacf3fb07..a75cb88fb 100644
--- a/packages/ws/src/throttling/SimpleIdentifyThrottler.ts
+++ b/packages/ws/src/throttling/SimpleIdentifyThrottler.ts
@@ -36,7 +36,7 @@ export class SimpleIdentifyThrottler implements IIdentifyThrottler {
 
 		try {
 			const diff = state.resetsAt - Date.now();
-			if (diff <= 5_000) {
+			if (diff > 0 && diff <= 5_000) {
 				// To account for the latency the IDENTIFY payload goes through, we add a bit more wait time
 				const time = diff + Math.random() * 1_500;
 				await sleep(time);