Allow setting a specific impersonation device for appservices (#43)

<!-- Thanks for submitting a PR! Please ensure the following
requirements are met in order for us to review your PR -->

## Checklist

* [ ] Tests written for all new code
* [ ] Linter has been satisfied
* [ ] Sign-off given on the changes (see CONTRIBUTING.md)

Author: Half-Shot

src/appservice/Intent.ts

@@ -97,10 +97,12 @@ export class Intent {
 
     /**
      * Sets up crypto on the client if it hasn't already been set up.
+     * @param providedDeviceId Optional device ID. If given, this will used instead of trying to
+     * masquerade as the first non-key enabled device.
      * @returns {Promise<void>} Resolves when complete.
      */
     @timedIntentFunctionCall()
-    public async enableEncryption(): Promise<void> {
+    public async enableEncryption(providedDeviceId?: string): Promise<void> {
         if (!this.cryptoSetupPromise) {
             // eslint-disable-next-line no-async-promise-executor
             this.cryptoSetupPromise = new Promise(async (resolve, reject) => {
@@ -116,24 +118,38 @@ export class Intent {
                         throw new Error("Failed to create crypto store");
                     }
 
-                    // Try to impersonate a device ID
-                    const ownDevices = await this.client.getOwnDevices();
                     let deviceId = await cryptoStore.getDeviceId();
-                    if (!deviceId || !ownDevices.some(d => d.device_id === deviceId)) {
-                        const deviceKeys = await this.client.getUserDevices([this.userId]);
-                        const userDeviceKeys = deviceKeys.device_keys[this.userId];
-                        if (userDeviceKeys) {
-                            // We really should be validating signatures here, but we're actively looking
-                            // for devices without keys to impersonate, so it should be fine. In theory,
-                            // those devices won't even be present but we're cautious.
-                            const devicesWithKeys = Array.from(Object.entries(userDeviceKeys))
-                                .filter(d => d[0] === d[1].device_id && !!d[1].keys?.[`${DeviceKeyAlgorithm.Curve25519}:${d[1].device_id}`])
-                                .map(t => t[0]); // grab device ID from tuple
-                            deviceId = ownDevices.find(d => !devicesWithKeys.includes(d.device_id))?.device_id;
+                    if (!providedDeviceId) {
+                        // Try to impersonate a device ID
+                        const ownDevices = await this.client.getOwnDevices();
+                        let deviceId = await cryptoStore.getDeviceId();
+                        if (!deviceId || !ownDevices.some(d => d.device_id === deviceId)) {
+                            const deviceKeys = await this.client.getUserDevices([this.userId]);
+                            const userDeviceKeys = deviceKeys.device_keys[this.userId];
+                            if (userDeviceKeys) {
+                                // We really should be validating signatures here, but we're actively looking
+                                // for devices without keys to impersonate, so it should be fine. In theory,
+                                // those devices won't even be present but we're cautious.
+                                const devicesWithKeys = Array.from(Object.entries(userDeviceKeys))
+                                    .filter(d => d[0] === d[1].device_id && !!d[1].keys?.[`${DeviceKeyAlgorithm.Curve25519}:${d[1].device_id}`])
+                                    .map(t => t[0]); // grab device ID from tuple
+                                deviceId = ownDevices.find(d => !devicesWithKeys.includes(d.device_id))?.device_id;
+                            }
                         }
+                    } else {
+                        if (deviceId && deviceId !== providedDeviceId) {
+                            LogService.warn(`Storage already configured with an existing device ${deviceId}. Old storage will be cleared.`);
+                        }
+                        deviceId = providedDeviceId;
                     }
                     let prepared = false;
+
                     if (deviceId) {
+                        const cryptoStore = this.cryptoStorage?.storageForUser(this.userId);
+                        const existingDeviceId = await cryptoStore.getDeviceId();
+                        if (existingDeviceId && existingDeviceId !== deviceId) {
+                            LogService.warn("Intent", `Device ID has changed for user ${this.userId} from ${existingDeviceId} to ${deviceId}`);
+                        }
                         this.makeClient(true);
                         this.client.impersonateUserId(this.userId, deviceId);
 

src/e2ee/CryptoClient.ts

@@ -72,23 +72,25 @@ export class CryptoClient {
         if (this.ready) return; // stop re-preparing here
 
         const storedDeviceId = await this.client.cryptoStore.getDeviceId();
-        if (storedDeviceId) {
-            this.deviceId = storedDeviceId;
-        } else {
-            const deviceId = (await this.client.getWhoAmI())['device_id'];
-            if (!deviceId) {
-                throw new Error("Encryption not possible: server not revealing device ID");
-            }
-            this.deviceId = deviceId;
-            await this.client.cryptoStore.setDeviceId(this.deviceId);
+        const { user_id: userId, device_id: deviceId } = (await this.client.getWhoAmI());
+
+        if (!deviceId) {
+            throw new Error("Encryption not possible: server not revealing device ID");
+        }
+
+        const storagePath = await this.storage.getMachineStoragePath(deviceId);
+
+        if (storedDeviceId !== deviceId) {
+            this.client.cryptoStore.setDeviceId(deviceId);
         }
+        this.deviceId = deviceId;
 
-        LogService.debug("CryptoClient", "Starting with device ID:", this.deviceId);
+        LogService.debug("CryptoClient", `Starting ${userId} with device ID:`, this.deviceId);
 
         const machine = await OlmMachine.initialize(
-            new UserId(await this.client.getUserId()),
+            new UserId(userId),
             new DeviceId(this.deviceId),
-            this.storage.storagePath, "",
+            storagePath, "",
             this.storage.storageType,
         );
         this.engine = new RustEngine(machine, this.client);

src/storage/IAppserviceStorageProvider.ts

@@ -47,8 +47,8 @@ export interface IAppserviceStorageProvider {
 export interface IAppserviceCryptoStorageProvider {
     /**
      * Gets a storage provider to use for the given user ID.
-     * @param {string} userId The user ID.
-     * @returns {ICryptoStorageProvider} The storage provider.
+     * @param userId The user ID.
+     * @returns The storage provider.
      */
     storageForUser(userId: string): ICryptoStorageProvider;
 }

src/storage/RustSdkCryptoStorageProvider.ts

@@ -2,16 +2,23 @@ import * as lowdb from "lowdb";
 import * as FileSync from "lowdb/adapters/FileSync";
 import * as mkdirp from "mkdirp";
 import * as path from "path";
+import { stat, rename, mkdir } from "fs/promises";
+import { PathLike } from "fs";
 import * as sha512 from "hash.js/lib/hash/sha/512";
 import * as sha256 from "hash.js/lib/hash/sha/256";
 import { StoreType as RustSdkCryptoStoreType } from "@matrix-org/matrix-sdk-crypto-nodejs";
 
 import { ICryptoStorageProvider } from "./ICryptoStorageProvider";
 import { IAppserviceCryptoStorageProvider } from "./IAppserviceStorageProvider";
 import { ICryptoRoomInformation } from "../e2ee/ICryptoRoomInformation";
+import { LogService } from "../logging/LogService";
 
 export { RustSdkCryptoStoreType };
 
+async function doesFileExist(path: PathLike) {
+    return stat(path).then(() => true).catch(() => false);
+}
+
 /**
  * A crypto storage provider for the file-based rust-sdk store.
  * @category Storage providers
@@ -40,6 +47,37 @@ export class RustSdkCryptoStorageProvider implements ICryptoStorageProvider {
         });
     }
 
+    public async getMachineStoragePath(deviceId: string): Promise<string> {
+        const newPath = path.join(this.storagePath, sha256().update(deviceId).digest('hex'));
+        if (await doesFileExist(newPath)) {
+            // Already exists, short circuit.
+            return newPath;
+        } // else: If the path does NOT exist we might need to perform a migration.
+
+        const legacyFilePath = path.join(this.storagePath, 'matrix-sdk-crypto.sqlite3');
+        // XXX: Slightly gross cross-dependency file name expectations.
+        if (await doesFileExist(legacyFilePath) === false) {
+            // No machine files at all, we can skip.
+            return newPath;
+        }
+        const legacyDeviceId = await this.getDeviceId();
+        // We need to move the file.
+        const previousDevicePath = path.join(this.storagePath, sha256().update(legacyDeviceId).digest('hex'));
+        LogService.warn("RustSdkCryptoStorageProvider", `Migrating path for SDK database for legacy device ${legacyDeviceId}`);
+        await mkdir(previousDevicePath);
+        await rename(legacyFilePath, path.join(previousDevicePath, 'matrix-sdk-crypto.sqlite3')).catch((ex) =>
+            LogService.warn("RustSdkCryptoStorageProvider", `Could not migrate matrix-sdk-crypto.sqlite3`, ex),
+        );
+        await rename(legacyFilePath, path.join(previousDevicePath, 'matrix-sdk-crypto.sqlite3-shm')).catch((ex) =>
+            LogService.warn("RustSdkCryptoStorageProvider", `Could not migrate matrix-sdk-crypto.sqlite3-shm`, ex),
+        );
+        await rename(legacyFilePath, path.join(previousDevicePath, 'matrix-sdk-crypto.sqlite3-wal')).catch((ex) =>
+            LogService.warn("RustSdkCryptoStorageProvider", `Could not migrate matrix-sdk-crypto.sqlite3-wal`, ex),
+        );
+
+        return newPath;
+    }
+
     public async getDeviceId(): Promise<string> {
         return this.db.get('deviceId').value();
     }
@@ -75,7 +113,7 @@ export class RustSdkAppserviceCryptoStorageProvider extends RustSdkCryptoStorage
 
     public storageForUser(userId: string): ICryptoStorageProvider {
         // sha256 because sha512 is a bit big for some operating systems
-        const key = sha256().update(userId).digest('hex');
-        return new RustSdkCryptoStorageProvider(path.join(this.baseStoragePath, key), this.storageType);
+        const storagePath = path.join(this.baseStoragePath, sha256().update(userId).digest('hex'));
+        return new RustSdkCryptoStorageProvider(storagePath, this.storageType);
     }
 }

test/MatrixClientTest.ts

@@ -1043,7 +1043,7 @@ describe('MatrixClient', () => {
         });
 
         it('should request the user ID if it is not known', async () => {
-            const { client, http } = createTestClient();
+            const { client, http } = createTestClient(undefined, undefined, undefined, { handleWhoAmI: false });
 
             const userId = "@example:example.org";
             const response = {
@@ -1061,7 +1061,7 @@ describe('MatrixClient', () => {
 
     describe('getWhoAmI', () => {
         it('should call the right endpoint', async () => {
-            const { client, http } = createTestClient();
+            const { client, http } = createTestClient(undefined, undefined, undefined, { handleWhoAmI: false });
 
             const response = {
                 user_id: "@user:example.org",

test/SynapseAdminApisTest.ts

@@ -24,7 +24,7 @@ export function createTestSynapseAdminClient(
     hsUrl: string;
     accessToken: string;
 } {
-    const result = createTestClient(storage);
+    const result = createTestClient(storage, undefined, undefined, { handleWhoAmI: false });
     const mxClient = result.client;
     const client = new SynapseAdminApis(mxClient);
 

test/TestUtils.ts

@@ -31,6 +31,7 @@ export function createTestClient(
     storage: IStorageProvider = null,
     userId: string = null,
     cryptoStoreType?: StoreType,
+    opts = { handleWhoAmI: true },
 ): {
     client: MatrixClient;
     http: HttpBackend;
@@ -44,6 +45,11 @@ export function createTestClient(
     (<any>client).userId = userId; // private member access
     setRequestFn(http.requestFn);
 
+    if (opts.handleWhoAmI) {
+        // Ensure we always respond to a whoami
+        client.getWhoAmI = () => Promise.resolve({ user_id: userId, device_id: TEST_DEVICE_ID });
+    }
+
     return { http, hsUrl, accessToken, client };
 }
 

test/encryption/CryptoClientTest.ts

@@ -8,7 +8,6 @@ describe('CryptoClient', () => {
     it('should not have a device ID or be ready until prepared', () => testCryptoStores(async (cryptoStoreType) => {
         const userId = "@alice:example.org";
         const { client, http } = createTestClient(null, userId, cryptoStoreType);
-
         client.getWhoAmI = () => Promise.resolve({ user_id: userId, device_id: TEST_DEVICE_ID });
 
         expect(client.crypto).toBeDefined();
@@ -46,17 +45,20 @@ describe('CryptoClient', () => {
             const { client, http } = createTestClient(null, userId, cryptoStoreType);
 
             await client.cryptoStore.setDeviceId(TEST_DEVICE_ID);
+            const CORRECT_DEVICE = "new_device";
 
-            const whoamiSpy = simple.stub().callFn(() => Promise.resolve({ user_id: userId, device_id: "wrong" }));
+            const whoamiSpy = simple.stub().callFn(() => Promise.resolve({ user_id: userId, device_id: CORRECT_DEVICE }));
             client.getWhoAmI = whoamiSpy;
 
             bindNullEngine(http);
             await Promise.all([
                 client.crypto.prepare(),
                 http.flushAllExpected(),
             ]);
-            expect(whoamiSpy.callCount).toEqual(0);
-            expect(client.crypto.clientDeviceId).toEqual(TEST_DEVICE_ID);
+            // This should be called to check
+            expect(whoamiSpy.callCount).toEqual(1);
+            expect(client.crypto.clientDeviceId).toEqual(CORRECT_DEVICE);
+            expect(await client.cryptoStore.getDeviceId()).toEqual(CORRECT_DEVICE);
         }));
     });