chore: warmup middleware for authorization lambdas

packages/daemon/__tests__/integration/config.ts

@@ -11,7 +11,7 @@ export const UNVOIDED_SCENARIO_PORT = 8081;
 // Last event is actually 39, but event 39 is ignored by the machine (because
 // the transaction is already added), and when we ignore an event, we don't store
 // it in the database.
-export const UNVOIDED_SCENARIO_LAST_EVENT = 38;
+export const UNVOIDED_SCENARIO_LAST_EVENT = 39;
 
 // reorg
 export const REORG_SCENARIO_PORT = 8082;

packages/daemon/src/types/event.ts

@@ -103,8 +103,11 @@ export const EventTxOutputSchema = z.object({
       type: z.string(),
       address: z.string(),
       timelock: z.number().nullable(),
-    }).nullable(),
-    z.object({}).strict()
+    }).passthrough().nullable(),
+    z.object({
+      token_data: z.number(),
+    }),
+    z.object({}).strict(),
   ]),
 });
 export type EventTxOutput = z.infer<typeof EventTxOutputSchema>;

packages/wallet-service/serverless.yml

@@ -588,12 +588,12 @@ functions:
           cors: true
     warmup:
       walletWarmer:
-        enabled: false
+        enabled: true
   bearerAuthorizer:
     handler: src/api/auth.bearerAuthorizer
     warmup:
       walletWarmer:
-        enabled: false
+        enabled: true
   metrics:
     handler: src/metrics.getMetrics
     events:

packages/wallet-service/src/api/auth.ts

@@ -26,6 +26,7 @@ import {
   validateAuthTimestamp,
   AUTH_MAX_TIMESTAMP_SHIFT_IN_SECONDS,
 } from '@src/utils';
+import { warmupMiddleware } from '@src/api/utils';
 import middy from '@middy/core';
 import cors from '@middy/http-cors';
 import createDefaultLogger from '@src/logger';
@@ -83,6 +84,17 @@ export const tokenHandler: APIGatewayProxyHandler = middy(async (event) => {
   const authXpubStr = value.xpub;
   const wallet: Wallet = await getWallet(mysql, value.walletId);
 
+  if (!wallet) {
+    await closeDbConnection(mysql);
+    return {
+      statusCode: 400,
+      body: JSON.stringify({
+        success: false,
+        error: ApiError.WALLET_NOT_FOUND,
+      }),
+    };
+  }
+
   const [validTimestamp, timestampShift] = validateAuthTimestamp(timestamp, Date.now() / 1000);
 
   if (!validTimestamp) {
@@ -155,6 +167,7 @@ export const tokenHandler: APIGatewayProxyHandler = middy(async (event) => {
     body: JSON.stringify({ success: true, token }),
   };
 }).use(cors())
+  .use(warmupMiddleware())
   .use(errorHandler());
 
 /**
@@ -235,4 +248,5 @@ export const bearerAuthorizer: APIGatewayTokenAuthorizerHandler = middy(async (e
   }
 
   return _generatePolicy(walletId, 'Deny', event.methodArn, logger);
-}).use(cors());
+}).use(cors())
+  .use(warmupMiddleware());

packages/wallet-service/src/api/newAddresses.ts

@@ -30,16 +30,16 @@ const mysql = getDbConnection();
  * This lambda is called by API Gateway on GET /addresses/new
  */
 export const get: APIGatewayProxyHandler = middy(walletIdProxyHandler(async (walletId) => {
-  const status = await getWallet(mysql, walletId);
+  const wallet = await getWallet(mysql, walletId);
 
-  if (!status) {
+  if (!wallet) {
     return closeDbAndGetError(mysql, ApiError.WALLET_NOT_FOUND);
   }
-  if (!status.readyAt) {
+  if (!wallet.readyAt) {
     return closeDbAndGetError(mysql, ApiError.WALLET_NOT_READY);
   }
 
-  const addresses = await getNewAddresses(mysql, walletId);
+  const addresses = await getNewAddresses(mysql, wallet);
 
   await closeDbConnection(mysql);
 

packages/wallet-service/src/db/index.ts

@@ -289,6 +289,7 @@ export const createWallet = async (
     status: WalletStatus.CREATING,
     created_at: ts,
     max_gap: maxGap,
+    last_used_address_index: -1,
   };
   await mysql.query(
     `INSERT INTO \`wallet\`
@@ -304,6 +305,7 @@ export const createWallet = async (
     status: WalletStatus.CREATING,
     createdAt: ts,
     readyAt: null,
+    lastUsedAddressIndex: -1,
   };
 };
 
@@ -1322,32 +1324,27 @@ export const getWalletAddresses = async (mysql: ServerlessMysql, walletId: strin
  * @param walletId - Wallet id
  * @returns A list of addresses and their indexes
  */
-export const getNewAddresses = async (mysql: ServerlessMysql, walletId: string): Promise<ShortAddressInfo[]> => {
+export const getNewAddresses = async (mysql: ServerlessMysql, wallet: Wallet): Promise<ShortAddressInfo[]> => {
   const addresses: ShortAddressInfo[] = [];
-  const resultsWallet: DbSelectResult = await mysql.query('SELECT * FROM `wallet` WHERE `id` = ?', walletId);
-  if (resultsWallet.length) {
-    const gapLimit = resultsWallet[0].max_gap as number;
-    const latestUsedIndex = resultsWallet[0].last_used_address_index as number;
-    // Select all addresses that are empty and the index is bigger than the last used address index
-    const results: DbSelectResult = await mysql.query(`
-      SELECT *
-        FROM \`address\`
-       WHERE \`wallet_id\` = ?
-         AND \`transactions\` = 0
-         AND \`index\` > ?
-    ORDER BY \`index\`
-         ASC
-    LIMIT ?`, [walletId, latestUsedIndex, gapLimit]);
+  // Select all addresses that are empty and the index is bigger than the last used address index
+  const results: DbSelectResult = await mysql.query(`
+    SELECT *
+      FROM \`address\`
+     WHERE \`wallet_id\` = ?
+       AND \`transactions\` = 0
+       AND \`index\` > ?
+  ORDER BY \`index\`
+       ASC
+  LIMIT ?`, [wallet.walletId, wallet.lastUsedAddressIndex, wallet.maxGap]);
 
-    for (const result of results) {
-      const index = result.index as number;
-      const address = {
-        address: result.address as string,
-        index,
-        addressPath: getAddressPath(index),
-      };
-      addresses.push(address);
-    }
+  for (const result of results) {
+    const index = result.index as number;
+    const address = {
+      address: result.address as string,
+      index,
+      addressPath: getAddressPath(index),
+    };
+    addresses.push(address);
   }
   return addresses;
 };

packages/wallet-service/src/db/utils.ts

@@ -75,14 +75,15 @@ export async function transactionDecorator(_mysql: ServerlessMysql, wrapped: Fun
  * @param result - The result row to map to a Wallet object
  */
 export const getWalletFromDbEntry = (entry: Record<string, unknown>): Wallet => ({
-  walletId: getWalletId(entry.xpubkey as string),
+  walletId: entry.id as string,
   xpubkey: entry.xpubkey as string,
   authXpubkey: entry.auth_xpubkey as string,
   status: entry.status as WalletStatus,
   retryCount: entry.retry_count as number,
   maxGap: entry.max_gap as number,
   createdAt: entry.created_at as number,
   readyAt: entry.ready_at as number,
+  lastUsedAddressIndex: entry.last_used_address_index as number,
 });
 
 /**

packages/wallet-service/src/types.ts

@@ -144,6 +144,7 @@ export interface Wallet {
   retryCount?: number;
   createdAt?: number;
   readyAt?: number;
+  lastUsedAddressIndex?: number;
 }
 
 export interface AddressInfo {

packages/wallet-service/tests/api.test.ts

@@ -750,14 +750,15 @@ test('GET /wallet', async () => {
   expect(result.statusCode).toBe(200);
   expect(returnBody.success).toBe(true);
   expect(returnBody.status).toStrictEqual({
-    walletId: getWalletId(XPUBKEY),
+    walletId: 'my-wallet',
     xpubkey: XPUBKEY,
     authXpubkey: AUTH_XPUBKEY,
     status: 'ready',
     maxGap: 5,
     retryCount: 0,
     createdAt: 10000,
     readyAt: 10001,
+    lastUsedAddressIndex: -1,
   });
 });