Merge pull request #5855 from BitGo/ME-178

feat(sdk-api): add dynamic headers in sdk requests

Author: jakeyangbitgo

modules/bitgo/test/v2/integration/bitgo.ts

@@ -634,4 +634,48 @@ describe('BitGo', function () {
       return bitgo.authenticateKnownBalanceTestUser(bitgo.testUserOTP());
     });
   });
+
+  describe('getAdditionalHeadersCb', function () {
+    let bitgo;
+    let additionalHeadersCb;
+
+    before(function () {
+      additionalHeadersCb = sinon.spy((method, url, data) => {
+        return [
+          { key: 'X-Custom-Header', value: 'CustomValue' },
+          { key: 'X-Another-Header', value: 'AnotherValue' },
+        ];
+      });
+
+      bitgo = new TestBitGo();
+      bitgo.initializeTestVars();
+      bitgo = new BitGoJS.BitGo({ getAdditionalHeadersCb: additionalHeadersCb });
+    });
+
+    it('should invoke getAdditionalHeadersCb and apply headers to the request', async function () {
+      const url = bitgo.url('/test-endpoint');
+      const method = 'get';
+
+      // Mock the request to avoid actual network calls
+      const mockRequest = sinon.stub(bitgo, 'getAgentRequest').returns({
+        set: sinon.stub().returnsThis(),
+        then: sinon.stub().resolves({}),
+      });
+
+      await bitgo.get(url);
+
+      // Verify that the callback was called with the correct arguments
+      sinon.assert.calledOnce(additionalHeadersCb);
+      sinon.assert.calledWith(additionalHeadersCb, method, url, undefined);
+
+      // Verify that the headers were applied to the request
+      const headers = additionalHeadersCb.returnValues[0];
+      headers.forEach(({ key, value }) => {
+        sinon.assert.calledWith(mockRequest().set, key, value);
+      });
+
+      // Restore the stubbed method
+      mockRequest.restore();
+    });
+  });
 });

modules/sdk-api/src/bitgoAPI.ts

@@ -44,6 +44,7 @@ import {
   AccessTokenOptions,
   AddAccessTokenOptions,
   AddAccessTokenResponse,
+  AdditionalHeadersCallback,
   AuthenticateOptions,
   AuthenticateWithAuthCodeOptions,
   BitGoAPIOptions,
@@ -67,6 +68,7 @@ import {
   RegisterPushTokenOptions,
   RemoveAccessTokenOptions,
   RequestHeaders,
+  RequestMethods,
   SplitSecret,
   SplitSecretOptions,
   TokenIssuance,
@@ -94,8 +96,6 @@ const PendingApprovals = require('./v1/pendingapprovals');
 const TravelRule = require('./v1/travelRule');
 const TransactionBuilder = require('./v1/transactionBuilder');
 
-const patchedRequestMethods = ['get', 'post', 'put', 'del', 'patch', 'options'] as const;
-
 export class BitGoAPI implements BitGoBase {
   // v1 types
   protected _keychains: any;
@@ -130,8 +130,10 @@ export class BitGoAPI implements BitGoBase {
   protected _validate: boolean;
   public readonly cookiesPropagationEnabled: boolean;
   private _customProxyAgent?: Agent;
+  private getAdditionalHeadersCb?: AdditionalHeadersCallback;
 
   constructor(params: BitGoAPIOptions = {}) {
+    this.getAdditionalHeadersCb = params.getAdditionalHeadersCb;
     this.cookiesPropagationEnabled = false;
     if (
       !common.validateParams(
@@ -303,7 +305,7 @@ export class BitGoAPI implements BitGoBase {
    * @param method - http method for the new request
    * @param url - URL for the new request
    */
-  protected getAgentRequest(method: (typeof patchedRequestMethods)[number], url: string): superagent.SuperAgentRequest {
+  protected getAgentRequest(method: RequestMethods, url: string): superagent.SuperAgentRequest {
     let req: superagent.SuperAgentRequest = superagent[method](url);
     if (this.cookiesPropagationEnabled) {
       req = req.withCredentials();
@@ -337,7 +339,7 @@ export class BitGoAPI implements BitGoBase {
    * headers to any outbound request.
    * @param method
    */
-  private requestPatch(method: (typeof patchedRequestMethods)[number], url: string) {
+  private requestPatch(method: RequestMethods, url: string) {
     const req = this.getAgentRequest(method, url);
     if (this._customProxyAgent) {
       debug('using custom proxy agent');
@@ -394,8 +396,8 @@ export class BitGoAPI implements BitGoBase {
 
       req.set('BitGo-Auth-Version', this._authVersion === 3 ? '3.0' : '2.0');
 
+      const data = serializeRequestData(req);
       if (this._token) {
-        const data = serializeRequestData(req);
         setRequestQueryString(req);
 
         const requestProperties = this.calculateRequestHeaders({
@@ -415,6 +417,13 @@ export class BitGoAPI implements BitGoBase {
         req.set('HMAC', requestProperties.hmac);
       }
 
+      if (this.getAdditionalHeadersCb) {
+        const additionalHeaders = this.getAdditionalHeadersCb(method, url, data);
+        for (const { key, value } of additionalHeaders) {
+          req.set(key, value);
+        }
+      }
+
       /**
        * Verify the response before calling the original onfulfilled handler,
        * and make sure onrejected is called if a verification error is encountered

modules/sdk-api/src/types.ts

@@ -1,6 +1,15 @@
 import { EnvironmentName, IRequestTracer, V1Network } from '@bitgo/sdk-core';
 import { ECPairInterface } from '@bitgo/utxo-lib';
 import { type Agent } from 'http';
+
+const patchedRequestMethods = ['get', 'post', 'put', 'del', 'patch', 'options'] as const;
+export type RequestMethods = (typeof patchedRequestMethods)[number];
+export type AdditionalHeadersCallback = (
+  method: RequestMethods,
+  url: string,
+  data?: string
+) => Array<{ key: string; value: string }>;
+
 export {
   supportedRequestMethods,
   AuthVersion,
@@ -36,6 +45,7 @@ export interface BitGoAPIOptions {
   userAgent?: string;
   validate?: boolean;
   cookiesPropagationEnabled?: boolean;
+  getAdditionalHeadersCb?: AdditionalHeadersCallback;
 }
 
 export interface AccessTokenOptions {