Merge pull request #66 from BitGo/BTC-2786.parse-input-type

feat(wasm-utxo): add input script type detection

Author: lcovar

packages/wasm-utxo/js/fixedScriptWallet.ts

@@ -56,11 +56,21 @@ type ReplayProtection =
 
 export type ScriptId = { chain: number; index: number };
 
+export type InputScriptType =
+  | "p2shP2pk"
+  | "p2sh"
+  | "p2shP2wsh"
+  | "p2wsh"
+  | "p2trLegacy"
+  | "p2trMusig2ScriptPath"
+  | "p2trMusig2KeyPath";
+
 export type ParsedInput = {
   address: string;
   script: Uint8Array;
   value: bigint;
   scriptId: ScriptId | null;
+  scriptType: InputScriptType;
 };
 
 export type ParsedOutput = {

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

@@ -97,7 +97,7 @@ pub enum BitGoPsbt {
 }
 
 // Re-export types from submodules for convenience
-pub use psbt_wallet_input::{ParsedInput, ScriptId};
+pub use psbt_wallet_input::{InputScriptType, ParsedInput, ScriptId};
 pub use psbt_wallet_output::ParsedOutput;
 
 /// Parsed transaction with wallet information

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

@@ -518,13 +518,75 @@ pub struct ScriptId {
     pub index: u32,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum InputScriptType {
+    P2shP2pk,
+    P2sh,
+    P2shP2wsh,
+    P2wsh,
+    P2trLegacy,
+    P2trMusig2ScriptPath,
+    P2trMusig2KeyPath,
+}
+
+impl InputScriptType {
+    pub fn from_script_id(script_id: ScriptId, psbt_input: &Input) -> Result<Self, String> {
+        let chain = Chain::try_from(script_id.chain).map_err(|e| e.to_string())?;
+        match chain {
+            Chain::P2shExternal | Chain::P2shInternal => Ok(InputScriptType::P2sh),
+            Chain::P2shP2wshExternal | Chain::P2shP2wshInternal => Ok(InputScriptType::P2shP2wsh),
+            Chain::P2wshExternal | Chain::P2wshInternal => Ok(InputScriptType::P2wsh),
+            Chain::P2trInternal | Chain::P2trExternal => Ok(InputScriptType::P2trLegacy),
+            Chain::P2trMusig2Internal | Chain::P2trMusig2External => {
+                // check if tap_script_sigs or tap_scripts are set
+                if !psbt_input.tap_script_sigs.is_empty() || !psbt_input.tap_scripts.is_empty() {
+                    Ok(InputScriptType::P2trMusig2ScriptPath)
+                } else {
+                    Ok(InputScriptType::P2trMusig2KeyPath)
+                }
+            }
+        }
+    }
+
+    /// Detects the script type from a script_id chain and PSBT input metadata
+    ///
+    /// # Arguments
+    /// - `script_id`: Optional script ID containing chain information (None for replay protection inputs)
+    /// - `psbt_input`: The PSBT input containing signature metadata
+    /// - `output_script`: The output script being spent
+    /// - `replay_protection`: Replay protection configuration
+    ///
+    /// # Returns
+    /// - `Ok(InputScriptType)` with the detected script type
+    /// - `Err(String)` if the script type cannot be determined
+    pub fn detect(
+        script_id: Option<ScriptId>,
+        psbt_input: &Input,
+        output_script: &ScriptBuf,
+        replay_protection: &ReplayProtection,
+    ) -> Result<Self, String> {
+        // For replay protection inputs (no script_id), detect from output script
+        match script_id {
+            Some(id) => Self::from_script_id(id, psbt_input),
+            None => {
+                if replay_protection.is_replay_protection_input(output_script) {
+                    Ok(InputScriptType::P2shP2pk)
+                } else {
+                    Err("Input without script_id is not a replay protection input".to_string())
+                }
+            }
+        }
+    }
+}
+
 /// Parsed input from a PSBT transaction
 #[derive(Debug, Clone)]
 pub struct ParsedInput {
     pub address: String,
     pub script: Vec<u8>,
     pub value: u64,
     pub script_id: Option<ScriptId>,
+    pub script_type: InputScriptType,
 }
 
 impl ParsedInput {
@@ -576,11 +638,17 @@ impl ParsedInput {
         )
         .map_err(ParseInputError::Address)?;
 
+        // Detect the script type using script_id chain information
+        let script_type =
+            InputScriptType::detect(script_id, psbt_input, output_script, replay_protection)
+                .map_err(ParseInputError::ScriptTypeDetection)?;
+
         Ok(Self {
             address,
             script: output_script.to_bytes(),
             value: value.to_sat(),
             script_id,
+            script_type,
         })
     }
 }
@@ -598,6 +666,8 @@ pub enum ParseInputError {
     WalletValidation(String),
     /// Failed to generate address for input
     Address(crate::address::AddressError),
+    /// Failed to detect script type for input
+    ScriptTypeDetection(String),
 }
 
 impl std::fmt::Display for ParseInputError {
@@ -618,6 +688,9 @@ impl std::fmt::Display for ParseInputError {
             ParseInputError::Address(error) => {
                 write!(f, "failed to generate address: {}", error)
             }
+            ParseInputError::ScriptTypeDetection(error) => {
+                write!(f, "failed to detect script type: {}", error)
+            }
         }
     }
 }

packages/wasm-utxo/src/wasm/try_into_js_value.rs

@@ -317,12 +317,29 @@ impl TryIntoJsValue for crate::fixed_script_wallet::bitgo_psbt::ScriptId {
     }
 }
 
+impl TryIntoJsValue for crate::fixed_script_wallet::bitgo_psbt::InputScriptType {
+    fn try_to_js_value(&self) -> Result<JsValue, WasmUtxoError> {
+        use crate::fixed_script_wallet::bitgo_psbt::InputScriptType;
+        let script_type = match self {
+            InputScriptType::P2shP2pk => "p2shP2pk",
+            InputScriptType::P2sh => "p2sh",
+            InputScriptType::P2shP2wsh => "p2shP2wsh",
+            InputScriptType::P2wsh => "p2wsh",
+            InputScriptType::P2trLegacy => "p2trLegacy",
+            InputScriptType::P2trMusig2ScriptPath => "p2trMusig2ScriptPath",
+            InputScriptType::P2trMusig2KeyPath => "p2trMusig2KeyPath",
+        };
+        Ok(JsValue::from_str(script_type))
+    }
+}
+
 impl TryIntoJsValue for crate::fixed_script_wallet::bitgo_psbt::ParsedInput {
     fn try_to_js_value(&self) -> Result<JsValue, WasmUtxoError> {
         js_obj!(
             "address" => self.address.clone(),
             "value" => self.value,
-            "scriptId" => self.script_id
+            "scriptId" => self.script_id,
+            "scriptType" => self.script_type
         )
     }
 }

packages/wasm-utxo/test/fixedScript/parseTransactionWithWalletKeys.ts

@@ -1,9 +1,29 @@
 import assert from "node:assert";
 import * as utxolib from "@bitgo/utxo-lib";
 import { fixedScriptWallet } from "../../js/index.js";
-import { BitGoPsbt } from "../../js/fixedScriptWallet.js";
+import { BitGoPsbt, InputScriptType } from "../../js/fixedScriptWallet.js";
 import { loadPsbtFixture, loadWalletKeysFromFixture, getPsbtBuffer } from "./fixtureUtil.js";
 
+function getExpectedInputScriptType(fixtureScriptType: string): InputScriptType {
+  // Map fixture types to InputScriptType values
+  // Based on the Rust mapping in src/fixed_script_wallet/test_utils/fixtures.rs
+  switch (fixtureScriptType) {
+    case "p2shP2pk":
+    case "p2sh":
+    case "p2shP2wsh":
+    case "p2wsh":
+      return fixtureScriptType;
+    case "p2tr":
+      return "p2trLegacy";
+    case "p2trMusig2":
+      return "p2trMusig2ScriptPath";
+    case "taprootKeyPathSpend":
+      return "p2trMusig2KeyPath";
+    default:
+      throw new Error(`Unknown fixture script type: ${fixtureScriptType}`);
+  }
+}
+
 function getOtherWalletKeys(): utxolib.bitgo.RootWalletKeys {
   const otherWalletKeys = utxolib.testutil.getKeyTriple("too many secrets");
   return new utxolib.bitgo.RootWalletKeys(otherWalletKeys);
@@ -34,9 +54,11 @@ describe("parseTransactionWithWalletKeys", function () {
       let fullsignedPsbtBytes: Buffer;
       let bitgoPsbt: BitGoPsbt;
       let rootWalletKeys: utxolib.bitgo.RootWalletKeys;
+      let fixture: ReturnType<typeof loadPsbtFixture>;
 
       before(function () {
-        fullsignedPsbtBytes = getPsbtBuffer(loadPsbtFixture(networkName, "fullsigned"));
+        fixture = loadPsbtFixture(networkName, "fullsigned");
+        fullsignedPsbtBytes = getPsbtBuffer(fixture);
         bitgoPsbt = fixedScriptWallet.BitGoPsbt.fromBytes(fullsignedPsbtBytes, networkName);
         rootWalletKeys = loadWalletKeysFromFixture(networkName);
       });
@@ -117,6 +139,23 @@ describe("parseTransactionWithWalletKeys", function () {
         assert.ok(parsed.virtualSize > 0, "Virtual size should be > 0");
       });
 
+      it("should parse inputs with correct scriptType", function () {
+        const parsed = bitgoPsbt.parseTransactionWithWalletKeys(rootWalletKeys, {
+          outputScripts: [replayProtectionScript],
+        });
+
+        // Verify all inputs have scriptType matching fixture
+        parsed.inputs.forEach((input, i) => {
+          const fixtureInput = fixture.psbtInputs[i];
+          const expectedScriptType = getExpectedInputScriptType(fixtureInput.type);
+          assert.strictEqual(
+            input.scriptType,
+            expectedScriptType,
+            `Input ${i} scriptType should be ${expectedScriptType}, got ${input.scriptType}`,
+          );
+        });
+      });
+
       it("should fail to parse with other wallet keys", function () {
         assert.throws(
           () => {