Improve mock-server orphaned storage key check

Co-authored-by: trevor-signal <131492920+trevor-signal@users.noreply.github.com>

.changeset/olive-poems-repeat.md

@@ -0,0 +1,5 @@
+---
+'@signalapp/mock-server': minor
+---
+
+Make the storage item/manifest accessors synchronous so that writes can't interleave with reads

packages/mock-server/src/api/primary-device.ts

@@ -117,11 +117,9 @@ export type Config = Readonly<{
   modifyGroup: (options: ModifyGroupOptions) => Promise<ModifyGroupResult>;
   waitForGroupUpdate: (group: GroupData) => Promise<void>;
 
-  getStorageManifest: () => Promise<Proto.StorageManifest.Params | undefined>;
-  getStorageItem: (
-    key: Buffer<ArrayBuffer>,
-  ) => Promise<Buffer<ArrayBuffer> | undefined>;
-  getAllStorageKeys: () => Promise<Array<Buffer<ArrayBuffer>>>;
+  getStorageManifest: () => Proto.StorageManifest.Params | undefined;
+  getStorageItem: (key: Buffer<ArrayBuffer>) => Buffer<ArrayBuffer> | undefined;
+  getAllStorageKeys: () => Array<Buffer<ArrayBuffer>>;
   waitForStorageManifest: (afterVersion?: bigint) => Promise<void>;
   applyStorageWrite: (
     operation: Proto.WriteOperation.Params,
@@ -1187,7 +1185,7 @@ export class PrimaryDevice {
   }
 
   public async getStorageState(): Promise<StorageState | undefined> {
-    const manifest = await this.config.getStorageManifest();
+    const manifest = this.config.getStorageManifest();
     if (!manifest) {
       return undefined;
     }
@@ -1228,13 +1226,13 @@ export class PrimaryDevice {
   }
 
   public async getOrphanedStorageKeys(): Promise<Array<Buffer<ArrayBuffer>>> {
-    const manifest = await this.config.getStorageManifest();
+    const manifest = this.config.getStorageManifest();
     if (!manifest) {
       return [];
     }
 
-    const state = await this.convertManifestToStorageState(manifest);
-    const keys = await this.config.getAllStorageKeys();
+    const state = this.convertManifestToStorageState(manifest);
+    const keys = this.config.getAllStorageKeys();
 
     return keys.filter((key) => !state.hasKey(key));
   }
@@ -2453,41 +2451,37 @@ export class PrimaryDevice {
     );
   }
 
-  private async convertManifestToStorageState(
+  private convertManifestToStorageState(
     manifest: Proto.StorageManifest.Params,
-  ): Promise<StorageState> {
+  ): StorageState {
     const decryptedManifest = decryptStorageManifest(this.storageKey, manifest);
     assert(decryptedManifest.version, 'Consistency check');
 
     const version = decryptedManifest.version;
-    const items = await Promise.all(
-      decryptedManifest.identifiers.map(async ({ type, raw: key }) => {
-        const keyBuffer = Buffer.from(key);
-        const item = await this.config.getStorageItem(keyBuffer);
-        if (!item) {
-          throw new Error(`Missing item ${keyBuffer.toString('base64')}`);
-        }
+    const items = decryptedManifest.identifiers.map(({ type, raw: key }) => {
+      const keyBuffer = Buffer.from(key);
+      const item = this.config.getStorageItem(keyBuffer);
+      if (!item) {
+        throw new Error(`Missing item ${keyBuffer.toString('base64')}`);
+      }
 
-        const decrypted = decryptStorageItem({
-          storageKey: this.storageKey,
-          recordIkm: this.storageRecordIkm,
-          item: {
-            key,
-            value: item,
-          },
-        });
-        if (!decrypted.record) {
-          throw new Error(
-            `Missing item record ${keyBuffer.toString('base64')}`,
-          );
-        }
-        return {
-          type: type as Proto.ManifestRecord.Identifier.Type,
-          key: keyBuffer,
-          record: decrypted.record,
-        };
-      }),
-    );
+      const decrypted = decryptStorageItem({
+        storageKey: this.storageKey,
+        recordIkm: this.storageRecordIkm,
+        item: {
+          key,
+          value: item,
+        },
+      });
+      if (!decrypted.record) {
+        throw new Error(`Missing item record ${keyBuffer.toString('base64')}`);
+      }
+      return {
+        type: type as Proto.ManifestRecord.Identifier.Type,
+        key: keyBuffer,
+        record: decrypted.record,
+      };
+    });
 
     return new StorageState(version, items);
   }

packages/mock-server/src/api/server.ts

@@ -856,10 +856,10 @@ export class Server extends BaseServer {
 
   // Override `getStorageItems` to provide configurable limit for maximum
   // storage read keys.
-  public override async getStorageItems(
+  public override getStorageItems(
     device: Device,
     keys: ReadonlyArray<Buffer<ArrayBuffer>>,
-  ): Promise<Array<Proto.StorageItem.Params> | undefined> {
+  ): Array<Proto.StorageItem.Params> | undefined {
     if (
       this.config.maxStorageReadKeys !== undefined &&
       keys.length > this.config.maxStorageReadKeys

packages/mock-server/src/server/base.ts

@@ -1038,9 +1038,9 @@ export abstract class Server {
     return auth.device;
   }
 
-  public async getStorageManifest(
+  public getStorageManifest(
     device: Device,
-  ): Promise<Proto.StorageManifest.Params | undefined> {
+  ): Proto.StorageManifest.Params | undefined {
     return this.storageManifestByAci.get(device.aci);
   }
 
@@ -1056,7 +1056,7 @@ export abstract class Server {
       return { error: 'missing `writeOperation.manifest.version`' };
     }
 
-    const existing = await this.getStorageManifest(device);
+    const existing = this.getStorageManifest(device);
     if (existing) {
       // Atomicity
       assert(existing.version, 'consistency check');
@@ -1072,27 +1072,25 @@ export abstract class Server {
 
     if (clearAll) {
       debug('clearing storage items for=%j', device.debugId);
-      await this.clearStorageItems(device);
+      this.clearStorageItems(device);
     }
 
-    const inserts = (insertItem ?? []).map(async (item) => {
+    for (const item of insertItem ?? []) {
       assert(item.key instanceof Uint8Array, 'insertItem.key must be a Buffer');
       assert(
         item.value instanceof Uint8Array,
         'insertItem.value must be a Buffer',
       );
-      return this.setStorageItem(
+      this.setStorageItem(
         device,
         Buffer.from(item.key),
         Buffer.from(item.value),
       );
-    });
-    await Promise.all(inserts);
+    }
 
-    const deletes = (deleteKey ?? []).map(async (key) => {
-      return this.deleteStorageItem(device, Buffer.from(key));
-    });
-    await Promise.all(deletes);
+    for (const key of deleteKey ?? []) {
+      this.deleteStorageItem(device, Buffer.from(key));
+    }
 
     debug(
       'updating storage manifest to version=%d for=%j',
@@ -1108,15 +1106,15 @@ export abstract class Server {
     return { updated: true };
   }
 
-  private async clearStorageItems(device: Device): Promise<void> {
+  private clearStorageItems(device: Device): void {
     this.storageItemsByAci.get(device.aci)?.clear();
   }
 
-  private async setStorageItem(
+  private setStorageItem(
     device: Device,
     key: Buffer<ArrayBuffer>,
     value: Buffer<ArrayBuffer>,
-  ): Promise<void> {
+  ): void {
     let map = this.storageItemsByAci.get(device.aci);
     if (!map) {
       map = new Map();
@@ -1126,10 +1124,10 @@ export abstract class Server {
     map.set(key.toString('hex'), value);
   }
 
-  public async getStorageItem(
+  public getStorageItem(
     device: Device,
     key: Buffer<ArrayBuffer>,
-  ): Promise<Buffer<ArrayBuffer> | undefined> {
+  ): Buffer<ArrayBuffer> | undefined {
     const map = this.storageItemsByAci.get(device.aci);
     if (!map) {
       return undefined;
@@ -1138,9 +1136,7 @@ export abstract class Server {
     return map.get(key.toString('hex'));
   }
 
-  public async getAllStorageKeys(
-    device: Device,
-  ): Promise<Array<Buffer<ArrayBuffer>>> {
+  public getAllStorageKeys(device: Device): Array<Buffer<ArrayBuffer>> {
     const map = this.storageItemsByAci.get(device.aci);
     if (!map) {
       return [];
@@ -1149,28 +1145,23 @@ export abstract class Server {
     return Array.from(map.keys()).map((hex) => Buffer.from(hex, 'hex'));
   }
 
-  public async getStorageItems(
+  public getStorageItems(
     device: Device,
     keys: ReadonlyArray<Buffer<ArrayBuffer>>,
-  ): Promise<Array<Proto.StorageItem.Params> | undefined> {
+  ): Array<Proto.StorageItem.Params> | undefined {
     const result = new Array<Proto.StorageItem.Params>();
 
-    await Promise.all(
-      keys.map(async (key) => {
-        const value = await this.getStorageItem(device, key);
-        if (value !== undefined) {
-          result.push({ key, value });
-        }
-      }),
-    );
+    for (const key of keys) {
+      const value = this.getStorageItem(device, key);
+      if (value !== undefined) {
+        result.push({ key, value });
+      }
+    }
 
     return result;
   }
 
-  public async deleteStorageItem(
-    device: Device,
-    key: Buffer<ArrayBuffer>,
-  ): Promise<void> {
+  public deleteStorageItem(device: Device, key: Buffer<ArrayBuffer>): void {
     const map = this.storageItemsByAci.get(device.aci);
     if (!map) {
       return;

packages/mock-server/src/server/http.ts

@@ -848,7 +848,7 @@ export const createHandler = (
       return;
     }
 
-    const manifest = await server.getStorageManifest(device);
+    const manifest = server.getStorageManifest(device);
     if (!manifest) {
       return send(res, 404, { error: 'Manifest not found' });
     }
@@ -866,7 +866,7 @@ export const createHandler = (
 
       assert(req.params.after != null, 'Missing after param');
       const after = BigInt(req.params.after);
-      const manifest = await server.getStorageManifest(device);
+      const manifest = server.getStorageManifest(device);
       if (manifest === undefined) {
         return send(res, 404);
       }
@@ -912,7 +912,7 @@ export const createHandler = (
 
     const keys = readOperation.readKey.map((key) => Buffer.from(key));
 
-    const items = await server.getStorageItems(device, keys);
+    const items = server.getStorageItems(device, keys);
     if (!items) {
       return send(res, 413, { error: 'Requested too many items' });
     }

packages/mock-server/test/primary-device-test.ts

@@ -103,13 +103,13 @@ async function createPrimaryDevice(name: string): Promise<PrimaryDevice> {
     async waitForGroupUpdate() {
       throw new Error('Not implemented');
     },
-    async getStorageManifest() {
+    getStorageManifest() {
       throw new Error('Not implemented');
     },
-    async getStorageItem() {
+    getStorageItem() {
       throw new Error('Not implemented');
     },
-    async getAllStorageKeys() {
+    getAllStorageKeys() {
       throw new Error('Not implemented');
     },
     async waitForStorageManifest() {