Merge pull request #248 from BitGo/BTC-2768.hydrate-with-rp

feat(wasm-utxo): support P2SH-P2PK replay protection inputs in legacy format hydration

Author: OttoAllmendinger

packages/wasm-utxo/js/fixedScriptWallet/BitGoPsbt.ts

@@ -124,11 +124,9 @@ export type ParseOutputsOptions = {
   payGoPubkeys?: ECPairArg[];
 };
 
-export type HydrationUnspent = {
-  chain: number;
-  index: number;
-  value: bigint;
-};
+export type HydrationUnspent =
+  | { chain: number; index: number; value: bigint } // wallet input
+  | { pubkey: Uint8Array; value: bigint }; // P2SH-P2PK replay protection input
 
 export class BitGoPsbt extends PsbtBase<WasmBitGoPsbt> implements IPsbtWithAddress {
   protected constructor(wasm: WasmBitGoPsbt) {

packages/wasm-utxo/src/fixed_script_wallet/bitgo_psbt/legacy_txformat.rs

@@ -120,6 +120,18 @@ pub use psbt_wallet_input::{
 };
 pub use psbt_wallet_output::ParsedOutput;
 
+/// Describes a single input for `from_half_signed_legacy_transaction`.
+pub enum HydrationUnspentInput {
+    /// A regular wallet input with derivation chain, index, and value.
+    Wallet(psbt_wallet_input::ScriptIdWithValue),
+    /// A P2SH-P2PK replay protection input. The caller provides the expected pubkey so it can be
+    /// validated against the redeemScript embedded in the legacy transaction.
+    ReplayProtection {
+        pubkey: miniscript::bitcoin::CompressedPublicKey,
+        value: u64,
+    },
+}
+
 /// Parsed transaction with wallet information
 #[derive(Debug, Clone)]
 pub struct ParsedTransaction {
@@ -505,12 +517,12 @@ impl BitGoPsbt {
     /// creates a PSBT with proper wallet metadata (bip32Derivation, scripts,
     /// witnessUtxo), and inserts the extracted signatures.
     ///
-    /// Only supports p2sh, p2shP2wsh, and p2wsh inputs (not taproot).
+    /// Supports p2sh, p2shP2wsh, p2wsh, and P2SH-P2PK (replay protection) inputs.
     pub fn from_half_signed_legacy_transaction(
         tx_bytes: &[u8],
         network: Network,
         wallet_keys: &crate::fixed_script_wallet::RootWalletKeys,
-        unspents: &[psbt_wallet_input::ScriptIdWithValue],
+        unspents: &[HydrationUnspentInput],
     ) -> Result<Self, String> {
         use miniscript::bitcoin::consensus::Decodable;
         use miniscript::bitcoin::{PublicKey, Transaction};
@@ -531,43 +543,113 @@ impl BitGoPsbt {
 
         let mut psbt = Self::new(network, wallet_keys, Some(version), Some(lock_time));
 
-        // Extract signatures before adding inputs (we need the raw tx_in data)
-        let partial_sigs: Vec<legacy_txformat::LegacyPartialSig> = tx
+        // Parse each input from the legacy tx
+        let input_results: Vec<legacy_txformat::LegacyInputResult> = tx
             .input
             .iter()
             .enumerate()
             .map(|(i, tx_in)| {
-                legacy_txformat::unsign_legacy_input(tx_in)
+                legacy_txformat::parse_legacy_input(tx_in)
                     .map_err(|e| format!("Input {}: {}", i, e))
             })
             .collect::<Result<Vec<_>, _>>()?;
 
-        // Add wallet inputs (populates bip32Derivation, scripts, witnessUtxo)
         for (i, (tx_in, unspent)) in tx.input.iter().zip(unspents.iter()).enumerate() {
-            let script_id = psbt_wallet_input::ScriptId {
-                chain: unspent.chain,
-                index: unspent.index,
-            };
-            psbt.add_wallet_input(
-                tx_in.previous_output.txid,
-                tx_in.previous_output.vout,
-                unspent.value,
-                wallet_keys,
-                script_id,
-                psbt_wallet_input::WalletInputOptions {
-                    sign_path: None,
-                    sequence: Some(tx_in.sequence.0),
-                    prev_tx: None, // psbt-lite: no nonWitnessUtxo
-                },
-            )
-            .map_err(|e| format!("Input {}: failed to add wallet input: {}", i, e))?;
+            match (&input_results[i], unspent) {
+                (
+                    legacy_txformat::LegacyInputResult::Multisig(sig),
+                    HydrationUnspentInput::Wallet(sv),
+                ) => {
+                    let script_id = psbt_wallet_input::ScriptId {
+                        chain: sv.chain,
+                        index: sv.index,
+                    };
+                    psbt.add_wallet_input(
+                        tx_in.previous_output.txid,
+                        tx_in.previous_output.vout,
+                        sv.value,
+                        wallet_keys,
+                        script_id,
+                        psbt_wallet_input::WalletInputOptions {
+                            sign_path: None,
+                            sequence: Some(tx_in.sequence.0),
+                            prev_tx: None, // psbt-lite: no nonWitnessUtxo
+                        },
+                    )
+                    .map_err(|e| format!("Input {}: failed to add wallet input: {}", i, e))?;
 
-            // Insert the extracted partial signature
-            let sig = &partial_sigs[i];
-            let pubkey = PublicKey::from(sig.pubkey);
-            psbt.psbt_mut().inputs[i]
-                .partial_sigs
-                .insert(pubkey, sig.sig);
+                    let pubkey = PublicKey::from(sig.pubkey);
+                    psbt.psbt_mut().inputs[i]
+                        .partial_sigs
+                        .insert(pubkey, sig.sig);
+                }
+                (
+                    legacy_txformat::LegacyInputResult::ReplayProtection {
+                        pubkey: tx_pubkey,
+                        sig,
+                    },
+                    HydrationUnspentInput::ReplayProtection {
+                        pubkey: expected_pubkey,
+                        value,
+                    },
+                ) => {
+                    if tx_pubkey.to_bytes() != expected_pubkey.to_bytes() {
+                        return Err(format!(
+                            "Input {}: replay protection pubkey mismatch: \
+                             tx has {}, expected {}",
+                            i,
+                            tx_pubkey
+                                .to_bytes()
+                                .iter()
+                                .map(|b| format!("{:02x}", b))
+                                .collect::<String>(),
+                            expected_pubkey
+                                .to_bytes()
+                                .iter()
+                                .map(|b| format!("{:02x}", b))
+                                .collect::<String>(),
+                        ));
+                    }
+                    psbt.add_replay_protection_input_at_index(
+                        i,
+                        *tx_pubkey,
+                        tx_in.previous_output.txid,
+                        tx_in.previous_output.vout,
+                        *value,
+                        ReplayProtectionOptions {
+                            sequence: Some(tx_in.sequence.0),
+                            prev_tx: None,
+                            sighash_type: None,
+                        },
+                    )
+                    .map_err(|e| {
+                        format!("Input {}: failed to add replay protection input: {}", i, e)
+                    })?;
+
+                    if let Some(ecdsa_sig) = sig {
+                        let pk = PublicKey::from(*tx_pubkey);
+                        psbt.psbt_mut().inputs[i]
+                            .partial_sigs
+                            .insert(pk, *ecdsa_sig);
+                    }
+                }
+                _ => {
+                    return Err(format!(
+                        "Input {}: mismatch between tx input type and provided unspent type \
+                         (tx has {}, unspent is {})",
+                        i,
+                        match &input_results[i] {
+                            legacy_txformat::LegacyInputResult::Multisig(_) => "multisig",
+                            legacy_txformat::LegacyInputResult::ReplayProtection { .. } =>
+                                "replay protection",
+                        },
+                        match unspent {
+                            HydrationUnspentInput::Wallet(_) => "wallet",
+                            HydrationUnspentInput::ReplayProtection { .. } => "replay protection",
+                        }
+                    ));
+                }
+            }
         }
 
         // Add outputs (plain script+value, no wallet metadata)
@@ -4197,7 +4279,7 @@ mod tests {
 
         // Step 2: Build unspents from bip32 derivation paths in the PSBT
         // The derivation path is m/<chain>/<index>
-        let unspents: Vec<ScriptIdWithValue> = psbt
+        let unspents: Vec<HydrationUnspentInput> = psbt
             .inputs
             .iter()
             .enumerate()
@@ -4219,11 +4301,11 @@ mod tests {
                     .ok_or_else(|| format!("Input {} has no witnessUtxo", i))?
                     .value
                     .to_sat();
-                Ok(ScriptIdWithValue {
+                Ok(HydrationUnspentInput::Wallet(ScriptIdWithValue {
                     chain,
                     index,
                     value,
-                })
+                }))
             })
             .collect::<Result<Vec<_>, String>>()?;
 

packages/wasm-utxo/src/fixed_script_wallet/bitgo_psbt/mod.rs

@@ -13,7 +13,7 @@ pub use checksigverify::{
     build_p2tr_ns_script, build_tap_tree_for_output, create_tap_bip32_derivation_for_output,
     ScriptP2tr,
 };
-pub use singlesig::{build_p2pk_script, ScriptP2shP2pk};
+pub use singlesig::{build_p2pk_script, parse_p2pk_script, ScriptP2shP2pk};
 
 use crate::address::networks::OutputScriptSupport;
 use crate::bitcoin::bip32::{ChildNumber, DerivationPath};

packages/wasm-utxo/src/fixed_script_wallet/wallet_scripts/mod.rs

@@ -12,6 +12,19 @@ pub fn build_p2pk_script(key: CompressedPublicKey) -> ScriptBuf {
         .into_script()
 }
 
+/// Parse a bare P2PK script (`<pubkey> OP_CHECKSIG`) and return the pubkey if valid.
+///
+/// P2PK format: `0x21 <33-byte compressed pubkey> 0xac`
+pub fn parse_p2pk_script(script: &ScriptBuf) -> Option<CompressedPublicKey> {
+    let b = script.as_bytes();
+    // 0x21 = push 33 bytes, 0xac = OP_CHECKSIG
+    if b.len() == 35 && b[0] == 0x21 && b[34] == 0xac {
+        CompressedPublicKey::from_slice(&b[1..34]).ok()
+    } else {
+        None
+    }
+}
+
 #[derive(Debug)]
 pub struct ScriptP2shP2pk {
     pub redeem_script: ScriptBuf,

packages/wasm-utxo/src/fixed_script_wallet/wallet_scripts/singlesig.rs

@@ -427,34 +427,58 @@ impl BitGoPsbt {
         unspents: JsValue,
     ) -> Result<BitGoPsbt, WasmUtxoError> {
         use crate::fixed_script_wallet::bitgo_psbt::psbt_wallet_input::ScriptIdWithValue;
+        use crate::fixed_script_wallet::bitgo_psbt::HydrationUnspentInput;
 
         let network = parse_network(network)?;
         let wallet_keys = wallet_keys.inner();
 
-        // Parse the unspents array from JsValue
+        // Parse the unspents array from JsValue.
+        // Each element is either:
+        //   { chain: number, index: number, value: bigint }  → wallet input
+        //   { value: bigint }                                 → replay protection input
+        // The presence of `chain` is used to distinguish the two.
         let arr = js_sys::Array::from(&unspents);
         let mut parsed_unspents = Vec::with_capacity(arr.length() as usize);
         for i in 0..arr.length() {
             let item = arr.get(i);
-            let chain = js_sys::Reflect::get(&item, &"chain".into())
-                .map_err(|_| WasmUtxoError::new("Missing 'chain' field on unspent"))?
-                .as_f64()
-                .ok_or_else(|| WasmUtxoError::new("'chain' must be a number"))?
-                as u32;
-            let index = js_sys::Reflect::get(&item, &"index".into())
-                .map_err(|_| WasmUtxoError::new("Missing 'index' field on unspent"))?
-                .as_f64()
-                .ok_or_else(|| WasmUtxoError::new("'index' must be a number"))?
-                as u32;
             let value_js = js_sys::Reflect::get(&item, &"value".into())
                 .map_err(|_| WasmUtxoError::new("Missing 'value' field on unspent"))?;
             let value = u64::try_from(js_sys::BigInt::unchecked_from_js(value_js))
                 .map_err(|_| WasmUtxoError::new("'value' must be a bigint convertible to u64"))?;
-            parsed_unspents.push(ScriptIdWithValue {
-                chain,
-                index,
-                value,
-            });
+
+            let chain_val =
+                js_sys::Reflect::get(&item, &"chain".into()).unwrap_or(JsValue::UNDEFINED);
+
+            if chain_val.is_undefined() {
+                // No 'chain' property → replay protection input; require pubkey
+                let pubkey_val = js_sys::Reflect::get(&item, &"pubkey".into()).map_err(|_| {
+                    WasmUtxoError::new("Missing 'pubkey' on replay protection unspent")
+                })?;
+                let pubkey_bytes = js_sys::Uint8Array::new(&pubkey_val).to_vec();
+                let pubkey = miniscript::bitcoin::CompressedPublicKey::from_slice(&pubkey_bytes)
+                    .map_err(|_| {
+                        WasmUtxoError::new(
+                            "'pubkey' is not a valid compressed public key (33 bytes)",
+                        )
+                    })?;
+                parsed_unspents.push(HydrationUnspentInput::ReplayProtection { pubkey, value });
+            } else {
+                // Has 'chain' → wallet input; also parse 'index'
+                let chain = chain_val
+                    .as_f64()
+                    .ok_or_else(|| WasmUtxoError::new("'chain' must be a number"))?
+                    as u32;
+                let index = js_sys::Reflect::get(&item, &"index".into())
+                    .map_err(|_| WasmUtxoError::new("Missing 'index' field on wallet unspent"))?
+                    .as_f64()
+                    .ok_or_else(|| WasmUtxoError::new("'index' must be a number"))?
+                    as u32;
+                parsed_unspents.push(HydrationUnspentInput::Wallet(ScriptIdWithValue {
+                    chain,
+                    index,
+                    value,
+                }));
+            }
         }
 
         let psbt =

packages/wasm-utxo/src/wasm/fixed_script_wallet/mod.rs

@@ -15,6 +15,7 @@ import * as utxolib from "@bitgo/utxo-lib";
 import { BitGoPsbt, type HydrationUnspent } from "../../js/fixedScriptWallet/BitGoPsbt.js";
 import { ZcashBitGoPsbt } from "../../js/fixedScriptWallet/ZcashBitGoPsbt.js";
 import { ChainCode } from "../../js/fixedScriptWallet/chains.js";
+import { ECPair } from "../../js/ecpair.js";
 import { getDefaultWalletKeys, getKeyTriple } from "../../js/testutils/keys.js";
 import { getCoinNameForNetwork } from "../networks.js";
 
@@ -128,4 +129,46 @@ describe("BitGoPsbt.fromHalfSignedLegacyTransaction", function () {
       });
     }
   });
+
+  describe("Round-trip with replay protection input", function () {
+    it("reconstructs PSBT from legacy tx with wallet + replay protection input", function () {
+      const rootWalletKeys = getDefaultWalletKeys();
+      const [userXprv] = getKeyTriple("default");
+      const ecpair = ECPair.fromPublicKey(rootWalletKeys.userKey().publicKey);
+
+      const psbt = BitGoPsbt.createEmpty("btc", rootWalletKeys, { version: 2, lockTime: 0 });
+      psbt.addWalletInput(
+        { txid: "00".repeat(32), vout: 0, value: BigInt(10000), sequence: 0xfffffffd },
+        rootWalletKeys,
+        { scriptId: { chain: 0, index: 0 } },
+      );
+      psbt.addReplayProtectionInput(
+        { txid: "aa".repeat(32), vout: 0, value: BigInt(1000), sequence: 0xfffffffd },
+        ecpair,
+      );
+      psbt.addWalletOutput(rootWalletKeys, { chain: 0, index: 100, value: BigInt(5000) });
+      // sign() only signs wallet inputs; replay protection input gets 0 sigs
+      psbt.sign(userXprv);
+
+      const txBytes = psbt.getHalfSignedLegacyFormat();
+
+      const unspents: HydrationUnspent[] = [
+        { chain: 0, index: 0, value: BigInt(10000) }, // wallet
+        { pubkey: ecpair.publicKey, value: BigInt(1000) }, // replay protection
+      ];
+      const reconstructed = BitGoPsbt.fromHalfSignedLegacyTransaction(
+        txBytes,
+        "btc",
+        rootWalletKeys,
+        unspents,
+      );
+
+      assert.ok(reconstructed.serialize().length > 0, "Reconstructed PSBT serializes");
+      assert.strictEqual(reconstructed.inputCount(), 2, "Both inputs present");
+      assert.ok(
+        reconstructed.verifySignature(0, rootWalletKeys.userKey().neutered().toBase58()),
+        "Wallet input signature preserved",
+      );
+    });
+  });
 });