From 85d65380177661865b8ba5ea7dc14b35d06e4e03 Mon Sep 17 00:00:00 2001
From: Donovan Daniels <hewwo@yiff.rocks>
Date: Sat, 17 Aug 2024 19:35:05 -0500
Subject: [PATCH] Add option to limit ratelimit retry window

---
 lib/gateway/Shard.ts                 | 2 +-
 lib/gateway/compression/pako.ts      | 2 +-
 lib/gateway/compression/zlib-sync.ts | 2 +-
 lib/rest/RequestHandler.ts           | 5 +++++
 lib/types/client.d.ts                | 6 ++++++
 lib/util/Errors.ts                   | 9 ++++++++-
 6 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/lib/gateway/Shard.ts b/lib/gateway/Shard.ts
index 04e35ade..5b4289c5 100644
--- a/lib/gateway/Shard.ts
+++ b/lib/gateway/Shard.ts
@@ -21,7 +21,7 @@ import type { RawGuild } from "../types/guilds";
 import ExtendedUser from "../structures/ExtendedUser";
 import type Guild from "../structures/Guild";
 import type { ShardEvents } from "../types/events";
-import GatewayError, { DependencyError } from "../util/Errors";
+import { GatewayError, DependencyError } from "../util/Errors";
 import ClientApplication from "../structures/ClientApplication";
 import WebSocket, { type Data } from "ws";
 import { randomBytes } from "node:crypto";
diff --git a/lib/gateway/compression/pako.ts b/lib/gateway/compression/pako.ts
index f1f03dd8..c289bcd5 100644
--- a/lib/gateway/compression/pako.ts
+++ b/lib/gateway/compression/pako.ts
@@ -1,6 +1,6 @@
 import Compression from "./base";
 import type Shard from "../Shard";
-import GatewayError from "../../util/Errors";
+import { GatewayError } from "../../util/Errors";
 import { Inflate, constants } from "pako";
 
 interface PakoExtra {
diff --git a/lib/gateway/compression/zlib-sync.ts b/lib/gateway/compression/zlib-sync.ts
index 008b7f66..33d44caa 100644
--- a/lib/gateway/compression/zlib-sync.ts
+++ b/lib/gateway/compression/zlib-sync.ts
@@ -1,6 +1,6 @@
 import Compression from "./base";
 import type Shard from "../Shard";
-import GatewayError from "../../util/Errors";
+import { GatewayError } from "../../util/Errors";
 import ZlibSync from "zlib-sync";
 
 export default class ZlibSyncCompression extends Compression {
diff --git a/lib/rest/RequestHandler.ts b/lib/rest/RequestHandler.ts
index 3b6644b5..fc24cce4 100644
--- a/lib/rest/RequestHandler.ts
+++ b/lib/rest/RequestHandler.ts
@@ -7,6 +7,7 @@ import { API_URL, RESTMethods, USER_AGENT, type RESTMethod } from "../Constants"
 import Base from "../structures/Base";
 import type { LatencyRef, RequestHandlerInstanceOptions, RequestOptions } from "../types/request-handler";
 import type { RESTOptions } from "../types/client";
+import { RateLimitedError } from "../util/Errors";
 
 /**
  * Latency & ratelimit related things lovingly borrowed from eris
@@ -33,6 +34,7 @@ export default class RequestHandler {
             followRedirects:            !!options.followRedirects,
             host:                       options.host ?? (options.baseURL ? new URL(options.baseURL).host : new URL(API_URL).host),
             latencyThreshold:           options.latencyThreshold ?? 30000,
+            maxRatelimitRetryWindow:    options.maxRatelimitRetryWindow ?? Infinity,
             ratelimiterOffset:          options.ratelimiterOffset ?? 0,
             requestTimeout:             options.requestTimeout ?? 15000,
             superProperties:            options.superProperties ?? null,
@@ -262,6 +264,9 @@ export default class RequestHandler {
 
                             this._manager.client.emit("debug", `${res.headers.has("x-ratelimit-global") ? "Global" : "Unexpected"} RateLimit: ${JSON.stringify(resBody)}\n${now} ${route} ${res.status}: ${latency}ms (${this.latencyRef.latency}ms avg) | ${this.ratelimits[route].remaining}/${this.ratelimits[route].limit} left | Reset ${delay} (${this.ratelimits[route].reset - now}ms left) | Scope ${res.headers.get("x-ratelimit-scope")!}`);
                             if (delay) {
+                                if (delay > this.options.maxRatelimitRetryWindow) {
+                                    throw new RateLimitedError(`Ratelimit on "${options.method} ${route}" exceeds the maximum retry window (${delay} > ${this.options.maxRatelimitRetryWindow})`);
+                                }
                                 setTimeout(() => {
                                     cb();
                                     // eslint-disable-next-line prefer-rest-params, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call, prefer-spread
diff --git a/lib/types/client.d.ts b/lib/types/client.d.ts
index 098e0393..ec328376 100644
--- a/lib/types/client.d.ts
+++ b/lib/types/client.d.ts
@@ -77,6 +77,12 @@ export interface RESTOptions {
      * @defaultValue 30000
      */
     latencyThreshold?: number;
+    /**
+     * In milliseconds, the maximum ratelimit delay (in seconds) the lib will internally wait for to retry the request. If a ratelimit resets after this window, an error will be thrown instead.
+     * @note This currently defaults to Infinity for backwards compatibility, but this will be changed in 1.12.0.
+     * @defaultValue Infinity
+     */
+    maxRatelimitRetryWindow?: number;
     /**
      * In milliseconds, the time to offset ratelimit calculations by.
      * @defaultValue 0
diff --git a/lib/util/Errors.ts b/lib/util/Errors.ts
index 7b90de4c..7485c1d9 100644
--- a/lib/util/Errors.ts
+++ b/lib/util/Errors.ts
@@ -47,7 +47,7 @@ export class DependencyError extends Error {
 
 
 /** A gateway error. */
-export default class GatewayError extends Error {
+export class GatewayError extends Error {
     code: number;
     override name = "GatewayError";
     constructor(message: string, code: number) {
@@ -55,3 +55,10 @@ export default class GatewayError extends Error {
         this.code = code;
     }
 }
+
+export class RateLimitedError extends Error {
+    override name = "RateLimitedError";
+    constructor(message: string) {
+        super(message);
+    }
+}