tests(sighash): add sighash bitmask tests [part 2/7]

tests/tx/scripts/test_p2pkh.py

@@ -0,0 +1,105 @@
+#  Copyright 2023 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from hathor.transaction.scripts import P2PKH, Opcode
+from hathor.transaction.scripts.sighash import InputsOutputsLimit, SighashBitmask
+
+
+def test_create_input_data_simple() -> None:
+    pub_key = b'my_pub_key'
+    signature = b'my_signature'
+    data = P2PKH.create_input_data(public_key_bytes=pub_key, signature=signature)
+
+    assert data == bytes([
+        len(signature),
+        *signature,
+        len(pub_key),
+        *pub_key
+    ])
+
+
+def test_create_input_data_with_sighash_bitmask() -> None:
+    pub_key = b'my_pub_key'
+    signature = b'my_signature'
+    inputs_bitmask = 0b111
+    outputs_bitmask = 0b101
+    sighash = SighashBitmask(inputs=inputs_bitmask, outputs=outputs_bitmask)
+    data = P2PKH.create_input_data(public_key_bytes=pub_key, signature=signature, sighash=sighash)
+
+    assert data == bytes([
+        1,
+        inputs_bitmask,
+        1,
+        outputs_bitmask,
+        Opcode.OP_SIGHASH_BITMASK,
+        len(signature),
+        *signature,
+        len(pub_key),
+        *pub_key
+    ])
+
+
+def test_create_input_data_with_inputs_outputs_limit() -> None:
+    pub_key = b'my_pub_key'
+    signature = b'my_signature'
+    max_inputs = 2
+    max_outputs = 3
+    limit = InputsOutputsLimit(max_inputs=max_inputs, max_outputs=max_outputs)
+    data = P2PKH.create_input_data(public_key_bytes=pub_key, signature=signature, inputs_outputs_limit=limit)
+
+    assert data == bytes([
+        1,
+        max_inputs,
+        1,
+        max_outputs,
+        Opcode.OP_MAX_INPUTS_OUTPUTS,
+        len(signature),
+        *signature,
+        len(pub_key),
+        *pub_key
+    ])
+
+
+def test_create_input_data_with_sighash_bitmask_and_inputs_outputs_limit() -> None:
+    pub_key = b'my_pub_key'
+    signature = b'my_signature'
+    inputs_bitmask = 0b111
+    outputs_bitmask = 0b101
+    max_inputs = 2
+    max_outputs = 3
+    sighash = SighashBitmask(inputs=inputs_bitmask, outputs=outputs_bitmask)
+    limit = InputsOutputsLimit(max_inputs=max_inputs, max_outputs=max_outputs)
+    data = P2PKH.create_input_data(
+        public_key_bytes=pub_key,
+        signature=signature,
+        sighash=sighash,
+        inputs_outputs_limit=limit
+    )
+
+    assert data == bytes([
+        1,
+        inputs_bitmask,
+        1,
+        outputs_bitmask,
+        Opcode.OP_SIGHASH_BITMASK,
+        1,
+        max_inputs,
+        1,
+        max_outputs,
+        Opcode.OP_MAX_INPUTS_OUTPUTS,
+        len(signature),
+        *signature,
+        len(pub_key),
+        *pub_key
+    ])

tests/tx/scripts/test_script_context.py

@@ -0,0 +1,97 @@
+#  Copyright 2023 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import hashlib
+from unittest.mock import Mock
+
+import pytest
+
+from hathor.conf.settings import HathorSettings
+from hathor.transaction import Transaction, TxInput, TxOutput
+from hathor.transaction.exceptions import ScriptError
+from hathor.transaction.scripts.script_context import ScriptContext
+from hathor.transaction.scripts.sighash import SighashAll, SighashBitmask
+
+
+@pytest.mark.parametrize(['max_num_outputs'], [(99,), (255,)])
+def test_defaults(max_num_outputs: int) -> None:
+    settings = Mock(spec_set=HathorSettings)
+    settings.MAX_NUM_OUTPUTS = max_num_outputs
+
+    context = ScriptContext(settings=settings, stack=Mock(), logs=[], extras=Mock())
+
+    tx = Transaction(
+        inputs=[
+            TxInput(tx_id=b'tx1', index=0, data=b''),
+            TxInput(tx_id=b'tx2', index=1, data=b''),
+        ],
+        outputs=[
+            TxOutput(value=11, script=b''),
+            TxOutput(value=22, script=b''),
+        ]
+    )
+
+    assert context._sighash == SighashAll()
+    assert context.get_tx_sighash_data(tx) == tx.get_sighash_all_data()
+    assert context.get_selected_outputs() == set(range(max_num_outputs))
+
+
+def test_set_sighash() -> None:
+    context = ScriptContext(settings=Mock(), stack=Mock(), logs=[], extras=Mock())
+
+    sighash = SighashBitmask(inputs=0b111, outputs=0b101)
+    context.set_sighash(sighash)
+    assert context._sighash == sighash
+
+    with pytest.raises(ScriptError):
+        context.set_sighash(sighash)
+
+
+@pytest.mark.parametrize(
+    ['outputs_bitmask', 'selected_outputs'],
+    [
+        (0b00, set()),
+        (0b01, {0}),
+        (0b10, {1}),
+        (0b11, {0, 1}),
+    ]
+)
+def test_sighash_bitmask(outputs_bitmask: int, selected_outputs: set[int]) -> None:
+    settings = Mock()
+    settings.MAX_NUM_INPUTS = 88
+    settings.MAX_NUM_OUTPUTS = 99
+
+    context = ScriptContext(settings=settings, stack=Mock(), logs=[], extras=Mock())
+    tx = Transaction(
+        inputs=[
+            TxInput(tx_id=b'tx1', index=0, data=b''),
+            TxInput(tx_id=b'tx2', index=1, data=b''),
+        ],
+        outputs=[
+            TxOutput(value=11, script=b''),
+            TxOutput(value=22, script=b''),
+        ]
+    )
+
+    sighash_bitmask = SighashBitmask(inputs=0b11, outputs=outputs_bitmask)
+    context.set_sighash(sighash_bitmask)
+
+    data = tx.get_custom_sighash_data(sighash_bitmask)
+    assert context.get_tx_sighash_data(tx) == hashlib.sha256(data).digest()
+
+    with pytest.raises(ScriptError) as e:
+        context.set_sighash(Mock())
+
+    assert str(e.value) == 'Cannot modify sighash after it is already set.'
+    assert context.get_selected_outputs() == selected_outputs