nostr-double-ratchet 0.0.21 → 0.0.23

package/src/Session.ts

@@ -10,7 +10,7 @@ import {
   Rumor,
   CHAT_MESSAGE_KIND,
 } from "./types";
-import { kdf, skippedMessageIndexKey } from "./utils";
+import { kdf } from "./utils";
 
 const MAX_SKIP = 1000;
 
@@ -66,8 +66,7 @@ export class Session {
       sendingChainMessageNumber: 0,
       receivingChainMessageNumber: 0,
       previousSendingChainMessageCount: 0,
-      skippedMessageKeys: {},
-      skippedHeaderKeys: {},
+      skippedKeys: {},
     };
     const session = new Session(nostrSubscribe, state);
     if (name) session.name = name;
@@ -220,43 +219,45 @@ export class Session {
     if (this.state.receivingChainMessageNumber + MAX_SKIP < until) {
       throw new Error("Too many skipped messages");
     }
+
+    if (!this.state.skippedKeys[nostrSender]) {
+      this.state.skippedKeys[nostrSender] = {
+        headerKeys: [],
+        messageKeys: {}
+      };
+      
+      // Store header keys
+      if (this.state.ourCurrentNostrKey) {
+        const currentSecret = nip44.getConversationKey(this.state.ourCurrentNostrKey.privateKey, nostrSender);
+        this.state.skippedKeys[nostrSender].headerKeys.push(currentSecret);
+      }
+      const nextSecret = nip44.getConversationKey(this.state.ourNextNostrKey.privateKey, nostrSender);
+      this.state.skippedKeys[nostrSender].headerKeys.push(nextSecret);
+    }
+
     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;
-      
-      if (!this.state.skippedHeaderKeys[nostrSender]) {
-        const secrets: Uint8Array[] = [];
-        if (this.state.ourCurrentNostrKey) {
-          const currentSecret = nip44.getConversationKey(this.state.ourCurrentNostrKey.privateKey, nostrSender);
-          secrets.push(currentSecret);
-        }
-        const nextSecret = nip44.getConversationKey(this.state.ourNextNostrKey.privateKey, nostrSender);
-        secrets.push(nextSecret);
-        this.state.skippedHeaderKeys[nostrSender] = secrets;
-      }
-      
+      this.state.skippedKeys[nostrSender].messageKeys[this.state.receivingChainMessageNumber] = messageKey;
       this.state.receivingChainMessageNumber++;
     }
   }
 
   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];
-      
-      // Check if we have any remaining skipped messages from this sender
-      const hasMoreSkippedMessages = Object.keys(this.state.skippedMessageKeys).some(k => k.startsWith(`${nostrSender}:`));
-      if (!hasMoreSkippedMessages) {
-        // Clean up header keys as no more skipped messages from this sender
-        delete this.state.skippedHeaderKeys[nostrSender];
-      }
-      
-      return nip44.decrypt(ciphertext, mk);
+    const skippedKeys = this.state.skippedKeys[nostrSender];
+    if (!skippedKeys) return null;
+
+    const messageKey = skippedKeys.messageKeys[header.number];
+    if (!messageKey) return null;
+
+    delete skippedKeys.messageKeys[header.number];
+    
+    // Clean up if no more skipped messages
+    if (Object.keys(skippedKeys.messageKeys).length === 0) {
+      delete this.state.skippedKeys[nostrSender];
     }
-    return null;
+    
+    return nip44.decrypt(ciphertext, messageKey);
   }
 
   // 4. NOSTR EVENT HANDLING
@@ -280,9 +281,9 @@ export class Session {
       // Decryption with nextSecret also failed
     }
 
-    const keys = this.state.skippedHeaderKeys[event.pubkey];
-    if (keys) {
-      for (const key of keys) {
+    const skippedKeys = this.state.skippedKeys[event.pubkey];
+    if (skippedKeys?.headerKeys) {
+      for (const key of skippedKeys.headerKeys) {
         try {
           const header = JSON.parse(nip44.decrypt(encryptedHeader, key)) as Header;
           return [header, false, true];
@@ -302,7 +303,7 @@ export class Session {
       if (this.state.theirNextNostrPublicKey !== header.nextPublicKey) {
         this.state.theirCurrentNostrPublicKey = this.state.theirNextNostrPublicKey;
         this.state.theirNextNostrPublicKey = header.nextPublicKey;
-        this.nostrUnsubscribe?.(); // should we keep this open for a while? maybe as long as we have skipped messages?
+        this.nostrUnsubscribe?.();
         this.nostrUnsubscribe = this.nostrNextUnsubscribe;
         this.nostrNextUnsubscribe = this.nostrSubscribe(
           {authors: [this.state.theirNextNostrPublicKey], kinds: [MESSAGE_EVENT_KIND]},
@@ -314,11 +315,6 @@ export class Session {
         this.skipMessageKeys(header.previousChainLength, e.pubkey);
         this.ratchetStep(header.nextPublicKey);
       }
-    } else {
-      const key = skippedMessageIndexKey(e.pubkey, header.number);
-      if (!(key in this.state.skippedMessageKeys)) {
-        return // maybe we already processed this message
-      }
     }
 
     const text = this.ratchetDecrypt(header, e.content, e.pubkey);
@@ -343,15 +339,19 @@ export class Session {
       (e) => this.handleNostrEvent(e)
     );
 
-    const authors = Object.keys(this.state.skippedHeaderKeys);
-    if (this.state.theirCurrentNostrPublicKey && !authors.includes(this.state.theirCurrentNostrPublicKey)) {
-      authors.push(this.state.theirCurrentNostrPublicKey)
+    if (this.state.theirCurrentNostrPublicKey) {
+      this.nostrUnsubscribe = this.nostrSubscribe(
+        {authors: [this.state.theirCurrentNostrPublicKey], kinds: [MESSAGE_EVENT_KIND]},
+        (e) => this.handleNostrEvent(e)
+      );  
+    }
+
+    const skippedAuthors = Object.keys(this.state.skippedKeys);
+    if (skippedAuthors.length) {
+      this.nostrSubscribe(
+        {authors: skippedAuthors, kinds: [MESSAGE_EVENT_KIND]},
+        (e) => this.handleNostrEvent(e)
+      );  
     }
-    // do we want this unsubscribed on rotation or should we keep it open
-    // in case more skipped messages are found by relays or peers?
-    this.nostrUnsubscribe = this.nostrSubscribe(
-      {authors, kinds: [MESSAGE_EVENT_KIND]},
-      (e) => this.handleNostrEvent(e)
-    );
   }
 }

package/src/types.ts

@@ -48,11 +48,13 @@ export interface SessionState {
   /** Number of messages sent in previous sending chain */
   previousSendingChainMessageCount: number;
   
-  /** Cache of message keys for handling out-of-order messages */
-  skippedMessageKeys: Record<string, Uint8Array>;
-
-  /** Cache of header keys for handling out-of-order messages */
-  skippedHeaderKeys: Record<string, Uint8Array[]>;
+  /** Cache of message & header keys for handling out-of-order messages */
+  skippedKeys: {
+    [pubKey: string]: {
+      headerKeys: Uint8Array[],
+      messageKeys: {[msgIndex: number]: Uint8Array}
+    };
+  };
 }
 
 /**

package/src/utils.ts

@@ -4,8 +4,11 @@ import { Session } from "./Session.ts";
 import { extract as hkdf_extract, expand as hkdf_expand } from '@noble/hashes/hkdf';
 import { sha256 } from '@noble/hashes/sha256';
 
+const VERSION_NUMBER = 1;
+
 export function serializeSessionState(state: SessionState): string {
   return JSON.stringify({
+    version: VERSION_NUMBER,
     rootKey: bytesToHex(state.rootKey),
     theirCurrentNostrPublicKey: state.theirCurrentNostrPublicKey,
     theirNextNostrPublicKey: state.theirNextNostrPublicKey,
@@ -22,16 +25,18 @@ export function serializeSessionState(state: SessionState): string {
     sendingChainMessageNumber: state.sendingChainMessageNumber,
     receivingChainMessageNumber: state.receivingChainMessageNumber,
     previousSendingChainMessageCount: state.previousSendingChainMessageCount,
-    skippedMessageKeys: Object.fromEntries(
-      Object.entries(state.skippedMessageKeys).map(([key, value]) => [
-        key,
-        bytesToHex(value),
-      ])
-    ),
-    skippedHeaderKeys: Object.fromEntries(
-      Object.entries(state.skippedHeaderKeys).map(([key, value]) => [
-        key,
-        value.map(bytes => bytesToHex(bytes))
+    skippedKeys: Object.fromEntries(
+      Object.entries(state.skippedKeys).map(([pubKey, value]) => [
+        pubKey,
+        {
+          headerKeys: value.headerKeys.map(bytes => bytesToHex(bytes)),
+          messageKeys: Object.fromEntries(
+            Object.entries(value.messageKeys).map(([msgIndex, bytes]) => [
+              msgIndex,
+              bytesToHex(bytes)
+            ])
+          )
+        }
       ])
     ),
   });
@@ -39,6 +44,43 @@ export function serializeSessionState(state: SessionState): string {
 
 export function deserializeSessionState(data: string): SessionState {
   const state = JSON.parse(data);
+  
+  // Handle version 0 (legacy format)
+  if (!state.version) {
+    const skippedKeys: SessionState['skippedKeys'] = {};
+    
+    // Migrate old skipped keys format to new structure
+    if (state.skippedMessageKeys) {
+      Object.entries(state.skippedMessageKeys).forEach(([pubKey, messageKeys]: [string, any]) => {
+        skippedKeys[pubKey] = {
+          headerKeys: state.skippedHeaderKeys?.[pubKey] || [],
+          messageKeys: messageKeys
+        };
+      });
+    }
+    
+    return {
+      rootKey: hexToBytes(state.rootKey),
+      theirCurrentNostrPublicKey: state.theirCurrentNostrPublicKey,
+      theirNextNostrPublicKey: state.theirNextNostrPublicKey,
+      ourCurrentNostrKey: state.ourCurrentNostrKey ? {
+        publicKey: state.ourCurrentNostrKey.publicKey,
+        privateKey: hexToBytes(state.ourCurrentNostrKey.privateKey),
+      } : undefined,
+      ourNextNostrKey: {
+        publicKey: state.ourNextNostrKey.publicKey,
+        privateKey: hexToBytes(state.ourNextNostrKey.privateKey),
+      },
+      receivingChainKey: state.receivingChainKey ? hexToBytes(state.receivingChainKey) : undefined,
+      sendingChainKey: state.sendingChainKey ? hexToBytes(state.sendingChainKey) : undefined,
+      sendingChainMessageNumber: state.sendingChainMessageNumber,
+      receivingChainMessageNumber: state.receivingChainMessageNumber,
+      previousSendingChainMessageCount: state.previousSendingChainMessageCount,
+      skippedKeys
+    };
+  }
+  
+  // Handle current version
   return {
     rootKey: hexToBytes(state.rootKey),
     theirCurrentNostrPublicKey: state.theirCurrentNostrPublicKey,
@@ -56,16 +98,18 @@ export function deserializeSessionState(data: string): SessionState {
     sendingChainMessageNumber: state.sendingChainMessageNumber,
     receivingChainMessageNumber: state.receivingChainMessageNumber,
     previousSendingChainMessageCount: state.previousSendingChainMessageCount,
-    skippedMessageKeys: Object.fromEntries(
-      Object.entries(state.skippedMessageKeys).map(([key, value]) => [
-        key,
-        hexToBytes(value as string),
-      ])
-    ),
-    skippedHeaderKeys: Object.fromEntries(
-      Object.entries(state.skippedHeaderKeys || {}).map(([key, value]) => [
-        key,
-        (value as string[]).map(hex => hexToBytes(hex))
+    skippedKeys: Object.fromEntries(
+      Object.entries(state.skippedKeys || {}).map(([pubKey, value]) => [
+        pubKey,
+        {
+          headerKeys: (value as any).headerKeys.map((hex: string) => hexToBytes(hex)),
+          messageKeys: Object.fromEntries(
+            Object.entries((value as any).messageKeys).map(([msgIndex, hex]) => [
+              msgIndex,
+              hexToBytes(hex as string)
+            ])
+          )
+        }
       ])
     ),
   };