Fix logout can take ages (#12191)

* Fix logout can take ages

* fix for of loop

* Add logout tests

* Unit test for logout behavior

* UserMenu tests update snapshot

playwright/e2e/crypto/logout.spec.ts

@@ -0,0 +1,110 @@
+/*
+Copyright 2024 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import { Page } from "@playwright/test";
+
+import { test, expect } from "../../element-web-test";
+import { logIntoElement } from "./utils";
+import { ElementAppPage } from "../../pages/ElementAppPage";
+
+test.describe("Logout tests", () => {
+    test.beforeEach(async ({ page, homeserver, credentials }) => {
+        await logIntoElement(page, homeserver, credentials);
+    });
+
+    async function createRoom(page: Page, roomName: string, isEncrypted: boolean): Promise<void> {
+        await page.getByRole("button", { name: "Add room" }).click();
+        await page.locator(".mx_IconizedContextMenu").getByRole("menuitem", { name: "New room" }).click();
+
+        const dialog = page.locator(".mx_Dialog");
+
+        await dialog.getByLabel("Name").fill(roomName);
+
+        if (!isEncrypted) {
+            // it's enabled by default
+            await page.getByLabel("Enable end-to-end encryption").click();
+        }
+
+        await dialog.getByRole("button", { name: "Create room" }).click();
+    }
+
+    async function sendMessageInCurrentRoom(page: Page, message: string): Promise<void> {
+        await page.locator(".mx_MessageComposer").getByRole("textbox").fill(message);
+        await page.getByTestId("sendmessagebtn").click();
+    }
+    async function setupRecovery(app: ElementAppPage, page: Page): Promise<void> {
+        const securityTab = await app.settings.openUserSettings("Security & Privacy");
+
+        await expect(securityTab.getByRole("heading", { name: "Secure Backup" })).toBeVisible();
+        await securityTab.getByRole("button", { name: "Set up", exact: true }).click();
+
+        const currentDialogLocator = page.locator(".mx_Dialog");
+
+        // It's the first time and secure storage is not set up, so it will create one
+        await expect(currentDialogLocator.getByRole("heading", { name: "Set up Secure Backup" })).toBeVisible();
+        await currentDialogLocator.getByRole("button", { name: "Continue", exact: true }).click();
+        await expect(currentDialogLocator.getByRole("heading", { name: "Save your Security Key" })).toBeVisible();
+        await currentDialogLocator.getByRole("button", { name: "Copy", exact: true }).click();
+        await currentDialogLocator.getByRole("button", { name: "Continue", exact: true }).click();
+
+        await expect(currentDialogLocator.getByRole("heading", { name: "Secure Backup successful" })).toBeVisible();
+        await currentDialogLocator.getByRole("button", { name: "Done", exact: true }).click();
+    }
+
+    test("Ask to set up recovery on logout if not setup", async ({ page, app }) => {
+        await createRoom(page, "E2e room", true);
+
+        // send a message (will be the first one so will create a new megolm session)
+        await sendMessageInCurrentRoom(page, "Hello secret world");
+
+        const locator = await app.settings.openUserMenu();
+        await locator.getByRole("menuitem", { name: "Sign out", exact: true }).click();
+
+        const currentDialogLocator = page.locator(".mx_Dialog");
+
+        await expect(
+            currentDialogLocator.getByRole("heading", { name: "You'll lose access to your encrypted messages" }),
+        ).toBeVisible();
+    });
+
+    test("If backup is set up show standard confirm", async ({ page, app }) => {
+        await setupRecovery(app, page);
+
+        await createRoom(page, "E2e room", true);
+
+        // send a message (will be the first one so will create a new megolm session)
+        await sendMessageInCurrentRoom(page, "Hello secret world");
+
+        const locator = await app.settings.openUserMenu();
+        await locator.getByRole("menuitem", { name: "Sign out", exact: true }).click();
+
+        const currentDialogLocator = page.locator(".mx_Dialog");
+
+        await expect(currentDialogLocator.getByText("Are you sure you want to sign out?")).toBeVisible();
+    });
+
+    test("Logout directly if the user has no room keys", async ({ page, app }) => {
+        await createRoom(page, "Clear room", false);
+
+        await sendMessageInCurrentRoom(page, "Hello public world!");
+
+        const locator = await app.settings.openUserMenu();
+        await locator.getByRole("menuitem", { name: "Sign out", exact: true }).click();
+
+        // Should have logged out directly
+        await expect(page.getByRole("heading", { name: "Sign in" })).toBeVisible();
+    });
+});

src/components/structures/UserMenu.tsx

@@ -258,17 +258,36 @@ export default class UserMenu extends React.Component<IProps, IState> {
         ev.preventDefault();
         ev.stopPropagation();
 
-        const cli = MatrixClientPeg.get();
+        if (await this.shouldShowLogoutDialog()) {
-        if (!cli || !cli.isCryptoEnabled() || !(await cli.exportRoomKeys())?.length) {
-            // log out without user prompt if they have no local megolm sessions
-            defaultDispatcher.dispatch({ action: "logout" });
-        } else {
             Modal.createDialog(LogoutDialog);
+        } else {
+            defaultDispatcher.dispatch({ action: "logout" });
         }
 
         this.setState({ contextMenuPosition: null }); // also close the menu
     };
 
+    /**
+     * Checks if the `LogoutDialog` should be shown instead of the simple logout flow.
+     * The `LogoutDialog` will check the crypto recovery status of the account and
+     * help the user setup recovery properly if needed.
+     * @private
+     */
+    private async shouldShowLogoutDialog(): Promise<boolean> {
+        const cli = MatrixClientPeg.get();
+        const crypto = cli?.getCrypto();
+        if (!crypto) return false;
+
+        // If any room is encrypted, we need to show the advanced logout flow
+        const allRooms = cli!.getRooms();
+        for (const room of allRooms) {
+            const isE2e = await crypto.isEncryptionEnabledInRoom(room.roomId);
+            if (isE2e) return true;
+        }
+
+        return false;
+    }
+
     private onSignInClick = (): void => {
         defaultDispatcher.dispatch({ action: "start_login" });
         this.setState({ contextMenuPosition: null }); // also close the menu

test/components/structures/UserMenu-test.tsx

@@ -15,8 +15,9 @@ limitations under the License.
 */
 
 import React from "react";
-import { act, render, RenderResult } from "@testing-library/react";
+import { act, render, RenderResult, screen, waitFor } from "@testing-library/react";
-import { MatrixClient, MatrixEvent } from "matrix-js-sdk/src/matrix";
+import { MatrixClient, MatrixEvent, Room } from "matrix-js-sdk/src/matrix";
+import { mocked } from "jest-mock";
 
 import UnwrappedUserMenu from "../../../src/components/structures/UserMenu";
 import { stubClient, wrapInSdkContext } from "../../test-utils";
@@ -27,11 +28,20 @@ import {
 } from "../../../src/voice-broadcast";
 import { mkVoiceBroadcastInfoStateEvent } from "../../voice-broadcast/utils/test-utils";
 import { TestSdkContext } from "../../TestSdkContext";
+import defaultDispatcher from "../../../src/dispatcher/dispatcher";
+import LogoutDialog from "../../../src/components/views/dialogs/LogoutDialog";
+import Modal from "../../../src/Modal";
 
 describe("<UserMenu>", () => {
     let client: MatrixClient;
     let renderResult: RenderResult;
     let sdkContext: TestSdkContext;
+
+    beforeEach(() => {
+        sdkContext = new TestSdkContext();
+    });
+
+    describe("<UserMenu> when video broadcast", () => {
         let voiceBroadcastInfoEvent: MatrixEvent;
         let voiceBroadcastRecording: VoiceBroadcastRecording;
         let voiceBroadcastRecordingsStore: VoiceBroadcastRecordingsStore;
@@ -47,7 +57,6 @@ describe("<UserMenu>", () => {
         });
 
         beforeEach(() => {
-        sdkContext = new TestSdkContext();
             voiceBroadcastRecordingsStore = new VoiceBroadcastRecordingsStore();
             sdkContext._VoiceBroadcastRecordingsStore = voiceBroadcastRecordingsStore;
 
@@ -89,3 +98,83 @@ describe("<UserMenu>", () => {
             });
         });
     });
+
+    describe("<UserMenu> logout", () => {
+        beforeEach(() => {
+            client = stubClient();
+        });
+
+        it("should logout directly if no crypto", async () => {
+            const UserMenu = wrapInSdkContext(UnwrappedUserMenu, sdkContext);
+            renderResult = render(<UserMenu isPanelCollapsed={true} />);
+
+            mocked(client.getRooms).mockReturnValue([
+                {
+                    roomId: "!room0",
+                } as unknown as Room,
+                {
+                    roomId: "!room1",
+                } as unknown as Room,
+            ]);
+            jest.spyOn(client, "getCrypto").mockReturnValue(undefined);
+
+            const spy = jest.spyOn(defaultDispatcher, "dispatch");
+            screen.getByRole("button", { name: /User menu/i }).click();
+            screen.getByRole("menuitem", { name: /Sign out/i }).click();
+            await waitFor(() => {
+                expect(spy).toHaveBeenCalledWith({ action: "logout" });
+            });
+        });
+
+        it("should logout directly if no encrypted rooms", async () => {
+            const UserMenu = wrapInSdkContext(UnwrappedUserMenu, sdkContext);
+            renderResult = render(<UserMenu isPanelCollapsed={true} />);
+
+            mocked(client.getRooms).mockReturnValue([
+                {
+                    roomId: "!room0",
+                } as unknown as Room,
+                {
+                    roomId: "!room1",
+                } as unknown as Room,
+            ]);
+            const crypto = client.getCrypto()!;
+
+            jest.spyOn(crypto, "isEncryptionEnabledInRoom").mockResolvedValue(false);
+
+            const spy = jest.spyOn(defaultDispatcher, "dispatch");
+            screen.getByRole("button", { name: /User menu/i }).click();
+            screen.getByRole("menuitem", { name: /Sign out/i }).click();
+            await waitFor(() => {
+                expect(spy).toHaveBeenCalledWith({ action: "logout" });
+            });
+        });
+
+        it("should show dialog if some encrypted rooms", async () => {
+            const UserMenu = wrapInSdkContext(UnwrappedUserMenu, sdkContext);
+            renderResult = render(<UserMenu isPanelCollapsed={true} />);
+
+            mocked(client.getRooms).mockReturnValue([
+                {
+                    roomId: "!room0",
+                } as unknown as Room,
+                {
+                    roomId: "!room1",
+                } as unknown as Room,
+            ]);
+            const crypto = client.getCrypto()!;
+
+            jest.spyOn(crypto, "isEncryptionEnabledInRoom").mockImplementation(async (roomId: string) => {
+                return roomId === "!room0";
+            });
+
+            const spy = jest.spyOn(Modal, "createDialog");
+            screen.getByRole("button", { name: /User menu/i }).click();
+            screen.getByRole("menuitem", { name: /Sign out/i }).click();
+
+            await waitFor(() => {
+                expect(spy).toHaveBeenCalledWith(LogoutDialog);
+            });
+        });
+    });
+});

test/components/structures/__snapshots__/UserMenu-test.tsx.snap

@@ -1,6 +1,6 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`<UserMenu> when rendered should render as expected 1`] = `
+exports[`<UserMenu> <UserMenu> when video broadcast when rendered should render as expected 1`] = `
 <div>
   <div
     class="mx_UserMenu"

test/test-utils/test-utils.ts

@@ -137,6 +137,7 @@ export function createTestClient(): MatrixClient {
             getUserVerificationStatus: jest.fn(),
             getDeviceVerificationStatus: jest.fn(),
             resetKeyBackup: jest.fn(),
+            isEncryptionEnabledInRoom: jest.fn(),
         }),
 
         getPushActionsForEvent: jest.fn(),