utils.ts

import { browser } from "@wdio/globals";
import type { Bench } from "tinybench";
import type { local } from "webdriver";
import {
    Ciphersuite,
    CoreCrypto,
    CoreCryptoLogLevel,
    type CommitBundle,
    type HistorySecret,
    type MlsTransport,
    type MlsTransportData,
} from "../src/browser/CoreCrypto";
type ccModuleType = typeof import("../src/browser/CoreCrypto");
declare global {
    interface Window {
        ccModule: ccModuleType;
        deliveryService: DeliveryService;
        _latestCommitBundle: CommitBundle;
        tinybench: typeof import("tinybench");
        benchRunning: boolean;
        bench: Bench;
        helpers: BenchmarkHelpers;
    }
}

const logLevel = Number(process.env["CC_TEST_LOG_LEVEL"] || "0");

export async function setup() {
    if ((await browser.getUrl()) === "about:blank") {
        await browser.url("/html/index.html");
    }

    browser.on("log.entryAdded", logEvents);
    await browser.execute(async (loglevel) => {
        if (window.tinybench === undefined) {
            window.tinybench =
                // @ts-expect-error TS2307: Cannot find module ./corecrypto.js or its corresponding type declarations.
                await import("/node_modules/tinybench/dist/index.js");
        }
        if (window.ccModule === undefined) {
            // @ts-expect-error TS2307: Cannot find module ./corecrypto.js or its corresponding type declarations.
            window.ccModule = await import("/corecrypto.js");
            await window.ccModule.initWasmModule(
                "/autogenerated/wasm-bindgen/"
            );

            if (loglevel >= 2) {
                window.ccModule.setLogger({
                    log: (
                        _level: CoreCryptoLogLevel,
                        message: string,
                        context: string | undefined
                    ) => {
                        console.log(message, context);
                    },
                });
                window.ccModule.setMaxLogLevel(
                    window.ccModule.CoreCryptoLogLevel.Debug
                );
            }
        }

        if (window.deliveryService === undefined) {
            window.deliveryService = {
                async sendCommitBundle(commitBundle: CommitBundle) {
                    window._latestCommitBundle = commitBundle;
                    return window.ccModule.MlsTransportResponse.Success.new();
                },
                async prepareForTransport(
                    secret: HistorySecret
                ): Promise<MlsTransportData> {
                    return Promise.resolve(secret.clientId.copyBytes());
                },
                async getLatestCommitBundle() {
                    return window._latestCommitBundle;
                },
            };
        }

        class BenchmarkHelpers {
            static async setupCc(
                cipherSuite: Ciphersuite
            ): Promise<CoreCrypto> {
                const encoder = new TextEncoder();

                const clientIdStr = window.crypto.randomUUID();
                const id = new window.ccModule.ClientId(
                    encoder.encode(clientIdStr).buffer
                );

                const key = new Uint8Array(32);
                crypto.getRandomValues(key);

                const db = await window.ccModule.Database.open(
                    clientIdStr,
                    new window.ccModule.DatabaseKey(key.buffer)
                );

                const cc = window.ccModule.CoreCrypto.new(db);

                await cc.newTransaction(async (ctx) => {
                    await ctx.mlsInit(id, window.deliveryService);
                    await ctx.addCredential(
                        window.ccModule.Credential.basic(cipherSuite, id)
                    );
                });
                return cc;
            }
        }
        window.helpers = BenchmarkHelpers;
    }, logLevel);
}

interface DeliveryService extends MlsTransport {
    getLatestCommitBundle: () => Promise<CommitBundle>;
}

interface BenchmarkHelpers {
    setupCc: (cipherSuite: Ciphersuite) => Promise<CoreCrypto>;
}

type ParameterSet = {
    count: number;
    size: number;
    cipherSuite: number;
};

export async function benchmarkParameters(): Promise<ParameterSet[]> {
    const messageCounts = [1, 10, 100];
    const messageSizes = [16, 1024, 65536];
    const cipherSuites = [
        Ciphersuite.Mls128Dhkemx25519Aes128gcmSha256Ed25519,
        Ciphersuite.Mls128Dhkemx25519Chacha20poly1305Sha256Ed25519,
        Ciphersuite.Mls128Dhkemp256Aes128gcmSha256P256,
        Ciphersuite.Mls256Dhkemp384Aes256gcmSha384P384,
        Ciphersuite.Mls256Dhkemp521Aes256gcmSha512P521,
    ];

    function* benchmarkCombinations() {
        for (const count of messageCounts) {
            for (const size of messageSizes) {
                for (const cipherSuite of cipherSuites) {
                    yield { count, size, cipherSuite };
                }
            }
        }
    }

    return Array.from(benchmarkCombinations()) as ParameterSet[]; // return as plain array
}

function logEvents(entry: local.LogEntry) {
    if (logLevel >= 1) {
        console.log(`[${entry.level}] ${entry.text}`);
    }
}