Add P2TR detection in wally_scriptpubkey_get_type()

Author: jgriffiths

include/wally_script.h

@@ -16,12 +16,14 @@ extern "C" {
 #define WALLY_SCRIPT_TYPE_P2WPKH    0x8
 #define WALLY_SCRIPT_TYPE_P2WSH     0x10
 #define WALLY_SCRIPT_TYPE_MULTISIG  0x20
+#define WALLY_SCRIPT_TYPE_P2TR      0x40
 
 /* Standard script lengths */
 #define WALLY_SCRIPTPUBKEY_P2PKH_LEN  25 /** OP_DUP OP_HASH160 [HASH160] OP_EQUALVERIFY OP_CHECKSIG */
 #define WALLY_SCRIPTPUBKEY_P2SH_LEN   23 /** OP_HASH160 [HASH160] OP_EQUAL */
 #define WALLY_SCRIPTPUBKEY_P2WPKH_LEN 22 /** OP_0 [HASH160] */
 #define WALLY_SCRIPTPUBKEY_P2WSH_LEN  34 /** OP_0 [SHA256] */
+#define WALLY_SCRIPTPUBKEY_P2TR_LEN   34 /** OP_1 [X-ONLY-PUBKEY] */
 
 #define WALLY_SCRIPTPUBKEY_OP_RETURN_MAX_LEN 83 /** OP_RETURN [80 bytes of data] */
 

src/script.c

@@ -312,30 +312,38 @@ static bool scriptpubkey_is_p2pkh(const unsigned char *bytes, size_t bytes_len)
 {
     return bytes_len == WALLY_SCRIPTPUBKEY_P2PKH_LEN &&
            bytes[0] == OP_DUP && bytes[1] == OP_HASH160 &&
-           bytes[2] == 20 && bytes[23] == OP_EQUALVERIFY &&
+           bytes[2] == 20 && /* HASH160 */
+           bytes[23] == OP_EQUALVERIFY &&
            bytes[24] == OP_CHECKSIG;
 }
 
 static bool scriptpubkey_is_p2sh(const unsigned char *bytes, size_t bytes_len)
 {
     return bytes_len == WALLY_SCRIPTPUBKEY_P2SH_LEN &&
            bytes[0] == OP_HASH160 &&
-           bytes[1] == 20 &&
+           bytes[1] == 20 && /* HASH160 */
            bytes[22] == OP_EQUAL;
 }
 
 static bool scriptpubkey_is_p2wpkh(const unsigned char *bytes, size_t bytes_len)
 {
     return bytes_len == WALLY_SCRIPTPUBKEY_P2WPKH_LEN &&
-           bytes[0] == 0 &&
-           bytes[1] == 20;
+           bytes[0] == OP_0 && /* Segwit v0 */
+           bytes[1] == 20; /* HASH160 */
 }
 
 static bool scriptpubkey_is_p2wsh(const unsigned char *bytes, size_t bytes_len)
 {
     return bytes_len == WALLY_SCRIPTPUBKEY_P2WSH_LEN &&
-           bytes[0] == 0 &&
-           bytes[1] == 32;
+           bytes[0] == OP_0 && /* Segwit v0 */
+           bytes[1] == 32; /* SHA256 */
+}
+
+static bool scriptpubkey_is_p2tr(const unsigned char *bytes, size_t bytes_len)
+{
+    return bytes_len == WALLY_SCRIPTPUBKEY_P2TR_LEN &&
+           bytes[0] == OP_1 && /* Segwit v1 */
+           bytes[1] == 32; /* X-ONLY-PUBKEY */
 }
 
 static bool scriptpubkey_is_multisig(const unsigned char *bytes, size_t bytes_len)
@@ -400,10 +408,13 @@ int wally_scriptpubkey_get_type(const unsigned char *bytes, size_t bytes_len,
             return WALLY_OK;
         }
         break;
-    case WALLY_SCRIPTPUBKEY_P2WSH_LEN:
+    case WALLY_SCRIPTPUBKEY_P2WSH_LEN: /* Also WALLY_SCRIPTPUBKEY_P2TR_LEN */
         if (scriptpubkey_is_p2wsh(bytes, bytes_len)) {
             *written = WALLY_SCRIPT_TYPE_P2WSH;
             return WALLY_OK;
+        } else if (scriptpubkey_is_p2tr(bytes, bytes_len)) {
+            *written = WALLY_SCRIPT_TYPE_P2TR;
+            return WALLY_OK;
         }
         break;
     }

src/test/test_script.py

@@ -1,10 +1,14 @@
 import unittest
 from util import *
 
+SCRIPT_TYPE_UNKNOWN = 0x0
 SCRIPT_TYPE_OP_RETURN = 0x1
 SCRIPT_TYPE_P2PKH = 0x2
 SCRIPT_TYPE_P2SH = 0x4
+SCRIPT_TYPE_P2WPKH = 0x8
+SCRIPT_TYPE_P2WSH = 0x10
 SCRIPT_TYPE_MULTISIG = 0x20
+SCRIPT_TYPE_P2TR = 0x40
 
 SCRIPT_MULTISIG_SORTED = 0x8
 
@@ -45,9 +49,20 @@ def test_scriptpubkey_get_type(self):
         # Test invalid args, we test results with the functions that make scripts
         in_, in_len = make_cbuffer('00' * 16)
         for b, b_len in [(None, in_len), (in_, 0)]:
-            ret, written = wally_scriptpubkey_get_type(b, b_len)
+            ret, typ = wally_scriptpubkey_get_type(b, b_len)
             self.assertEqual(ret, WALLY_EINVAL)
-            self.assertEqual(written, 0)
+            self.assertEqual(typ, SCRIPT_TYPE_UNKNOWN)
+
+        # Segwit scripts
+        segwit_cases = [
+            [ '0014' + ('33' * 20), SCRIPT_TYPE_P2WPKH ],
+            [ '0020' + ('33' * 32), SCRIPT_TYPE_P2WSH ],
+            [ '5120' + ('33' * 32), SCRIPT_TYPE_P2TR ],
+        ]
+        for script_hex, expected_type in segwit_cases:
+            script, script_len = make_cbuffer(script_hex)
+            ret, typ = wally_scriptpubkey_get_type(script, script_len)
+            self.assertEqual((ret, typ), (WALLY_OK, expected_type))
 
     def test_scriptpubkey_op_return_from_bytes(self):
         """Tests for creating OP_RETURN scriptPubKeys"""