e2ee.ts

export type Deps = { crypto: Crypto; TransformStream: typeof TransformStream };
export type Params = {
    counterLength: 64 | 128;
    namedCurve: "P-256" | "P-384" | "P-521";
    keyLength: 128 | 192 | 256;
};
export type Options = { deps?: Deps; params?: Params };
export type Marshalled = { params: Params; keyPair: CryptoKeyPair | null };
export type KeyGenOptions = {
    extractable?: boolean;
    additionalUsages?: KeyUsage[];
};
export type UnmarshalOptions = { marshalled: Marshalled; deps?: Deps };

export class E2EE {
    #deps: Deps = { crypto: globalThis.crypto, TransformStream: globalThis.TransformStream };

    #params: Params = {
        counterLength: 64,
        namedCurve: "P-256",
        keyLength: 256,
    };

    #keyPair: CryptoKeyPair | null = null;

    #sharedSecrets: Record<string | symbol, CryptoKey> = {};

    #unicast = Symbol("unicast");

    /**
     * @param options
     * @param options.deps Optionally inject your own platform dependencies.
     * @param options.deps.crypto An implementation of the WebCrypto API. Defaults to `globalThis.crypto`.
     * @param options.params Optionally override the default security parameters.
     * @param options.params.counterLength The length of the counter used in AES-CTR. Defaults to `64`.
     * @param options.params.namedCurve The named curve used in ECDH. Defaults to `P-256`.
     * @param options.params.keyLength The length of the key used in AES-CTR. Defaults to `256`.
     */
    constructor(options?: Options) {
        Object.assign(this.#deps, options?.deps);
        Object.assign(this.#params, options?.params);
    }

    /**
     * Generates an ECDH key pair for this party.
     * @param options
     * @param options.extractable Whether the key should be extractable. Defaults to `false`. SECURITY: Leave it to false, unless you know what you're doing.
     * @param options.additionalUsages Additional usages for the key. Defaults to `[]`.
     */
    async generateKeyPair({ extractable = false, additionalUsages = [] }: KeyGenOptions = {}) {
        if (this.#keyPair) throw new Error("Key pair already exists");

        this.#keyPair = await this.#deps.crypto.subtle.generateKey(
            { name: "ECDH", namedCurve: this.#params.namedCurve },
            extractable,
            ["deriveKey", ...additionalUsages] as ReadonlyArray<KeyUsage>
        );
    }

    /**
     * @returns The public key of this party, serialised as a JWK.
     */
    async exportPublicKey() {
        if (!this.#keyPair) throw new Error("Key pair not generated");

        return this.#marshalKey(this.#keyPair.publicKey);
    }

    /**
     * Set the public key generated by calling `exportPublicKey()` on another party. Provide an identifier
     * if you want to communicate with multiple parties from this one instance.
     * @param remotePublicKey The public key that was exported from another party, serialised as a JWK.
     * @param identifier Optional identifier for the other party, used for multi-cast communication. You can refer to this
     * identifier in future `encrypt()` and `decrypt()` calls to specify which party you want to communicate with.
     */
    async setRemotePublicKey(remotePublicKey: string, identifier: string | symbol = this.#unicast) {
        if (!this.#keyPair) throw new Error("Key pair not generated");

        const unmarshalled = await this.#unmarshalPublicKey(remotePublicKey);
        this.#sharedSecrets[identifier] = await this.#deps.crypto.subtle.deriveKey(
            { name: "ECDH", public: unmarshalled },
            this.#keyPair!.privateKey,
            { name: "AES-CTR", length: this.#params.keyLength },
            false,
            ["encrypt", "decrypt"]
        );
    }

    /**
     * Encrypts the provided string with AES-CTR using the shared secret generated via the key exchange as the key.
     * @param plaintext Any JS string.
     * @param identifier Optional identifier for the other party, used for multi-cast communication.
     * @returns A string representing the ciphertext.
     */
    async encrypt(plaintext: string, identifier: string | symbol = this.#unicast) {
        if (!this.#sharedSecrets[identifier]) throw new Error("Shared secret not set");

        return this.#encryptRaw(this.#UTF8ToUint8Array(plaintext), identifier);
    }

    /**
     * A way to encrypt streaming Uint8Array data.
     * @param identifier Optional identifier for the other party, used for multi-cast communication
     * @returns A TransformStream that can be fitted in any pipeline to encrypt the data flowing through it.
     */
    encryptStream(identifier: string | symbol = this.#unicast) {
        if (!this.#sharedSecrets[identifier]) throw new Error("Shared secret not set");

        return new this.#deps.TransformStream<Uint8Array, string>({
            transform: async (chunk, controller) => {
                controller.enqueue(await this.#encryptRaw(chunk, identifier));
            },
        });
    }

    /**
     * Decrypts the provided string with AES-CTR using the shared secret generated via the key exchange as the key.
     * @param ciphertext A string representing the ciphertext.
     * @param identifier Optional identifier for the other party, used for multi-cast communication.
     * @returns The decrypted string.
     */
    async decrypt(ciphertext: string, identifier: string | symbol = this.#unicast) {
        if (!this.#sharedSecrets[identifier]) throw new Error("Shared secret not set");

        const decrypted = await this.#decryptRaw(ciphertext, identifier);
        return this.#Uint8ArrayToUTF8(decrypted);
    }

    /**
     * Opposite of `encryptStream()`
     * @param identifier Optional identifier for the other party, used for multi-cast communication.
     * @returns A TransformStream that can be fitted in any pipeline to decrypt the data flowing through it.
     */
    decryptStream(identifier: string | symbol = this.#unicast) {
        if (!this.#sharedSecrets[identifier]) throw new Error("Shared secret not set");

        return new this.#deps.TransformStream<string, Uint8Array>({
            transform: async (chunk, controller) => {
                controller.enqueue(await this.#decryptRaw(chunk, identifier));
            },
        });
    }

    /**
     *
     * @returns The parameters used to create this instance of E2EE.
     */
    exportParams() {
        return this.#params;
    }

    /**
     * A secure way to marshal this party's security parameters and key pair.
     * The key pair is returned as a `CryptoKeyPair` object, which means that the `CryptoKey` objects corresponding
     * to the two keys are facades. In particular, the private key cannot be viewed from JavaScript. As such, you cannot
     * serialise the key pair. Instead, store it in IndexedDB.
     * @returns An object with the security parameters of this party, and the key pair if it has been generated.
     */
    marshal() {
        return {
            keyPair: this.#keyPair,
            params: this.exportParams(),
        };
    }

    /**
     * Restore a previously marshalled instance of `E2EE`.
     * @param options
     * @param options.marshalled The object returned by `marshal()`.
     * @param options.deps Optionally inject your own platform dependencies.
     * @returns A new instance of `E2EE` with the same security parameters and key pair as the marshalled instance.
     */
    static unmarshal({ marshalled: { keyPair, params }, deps }) {
        const e2ee = new E2EE({ params, deps });
        e2ee.#restoreKeyPairObject(keyPair);
        return e2ee;
    }

    /**
     * In case you want to share the identity of this party across multiple instances of this class (maybe you want
     * to support multiple devices), use this method to securely share the private key between the devices.
     *
     * SECURITY: Make sure you handle the private key securely. It will be visible to JavaScript.
     * @returns  The private key of this party, serialised as a JWK.
     */
    async exportPrivateKey() {
        if (!this.#keyPair) throw new Error("Key pair not generated");
        if (!this.#keyPair.privateKey.extractable) throw new Error("Private key is not extractable");
        return this.#marshalKey(this.#keyPair!.privateKey);
    }

    /**
     * Used in conjunction with `exportPrivateKey()` to share the private key between multiple instances of this class.
     *
     * SECURITY: Make sure you handle the private key securely. It will be visible to JavaScript.
     *
     * You can also use this to migrate your existing identities into this library.
     * @param options
     * @param options.privateKey The JWK of the private key you want to import.
     * @param options.publicKey The JWK of the public key you want to import.
     */
    async importKeyPair({ privateKey, publicKey }: { privateKey: string; publicKey: string }) {
        if (this.#keyPair) throw new Error("Key pair already exists");

        const unmarshalledPrivateKey = await this.#unmarshalPrivateKey(privateKey);
        const unmarshalledPublicKey = await this.#unmarshalPublicKey(publicKey);
        this.#keyPair = { privateKey: unmarshalledPrivateKey, publicKey: unmarshalledPublicKey };
    }

    #restoreKeyPairObject(keyPair: CryptoKeyPair) {
        this.#keyPair = keyPair;
    }

    async #encryptRaw(data: Uint8Array, identifier: string | symbol) {
        const counter = this.#generateIv();
        const buffer = new Uint8Array(
            await this.#deps.crypto.subtle.encrypt(
                {
                    name: "AES-CTR",
                    counter: counter,
                    length: this.#params.counterLength,
                },
                this.#sharedSecrets[identifier],
                data
            )
        );
        return this.#marshalCiphertext({ buffer, counter });
    }

    async #decryptRaw(data: string, identifier: string | symbol) {
        const { buffer, counter } = this.#unmarshalCiphertext(data);
        const decryptedBuffer = await this.#deps.crypto.subtle.decrypt(
            {
                name: "AES-CTR",
                counter: counter,
                length: this.#params.counterLength,
            },
            this.#sharedSecrets[identifier],
            buffer
        );
        return new Uint8Array(decryptedBuffer);
    }

    async #marshalKey(key: CryptoKey) {
        const exported = await this.#deps.crypto.subtle.exportKey("jwk", key);
        const marshalled = JSON.stringify(exported);
        return marshalled;
    }

    async #unmarshalPublicKey(marshalled: string) {
        const unmarshalled = JSON.parse(marshalled);
        const key = await this.#deps.crypto.subtle.importKey(
            "jwk",
            unmarshalled,
            { name: "ECDH", namedCurve: this.#params.namedCurve },
            true,
            []
        );
        return key;
    }

    async #unmarshalPrivateKey(marshalled: string) {
        const unmarshalled = JSON.parse(marshalled);
        const key = await this.#deps.crypto.subtle.importKey(
            "jwk",
            unmarshalled,
            { name: "ECDH", namedCurve: this.#params.namedCurve },
            true,
            unmarshalled.key_ops
        );
        return key;
    }

    #marshalCiphertext({ buffer, counter }: { buffer: Uint8Array; counter: Uint8Array }) {
        const marshalled = JSON.stringify({
            buffer: this.#Uint8ArrayToAscii(buffer),
            counter: this.#Uint8ArrayToAscii(counter),
        });
        return marshalled;
    }

    #unmarshalCiphertext(marshalled: string) {
        const unmarshalled = JSON.parse(marshalled);
        const buffer = this.#asciiToUint8Array(unmarshalled.buffer);
        const counter = this.#asciiToUint8Array(unmarshalled.counter);
        return { buffer, counter };
    }

    #Uint8ArrayToUTF8(buffer: Uint8Array) {
        return new TextDecoder().decode(buffer);
    }

    #UTF8ToUint8Array(text: string) {
        return new TextEncoder().encode(text);
    }

    #asciiToUint8Array(text: string) {
        return new Uint8Array([...text].map(c => c.charCodeAt(0)));
    }

    #Uint8ArrayToAscii(buffer: Uint8Array) {
        const CHUNK_SIZE = 0x2000;

        if (buffer.length <= CHUNK_SIZE) {
            return String.fromCharCode(...buffer);
        }

        let accum = "";
        for (let i = 0; i < buffer.length; i += CHUNK_SIZE) {
            const chunk = buffer.subarray(i, i + CHUNK_SIZE);
            accum += String.fromCharCode(...chunk);
        }

        return accum;
    }

    #generateIv() {
        return this.#deps.crypto.getRandomValues(new Uint8Array(16));
    }
}