Eve-Relay/mls.ts

import type { NPub, NSec } from "@nostr/tools/nip19";
import * as nip19 from "@nostr/tools/nip19";
import * as nostrTools from "@nostr/tools";
import { hkdf } from "npm:@noble/hashes@1.3.1/hkdf";
import { sha256 } from "npm:@noble/hashes@1.3.1/sha256";
import * as secp256k1 from "@noble/secp256k1";
import { randomBytes } from "@noble/ciphers/webcrypto";
import { xchacha20poly1305 } from "@noble/ciphers/chacha";
import { bytesEqual, getCCNPrivateKey } from "./utils.ts";
import { decodeBase64, encodeBase64 } from "jsr:@std/encoding/base64";
import {
  decryptUint8Array,
  encryptionKey,
  encryptUint8Array,
} from "./utils/encryption.ts";

type CCNId = string;
type Epoch = number;

interface CCNState {
  id: CCNId;
  epoch: Epoch;
  members: Map<NPub, NPub>;
  ccnKey: Uint8Array;
  privateKey: NSec;
  ratchetTree: RatchetNode[];
}

interface RatchetNode {
  publicKey: NPub | null;
  path: number[];
  parent: number | null;
  children: number[];
}

interface EncryptedMessage {
  sender: NPub;
  epoch: Epoch;
  ciphertext: Uint8Array;
  nonce: Uint8Array;
  signature: Uint8Array;
}

interface Welcome {
  ccnId: CCNId;
  epoch: Epoch;
  encryptedCCNState: Uint8Array;
  pathSecret: Uint8Array;
  inviterPublicKey: NPub;
  nonce: Uint8Array;
}

interface SerializedCCN {
  version: number;
  id: CCNId;
  epoch: Epoch;
  members: [NPub, NPub][];
  ccnKey: string;
  privateKey: NSec;
  ratchetTree: SerializedRatchetNode[];
}

interface SerializedRatchetNode {
  publicKey: NPub | null;
  path: number[];
  parent: number | null;
  children: number[];
}

const textEncoder = new TextEncoder();
const textDecoder = new TextDecoder();

function generateNonce(): Uint8Array {
  return randomBytes(24);
}

function deriveCCNKey(
  secret: Uint8Array,
  ccnId: CCNId,
  epoch: Epoch,
): Uint8Array {
  const info = textEncoder.encode(`MLS CCN Key ${ccnId} ${epoch}`);
  return hkdf(sha256, secret, new Uint8Array(0), info, 32);
}

function derivePathSecret(secret: Uint8Array, path: number[]): Uint8Array {
  const pathStr = path.join("/");
  const info = textEncoder.encode(`MLS Path Secret ${pathStr}`);
  return hkdf(sha256, secret, new Uint8Array(0), info, 32);
}

interface UpdateMessage {
  sender: NPub;
  ccnId: CCNId;
  epoch: Epoch;
  pathSecrets: Map<NPub, EncryptedPathSecret>;
  updatedPublicKeys: Map<number, NPub>;
  signature: Uint8Array;
}

interface EncryptedPathSecret {
  ciphertext: Uint8Array;
  nonce: Uint8Array;
}

export class MLS {
  private ccns: Map<CCNId, CCNState> = new Map();
  private identity: { publicKey: NPub; privateKey: NSec };

  constructor(privateKey?: NSec) {
    const privateKeyBytes = privateKey
      ? nip19.decode(privateKey).data
      : nostrTools.generateSecretKey();
    if (!privateKey) {
      privateKey = nip19.nsecEncode(privateKeyBytes);
    }
    this.identity = {
      privateKey: privateKey,
      publicKey: nip19.npubEncode(
        nostrTools.getPublicKey(privateKeyBytes),
      ),
    };
  }

  get memberId(): NPub {
    return nip19.npubEncode(nostrTools.getPublicKey(
      nip19.decode(this.identity.privateKey).data,
    ));
  }

  createCCN(ccnId: CCNId): CCNId {
    const secret = randomBytes(32);
    const epoch = 0;
    const ccnKey = deriveCCNKey(secret, ccnId, epoch);

    const ratchetTree: RatchetNode[] = [
      {
        publicKey: this.identity.publicKey,
        path: [0],
        parent: null,
        children: [],
      },
    ];

    const ccnState: CCNState = {
      id: ccnId,
      epoch,
      members: new Map([[this.memberId, this.identity.publicKey]]),
      ccnKey,
      privateKey: this.identity.privateKey,
      ratchetTree,
    };

    this.ccns.set(ccnId, ccnState);
    return ccnId;
  }

  serializeCCN(ccnId: CCNId) {
    const CCN = this.ccns.get(ccnId);
    if (!CCN) {
      throw new Error(`CCN ${ccnId} not found`);
    }
    const serializedCCN: SerializedCCN = {
      version: 1, // For future compatibility
      id: CCN.id,
      epoch: CCN.epoch,
      members: Array.from(CCN.members.entries()),
      ccnKey: encodeBase64(CCN.ccnKey),
      privateKey: CCN.privateKey,
      ratchetTree: CCN.ratchetTree,
    };
    const json = JSON.stringify(serializedCCN);
    const encrypted = encryptUint8Array(
      textEncoder.encode(json),
      encryptionKey,
    );
    return encodeBase64(encrypted);
  }

  deserializeCCN(serialized: string) {
    const encrypted = decodeBase64(serialized);
    const decrypted = decryptUint8Array(encrypted, encryptionKey);
    const json = textDecoder.decode(decrypted);
    const serializedCCN = JSON.parse(json) as SerializedCCN;

    if (serializedCCN.version !== 1) {
      throw new Error(
        `Unknown CCN serialization version: ${serializedCCN.version}`,
      );
    }

    this.ccns.set(
      serializedCCN.id,
      {
        id: serializedCCN.id,
        epoch: serializedCCN.epoch,
        members: new Map(serializedCCN.members),
        ccnKey: decodeBase64(serializedCCN.ccnKey),
        privateKey: serializedCCN.privateKey,
        ratchetTree: serializedCCN.ratchetTree,
      },
    );
  }

  addMember(
    ccnId: CCNId,
    memberId: NPub,
  ): Welcome | null {
    const CCN = this.ccns.get(ccnId);
    if (!CCN) return null;

    CCN.epoch += 1;
    CCN.members.set(memberId, memberId);

    const newNodeIndex = CCN.ratchetTree.length;
    CCN.ratchetTree.push({
      publicKey: memberId,
      path: [newNodeIndex],
      parent: null,
      children: [],
    });

    const newSecret = randomBytes(32);
    CCN.ccnKey = deriveCCNKey(newSecret, ccnId, CCN.epoch);

    const pathSecret = derivePathSecret(newSecret, [newNodeIndex]);
    const serializedState = JSON.stringify({
      id: CCN.id,
      epoch: CCN.epoch,
      members: Array.from(CCN.members.entries()),
      ratchetTree: CCN.ratchetTree,
    });
    const nonce = generateNonce();
    const sharedSecret = secp256k1.getSharedSecret(
      this.identity.privateKey,
      memberId,
      true,
    ).slice(1); // Remove the prefix byte

    const cipher = xchacha20poly1305(sharedSecret, nonce);
    const encryptedCCNState = cipher.encrypt(
      textEncoder.encode(serializedState),
    );

    const welcome: Welcome = {
      ccnId,
      epoch: CCN.epoch,
      encryptedCCNState,
      pathSecret,
      inviterPublicKey: this.identity.publicKey,
      nonce,
    };

    return welcome;
  }

  joinCCN(welcome: Welcome): CCNId | null {
    try {
      const {
        ccnId,
        epoch,
        encryptedCCNState,
        pathSecret,
        inviterPublicKey,
        nonce,
      } = welcome;

      const ccnKey = deriveCCNKey(pathSecret, ccnId, epoch);

      const sharedSecret = secp256k1.getSharedSecret(
        this.identity.privateKey,
        inviterPublicKey,
        true,
      ).slice(1); // Remove the prefix byte

      const decipher = xchacha20poly1305(sharedSecret, nonce);
      const decryptedStateBytes = decipher.decrypt(encryptedCCNState);
      const decryptedState = JSON.parse(
        textDecoder.decode(decryptedStateBytes),
      );

      const members = new Map<NPub, NPub>(decryptedState.members);
      const ratchetTree = decryptedState.ratchetTree;

      const ccnState: CCNState = {
        id: ccnId,
        epoch,
        members,
        ccnKey,
        privateKey: this.identity.privateKey,
        ratchetTree,
      };

      this.ccns.set(ccnId, ccnState);
      return ccnId;
    } catch (error) {
      console.error("Failed to join ccn:", error);
      return null;
    }
  }

  encryptMessage(ccnId: CCNId, plaintext: string): EncryptedMessage | null {
    const CCN = this.ccns.get(ccnId);
    if (!CCN) return null;

    const nonce = generateNonce();
    const cipher = xchacha20poly1305(CCN.ccnKey, nonce);
    const ciphertext = cipher.encrypt(textEncoder.encode(plaintext));

    const messageToSign = new Uint8Array([
      ...textEncoder.encode(ccnId),
      ...new Uint8Array(new BigUint64Array([BigInt(CCN.epoch)]).buffer),
      ...ciphertext,
      ...nonce,
    ]);

    const messageHash = sha256(messageToSign);
    const signature = secp256k1.sign(messageHash, this.identity.privateKey);

    return {
      sender: this.memberId,
      epoch: CCN.epoch,
      ciphertext,
      nonce,
      signature: signature.toCompactRawBytes(),
    };
  }

  decryptMessage(ccnId: CCNId, message: EncryptedMessage): string | null {
    const CCN = this.ccns.get(ccnId);
    if (!CCN) return null;

    if (message.epoch !== CCN.epoch) {
      console.error("Message from wrong epoch");
      return null;
    }

    const senderPublicKey = CCN.members.get(message.sender);
    if (!senderPublicKey) {
      console.error("Message from unknown sender");
      return null;
    }

    // Verify signature
    const messageToVerify = new Uint8Array([
      ...textEncoder.encode(ccnId),
      ...new Uint8Array(new BigUint64Array([BigInt(message.epoch)]).buffer),
      ...message.ciphertext,
      ...message.nonce,
    ]);

    const messageHash = sha256(messageToVerify);
    const isValid = secp256k1.verify(
      message.signature,
      messageHash,
      senderPublicKey,
    );

    if (!isValid) {
      console.error("Invalid message signature");
      return null;
    }

    try {
      const decipher = xchacha20poly1305(CCN.ccnKey, message.nonce);
      const plaintext = decipher.decrypt(message.ciphertext);
      return textDecoder.decode(plaintext);
    } catch (error) {
      console.error("Failed to decrypt message:", error);
      return null;
    }
  }

  updateCCNKey(ccnId: CCNId): UpdateMessage | null {
    const CCN = this.ccns.get(ccnId);
    if (!CCN) return null;

    const newEpoch = CCN.epoch + 1;
    const newRootSecret = randomBytes(32);
    const updatedTree = [...CCN.ratchetTree];
    const myLeafIndex = updatedTree.findIndex((node) =>
      node.publicKey &&
      bytesEqual(node.publicKey, this.identity.publicKey)
    );

    if (myLeafIndex === -1) {
      console.error("Cannot find self in ratchet tree");
      return null;
    }

    const myPath: number[] = [myLeafIndex];
    let currentNode = myLeafIndex;

    while (updatedTree[currentNode].parent !== null) {
      currentNode = updatedTree[currentNode].parent!;
      myPath.unshift(currentNode);
    }

    const pathSecrets: Uint8Array[] = [newRootSecret];
    const updatedPublicKeys = new Map<number, NPub>();

    for (let i = 0; i < myPath.length; i++) {
      const nodeIndex = myPath[i];

      let newPublicKey: NPub;

      if (i === 0) {
        newPublicKey = nip19.npubEncode(nostrTools.getPublicKey(newRootSecret));
      } else {
        const parentIndex = myPath[i - 1];
        const childSecret = derivePathSecret(pathSecrets[i - 1], [
          parentIndex,
          nodeIndex,
        ]);
        pathSecrets.push(childSecret);

        newPublicKey = nip19.npubEncode(nostrTools.getPublicKey(childSecret));
      }

      updatedTree[nodeIndex].publicKey = newPublicKey;
      updatedPublicKeys.set(nodeIndex, newPublicKey);
    }

    const encryptedPathSecrets = new Map<NPub, EncryptedPathSecret>();

    for (const [memberId, memberPublicKey] of CCN.members.entries()) {
      if (memberId === this.memberId) continue;

      const memberLeafIndex = updatedTree.findIndex((node) =>
        node.publicKey &&
        bytesEqual(node.publicKey, memberPublicKey)
      );

      if (memberLeafIndex === -1) continue;

      let memberNode = memberLeafIndex;
      const memberPath: number[] = [memberLeafIndex];

      while (updatedTree[memberNode].parent !== null) {
        memberNode = updatedTree[memberNode].parent!;
        memberPath.unshift(memberNode);
      }

      let commonAncestorIndex = 0;
      while (
        commonAncestorIndex < myPath.length &&
        commonAncestorIndex < memberPath.length &&
        myPath[commonAncestorIndex] === memberPath[commonAncestorIndex]
      ) {
        commonAncestorIndex++;
      }

      const secretToShare = pathSecrets[commonAncestorIndex - 1];

      const nonce = generateNonce();
      const sharedSecret = secp256k1.getSharedSecret(
        this.identity.privateKey,
        memberPublicKey,
        true,
      ).slice(1);

      const cipher = xchacha20poly1305(sharedSecret, nonce);
      const ciphertext = cipher.encrypt(secretToShare);

      encryptedPathSecrets.set(memberId, { ciphertext, nonce });
    }

    const updateData = new Uint8Array([
      ...textEncoder.encode(ccnId),
      ...new Uint8Array(new BigUint64Array([BigInt(newEpoch)]).buffer),
      ...textEncoder.encode(
        JSON.stringify(Array.from(encryptedPathSecrets.entries())),
      ),
      ...textEncoder.encode(
        JSON.stringify(Array.from(updatedPublicKeys.entries())),
      ),
    ]);

    const updateHash = sha256(updateData);
    const signature = secp256k1.sign(updateHash, this.identity.privateKey);

    CCN.epoch = newEpoch;
    CCN.ratchetTree = updatedTree;
    CCN.ccnKey = deriveCCNKey(newRootSecret, ccnId, newEpoch);

    const updateMessage: UpdateMessage = {
      sender: this.memberId,
      ccnId,
      epoch: newEpoch,
      pathSecrets: encryptedPathSecrets,
      updatedPublicKeys,
      signature: signature.toCompactRawBytes(),
    };

    return updateMessage;
  }

  processUpdate(updateMessage: UpdateMessage): boolean {
    const {
      sender,
      ccnId,
      epoch,
      pathSecrets,
      updatedPublicKeys,
      signature,
    } = updateMessage;
    const CCN = this.ccns.get(ccnId);
    if (!CCN) {
      console.error("Unknown CCN");
      return false;
    }
    if (epoch !== CCN.epoch + 1) {
      console.error("Invalid epoch in update");
      return false;
    }
    const senderPublicKey = CCN.members.get(sender);
    if (!senderPublicKey) {
      console.error("Update from unknown sender");
      return false;
    }

    const updateData = new Uint8Array([
      ...textEncoder.encode(ccnId),
      ...new Uint8Array(new BigUint64Array([BigInt(epoch)]).buffer),
      ...textEncoder.encode(JSON.stringify(Array.from(pathSecrets.entries()))),
      ...textEncoder.encode(
        JSON.stringify(Array.from(updatedPublicKeys.entries())),
      ),
    ]);

    const updateHash = sha256(updateData);
    const isValid = secp256k1.verify(
      signature,
      updateHash,
      senderPublicKey,
    );

    if (!isValid) {
      console.error("Invalid update signature");
      return false;
    }

    const myPathSecret = pathSecrets.get(this.memberId);
    if (!myPathSecret) {
      console.error("No path secret for me in update");
      return false;
    }

    try {
      const sharedSecret = secp256k1.getSharedSecret(
        this.identity.privateKey,
        senderPublicKey,
        true,
      ).slice(1);

      const decipher = xchacha20poly1305(sharedSecret, myPathSecret.nonce);
      const decryptedSecret = decipher.decrypt(myPathSecret.ciphertext);

      const myLeafIndex = CCN.ratchetTree.findIndex((node) =>
        node.publicKey &&
        bytesEqual(node.publicKey, this.identity.publicKey)
      );

      if (myLeafIndex === -1) {
        console.error("Cannot find self in ratchet tree");
        return false;
      }

      const myPath: number[] = [myLeafIndex];
      let currentNode = myLeafIndex;

      while (CCN.ratchetTree[currentNode].parent !== null) {
        currentNode = CCN.ratchetTree[currentNode].parent!;
        myPath.unshift(currentNode);
      }

      const updatedTree = [...CCN.ratchetTree];
      for (const [nodeIndex, newPublicKey] of updatedPublicKeys.entries()) {
        const index = Number(nodeIndex);
        if (index >= 0 && index < updatedTree.length) {
          updatedTree[index].publicKey = newPublicKey;
        }
      }

      CCN.epoch = epoch;
      CCN.ratchetTree = updatedTree;
      CCN.ccnKey = deriveCCNKey(decryptedSecret, ccnId, epoch);

      return true;
    } catch (error) {
      console.error("Failed to process update:", error);
      return false;
    }
  }

  generateInvite(ccnId: CCNId, recipientPubkey: NPub): string {
    const welcome = this.addMember(
      ccnId,
      recipientPubkey,
    );

    if (!welcome) {
      throw new Error("Failed to create welcome message");
    }

    const ccn = this.ccns.get(ccnId);
    if (!ccn) {
      throw new Error("CCN not found");
    }

    const ccnPrivateKey = ccn.ccnKey;
    const conversationKey = nostrTools.nip44.getConversationKey(
      ccnPrivateKey,
      recipientPubkey,
    );

    const invite = {
      type: "ccn-invite",
      version: 1,
      ccnId,
      welcome,
      metadata: {
        createdAt: Math.floor(Date.now() / 1000),
        memberCount: ccn.members.size,
      },
    };

    const encryptedInvite = nostrTools.nip44.encrypt(
      JSON.stringify(invite),
      conversationKey,
    );

    return `ccn:${ccnId}:${btoa(encryptedInvite)}`;
  }

  acceptInvite(inviteString: string): CCNId | null {
    try {
      // Parse the invite string format: ccn:ccnId:encryptedData
      const parts = inviteString.split(":");
      if (parts.length !== 3 || parts[0] !== "ccn") {
        throw new Error("Invalid invite format");
      }

      const ccnId = parts[1];
      const encryptedInvite = atob(parts[2]);

      const parsedData = JSON.parse(nostrTools.nip44.decrypt(
        encryptedInvite,
        nip19.decode<"nsec">(this.identity.privateKey).data,
      ));

      if (parsedData.type !== "ccn-invite" || parsedData.version !== 1) {
        throw new Error("Invalid invite type or version");
      }

      const welcome = parsedData.welcome;

      const joinResult = this.joinCCN(welcome);

      if (!joinResult) {
        throw new Error("Failed to join CCN");
      }

      if (!this.ccns.has(ccnId)) {
        throw new Error("CCN was not properly added to the ccns map");
      }

      return ccnId;
    } catch (error) {
      console.error("Failed to accept invite:", error);
      return null;
    }
  }
}