Merge pull request #192 from BitGo/BTC-3000.add-get-wallet-keys

feat(wasm-utxo): add PSBT global xpub extraction and sorting

Author: OttoAllmendinger

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

@@ -1,5 +1,7 @@
 import {
   BitGoPsbt as WasmBitGoPsbt,
+  FixedScriptWalletNamespace,
+  WasmBIP32,
   type PsbtInputData,
   type PsbtOutputData,
   type PsbtOutputDataWithAddress,
@@ -972,4 +974,32 @@ export class BitGoPsbt implements IPsbtWithAddress {
   getOutputsWithAddress(): PsbtOutputDataWithAddress[] {
     return this._wasm.get_outputs_with_address() as PsbtOutputDataWithAddress[];
   }
+
+  /**
+   * Returns the unordered global xpubs from this PSBT as BIP32 instances.
+   */
+  getGlobalXpubs(): BIP32[] {
+    const result = this._wasm.get_global_xpubs() as WasmBIP32[];
+    return result.map((w) => BIP32.fromWasm(w));
+  }
+}
+
+/**
+ * Extract sorted wallet keys from a PSBT's global xpub fields.
+ *
+ * This should only be used in exceptional circumstances where the real wallet
+ * keys are not available — for example, legacy cold wallets where the PSBT
+ * was built with derived keys (from coldDerivationSeed) but the caller only
+ * has root xpubs. Prefer passing wallet keys explicitly wherever possible.
+ *
+ * @returns Sorted [user, backup, bitgo] RootWalletKeys
+ */
+export function getWalletKeysFromPsbt(psbt: BitGoPsbt, xpubs: BIP32[]): RootWalletKeys {
+  const wasmKeys = FixedScriptWalletNamespace.to_wallet_keys(
+    psbt.wasm,
+    xpubs[0].wasm,
+    xpubs[1].wasm,
+    xpubs[2].wasm,
+  );
+  return RootWalletKeys.fromWasm(wasmKeys);
 }

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

@@ -54,6 +54,14 @@ function extractDerivationPrefixes(keys: WalletKeysArg): Triple<string> | null {
 export class RootWalletKeys {
   private constructor(private _wasm: WasmRootWalletKeys) {}
 
+  /**
+   * Create a RootWalletKeys instance from a WasmRootWalletKeys instance (internal use)
+   * @internal
+   */
+  static fromWasm(wasm: WasmRootWalletKeys): RootWalletKeys {
+    return new RootWalletKeys(wasm);
+  }
+
   /**
    * Create a RootWalletKeys from various input formats
    * @param keys - Can be a triple of xpub strings, an IWalletKeys object, or another RootWalletKeys instance

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

@@ -17,6 +17,7 @@ export { ChainCode, chainCodes, assertChainCode, type Scope } from "./chains.js"
 // Bitcoin-like PSBT (for all non-Zcash networks)
 export {
   BitGoPsbt,
+  getWalletKeysFromPsbt,
   type NetworkName,
   type ScriptId,
   type ParsedInput,

packages/wasm-utxo/js/index.ts

@@ -106,6 +106,7 @@ declare module "./wasm/wasm_utxo.js" {
     getInputs(): PsbtInputData[];
     getOutputs(): PsbtOutputData[];
     getOutputsWithAddress(coin: import("./coinName.js").CoinName): PsbtOutputDataWithAddress[];
+    getGlobalXpubs(): WasmBIP32[];
     getPartialSignatures(inputIndex: number): Array<{
       pubkey: Uint8Array;
       signature: Uint8Array;

packages/wasm-utxo/js/psbt.ts

@@ -1,11 +1,13 @@
 import type { PsbtInputData, PsbtOutputData, PsbtOutputDataWithAddress } from "./wasm/wasm_utxo.js";
+import type { BIP32 } from "./bip32.js";
 
 /** Common interface for PSBT types */
 export interface IPsbt {
   inputCount(): number;
   outputCount(): number;
   getInputs(): PsbtInputData[];
   getOutputs(): PsbtOutputData[];
+  getGlobalXpubs(): BIP32[];
   version(): number;
   lockTime(): number;
   unsignedTxId(): string;

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

@@ -1292,6 +1292,22 @@ impl BitGoPsbt {
         }
     }
 
+    /// Returns the global xpubs from the PSBT, or None if the PSBT has no global xpubs.
+    ///
+    /// # Panics
+    /// Panics if the PSBT has global xpubs but not exactly 3.
+    pub fn get_global_xpubs(&self) -> Option<crate::fixed_script_wallet::XpubTriple> {
+        let xpubs: Vec<_> = self.psbt().xpub.keys().copied().collect();
+        if xpubs.is_empty() {
+            return None;
+        }
+        Some(
+            xpubs
+                .try_into()
+                .expect("expected exactly 3 global xpubs in PSBT"),
+        )
+    }
+
     /// Set version information in the PSBT's proprietary fields
     ///
     /// This embeds the wasm-utxo version and git hash into the PSBT's global
@@ -3003,6 +3019,71 @@ impl BitGoPsbt {
     }
 }
 
+/// All 6 orderings of a 3-element array, used to brute-force the
+/// [user, backup, bitgo] assignment from an unordered xpub triple.
+const XPUB_TRIPLE_PERMUTATIONS: [[usize; 3]; 6] = [
+    [0, 1, 2],
+    [0, 2, 1],
+    [1, 0, 2],
+    [1, 2, 0],
+    [2, 0, 1],
+    [2, 1, 0],
+];
+
+/// Sort an xpub triple into `[user, backup, bitgo]` order by trying all permutations
+/// against the PSBT's wallet inputs.
+///
+/// For each permutation, constructs `RootWalletKeys` and validates every non-replay-protection
+/// input against it. The first permutation where all inputs pass validation is returned.
+/// Works for all script types including p2tr.
+pub fn to_wallet_keys(
+    psbt: &BitGoPsbt,
+    xpubs: crate::fixed_script_wallet::XpubTriple,
+) -> Result<crate::fixed_script_wallet::RootWalletKeys, String> {
+    use crate::fixed_script_wallet::RootWalletKeys;
+
+    let inner_psbt = psbt.psbt();
+
+    // Collect non-replay-protection inputs (those with derivation info)
+    let wallet_inputs: Vec<_> = inner_psbt
+        .unsigned_tx
+        .input
+        .iter()
+        .zip(inner_psbt.inputs.iter())
+        .filter(|(_tx_input, psbt_input)| {
+            !psbt_input.bip32_derivation.is_empty() || !psbt_input.tap_key_origins.is_empty()
+        })
+        .collect();
+
+    if wallet_inputs.is_empty() {
+        return Err("no wallet inputs found in PSBT".to_string());
+    }
+
+    for perm in &XPUB_TRIPLE_PERMUTATIONS {
+        let permuted = [xpubs[perm[0]], xpubs[perm[1]], xpubs[perm[2]]];
+        let wallet_keys = RootWalletKeys::new(permuted);
+
+        let all_match = wallet_inputs.iter().all(|(tx_input, psbt_input)| {
+            let output_script = psbt_wallet_input::get_output_script_and_value(
+                psbt_input,
+                tx_input.previous_output,
+            );
+            match output_script {
+                Ok((script, _value)) => {
+                    psbt_wallet_input::assert_wallet_input(&wallet_keys, psbt_input, script).is_ok()
+                }
+                Err(_) => false,
+            }
+        });
+
+        if all_match {
+            return Ok(wallet_keys);
+        }
+    }
+
+    Err("no permutation of xpubs matches the PSBT wallet inputs".to_string())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -4812,4 +4893,78 @@ mod tests {
         assert!(!version_info.version.is_empty());
         assert!(!version_info.git_hash.is_empty());
     }
+
+    #[test]
+    fn test_get_global_xpubs() {
+        use crate::fixed_script_wallet::test_utils::get_test_wallet_keys;
+
+        let xpubs = get_test_wallet_keys("test_global_xpubs");
+        let wallet_keys = RootWalletKeys::new(xpubs);
+        let psbt = BitGoPsbt::new(Network::Bitcoin, &wallet_keys, Some(2), Some(0));
+
+        let global = psbt.get_global_xpubs().expect("should have global xpubs");
+        // The xpubs may be in BTreeMap order, not insertion order
+        let mut sorted_input: Vec<_> = xpubs.iter().map(|x| x.to_string()).collect();
+        sorted_input.sort();
+        let mut sorted_output: Vec<_> = global.iter().map(|x| x.to_string()).collect();
+        sorted_output.sort();
+        assert_eq!(sorted_input, sorted_output);
+    }
+
+    #[test]
+    fn test_to_wallet_keys_canonical_order() {
+        use crate::fixed_script_wallet::test_utils::get_test_wallet_keys;
+        use miniscript::bitcoin::hashes::Hash;
+
+        let xpubs = get_test_wallet_keys("test_to_wallet_keys");
+        let wallet_keys = RootWalletKeys::new(xpubs);
+        let mut psbt = BitGoPsbt::new(Network::Bitcoin, &wallet_keys, Some(2), Some(0));
+
+        let txid = Txid::all_zeros();
+        psbt.add_wallet_input(
+            txid,
+            0,
+            100_000,
+            &wallet_keys,
+            ScriptId {
+                chain: 10,
+                index: 0,
+            },
+            WalletInputOptions::default(),
+        )
+        .expect("add_wallet_input");
+
+        let result = to_wallet_keys(&psbt, xpubs).expect("should find correct order");
+        assert_eq!(result.xpubs, xpubs);
+    }
+
+    #[test]
+    fn test_to_wallet_keys_shuffled_order() {
+        use crate::fixed_script_wallet::test_utils::get_test_wallet_keys;
+        use miniscript::bitcoin::hashes::Hash;
+
+        let xpubs = get_test_wallet_keys("test_to_wallet_keys_shuffled");
+        let wallet_keys = RootWalletKeys::new(xpubs);
+        let mut psbt = BitGoPsbt::new(Network::Bitcoin, &wallet_keys, Some(2), Some(0));
+
+        let txid = Txid::all_zeros();
+        psbt.add_wallet_input(
+            txid,
+            0,
+            100_000,
+            &wallet_keys,
+            ScriptId {
+                chain: 10,
+                index: 0,
+            },
+            WalletInputOptions::default(),
+        )
+        .expect("add_wallet_input");
+
+        // Shuffle the xpubs: [bitgo, user, backup] instead of [user, backup, bitgo]
+        let shuffled = [xpubs[2], xpubs[0], xpubs[1]];
+        let result = to_wallet_keys(&psbt, shuffled).expect("should find correct order");
+        // Result should be sorted back to [user, backup, bitgo]
+        assert_eq!(result.xpubs, xpubs);
+    }
 }

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

@@ -224,7 +224,28 @@ impl FixedScriptWalletNamespace {
 
         result.into()
     }
+
+    /// Sort an xpub triple into [user, backup, bitgo] order by validating
+    /// against the PSBT's wallet inputs. Returns a RootWalletKeys with the
+    /// correct ordering.
+    #[wasm_bindgen]
+    pub fn to_wallet_keys(
+        psbt: &BitGoPsbt,
+        user_or_a: &WasmBIP32,
+        backup_or_b: &WasmBIP32,
+        bitgo_or_c: &WasmBIP32,
+    ) -> Result<WasmRootWalletKeys, WasmUtxoError> {
+        let xpubs = [
+            user_or_a.to_xpub()?,
+            backup_or_b.to_xpub()?,
+            bitgo_or_c.to_xpub()?,
+        ];
+        let wallet_keys = crate::fixed_script_wallet::bitgo_psbt::to_wallet_keys(&psbt.psbt, xpubs)
+            .map_err(|e| WasmUtxoError::new(&e))?;
+        Ok(WasmRootWalletKeys::from_inner(wallet_keys))
+    }
 }
+
 #[wasm_bindgen]
 pub struct BitGoPsbt {
     pub(crate) psbt: crate::fixed_script_wallet::bitgo_psbt::BitGoPsbt,
@@ -730,6 +751,11 @@ impl BitGoPsbt {
         crate::wasm::psbt::get_outputs_with_address_from_psbt(self.psbt.psbt(), self.psbt.network())
     }
 
+    /// Returns the global xpubs from the PSBT as an array of WasmBIP32 instances.
+    pub fn get_global_xpubs(&self) -> JsValue {
+        crate::wasm::psbt::get_global_xpubs_from_psbt(self.psbt.psbt())
+    }
+
     /// Parse transaction with wallet keys to identify wallet inputs/outputs
     pub fn parse_transaction_with_wallet_keys(
         &self,

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

@@ -225,6 +225,15 @@ pub fn get_outputs_from_psbt(psbt: &Psbt) -> Result<JsValue, WasmUtxoError> {
     outputs.try_to_js_value()
 }
 
+/// Get global xpubs from a PSBT as an array of WasmBIP32 instances
+pub fn get_global_xpubs_from_psbt(psbt: &Psbt) -> JsValue {
+    let arr = js_sys::Array::new();
+    for xpub in psbt.xpub.keys() {
+        arr.push(&WasmBIP32::from_xpub_internal(*xpub).into());
+    }
+    arr.into()
+}
+
 /// Get all PSBT outputs with resolved address strings
 pub fn get_outputs_with_address_from_psbt(
     psbt: &Psbt,
@@ -674,6 +683,12 @@ impl WrapPsbt {
         get_outputs_with_address_from_psbt(&self.0, network)
     }
 
+    /// Get global xpubs from the PSBT as an array of WasmBIP32 instances.
+    #[wasm_bindgen(js_name = getGlobalXpubs)]
+    pub fn get_global_xpubs(&self) -> JsValue {
+        get_global_xpubs_from_psbt(&self.0)
+    }
+
     /// Get partial signatures for an input
     /// Returns array of { pubkey: Uint8Array, signature: Uint8Array }
     #[wasm_bindgen(js_name = getPartialSignatures)]

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

@@ -19,6 +19,11 @@ impl WasmRootWalletKeys {
     pub(crate) fn inner(&self) -> &RootWalletKeys {
         &self.inner
     }
+
+    /// Create from an inner RootWalletKeys
+    pub(crate) fn from_inner(inner: RootWalletKeys) -> Self {
+        Self { inner }
+    }
 }
 
 #[wasm_bindgen]