From 6943f02848544d02d366c44a578bee6a1d8b559c Mon Sep 17 00:00:00 2001
From: michael1011 <me@michael1011.at>
Date: Thu, 7 Dec 2023 20:35:09 +0100
Subject: [PATCH] fix tweaking by returning tweaked key cache

---
 src/lib/interface.ts   |  5 +++-
 src/lib/musig.ts       | 16 +++++++++---
 src/test/musig.spec.ts | 55 ++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 70 insertions(+), 6 deletions(-)

diff --git a/src/lib/interface.ts b/src/lib/interface.ts
index a1214b2..ce00255 100644
--- a/src/lib/interface.ts
+++ b/src/lib/interface.ts
@@ -164,7 +164,10 @@ export interface Musig {
     keyaggCache: Uint8Array,
     tweak: Uint8Array,
     compress?: boolean
-  ): Uint8Array;
+  ): {
+    pubkey: Uint8Array;
+    keyaggCache: Uint8Array;
+  };
 }
 
 export interface Secp256k1ZKP {
diff --git a/src/lib/musig.ts b/src/lib/musig.ts
index 33a6e8a..b813104 100644
--- a/src/lib/musig.ts
+++ b/src/lib/musig.ts
@@ -303,6 +303,8 @@ function pubkeyXonlyTweakAdd(
     const outputLen = memory.malloc(8);
     cModule.setValue(outputLen, 65, 'i64');
 
+    const keyaggCacheTweaked = memory.charStar(keyaggCache);
+
     const ret = cModule.ccall(
       'musig_pubkey_xonly_tweak_add',
       'number',
@@ -311,7 +313,7 @@ function pubkeyXonlyTweakAdd(
         output,
         outputLen,
         compress ? 1 : 0,
-        memory.charStar(keyaggCache),
+        keyaggCacheTweaked,
         memory.charStar(tweak),
       ]
     );
@@ -321,12 +323,20 @@ function pubkeyXonlyTweakAdd(
       throw new Error('musig_pubkey_xonly_tweak_add');
     }
 
-    const res = memory.charStarToUint8(
+    const pubkey = memory.charStarToUint8(
       output,
       cModule.getValue(outputLen, 'i64')
     );
+    const keyaggCacheTweakedRes = memory.charStarToUint8(
+      keyaggCacheTweaked,
+      165
+    );
+
     memory.free();
-    return res;
+    return {
+      pubkey,
+      keyaggCache: keyaggCacheTweakedRes,
+    };
   };
 }
 
diff --git a/src/test/musig.spec.ts b/src/test/musig.spec.ts
index 959bc0d..ea894fb 100644
--- a/src/test/musig.spec.ts
+++ b/src/test/musig.spec.ts
@@ -136,8 +136,8 @@ test('pubkeyXonlyTweakAdd', (t) => {
       f.compress
     );
 
-    t.is(tweaked.length, f.tweakedLength);
-    t.is(uintToString(tweaked), f.tweaked);
+    t.is(tweaked.pubkey.length, f.tweakedLength);
+    t.is(uintToString(tweaked.pubkey), f.tweaked);
   });
 });
 
@@ -186,3 +186,54 @@ test('full example', (t) => {
   const sig = musig.partialSigAgg(session, partialSigs);
   t.true(musig.ecc.verifySchnorr(message, pubkeyAgg.aggPubkey, sig));
 });
+
+test('full example tweaked', (t) => {
+  const musig = t.context;
+
+  const privateKeys = fixtures.fullExample.privateKeys.map((key) =>
+    fromHex(key)
+  );
+  const publicKeys = privateKeys.map((key) =>
+    musig.ec.fromPrivateKey(key).publicKey.subarray(1)
+  );
+  t.is(publicKeys.length, privateKeys.length);
+
+  const pubkeyAgg = musig.pubkeyAgg(publicKeys);
+  const tweak = musig.pubkeyXonlyTweakAdd(
+    pubkeyAgg.keyaggCache,
+    randomBytes(32),
+    true
+  );
+
+  const nonces = publicKeys.map(() => musig.nonceGen(randomBytes(32)));
+  const nonceAgg = musig.nonceAgg(nonces.map((nonce) => nonce.pubNonce));
+
+  const message = randomBytes(32);
+  const session = musig.nonceProcess(nonceAgg, message, tweak.keyaggCache);
+
+  const partialSigs = privateKeys.map((privateKey, i) =>
+    musig.partialSign(
+      nonces[i].secNonce,
+      privateKey,
+      tweak.keyaggCache,
+      session
+    )
+  );
+
+  // Verify each partial signature individually, to make sure they are fine on their own
+  partialSigs.forEach((sig, i) =>
+    t.true(
+      musig.partialVerify(
+        sig,
+        nonces[i].pubNonce,
+        publicKeys[i],
+        tweak.keyaggCache,
+        session
+      )
+    )
+  );
+
+  // Combine the partial signatures into one and verify it
+  const sig = musig.partialSigAgg(session, partialSigs);
+  t.true(musig.ecc.verifySchnorr(message, tweak.pubkey.slice(1), sig));
+});