Implemented tests and validation for P2MS (#257)

- [x] issue # - [x] follows contribution [guide](https://github.com/keep-starknet-strange/shinigami/blob/main/CONTRIBUTING.md) - [x] code change includes tests  Closes: #248 --------- Co-authored-by: Victor Omorogbe <[email protected]> Co-authored-by: Brandon Roberts <[email protected]>
keep-starknet-strange · Oct 19, 2024 · 20915c8 · 20915c8
1 parent 926e785
commit 20915c8
Show file tree

Hide file tree

Showing 11 changed files with 352 additions and 61 deletions.
diff --git a/packages/engine/src/errors.cairo b/packages/engine/src/errors.cairo
@@ -33,6 +33,7 @@ pub mod Error {
     pub const WITNESS_PROGRAM_WRONG_LENGTH: felt252 = 'Witness program wrong length';
     pub const WITNESS_PROGRAM_EMPTY: felt252 = 'Empty witness program';
     pub const SCRIPT_TOO_LARGE: felt252 = 'Script is too large';
+    pub const INVALID_P2MS: felt252 = 'Invalid P2MS transaction';
     pub const SCRIPT_UNFINISHED: felt252 = 'Script unfinished';
     pub const SCRIPT_ERR_SIG_DER: felt252 = 'Signature DER error';
 }

diff --git a/packages/engine/src/hash_cache.cairo b/packages/engine/src/hash_cache.cairo
@@ -1,7 +1,7 @@
 use crate::transaction::{
     EngineTransactionInputTrait, EngineTransactionOutputTrait, EngineTransactionTrait
 };
-use shinigami_utils::bytecode::int_size_in_bytes;
+use shinigami_utils::bytecode::write_var_int;
 use shinigami_utils::hash::double_sha256;
 
 #[derive(Clone, Copy, Drop)]
@@ -49,11 +49,7 @@ pub impl SigHashMidstateImpl<
         let outputs = transaction.get_transaction_outputs();
         for output in outputs {
             outputs_v0_bytes.append_word_rev(output.get_value().into(), 8);
-            outputs_v0_bytes
-                .append_word_rev(
-                    output.get_publickey_script().len().into(),
-                    int_size_in_bytes(output.get_publickey_script().len())
-                );
+            write_var_int(ref outputs_v0_bytes, output.get_publickey_script().len().into());
             outputs_v0_bytes.append(output.get_publickey_script());
         };
         SegwitSigHashMidstate {

diff --git a/packages/engine/src/signature/sighash.cairo b/packages/engine/src/signature/sighash.cairo
@@ -7,7 +7,7 @@ use crate::signature::utils::{
     remove_opcodeseparator, transaction_procedure, is_witness_pub_key_hash
 };
 use crate::hash_cache::SegwitSigHashMidstate;
-use shinigami_utils::bytecode::int_size_in_bytes;
+use shinigami_utils::bytecode::write_var_int;
 use shinigami_utils::hash::double_sha256;
 use crate::opcodes::opcodes::Opcode;
 
@@ -119,9 +119,7 @@ pub fn calc_witness_signature_hash<
         sig_hash_bytes.append_byte(Opcode::OP_EQUALVERIFY);
         sig_hash_bytes.append_byte(Opcode::OP_CHECKSIG);
     } else {
-        // TODO: VarIntBuf
-        sig_hash_bytes
-            .append_word_rev(sub_script.len().into(), int_size_in_bytes(sub_script.len()));
+        write_var_int(ref sig_hash_bytes, sub_script.len().into());
         sig_hash_bytes.append(sub_script);
     }
 
@@ -138,11 +136,7 @@ pub fn calc_witness_signature_hash<
         let output = transaction.get_transaction_outputs().at(tx_idx);
         let mut output_bytes: ByteArray = "";
         output_bytes.append_word_rev(output.get_value().into(), 8);
-        output_bytes
-            .append_word_rev(
-                output.get_publickey_script().len().into(),
-                int_size_in_bytes(output.get_publickey_script().len())
-            );
+        write_var_int(ref output_bytes, output.get_publickey_script().len().into());
         output_bytes.append(output.get_publickey_script());
         let hashed_output: u256 = double_sha256(@output_bytes);
         sig_hash_bytes.append_word(hashed_output.high.into(), 16);

diff --git a/packages/engine/src/signature/signature.cairo b/packages/engine/src/signature/signature.cairo
@@ -10,6 +10,7 @@ use crate::hash_cache::SigHashMidstateTrait;
 use shinigami_utils::byte_array::u256_from_byte_array_with_offset;
 use crate::signature::{sighash, constants};
 use crate::errors::Error;
+use crate::parser;
 
 //`BaseSigVerifier` is used to verify ECDSA signatures encoded in DER or BER format (pre-SegWit sig)
 #[derive(Drop)]
@@ -531,33 +532,20 @@ pub fn remove_signature(script: @ByteArray, sig_bytes: @ByteArray) -> @ByteArray
 
     let script_len = script.len();
     while i < script_len {
-        let push_data: u8 = script[i];
-        if push_data >= 8 && push_data <= 72 {
-            let mut len: usize = push_data.into();
-            let mut found: bool = false;
-
-            if len == sig_bytes.len() {
-                found = compare_data(script, sig_bytes, i, push_data);
-            }
-
-            if i + len <= script.len() {
-                i += len;
-            } else {
-                i += 1;
-            }
+        let opcode = script[i];
+        let data_len = parser::data_len(script, i).unwrap();
+        let end = i + data_len + 1;
+        if data_len == sig_bytes.len() {
+            let mut found = compare_data(script, sig_bytes, i, opcode);
             if found {
-                i += 1;
+                i = end;
                 continue;
             }
-            processed_script.append_byte(push_data);
-            while len != 0 && i - len < script_len {
-                processed_script.append_byte(script[i - len + 1]);
-                len -= 1;
-            };
-        } else {
-            processed_script.append_byte(push_data);
         }
-        i += 1;
+        while i != end {
+            processed_script.append_byte(script[i]);
+            i += 1;
+        };
     };
 
     @processed_script

diff --git a/packages/engine/src/transaction.cairo b/packages/engine/src/transaction.cairo
@@ -1,6 +1,6 @@
 use crate::errors::Error;
 use shinigami_utils::byte_array::{byte_array_value_at_le, byte_array_value_at_be, sub_byte_array};
-use shinigami_utils::bytecode::{int_size_in_bytes, bytecode_to_hex};
+use shinigami_utils::bytecode::{bytecode_to_hex, read_var_int, write_var_int};
 use shinigami_utils::bit_shifts::shr;
 use shinigami_utils::hash::double_sha256;
 
@@ -177,8 +177,7 @@ pub impl EngineInternalTransactionImpl of EngineInternalTransactionTrait {
             // consume flags
             offset += 2;
         }
-        // TODO: ReadVerIntBuf
-        let input_len: u8 = byte_array_value_at_le(@raw, ref offset, 1).try_into().unwrap();
+        let input_len = read_var_int(@raw, ref offset);
         // TODO: Error handling and bounds checks
         // TODO: Byte orderings
         let mut i = 0;
@@ -189,7 +188,7 @@ pub impl EngineInternalTransactionImpl of EngineInternalTransactionTrait {
                 low: byte_array_value_at_be(@raw, ref offset, 16).try_into().unwrap(),
             };
             let vout: u32 = byte_array_value_at_le(@raw, ref offset, 4).try_into().unwrap();
-            let script_len = byte_array_value_at_le(@raw, ref offset, 1).try_into().unwrap();
+            let script_len = read_var_int(@raw, ref offset).try_into().unwrap();
             let script = sub_byte_array(@raw, ref offset, script_len);
             let sequence: u32 = byte_array_value_at_le(@raw, ref offset, 4).try_into().unwrap();
             let input = EngineTransactionInput {
@@ -202,13 +201,13 @@ pub impl EngineInternalTransactionImpl of EngineInternalTransactionTrait {
             i += 1;
         };
 
-        let output_len: u8 = byte_array_value_at_le(@raw, ref offset, 1).try_into().unwrap();
+        let output_len = read_var_int(@raw, ref offset);
         let mut i = 0;
         let mut outputs: Array<EngineTransactionOutput> = array![];
         while i != output_len {
             // TODO: negative values
             let value: i64 = byte_array_value_at_le(@raw, ref offset, 8).try_into().unwrap();
-            let script_len = byte_array_value_at_le(@raw, ref offset, 1).try_into().unwrap();
+            let script_len = read_var_int(@raw, ref offset).try_into().unwrap();
             let script = sub_byte_array(@raw, ref offset, script_len);
             let output = EngineTransactionOutput { value: value, publickey_script: script, };
             outputs.append(output);
@@ -221,21 +220,17 @@ pub impl EngineInternalTransactionImpl of EngineInternalTransactionTrait {
             // one witness for each input
             i = 0;
             while i != input_len {
-                let witness_count: u8 = byte_array_value_at_le(@raw, ref offset, 1)
-                    .try_into()
-                    .unwrap();
+                let witness_count = read_var_int(@raw, ref offset);
                 let mut j = 0;
                 let mut witness: Array<ByteArray> = array![];
                 while j != witness_count {
-                    let script_len = byte_array_value_at_le(@raw, ref offset, 1)
-                        .try_into()
-                        .unwrap();
+                    let script_len = read_var_int(@raw, ref offset).try_into().unwrap();
                     let script = sub_byte_array(@raw, ref offset, script_len);
                     witness.append(script);
                     j += 1;
                 };
                 // update Transaction Input
-                let input = inputs.at(i.into());
+                let input = inputs.at(i.try_into().unwrap());
                 let mut input_with_witness = input.clone();
                 input_with_witness.witness = witness;
                 inputs_with_witness.append(input_with_witness);
@@ -284,7 +279,7 @@ pub impl EngineInternalTransactionImpl of EngineInternalTransactionTrait {
 
         // Serialize each input in the transaction.
         let input_len: usize = self.transaction_inputs.len();
-        bytes.append_word_rev(input_len.into(), int_size_in_bytes(input_len));
+        write_var_int(ref bytes, input_len.into());
         let mut i: usize = 0;
         while i != input_len {
             let input: @EngineTransactionInput = self.transaction_inputs.at(i);
@@ -297,7 +292,7 @@ pub impl EngineInternalTransactionImpl of EngineInternalTransactionTrait {
             bytes.append_word(input_txid.high.into(), 16);
             bytes.append_word(input_txid.low.into(), 16);
             bytes.append_word_rev(vout.into(), 4);
-            bytes.append_word_rev(script_len.into(), int_size_in_bytes(script_len));
+            write_var_int(ref bytes, script_len.into());
             bytes.append(script);
             bytes.append_word_rev(sequence.into(), 4);
 
@@ -306,7 +301,7 @@ pub impl EngineInternalTransactionImpl of EngineInternalTransactionTrait {
 
         // Serialize each output in the transaction.
         let output_len: usize = self.transaction_outputs.len();
-        bytes.append_word_rev(output_len.into(), int_size_in_bytes(output_len));
+        write_var_int(ref bytes, output_len.into());
         i = 0;
         while i != output_len {
             let output: @EngineTransactionOutput = self.transaction_outputs.at(i);
@@ -315,7 +310,7 @@ pub impl EngineInternalTransactionImpl of EngineInternalTransactionTrait {
             let script_len: usize = script.len();
 
             bytes.append_word_rev(value.into(), 8);
-            bytes.append_word_rev(script_len.into(), int_size_in_bytes(script_len));
+            write_var_int(ref bytes, script_len.into());
             bytes.append(script);
 
             i += 1;

diff --git a/packages/tests/src/lib.cairo b/packages/tests/src/lib.cairo
@@ -19,4 +19,5 @@ pub mod tests {
     mod test_transactions;
     mod test_p2pk;
     mod test_p2pkh;
+    mod test_p2ms;
 }