1 files changed, 28 insertions, 2 deletions

diff --git a/js/axon.js b/js/axon.js
index 9f50340..761910c 100644
--- a/js/axon.js
+++ b/js/axon.js
@@ -8,6 +8,7 @@
 
 import { ed25519, x25519, edwardsToMontgomeryPriv, edwardsToMontgomeryPub } from '@noble/curves/ed25519';
 import { sha256 } from '@noble/hashes/sha256';
+import { hkdf } from '@noble/hashes/hkdf';
 import { chacha20poly1305 } from '@noble/ciphers/chacha';
 import { encode, decode } from '@msgpack/msgpack';
 
@@ -287,6 +288,25 @@ function dhSharedSecret(localSeed, remoteEd25519Pub) {
 }
 
 /**
+ * dmKey derives a 32-byte symmetric key from the raw DH shared secret
+ * via HKDF-SHA256, matching Go's dmKey().
+ */
+function dmKey(shared) {
+  return hkdf(sha256, shared, undefined, 'axon-dm-v1', 32);
+}
+
+/**
+ * dmAD builds associated data by concatenating sender + recipient Ed25519
+ * public keys, matching Go's dmAD().
+ */
+function dmAD(senderPub, recipientPub) {
+  const ad = new Uint8Array(senderPub.length + recipientPub.length);
+  ad.set(senderPub);
+  ad.set(recipientPub, senderPub.length);
+  return ad;
+}
+
+/**
  * encryptDM(senderSeed, recipientPubkey, plaintext) → Uint8Array
  *
  * Returns nonce(12) || ciphertext+tag, matching Go's EncryptDM.
@@ -296,8 +316,11 @@ export function encryptDM(senderSeed, recipientPubkey, plaintext) {
   const enc = new TextEncoder();
   const pt = typeof plaintext === 'string' ? enc.encode(plaintext) : plaintext;
   const shared = dhSharedSecret(senderSeed, recipientPubkey);
+  const key = dmKey(shared);
   const nonce = crypto.getRandomValues(new Uint8Array(12));
-  const ct = chacha20poly1305(shared, nonce).encrypt(pt);
+  const senderPub = ed25519.getPublicKey(senderSeed);
+  const ad = dmAD(senderPub, recipientPubkey);
+  const ct = chacha20poly1305(key, nonce, ad).encrypt(pt);
   const out = new Uint8Array(12 + ct.length);
   out.set(nonce);
   out.set(ct, 12);
@@ -312,9 +335,12 @@ export function encryptDM(senderSeed, recipientPubkey, plaintext) {
 export function decryptDM(recipientSeed, senderPubkey, blob) {
   if (blob.length < 12) throw new Error('axon: dm blob too short');
   const shared = dhSharedSecret(recipientSeed, senderPubkey);
+  const key = dmKey(shared);
   const nonce = blob.slice(0, 12);
   const ct    = blob.slice(12);
-  return chacha20poly1305(shared, nonce).decrypt(ct);
+  const recipientPub = ed25519.getPublicKey(recipientSeed);
+  const ad = dmAD(senderPubkey, recipientPub);
+  return chacha20poly1305(key, nonce, ad).decrypt(ct);
 }
 
 // ---------------------------------------------------------------------------