From 0a5353f287a14829f3803cb8aeac4fbf4fe2cec0 Mon Sep 17 00:00:00 2001
From: Alec Schaefer <alec@cLabs.co>
Date: Tue, 6 Apr 2021 12:30:16 -0400
Subject: [PATCH] adds whitelist logic to ODIS signers for load testing (#7624)

* adds requestID to logger in common pkg

* assign sessionID in combiner if not provided by client

* adds request whitelisting
---
 .env.alfajores                                |  1 +
 .../signer/src/common/metrics.ts              |  4 +++
 .../phone-number-privacy/signer/src/config.ts |  2 ++
 .../src/signing/get-partial-signature.ts      | 32 ++++++++++++-------
 4 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/.env.alfajores b/.env.alfajores
index e210742744a..ec7f6245880 100644
--- a/.env.alfajores
+++ b/.env.alfajores
@@ -161,6 +161,7 @@ LOG_FORMAT=stackdriver
 # Options: fatal, error, warn, info (default), debug, trace
 LOG_LEVEL=info
 
+WHITELIST_PERCENTAGE=50
 
 # ODIS signer 1 Azure info
 AZURE_ODIS_EASTUS_1_AZURE_SUBSCRIPTION_ID=97e2b592-255b-4f92-bce0-127257163c36
diff --git a/packages/phone-number-privacy/signer/src/common/metrics.ts b/packages/phone-number-privacy/signer/src/common/metrics.ts
index 76cca86615a..4e8dc5bd44f 100644
--- a/packages/phone-number-privacy/signer/src/common/metrics.ts
+++ b/packages/phone-number-privacy/signer/src/common/metrics.ts
@@ -52,6 +52,10 @@ export const Counters = {
     help:
       'Counter for the number of requests in which the account is not verified but meets min balance',
   }),
+  whitelistedRequests: new Counter({
+    name: 'whitelisted_requests',
+    help: 'Counter for the number of whitelisted requests not requiring quota (testing only)',
+  }),
 }
 const buckets = [
   0.001,
diff --git a/packages/phone-number-privacy/signer/src/config.ts b/packages/phone-number-privacy/signer/src/config.ts
index 8a1e3dfc2aa..2cc3a3e7bd6 100644
--- a/packages/phone-number-privacy/signer/src/config.ts
+++ b/packages/phone-number-privacy/signer/src/config.ts
@@ -77,6 +77,7 @@ interface Config {
       secretKey: string
     }
   }
+  whitelist_percentage: number
 }
 
 const env = process.env as any
@@ -131,5 +132,6 @@ const config: Config = {
       secretKey: env.KEYSTORE_AWS_SECRET_KEY,
     },
   },
+  whitelist_percentage: Number(env.WHITELIST_PERCENTAGE) || 0,
 }
 export default config
diff --git a/packages/phone-number-privacy/signer/src/signing/get-partial-signature.ts b/packages/phone-number-privacy/signer/src/signing/get-partial-signature.ts
index 4c883236523..312735c43a2 100644
--- a/packages/phone-number-privacy/signer/src/signing/get-partial-signature.ts
+++ b/packages/phone-number-privacy/signer/src/signing/get-partial-signature.ts
@@ -16,7 +16,7 @@ import allSettled from 'promise.allsettled'
 import { computeBlindedSignature } from '../bls/bls-cryptography-client'
 import { respondWithError } from '../common/error-utils'
 import { Counters, Histograms, Labels } from '../common/metrics'
-import { getVersion } from '../config'
+import config, { getVersion } from '../config'
 import { incrementQueryCount } from '../database/wrappers/account'
 import { getRequestExists, storeRequest } from '../database/wrappers/request'
 import { getKeyProvider } from '../key-management/key-provider'
@@ -111,16 +111,21 @@ export async function handleGetBlindedMessagePartialSig(
 
     if (_queryCount.status === 'fulfilled' && performedQueryCount >= totalQuota) {
       logger.debug('No remaining query count')
-      respondWithError(
-        Endpoints.GET_BLINDED_MESSAGE_PARTIAL_SIG,
-        response,
-        403,
-        WarningMessage.EXCEEDED_QUOTA,
-        performedQueryCount,
-        totalQuota,
-        blockNumber
-      )
-      return
+      if (isWhitelisted(request.body)) {
+        Counters.whitelistedRequests.inc()
+        logger.info({ request: request.body }, 'Request whitelisted')
+      } else {
+        respondWithError(
+          Endpoints.GET_BLINDED_MESSAGE_PARTIAL_SIG,
+          response,
+          403,
+          WarningMessage.EXCEEDED_QUOTA,
+          performedQueryCount,
+          totalQuota,
+          blockNumber
+        )
+        return
+      }
     }
 
     const meterGenerateSignature = Histograms.getBlindedSigInstrumentation
@@ -195,3 +200,8 @@ function isValidGetSignatureInput(requestBody: GetBlindedMessagePartialSigReques
     hasValidTimestamp(requestBody)
   )
 }
+
+function isWhitelisted(requestBody: GetBlindedMessagePartialSigRequest) {
+  const sessionID = Number(requestBody.sessionID)
+  return sessionID && sessionID % 100 < config.whitelist_percentage
+}