Support staged rollout of migration to Rust Crypto (#12184)

* Rust migration staged rollout

* Phased rollout unit tests

src/MatrixClientPeg.ts

@@ -18,15 +18,15 @@ limitations under the License.
 */
 
 import {
-    ICreateClientOpts,
-    PendingEventOrdering,
-    RoomNameState,
-    RoomNameType,
     EventTimeline,
     EventTimelineSet,
+    ICreateClientOpts,
     IStartClientOpts,
     MatrixClient,
     MemoryStore,
+    PendingEventOrdering,
+    RoomNameState,
+    RoomNameType,
     TokenRefreshFunction,
 } from "matrix-js-sdk/src/matrix";
 import * as utils from "matrix-js-sdk/src/utils";
@@ -53,6 +53,7 @@ import PlatformPeg from "./PlatformPeg";
 import { formatList } from "./utils/FormattingUtils";
 import SdkConfig from "./SdkConfig";
 import { Features } from "./settings/Settings";
+import { PhasedRolloutFeature } from "./utils/PhasedRolloutFeature";
 
 export interface IMatrixClientCreds {
     homeserverUrl: string;
@@ -302,13 +303,34 @@ class MatrixClientPegClass implements IMatrixClientPeg {
             throw new Error("createClient must be called first");
         }
 
-        const useRustCrypto = SettingsStore.getValue(Features.RustCrypto);
+        let useRustCrypto = SettingsStore.getValue(Features.RustCrypto);
+
+        // We want the value that is set in the config.json for that web instance
+        const defaultUseRustCrypto = SettingsStore.getValueAt(SettingLevel.CONFIG, Features.RustCrypto);
+        const migrationPercent = SettingsStore.getValueAt(SettingLevel.CONFIG, "RustCrypto.staged_rollout_percent");
+
+        // If the default config is to use rust crypto, and the user is on legacy crypto,
+        // we want to check if we should migrate the current user.
+        if (!useRustCrypto && defaultUseRustCrypto && Number.isInteger(migrationPercent)) {
+            // The user is not on rust crypto, but the default stack is now rust; Let's check if we should migrate
+            // the current user to rust crypto.
+            try {
+                const stagedRollout = new PhasedRolloutFeature("RustCrypto.staged_rollout_percent", migrationPercent);
+                // Device id should not be null at that point, or init crypto will fail anyhow
+                const deviceId = this.matrixClient.getDeviceId()!;
+                // we use deviceId rather than userId because we don't particularly want all devices
+                // of a user to be migrated at the same time.
+                useRustCrypto = stagedRollout.isFeatureEnabled(deviceId);
+            } catch (e) {
+                logger.warn("Failed to create staged rollout feature for rust crypto migration", e);
+            }
+        }
 
         // we want to make sure that the same crypto implementation is used throughout the lifetime of a device,
         // so persist the setting at the device layer
         // (At some point, we'll allow the user to *enable* the setting via labs, which will migrate their existing
         // device to the rust-sdk implementation, but that won't change anything here).
-        await SettingsStore.setValue("feature_rust_crypto", null, SettingLevel.DEVICE, useRustCrypto);
+        await SettingsStore.setValue(Features.RustCrypto, null, SettingLevel.DEVICE, useRustCrypto);
 
         // Now we can initialise the right crypto impl.
         if (useRustCrypto) {

src/settings/Settings.tsx

@@ -96,6 +96,7 @@ export enum Features {
     VoiceBroadcastForceSmallChunks = "feature_voice_broadcast_force_small_chunks",
     NotificationSettings2 = "feature_notification_settings2",
     OidcNativeFlow = "feature_oidc_native_flow",
+    // If true, every new login will use the new rust crypto implementation
     RustCrypto = "feature_rust_crypto",
 }
 
@@ -503,6 +504,13 @@ export const SETTINGS: { [setting: string]: ISetting } = {
         default: false,
         controller: new RustCryptoSdkController(),
     },
+    // Must be set under `setting_defaults` in config.json.
+    // If set to 100 in conjunction with `feature_rust_crypto`, all existing users will migrate to the new crypto.
+    // Default is 0, meaning no existing users on legacy crypto will migrate.
+    "RustCrypto.staged_rollout_percent": {
+        supportedLevels: [SettingLevel.CONFIG],
+        default: 0,
+    },
     "baseFontSize": {
         displayName: _td("settings|appearance|font_size"),
         supportedLevels: LEVELS_ACCOUNT_SETTINGS,

src/utils/PhasedRolloutFeature.ts

@@ -0,0 +1,63 @@
+/*
+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 { xxHash32 } from "js-xxhash";
+
+/**
+ * The PhasedRolloutFeature class is used to manage the phased rollout of a new feature.
+ *
+ * It uses a hash of the user's identifier and the feature name to determine if a feature is enabled for a specific user.
+ * The rollout percentage determines the probability that a user will be enabled for the feature.
+ * The feature will be enabled for all users if the rollout percentage is 100, and for no users if the percentage is 0.
+ * If a user is enabled for a feature at x% rollout, it will also be for any greater than x percent.
+ *
+ * The process ensures a uniform distribution of enabled features across users.
+ *
+ * @property featureName - The name of the feature to be rolled out.
+ * @property rolloutPercentage - The int percentage (0..100) of users for whom the feature should be enabled.
+ */
+export class PhasedRolloutFeature {
+    public readonly featureName: string;
+    private readonly rolloutPercentage: number;
+    private readonly seed: number;
+
+    public constructor(featureName: string, rolloutPercentage: number) {
+        this.featureName = featureName;
+        if (!Number.isInteger(rolloutPercentage) || rolloutPercentage < 0 || rolloutPercentage > 100) {
+            throw new Error("Rollout percentage must be an integer between 0 and 100");
+        }
+        this.rolloutPercentage = rolloutPercentage;
+        // We add the feature name for the seed to ensure that the hash is different for each feature
+        this.seed = Array.from(featureName).reduce((sum, char) => sum + char.charCodeAt(0), 0);
+    }
+
+    /**
+     * Returns true if the feature should be enabled for the given user.
+     * @param userIdentifier - Some unique identifier for the user, e.g. their user ID or device ID.
+     */
+    public isFeatureEnabled(userIdentifier: string): boolean {
+        /*
+         * We use a hash function to convert the unique user ID string into an integer.
+         * This integer can then be used as a basis for deciding whether the user should have access to the new feature.
+         * We need some hash with good uniform distribution properties, security is not a concern here.
+         * We use xxHash32, which is fast and has good distribution properties.
+         */
+        const hash = xxHash32(userIdentifier, this.seed);
+        // We use the hash modulo 100 to get a number between 0 and 99.
+        // Modulo is simple and effective and the distribution should be uniform enough for our purposes.
+        return hash % 100 < this.rolloutPercentage;
+    }
+}