Modification of OP_CHECKMULTISIG to remove some FPs and FNs

symbolic_execution/hardware/machines/btc_machine.py

@@ -96,6 +96,7 @@ def int_to_bytes(val: int) -> str:
 
 # Set of functions specific to BTC
 CHECK_SINGLE_SIG = Function("check_sig", IntSort(), IntSort(), BoolSort())
+CHECK_MULTI_SIG = Function("check_multi_sig", IntSort(), IntSort(), BoolSort())
 HASH_FUNCTIONS = {
     "ripemd160": [Function("ripemd160", IntSort(), IntSort()), RIPEMD160.new, 20],
     "sha1": [Function("sha1", IntSort(), IntSort()), hashlib.sha1, 20],
@@ -317,24 +318,39 @@ class BTCVector:
             If(ret == BoolVal(False), IntVal(0), IntVal(1)),
         )
 
-    def check_multisig(pubkeys_cnt, sigs_cnt):
+    def check_multisig(pubkeys_cnt, sigs_cnt, pubkeys, sigs):
         """Check a list of signatures against a list of public keys."""
-        # TODO Could be refined. We could compute the signature in case both signature
-        # and key are constants.
+        if pubkeys_cnt.val == 0:
+            return BTCVector(BTCVector.BTC_TRUE, IntVal(1))
+
+        pubkeys_unset = 0
+        for pubkey in pubkeys:
+            const_values = pubkey.fetch_const_values()
+            pubkeys_unset = (
+                pubkeys_unset + 1 if len(const_values) == 0 else pubkeys_unset
+            )
 
-        # If the number of keys is not hardcoded, it is possible to alter it and pass
-        # more keys
-        if not is_int_value(pubkeys_cnt):
+        unsafe = False
+
+        # If the number of public keys is not hardcoded, the attacker could forge it
+        # and pass controlled keys
+        if len(pubkeys_cnt.fetch_const_values()) == 0:
+            unsafe = True
+        #  If the number of required signatures is not set or lower than the number of
+        # hardcoded public keys, the attacked could sign it correctly.
+        if len(sigs_cnt.fetch_const_values()) == 0 or (
+            is_int_value(sigs_cnt.val) and sigs_cnt.val.as_long() <= pubkeys_unset
+        ):
+            unsafe = True
+
+        if unsafe:
+            ret = CHECK_MULTI_SIG(pubkeys_cnt.val, sigs_cnt.val)
             return BTCVector(
-                BTCVector.BTC_TRUE,
-                IntVal(1)
+                If(ret == BoolVal(False), BTCVector.BTC_FALSE, BTCVector.BTC_TRUE),
+                If(ret == BoolVal(False), IntVal(0), IntVal(1)),
             )
 
-        return BTCVector(
-            BTCVector.BTC_FALSE,
-            IntVal(0)
-        )
-
+        return BTCVector(BTCVector.BTC_FALSE, IntVal(0))
 
     def is_result_of_hash(self):
         def parse_ast(e):

symbolic_execution/language/bitcoin/op_codes.py

@@ -8,19 +8,25 @@ from symbolic_execution.program.instruction.general import (
     InstructionTag,
     UnconditionalJump,
 )
-from symbolic_execution.program.instruction.hybrid import PickInstruction
+from symbolic_execution.program.instruction.hybrid import (
+    CheckMultiSigInstruction,
+    PickInstruction,
+)
 
 from symbolic_execution.program.instruction.interfaces import IInstruction
 from symbolic_execution.program.instruction.stack_based import (
     ArithmeticOperator,
-    CheckMultiSigInstruction,
     CheckSigInstruction,
     ConditionalJump,
     PushDataOperator,
     StackAlterationOperator,
     SwapDataOperator,
 )
-from symbolic_execution.hardware.machines.btc_machine import CHECK_SINGLE_SIG, BTCVector, BTCMachine
+from symbolic_execution.hardware.machines.btc_machine import (
+    CHECK_SINGLE_SIG,
+    BTCVector,
+    BTCMachine,
+)
 from z3 import And, Int, Or, Function, IntSort, BoolVal
 
 # Special Bitcoin Fields
@@ -128,12 +134,15 @@ def signature_constraints(ops: List[BTCVector]):
     if len(ops) != 2:
         raise ValueError("Exactly one pubkey and one signature are expected.")
 
-    pubkey = ops[0]
-    sig = ops[1]
-
+    pubkeys, sigs = ops
     conditions = []
-    for pk_const_val in pubkey.fetch_const_values():
-        conditions.append(CHECK_SINGLE_SIG(pk_const_val, sig.val) == BoolVal(False))
+
+    for pubkey in pubkeys:
+        for sig in sigs:
+            for pk_const_val in pubkey.fetch_const_values():
+                conditions.append(
+                    CHECK_SINGLE_SIG(pk_const_val, sig.val) == BoolVal(False)
+                )
 
     return conditions
 
@@ -736,7 +745,10 @@ class OpCodes:
         "OP_CHECKSIG": OpCode(
             "OP_CHECKSIG",
             instruction_constructor=CheckSigInstruction,
-            default_args={"sig_function": BTCVector.check_sig, "constraint_on_operands": signature_constraints},
+            default_args={
+                "sig_function": BTCVector.check_sig,
+                "constraint_on_operands": signature_constraints,
+            },
         ),
         "OP_CHECKMULTISIG": OpCode(
             "OP_CHECKMULTISIG",
@@ -809,7 +821,7 @@ class OpCodes:
         if opcode in OpCodes.__int_to_tokens:
             return OpCodes.__int_to_tokens[opcode]
         else:
-            raise InvalidOpCode(f"No known token for {opcode}")
+            raise InvalidOpCode(f"No known token for {hex(opcode)}")
 
     def token_to_opcode(token: str) -> OpCode:
         """Return the Opcode (object) associated to a token."""

symbolic_execution/program/instruction/hybrid.py

@@ -37,3 +37,56 @@ class PickInstruction(IInstruction):
     def __repr__(self) -> str:
         """Return the string representation of the instruction."""
         return f"{super().__repr__()}(top_stack_elt)"
+
+
+class CheckMultiSigInstruction(IInstruction):
+    """IInstruction that checks several signatures against public keys."""
+
+    def __init__(
+        self,
+        pc: int,
+        sig_function: Callable,
+        constraint_on_operands: Callable = None,
+        **args,
+    ):
+        """Initialize the check sig instruction."""
+        super().__init__(pc, **args)
+
+        if sig_function is None:
+            raise InstructionInitializationError("The signature function is None.")
+        self._sig_function = sig_function
+
+        self._constraint_on_operands = constraint_on_operands
+
+    def run(self, machine: IStackBasedMachine):
+        """Check several signatures against several public keys."""
+        super().run(machine)
+
+        pubkeys_count = machine.pop_from_stack()
+        if is_int_value(pubkeys_count.val):
+            pubkeys = [
+                machine.pop_from_stack() for _ in range(pubkeys_count.val.as_long())
+            ]
+        else:
+            machine.drop_from_stack(pubkeys_count)
+            pubkeys = []
+
+        sig_count = machine.pop_from_stack()
+        if is_int_value(sig_count.val):
+            sigs = [machine.pop_from_stack() for _ in range(sig_count.val.as_long())]
+        else:
+            machine.drop_from_stack(sig_count)
+            sigs = []
+        # Due to the bug in the BTC implementation, an extra element is dropped
+        machine.pop_from_stack()
+
+        ret = self._sig_function(pubkeys_count, sig_count, pubkeys, sigs)
+        machine.push_onto_stack(ret)
+
+        if self._constraint_on_operands and isinstance(machine, ISymbolicMachine):
+            constraints = self._constraint_on_operands([pubkeys, sigs])
+            machine.add_constraints(constraints)
+
+    def __repr__(self) -> str:
+        """Return the string representation of the instruction."""
+        return f"{super().__repr__()}({'multi_sig'})"

symbolic_execution/program/instruction/stack_based.py

@@ -272,32 +272,3 @@ class CheckSigInstruction(IInstruction):
     def __repr__(self) -> str:
         """Return the string representation of the instruction."""
         return f"{super().__repr__()}({'single_sig'})"
-
-
-class CheckMultiSigInstruction(IInstruction):
-    """IInstruction that checks several signatures against public keys."""
-
-    def __init__(self, pc: int, sig_function: Callable, **args):
-        """Initialize the check sig instruction."""
-        super().__init__(pc, **args)
-
-        if sig_function is None:
-            raise InstructionInitializationError("The signature function is None.")
-        self._sig_function = sig_function
-
-    def run(self, machine: IStackBasedMachine):
-        """Check several signatures against several public keys."""
-        super().run(machine)
-
-        # NOTE this code will crash against symbolic keys and signatures count
-        pubkeys_count = machine.pop_from_stack()
-        machine.drop_from_stack(pubkeys_count)
-        sig_count = machine.pop_from_stack()
-        machine.drop_from_stack(sig_count)
-
-        ret = self._sig_function(pubkeys_count, sig_count)
-        machine.push_onto_stack(ret)
-
-    def __repr__(self) -> str:
-        """Return the string representation of the instruction."""
-        return f"{super().__repr__()}({'multi_sig'})"
\ No newline at end of file