WCN-3934 Path-agnostic wallet PSBT parsing and wallet prefix helpers

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

@@ -40,7 +40,11 @@ export type ParsedInput = {
   scriptId: ScriptId | null;
   scriptType: InputScriptType;
   sequence: number;
-  /** Full BIP32 derivation path from the wallet xpub (e.g. "0/1"). Null for replay-protection inputs. */
+  /**
+   * Full BIP32 derivation path recorded in the PSBT, from the wallet root key: the prefix followed
+   * by `[chain, index]` (e.g. "0/0/20/5" for legacy, "44/0/<h1>/<h2>/20/5" for BIP44-style).
+   * The prefix is everything except the last two components. Null for replay-protection inputs.
+   */
   derivationPath: string | null;
 };
 
@@ -51,7 +55,11 @@ export type ParsedOutput = {
   /** Set only when the derivation path is chain-standard (chain code encodes script type per BitGo convention). */
   scriptId: ScriptId | null;
   paygo: boolean;
-  /** Full BIP32 derivation path from the wallet xpub (e.g. "0/1"). Null for external outputs. */
+  /**
+   * Full BIP32 derivation path recorded in the PSBT, from the wallet root key: the prefix followed
+   * by `[chain, index]` (e.g. "0/0/20/5" for legacy, "44/0/<h1>/<h2>/20/5" for BIP44-style).
+   * The prefix is everything except the last two components. Null for external outputs.
+   */
   derivationPath: string | null;
 };
 
@@ -114,6 +122,14 @@ export type AddWalletInputOptions = {
   scriptId: ScriptId;
   /** Sign path - required for p2tr/p2trMusig2 (chains 30-41) */
   signPath?: SignPath;
+  /**
+   * The output script the input is expected to spend (e.g. the script of the unspent as indexed).
+   * If provided, it is compared with the script derived for `scriptId` from the wallet keys, and
+   * a difference throws a `WasmUtxoError` with code `AddWalletInputError.WalletInputScriptMismatch`
+   * (the PSBT is left unchanged). Catches wallet keys whose derivation prefix does not match the
+   * one the unspent was created under. If omitted, nothing is compared.
+   */
+  expectedScript?: Uint8Array;
 };
 
 export type AddWalletOutputOptions = {
@@ -493,6 +509,7 @@ export class BitGoPsbt<
       walletOptions.signPath?.cosigner,
       inputOptions.sequence,
       inputOptions.prevTx,
+      walletOptions.expectedScript,
     );
   }
 
@@ -513,6 +530,7 @@ export class BitGoPsbt<
       walletOptions.signPath?.cosigner,
       inputOptions.sequence,
       inputOptions.prevTx,
+      walletOptions.expectedScript,
     );
   }
 
@@ -1098,6 +1116,11 @@ export class BitGoPsbt<
  * was built with derived keys (from coldDerivationSeed) but the caller only
  * has root xpubs. Prefer passing wallet keys explicitly wherever possible.
  *
+ * The PSBT's key origins are matched by fingerprint along their full recorded path, so legacy
+ * (`0/0`), BIP44-style (`44/<coinType>/<h1>/<h2>`) and mixed PSBTs are all resolved. The
+ * returned keys carry the derivation prefix recorded in the PSBT's wallet inputs; if the inputs
+ * record more than one prefix (or none), the default `0/0` prefix is used.
+ *
  * @returns Sorted [user, backup, bitgo] RootWalletKeys
  */
 export function getWalletKeysFromPsbt(psbt: BitGoPsbt, xpubs: BIP32[]): RootWalletKeys {

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

@@ -4,6 +4,11 @@ import type { CoinName } from "../coinName.js";
 export { RootWalletKeys, type WalletKeysArg, type IWalletKeys } from "./RootWalletKeys.js";
 export { ReplayProtection, type ReplayProtectionArg } from "./ReplayProtection.js";
 export { outputScript, address } from "./address.js";
+export {
+  deriveWalletAccount,
+  deriveWalletPrefix,
+  type WalletPrefixKeysArg,
+} from "./walletPrefix.js";
 export { Dimensions } from "./Dimensions.js";
 export {
   outputScriptTypes,

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

@@ -0,0 +1,74 @@
+import { BIP32, type BIP32Arg } from "../bip32.js";
+import type { Triple } from "../triple.js";
+import { FixedScriptWalletNamespace } from "../wasm/wasm_utxo.js";
+import { RootWalletKeys, type WalletKeysArg } from "./RootWalletKeys.js";
+
+/**
+ * The three wallet root keys in `[user, backup, bitgo]` order.
+ *
+ * Accepts everything {@link RootWalletKeys.from} accepts, plus a triple of any {@link BIP32Arg}
+ * (xpub/xprv strings, `BIP32` instances, `BIP32Interface` objects). Private keys are neutered
+ * before use: the result is the same for an xprv and the matching xpub.
+ */
+export type WalletPrefixKeysArg = WalletKeysArg | Triple<BIP32Arg>;
+
+function toBIP32Triple(keys: WalletPrefixKeysArg): Triple<BIP32> {
+  if (keys instanceof RootWalletKeys) {
+    return [keys.userKey(), keys.backupKey(), keys.bitgoKey()];
+  }
+  if (Array.isArray(keys)) {
+    return (keys as Triple<BIP32Arg>).map((key) => BIP32.from(key)) as Triple<BIP32>;
+  }
+  return keys.triple.map((key) => BIP32.from(key)) as Triple<BIP32>;
+}
+
+/**
+ * Compute the `[h1, h2]` account indices of a wallet from its three root keys.
+ *
+ * ```
+ * d  = sha256("BitGo/utxo/wallet-prefix/v1" || K_user || K_backup || K_bitgo)
+ * K  = compressed public key (33 bytes) || chain code (32 bytes)    (of the neutered root key)
+ * h1 = u32be(d[0..4]) & 0x7fffffff
+ * h2 = u32be(d[4..8]) & 0x7fffffff
+ * ```
+ *
+ * Both are valid non-hardened BIP32 indices (`0..=2^31-1`). The result depends on the key order
+ * `[user, backup, bitgo]`. These are indices, not amounts, hence `number`.
+ *
+ * @param keys - The three wallet root keys `[user, backup, bitgo]`
+ */
+export function deriveWalletAccount(keys: WalletPrefixKeysArg): [number, number] {
+  const [user, backup, bitgo] = toBIP32Triple(keys);
+  const account = FixedScriptWalletNamespace.derive_wallet_account(
+    user.wasm,
+    backup.wasm,
+    bitgo.wasm,
+  );
+  return [account[0], account[1]];
+}
+
+/**
+ * Compute the relative per-wallet derivation prefix `44/<coinType>/<h1>/<h2>` (no leading `/` or
+ * `m/`), see {@link deriveWalletAccount}. Addresses of a wallet using this prefix derive at
+ * `<prefix>/<chain>/<index>` from each of the three root keys.
+ *
+ * @param keys - The three wallet root keys `[user, backup, bitgo]`
+ * @param coinType - BIP44 coin type of the coin family, an integer in `0..=2^31-1`. wasm-utxo has
+ *   no coin registry: the caller supplies it (e.g. `coin.bip44CoinType` from `@bitgo/statics`).
+ * @throws a `WasmUtxoError` with code `WalletPrefixError.CoinTypeOutOfRange` if `coinType` is
+ *   larger than `2^31-1`
+ */
+export function deriveWalletPrefix(keys: WalletPrefixKeysArg, coinType: number): string {
+  // wasm-bindgen wraps numbers into u32, so reject what cannot be a u32 before it gets there;
+  // values in (2^31-1, 2^32-1] are rejected by the Rust side with a typed error.
+  if (!Number.isInteger(coinType) || coinType < 0 || coinType > 0xffffffff) {
+    throw new RangeError(`coinType must be an integer in 0..=2^31-1, got ${coinType}`);
+  }
+  const [user, backup, bitgo] = toBIP32Triple(keys);
+  return FixedScriptWalletNamespace.derive_wallet_prefix(
+    user.wasm,
+    backup.wasm,
+    bitgo.wasm,
+    coinType,
+  );
+}

packages/wasm-utxo/src/error.rs

@@ -1,7 +1,10 @@
 use core::fmt;
 
+use crate::fixed_script_wallet::bitgo_psbt::psbt_wallet_input::AddWalletInputError;
 use crate::fixed_script_wallet::bitgo_psbt::zcash_psbt::VerifyV6SignatureError;
 use crate::fixed_script_wallet::bitgo_psbt::ParseTransactionError;
+use crate::fixed_script_wallet::wallet_prefix::WalletPrefixError;
+use crate::fixed_script_wallet::WalletKeysError;
 
 pub trait WasmErrorCode {
     fn code(&self) -> String;
@@ -32,6 +35,9 @@ pub enum WasmUtxoError {
     VerifyIronwood(crate::zcash::ironwood_build::IronwoodVerifyError),
     IronwoodPsbt(crate::fixed_script_wallet::bitgo_psbt::zcash::error::IronwoodPsbtError),
     VerifyV6Signature(VerifyV6SignatureError),
+    WalletKeys(WalletKeysError),
+    WalletPrefix(WalletPrefixError),
+    AddWalletInput(AddWalletInputError),
 }
 
 impl std::error::Error for WasmUtxoError {}
@@ -49,6 +55,9 @@ impl fmt::Display for WasmUtxoError {
             WasmUtxoError::VerifyIronwood(e) => write!(f, "{}", e),
             WasmUtxoError::IronwoodPsbt(e) => write!(f, "{}", e),
             WasmUtxoError::VerifyV6Signature(e) => write!(f, "{}", e),
+            WasmUtxoError::WalletKeys(e) => write!(f, "{}", e),
+            WasmUtxoError::WalletPrefix(e) => write!(f, "{}", e),
+            WasmUtxoError::AddWalletInput(e) => write!(f, "{}", e),
         }
     }
 }
@@ -76,6 +85,9 @@ impl WasmErrorCode for WasmUtxoError {
             WasmUtxoError::VerifyIronwood(e) => e.code(),
             WasmUtxoError::IronwoodPsbt(e) => e.code(),
             WasmUtxoError::VerifyV6Signature(e) => e.code(),
+            WasmUtxoError::WalletKeys(e) => e.code(),
+            WasmUtxoError::WalletPrefix(e) => e.code(),
+            WasmUtxoError::AddWalletInput(e) => e.code(),
         }
     }
 }
@@ -159,6 +171,24 @@ impl From<VerifyV6SignatureError> for WasmUtxoError {
     }
 }
 
+impl From<WalletKeysError> for WasmUtxoError {
+    fn from(err: WalletKeysError) -> Self {
+        WasmUtxoError::WalletKeys(err)
+    }
+}
+
+impl From<WalletPrefixError> for WasmUtxoError {
+    fn from(err: WalletPrefixError) -> Self {
+        WasmUtxoError::WalletPrefix(err)
+    }
+}
+
+impl From<AddWalletInputError> for WasmUtxoError {
+    fn from(err: AddWalletInputError) -> Self {
+        WasmUtxoError::AddWalletInput(err)
+    }
+}
+
 impl WasmUtxoError {
     pub fn new(s: &str) -> WasmUtxoError {
         WasmUtxoError::StringError(s.to_string())