Recreate KMSClient on each request (#358)

Uniswap · Jul 24, 2024 · bf87dcc · bf87dcc
1 parent 7cb7535
commit bf87dcc
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/lib/handlers/hard-quote/handler.ts b/lib/handlers/hard-quote/handler.ts
@@ -1,3 +1,5 @@
+import { KMSClient } from '@aws-sdk/client-kms';
+import { KmsSigner } from '@uniswap/signer';
 import { TradeType } from '@uniswap/sdk-core';
 import { MetricLoggerUnit } from '@uniswap/smart-order-router';
 import { CosignedV2DutchOrder, CosignerData } from '@uniswap/uniswapx-sdk';
@@ -18,6 +20,7 @@ import {
   HardQuoteResponseDataJoi,
 } from './schema';
 import { ChainId } from '../../util/chains';
+import { checkDefined } from '../../preconditions/preconditions';
 
 const DEFAULT_EXCLUSIVITY_OVERRIDE_BPS = BigNumber.from(100); // non-exclusive fillers must override price by this much
 const RESPONSE_LOG_TYPE = 'HardResponse';
@@ -33,7 +36,7 @@ export class QuoteHandler extends APIGLambdaHandler<
   ): Promise<ErrorResponse | Response<HardQuoteResponseData>> {
     const {
       requestInjected: { log, metric },
-      containerInjected: { quoters, orderServiceProvider, cosigner, cosignerAddress },
+      containerInjected: { quoters, orderServiceProvider, cosignerAddress },
       requestBody,
     } = params;
     const start = Date.now();
@@ -89,7 +92,9 @@ export class QuoteHandler extends APIGLambdaHandler<
       log.info({ cosignerData: cosignerData }, 'open order with default cosignerData');
     }
 
-    // TODO: use server key to cosign instead of local wallet
+    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
@@ -20,14 +20,9 @@ import { STAGE } from '../../util/stage';
 import { ApiInjector, ApiRInj } from '../base/api-handler';
 import { HardQuoteRequestBody } from './schema';
 
-interface Cosigner {
-  signDigest(digest: Buffer | string): Promise<string>;
-}
-
 export interface ContainerInjected {
   quoters: Quoter[];
   firehose: FirehoseLogger;
-  cosigner: Cosigner;
   cosignerAddress: string;
   orderServiceProvider: OrderServiceProvider;
 }
@@ -91,7 +86,6 @@ export class QuoteInjector extends ApiInjector<ContainerInjected, RequestInjecte
     return {
       quoters: quoters,
       firehose: firehose,
-      cosigner,
       cosignerAddress,
       orderServiceProvider,
     };