From 600c28887a53d66b2835fb412f585a938bbd588e Mon Sep 17 00:00:00 2001
From: callebtc <93376500+callebtc@users.noreply.github.com>
Date: Fri, 22 Sep 2023 16:31:26 +0200
Subject: [PATCH] add refund signature checks

---
 cashu/mint/conditions.py  | 10 +++---
 tests/test_wallet_htlc.py | 72 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 6 deletions(-)

diff --git a/cashu/mint/conditions.py b/cashu/mint/conditions.py
index bdcb99e8..830489c5 100644
--- a/cashu/mint/conditions.py
+++ b/cashu/mint/conditions.py
@@ -149,19 +149,17 @@ def _verify_input_spending_conditions(self, proof: Proof) -> bool:
             if htlc_secret.locktime and htlc_secret.locktime < time.time():
                 refund_pubkeys = htlc_secret.tags.get_tag_all("refund")
                 if refund_pubkeys:
-                    assert proof.htlcpreimage, "no HTLC preimage provided"
+                    assert proof.htlcsignature, TransactionError(
+                        "no HTLC refund signature provided"
+                    )
                     for pubkey in refund_pubkeys:
                         if verify_p2pk_signature(
                             message=htlc_secret.serialize().encode("utf-8"),
                             pubkey=PublicKey(bytes.fromhex(pubkey), raw=True),
-                            signature=bytes.fromhex(proof.htlcpreimage),
+                            signature=bytes.fromhex(proof.htlcsignature),
                         ):
                             # a signature matches
                             return True
-                    if refund_pubkeys:
-                        # none of the pubkeys had a match
-                        return False
-                    # none of the signatures match
                     raise TransactionError("HTLC refund signatures did not match.")
                 # no pubkeys given in secret, anyone can spend
                 return True
diff --git a/tests/test_wallet_htlc.py b/tests/test_wallet_htlc.py
index efc57b84..3303e099 100644
--- a/tests/test_wallet_htlc.py
+++ b/tests/test_wallet_htlc.py
@@ -1,3 +1,4 @@
+import asyncio
 import hashlib
 import secrets
 from typing import List
@@ -164,3 +165,74 @@ async def test_htlc_redeem_with_correct_signature(wallet1: Wallet, wallet2: Wall
         p.htlcsignature = s
 
     await wallet2.redeem(send_proofs)
+
+
+@pytest.mark.asyncio
+async def test_htlc_redeem_hashlock_wrong_signature_timelock_correct_signature(
+    wallet1: Wallet, wallet2: Wallet
+):
+    await wallet1.mint(64)
+    preimage = "00000000000000000000000000000000"
+    pubkey_wallet1 = await wallet1.create_p2pk_pubkey()
+    pubkey_wallet2 = await wallet2.create_p2pk_pubkey()
+    # preimage_hash = hashlib.sha256(bytes.fromhex(preimage)).hexdigest()
+    secret = await wallet1.create_htlc_lock(
+        preimage=preimage,
+        hacklock_pubkey=pubkey_wallet2,
+        locktime_seconds=5,
+        locktime_pubkey=pubkey_wallet1,
+    )
+    # p2pk test
+    _, send_proofs = await wallet1.split_to_send(wallet1.proofs, 8, secret_lock=secret)
+
+    signatures = await wallet1.sign_p2pk_proofs(send_proofs)
+    for p, s in zip(send_proofs, signatures):
+        p.htlcpreimage = preimage
+        p.htlcsignature = s
+
+    # should error because we used wallet2 signatures for the hash lock
+    await assert_err(
+        wallet1.redeem(send_proofs),
+        "Mint Error: HTLC hash lock signatures did not match.",
+    )
+
+    await asyncio.sleep(5)
+    # should succeed since lock time has passed and we provided wallet1 signature for timelock
+    await wallet1.redeem(send_proofs)
+
+
+@pytest.mark.asyncio
+async def test_htlc_redeem_hashlock_wrong_signature_timelock_wrong_signature(
+    wallet1: Wallet, wallet2: Wallet
+):
+    await wallet1.mint(64)
+    preimage = "00000000000000000000000000000000"
+    pubkey_wallet1 = await wallet1.create_p2pk_pubkey()
+    pubkey_wallet2 = await wallet2.create_p2pk_pubkey()
+    # preimage_hash = hashlib.sha256(bytes.fromhex(preimage)).hexdigest()
+    secret = await wallet1.create_htlc_lock(
+        preimage=preimage,
+        hacklock_pubkey=pubkey_wallet2,
+        locktime_seconds=5,
+        locktime_pubkey=pubkey_wallet1,
+    )
+    # p2pk test
+    _, send_proofs = await wallet1.split_to_send(wallet1.proofs, 8, secret_lock=secret)
+
+    signatures = await wallet1.sign_p2pk_proofs(send_proofs)
+    for p, s in zip(send_proofs, signatures):
+        p.htlcpreimage = preimage
+        p.htlcsignature = s[:-1] + "1"  # wrong signature
+
+    # should error because we used wallet2 signatures for the hash lock
+    await assert_err(
+        wallet1.redeem(send_proofs),
+        "Mint Error: HTLC hash lock signatures did not match.",
+    )
+
+    await asyncio.sleep(5)
+    # should fail since lock time has passed and we provided a wrong signature for timelock
+    await assert_err(
+        wallet1.redeem(send_proofs),
+        "Mint Error: HTLC refund signatures did not match.",
+    )