From fb5a8a5c5fbfd141f5c119a283610c744865f28a Mon Sep 17 00:00:00 2001
From: Ouziel Slama <ouziel@gmail.com>
Date: Mon, 9 Dec 2024 12:14:17 +0000
Subject: [PATCH] Improve DER and schnorr signature detection

---
 .../counterpartycore/lib/gettxinfo.py         | 75 ++++++++++++++-----
 .../fixtures/contract_vectors/gettxinfo.py    | 12 +--
 2 files changed, 61 insertions(+), 26 deletions(-)

diff --git a/counterparty-core/counterpartycore/lib/gettxinfo.py b/counterparty-core/counterpartycore/lib/gettxinfo.py
index e0590fa62..87544521a 100644
--- a/counterparty-core/counterpartycore/lib/gettxinfo.py
+++ b/counterparty-core/counterpartycore/lib/gettxinfo.py
@@ -1,6 +1,7 @@
 import binascii
 import logging
 import struct
+from io import BytesIO
 
 from counterpartycore.lib import arc4, backend, config, ledger, message_type, script, util
 from counterpartycore.lib.exceptions import BTCOnlyError, DecodeError
@@ -156,31 +157,65 @@ def get_vin_info(vin):
     return vout["value"], vout["script_pub_key"], is_segwit
 
 
+def is_valid_der(der):
+    if not isinstance(der, bytes):
+        return False
+    try:
+        s = BytesIO(der)
+        compound = s.read(1)[0]
+        if compound != 0x30:
+            return False
+        length = s.read(1)[0]
+        if length + 2 != len(der):
+            return False
+        marker = s.read(1)[0]
+        if marker != 0x02:
+            return False
+        rlength = s.read(1)[0]
+        _r = int(s.read(rlength).hex(), 16)
+        marker = s.read(1)[0]
+        if marker != 0x02:
+            return False
+        slength = s.read(1)[0]
+        s = int(s.read(slength).hex(), 16)
+        if len(der) != 6 + rlength + slength:
+            return False
+        return True
+    except Exception:
+        return False
+
+
+def is_valid_schnorr(schnorr):
+    p = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F
+    n = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
+
+    if not isinstance(schnorr, bytes):
+        return False
+    if len(schnorr) not in [64, 65]:
+        return False
+    if len(schnorr) == 65:
+        schnorr = schnorr[:-1]
+    try:
+        r = int.from_bytes(schnorr[0:32], byteorder="big")
+        s = int.from_bytes(schnorr[32:64], byteorder="big")
+    except Exception:
+        return False
+    if (r >= p) or (s >= n):
+        return False
+    return True
+
+
 def get_der_signature_sighash_flag(value):
-    if not isinstance(value, bytes):
-        return None
-    lenght_by_prefix = {
-        "3044": 71,
-        "3045": 72,
-        "3046": 73,
-        "3041": 68,
-        "3042": 69,
-        "3043": 70,
-    }
-    for prefix, length in lenght_by_prefix.items():
-        if value.startswith(binascii.unhexlify(prefix)) and len(value) == length:
-            return value[-1:]
+    if is_valid_der(value[:-1]):
+        return value[-1:]
     return None
 
 
 def get_schnorr_signature_sighash_flag(value):
-    if not isinstance(value, bytes):
-        return None
-    if len(value) not in [64, 65]:
-        return None
-    if len(value) == 65:
-        return value[-1:]
-    return b"\x01"  # SIGHASH_ALL by default
+    if is_valid_schnorr(value):
+        if len(value) == 65:
+            return value[-1:]
+        return b"\x01"  # SIGHASH_ALL
 
 
 def collect_sighash_flags(script_sig, witnesses):
diff --git a/counterparty-core/counterpartycore/test/fixtures/contract_vectors/gettxinfo.py b/counterparty-core/counterpartycore/test/fixtures/contract_vectors/gettxinfo.py
index 0508752ed..aa9cbf65a 100644
--- a/counterparty-core/counterpartycore/test/fixtures/contract_vectors/gettxinfo.py
+++ b/counterparty-core/counterpartycore/test/fixtures/contract_vectors/gettxinfo.py
@@ -370,19 +370,19 @@
         ],
         "get_der_signature_sighash_flag": [
             {
-                "comment": "69 length",
+                "comment": "68 length",
                 "in": (
                     binascii.unhexlify(
-                        "30420229a522e65ca8500ebe05a70d5a49d840ccc15f2afa4ee9df783f06b2a322310220489a46c37feb33f52c586da25c70113b8eea41216440eb84771cb67a67fdb68c01"
+                        "3041021d067f716293ffdaa144d5281185f5f9ca35f09e24f10af0ba7aa8e9ae8f0220555f644564e472d835d707200e03ae7c3728d5e2fee00eb7c351d025d8c6e16701"
                     ),
                 ),
                 "out": b"\x01",
             },
             {
-                "comment": "70 length",
+                "comment": "69 length",
                 "in": (
                     binascii.unhexlify(
-                        "304302c219a522e65ca8500ebe05a70d5a49d840ccc15f2afa4ee9df783f06b2a322310220489a46c37feb33f52c586da25c70113b8eea41216440eb84771cb67a67fdb68c01"
+                        "3042021e78565f3de8b3b526ee9c3a07b396f4e18e082f77175f4aa444b63f66c9ad02205813ea177ac1d9b2dbb43f63cd3a00d47a3a525f3030c5233b8240239f9c15e401"
                     ),
                 ),
                 "out": b"\x01",
@@ -391,7 +391,7 @@
                 "comment": "71 length",
                 "in": (
                     binascii.unhexlify(
-                        "30442100c219a522e65ca8500ebe05a70d5a49d840ccc15f2afa4ee9df783f06b2a322310220489a46c37feb33f52c586da25c70113b8eea41216440eb84771cb67a67fdb68c01"
+                        "3044022063c96d6644f7d325bc7fed3362fd6cc81d81bf4a4af8df8d5f13147d6c74267a02201340b3b01b1f29d2d0e180abf5e3d14cc832b412cb27a5a68135f98493bb006e01"
                     ),
                 ),
                 "out": b"\x01",
@@ -409,7 +409,7 @@
                 "comment": "73 length",
                 "in": (
                     binascii.unhexlify(
-                        "304600022100c219a522e65ca8500ebe05a70d5a49d840ccc15f2afa4ee9df783f06b2a322310220489a46c37feb33f52c586da25c70113b8eea41216440eb84771cb67a67fdb68c01"
+                        "3046022100d8697a511eea7c0949f4295dc185d3e19cbd80aed547a2c0d29dd635b04430ef022100b334f949949ad19a3e78c9c3eb82320d877a5acc094e7ed3b15987a45ffca89101"
                     ),
                 ),
                 "out": b"\x01",