From e226a0049c808f858f417cabe92ca96e28a45697 Mon Sep 17 00:00:00 2001
From: dydxwill <119354122+dydxwill@users.noreply.github.com>
Date: Tue, 5 Mar 2024 15:03:28 -0500
Subject: [PATCH] [OTE-141] implement post /compliance/geoblock (#1129)

---
 .../api/v4/compliance-v2-controller.test.ts   | 255 +++++++++++++++---
 .../api/v4/compliance-v2-controller.ts        |  88 +++++-
 .../helpers/compliance/compliance-utils.ts    |  18 ++
 3 files changed, 328 insertions(+), 33 deletions(-)
 create mode 100644 indexer/services/comlink/src/helpers/compliance/compliance-utils.ts

diff --git a/indexer/services/comlink/__tests__/controllers/api/v4/compliance-v2-controller.test.ts b/indexer/services/comlink/__tests__/controllers/api/v4/compliance-v2-controller.test.ts
index ec53cd0dd8..03d7e4ab55 100644
--- a/indexer/services/comlink/__tests__/controllers/api/v4/compliance-v2-controller.test.ts
+++ b/indexer/services/comlink/__tests__/controllers/api/v4/compliance-v2-controller.test.ts
@@ -1,16 +1,16 @@
 import {
+  ComplianceReason,
   ComplianceStatus,
+  ComplianceStatusFromDatabase,
+  ComplianceStatusTable,
   dbHelpers,
-  testMocks,
   testConstants,
-  ComplianceReason,
-  ComplianceStatusTable,
-  ComplianceStatusFromDatabase,
+  testMocks,
 } from '@dydxprotocol-indexer/postgres';
 import { getIpAddr } from '../../../../src/lib/utils';
 import { sendRequest } from '../../../helpers/helpers';
 import { RequestMethod } from '../../../../src/types';
-import { stats } from '@dydxprotocol-indexer/base';
+import { logger, stats } from '@dydxprotocol-indexer/base';
 import { redis } from '@dydxprotocol-indexer/redis';
 import { ratelimitRedis } from '../../../../src/caches/rate-limiters';
 import { ComplianceControllerHelper } from '../../../../src/controllers/api/v4/compliance-controller';
@@ -18,6 +18,11 @@ import config from '../../../../src/config';
 import { DateTime } from 'luxon';
 import { ComplianceAction } from '../../../../src/controllers/api/v4/compliance-v2-controller';
 import { ExtendedSecp256k1Signature, Secp256k1, sha256 } from '@cosmjs/crypto';
+import { getGeoComplianceReason } from '../../../../src/helpers/compliance/compliance-utils';
+import { isRestrictedCountryHeaders } from '@dydxprotocol-indexer/compliance';
+
+jest.mock('@dydxprotocol-indexer/compliance');
+jest.mock('../../../../src/helpers/compliance/compliance-utils');
 
 jest.mock('../../../../src/lib/utils', () => ({
   ...jest.requireActual('../../../../src/lib/utils'),
@@ -231,7 +236,21 @@ describe('ComplianceV2Controller', () => {
   });
 
   describe('POST /geoblock', () => {
+    let getGeoComplianceReasonSpy: jest.SpyInstance;
+    let isRestrictedCountryHeadersSpy: jest.SpyInstance;
+
+    const body: any = {
+      address: testConstants.defaultAddress,
+      message: 'Test message',
+      action: ComplianceAction.ONBOARD,
+      signedMessage: sha256(Buffer.from('msg')),
+      pubkey: new Uint8Array([/* public key bytes */]),
+      timestamp: 1620000000,
+    };
+
     beforeEach(async () => {
+      getGeoComplianceReasonSpy = getGeoComplianceReason as unknown as jest.Mock;
+      isRestrictedCountryHeadersSpy = isRestrictedCountryHeaders as unknown as jest.Mock;
       ipAddrMock.mockReturnValue(ipAddr);
       await testMocks.seedData();
       jest.mock('@cosmjs/crypto', () => ({
@@ -257,12 +276,8 @@ describe('ComplianceV2Controller', () => {
         type: RequestMethod.POST,
         path: '/v4/compliance/geoblock',
         body: {
+          ...body,
           address: '0x123', // Non-dYdX address
-          message: 'Test message',
-          action: ComplianceAction.ONBOARD,
-          signedMessage: sha256(Buffer.from('msg')),
-          pubkey: new Uint8Array([/* public key bytes */]),
-          timestamp: 1620000000,
         },
         expectedStatus: 400,
       });
@@ -273,11 +288,7 @@ describe('ComplianceV2Controller', () => {
         type: RequestMethod.POST,
         path: '/v4/compliance/geoblock',
         body: {
-          address: testConstants.defaultAddress,
-          message: 'Test message',
-          action: ComplianceAction.ONBOARD,
-          signedMessage: sha256(Buffer.from('msg')),
-          pubkey: new Uint8Array([/* public key bytes */]),
+          ...body,
           timestamp: 1619996600, // More than 30 seconds difference
         },
         expectedStatus: 400,
@@ -291,14 +302,7 @@ describe('ComplianceV2Controller', () => {
       await sendRequest({
         type: RequestMethod.POST,
         path: '/v4/compliance/geoblock',
-        body: {
-          address: testConstants.defaultAddress,
-          message: 'Test message',
-          action: ComplianceAction.ONBOARD,
-          signedMessage: sha256(Buffer.from('msg')),
-          pubkey: new Uint8Array([/* public key bytes */]),
-          timestamp: 1620000000,
-        },
+        body,
         expectedStatus: 400,
       });
     });
@@ -306,21 +310,212 @@ describe('ComplianceV2Controller', () => {
     it('should process valid request', async () => {
       (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
 
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body,
+      });
+
+      expect(response.status).toEqual(200);
+      expect(response.body.status).toEqual(ComplianceStatus.COMPLIANT);
+    });
+
+    it('should set status to BLOCKED for ONBOARD action from a restricted country with no existing compliance status', async () => {
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      getGeoComplianceReasonSpy.mockReturnValueOnce(ComplianceReason.US_GEO);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body,
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.BLOCKED,
+        reason: ComplianceReason.US_GEO,
+      }));
+
+      expect(response.body.status).toEqual(ComplianceStatus.BLOCKED);
+      expect(response.body.reason).toEqual(ComplianceReason.US_GEO);
+    });
+
+    it('should set status to FIRST_STRIKE for CONNECT action from a restricted country with no existing compliance status', async () => {
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      getGeoComplianceReasonSpy.mockReturnValueOnce(ComplianceReason.US_GEO);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
       const response: any = await sendRequest({
         type: RequestMethod.POST,
         path: '/v4/compliance/geoblock',
         body: {
-          address: testConstants.defaultAddress,
-          message: 'Test message',
-          action: ComplianceAction.ONBOARD,
-          signedMessage: sha256(Buffer.from('msg')),
-          pubkey: new Uint8Array([/* public key bytes */]),
-          timestamp: 1620000000, // Valid timestamp
+          ...body,
+          action: ComplianceAction.CONNECT,
         },
+        expectedStatus: 200,
       });
 
-      expect(response.status).toEqual(200);
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.FIRST_STRIKE,
+        reason: ComplianceReason.US_GEO,
+      }));
+
+      expect(response.body.status).toEqual(ComplianceStatus.FIRST_STRIKE);
+      expect(response.body.reason).toEqual(ComplianceReason.US_GEO);
+    });
+
+    it('should set status to COMPLIANT for any action from a non-restricted country with no existing compliance status', async () => {
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      isRestrictedCountryHeadersSpy.mockReturnValue(false);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body,
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.COMPLIANT,
+      }));
+
       expect(response.body.status).toEqual(ComplianceStatus.COMPLIANT);
     });
+
+    it('should update status to FIRST_STRIKE for CONNECT action from a restricted country with existing COMPLIANT status', async () => {
+      await ComplianceStatusTable.create({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.COMPLIANT,
+      });
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      getGeoComplianceReasonSpy.mockReturnValueOnce(ComplianceReason.US_GEO);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body: {
+          ...body,
+          action: ComplianceAction.CONNECT,
+        },
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.FIRST_STRIKE,
+        reason: ComplianceReason.US_GEO,
+      }));
+
+      expect(response.body.status).toEqual(ComplianceStatus.FIRST_STRIKE);
+      expect(response.body.reason).toEqual(ComplianceReason.US_GEO);
+    });
+
+    it('should be a no-op for ONBOARD action with existing COMPLIANT status', async () => {
+      const loggerError = jest.spyOn(logger, 'error');
+      await ComplianceStatusTable.create({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.COMPLIANT,
+      });
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body,
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.COMPLIANT,
+      }));
+
+      expect(loggerError).toHaveBeenCalledWith(expect.objectContaining({
+        at: 'ComplianceV2Controller POST /geoblock',
+        message: 'Invalid action for current compliance status',
+      }));
+      expect(response.body.status).toEqual(ComplianceStatus.COMPLIANT);
+    });
+
+    it('should be a no-op for ONBOARD action with existing FIRST_STRIKE status', async () => {
+      const loggerError = jest.spyOn(logger, 'error');
+      await ComplianceStatusTable.create({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.FIRST_STRIKE,
+        reason: ComplianceReason.US_GEO,
+      });
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body,
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.FIRST_STRIKE,
+        reason: ComplianceReason.US_GEO,
+      }));
+
+      expect(loggerError).toHaveBeenCalledWith(expect.objectContaining({
+        at: 'ComplianceV2Controller POST /geoblock',
+        message: 'Invalid action for current compliance status',
+      }));
+      expect(response.body.status).toEqual(ComplianceStatus.FIRST_STRIKE);
+      expect(response.body.reason).toEqual(ComplianceReason.US_GEO);
+    });
+
+    it('should update status to CLOSE_ONLY for CONNECT action from a restricted country with existing FIRST_STRIKE status', async () => {
+      await ComplianceStatusTable.create({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.FIRST_STRIKE,
+        reason: ComplianceReason.US_GEO,
+      });
+      (Secp256k1.verifySignature as jest.Mock).mockResolvedValueOnce(true);
+      getGeoComplianceReasonSpy.mockReturnValueOnce(ComplianceReason.US_GEO);
+      isRestrictedCountryHeadersSpy.mockReturnValue(true);
+
+      const response: any = await sendRequest({
+        type: RequestMethod.POST,
+        path: '/v4/compliance/geoblock',
+        body: {
+          ...body,
+          action: ComplianceAction.CONNECT,
+        },
+        expectedStatus: 200,
+      });
+
+      const data: ComplianceStatusFromDatabase[] = await ComplianceStatusTable.findAll({}, [], {});
+      expect(data).toHaveLength(1);
+      expect(data[0]).toEqual(expect.objectContaining({
+        address: testConstants.defaultAddress,
+        status: ComplianceStatus.CLOSE_ONLY,
+        reason: ComplianceReason.US_GEO,
+      }));
+
+      expect(response.body.status).toEqual(ComplianceStatus.CLOSE_ONLY);
+      expect(response.body.reason).toEqual(ComplianceReason.US_GEO);
+    });
   });
 });
diff --git a/indexer/services/comlink/src/controllers/api/v4/compliance-v2-controller.ts b/indexer/services/comlink/src/controllers/api/v4/compliance-v2-controller.ts
index df0d4e745c..fb9307cb5d 100644
--- a/indexer/services/comlink/src/controllers/api/v4/compliance-v2-controller.ts
+++ b/indexer/services/comlink/src/controllers/api/v4/compliance-v2-controller.ts
@@ -1,5 +1,6 @@
-import { Secp256k1, sha256, ExtendedSecp256k1Signature } from '@cosmjs/crypto';
+import { ExtendedSecp256k1Signature, Secp256k1, sha256 } from '@cosmjs/crypto';
 import { logger, stats, TooManyRequestsError } from '@dydxprotocol-indexer/base';
+import { CountryHeaders, isRestrictedCountryHeaders } from '@dydxprotocol-indexer/compliance';
 import {
   ComplianceReason,
   ComplianceStatus,
@@ -16,6 +17,7 @@ import {
 import { getReqRateLimiter } from '../../../caches/rate-limiters';
 import config from '../../../config';
 import { complianceProvider } from '../../../helpers/compliance/compliance-clients';
+import { getGeoComplianceReason } from '../../../helpers/compliance/compliance-utils';
 import { DYDX_ADDRESS_PREFIX, GEOBLOCK_REQUEST_TTL_SECONDS } from '../../../lib/constants';
 import { create4xxResponse, handleControllerError } from '../../../lib/helpers';
 import { rateLimiterMiddleware } from '../../../lib/rate-limit';
@@ -34,6 +36,11 @@ export enum ComplianceAction {
   CONNECT = 'CONNECT',
 }
 
+const COMPLIANCE_PROGRESSION: Partial<Record<ComplianceStatus, ComplianceStatus>> = {
+  [ComplianceStatus.COMPLIANT]: ComplianceStatus.FIRST_STRIKE,
+  [ComplianceStatus.FIRST_STRIKE]: ComplianceStatus.CLOSE_ONLY,
+};
+
 @Route('compliance')
 class ComplianceV2Controller extends Controller {
   private ipAddress: string;
@@ -209,9 +216,84 @@ router.post(
         );
       }
 
-      // TODO(OTE-141): Implement logic.
+      /**
+       * If the address doesn't exist in the compliance table:
+       * - if the request is from a restricted country:
+       *  - if the action is ONBOARD, set the status to BLOCKED
+       *  - if the action is CONNECT, set the status to FIRST_STRIKE
+       * - else if the request is from a non-restricted country:
+       *  - set the status to COMPLIANT
+       *
+       * if the address is COMPLIANT:
+       * - the ONLY action should be CONNECT. ONBOARD is a no-op.
+       * - if the request is from a restricted country:
+       *  - set the status to FIRST_STRIKE
+       *
+       * if the address is FIRST_STRIKE:
+       * - the ONLY action should be CONNECT. ONBOARD is a no-op.
+       * - if the request is from a restricted country:
+       *  - set the status to CLOSE_ONLY
+       */
+      const complianceStatus: ComplianceStatusFromDatabase[] = await
+      ComplianceStatusTable.findAll(
+        { address: [address] },
+        [],
+      );
+      let complianceStatusFromDatabase: ComplianceStatusFromDatabase | undefined;
+      if (complianceStatus.length === 0) {
+        if (isRestrictedCountryHeaders(req.headers as CountryHeaders)) {
+          if (action === ComplianceAction.ONBOARD) {
+            complianceStatusFromDatabase = await ComplianceStatusTable.upsert({
+              address,
+              status: ComplianceStatus.BLOCKED,
+              reason: getGeoComplianceReason(req.headers as CountryHeaders)!,
+              updatedAt: DateTime.utc().toISO(),
+            });
+          } else if (action === ComplianceAction.CONNECT) {
+            complianceStatusFromDatabase = await ComplianceStatusTable.upsert({
+              address,
+              status: ComplianceStatus.FIRST_STRIKE,
+              reason: getGeoComplianceReason(req.headers as CountryHeaders)!,
+              updatedAt: DateTime.utc().toISO(),
+            });
+          }
+        } else {
+          complianceStatusFromDatabase = await ComplianceStatusTable.upsert({
+            address,
+            status: ComplianceStatus.COMPLIANT,
+            updatedAt: DateTime.utc().toISO(),
+          });
+        }
+      } else {
+        complianceStatusFromDatabase = complianceStatus[0];
+        if (
+          complianceStatus[0].status === ComplianceStatus.FIRST_STRIKE ||
+          complianceStatus[0].status === ComplianceStatus.COMPLIANT
+        ) {
+          if (action === ComplianceAction.ONBOARD) {
+            logger.error({
+              at: 'ComplianceV2Controller POST /geoblock',
+              message: 'Invalid action for current compliance status',
+              address,
+              action,
+              complianceStatus: complianceStatus[0],
+            });
+          } else if (
+            isRestrictedCountryHeaders(req.headers as CountryHeaders) &&
+            action === ComplianceAction.CONNECT
+          ) {
+            complianceStatusFromDatabase = await ComplianceStatusTable.update({
+              address,
+              status: COMPLIANCE_PROGRESSION[complianceStatus[0].status],
+              reason: getGeoComplianceReason(req.headers as CountryHeaders)!,
+              updatedAt: DateTime.utc().toISO(),
+            });
+          }
+        }
+      }
       const response = {
-        status: ComplianceStatus.COMPLIANT,
+        status: complianceStatusFromDatabase!.status,
+        reason: complianceStatusFromDatabase!.reason,
       };
 
       return res.send(response);
diff --git a/indexer/services/comlink/src/helpers/compliance/compliance-utils.ts b/indexer/services/comlink/src/helpers/compliance/compliance-utils.ts
new file mode 100644
index 0000000000..7aead50b20
--- /dev/null
+++ b/indexer/services/comlink/src/helpers/compliance/compliance-utils.ts
@@ -0,0 +1,18 @@
+import { CountryHeaders, isRestrictedCountryHeaders } from '@dydxprotocol-indexer/compliance';
+import { ComplianceReason } from '@dydxprotocol-indexer/postgres';
+
+export function getGeoComplianceReason(
+  headers: CountryHeaders,
+): ComplianceReason | undefined {
+  if (isRestrictedCountryHeaders(headers)) {
+    const country: string | undefined = headers['cf-ipcountry'];
+    if (country === 'US') {
+      return ComplianceReason.US_GEO;
+    } else if (country === 'CA') {
+      return ComplianceReason.CA_GEO;
+    } else {
+      return ComplianceReason.SANCTIONED_GEO;
+    }
+  }
+  return undefined;
+}