Get the access token directly

src/serviceworker/index.ts

@@ -1,24 +1,40 @@
+/*
+Copyright 2024 New Vector Ltd
+
+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 { idbLoad } from "matrix-react-sdk/src/utils/StorageAccess";
+import { ACCESS_TOKEN_IV, tryDecryptToken } from "matrix-react-sdk/src/utils/tokens/tokens";
+import { encodeUnpaddedBase64 } from "matrix-js-sdk/src/base64";
+
 const serverSupportMap: {
     [serverUrl: string]: {
         supportsMSC3916: boolean;
         cacheExpires: number;
     };
 } = {};
 
-const credentialStore: {
-    [serverUrl: string]: string;
-} = {};
-
-// We skipWaiting() to update the service worker more frequently, particularly in development environments.
-// @ts-expect-error - service worker types are not available. See 'fetch' event handler.
-skipWaiting();
+self.addEventListener("install", (event) => {
+    // We skipWaiting() to update the service worker more frequently, particularly in development environments.
+    // @ts-expect-error - service worker types are not available. See 'fetch' event handler.
+    event.waitUntil(skipWaiting());
+});
 
-self.addEventListener("message", (event) => {
-    if (event.data?.type !== "credentials") return; // ignore
-    credentialStore[event.data.homeserverUrl] = event.data.accessToken;
-    console.log(
-        `[Service Worker] Updated access token for ${event.data.homeserverUrl} (accessToken? ${Boolean(event.data.accessToken)})`,
-    );
+self.addEventListener("activate", (event) => {
+    // We force all clients to be under our control, immediately. This could be old tabs.
+    // @ts-expect-error - service worker types are not available. See 'fetch' event handler.
+    event.waitUntil(clients.claim());
 });
 
 // @ts-expect-error - getting types to work for this is difficult, so we anticipate that "addEventListener" doesn't
@@ -42,34 +58,40 @@ self.addEventListener("fetch", (event: FetchEvent) => {
         // later on we need to proxy the request through if it turns out the server doesn't support authentication.
         event.respondWith(
             (async (): Promise<Response> => {
-                // Figure out which homeserver we're communicating with
-                const csApi = url.substring(0, url.indexOf("/_matrix/media/v3"));
-
-                // Locate our access token, and populate the fetchConfig with the authentication header.
-                const accessToken = credentialStore[csApi];
                 let fetchConfig: { headers?: { [key: string]: string } } = {};
-                if (accessToken) {
-                    fetchConfig = {
-                        headers: {
-                            Authorization: `Bearer ${accessToken}`,
-                        },
-                    };
-                }
+                try {
+                    // Figure out which homeserver we're communicating with
+                    const csApi = url.substring(0, url.indexOf("/_matrix/media/v3"));
 
-                // Update or populate the server support map using a (usually) authenticated `/versions` call.
-                if (!serverSupportMap[csApi] || serverSupportMap[csApi].cacheExpires <= new Date().getTime()) {
-                    const versions = await (await fetch(`${csApi}/_matrix/client/versions`, fetchConfig)).json();
-                    serverSupportMap[csApi] = {
-                        supportsMSC3916: Boolean(versions?.unstable_features?.["org.matrix.msc3916"]),
-                        cacheExpires: new Date().getTime() + 2 * 60 * 60 * 1000, // 2 hours from now
-                    };
-                }
+                    // Locate our access token, and populate the fetchConfig with the authentication header.
+                    // @ts-expect-error - service worker types are not available. See 'fetch' event handler.
+                    const client = await self.clients.get(event.clientId);
+                    const accessToken = await getAccessToken(client);
+                    if (accessToken) {
+                        fetchConfig = {
+                            headers: {
+                                Authorization: `Bearer ${accessToken}`,
+                            },
+                        };
+                    }
+
+                    // Update or populate the server support map using a (usually) authenticated `/versions` call.
+                    if (!serverSupportMap[csApi] || serverSupportMap[csApi].cacheExpires <= new Date().getTime()) {
+                        const versions = await (await fetch(`${csApi}/_matrix/client/versions`, fetchConfig)).json();
+                        serverSupportMap[csApi] = {
+                            supportsMSC3916: Boolean(versions?.unstable_features?.["org.matrix.msc3916"]),
+                            cacheExpires: new Date().getTime() + 2 * 60 * 60 * 1000, // 2 hours from now
+                        };
+                    }
 
-                // If we have server support (and a means of authentication), rewrite the URL to use MSC3916 endpoints.
-                if (serverSupportMap[csApi].supportsMSC3916 && accessToken) {
-                    // Currently unstable only.
-                    url = url.replace(/\/media\/v3\/(.*)\//, "/client/unstable/org.matrix.msc3916/media/$1/");
-                } // else by default we make no changes
+                    // If we have server support (and a means of authentication), rewrite the URL to use MSC3916 endpoints.
+                    if (serverSupportMap[csApi].supportsMSC3916 && accessToken) {
+                        // Currently unstable only.
+                        url = url.replace(/\/media\/v3\/(.*)\//, "/client/unstable/org.matrix.msc3916/media/$1/");
+                    } // else by default we make no changes
+                } catch (err) {
+                    console.error("SW: Error in request rewrite.", err);
+                }
 
                 // Add authentication and send the request. We add authentication even if MSC3916 endpoints aren't
                 // being used to ensure patches like this work:
@@ -79,3 +101,69 @@ self.addEventListener("fetch", (event: FetchEvent) => {
         );
     }
 });
+
+// Ideally we'd use the `Client` interface for `client`, but since it's not available (see 'fetch' listener), we use
+// unknown for now and force-cast it to something close enough later.
+async function getAccessToken(client: unknown): Promise<string | undefined> {
+    // Access tokens are encrypted at rest, so while we can grab the "access token", we'll need to do work to get the
+    // real thing.
+    const encryptedAccessToken = await idbLoad("account", "mx_access_token");
+
+    // We need to extract a user ID and device ID from localstorage, which means calling WebPlatform for the
+    // read operation. Service workers can't access localstorage.
+    const { userId, deviceId } = await new Promise<{ userId: string; deviceId: string }>((resolve, reject) => {
+        // Avoid stalling the tab in case something goes wrong.
+        const timeoutId = setTimeout(() => reject(new Error("timeout in postMessage")), 1000);
+
+        // We don't need particularly good randomness here - we just use this to generate a request ID, so we know
+        // which postMessage reply is for our active request.
+        const responseKey = Math.random().toString(36);
+
+        // Add the listener first, just in case the tab is *really* fast.
+        const listener = (event: MessageEvent): void => {
+            if (event.data?.responseKey !== responseKey) return; // not for us
+            clearTimeout(timeoutId); // do this as soon as possible, avoiding a race between resolve and reject.
+            resolve(event.data); // "unblock" the remainder of the thread, if that were such a thing in JavaScript.
+            self.removeEventListener("message", listener); // cleanup, since we're not going to do anything else.
+        };
+        self.addEventListener("message", listener);
+
+        // Ask the tab for the information we need. This is handled by WebPlatform.
+        (client as Window).postMessage({ responseKey, type: "userinfo" });
+    });
+
+    // ... and this is why we need the user ID and device ID: they're index keys for the pickle key table.
+    const pickleKeyData = await idbLoad("pickleKey", [userId, deviceId]);
+    if (pickleKeyData && (!pickleKeyData.encrypted || !pickleKeyData.iv || !pickleKeyData.cryptoKey)) {
+        console.error("SW: Invalid pickle key loaded - ignoring");
+        return undefined;
+    }
+
+    let pickleKey: string | undefined;
+
+    // Extract a useful pickle key out of our queries.
+    if (pickleKeyData) {
+        // We also need to generate the additional data needed for the key
+        const additionalData = new Uint8Array(userId.length + deviceId.length + 1);
+        for (let i = 0; i < userId.length; i++) {
+            additionalData[i] = userId.charCodeAt(i);
+        }
+        additionalData[userId.length] = 124; // "|"
+        for (let i = 0; i < deviceId.length; i++) {
+            additionalData[userId.length + 1 + i] = deviceId.charCodeAt(i);
+        }
+
+        // Convert pickle key to a base64 key we can use
+        const pickleKeyBuf = await crypto.subtle.decrypt(
+            { name: "AES-GCM", iv: pickleKeyData.iv, additionalData },
+            pickleKeyData.cryptoKey,
+            pickleKeyData.encrypted,
+        );
+        if (pickleKeyBuf) {
+            pickleKey = encodeUnpaddedBase64(pickleKeyBuf);
+        }
+    }
+
+    // Finally, try decrypting the thing and return that. This may fail, but that's okay.
+    return tryDecryptToken(pickleKey, encryptedAccessToken, ACCESS_TOKEN_IV);
+}

src/vector/platform/WebPlatform.ts

@@ -1,7 +1,7 @@
 /*
 Copyright 2016 Aviral Dasgupta
 Copyright 2016 OpenMarket Ltd
-Copyright 2017-2020 New Vector Ltd
+Copyright 2017-2020, 2024 New Vector Ltd
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -52,7 +52,27 @@ export default class WebPlatform extends VectorBasePlatform {
             // Jest causes `register()` to return undefined, so swallow that case.
             if (swPromise) {
                 swPromise
-                    .then((r) => r.update())
+                    .then(async (r) => {
+                        await r.update();
+                        return r;
+                    })
+                    .then((r) => {
+                        navigator.serviceWorker.addEventListener("message", (e) => {
+                            try {
+                                if (e.data?.["type"] === "userinfo" && e.data?.["responseKey"]) {
+                                    const userId = localStorage.getItem("mx_user_id");
+                                    const deviceId = localStorage.getItem("mx_device_id");
+                                    r.active!.postMessage({
+                                        responseKey: e.data["responseKey"],
+                                        userId,
+                                        deviceId,
+                                    });
+                                }
+                            } catch (e) {
+                                console.error("Error responding to service worker: ", e);
+                            }
+                        });
+                    })
                     .catch((e) => console.error("Error registering/updating service worker:", e));
             }
         }