Simplify key local storage management.

Author: toger5

src/e2ee/sharedKeyManagement.ts

@@ -8,56 +8,52 @@ Please see LICENSE in the repository root for full details.
 import { useEffect, useMemo } from "react";
 
 import { setLocalStorageItem, useLocalStorage } from "../useLocalStorage";
-import { type UrlParams, getUrlParams, useUrlParams } from "../UrlParams";
+import { getUrlParams } from "../UrlParams";
 import { E2eeType } from "./e2eeType";
 import { useClient } from "../ClientContext";
 
+/**
+ * This setter will update the state for all `useRoomSharedKey` hooks
+ */
 export function saveKeyForRoom(roomId: string, password: string): void {
   setLocalStorageItem(getRoomSharedKeyLocalStorageKey(roomId), password);
 }
 
 const getRoomSharedKeyLocalStorageKey = (roomId: string): string =>
   `room-shared-key-${roomId}`;
 
-const useInternalRoomSharedKey = (roomId: string): string | null => {
-  const key = getRoomSharedKeyLocalStorageKey(roomId);
-  const [roomSharedKey] = useLocalStorage(key);
-
-  return roomSharedKey;
+/**
+ * An upto-date shared key for the room. Either from local storage or the value from `setInitialValue`.
+ * @param roomId
+ * @param setInitialValue The value we get from the URL. The hook will overwrite the local storage value with this.
+ * @returns [roomSharedKey, setRoomSharedKey] like a react useState hook.
+ */
+const useRoomSharedKey = (
+  roomId: string,
+  setInitialValue?: string,
+): [string | null, setKey: (key: string) => void] => {
+  const [roomSharedKey, setRoomSharedKey] = useLocalStorage(
+    getRoomSharedKeyLocalStorageKey(roomId),
+  );
+  useEffect(() => {
+    // If setInitialValue is available, update the local storage (usually the password from the url).
+    // This will update roomSharedKey but wont update the returned value since
+    // that already defaults to setInitialValue.
+    if (setInitialValue) setRoomSharedKey(setInitialValue);
+  }, [setInitialValue, setRoomSharedKey]);
+
+  // make sure we never return the initial null value from `useLocalStorage`
+  return [setInitialValue ?? roomSharedKey, setRoomSharedKey];
 };
 
 export function getKeyForRoom(roomId: string): string | null {
-  saveKeyFromUrlParams(getUrlParams());
-  const key = getRoomSharedKeyLocalStorageKey(roomId);
-  return localStorage.getItem(key);
+  const { roomId: urlRoomId, password } = getUrlParams();
+  if (roomId !== urlRoomId) return null;
+  return (
+    password ?? localStorage.getItem(getRoomSharedKeyLocalStorageKey(roomId))
+  );
 }
 
-function saveKeyFromUrlParams(urlParams: UrlParams): void {
-  if (!urlParams.password || !urlParams.roomId) return;
-
-  // Take the key from the URL and save it.
-  // It's important to always use the room ID specified in the URL
-  // when saving keys rather than whatever the current room ID might be,
-  // in case we've moved to a different room but the URL hasn't changed.
-  saveKeyForRoom(urlParams.roomId, urlParams.password);
-}
-
-/**
- * Extracts the room password from the URL if one is present, saving it in localstorage
- * and returning it in a tuple with the corresponding room ID from the URL.
- * @returns A tuple of the roomId and password from the URL if the URL has both,
- *          otherwise [undefined, undefined]
- */
-const useKeyFromUrl = (): [string, string] | [undefined, undefined] => {
-  const urlParams = useUrlParams();
-
-  useEffect(() => saveKeyFromUrlParams(urlParams), [urlParams]);
-
-  return urlParams.roomId && urlParams.password
-    ? [urlParams.roomId, urlParams.password]
-    : [undefined, undefined];
-};
-
 export type Unencrypted = { kind: E2eeType.NONE };
 export type SharedSecret = { kind: E2eeType.SHARED_KEY; secret: string };
 export type PerParticipantE2EE = { kind: E2eeType.PER_PARTICIPANT };
@@ -66,12 +62,11 @@ export type EncryptionSystem = Unencrypted | SharedSecret | PerParticipantE2EE;
 export function useRoomEncryptionSystem(roomId: string): EncryptionSystem {
   const { client } = useClient();
 
-  // make sure we've extracted the key from the URL first
-  // (and we still need to take the value it returns because
-  // the effect won't run in time for it to save to localstorage in
-  // time for us to read it out again).
-  const [urlRoomId, passwordFromUrl] = useKeyFromUrl();
-  const storedPassword = useInternalRoomSharedKey(roomId);
+  const [storedPassword] = useRoomSharedKey(
+    getRoomSharedKeyLocalStorageKey(roomId),
+    getKeyForRoom(roomId) ?? undefined,
+  );
+
   const room = client?.getRoom(roomId);
   const e2eeSystem = <EncryptionSystem>useMemo(() => {
     if (!room) return { kind: E2eeType.NONE };
@@ -80,15 +75,10 @@ export function useRoomEncryptionSystem(roomId: string): EncryptionSystem {
         kind: E2eeType.SHARED_KEY,
         secret: storedPassword,
       };
-    if (urlRoomId === roomId)
-      return {
-        kind: E2eeType.SHARED_KEY,
-        secret: passwordFromUrl,
-      };
     if (room.hasEncryptionStateEvent()) {
       return { kind: E2eeType.PER_PARTICIPANT };
     }
     return { kind: E2eeType.NONE };
-  }, [passwordFromUrl, room, roomId, storedPassword, urlRoomId]);
+  }, [room, storedPassword]);
   return e2eeSystem;
 }

src/home/useGroupCallRooms.ts

@@ -21,7 +21,7 @@ import {
   type MatrixRTCSession,
 } from "matrix-js-sdk/lib/matrixrtc";
 
-import { getKeyForRoom } from "../e2ee/sharedKeyManagement";
+import { getKeyForRoom, saveKeyForRoom } from "../e2ee/sharedKeyManagement";
 
 export interface GroupCallRoom {
   roomAlias?: string;
@@ -81,11 +81,13 @@ function sortRooms(client: MatrixClient, rooms: Room[]): Room[] {
 }
 
 const roomIsJoinable = (room: Room): boolean => {
-  if (!room.hasEncryptionStateEvent() && !getKeyForRoom(room.roomId)) {
-    // if we have an non encrypted room (no encryption state event) we need a locally stored shared key.
+  const password = getKeyForRoom(room.roomId);
+  if (!room.hasEncryptionStateEvent() && !password) {
+    // if we have a non encrypted room (no encryption state event) we need a locally stored shared key.
     // in case this key also does not exists we cannot join the room.
     return false;
   }
+  if (password) saveKeyForRoom(room.roomId, password);
   // otherwise we can always join rooms because we will automatically decide if we want to use perParticipant or password
   switch (room.getJoinRule()) {
     case JoinRule.Public:

src/useLocalStorage.ts

@@ -13,7 +13,11 @@ type LocalStorageItem = ReturnType<typeof localStorage.getItem>;
 // Bus to notify other useLocalStorage consumers when an item is changed
 export const localStorageBus = new EventEmitter();
 
-// Like useState, but reads from and persists the value to localStorage
+/**
+ * Like useState, but reads from and persists the value to localStorage
+ * this hook will not update when we write to localStorage.setItem(key, value) directly.
+ * For the hook to react either use the returned setter or `saveKeyForRoom`.
+ */
 export const useLocalStorage = (
   key: string,
 ): [LocalStorageItem, (value: string) => void] => {
@@ -42,14 +46,6 @@ export const useLocalStorage = (
 };
 
 export const setLocalStorageItem = (key: string, value: string): void => {
-  // Avoid unnecessary updates. Not avoiding them so can cause unexpected state updates across hooks.
-  // For instance:
-  // - In call view uses useRoomEncryptionSystem
-  // - This will set the key again.
-  // - All other instances of useRoomEncryptionSystem will now do a useMemo update of the e2eeSystem
-  //   - because the dependency `storedPassword = useInternalRoomSharedKey(roomId);` would change.
-  if (localStorage.getItem(key) === value) return;
-
   localStorage.setItem(key, value);
   localStorageBus.emit(key, value);
 };