From f9453c64dd1fc2324cb00ddc53919973b41540b1 Mon Sep 17 00:00:00 2001
From: Fedor Indutny <79877362+indutny-signal@users.noreply.github.com>
Date: Mon, 19 Sep 2022 15:08:55 -0700
Subject: [PATCH] Test rate-limiting, stories in mock server

---
 package.json                          |   2 +-
 ts/CI.ts                              |   5 +
 ts/background.ts                      |   4 +
 ts/components/SendStoryModal.tsx      |   4 +-
 ts/services/storage.ts                |   1 +
 ts/test-mock/bootstrap.ts             |  27 ++++-
 ts/test-mock/playwright.ts            |  17 +++
 ts/test-mock/rate-limit/story_test.ts | 154 ++++++++++++++++++++++++++
 yarn.lock                             |   8 +-
 9 files changed, 212 insertions(+), 10 deletions(-)
 create mode 100644 ts/test-mock/rate-limit/story_test.ts

diff --git a/package.json b/package.json
index ae59f4f74..128ba4b87 100644
--- a/package.json
+++ b/package.json
@@ -193,7 +193,7 @@
     "@babel/preset-typescript": "7.17.12",
     "@electron/fuses": "1.5.0",
     "@mixer/parallel-prettier": "2.0.1",
-    "@signalapp/mock-server": "2.8.1",
+    "@signalapp/mock-server": "2.9.0",
     "@storybook/addon-a11y": "6.5.6",
     "@storybook/addon-actions": "6.5.6",
     "@storybook/addon-controls": "6.5.6",
diff --git a/ts/CI.ts b/ts/CI.ts
index 145799cd7..63d8bfc30 100644
--- a/ts/CI.ts
+++ b/ts/CI.ts
@@ -6,6 +6,7 @@ import { ipcRenderer } from 'electron';
 import { explodePromise } from './util/explodePromise';
 import { SECOND } from './util/durations';
 import * as log from './logging/log';
+import type { IPCResponse as ChallengeResponseType } from './challenge';
 
 type ResolveType = (data: unknown) => void;
 
@@ -84,4 +85,8 @@ export class CI {
     }
     resultList.push(data);
   }
+
+  public solveChallenge(response: ChallengeResponseType): void {
+    window.Signal.challengeHandler?.onResponse(response);
+  }
 }
diff --git a/ts/background.ts b/ts/background.ts
index fdbe76eb2..82fd991e5 100644
--- a/ts/background.ts
+++ b/ts/background.ts
@@ -238,6 +238,10 @@ export async function startApp(): Promise<void> {
       },
 
       requestChallenge(request) {
+        if (window.CI) {
+          window.CI.handleEvent('challenge', request);
+          return;
+        }
         window.sendChallengeRequest(request);
       },
 
diff --git a/ts/components/SendStoryModal.tsx b/ts/components/SendStoryModal.tsx
index 5ec81ec07..56275ba66 100644
--- a/ts/components/SendStoryModal.tsx
+++ b/ts/components/SendStoryModal.tsx
@@ -692,7 +692,7 @@ export const SendStoryModal = ({
         )}
         {page === Page.ChooseGroups && (
           <button
-            aria-label="SendStoryModal__ok"
+            aria-label={i18n('SendStoryModal__ok')}
             className="SendStoryModal__ok"
             disabled={!chosenGroupIds.size}
             onClick={() => {
@@ -705,7 +705,7 @@ export const SendStoryModal = ({
         )}
         {page === Page.SendStory && (
           <button
-            aria-label="SendStoryModal__send"
+            aria-label={i18n('SendStoryModal__send')}
             className="SendStoryModal__send"
             disabled={!selectedListIds.size && !selectedGroupIds.size}
             onClick={() => {
diff --git a/ts/services/storage.ts b/ts/services/storage.ts
index db6789392..b03f32d0c 100644
--- a/ts/services/storage.ts
+++ b/ts/services/storage.ts
@@ -1270,6 +1270,7 @@ async function processManifest(
     );
 
     if (!myStories) {
+      log.info(`storageService.process(${version}): creating my stories`);
       const storyDistribution: StoryDistributionWithMembersType = {
         allowsReplies: true,
         id: MY_STORIES_ID,
diff --git a/ts/test-mock/bootstrap.ts b/ts/test-mock/bootstrap.ts
index d04e45b78..df6e233c1 100644
--- a/ts/test-mock/bootstrap.ts
+++ b/ts/test-mock/bootstrap.ts
@@ -62,6 +62,7 @@ export type BootstrapOptions = Readonly<{
 
   linkedDevices?: number;
   contactCount?: number;
+  contactsWithoutProfileKey?: number;
   contactNames?: ReadonlyArray<string>;
   contactPreKeyCount?: number;
 }>;
@@ -71,6 +72,7 @@ type BootstrapInternalOptions = Pick<BootstrapOptions, 'extraConfig'> &
     benchmark: boolean;
     linkedDevices: number;
     contactCount: number;
+    contactsWithoutProfileKey: number;
     contactNames: ReadonlyArray<string>;
   }>;
 
@@ -104,6 +106,7 @@ export class Bootstrap {
 
   private readonly options: BootstrapInternalOptions;
   private privContacts?: ReadonlyArray<PrimaryDevice>;
+  private privContactsWithoutProfileKey?: ReadonlyArray<PrimaryDevice>;
   private privPhone?: PrimaryDevice;
   private privDesktop?: Device;
   private storagePath?: string;
@@ -118,13 +121,17 @@ export class Bootstrap {
     this.options = {
       linkedDevices: 5,
       contactCount: MAX_CONTACTS,
+      contactsWithoutProfileKey: 0,
       contactNames: CONTACT_NAMES,
       benchmark: false,
 
       ...options,
     };
 
-    assert(this.options.contactCount <= this.options.contactNames.length);
+    assert(
+      this.options.contactCount + this.options.contactsWithoutProfileKey <=
+        this.options.contactNames.length
+    );
   }
 
   public async init(): Promise<void> {
@@ -137,10 +144,10 @@ export class Bootstrap {
 
     const contactNames = this.options.contactNames.slice(
       0,
-      this.options.contactCount
+      this.options.contactCount + this.options.contactsWithoutProfileKey
     );
 
-    this.privContacts = await Promise.all(
+    const allContacts = await Promise.all(
       contactNames.map(async profileName => {
         const primary = await this.server.createPrimaryDevice({
           profileName,
@@ -155,9 +162,15 @@ export class Bootstrap {
       })
     );
 
+    this.privContacts = allContacts.slice(0, this.options.contactCount);
+    this.privContactsWithoutProfileKey = allContacts.slice(
+      this.contacts.length
+    );
+
     this.privPhone = await this.server.createPrimaryDevice({
       profileName: 'Myself',
       contacts: this.contacts,
+      contactsWithoutProfileKey: this.contactsWithoutProfileKey,
     });
 
     this.storagePath = await fs.mkdtemp(path.join(os.tmpdir(), 'mock-signal-'));
@@ -309,6 +322,14 @@ export class Bootstrap {
     return this.privContacts;
   }
 
+  public get contactsWithoutProfileKey(): ReadonlyArray<PrimaryDevice> {
+    assert(
+      this.privContactsWithoutProfileKey,
+      'Bootstrap has to be initialized first, see: bootstrap.init()'
+    );
+    return this.privContactsWithoutProfileKey;
+  }
+
   //
   // Private
   //
diff --git a/ts/test-mock/playwright.ts b/ts/test-mock/playwright.ts
index c991f4214..e860d142c 100644
--- a/ts/test-mock/playwright.ts
+++ b/ts/test-mock/playwright.ts
@@ -4,6 +4,11 @@
 import type { ElectronApplication, Page } from 'playwright';
 import { _electron as electron } from 'playwright';
 
+import type {
+  IPCRequest as ChallengeRequestType,
+  IPCResponse as ChallengeResponseType,
+} from '../challenge';
+
 export type AppLoadedInfoType = Readonly<{
   loadTime: number;
   messagesPerSec: number;
@@ -57,6 +62,18 @@ export class App {
     return this.waitForEvent('conversation:open');
   }
 
+  public async waitForChallenge(): Promise<ChallengeRequestType> {
+    return this.waitForEvent('challenge');
+  }
+
+  public async solveChallenge(response: ChallengeResponseType): Promise<void> {
+    const window = await this.getWindow();
+
+    await window.evaluate(
+      `window.CI.solveChallenge(${JSON.stringify(response)})`
+    );
+  }
+
   public async close(): Promise<void> {
     await this.app.close();
   }
diff --git a/ts/test-mock/rate-limit/story_test.ts b/ts/test-mock/rate-limit/story_test.ts
new file mode 100644
index 000000000..f046ab5b1
--- /dev/null
+++ b/ts/test-mock/rate-limit/story_test.ts
@@ -0,0 +1,154 @@
+// Copyright 2022 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { assert } from 'chai';
+import createDebug from 'debug';
+import { Proto, StorageState } from '@signalapp/mock-server';
+
+import * as durations from '../../util/durations';
+import { uuidToBytes } from '../../util/uuidToBytes';
+import { MY_STORIES_ID } from '../../types/Stories';
+import { Bootstrap } from '../bootstrap';
+import type { App } from '../bootstrap';
+
+export const debug = createDebug('mock:test:rate-limit');
+
+const IdentifierType = Proto.ManifestRecord.Identifier.Type;
+
+describe('rate-limit/story', function needsName() {
+  this.timeout(durations.MINUTE);
+
+  let bootstrap: Bootstrap;
+  let app: App;
+
+  beforeEach(async () => {
+    bootstrap = new Bootstrap({
+      contactCount: 0,
+      contactsWithoutProfileKey: 40,
+    });
+    await bootstrap.init();
+
+    const { phone } = bootstrap;
+
+    let state = StorageState.getEmpty();
+
+    state = state.updateAccount({
+      profileKey: phone.profileKey.serialize(),
+      e164: phone.device.number,
+      givenName: phone.profileName,
+      hasSetMyStoriesPrivacy: true,
+    });
+
+    state = state.addRecord({
+      type: IdentifierType.STORY_DISTRIBUTION_LIST,
+      record: {
+        storyDistributionList: {
+          allowsReplies: true,
+          identifier: uuidToBytes(MY_STORIES_ID),
+          isBlockList: true,
+          name: MY_STORIES_ID,
+          recipientUuids: [],
+        },
+      },
+    });
+    phone.setStorageState(state);
+
+    app = await bootstrap.link();
+  });
+
+  afterEach(async function after() {
+    if (this.currentTest?.state !== 'passed') {
+      await bootstrap.saveLogs(app);
+    }
+
+    await app.close();
+    await bootstrap.teardown();
+  });
+
+  it('should request challenge and accept solution', async () => {
+    const {
+      server,
+      contactsWithoutProfileKey: contacts,
+      phone,
+      desktop,
+    } = bootstrap;
+
+    for (const contact of contacts) {
+      server.rateLimit({ source: desktop.uuid, target: contact.device.uuid });
+    }
+
+    const window = await app.getWindow();
+
+    debug('Posting a new story');
+    {
+      const storiesPane = window.locator('.Stories');
+
+      await window.locator('button.module-main-header__stories-icon').click();
+
+      await storiesPane
+        .locator('button.Stories__pane__add-story__button')
+        .click();
+      await storiesPane
+        .locator(
+          '.ContextMenu__popper .Stories__pane__add-story__option--title ' +
+            '>> "Text story"'
+        )
+        .click();
+
+      debug('Focusing textarea');
+      await storiesPane.locator('.TextAttachment__story').click();
+
+      debug('Entering text');
+      await storiesPane.locator('.TextAttachment__text__textarea').type('123');
+
+      debug('Clicking "Next"');
+      await storiesPane
+        .locator('.StoryCreator__toolbar button >> "Next"')
+        .click();
+
+      debug('Selecting "My Stories"');
+      await window
+        .locator('.SendStoryModal__distribution-list__name >> "My Stories"')
+        .click();
+
+      debug('Hitting Send');
+      await window.locator('button.SendStoryModal__send').click();
+    }
+
+    debug('Waiting for challenge');
+    const request = await app.waitForChallenge();
+
+    debug('Checking for presence of captcha modal');
+    await window
+      .locator('.module-Modal__title >> "Verify to continue messaging"')
+      .waitFor();
+
+    debug('Removing rate-limiting');
+    for (const contact of contacts) {
+      const failedMessages = server.stopRateLimiting({
+        source: desktop.uuid,
+        target: contact.device.uuid,
+      });
+      assert.isAtMost(failedMessages ?? 0, 1);
+    }
+
+    debug('Solving challenge');
+    app.solveChallenge({
+      seq: request.seq,
+      data: { captcha: 'anything' },
+    });
+
+    debug('Verifying that all contacts received story');
+    await Promise.all(
+      contacts.map(async contact => {
+        const { storyMessage } = await contact.waitForStory();
+        assert.isTrue(
+          phone.profileKey
+            .serialize()
+            .equals(storyMessage.profileKey ?? new Uint8Array(0))
+        );
+        assert.strictEqual(storyMessage.textAttachment?.text, '123');
+      })
+    );
+  });
+});
diff --git a/yarn.lock b/yarn.lock
index 2027cc6e1..7ceb066f1 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -1893,10 +1893,10 @@
     node-gyp-build "^4.2.3"
     uuid "^8.3.0"
 
-"@signalapp/mock-server@2.8.1":
-  version "2.8.1"
-  resolved "https://registry.yarnpkg.com/@signalapp/mock-server/-/mock-server-2.8.1.tgz#f09c35de6fd084b0a654e2d20fd77095f7e5fca0"
-  integrity sha512-kQkzgtnVpFO9J00MpiTHBtKsyQhe/8PHy8nxJmIR3L2IQvZ8glbX4GCVh7dRIASY3IIJ8PpiAoYgePkL0DH8Gw==
+"@signalapp/mock-server@2.9.0":
+  version "2.9.0"
+  resolved "https://registry.yarnpkg.com/@signalapp/mock-server/-/mock-server-2.9.0.tgz#7b86b7843cd828db5167d9152214646cabe76724"
+  integrity sha512-rowEcBwGb53C7gQTR9WTQrS3kO4WB43qTVlA2qeEUIzMruJb44klgrmo1QUiMTtGYuByDApbAC4yOcVNYmae4w==
   dependencies:
     "@signalapp/libsignal-client" "^0.20.0"
     debug "^4.3.2"