feat: Add nested-segwit support

feat: add nested segwit support test nested segwit digest for 1 inputs and 2 output transactions add tests to verify nested segwit transactions test nested segwit digest for 2 inputs and 2 output transactions
Cypherock · Jul 12, 2024 · 1d49b1a · 1d49b1a
1 parent 38f2eff
commit 1d49b1a
Show file tree

Hide file tree

Showing 5 changed files with 391 additions and 13 deletions.
diff --git a/apps/btc_family/btc_pub_key.c b/apps/btc_family/btc_pub_key.c
@@ -204,7 +204,17 @@ static size_t btc_get_address(const uint8_t *seed,
     case NON_SEGWIT:
       hdnode_get_address(&node, g_btc_app->p2pkh_addr_ver, addr, 35);
       break;
-    // TODO: add support for taproot and segwit
+    case PURPOSE_SEGWIT:
+      ecdsa_get_address_segwit_p2sh(
+              node.public_key,
+              g_btc_app->p2sh_addr_ver,
+              node.curve->hasher_pubkey,
+              node.curve->hasher_base58,
+              addr,
+              36
+            );
+      break;
+    // TODO: add support for taproot
     default:
       break;
   }

diff --git a/apps/btc_family/btc_txn.c b/apps/btc_family/btc_txn.c
@@ -648,11 +648,26 @@ static bool sign_input(scrip_sig_t *signatures) {
   status = true;
   for (int idx = 0; idx < btc_txn_context->metadata.input_count; idx++) {
     // generate the input digest and respective private key
-    status = btc_digest_input(btc_txn_context, idx, buffer);
     memcpy(&t_node, &node, sizeof(HDNode));
     hdnode_private_ckd(&t_node, btc_txn_context->inputs[idx].change_index);
     hdnode_private_ckd(&t_node, btc_txn_context->inputs[idx].address_index);
     hdnode_fill_public_key(&t_node);
+
+    // detect input type
+    btc_sign_txn_input_script_pub_key_t *script =
+        &btc_txn_context->inputs[idx].script_pub_key;
+    btc_script_type_e type = btc_get_script_type(script->bytes, script->size);
+    if (SCRIPT_TYPE_P2SH == type) {
+      // replace BIP16 scriptpubkey with redeemscript(P2WPKH)
+      uint8_t buf[22] = {0};
+      buf[0] = 0;   // version byte
+      buf[1] = 20;  // push 20 bytes
+      ecdsa_get_pubkeyhash(t_node.public_key, t_node.curve->hasher_pubkey, buf + 2);
+      memcpy(btc_txn_context->inputs[idx].script_pub_key.bytes, buf, 22);
+      btc_txn_context->inputs[idx].script_pub_key.size = 22;
+    }
+
+    status = btc_digest_input(btc_txn_context, idx, buffer);
     ecdsa_sign_digest(
         curve, t_node.private_key, buffer, signatures[idx].bytes, NULL, NULL);
     signatures[idx].size = btc_sig_to_script_sig(

diff --git a/apps/btc_family/btc_txn_helpers.c b/apps/btc_family/btc_txn_helpers.c
@@ -390,6 +390,23 @@ static void update_locktime(btc_verify_input_t *verify_input_data,
     return;
   }
 }
+
+// TODO: Add chunking condition for varint decode
+// refer: https://app.clickup.com/t/9002019994/PRF-7288
+static int64_t varint_decode(const uint8_t *raw_txn_chunk, int32_t *offset) {
+  uint8_t first_byte = raw_txn_chunk[*offset];
+  if (first_byte < 0xFD) {
+    return first_byte;
+  } else {
+    // TODO: var-int varies between 1-9 bytes
+    // current implementation supports decoding
+    // upto 3 bytes only
+    uint8_t result[2];
+    memcpy(result, raw_txn_chunk + *offset + 1, 2);
+    *offset += 2;
+    return U16_READ_LE_ARRAY(result);
+  }
+}
 /*****************************************************************************
  * GLOBAL FUNCTIONS
  *****************************************************************************/
@@ -414,7 +431,7 @@ int btc_verify_input(const uint8_t *raw_txn_chunk,
     // store the number of inputs in the raw_txn
     verify_input_data->count = raw_txn_chunk[offset++];
     // TODO: Improve varint decode.
-    // size of variable containing ip-count/op-count
+    // size of variable containing script size and ip-count/op-count
     // varies (1-9 Bytes) depending on its value.
     // refer:
     // https://en.bitcoin.it/wiki/Protocol_documentation#Variable_length_integer
@@ -442,17 +459,18 @@ int btc_verify_input(const uint8_t *raw_txn_chunk,
           }
 
           case SCRIPT_LENGTH_CASE: {
-            if (offset + raw_txn_chunk[offset] + 1 + 4 > CHUNK_SIZE) {
+            int64_t script_length = varint_decode(raw_txn_chunk, &offset);
+            if (offset + script_length + 1 + 4 > CHUNK_SIZE) {
               verify_input_data->prev_offset =
-                  (offset + raw_txn_chunk[offset] + 1 + 4) - CHUNK_SIZE;
+                  (offset + script_length + 1 + 4) - CHUNK_SIZE;
               update_hash(
                   verify_input_data, raw_txn_chunk, chunk_index, CHUNK_SIZE);
               verify_input_data->input_parse =
                   PREVIOUS_TX_HASH_PLUS_OP_INDEX_CASE;
               verify_input_data->input_index++;
               return 4;
             } else {
-              offset += (raw_txn_chunk[offset] + 1 + 4);
+              offset += (script_length + 1 + 4);
             }
             break;
           }
@@ -652,7 +670,6 @@ bool btc_digest_input(const btc_txn_context_t *context,
   btc_sign_txn_input_script_pub_key_t *script =
       &context->inputs[index].script_pub_key;
   btc_script_type_e type = btc_get_script_type(script->bytes, script->size);
-
   if (SCRIPT_TYPE_P2WPKH == type) {
     // segwit digest calculation; could fail if segwit_cache not filled
     status = calculate_p2wpkh_digest(context, index, digest);