Merge pull request #378 from Uniswap/recreating-kms

fix: remove kms client initialization during init stage (in injector)
Uniswap · Sep 19, 2024 · 1b558ef · 1b558ef
2 parents b1044db + 395c9c6
commit 1b558ef
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 19 deletions.
diff --git a/lib/handlers/hard-quote/handler.ts b/lib/handlers/hard-quote/handler.ts
@@ -37,16 +37,22 @@ export class QuoteHandler extends APIGLambdaHandler<
   ): Promise<ErrorResponse | Response<HardQuoteResponseData>> {
     const {
       requestInjected: { log, metric },
-      containerInjected: { quoters, orderServiceProvider, cosignerAddress },
+      containerInjected: { quoters, orderServiceProvider },
       requestBody,
     } = params;
     const start = Date.now();
 
     metric.putMetric(Metric.QUOTE_REQUESTED, 1, MetricLoggerUnit.Count);
 
-    log.info({ cosignerAddress: cosignerAddress }, 'cosignerAddress');
     const request = HardQuoteRequest.fromHardRequestBody(requestBody);
 
+    // re-create KmsClient every call to avoid clock skew issue
+    // https://github.com/aws/aws-sdk-js-v3/issues/6400
+    const kmsKeyId = checkDefined(process.env.KMS_KEY_ID, 'KMS_KEY_ID is not defined');
+    const awsRegion = checkDefined(process.env.REGION, 'REGION is not defined');
+    const cosigner = new KmsSigner(new KMSClient({ region: awsRegion }), kmsKeyId);
+    const cosignerAddress = await cosigner.getAddress();
+
     // we dont have access to the cosigner key, throw
     if (request.order.info.cosigner !== cosignerAddress) {
       log.error({ cosignerInReq: request.order.info.cosigner, expected: cosignerAddress }, 'Unknown cosigner');
@@ -92,10 +98,6 @@ export class QuoteHandler extends APIGLambdaHandler<
       cosignerData = getDefaultCosignerData(request);
       log.info({ cosignerData: cosignerData }, 'open order with default cosignerData');
     }
-
-    const kmsKeyId = checkDefined(process.env.KMS_KEY_ID, 'KMS_KEY_ID is not defined');
-    const awsRegion = checkDefined(process.env.REGION, 'REGION is not defined');
-    const cosigner = new KmsSigner(new KMSClient({ region: awsRegion }), kmsKeyId);
     const cosignature = await cosigner.signDigest(request.order.cosignatureHash(cosignerData));
     const cosignedOrder = CosignedV2DutchOrder.fromUnsignedOrder(request.order, cosignerData, cosignature);
 

diff --git a/lib/handlers/hard-quote/injector.ts b/lib/handlers/hard-quote/injector.ts
@@ -1,5 +1,3 @@
-import { KMSClient } from '@aws-sdk/client-kms';
-import { KmsSigner } from '@uniswap/signer';
 import { IMetric, setGlobalLogger, setGlobalMetric } from '@uniswap/smart-order-router';
 import { MetricsLogger } from 'aws-embedded-metrics';
 import { APIGatewayProxyEvent, Context } from 'aws-lambda';
@@ -23,7 +21,6 @@ import { HardQuoteRequestBody } from './schema';
 export interface ContainerInjected {
   quoters: Quoter[];
   firehose: FirehoseLogger;
-  cosignerAddress: string;
   orderServiceProvider: OrderServiceProvider;
 }
 
@@ -44,11 +41,6 @@ export class QuoteInjector extends ApiInjector<ContainerInjected, RequestInjecte
 
     const orderServiceUrl = checkDefined(process.env.ORDER_SERVICE_URL, 'ORDER_SERVICE_URL is not defined');
 
-    const kmsKeyId = checkDefined(process.env.KMS_KEY_ID, 'KMS_KEY_ID is not defined');
-    const awsRegion = checkDefined(process.env.REGION, 'REGION is not defined');
-    const cosigner = new KmsSigner(new KMSClient({ region: awsRegion }), kmsKeyId);
-    const cosignerAddress = await cosigner.getAddress();
-
     const webhookProvider = new S3WebhookConfigurationProvider(log, `${WEBHOOK_CONFIG_BUCKET}-${stage}-1`, s3Key);
     await webhookProvider.fetchEndpoints();
     const circuitBreakerProvider = new DynamoCircuitBreakerConfigurationProvider(
@@ -82,11 +74,9 @@ export class QuoteInjector extends ApiInjector<ContainerInjected, RequestInjecte
     const quoters: Quoter[] = [
       new WebhookQuoter(log, firehose, webhookProvider, circuitBreakerProvider, fillerComplianceProvider, repository),
     ];
-    log.info({ cosignerAddress }, 'Cosigner address from KMS Signer');
     return {
       quoters: quoters,
       firehose: firehose,
-      cosignerAddress,
       orderServiceProvider,
     };
   }

diff --git a/test/handlers/hard-quote/handler.test.ts b/test/handlers/hard-quote/handler.test.ts
@@ -1,3 +1,4 @@
+import { KMSClient } from '@aws-sdk/client-kms';
 import { TradeType } from '@uniswap/sdk-core';
 import { CosignedV2DutchOrder, UnsignedV2DutchOrder, UnsignedV2DutchOrderInfo } from '@uniswap/uniswapx-sdk';
 import { createMetricsLogger } from 'aws-embedded-metrics';
@@ -19,8 +20,11 @@ import {
 import { getCosignerData } from '../../../lib/handlers/hard-quote/handler';
 import { MockOrderServiceProvider } from '../../../lib/providers';
 import { MockQuoter, MOCK_FILLER_ADDRESS, Quoter } from '../../../lib/quoters';
+import { KmsSigner } from '@uniswap/signer';
 
 jest.mock('axios');
+jest.mock('@aws-sdk/client-kms');
+jest.mock('@uniswap/signer');
 
 const QUOTE_ID = 'a83f397c-8ef4-4801-a9b7-6e79155049f6';
 const REQUEST_ID = 'a83f397c-8ef4-4801-a9b7-6e79155049f6';
@@ -33,6 +37,9 @@ const CHAIN_ID = 1;
 const logger = Logger.createLogger({ name: 'test' });
 logger.level(Logger.FATAL);
 
+process.env.KMS_KEY_ID = 'test-key-id';
+process.env.REGION = 'us-east-2';
+
 export const getOrder = (data: Partial<UnsignedV2DutchOrderInfo>): UnsignedV2DutchOrder => {
   const now = Math.floor(new Date().getTime() / 1000);
   return new UnsignedV2DutchOrder(
@@ -70,7 +77,17 @@ export const getOrder = (data: Partial<UnsignedV2DutchOrderInfo>): UnsignedV2Dut
 describe('Quote handler', () => {
   const swapperWallet = Wallet.createRandom();
   const cosignerWallet = Wallet.createRandom();
-
+
+  const mockGetAddress = jest.fn().mockResolvedValue(cosignerWallet.address);
+  const mockSignDigest = jest.fn().mockImplementation((digest) => cosignerWallet.signMessage(ethers.utils.arrayify(digest)));
+
+  (KmsSigner as jest.Mock).mockImplementation(() => ({
+    getAddress: mockGetAddress,
+    signDigest: mockSignDigest,
+  }));
+  (KMSClient as jest.Mock).mockImplementation(() => jest.fn());
+
+
   // Creating mocks for all the handler dependencies.
   const requestInjectedMock: Promise<RequestInjected> = new Promise(
     (resolve) =>
@@ -89,8 +106,6 @@ describe('Quote handler', () => {
         getContainerInjected: () => {
           return {
             quoters,
-            cosigner: cosignerWallet._signingKey(),
-            cosignerAddress: cosignerWallet.address,
             orderServiceProvider: new MockOrderServiceProvider(),
           };
         },