From a790fcaf049884ff700bcd313b3c8186b36f7a22 Mon Sep 17 00:00:00 2001
From: tiagosiebler <tiagosiebler@github.com>
Date: Mon, 30 Oct 2023 13:47:40 +0000
Subject: [PATCH] feat(#251): add optional bapi rate limit parsing to REST
 clients

---
 src/types/request/v5-market.ts |  4 +--
 src/types/shared.ts            | 27 ++++++++++++----
 src/util/BaseRestClient.ts     | 13 +++++++-
 src/util/requestUtils.ts       | 58 ++++++++++++++++++++++++++++++++++
 4 files changed, 93 insertions(+), 9 deletions(-)

diff --git a/src/types/request/v5-market.ts b/src/types/request/v5-market.ts
index 0fe9f0c..9e836b7 100644
--- a/src/types/request/v5-market.ts
+++ b/src/types/request/v5-market.ts
@@ -52,8 +52,8 @@ export interface GetOrderbookParamsV5 {
   limit?: number;
 }
 
-export interface GetTickersParamsV5 {
-  category: CategoryV5;
+export interface GetTickersParamsV5<TCategory = CategoryV5> {
+  category: TCategory;
   symbol?: string;
   baseCoin?: string;
   expDate?: string;
diff --git a/src/types/shared.ts b/src/types/shared.ts
index 3e13c34..9be7676 100644
--- a/src/types/shared.ts
+++ b/src/types/shared.ts
@@ -61,18 +61,33 @@ export interface APIResponse<T> {
   result: T;
 }
 
+export interface APIRateLimit {
+  /** Remaining requests to this endpoint before the next reset */
+  remainingRequests: number;
+  /** Max requests for this endpoint per rollowing window (before next reset) */
+  maxRequests: number;
+  /**
+   * Timestamp when the rate limit resets if you have exceeded your current maxRequests.
+   * Otherwise, this is approximately your current timestamp.
+   */
+  resetAtTimestamp: number;
+}
+
 export interface APIResponseV3<T> {
   retCode: number;
   retMsg: 'OK' | string;
   result: T;
+  /**
+   * These are per-UID per-endpoint rate limits, automatically parsed from response headers if available.
+   *
+   * Note:
+   * - this is primarily for V5 (or newer) APIs.
+   * - these rate limits are per-endpoint per-account, so will not appear for public API calls
+   */
+  rateLimitApi?: APIRateLimit;
 }
 
-export interface APIResponseV3WithTime<T> {
-  retCode: number;
-  retMsg: 'OK' | string;
-  result: T;
-  time: number;
-}
+export type APIResponseV3WithTime<T> = APIResponseV3<T> & { time: number };
 
 export interface APIResponseWithTime<T = {}> extends APIResponse<T> {
   /** UTC timestamp */
diff --git a/src/util/BaseRestClient.ts b/src/util/BaseRestClient.ts
index fdbb3a8..d7682fd 100644
--- a/src/util/BaseRestClient.ts
+++ b/src/util/BaseRestClient.ts
@@ -7,6 +7,7 @@ import {
   RestClientOptions,
   RestClientType,
   getRestBaseUrl,
+  parseRateLimitHeaders,
   serializeParams,
 } from './requestUtils';
 import { signMessage } from './node-support';
@@ -323,7 +324,17 @@ export default abstract class BaseRestClient {
     return axios(options)
       .then((response) => {
         if (response.status == 200) {
-          return response.data;
+          const perAPIRateLimits = this.options.parseAPIRateLimits
+            ? parseRateLimitHeaders(
+                response.headers,
+                this.options.throwOnFailedRateLimitParse === true,
+              )
+            : undefined;
+
+          return {
+            rateLimitApi: perAPIRateLimits,
+            ...response.data,
+          };
         }
 
         throw response;
diff --git a/src/util/requestUtils.ts b/src/util/requestUtils.ts
index 8699ab6..3297ea1 100644
--- a/src/util/requestUtils.ts
+++ b/src/util/requestUtils.ts
@@ -1,3 +1,4 @@
+import { APIRateLimit } from '../types';
 import { WebsocketSucceededTopicSubscriptionConfirmationEvent } from '../types/ws-events/succeeded-topic-subscription-confirmation';
 import { WebsocketTopicSubscriptionConfirmationEvent } from '../types/ws-events/topic-subscription-confirmation';
 
@@ -46,6 +47,12 @@ export interface RestClientOptions {
 
   /** Default: true. whether to try and post-process request exceptions. */
   parse_exceptions?: boolean;
+
+  /** Default: false. Enable to parse/include per-API/endpoint rate limits in responses. */
+  parseAPIRateLimits?: boolean;
+
+  /** Default: false. Enable to throw error if rate limit parser fails */
+  throwOnFailedRateLimitParse?: boolean;
 }
 
 /**
@@ -169,3 +176,54 @@ export const REST_CLIENT_TYPE_ENUM = {
 
 export type RestClientType =
   (typeof REST_CLIENT_TYPE_ENUM)[keyof typeof REST_CLIENT_TYPE_ENUM];
+
+/** Parse V5 rate limit response headers, if enabled */
+export function parseRateLimitHeaders(
+  headers: Record<string, string | undefined> = {},
+  throwOnFailedRateLimitParse: boolean,
+): APIRateLimit | undefined {
+  try {
+    const remaining = headers['x-bapi-limit-status'];
+    const max = headers['x-bapi-limit'];
+    const resetAt = headers['x-bapi-limit-reset-timestamp'];
+
+    if (
+      typeof remaining === 'undefined' ||
+      typeof max === 'undefined' ||
+      typeof resetAt === 'undefined'
+    ) {
+      return;
+    }
+
+    const result: APIRateLimit = {
+      remainingRequests: Number(remaining),
+      maxRequests: Number(max),
+      resetAtTimestamp: Number(resetAt),
+    };
+
+    if (
+      isNaN(result.remainingRequests) ||
+      isNaN(result.maxRequests) ||
+      isNaN(result.resetAtTimestamp)
+    ) {
+      return;
+    }
+
+    return result;
+  } catch (e) {
+    if (throwOnFailedRateLimitParse) {
+      console.log(
+        new Date(),
+        'parseRateLimitHeaders()',
+        'Failed to parse rate limit headers',
+        {
+          headers,
+          exception: e,
+        },
+      );
+      throw e;
+    }
+  }
+
+  return undefined;
+}