[OTE-130] implement new geo-blocking strategy for socks

indexer/services/socks/__tests__/lib/subscriptions.test.ts

@@ -10,7 +10,6 @@ import { btcTicker, invalidChannel, invalidTicker } from '../constants';
 import { axiosRequest } from '../../src/lib/axios';
 import { AxiosSafeServerError, makeAxiosSafeServerError } from '@dydxprotocol-indexer/base';
 import { BlockedError } from '../../src/lib/errors';
-import { isRestrictedCountry } from '@dydxprotocol-indexer/compliance';
 
 jest.mock('ws');
 jest.mock('../../src/helpers/wss');
@@ -83,9 +82,6 @@ describe('Subscriptions', () => {
     axiosRequestMock = (axiosRequest as jest.Mock);
     axiosRequestMock.mockClear();
     axiosRequestMock.mockImplementation(() => (JSON.stringify(initialMessage)));
-    (isRestrictedCountry as jest.Mock).mockImplementation((country: string): boolean => {
-      return country === restrictedCountry;
-    });
   });
 
   describe('subscribe', () => {
@@ -126,6 +122,9 @@ describe('Subscriptions', () => {
         for (const urlPattern of urlPatterns) {
           expect(axiosRequestMock).toHaveBeenCalledWith(expect.objectContaining({
             url: expect.stringMatching(RegExp(urlPattern)),
+            headers: {
+              'cf-ipcountry': nonRestrictedCountry,
+            },
           }));
         }
       } else {
@@ -150,7 +149,6 @@ describe('Subscriptions', () => {
           initialMsgId,
           id,
           false,
-          nonRestrictedCountry,
         );
 
         expect(sendMessageMock).toHaveBeenCalledTimes(1);
@@ -179,7 +177,6 @@ describe('Subscriptions', () => {
             initialMsgId,
             defaultId,
             false,
-            nonRestrictedCountry,
           );
         },
       ).rejects.toEqual(new Error(`Invalid channel: ${invalidChannel}`));
@@ -194,7 +191,6 @@ describe('Subscriptions', () => {
         initialMsgId,
         mockSubaccountId,
         false,
-        nonRestrictedCountry,
       );
 
       expect(sendMessageMock).toHaveBeenCalledTimes(1);
@@ -217,7 +213,6 @@ describe('Subscriptions', () => {
         initialMsgId,
         mockSubaccountId,
         false,
-        nonRestrictedCountry,
       );
 
       expect(sendMessageMock).toHaveBeenCalledTimes(1);
@@ -269,31 +264,6 @@ describe('Subscriptions', () => {
       expect(subscriptions.subscriptionLists[connectionId]).toBeUndefined();
     });
 
-    it('sends blocked error if subscribing to subaccount from restricted country', async () => {
-      const expectedError: BlockedError = new BlockedError();
-      await subscriptions.subscribe(
-        mockWs,
-        Channel.V4_ACCOUNTS,
-        connectionId,
-        initialMsgId,
-        mockSubaccountId,
-        false,
-        restrictedCountry,
-      );
-
-      expect(sendMessageMock).toHaveBeenCalledTimes(1);
-      expect(sendMessageMock).toHaveBeenCalledWith(
-        mockWs,
-        connectionId,
-        expect.objectContaining({
-          connection_id: connectionId,
-          type: 'error',
-          message: expectedError.message,
-        }));
-      expect(subscriptions.subscriptions[Channel.V4_ACCOUNTS]).toBeUndefined();
-      expect(subscriptions.subscriptionLists[connectionId]).toBeUndefined();
-    });
-
     it('sends empty contents if initial message request fails with 404 for accounts', async () => {
       axiosRequestMock.mockImplementation(() => {
         return Promise.reject(makeAxiosSafeServerError(404, '', ''));
@@ -386,7 +356,6 @@ describe('Subscriptions', () => {
           initialMsgId,
           validIds[channel],
           false,
-          nonRestrictedCountry,
         );
       }));
 

indexer/services/socks/__tests__/websocket/index.test.ts

@@ -14,15 +14,13 @@ import {
 } from '../../src/types';
 import { InvalidMessageHandler } from '../../src/lib/invalid-message';
 import { PingHandler } from '../../src/lib/ping';
-import config from '../../src/config';
-import { isRestrictedCountryHeaders, COUNTRY_HEADER_KEY } from '@dydxprotocol-indexer/compliance';
+import { COUNTRY_HEADER_KEY } from '@dydxprotocol-indexer/compliance';
 
 jest.mock('uuid');
 jest.mock('../../src/helpers/wss');
 jest.mock('../../src/lib/subscription');
 jest.mock('../../src/lib/invalid-message');
 jest.mock('../../src/lib/ping');
-jest.mock('@dydxprotocol-indexer/compliance');
 
 describe('Index', () => {
   let index: Index;
@@ -32,12 +30,10 @@ describe('Index', () => {
   let mockConnect: (ws: WebSocket, req: IncomingMessage) => void;
   let wsOnSpy: jest.SpyInstance;
   let wsPingSpy: jest.SpyInstance;
-  let wsTerminateSpy: jest.SpyInstance;
   let invalidMsgHandlerSpy: jest.SpyInstance;
   let pingHandlerSpy: jest.SpyInstance;
 
   const connectionId: string = 'conId';
-  const defaultGeoblockingEnabled: boolean = config.INDEXER_LEVEL_GEOBLOCKING_ENABLED;
   const countryCode: string = 'AR';
 
   beforeAll(() => {
@@ -58,7 +54,6 @@ describe('Index', () => {
     websocket = new WebSocket(null);
     wsOnSpy = jest.spyOn(websocket, 'on');
     wsPingSpy = jest.spyOn(websocket, 'ping').mockImplementation(jest.fn());
-    wsTerminateSpy = jest.spyOn(websocket, 'terminate').mockImplementation(jest.fn());
     mockWss.onConnection = jest.fn().mockImplementation(
       (cb: (ws: WebSocket, req: IncomingMessage) => void) => {
         mockConnect = cb;
@@ -97,46 +92,6 @@ describe('Index', () => {
         }),
       );
     });
-
-    describe('geoblocking', () => {
-      const isRestrictedCountrySpy: jest.Mock = isRestrictedCountryHeaders as unknown as jest.Mock;
-
-      beforeAll(() => {
-        config.INDEXER_LEVEL_GEOBLOCKING_ENABLED = true;
-      });
-
-      afterAll(() => {
-        config.INDEXER_LEVEL_GEOBLOCKING_ENABLED = defaultGeoblockingEnabled;
-      });
-
-      it('rejects connection if from restricted country', () => {
-        jest.spyOn(websocket, 'terminate').mockImplementation(jest.fn());
-        // restricted country headers
-        isRestrictedCountrySpy.mockReturnValue(true);
-
-        const message: IncomingMessage = new IncomingMessage(new Socket());
-        mockConnect(websocket, message);
-        expect(websocket.terminate).toHaveBeenCalled();
-        expect(Object.keys(index.connections)).toHaveLength(0);
-        expect(wsOnSpy).not.toHaveBeenCalled();
-        expect(wsTerminateSpy).toHaveBeenCalled();
-        expect(sendMessage).not.toHaveBeenCalled();
-      });
-
-      it('does not reject connection if from restricted country', () => {
-        (v4 as unknown as jest.Mock).mockReturnValueOnce(connectionId);
-        // non-restricted country headers
-        isRestrictedCountrySpy.mockReturnValue(false);
-
-        const message: IncomingMessage = new IncomingMessage(new Socket());
-        mockConnect(websocket, message);
-
-        // Test that the connection is tracked.
-        expect(index.connections[connectionId]).not.toBeUndefined();
-        expect(index.connections[connectionId].ws).toEqual(websocket);
-        expect(index.connections[connectionId].messageId).toEqual(0);
-      });
-    });
   });
 
   describe('handlers', () => {

indexer/services/socks/src/helpers/header-utils.ts

@@ -0,0 +1,8 @@
+import { CountryHeaders } from '@dydxprotocol-indexer/compliance';
+
+import { IncomingMessage } from '../types';
+
+export function getCountry(req: IncomingMessage): string | undefined {
+  const countryHeaders: CountryHeaders = req.headers as CountryHeaders;
+  return countryHeaders['cf-ipcountry'];
+}

indexer/services/socks/src/lib/subscription.ts

@@ -491,13 +491,6 @@ export class Subscriptions {
       throw new Error('Invalid undefined id');
     }
 
-    // TODO(IND-508): Change this to match technical spec for persistent geo-blocking. This may
-    // either have to replicate any blocking logic added on comlink, or re-direct to comlink to
-    // determine if subscribing to a specific subaccount is blocked.
-    if (country !== undefined && isRestrictedCountry(country)) {
-      throw new BlockedError();
-    }
-
     try {
       const {
         address,
@@ -518,13 +511,19 @@ export class Subscriptions {
           method: RequestMethod.GET,
           url: `${COMLINK_URL}/v4/addresses/${address}/subaccountNumber/${subaccountNumber}`,
           timeout: config.INITIAL_GET_TIMEOUT_MS,
+          headers: {
+            'cf-ipcountry': country,
+          },
           transformResponse: (res) => res,
         }),
         // TODO(DEC-1462): Use the /active-orders endpoint once it's added.
         axiosRequest({
           method: RequestMethod.GET,
           url: `${COMLINK_URL}/v4/orders?address=${address}&subaccountNumber=${subaccountNumber}&status=OPEN,UNTRIGGERED,BEST_EFFORT_OPENED`,
           timeout: config.INITIAL_GET_TIMEOUT_MS,
+          headers: {
+            'cf-ipcountry': country,
+          },
           transformResponse: (res) => res,
         }),
       ]);
@@ -597,6 +596,9 @@ export class Subscriptions {
       method: RequestMethod.GET,
       url: endpoint,
       timeout: config.INITIAL_GET_TIMEOUT_MS,
+      headers: {
+        'cf-ipcountry': country,
+      },
       transformResponse: (res) => res, // Disables JSON parsing
     });
   }

indexer/services/socks/src/websocket/index.ts

@@ -5,6 +5,7 @@ import { v4 as uuidv4 } from 'uuid';
 import WebSocket from 'ws';
 
 import config from '../config';
+import { getCountry } from '../helpers/header-utils';
 import {
   createErrorMessage,
   createConnectedMessage,
@@ -26,7 +27,6 @@ import {
   ALL_CHANNELS,
   WebsocketEvents,
 } from '../types';
-import { CountryRestrictor } from './restrict-countries';
 
 const HEARTBEAT_INTERVAL_MS: number = config.WS_HEARTBEAT_INTERVAL_MS;
 const HEARTBEAT_TIMEOUT_MS: number = config.WS_HEARTBEAT_TIMEOUT_MS;
@@ -42,15 +42,13 @@ export class Index {
   // Handlers for pings and invalid messages.
   private pingHandler: PingHandler;
   private invalidMessageHandler: InvalidMessageHandler;
-  private countryRestrictor: CountryRestrictor;
 
   constructor(wss: Wss, subscriptions: Subscriptions) {
     this.wss = wss;
     this.connections = {};
     this.subscriptions = subscriptions;
     this.pingHandler = new PingHandler();
     this.invalidMessageHandler = new InvalidMessageHandler();
-    this.countryRestrictor = new CountryRestrictor();
 
     // Attach the new connection handler to the websocket server.
     this.wss.onConnection((ws: WebSocket, req: IncomingMessage) => this.onConnection(ws, req));
@@ -99,17 +97,13 @@ export class Index {
    * @param req HTTP request accompanying new connection request.
    */
   private onConnection(ws: WebSocket, req: IncomingMessage): void {
-    // Terminate the connection if the connection requestion originated from a restricted country
-    if (this.countryRestrictor.isRestrictedCountry(req)) {
-      return ws.terminate();
-    }
 
     const connectionId: string = uuidv4();
 
     this.connections[connectionId] = {
       ws,
       messageId: 0,
-      countryCode: this.countryRestrictor.getCountry(req),
+      countryCode: getCountry(req),
     };
 
     const numConcurrentConnections: number = Object.keys(this.connections).length;