utils.ts

import { browser } from "@wdio/globals";
import type { local } from "webdriver";

import {
    Ciphersuite,
    CoreCrypto,
    type MlsTransport,
    type CommitBundle,
    type GroupInfoBundle,
    type HistorySecret,
    type MlsTransportData,
    CoreCryptoLogLevel,
    type PkiEnvironmentHooks,
    HttpMethod,
    HttpHeader,
} from "../../src/browser/CoreCrypto";

type ccModuleType = typeof import("../../src/browser/CoreCrypto");

// Logging can be adjusted via the CC_TEST_LOG_LEVEL variable:
// 0 = no logs
// 1 = browser logs
// 2 = browser logs + CoreCrypto logs
const logLevel = Number(process.env["CC_TEST_LOG_LEVEL"] || "0");

/**
 * Log entry from the core crypto logger
 */
export interface LogEntry {
    level: number;
    message: string;
    context: string;
}

declare global {
    interface Window {
        ccModule: ccModuleType;
        cc: Map<string, CoreCrypto>;
        defaultCipherSuite: Ciphersuite;
        deliveryService: DeliveryService;
        pkiEnvironmentHooks: PkiEnvironmentHooks;
        _latestCommitBundle: CommitBundle;
        recordedLogs: LogEntry[];

        // Helper functions that are used inside the browser context
        /**
         * Gets a {@link CoreCrypto} instance initialized previously via
         * {@link ccInit}.
         *
         * @param clientName The name the {@link ccInit} was called with.
         *
         * @returns {CoreCrypto} The {@link CoreCrypto} instance.
         *
         * @throws Error if no instance with the name has been initialized.
         */
        ensureCcDefined: (clientName: string) => CoreCrypto;
    }
}

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

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

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

    // Forward browser log events to the console.
    browser.on("log.entryAdded", logEvents);

    await browser.execute(async (logLevel) => {
        if (window.ccModule === undefined) {
            // This is imported in the browser context, where it is fetched from the static file server,
            // but typescript tries to resolve this in the local directory.
            // @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
                );
            }

            window.defaultCipherSuite =
                window.ccModule.Ciphersuite.Mls128Dhkemx25519Aes128gcmSha256Ed25519;
        }

        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;
            },
        };

        window.pkiEnvironmentHooks = {
            async httpRequest(
                _method: HttpMethod,
                _url: string,
                _headers: Array<HttpHeader>,
                _body: ArrayBuffer
            ) {
                // return a HttpResponse
                return {
                    status: 200,
                    headers: [],
                    body: new Uint8Array().buffer,
                };
            },

            async authenticate(
                _idp: string,
                _keyAuth: string,
                _acmeAud: string
            ) {
                return "dummy-id-token";
            },

            async getBackendNonce() {
                return "dummy-backend-nonce";
            },

            async fetchBackendAccessToken(_dpop) {
                return "dummy-backend-token";
            },
        };

        window.ensureCcDefined = (clientName: string) => {
            const cc = window.cc.get(clientName);
            if (cc === undefined) {
                throw new Error(
                    `Client with name '${clientName}' is not initialized in the browser context.`
                );
            }
            return cc;
        };
    }, logLevel);
}

export async function teardown() {
    browser.off("log.entryAdded", logEvents);
}

/**
 * Initialize a {@link CoreCrypto} instance that can be obtained inside the
 * browser context via {@link Window.ensureCcDefined}.
 *
 * @param clientName The client name used to initialize.
 * @param withCredential When set (default), adds a basic credential to the CC instance
 *
 * @returns {Promise<void>}
 */
export async function ccInit(
    clientName: string,
    withCredential: boolean = true
): Promise<void> {
    return await browser.execute(
        async (clientName, withCredential) => {
            const cipherSuite = window.defaultCipherSuite;
            const encoder = new TextEncoder();
            const clientId = new window.ccModule.ClientId(
                encoder.encode(clientName).buffer
            );

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

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

            const instance = window.ccModule.CoreCrypto.new(database);
            await instance.newTransaction(async (ctx) => {
                await ctx.mlsInit(clientId, window.deliveryService);
                if (withCredential) {
                    await ctx.addCredential(
                        window.ccModule.Credential.basic(cipherSuite, clientId)
                    );
                }
            });

            if (window.cc === undefined) {
                window.cc = new Map();
            }
            window.cc.set(clientName, instance);
        },
        clientName,
        withCredential
    );
}

/**
 * Records logs by setting a logger and maximum log level in the browser's context.
 * The logs are stored in a global variable `window.recordedLogs` for further retrieval.
 *
 * @return {Promise<void>}
 */
export async function recordLogs(): Promise<void> {
    await browser.execute(async () => {
        const { setMaxLogLevel, CoreCryptoLogLevel, setLogger } =
            window.ccModule;
        window.recordedLogs = [];

        setLogger({
            log: (level: number, message: string, context: string) => {
                console.log(message, context);
                window.recordedLogs.push({
                    level: level,
                    message: message,
                    context: context,
                });
            },
        });
        setMaxLogLevel(CoreCryptoLogLevel.Debug);
    });
}

/**
 * Retrieves the logs recorded on the browser-side.
 *
 * @return {Promise<LogEntry[]>} A promise that resolves to an array of log entries
 */
export async function retrieveLogs(): Promise<LogEntry[]> {
    return await browser.execute(async () => {
        return window.recordedLogs.map((entry) => {
            return entry;
        });
    });
}

/**
 * Create a conversation on a {@link CoreCrypto} instance that has
 * been initialized before via {@link ccInit}.
 *
 * @param clientName The name the {@link CoreCrypto} instance has been
 * initialized with.
 * @param conversationId The id that the conversation will be created with.
 *
 * @returns {Promise<void>}
 *
 * @throws Error if the instance with {@link clientName} cannot be found.
 */
export async function createConversation(
    clientName: string,
    conversationId: string
): Promise<void> {
    return await browser.execute(
        async (clientName, conversationId) => {
            const cc = window.ensureCcDefined(clientName);
            await cc.newTransaction(async (ctx) => {
                const conversationIdBytes = new window.ccModule.ConversationId(
                    new TextEncoder().encode(conversationId).buffer
                );
                const [credentialRef] = await ctx.getCredentials();
                await ctx.createConversation(
                    conversationIdBytes,
                    credentialRef!
                );
            });
        },
        clientName,
        conversationId
    );
}

/**
 * Invite {@link client2} to a previously created conversation on the
 * instance of {@link client1} (via {@link createConversation}).
 *
 * @param client1 The name of the {@link CoreCrypto} instance on which the
 * conversation was created previously.
 * @param client2 The name of the {@link CoreCrypto} instance that will be
 * invited.
 * @param conversationId The id of the previously created conversation.
 *
 * @returns {Promise<GroupInfoBundle>} The resulting group info.
 *
 * @throws Error if {@link client1} or {@link client2} instances cannot be found.
 */
export async function invite(
    client1: string,
    client2: string,
    conversationId: string
): Promise<GroupInfoBundle> {
    return await browser.execute(
        async (client1, client2, conversationId) => {
            const cc1 = window.ensureCcDefined(client1);
            const cc2 = window.ensureCcDefined(client2);
            const conversationIdBytes = new window.ccModule.ConversationId(
                new TextEncoder().encode(conversationId).buffer
            );
            const kp = await cc2.newTransaction(async (ctx) => {
                const [credentialRef] = await ctx.getFilteredCredentials({
                    ciphersuite: window.defaultCipherSuite,
                    credentialType: window.ccModule.CredentialType.Basic,
                });
                console.log("generated key package");
                return await ctx.generateKeypackage(credentialRef!);
            });

            const clients = await cc1.getClientIds(conversationIdBytes);
            console.log("clients");
            console.log(clients);

            console.log("inviting bob");
            await cc1.newTransaction((ctx) =>
                ctx.addClientsToConversation(conversationIdBytes, [kp])
            );
            console.log("processing welcome");
            const commitBundle =
                await window.deliveryService.getLatestCommitBundle();
            await cc2.newTransaction((ctx) =>
                ctx.processWelcomeMessage(
                    new window.ccModule.Welcome(
                        commitBundle.welcome!.copyBytes()
                    )
                )
            );

            return commitBundle.groupInfo;
        },
        client1,
        client2,
        conversationId
    );
}

/**
 * Remove {@link client2} from a previously created conversation on the
 * instance of {@link client1} (via {@link createConversation}).
 *
 * @param client1 The name of the {@link CoreCrypto} instance on which the
 * conversation was created previously.
 * @param client2 The name of the {@link CoreCrypto} instance that will be
 * removed.
 * @param conversationId The id of the previously created conversation.
 *
 * @returns {Promise<GroupInfoBundle>} The resulting group info.
 *
 * @throws Error if {@link client1} or {@link client2} instances cannot be found.
 */
export async function remove(
    client1: string,
    client2: string,
    conversationId: string
): Promise<GroupInfoBundle> {
    return await browser.execute(
        async (client1, client2, conversationId) => {
            const cc1 = window.ensureCcDefined(client1);
            const cid = new window.ccModule.ConversationId(
                new TextEncoder().encode(conversationId).buffer
            );
            const clientId = new window.ccModule.ClientId(
                new TextEncoder().encode(client2).buffer
            );
            await cc1.newTransaction((ctx) =>
                ctx.removeClientsFromConversation(cid, [clientId])
            );
            const commitBundle =
                await window.deliveryService.getLatestCommitBundle();

            return commitBundle.groupInfo;
        },
        client1,
        client2,
        conversationId
    );
}

/**
 * Consume the last commit message on {@link client1}
 *
 * @param client1 The name of the {@link CoreCrypto} instance on which to consume the commit.
 * @param conversationId The id of the previously created conversation.
 *
 * @returns {Promise<void>}
 *
 * @throws Error if {@link client1} instances cannot be found.
 */
export async function consumeLastestCommit(
    client1: string,
    conversationId: string
): Promise<void> {
    return await browser.execute(
        async (client1, conversationId) => {
            const cc1 = window.ensureCcDefined(client1);
            const cid = new window.ccModule.ConversationId(
                new TextEncoder().encode(conversationId).buffer
            );
            const commitBundle =
                await window.deliveryService.getLatestCommitBundle();
            await cc1.newTransaction((ctx) =>
                ctx.decryptMessage(cid, commitBundle.commit)
            );
        },
        client1,
        conversationId
    );
}

/**
 * Inside a previously created conversation, {@link client1} encrypts
 * {@link message}, sends it to {@link client2}, who then decrypts it.
 * This procedure is then repeated vice versa.
 *
 * @param client1 The first of the conversation.
 * @param client2 The second member of the conversation.
 * @param conversationId The id of the conversation.
 * @param message The message encrypted, sent, and decrypted once in each
 * direction.
 *
 * @returns {Promise<(string | null)[]>} A two-element list, containing the decrypted {@link message} by
 * {@link client1} and {@link client2}, in that order.
 */
export async function roundTripMessage(
    client1: string,
    client2: string,
    conversationId: string,
    message: string
): Promise<(string | null)[]> {
    const [decrypted1, decrypted2] = await browser.execute(
        async (client1, client2, conversationId, message) => {
            const cc1 = window.ensureCcDefined(client1);
            const cc2 = window.ensureCcDefined(client2);

            const encoder = new TextEncoder();
            const cid = new window.ccModule.ConversationId(
                encoder.encode(conversationId).buffer
            );
            const messageBytes = encoder.encode(message);

            const encryptedByClient1 = await cc1.newTransaction(async (ctx) => {
                return await ctx.encryptMessage(cid, messageBytes.buffer);
            });
            const decryptedByClient2 = await cc2.newTransaction(async (ctx) => {
                return await ctx.decryptMessage(cid, encryptedByClient1);
            });

            const encryptedByClient2 = await cc2.newTransaction(async (ctx) => {
                return await ctx.encryptMessage(cid, messageBytes.buffer);
            });
            const decryptedByClient1 = await cc1.newTransaction(async (ctx) => {
                return await ctx.decryptMessage(cid, encryptedByClient2);
            });

            const decoder = new TextDecoder();
            const result1 =
                decryptedByClient1.message !== undefined
                    ? decoder.decode(decryptedByClient1.message)
                    : null;
            const result2 =
                decryptedByClient2.message !== undefined
                    ? decoder.decode(decryptedByClient2.message)
                    : null;

            return [result1, result2];
        },
        client1,
        client2,
        conversationId,
        message
    );
    return [decrypted1, decrypted2];
}

/**
 * Initialize a {@link CoreCrypto} instance without initializing MLS.
 * Instead, initialize proteus.
 * It can be obtained inside the browser context via
 * {@link Window.ensureCcDefined}.
 *
 * @param clientName the client name used to initialize.
 *
 * @returns {Promise<void>}
 */
export async function proteusInit(clientName: string): Promise<void> {
    return await browser.execute(async (clientName) => {
        const key = new Uint8Array(32);
        window.crypto.getRandomValues(key);

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

        const instance = window.ccModule.CoreCrypto.new(database);
        await instance.newTransaction((ctx) => ctx.proteusInit());

        if (window.cc === undefined) {
            window.cc = new Map();
        }
        window.cc.set(clientName, instance);
    }, clientName);
}

/**
 * Create a proteus session on the {@link CoreCrypto} instance of
 * {@link client1}, with the prekey of {@link client2}.
 *
 * @param client1 The name of the {@link CoreCrypto} instance which will
 * create the session.
 * @param client2 The name of the {@link CoreCrypto} instance whose pre key will
 * be used.
 * @param sessionId The id of session that will be created.
 *
 * @returns {Promise<void>}
 *
 * @throws Error if {@link client1} or {@link client2} instances cannot be found.
 */
export async function newProteusSessionFromPrekey(
    client1: string,
    client2: string,
    sessionId: string
): Promise<void> {
    return await browser.execute(
        async (client1, client2, sessionId) => {
            const cc1 = window.ensureCcDefined(client1);
            const cc2 = window.ensureCcDefined(client2);

            const cc2Prekey = await cc2.newTransaction(async (ctx) => {
                return await ctx.proteusNewPrekey(10);
            });

            await cc1.newTransaction(async (ctx) => {
                return await ctx.proteusSessionFromPrekey(sessionId, cc2Prekey);
            });
        },
        client1,
        client2,
        sessionId
    );
}

/**
 * Create a proteus session on the {@link CoreCrypto} instance of
 * {@link client2}, from a message encrypted by {@link client1} in a session
 * created previously via {@link newProteusSessionFromPrekey}.
 *
 * @param client1 The name of the {@link CoreCrypto} instance which used its
 * existing session to encrypt the message.
 * @param client2 The name of the {@link CoreCrypto} instance whose session will
 * be created.
 * @param sessionId The id of session that will be created.
 * For simplicity, this must match the id of the previously created session.
 * @param message The message to encrypt and create the message from.
 *
 * @returns {Promise<string | null>} the decrypted {@link message}.
 *
 * @throws Error if {@link client1} or {@link client2} instances cannot be found.
 */
export async function newProteusSessionFromMessage(
    client1: string,
    client2: string,
    sessionId: string,
    message: string
): Promise<string | null> {
    const decrypted = (await browser.execute(
        async (client1, client2, sessionId, message) => {
            const cc1 = window.ensureCcDefined(client1);
            const cc2 = window.ensureCcDefined(client2);

            const encoder = new TextEncoder();
            const messageBytes = encoder.encode(message);
            const encrypted = await cc1.newTransaction(async (ctx) => {
                return await ctx.proteusEncrypt(sessionId, messageBytes.buffer);
            });

            const decrypted = await cc2.newTransaction(async (ctx) => {
                return await ctx.proteusSessionFromMessage(
                    sessionId,
                    encrypted
                );
            });

            const decoder = new TextDecoder();
            return decrypted !== null ? decoder.decode(decrypted) : null;
        },
        client1,
        client2,
        sessionId,
        message
    )) as string | null;
    return decrypted;
}