From 8cfdb3ed6f7a631d31302966724af7ed35e08782 Mon Sep 17 00:00:00 2001
From: Tomek Piotrowski <tomek.piotrowski@atteo.com>
Date: Tue, 6 Aug 2024 09:11:38 +0200
Subject: [PATCH] Add NSM transaction test vectors

---
 pyproject.toml                    |   1 +
 zcash_test_vectors/transaction.py |  26 ++++-
 zcash_test_vectors/zip_nsm.py     | 172 ++++++++++++++++++++++++++++++
 3 files changed, 196 insertions(+), 3 deletions(-)
 create mode 100755 zcash_test_vectors/zip_nsm.py

diff --git a/pyproject.toml b/pyproject.toml
index a0f9359..868e970 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -44,6 +44,7 @@ unified_incoming_viewing_keys = "zcash_test_vectors.unified_incoming_viewing_key
 zip_0143 = "zcash_test_vectors.zip_0143:main"
 zip_0243 = "zcash_test_vectors.zip_0243:main"
 zip_0244 = "zcash_test_vectors.zip_0244:main"
+zip_nsm = "zcash_test_vectors.zip_nsm:main"
 
 # Transparent test vectors
 bip_0032 = "zcash_test_vectors.transparent.bip_0032:main"
diff --git a/zcash_test_vectors/transaction.py b/zcash_test_vectors/transaction.py
index 01c2d06..801b8da 100644
--- a/zcash_test_vectors/transaction.py
+++ b/zcash_test_vectors/transaction.py
@@ -26,6 +26,9 @@
 NU5_VERSION_GROUP_ID = 0x26A7270A
 NU5_TX_VERSION = 5
 
+ZFUTURE_VERSION_GROUP_ID = 0xFFFFFFFF
+ZFUTURE_TX_VERSION = 0x0000FFFF
+
 # Sapling note magic values, copied from src/zcash/Zcash.h
 NOTEENCRYPTION_AUTH_BYTES = 16
 ZC_NOTEPLAINTEXT_LEADING = 1
@@ -413,7 +416,7 @@ def __bytes__(self):
 
 
 class TransactionV5(object):
-    def __init__(self, rand, consensus_branch_id):
+    def __init__(self, rand, consensus_branch_id, version_group_id):
         # Decide which transaction parts will be generated.
         flip_coins = rand.u8()
         have_transparent_in = (flip_coins >> 0) % 2
@@ -423,7 +426,7 @@ def __init__(self, rand, consensus_branch_id):
         is_coinbase = (not have_transparent_in) and (flip_coins >> 4) % 2
 
         # Common Transaction Fields
-        self.nVersionGroupId = NU5_VERSION_GROUP_ID
+        self.nVersionGroupId = version_group_id
         self.nConsensusBranchId = consensus_branch_id
         self.nLockTime = rand.u32()
         self.nExpiryHeight = rand.u32() % TX_EXPIRY_HEIGHT_THRESHOLD
@@ -549,12 +552,29 @@ def __bytes__(self):
 
         return ret
 
+class TransactionNSM(TransactionV5):
+    def __init__(self, rand, consensus_branch_id, version_group_id):
+        super().__init__(rand, consensus_branch_id, version_group_id)
+
+        self.burnAmount = rand.u64() % (MAX_MONEY + 1)
+
+    def version_bytes(self):
+        return ZFUTURE_TX_VERSION | (1 << 31)
+
+    def __bytes__(self):
+        ret = super().__bytes__()
+        ret += struct.pack('<Q', self.burnAmount)
+        return ret
+
 
 class Transaction(object):
     def __init__(self, rand, version, consensus_branch_id=None):
         if version == NU5_TX_VERSION:
             assert consensus_branch_id is not None
-            self.inner = TransactionV5(rand, consensus_branch_id)
+            self.inner = TransactionV5(rand, consensus_branch_id, NU5_VERSION_GROUP_ID)
+        elif version == ZFUTURE_TX_VERSION:
+            assert consensus_branch_id is not None
+            self.inner = TransactionNSM(rand, consensus_branch_id, ZFUTURE_VERSION_GROUP_ID)
         else:
             self.inner = LegacyTransaction(rand, version)
 
diff --git a/zcash_test_vectors/zip_nsm.py b/zcash_test_vectors/zip_nsm.py
new file mode 100755
index 0000000..cbf4a49
--- /dev/null
+++ b/zcash_test_vectors/zip_nsm.py
@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+import sys; assert sys.version_info[0] >= 3, "Python 3 required."
+
+from hashlib import blake2b
+import struct
+
+from .transaction import (
+    ZFUTURE_VERSION_GROUP_ID,
+    TransactionNSM,
+)
+from .output import render_args, render_tv, Some
+from .rand import Rand
+from .zip_0143 import (
+    SIGHASH_ALL,
+    SIGHASH_ANYONECANPAY,
+    SIGHASH_NONE,
+    SIGHASH_SINGLE,
+)
+
+from .zip_0244 import *
+
+def auth_digest(tx):
+    digest = blake2b(
+        digest_size=32,
+        person=b'ZTxAuthHash_' + struct.pack('<I', tx.nConsensusBranchId),
+    )
+
+    digest.update(transparent_scripts_digest(tx))
+    digest.update(sapling_auth_digest(tx))
+    digest.update(orchard_auth_digest(tx))
+    digest.update(burn_amount_digest(tx))
+
+    return digest.digest()
+
+# NSM
+def burn_amount_digest(tx):
+    digest = blake2b(digest_size=32, person=b'ZTxBurnAmnt_Hash')
+
+    digest.update(struct.pack('<Q', tx.burnAmount))
+
+    return digest.digest()
+
+def txid_digest(tx):
+    digest = blake2b(
+        digest_size=32,
+        person=b'ZcashTxHash_' + struct.pack('<I', tx.nConsensusBranchId),
+    )
+
+    digest.update(header_digest(tx))
+    digest.update(transparent_digest(tx))
+    digest.update(sapling_digest(tx))
+    digest.update(orchard_digest(tx))
+    digest.update(burn_amount_digest(tx))
+
+    return digest.digest()
+
+def signature_digest(tx, t_inputs, nHashType, txin):
+    digest = blake2b(
+        digest_size=32,
+        person=b'ZcashTxHash_' + struct.pack('<I', tx.nConsensusBranchId),
+    )
+
+    digest.update(header_digest(tx))
+    digest.update(transparent_sig_digest(tx, t_inputs, nHashType, txin))
+    digest.update(sapling_digest(tx))
+    digest.update(orchard_digest(tx))
+    digest.update(burn_amount_digest(tx))
+
+    return digest.digest()
+
+def txin_sig_digest(tx, txin):
+    digest = blake2b(digest_size=32, person=b'Zcash___TxInHash')
+    if txin is not None:
+        digest.update(bytes(tx.vin[txin.nIn].prevout))
+        digest.update(struct.pack('<Q', txin.amount))
+        digest.update(bytes(txin.scriptPubKey))
+        digest.update(struct.pack('<I', tx.vin[txin.nIn].nSequence))
+    return digest.digest()
+
+
+def main():
+    args = render_args()
+
+    from random import Random
+    rng = Random(0xB7D6_0F44)
+    def randbytes(l):
+        ret = []
+        while len(ret) < l:
+            ret.append(rng.randrange(0, 256))
+        return bytes(ret)
+    rand = Rand(randbytes)
+
+    consensusBranchId = 0xFFFF_FFFF # ZFUTURE
+
+    test_vectors = []
+    for _ in range(10):
+        tx = TransactionNSM(rand, consensusBranchId, ZFUTURE_VERSION_GROUP_ID)
+        txid = txid_digest(tx)
+        auth = auth_digest(tx)
+
+        # Generate amounts and scriptCodes for each non-dummy transparent input.
+        if tx.is_coinbase():
+            t_inputs = []
+        else:
+            t_inputs = [TransparentInput(nIn, rand) for nIn in range(len(tx.vin))]
+
+        # If there are any non-dummy transparent inputs, derive a corresponding transparent sighash.
+        if len(t_inputs) > 0:
+            txin = rand.a(t_inputs)
+        else:
+            txin = None
+
+        sighash_shielded = signature_digest(tx, t_inputs, SIGHASH_ALL, None)
+        other_sighashes = {
+            nHashType: None if txin is None else signature_digest(tx, t_inputs, nHashType, txin)
+            for nHashType in ([
+                SIGHASH_ALL,
+                SIGHASH_NONE,
+                SIGHASH_SINGLE,
+                SIGHASH_ALL | SIGHASH_ANYONECANPAY,
+                SIGHASH_NONE | SIGHASH_ANYONECANPAY,
+                SIGHASH_SINGLE | SIGHASH_ANYONECANPAY,
+            ] if txin is None or txin.nIn < len(tx.vout) else [
+                SIGHASH_ALL,
+                SIGHASH_NONE,
+                SIGHASH_ALL | SIGHASH_ANYONECANPAY,
+                SIGHASH_NONE | SIGHASH_ANYONECANPAY,
+            ])
+        }
+
+        test_vectors.append({
+            'tx': bytes(tx),
+            'txid': txid,
+            'auth_digest': auth,
+            'amounts': [x.amount for x in t_inputs],
+            'burn_amount': tx.burnAmount,
+            'script_pubkeys': [x.scriptPubKey.raw() for x in t_inputs],
+            'transparent_input': None if txin is None else txin.nIn,
+            'sighash_shielded': sighash_shielded,
+            'sighash_all': other_sighashes.get(SIGHASH_ALL),
+            'sighash_none': other_sighashes.get(SIGHASH_NONE),
+            'sighash_single': other_sighashes.get(SIGHASH_SINGLE),
+            'sighash_all_anyone': other_sighashes.get(SIGHASH_ALL | SIGHASH_ANYONECANPAY),
+            'sighash_none_anyone': other_sighashes.get(SIGHASH_NONE | SIGHASH_ANYONECANPAY),
+            'sighash_single_anyone': other_sighashes.get(SIGHASH_SINGLE | SIGHASH_ANYONECANPAY),
+        })
+
+    render_tv(
+        args,
+        'zip_nsm',
+        (
+            ('tx',                    {'rust_type': 'Vec<u8>', 'bitcoin_flavoured': False}),
+            ('txid',                  '[u8; 32]'),
+            ('auth_digest',           '[u8; 32]'),
+            ('amounts',               'Vec<i64>'),
+            ('burn_amount',           'u64'),
+            ('script_pubkeys',        {'rust_type': 'Vec<Vec<u8>>', 'bitcoin_flavoured': False}),
+            ('transparent_input',     'Option<u32>'),
+            ('sighash_shielded',      '[u8; 32]'),
+            ('sighash_all',           'Option<[u8; 32]>'),
+            ('sighash_none',          'Option<[u8; 32]>'),
+            ('sighash_single',        'Option<[u8; 32]>'),
+            ('sighash_all_anyone',    'Option<[u8; 32]>'),
+            ('sighash_none_anyone',   'Option<[u8; 32]>'),
+            ('sighash_single_anyone', 'Option<[u8; 32]>'),
+        ),
+        test_vectors,
+    )
+
+
+if __name__ == '__main__':
+    main()