Skip to content

Commit 19d2e0a

Browse files
achow101achow101
committed
psbt: Implement Taproot fields
1 parent fe85a31 commit 19d2e0a

File tree

1 file changed

+167
-0
lines changed

1 file changed

+167
-0
lines changed

‎hwilib/psbt.py

+167
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
Optional,
1616
Sequence,
1717
Set,
18+
Tuple,
1819
)
1920

2021
from .key import KeyOriginInfo
@@ -25,6 +26,7 @@
2526
CTxOut,
2627
)
2728
from ._serialize import (
29+
deser_compact_size,
2830
deser_string,
2931
Readable,
3032
ser_compact_size,
@@ -86,6 +88,12 @@ class PartiallySignedInput:
8688
PSBT_IN_BIP32_DERIVATION = 0x06
8789
PSBT_IN_FINAL_SCRIPTSIG = 0x07
8890
PSBT_IN_FINAL_SCRIPTWITNESS = 0x08
91+
PSBT_IN_TAP_KEY_SIG = 0x13
92+
PSBT_IN_TAP_SCRIPT_SIG = 0x14
93+
PSBT_IN_TAP_LEAF_SCRIPT = 0x15
94+
PSBT_IN_TAP_BIP32_DERIVATION = 0x16
95+
PSBT_IN_TAP_INTERNAL_KEY = 0x17
96+
PSBT_IN_TAP_MERKLE_ROOT = 0x18
8997

9098
def __init__(self) -> None:
9199
self.non_witness_utxo: Optional[CTransaction] = None
@@ -97,6 +105,12 @@ def __init__(self) -> None:
97105
self.hd_keypaths: Dict[bytes, KeyOriginInfo] = {}
98106
self.final_script_sig = b""
99107
self.final_script_witness = CTxInWitness()
108+
self.tap_key_sig = b""
109+
self.tap_script_sigs: Dict[Tuple[bytes, bytes], bytes] = {}
110+
self.tap_scripts: Dict[Tuple[bytes, int], Set[bytes]] = {}
111+
self.tap_bip32_paths: Dict[bytes, Tuple[Set[bytes], KeyOriginInfo]] = {}
112+
self.tap_internal_key = b""
113+
self.tap_merkle_root = b""
100114
self.unknown: Dict[bytes, bytes] = {}
101115

102116
def set_null(self) -> None:
@@ -112,6 +126,12 @@ def set_null(self) -> None:
112126
self.hd_keypaths.clear()
113127
self.final_script_sig = b""
114128
self.final_script_witness = CTxInWitness()
129+
self.tap_key_sig = b""
130+
self.tap_script_sigs.clear()
131+
self.tap_scripts.clear()
132+
self.tap_bip32_paths.clear()
133+
self.tap_internal_key = b""
134+
self.tap_merkle_root = b""
115135
self.unknown.clear()
116136

117137
def deserialize(self, f: Readable) -> None:
@@ -196,6 +216,72 @@ def deserialize(self, f: Readable) -> None:
196216
raise PSBTSerializationError("final scriptWitness key is more than one byte type")
197217
witness_bytes = BufferedReader(BytesIO(deser_string(f))) # type: ignore
198218
self.final_script_witness.deserialize(witness_bytes)
219+
elif key_type == PartiallySignedInput.PSBT_IN_TAP_KEY_SIG:
220+
if key in key_lookup:
221+
raise PSBTSerializationError("Duplicate key, input Taproot key signature already provided")
222+
elif len(key) != 1:
223+
raise PSBTSerializationError("Input Taproot key signature key is more than one byte type")
224+
self.tap_key_sig = deser_string(f)
225+
if len(self.tap_key_sig) < 64:
226+
raise PSBTSerializationError("Input Taproot key path signature is shorter than 64 bytes")
227+
elif len(self.tap_key_sig) > 65:
228+
raise PSBTSerializationError("Input Taproot key path signature is longer than 65 bytes")
229+
elif key_type == PartiallySignedInput.PSBT_IN_TAP_SCRIPT_SIG:
230+
if key in key_lookup:
231+
raise PSBTSerializationError("Duplicate key, input Taproot script signature already provided")
232+
elif len(key) != 65:
233+
raise PSBTSerializationError("Input Taproot script signature key is not 65 bytes")
234+
xonly = key[1:33]
235+
script_hash = key[33:65]
236+
sig = deser_string(f)
237+
if len(sig) < 64:
238+
raise PSBTSerializationError("Input Taproot script path signature is shorter than 64 bytes")
239+
elif len(sig) > 65:
240+
raise PSBTSerializationError("Input Taproot script path signature is longer than 65 bytes")
241+
self.tap_script_sigs[(xonly, script_hash)] = sig
242+
elif key_type == PartiallySignedInput.PSBT_IN_TAP_LEAF_SCRIPT:
243+
if key in key_lookup:
244+
raise PSBTSerializationError("Duplicate key, input Taproot leaf script already provided")
245+
elif len(key) < 34:
246+
raise PSBTSerializationError("Input Taproot leaf script key is not at least 34 bytes")
247+
elif (len(key) - 2) % 32 != 0:
248+
raise PSBTSerializationError("Input Taproot leaf script key's control block is not valid")
249+
script = deser_string(f)
250+
if len(script) == 0:
251+
raise PSBTSerializationError("Intput Taproot leaf script cannot be empty")
252+
leaf_script = (script[:-1], int(script[-1]))
253+
if leaf_script not in self.tap_scripts:
254+
self.tap_scripts[leaf_script] = set()
255+
self.tap_scripts[(script[:-1], int(script[-1]))].add(key[1:])
256+
elif key_type == PartiallySignedInput.PSBT_IN_TAP_BIP32_DERIVATION:
257+
if key in key_lookup:
258+
raise PSBTSerializationError("Duplicate key, input Taproot BIP 32 keypath already provided")
259+
elif len(key) != 33:
260+
raise PSBTSerializationError("Input Taproot BIP 32 keypath key is not 33 bytes")
261+
xonly = key[1:33]
262+
value = deser_string(f)
263+
vs = BytesIO(value)
264+
num_hashes = deser_compact_size(vs)
265+
leaf_hashes = set()
266+
for i in range(0, num_hashes):
267+
leaf_hashes.add(vs.read(32))
268+
self.tap_bip32_paths[xonly] = (leaf_hashes, KeyOriginInfo.deserialize(vs.read()))
269+
elif key_type == PartiallySignedInput.PSBT_IN_TAP_INTERNAL_KEY:
270+
if key in key_lookup:
271+
raise PSBTSerializationError("Duplicate key, input Taproot internal key already provided")
272+
elif len(key) != 1:
273+
raise PSBTSerializationError("Input Taproot internal key key is more than one byte type")
274+
self.tap_internal_key = deser_string(f)
275+
if len(self.tap_internal_key) != 32:
276+
raise PSBTSerializationError("Input Taproot internal key is not 32 bytes")
277+
elif key_type == PartiallySignedInput.PSBT_IN_TAP_MERKLE_ROOT:
278+
if key in key_lookup:
279+
raise PSBTSerializationError("Duplicate key, input Taproot merkle root already provided")
280+
elif len(key) != 1:
281+
raise PSBTSerializationError("Input Taproot merkle root key is more than one byte type")
282+
self.tap_merkle_root = deser_string(f)
283+
if len(self.tap_merkle_root) != 32:
284+
raise PSBTSerializationError("Input Taproot merkle root is not 32 bytes")
199285
else:
200286
if key in self.unknown:
201287
raise PSBTSerializationError("Duplicate key, key for unknown value already provided")
@@ -241,6 +327,35 @@ def serialize(self) -> bytes:
241327

242328
r += SerializeHDKeypath(self.hd_keypaths, ser_compact_size(PartiallySignedInput.PSBT_IN_BIP32_DERIVATION))
243329

330+
if len(self.tap_key_sig) != 0:
331+
r += ser_string(ser_compact_size(PartiallySignedInput.PSBT_IN_TAP_KEY_SIG))
332+
r += ser_string(self.tap_key_sig)
333+
334+
for (xonly, leaf_hash), sig in self.tap_script_sigs.items():
335+
r += ser_string(ser_compact_size(PartiallySignedInput.PSBT_IN_TAP_SCRIPT_SIG) + xonly + leaf_hash)
336+
r += ser_string(sig)
337+
338+
for (script, leaf_ver), control_blocks in self.tap_scripts.items():
339+
for control_block in control_blocks:
340+
r += ser_string(ser_compact_size(PartiallySignedInput.PSBT_IN_TAP_LEAF_SCRIPT) + control_block)
341+
r += ser_string(script + struct.pack("B", leaf_ver))
342+
343+
for xonly, (leaf_hashes, origin) in self.tap_bip32_paths.items():
344+
r += ser_string(ser_compact_size(PartiallySignedInput.PSBT_IN_TAP_BIP32_DERIVATION) + xonly)
345+
value = ser_compact_size(len(leaf_hashes))
346+
for lh in leaf_hashes:
347+
value += lh
348+
value += origin.serialize()
349+
r += ser_string(value)
350+
351+
if len(self.tap_internal_key) != 0:
352+
r += ser_string(ser_compact_size(PartiallySignedInput.PSBT_IN_TAP_INTERNAL_KEY))
353+
r += ser_string(self.tap_internal_key)
354+
355+
if len(self.tap_merkle_root) != 0:
356+
r += ser_string(ser_compact_size(PartiallySignedInput.PSBT_IN_TAP_MERKLE_ROOT))
357+
r += ser_string(self.tap_merkle_root)
358+
244359
if len(self.final_script_sig) != 0:
245360
r += ser_string(ser_compact_size(PartiallySignedInput.PSBT_IN_FINAL_SCRIPTSIG))
246361
r += ser_string(self.final_script_sig)
@@ -266,11 +381,17 @@ class PartiallySignedOutput:
266381
PSBT_OUT_REDEEM_SCRIPT = 0x00
267382
PSBT_OUT_WITNESS_SCRIPT = 0x01
268383
PSBT_OUT_BIP32_DERIVATION = 0x02
384+
PSBT_OUT_TAP_INTERNAL_KEY = 0x05
385+
PSBT_OUT_TAP_TREE = 0x06
386+
PSBT_OUT_TAP_BIP32_DERIVATION = 0x07
269387

270388
def __init__(self) -> None:
271389
self.redeem_script = b""
272390
self.witness_script = b""
273391
self.hd_keypaths: Dict[bytes, KeyOriginInfo] = {}
392+
self.tap_internal_key = b""
393+
self.tap_tree = b""
394+
self.tap_bip32_paths: Dict[bytes, Tuple[Set[bytes], KeyOriginInfo]] = {}
274395
self.unknown: Dict[bytes, bytes] = {}
275396

276397
def set_null(self) -> None:
@@ -280,6 +401,9 @@ def set_null(self) -> None:
280401
self.redeem_script = b""
281402
self.witness_script = b""
282403
self.hd_keypaths.clear()
404+
self.tap_internal_key = b""
405+
self.tap_tree = b""
406+
self.tap_bip32_paths.clear()
283407
self.unknown.clear()
284408

285409
def deserialize(self, f: Readable) -> None:
@@ -318,6 +442,33 @@ def deserialize(self, f: Readable) -> None:
318442
self.witness_script = deser_string(f)
319443
elif key_type == PartiallySignedOutput.PSBT_OUT_BIP32_DERIVATION:
320444
DeserializeHDKeypath(f, key, self.hd_keypaths, [34, 66])
445+
elif key_type == PartiallySignedOutput.PSBT_OUT_TAP_INTERNAL_KEY:
446+
if key in key_lookup:
447+
raise PSBTSerializationError("Duplicate key, output Taproot internal key already provided")
448+
elif len(key) != 1:
449+
raise PSBTSerializationError("Output Taproot internal key key is more than one byte type")
450+
self.tap_internal_key = deser_string(f)
451+
if len(self.tap_internal_key) != 32:
452+
raise PSBTSerializationError("Output Taproot internal key is not 32 bytes")
453+
elif key_type == PartiallySignedOutput.PSBT_OUT_TAP_TREE:
454+
if key in key_lookup:
455+
raise PSBTSerializationError("Duplicate key, output Taproot tree already provided")
456+
elif len(key) != 1:
457+
raise PSBTSerializationError("Output Taproot tree key is more than one byte type")
458+
self.tap_tree = deser_string(f)
459+
elif key_type == PartiallySignedOutput.PSBT_OUT_TAP_BIP32_DERIVATION:
460+
if key in key_lookup:
461+
raise PSBTSerializationError("Duplicate key, output Taproot BIP 32 keypath already provided")
462+
elif len(key) != 33:
463+
raise PSBTSerializationError("Output Taproot BIP 32 keypath key is not 33 bytes")
464+
xonly = key[1:33]
465+
value = deser_string(f)
466+
vs = BytesIO(value)
467+
num_hashes = deser_compact_size(vs)
468+
leaf_hashes = set()
469+
for i in range(0, num_hashes):
470+
leaf_hashes.add(vs.read(32))
471+
self.tap_bip32_paths[xonly] = (leaf_hashes, KeyOriginInfo.deserialize(vs.read()))
321472
else:
322473
if key in self.unknown:
323474
raise PSBTSerializationError("Duplicate key, key for unknown value already provided")
@@ -343,6 +494,22 @@ def serialize(self) -> bytes:
343494

344495
r += SerializeHDKeypath(self.hd_keypaths, ser_compact_size(PartiallySignedOutput.PSBT_OUT_BIP32_DERIVATION))
345496

497+
if len(self.tap_internal_key) != 0:
498+
r += ser_string(ser_compact_size(PartiallySignedOutput.PSBT_OUT_TAP_INTERNAL_KEY))
499+
r += ser_string(self.tap_internal_key)
500+
501+
if len(self.tap_tree) != 0:
502+
r += ser_string(ser_compact_size(PartiallySignedOutput.PSBT_OUT_TAP_TREE))
503+
r += ser_string(self.tap_tree)
504+
505+
for xonly, (leaf_hashes, origin) in self.tap_bip32_paths.items():
506+
r += ser_string(ser_compact_size(PartiallySignedOutput.PSBT_OUT_TAP_BIP32_DERIVATION) + xonly)
507+
value = ser_compact_size(len(leaf_hashes))
508+
for lh in leaf_hashes:
509+
value += lh
510+
value += origin.serialize()
511+
r += ser_string(value)
512+
346513
for key, value in sorted(self.unknown.items()):
347514
r += ser_string(key)
348515
r += ser_string(value)

0 commit comments

Comments
 (0)