feat(sighash): implement sighash range

glevco · glevco · commit d7aa643f4fd3 · 2025-02-25T18:57:08.000-03:00
diff --git a/hathor/transaction/scripts/opcode.py b/hathor/transaction/scripts/opcode.py
@@ -45,7 +45,7 @@
 )
 from hathor.transaction.scripts.execute import Stack, binary_to_int, decode_opn, get_data_value, get_script_op
 from hathor.transaction.scripts.script_context import ScriptContext
-from hathor.transaction.scripts.sighash import InputsOutputsLimit, SighashBitmask
+from hathor.transaction.scripts.sighash import InputsOutputsLimit, SighashBitmask, SighashRange
 from hathor.transaction.util import bytes_to_int
 
 
@@ -655,6 +655,38 @@ def op_sighash_bitmask(context: ScriptContext) -> None:
     context.set_sighash(sighash)
 
 
+def op_sighash_range(context: ScriptContext) -> None:
+    """Pop four items from the stack, constructing a sighash range and setting it in the script context."""
+    if len(context.stack) < 4:
+        raise MissingStackItems(f'OP_SIGHASH_RANGE: expected 4 elements on stack, has {len(context.stack)}')
+
+    output_end = context.stack.pop()
+    output_start = context.stack.pop()
+    input_end = context.stack.pop()
+    input_start = context.stack.pop()
+    assert isinstance(output_end, bytes)
+    assert isinstance(output_start, bytes)
+    assert isinstance(input_end, bytes)
+    assert isinstance(input_start, bytes)
+
+    try:
+        sighash = SighashRange(
+            input_start=bytes_to_int(input_start),
+            input_end=bytes_to_int(input_end),
+            output_start=bytes_to_int(output_start),
+            output_end=bytes_to_int(output_end),
+        )
+    except Exception as e:
+        raise CustomSighashModelInvalid('Could not construct sighash range.') from e
+
+    if context.extras.input_index not in sighash.get_input_indexes():
+        raise InputNotSelectedError(
+            f'Input at index {context.extras.input_index} must select itself when using a custom sighash.'
+        )
+
+    context.set_sighash(sighash)
+
+
 def op_max_inputs_outputs(context: ScriptContext) -> None:
     """Pop two items from the stack, constructing an inputs and outputs limit and setting it in the script context."""
     if len(context.stack) < 2:
@@ -707,6 +739,7 @@ def execute_op_code(opcode: Opcode, context: ScriptContext) -> None:
         case Opcode.OP_CHECKDATASIG: op_checkdatasig(context)
         case Opcode.OP_FIND_P2PKH: op_find_p2pkh(context)
         case Opcode.OP_SIGHASH_BITMASK: op_sighash_bitmask(context)
+        case Opcode.OP_SIGHASH_RANGE: op_sighash_range(context)
         case Opcode.OP_MAX_INPUTS_OUTPUTS: op_max_inputs_outputs(context)
         case _: raise ScriptError(f'unknown opcode: {opcode}')
 
diff --git a/hathor/transaction/scripts/p2pkh.py b/hathor/transaction/scripts/p2pkh.py
@@ -15,12 +15,20 @@
 import struct
 from typing import Any, Optional
 
+from typing_extensions import assert_never
+
 from hathor.crypto.util import decode_address, get_address_b58_from_public_key_hash
 from hathor.transaction.scripts.base_script import BaseScript
 from hathor.transaction.scripts.construct import get_pushdata, re_compile
 from hathor.transaction.scripts.hathor_script import HathorScript
 from hathor.transaction.scripts.opcode import Opcode
-from hathor.transaction.scripts.sighash import InputsOutputsLimit, SighashBitmask
+from hathor.transaction.scripts.sighash import (
+    InputsOutputsLimit,
+    SighashAll,
+    SighashBitmask,
+    SighashRange,
+    SighashType,
+)
 
 
 class P2PKH(BaseScript):
@@ -97,7 +105,7 @@ def create_input_data(
         public_key_bytes: bytes,
         signature: bytes,
         *,
-        sighash: SighashBitmask | None = None,
+        sighash: SighashType = SighashAll(),
         inputs_outputs_limit: InputsOutputsLimit | None = None
     ) -> bytes:
         """
@@ -108,10 +116,21 @@ def create_input_data(
         """
         s = HathorScript()
 
-        if sighash:
-            s.pushData(sighash.inputs)
-            s.pushData(sighash.outputs)
-            s.addOpcode(Opcode.OP_SIGHASH_BITMASK)
+        match sighash:
+            case SighashAll():
+                pass
+            case SighashBitmask():
+                s.pushData(sighash.inputs)
+                s.pushData(sighash.outputs)
+                s.addOpcode(Opcode.OP_SIGHASH_BITMASK)
+            case SighashRange():
+                s.pushData(sighash.input_start)
+                s.pushData(sighash.input_end)
+                s.pushData(sighash.output_start)
+                s.pushData(sighash.output_end)
+                s.addOpcode(Opcode.OP_SIGHASH_RANGE)
+            case _:
+                assert_never(sighash)
 
         if inputs_outputs_limit:
             s.pushData(inputs_outputs_limit.max_inputs)
diff --git a/hathor/transaction/scripts/script_context.py b/hathor/transaction/scripts/script_context.py
@@ -20,7 +20,7 @@
 from hathor.transaction import Transaction
 from hathor.transaction.exceptions import ScriptError
 from hathor.transaction.scripts.execute import ScriptExtras, Stack
-from hathor.transaction.scripts.sighash import SighashAll, SighashBitmask, SighashType
+from hathor.transaction.scripts.sighash import SighashAll, SighashBitmask, SighashRange, SighashType
 
 
 class ScriptContext:
@@ -52,7 +52,7 @@ def get_tx_sighash_data(self, tx: Transaction) -> bytes:
         match self._sighash:
             case SighashAll():
                 return tx.get_sighash_all_data()
-            case SighashBitmask():
+            case SighashBitmask() | SighashRange():
                 data = tx.get_custom_sighash_data(self._sighash)
                 return hashlib.sha256(data).digest()
             case _:
@@ -63,7 +63,7 @@ def get_selected_outputs(self) -> set[int]:
         match self._sighash:
             case SighashAll():
                 return set(range(self._settings.MAX_NUM_OUTPUTS))
-            case SighashBitmask():
+            case SighashBitmask() | SighashRange():
                 return set(self._sighash.get_output_indexes())
             case _:
                 assert_never(self._sighash)
diff --git a/hathor/transaction/scripts/sighash.py b/hathor/transaction/scripts/sighash.py
@@ -14,9 +14,9 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import TypeAlias
+from typing import Any, TypeAlias
 
-from pydantic import Field
+from pydantic import Field, validator
 from typing_extensions import override
 
 from hathor.utils.pydantic import BaseModel
@@ -60,7 +60,37 @@ def _get_indexes(bitmask: int) -> list[int]:
         return [index for index in range(8) if (bitmask >> index) & 1]
 
 
-SighashType: TypeAlias = SighashAll | SighashBitmask
+class SighashRange(CustomSighash):
+    """A model representing the sighash range type config. Range ends are not inclusive."""
+    input_start: int = Field(ge=0, le=255)
+    input_end: int = Field(ge=0, le=255)
+    output_start: int = Field(ge=0, le=255)
+    output_end: int = Field(ge=0, le=255)
+
+    @validator('input_end')
+    def _validate_input_end(cls, input_end: int, values: dict[str, Any]) -> int:
+        if input_end < values['input_start']:
+            raise ValueError('input_end must be greater than or equal to input_start.')
+
+        return input_end
+
+    @validator('output_end')
+    def _validate_output_end(cls, output_end: int, values: dict[str, Any]) -> int:
+        if output_end < values['output_start']:
+            raise ValueError('output_end must be greater than or equal to output_start.')
+
+        return output_end
+
+    @override
+    def get_input_indexes(self) -> list[int]:
+        return list(range(self.input_start, self.input_end))
+
+    @override
+    def get_output_indexes(self) -> list[int]:
+        return list(range(self.output_start, self.output_end))
+
+
+SighashType: TypeAlias = SighashAll | SighashBitmask | SighashRange
 
 
 class InputsOutputsLimit(BaseModel):
diff --git a/tests/tx/scripts/test_p2pkh.py b/tests/tx/scripts/test_p2pkh.py
@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from hathor.transaction.scripts import P2PKH, Opcode
-from hathor.transaction.scripts.sighash import InputsOutputsLimit, SighashBitmask
+from hathor.transaction.scripts.sighash import InputsOutputsLimit, SighashAll, SighashBitmask, SighashRange
 
 
 def test_create_input_data_simple() -> None:
@@ -29,6 +29,19 @@ def test_create_input_data_simple() -> None:
     ])
 
 
+def test_create_input_data_with_sighash_all() -> None:
+    pub_key = b'my_pub_key'
+    signature = b'my_signature'
+    data = P2PKH.create_input_data(public_key_bytes=pub_key, signature=signature, sighash=SighashAll())
+
+    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'
@@ -50,6 +63,38 @@ def test_create_input_data_with_sighash_bitmask() -> None:
     ])
 
 
+def test_create_input_data_with_sighash_range() -> None:
+    pub_key = b'my_pub_key'
+    signature = b'my_signature'
+    input_start = 123
+    input_end = 145
+    output_start = 10
+    output_end = 20
+    sighash = SighashRange(
+        input_start=input_start,
+        input_end=input_end,
+        output_start=output_start,
+        output_end=output_end,
+    )
+    data = P2PKH.create_input_data(public_key_bytes=pub_key, signature=signature, sighash=sighash)
+
+    assert data == bytes([
+        1,
+        input_start,
+        1,
+        input_end,
+        1,
+        output_start,
+        1,
+        output_end,
+        Opcode.OP_SIGHASH_RANGE,
+        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'
@@ -103,3 +148,48 @@ def test_create_input_data_with_sighash_bitmask_and_inputs_outputs_limit() -> No
         len(pub_key),
         *pub_key
     ])
+
+
+def test_create_input_data_with_sighash_range_and_inputs_outputs_limit() -> None:
+    pub_key = b'my_pub_key'
+    signature = b'my_signature'
+    input_start = 123
+    input_end = 145
+    output_start = 10
+    output_end = 20
+    max_inputs = 2
+    max_outputs = 3
+    sighash = SighashRange(
+        input_start=input_start,
+        input_end=input_end,
+        output_start=output_start,
+        output_end=output_end,
+    )
+    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,
+        input_start,
+        1,
+        input_end,
+        1,
+        output_start,
+        1,
+        output_end,
+        Opcode.OP_SIGHASH_RANGE,
+        1,
+        max_inputs,
+        1,
+        max_outputs,
+        Opcode.OP_MAX_INPUTS_OUTPUTS,
+        len(signature),
+        *signature,
+        len(pub_key),
+        *pub_key
+    ])
diff --git a/tests/tx/scripts/test_script_context.py b/tests/tx/scripts/test_script_context.py
@@ -21,7 +21,7 @@
 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
+from hathor.transaction.scripts.sighash import SighashAll, SighashBitmask, SighashRange
 
 
 @pytest.mark.parametrize(['max_num_outputs'], [(99,), (255,)])
@@ -95,3 +95,42 @@ def test_sighash_bitmask(outputs_bitmask: int, selected_outputs: set[int]) -> No
 
     assert str(e.value) == 'Cannot modify sighash after it is already set.'
     assert context.get_selected_outputs() == selected_outputs
+
+
+@pytest.mark.parametrize(
+    ['output_start', 'output_end', 'selected_outputs'],
+    [
+        (100, 100, set()),
+        (0, 1, {0}),
+        (1, 2, {1}),
+        (0, 2, {0, 1}),
+    ]
+)
+def test_sighash_range(output_start: int, output_end: 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_range = SighashRange(input_start=0, input_end=2, output_start=output_start, output_end=output_end)
+    context.set_sighash(sighash_range)
+
+    data = tx.get_custom_sighash_data(sighash_range)
+    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
diff --git a/tests/tx/scripts/test_sighash_bitmask.py b/tests/tx/scripts/test_sighash_bitmask.py
@@ -30,7 +30,7 @@
 from tests.utils import add_blocks_unlock_reward, create_tokens, get_genesis_key
 
 
-class SighashTest(unittest.TestCase):
+class SighashBitmaskTest(unittest.TestCase):
     def setUp(self) -> None:
         super().setUp()
         self.manager1: HathorManager = self.create_peer('testnet', unlock_wallet=True, wallet_index=True)
diff --git a/tests/tx/scripts/test_sighash_range.py b/tests/tx/scripts/test_sighash_range.py
diff --git a/tests/tx/scripts/test_tx_sighash.py b/tests/tx/scripts/test_tx_sighash.py
diff --git a/tests/tx/test_scripts.py b/tests/tx/test_scripts.py
