nostr-double-ratchet 0.0.1 → 0.0.3

package/src/Channel.ts

@@ -1,143 +1,231 @@
 import { generateSecretKey, getPublicKey, nip44, finalizeEvent, VerifiedEvent } from "nostr-tools";
-import { hexToBytes } from "@noble/hashes/utils";
+import { bytesToHex } from "@noble/hashes/utils";
 import {
   ChannelState,
-  RatchetMessage,
+  Header,
   Unsubscribe,
   NostrSubscribe,
   MessageCallback,
   EVENT_KIND,
-  KeyPair,
-  Sender,
-  KeyType,
 } from "./types";
-import { kdf } from "./utils";
+import { kdf, skippedMessageIndexKey } from "./utils";
 
+const MAX_SKIP = 1000;
+
+/**
+ * Similar to Signal's "Double Ratchet with header encryption"
+ * https://signal.org/docs/specifications/doubleratchet/
+ */
 export class Channel {
-  nostrUnsubscribe: Unsubscribe | undefined
-  nostrNextUnsubscribe: Unsubscribe | undefined
-  currentInternalSubscriptionId = 0
-  internalSubscriptions = new Map<number, MessageCallback>()
-  name = Math.random().toString(36).substring(2, 6)
+  private nostrUnsubscribe?: Unsubscribe;
+  private nostrNextUnsubscribe?: Unsubscribe;
+  private internalSubscriptions = new Map<number, MessageCallback>();
+  private currentInternalSubscriptionId = 0;
+  public name: string;
 
   constructor(private nostrSubscribe: NostrSubscribe, public state: ChannelState) {
-    this.name = Math.random().toString(36).substring(2, 6)
+    this.name = Math.random().toString(36).substring(2, 6);
   }
 
   /**
-   * To preserve forward secrecy, do not use long-term keys for channel initialization. Use e.g. InviteLink to exchange session keys.
+   * @param sharedSecret optional, but useful to keep the first chain of messages secure. Unlike the Nostr keys, it can be forgotten after the 1st message in the chain.
+   * @param isInitiator determines which chain key is used for sending vs receiving
    */
-  static init(nostrSubscribe: NostrSubscribe, theirCurrentNostrPublicKey: string, ourCurrentPrivateKey: Uint8Array, name?: string): Channel {
-    const ourNextPrivateKey = generateSecretKey()
+  static init(nostrSubscribe: NostrSubscribe, theirNostrPublicKey: string, ourCurrentPrivateKey: Uint8Array, sharedSecret = new Uint8Array(), name?: string, isInitiator = true): Channel {
+    const ourNextPrivateKey = generateSecretKey();
+    const [rootKey, sendingChainKey] = kdf(sharedSecret, nip44.getConversationKey(ourNextPrivateKey, theirNostrPublicKey), 2);
+    let ourCurrentNostrKey;
+    let ourNextNostrKey;
+    if (isInitiator) {
+      ourCurrentNostrKey = { publicKey: getPublicKey(ourCurrentPrivateKey), privateKey: ourCurrentPrivateKey };
+      ourNextNostrKey = { publicKey: getPublicKey(ourNextPrivateKey), privateKey: ourNextPrivateKey };
+    } else {
+      ourNextNostrKey = { publicKey: getPublicKey(ourCurrentPrivateKey), privateKey: ourCurrentPrivateKey };
+    }
     const state: ChannelState = {
-      theirCurrentNostrPublicKey,
-      ourCurrentNostrKey: { publicKey: getPublicKey(ourCurrentPrivateKey), privateKey: ourCurrentPrivateKey },
-      ourNextNostrKey: { publicKey: getPublicKey(ourNextPrivateKey), privateKey: ourNextPrivateKey },
-      receivingChainKey: new Uint8Array(),
-      nextReceivingChainKey: new Uint8Array(),
-      sendingChainKey: new Uint8Array(),
+      rootKey: isInitiator ? rootKey : sharedSecret,
+      theirNostrPublicKey,
+      ourCurrentNostrKey,
+      ourNextNostrKey,
+      receivingChainKey: undefined,
+      sendingChainKey: isInitiator ? sendingChainKey : undefined,
       sendingChainMessageNumber: 0,
       receivingChainMessageNumber: 0,
       previousSendingChainMessageCount: 0,
-      skippedMessageKeys: {}
+      skippedMessageKeys: {},
+    };
+    const channel = new Channel(nostrSubscribe, state);
+    if (name) channel.name = name;
+    console.log(channel.name, 'initial root key', bytesToHex(state.rootKey).slice(0,4))
+    return channel;
+  }
+
+  send(data: string): VerifiedEvent {
+    if (!this.state.theirNostrPublicKey || !this.state.ourCurrentNostrKey) {
+      throw new Error("we are not the initiator, so we can't send the first message");
     }
-    const channel = new Channel(nostrSubscribe, state)
-    channel.updateTheirCurrentNostrPublicKey(theirCurrentNostrPublicKey)
-    if (name) channel.name = name
-    return channel
+
+    const [header, encryptedData] = this.ratchetEncrypt(data);
+    
+    const sharedSecret = nip44.getConversationKey(this.state.ourCurrentNostrKey.privateKey, this.state.theirNostrPublicKey);
+    const encryptedHeader = nip44.encrypt(JSON.stringify(header), sharedSecret);
+    
+    const nostrEvent = finalizeEvent({
+      content: encryptedData,
+      kind: EVENT_KIND,
+      tags: [["header", encryptedHeader]],
+      created_at: Math.floor(Date.now() / 1000)
+    }, this.state.ourCurrentNostrKey.privateKey);
+
+    return nostrEvent;
+  }
+
+  onMessage(callback: MessageCallback): Unsubscribe {
+    const id = this.currentInternalSubscriptionId++
+    this.internalSubscriptions.set(id, callback)
+    this.subscribeToNostrEvents()
+    return () => this.internalSubscriptions.delete(id)
+  }
+
+  private ratchetEncrypt(plaintext: string): [Header, string] {
+    const [newSendingChainKey, messageKey] = kdf(this.state.sendingChainKey!, new Uint8Array([1]), 2);
+    this.state.sendingChainKey = newSendingChainKey;
+    const header: Header = {
+      number: this.state.sendingChainMessageNumber++,
+      nextPublicKey: this.state.ourNextNostrKey.publicKey,
+      time: Date.now(),
+      previousChainLength: this.state.previousSendingChainMessageCount
+    };
+    return [header, nip44.encrypt(plaintext, messageKey)];
   }
 
-  updateTheirCurrentNostrPublicKey(theirNewPublicKey: string) {
-    this.state.theirCurrentNostrPublicKey = theirNewPublicKey
-    this.state.previousSendingChainMessageCount = this.state.sendingChainMessageNumber
-    this.state.sendingChainMessageNumber = 0
-    this.state.receivingChainMessageNumber = 0
-    this.state.receivingChainKey = this.getNostrSenderKeypair(Sender.Them, KeyType.Current).privateKey
-    this.state.nextReceivingChainKey = this.getNostrSenderKeypair(Sender.Them, KeyType.Next).privateKey
-    this.state.sendingChainKey = this.getNostrSenderKeypair(Sender.Us, KeyType.Current).privateKey
+  private ratchetDecrypt(header: Header, ciphertext: string, nostrSender: string): string {
+    const plaintext = this.trySkippedMessageKeys(header, ciphertext, nostrSender);
+    if (plaintext) return plaintext;
+
+    this.skipMessageKeys(header.number, nostrSender);
+    
+    const [newReceivingChainKey, messageKey] = kdf(this.state.receivingChainKey!, new Uint8Array([1]), 2);
+    this.state.receivingChainKey = newReceivingChainKey;
+    this.state.receivingChainMessageNumber++;
+
+    try {
+      return nip44.decrypt(ciphertext, messageKey);
+    } catch (error) {
+      console.error(this.name, 'Decryption failed:', error, {
+        messageKey: bytesToHex(messageKey).slice(0, 4),
+        receivingChainKey: bytesToHex(this.state.receivingChainKey).slice(0, 4),
+        sendingChainKey: this.state.sendingChainKey && bytesToHex(this.state.sendingChainKey).slice(0, 4),
+        rootKey: bytesToHex(this.state.rootKey).slice(0, 4)
+      });
+      throw error;
+    }
   }
 
-  private rotateOurCurrentNostrKey() {
-    this.state.ourCurrentNostrKey = this.state.ourNextNostrKey
-    const ourNextSecretKey = generateSecretKey()
+  private ratchetStep(theirNostrPublicKey: string) {
+    this.state.previousSendingChainMessageCount = this.state.sendingChainMessageNumber;
+    this.state.sendingChainMessageNumber = 0;
+    this.state.receivingChainMessageNumber = 0;
+    this.state.theirNostrPublicKey = theirNostrPublicKey;
+
+    const conversationKey1 = nip44.getConversationKey(this.state.ourNextNostrKey.privateKey, this.state.theirNostrPublicKey!);
+    const [intermediateRootKey, receivingChainKey] = kdf(this.state.rootKey, conversationKey1, 3);
+
+    this.state.receivingChainKey = receivingChainKey;
+
+    this.state.ourCurrentNostrKey = this.state.ourNextNostrKey;
+    const ourNextSecretKey = generateSecretKey();
     this.state.ourNextNostrKey = {
       publicKey: getPublicKey(ourNextSecretKey),
       privateKey: ourNextSecretKey
-    }
+    };
+
+    const conversationKey2 = nip44.getConversationKey(this.state.ourNextNostrKey.privateKey, this.state.theirNostrPublicKey!);
+    const [rootKey2, sendingChainKey] = kdf(intermediateRootKey, conversationKey2, 3);
+    this.state.rootKey = rootKey2;
+    this.state.sendingChainKey = sendingChainKey;
   }
 
-  getNostrSenderKeypair(sender: Sender, keyType: KeyType): KeyPair {
-    if (sender === Sender.Us && keyType === KeyType.Next) {
-      throw new Error("We don't have their next key")
+  private skipMessageKeys(until: number, nostrSender: string) {
+    if (this.state.receivingChainMessageNumber + MAX_SKIP < until) {
+      throw new Error("Too many skipped messages");
     }
-    const ourPrivate = keyType === KeyType.Current ? this.state.ourCurrentNostrKey.privateKey : this.state.ourNextNostrKey.privateKey
-    const theirPublic = this.state.theirCurrentNostrPublicKey
-    const senderPubKey = sender === Sender.Us ? getPublicKey(ourPrivate) : theirPublic
-    const privateKey = kdf(nip44.getConversationKey(ourPrivate, theirPublic), hexToBytes(senderPubKey))
-    return {
-      publicKey: getPublicKey(privateKey),
-      privateKey
+    while (this.state.receivingChainMessageNumber < until) {
+      const [newReceivingChainKey, messageKey] = kdf(this.state.receivingChainKey!, new Uint8Array([1]), 2);
+      this.state.receivingChainKey = newReceivingChainKey;
+      const key = skippedMessageIndexKey(nostrSender, this.state.receivingChainMessageNumber);
+      this.state.skippedMessageKeys[key] = messageKey;
+      this.state.receivingChainMessageNumber++;
     }
   }
 
-  private nostrSubscribeNext() {
-    const nextReceivingPublicKey = this.getNostrSenderKeypair(Sender.Them, KeyType.Next).publicKey
-    const decryptKey = this.state.nextReceivingChainKey
-    this.nostrNextUnsubscribe = this.nostrSubscribe({authors: [nextReceivingPublicKey], kinds: [EVENT_KIND]}, (e) => {
-      // they acknowledged our next key and sent with the corresponding new nostr sender key
-      const msg = JSON.parse(nip44.decrypt(e.content, decryptKey)) as RatchetMessage
-      if (msg.nextPublicKey !== this.state.theirCurrentNostrPublicKey) {
-        this.rotateOurCurrentNostrKey()
-        this.updateTheirCurrentNostrPublicKey(msg.nextPublicKey)
-        this.nostrUnsubscribe?.()
-        this.nostrUnsubscribe = this.nostrNextUnsubscribe
-        this.nostrSubscribeNext()
-      }
-      this.internalSubscriptions.forEach(callback => callback({id: e.id, data: msg.data, pubkey: msg.nextPublicKey, time: msg.time}))
-    })
+  private trySkippedMessageKeys(header: Header, ciphertext: string, nostrSender: string): string | null {
+    const key = skippedMessageIndexKey(nostrSender, header.number);
+    if (key in this.state.skippedMessageKeys) {
+      const mk = this.state.skippedMessageKeys[key];
+      delete this.state.skippedMessageKeys[key];
+      return nip44.decrypt(ciphertext, mk);
+    }
+    return null;
   }
 
-  private subscribeToNostrEvents() {
-    if (this.nostrUnsubscribe) {
-      return
-    }
-    const receivingPublicKey = this.getNostrSenderKeypair(Sender.Them, KeyType.Current).publicKey
-    const decryptKey = this.state.receivingChainKey
-    this.nostrUnsubscribe = this.nostrSubscribe({authors: [receivingPublicKey], kinds: [EVENT_KIND]}, (e) => {
-      const msg = JSON.parse(nip44.decrypt(e.content, decryptKey)) as RatchetMessage
-      if (msg.nextPublicKey !== this.state.theirCurrentNostrPublicKey) {
-        // they announced their next key: we will use it to derive the next nostr sender key
-        this.updateTheirCurrentNostrPublicKey(msg.nextPublicKey)
+  private decryptHeader(e: any): [Header, boolean] {
+    const encryptedHeader = e.tags[0][1];
+    if (this.state.ourCurrentNostrKey) {
+      const currentSecret = nip44.getConversationKey(this.state.ourCurrentNostrKey.privateKey, e.pubkey);
+      try {
+        const header = JSON.parse(nip44.decrypt(encryptedHeader, currentSecret)) as Header;
+        return [header, false];
+      } catch (error) {
+        // Decryption with currentSecret failed, try with nextSecret
       }
-      this.internalSubscriptions.forEach(callback => callback({id: e.id, data: msg.data, pubkey: msg.nextPublicKey, time: msg.time}))
-    })
-    this.nostrSubscribeNext()
-  }
+    }
 
-  onMessage(callback: MessageCallback): Unsubscribe {
-    const id = this.currentInternalSubscriptionId++
-    this.internalSubscriptions.set(id, callback)
-    this.subscribeToNostrEvents()
-    return () => this.internalSubscriptions.delete(id)
+    const nextSecret = nip44.getConversationKey(this.state.ourNextNostrKey.privateKey, e.pubkey);
+    try {
+      const header = JSON.parse(nip44.decrypt(encryptedHeader, nextSecret)) as Header;
+      return [header, true];
+    } catch (error) {
+      // Decryption with nextSecret also failed
+    }
+
+    throw new Error("Failed to decrypt header with both current and next secrets");
   }
 
-  send(data: string): VerifiedEvent {
-    const message: RatchetMessage = {
-      number: this.state.sendingChainMessageNumber,
-      data: data,
-      nextPublicKey: this.state.ourNextNostrKey.publicKey,
-      time: Date.now()
+  private handleNostrEvent(e: any) {
+    const [header, shouldRatchet] = this.decryptHeader(e);
+
+    if (this.state.theirNostrPublicKey !== header.nextPublicKey) {
+      this.state.theirNostrPublicKey = header.nextPublicKey;
+      this.nostrUnsubscribe?.(); // should we keep this open for a while? maybe as long as we have skipped messages?
+      this.nostrUnsubscribe = this.nostrNextUnsubscribe;
+      this.nostrNextUnsubscribe = this.nostrSubscribe(
+        {authors: [this.state.theirNostrPublicKey], kinds: [EVENT_KIND]},
+        (e) => this.handleNostrEvent(e)
+      );
     }
-    this.state.sendingChainMessageNumber++
-    const sendingPrivateKey = this.getNostrSenderKeypair(Sender.Us, KeyType.Current).privateKey
-    const encryptedData = nip44.encrypt(JSON.stringify(message), this.state.sendingChainKey)
-    const nostrEvent = finalizeEvent({
-      content: encryptedData,
-      kind: EVENT_KIND,
-      tags: [],
-      created_at: Math.floor(Date.now() / 1000)
-    }, sendingPrivateKey)
 
-    return nostrEvent
+    if (shouldRatchet) {
+      this.skipMessageKeys(header.previousChainLength, e.pubkey);
+      this.ratchetStep(header.nextPublicKey);
+    }
+
+    const data = this.ratchetDecrypt(header, e.content, e.pubkey);
+
+    this.internalSubscriptions.forEach(callback => callback({id: e.id, data, pubkey: header.nextPublicKey, time: header.time}));  
+  }
+
+  private subscribeToNostrEvents() {
+    if (this.nostrNextUnsubscribe) return;
+    if (this.state.theirNostrPublicKey) {
+      this.nostrUnsubscribe = this.nostrSubscribe(
+        {authors: [this.state.theirNostrPublicKey], kinds: [EVENT_KIND]},
+        (e) => this.handleNostrEvent(e)
+      );
+    }
+    this.nostrNextUnsubscribe = this.nostrSubscribe(
+      {authors: [this.state.theirNostrPublicKey], kinds: [EVENT_KIND]},
+      (e) => this.handleNostrEvent(e)
+    );
   }
 }

package/src/InviteLink.ts

@@ -127,7 +127,7 @@ export class InviteLink {
         const inviteeSessionPublicKey = getPublicKey(inviteeSessionKey);
         const inviterPublicKey = this.inviter || this.inviterSessionPublicKey;
 
-        const channel = Channel.init(nostrSubscribe, this.inviterSessionPublicKey, inviteeSessionKey);
+        const channel = Channel.init(nostrSubscribe, this.inviterSessionPublicKey, inviteeSessionKey, new Uint8Array(), undefined, true);
 
         // Create a random keypair for the envelope sender
         const randomSenderKey = generateSecretKey();
@@ -181,7 +181,8 @@ export class InviteLink {
     
                 const inviteeSessionPublicKey = await innerDecrypt(innerEvent.content, innerEvent.pubkey);
 
-                const channel = Channel.init(nostrSubscribe, inviteeSessionPublicKey, this.inviterSessionPrivateKey!);
+                const name = event.id;
+                const channel = Channel.init(nostrSubscribe, inviteeSessionPublicKey, this.inviterSessionPrivateKey!, new Uint8Array(), name, false);
 
                 onChannel(channel, innerEvent.pubkey);
             } catch (error) {

package/src/types.ts

@@ -7,9 +7,9 @@ export type Message = {
   time: number; // unlike Nostr, we use milliseconds instead of seconds
 }
 
-export type RatchetMessage = {
+export type Header = {
   number: number;
-  data: string;
+  previousChainLength: number;
   nextPublicKey: string;
   time: number;
 }
@@ -24,17 +24,39 @@ export type KeyPair = {
   privateKey: Uint8Array;
 }
 
+/** 
+ * Represents the state of a Double Ratchet channel between two parties. Needed for persisting channels.
+ */
 export interface ChannelState {
-  theirCurrentNostrPublicKey: string;
-  ourCurrentNostrKey: KeyPair;
+  /** Root key used to derive new sending / receiving chain keys */
+  rootKey: Uint8Array;
+  
+  /** The other party's current Nostr public key */
+  theirNostrPublicKey: string;
+
+  /** Our current Nostr keypair used for this channel */
+  ourCurrentNostrKey?: KeyPair;
+  
+  /** Our next Nostr keypair, used when ratcheting forward. It is advertised in messages we send. */
   ourNextNostrKey: KeyPair;
-  receivingChainKey: Uint8Array;
-  nextReceivingChainKey: Uint8Array;
-  sendingChainKey: Uint8Array;
+  
+  /** Key for decrypting incoming messages in current chain */
+  receivingChainKey?: Uint8Array;
+  
+  /** Key for encrypting outgoing messages in current chain */
+  sendingChainKey?: Uint8Array;
+  
+  /** Number of messages sent in current sending chain */
   sendingChainMessageNumber: number;
+  
+  /** Number of messages received in current receiving chain */
   receivingChainMessageNumber: number;
+  
+  /** Number of messages sent in previous sending chain */
   previousSendingChainMessageCount: number;
-  skippedMessageKeys: Record<number, Uint8Array>;
+  
+  /** Cache of message keys for handling out-of-order messages */
+  skippedMessageKeys: Record<string, Uint8Array>;
 }
 
 export type Unsubscribe = () => void;
@@ -59,10 +81,5 @@ export enum Sender {
   Them
 }
 
-export enum KeyType {
-  Current,
-  Next
-}
-
 export type EncryptFunction = (plaintext: string, pubkey: string) => Promise<string>;
 export type DecryptFunction = (ciphertext: string, pubkey: string) => Promise<string>;

package/src/utils.ts

@@ -6,18 +6,18 @@ import { sha256 } from '@noble/hashes/sha256';
 
 export function serializeChannelState(state: ChannelState): string {
   return JSON.stringify({
-    theirCurrentNostrPublicKey: state.theirCurrentNostrPublicKey,
-    ourCurrentNostrKey: {
+    rootKey: bytesToHex(state.rootKey),
+    theirNostrPublicKey: state.theirNostrPublicKey,
+    ourCurrentNostrKey: state.ourCurrentNostrKey ? {
       publicKey: state.ourCurrentNostrKey.publicKey,
       privateKey: bytesToHex(state.ourCurrentNostrKey.privateKey),
-    },
+    } : undefined,
     ourNextNostrKey: {
       publicKey: state.ourNextNostrKey.publicKey,
       privateKey: bytesToHex(state.ourNextNostrKey.privateKey),
     },
-    receivingChainKey: bytesToHex(state.receivingChainKey),
-    nextReceivingChainKey: bytesToHex(state.nextReceivingChainKey),
-    sendingChainKey: bytesToHex(state.sendingChainKey),
+    receivingChainKey: state.receivingChainKey ? bytesToHex(state.receivingChainKey) : undefined,
+    sendingChainKey: state.sendingChainKey ? bytesToHex(state.sendingChainKey) : undefined,
     sendingChainMessageNumber: state.sendingChainMessageNumber,
     receivingChainMessageNumber: state.receivingChainMessageNumber,
     previousSendingChainMessageCount: state.previousSendingChainMessageCount,
@@ -33,18 +33,18 @@ export function serializeChannelState(state: ChannelState): string {
 export function deserializeChannelState(data: string): ChannelState {
   const state = JSON.parse(data);
   return {
-    theirCurrentNostrPublicKey: state.theirCurrentNostrPublicKey,
-    ourCurrentNostrKey: {
+    rootKey: hexToBytes(state.rootKey),
+    theirNostrPublicKey: state.theirNostrPublicKey,
+    ourCurrentNostrKey: state.ourCurrentNostrKey ? {
       publicKey: state.ourCurrentNostrKey.publicKey,
       privateKey: hexToBytes(state.ourCurrentNostrKey.privateKey),
-    },
+    } : undefined,
     ourNextNostrKey: {
       publicKey: state.ourNextNostrKey.publicKey,
       privateKey: hexToBytes(state.ourNextNostrKey.privateKey),
     },
-    receivingChainKey: hexToBytes(state.receivingChainKey),
-    nextReceivingChainKey: hexToBytes(state.nextReceivingChainKey),
-    sendingChainKey: hexToBytes(state.sendingChainKey),
+    receivingChainKey: state.receivingChainKey ? hexToBytes(state.receivingChainKey) : undefined,
+    sendingChainKey: state.sendingChainKey ? hexToBytes(state.sendingChainKey) : undefined,
     sendingChainMessageNumber: state.sendingChainMessageNumber,
     receivingChainMessageNumber: state.receivingChainMessageNumber,
     previousSendingChainMessageCount: state.previousSendingChainMessageCount,
@@ -85,7 +85,16 @@ export async function* createMessageStream(channel: Channel): AsyncGenerator<Mes
   }
 }
 
-export function kdf(input1: Uint8Array, input2: Uint8Array = new Uint8Array(32)) {
-  const prk = hkdf_extract(sha256, input1, input2)
-  return hkdf_expand(sha256, prk, new Uint8Array([1]), 32)
+export function kdf(input1: Uint8Array, input2: Uint8Array = new Uint8Array(32), numOutputs: number = 1): Uint8Array[] {
+  const prk = hkdf_extract(sha256, input1, input2);
+  
+  const outputs: Uint8Array[] = [];
+  for (let i = 1; i <= numOutputs; i++) {
+    outputs.push(hkdf_expand(sha256, prk, new Uint8Array([i]), 32));
+  }
+  return outputs;
 }
+
+export function skippedMessageIndexKey(nostrSender: string, number: number): string {
+  return `${nostrSender}:${number}`;
+}