Merge pull request #68 from BitGo/BTC-2786.cosigning-with-strong-key-types

feat(wasm-utxo): implement signing workflow

Author: OttoAllmendinger

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

@@ -142,6 +142,66 @@ export class BitGoPsbt {
     return this.wasm.verify_signature_with_pub(inputIndex, wasmECPair);
   }
 
+  /**
+   * Sign a single input with a private key
+   *
+   * This method signs a specific input using the provided key. It accepts either:
+   * - An xpriv (BIP32Arg: base58 string, BIP32 instance, or WasmBIP32) for wallet inputs - derives the key and signs
+   * - A raw privkey (ECPairArg: Buffer, ECPair instance, or WasmECPair) for replay protection inputs - signs directly
+   *
+   * This method automatically detects and handles different input types:
+   * - For regular inputs: uses standard PSBT signing
+   * - For MuSig2 inputs: uses the FirstRound state stored by generateMusig2Nonces()
+   * - For replay protection inputs: signs with legacy P2SH sighash
+   *
+   * @param inputIndex - The index of the input to sign (0-based)
+   * @param key - Either an xpriv (BIP32Arg) or a raw privkey (ECPairArg)
+   * @throws Error if signing fails, or if generateMusig2Nonces() was not called first for MuSig2 inputs
+   *
+   * @example
+   * ```typescript
+   * // Parse transaction to identify input types
+   * const parsed = psbt.parseTransactionWithWalletKeys(walletKeys, replayProtection);
+   *
+   * // Sign regular wallet inputs with xpriv
+   * for (let i = 0; i < parsed.inputs.length; i++) {
+   *   const input = parsed.inputs[i];
+   *   if (input.scriptId !== null && input.scriptType !== "p2shP2pk") {
+   *     psbt.sign(i, userXpriv);
+   *   }
+   * }
+   *
+   * // Sign replay protection inputs with raw privkey
+   * const userPrivkey = bip32.fromBase58(userXpriv).privateKey!;
+   * for (let i = 0; i < parsed.inputs.length; i++) {
+   *   const input = parsed.inputs[i];
+   *   if (input.scriptType === "p2shP2pk") {
+   *     psbt.sign(i, userPrivkey);
+   *   }
+   * }
+   * ```
+   */
+  sign(inputIndex: number, key: BIP32Arg | ECPairArg): void {
+    // Detect key type
+    // If string or has 'derive' method → BIP32Arg
+    // Otherwise → ECPairArg
+    if (
+      typeof key === "string" ||
+      (typeof key === "object" &&
+        key !== null &&
+        "derive" in key &&
+        typeof key.derive === "function")
+    ) {
+      // It's a BIP32Arg
+      const wasmKey = BIP32.from(key as BIP32Arg);
+      this.wasm.sign_with_xpriv(inputIndex, wasmKey.wasm);
+    } else {
+      // It's an ECPairArg
+      const wasmKey = ECPair.from(key as ECPairArg);
+      this.wasm.sign_with_privkey(inputIndex, wasmKey.wasm);
+    }
+  }
+
   /**
    * @deprecated - use verifySignature with the replay protection key instead
    *
@@ -179,6 +239,72 @@ export class BitGoPsbt {
     return this.wasm.serialize();
   }
 
+  /**
+   * Generate and store MuSig2 nonces for all MuSig2 inputs
+   *
+   * This method generates nonces using the State-Machine API and stores them in the PSBT.
+   * The nonces are stored as proprietary fields in the PSBT and will be included when serialized.
+   * After ALL participants have generated their nonces, you can sign MuSig2 inputs using
+   * sign().
+   *
+   * @param key - The extended private key (xpriv) for signing. Can be a base58 string, BIP32 instance, or WasmBIP32
+   * @param sessionId - Optional 32-byte session ID for nonce generation. **Only allowed on testnets**.
+   *                    On mainnets, a secure random session ID is always generated automatically.
+   *                    Must be unique per signing session.
+   * @throws Error if nonce generation fails, sessionId length is invalid, or custom sessionId is
+   *         provided on a mainnet (security restriction)
+   *
+   * @security The sessionId MUST be cryptographically random and unique for each signing session.
+   * Never reuse a sessionId with the same key! On mainnets, sessionId is always randomly
+   * generated for security. Custom sessionId is only allowed on testnets for testing purposes.
+   *
+   * @example
+   * ```typescript
+   * // Phase 1: Both parties generate nonces (with auto-generated session ID)
+   * psbt.generateMusig2Nonces(userXpriv);
+   * // Nonces are stored in the PSBT
+   * // Send PSBT to counterparty
+   *
+   * // Phase 2: After receiving counterparty PSBT with their nonces
+   * const counterpartyPsbt = BitGoPsbt.fromBytes(counterpartyPsbtBytes, network);
+   * psbt.combineMusig2Nonces(counterpartyPsbt);
+   * // Sign MuSig2 key path inputs
+   * const parsed = psbt.parseTransactionWithWalletKeys(walletKeys, replayProtection);
+   * for (let i = 0; i < parsed.inputs.length; i++) {
+   *   if (parsed.inputs[i].scriptType === "p2trMusig2KeyPath") {
+   *     psbt.sign(i, userXpriv);
+   *   }
+   * }
+   * ```
+   */
+  generateMusig2Nonces(key: BIP32Arg, sessionId?: Uint8Array): void {
+    const wasmKey = BIP32.from(key);
+    this.wasm.generate_musig2_nonces(wasmKey.wasm, sessionId);
+  }
+
+  /**
+   * Combine/merge data from another PSBT into this one
+   *
+   * This method copies MuSig2 nonces and signatures (proprietary key-value pairs) from the
+   * source PSBT to this PSBT. This is useful for merging PSBTs during the nonce exchange
+   * and signature collection phases.
+   *
+   * @param sourcePsbt - The source PSBT containing data to merge
+   * @throws Error if networks don't match
+   *
+   * @example
+   * ```typescript
+   * // After receiving counterparty's PSBT with their nonces
+   * const counterpartyPsbt = BitGoPsbt.fromBytes(counterpartyPsbtBytes, network);
+   * psbt.combineMusig2Nonces(counterpartyPsbt);
+   * // Now can sign with all nonces present
+   * psbt.sign(0, userXpriv);
+   * ```
+   */
+  combineMusig2Nonces(sourcePsbt: BitGoPsbt): void {
+    this.wasm.combine_musig2_nonces(sourcePsbt.wasm);
+  }
+
   /**
    * Finalize all inputs in the PSBT
    *

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

@@ -204,6 +204,71 @@ impl BitGoPsbt {
         }
     }
 
+    /// Combine/merge data from another PSBT into this one
+    ///
+    /// This method copies MuSig2 nonces and signatures (proprietary key-value pairs) from the
+    /// source PSBT to this PSBT. This is useful for merging PSBTs during the nonce exchange
+    /// and signature collection phases.
+    ///
+    /// # Arguments
+    /// * `source_psbt` - The source PSBT containing data to merge
+    ///
+    /// # Returns
+    /// Ok(()) if data was successfully merged
+    ///
+    /// # Errors
+    /// Returns error if networks don't match
+    pub fn combine_musig2_nonces(&mut self, source_psbt: &BitGoPsbt) -> Result<(), String> {
+        // Check network match
+        if self.network() != source_psbt.network() {
+            return Err(format!(
+                "Network mismatch: destination is {}, source is {}",
+                self.network(),
+                source_psbt.network()
+            ));
+        }
+
+        let source = source_psbt.psbt();
+        let dest = self.psbt_mut();
+
+        // Check that both PSBTs have the same number of inputs
+        if source.inputs.len() != dest.inputs.len() {
+            return Err(format!(
+                "PSBT input count mismatch: source has {} inputs, destination has {}",
+                source.inputs.len(),
+                dest.inputs.len()
+            ));
+        }
+
+        // Copy MuSig2 nonces and partial signatures (proprietary key-values with BITGO identifier)
+        for (source_input, dest_input) in source.inputs.iter().zip(dest.inputs.iter_mut()) {
+            // Only process if the input is a MuSig2 input
+            if !p2tr_musig2_input::Musig2Input::is_musig2_input(source_input) {
+                continue;
+            }
+
+            // Parse nonces from source input using native Musig2 functions
+            let nonces = p2tr_musig2_input::parse_musig2_nonces(source_input)
+                .map_err(|e| format!("Failed to parse MuSig2 nonces from source: {}", e))?;
+
+            // Copy each nonce to the destination input
+            for nonce in nonces {
+                let (key, value) = nonce.to_key_value().to_key_value();
+                dest_input.proprietary.insert(key, value);
+            }
+
+            // Also copy partial signatures if present
+            // Partial sigs are stored as tap_script_sigs in the PSBT input
+            for (control_block, leaf_script) in &source_input.tap_script_sigs {
+                dest_input
+                    .tap_script_sigs
+                    .insert(*control_block, *leaf_script);
+            }
+        }
+
+        Ok(())
+    }
+
     /// Serialize the PSBT to bytes, using network-specific logic
     pub fn serialize(&self) -> Result<Vec<u8>, SerializeError> {
         match self {
@@ -426,6 +491,146 @@ impl BitGoPsbt {
             .map_err(|e| e.to_string())
     }
 
+    /// Sign a single input with a raw private key
+    ///
+    /// This method signs a specific input using the provided private key. It automatically
+    /// detects the input type and uses the appropriate signing method:
+    /// - Replay protection inputs (P2SH-P2PK): Signs with legacy P2SH sighash
+    /// - Regular inputs: Uses standard PSBT signing
+    /// - MuSig2 inputs: Returns error (requires FirstRound state, use sign_with_first_round)
+    ///
+    /// # Arguments
+    /// - `input_index`: The index of the input to sign
+    /// - `privkey`: The private key to sign with
+    ///
+    /// # Returns
+    /// - `Ok(())` if signing was successful
+    /// - `Err(String)` if signing fails or input type is not supported
+    pub fn sign_with_privkey(
+        &mut self,
+        input_index: usize,
+        privkey: &secp256k1::SecretKey,
+    ) -> Result<(), String> {
+        use miniscript::bitcoin::{
+            ecdsa::Signature as EcdsaSignature, hashes::Hash, sighash::SighashCache, PublicKey,
+        };
+
+        // Get network before mutable borrow
+        let network = self.network();
+        let is_testnet = network.is_testnet();
+
+        let psbt = self.psbt_mut();
+
+        // Check bounds
+        if input_index >= psbt.inputs.len() {
+            return Err(format!(
+                "Input index {} out of bounds (total inputs: {})",
+                input_index,
+                psbt.inputs.len()
+            ));
+        }
+
+        // Check if this is a MuSig2 input
+        if p2tr_musig2_input::Musig2Input::is_musig2_input(&psbt.inputs[input_index]) {
+            return Err(
+                "MuSig2 inputs cannot be signed with raw privkey. Use sign_with_first_round instead."
+                    .to_string(),
+            );
+        }
+
+        let secp = secp256k1::Secp256k1::new();
+
+        // Derive public key from private key
+        let public_key = PublicKey::from_slice(
+            &secp256k1::PublicKey::from_secret_key(&secp, privkey).serialize(),
+        )
+        .map_err(|e| format!("Failed to derive public key: {}", e))?;
+
+        // Check if this is a replay protection input (P2SH-P2PK)
+        if let Some(redeem_script) = &psbt.inputs[input_index].redeem_script.clone() {
+            // Try to extract pubkey from redeem script
+            if let Ok(redeem_pubkey) = Self::extract_pubkey_from_p2pk_redeem_script(redeem_script) {
+                // This is a replay protection input - verify the derived pubkey matches
+                if public_key != redeem_pubkey {
+                    return Err(
+                        "Public key mismatch: derived pubkey does not match redeem_script pubkey"
+                            .to_string(),
+                    );
+                }
+
+                // Sign the replay protection input with legacy P2SH sighash
+                let sighash_type = miniscript::bitcoin::sighash::EcdsaSighashType::All;
+                let cache = SighashCache::new(&psbt.unsigned_tx);
+                let sighash = cache
+                    .legacy_signature_hash(input_index, redeem_script, sighash_type.to_u32())
+                    .map_err(|e| format!("Failed to compute sighash: {}", e))?;
+
+                // Create ECDSA signature
+                let message = secp256k1::Message::from_digest(sighash.to_byte_array());
+                let signature = secp.sign_ecdsa(&message, privkey);
+                let ecdsa_sig = EcdsaSignature {
+                    signature,
+                    sighash_type,
+                };
+
+                // Add signature to partial_sigs
+                psbt.inputs[input_index]
+                    .partial_sigs
+                    .insert(public_key, ecdsa_sig);
+
+                return Ok(());
+            }
+        }
+
+        // For regular inputs (non-RP, non-MuSig2), use standard signing via miniscript
+        // This will handle legacy, SegWit, and Taproot script path inputs
+        match self {
+            BitGoPsbt::BitcoinLike(ref mut psbt, _network) => {
+                // Create a key provider that returns our single key
+                // Convert SecretKey to PrivateKey for the GetKey trait
+                // Note: The network parameter is only used for WIF serialization, not for signing
+                let bitcoin_network = if is_testnet {
+                    miniscript::bitcoin::Network::Testnet
+                } else {
+                    miniscript::bitcoin::Network::Bitcoin
+                };
+                let private_key = miniscript::bitcoin::PrivateKey::new(*privkey, bitcoin_network);
+                let key_map = std::collections::BTreeMap::from_iter([(public_key, private_key)]);
+
+                // Sign the PSBT
+                let result = psbt.sign(&key_map, &secp);
+
+                // Check if our specific input was signed
+                match result {
+                    Ok(signing_keys) => {
+                        if signing_keys.contains_key(&input_index) {
+                            Ok(())
+                        } else {
+                            Err(format!(
+                                "Input {} was not signed (no key found or already signed)",
+                                input_index
+                            ))
+                        }
+                    }
+                    Err((partial_success, errors)) => {
+                        // Check if there's an error for our specific input
+                        if let Some(error) = errors.get(&input_index) {
+                            Err(format!("Failed to sign input {}: {:?}", input_index, error))
+                        } else if partial_success.contains_key(&input_index) {
+                            // Input was signed successfully despite other errors
+                            Ok(())
+                        } else {
+                            Err(format!("Input {} was not signed", input_index))
+                        }
+                    }
+                }
+            }
+            BitGoPsbt::Zcash(_zcash_psbt, _network) => {
+                Err("Zcash signing not yet implemented".to_string())
+            }
+        }
+    }
+
     /// Sign the PSBT with the provided key.
     /// Wraps the underlying PSBT's sign method from miniscript::psbt::PsbtExt.
     ///

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

@@ -334,4 +334,12 @@ impl WasmBIP32 {
     pub(crate) fn to_xpub(&self) -> Result<crate::bitcoin::bip32::Xpub, WasmUtxoError> {
         Ok(self.0.to_xpub())
     }
+
+    /// Convert to Xpriv (for internal Rust use, not exposed to JS)
+    pub(crate) fn to_xpriv(&self) -> Result<crate::bitcoin::bip32::Xpriv, WasmUtxoError> {
+        match &self.0 {
+            BIP32Key::Private(xpriv) => Ok(*xpriv),
+            BIP32Key::Public(_) => Err(WasmUtxoError::new("Cannot get xpriv from public key")),
+        }
+    }
 }

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

@@ -41,6 +41,13 @@ impl WasmECPair {
     pub(crate) fn get_public_key(&self) -> PublicKey {
         self.key.public_key()
     }
+
+    /// Get the private key as a secp256k1::SecretKey (for internal Rust use)
+    pub(crate) fn get_private_key(&self) -> Result<SecretKey, WasmUtxoError> {
+        self.key
+            .secret_key()
+            .ok_or_else(|| WasmUtxoError::new("Cannot get private key from public-only ECPair"))
+    }
 }
 
 #[wasm_bindgen]