Merge pull request #209 from BitGo/BTC-3033.add-bip322-storage

feat(wasm-utxo): add BIP322 message storage and retrieval

Author: lcovar

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

@@ -133,6 +133,14 @@ export function addBip322Input(psbt: BitGoPsbt, params: AddBip322InputParams): n
   );
 }
 
+/**
+ * Get the BIP322 message stored at a PSBT input index.
+ * Returns null if no message is stored.
+ */
+export function getBip322Message(psbt: BitGoPsbt, inputIndex: number): string | null {
+  return Bip322Namespace.get_bip322_message(psbt.wasm, inputIndex) ?? null;
+}
+
 /**
  * Verify a single input of a BIP-0322 transaction proof
  *

packages/wasm-utxo/src/bip322/bitgo_psbt.rs

@@ -4,7 +4,8 @@
 //! with BitGo fixed-script wallets.
 
 use crate::fixed_script_wallet::bitgo_psbt::{
-    create_bip32_derivation, create_tap_bip32_derivation, BitGoPsbt,
+    create_bip32_derivation, create_tap_bip32_derivation, find_kv, BitGoKeyValue, BitGoPsbt,
+    ProprietaryKeySubtype,
 };
 use crate::fixed_script_wallet::wallet_scripts::{
     build_multisig_script_2_of_3, build_p2tr_ns_script, ScriptP2tr,
@@ -194,9 +195,38 @@ pub fn add_bip322_input(
         }
     }
 
+    // Store the BIP322 message as a proprietary field for later extraction
+    let (prop_key, prop_value) = BitGoKeyValue::new(
+        ProprietaryKeySubtype::Bip322Message,
+        vec![],
+        message.as_bytes().to_vec(),
+    )
+    .to_key_value();
+    inner_psbt.inputs[input_index]
+        .proprietary
+        .insert(prop_key, prop_value);
+
     Ok(input_index)
 }
 
+/// Extract the BIP322 message stored in a PSBT input's proprietary fields.
+/// Returns None if no message is stored at that index.
+pub fn get_bip322_message(psbt: &BitGoPsbt, input_index: usize) -> Result<Option<String>, String> {
+    let input = psbt
+        .psbt()
+        .inputs
+        .get(input_index)
+        .ok_or_else(|| format!("Input index {} out of bounds", input_index))?;
+
+    let mut iter = find_kv(ProprietaryKeySubtype::Bip322Message, &input.proprietary);
+    match iter.next() {
+        Some(kv) => String::from_utf8(kv.value)
+            .map(Some)
+            .map_err(|e| format!("Invalid UTF-8 in BIP322 message: {e}")),
+        None => Ok(None),
+    }
+}
+
 /// Verify a single input of a BIP-0322 transaction proof
 ///
 /// # Arguments

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

@@ -17,7 +17,7 @@ pub mod zcash_psbt;
 use crate::Network;
 pub use dash_psbt::DashBitGoPsbt;
 use miniscript::bitcoin::{psbt::Psbt, secp256k1, CompressedPublicKey, Txid};
-pub use propkv::{BitGoKeyValue, ProprietaryKeySubtype, WasmUtxoVersionInfo, BITGO};
+pub use propkv::{find_kv, BitGoKeyValue, ProprietaryKeySubtype, WasmUtxoVersionInfo, BITGO};
 pub use sighash::validate_sighash_type;
 pub use zcash_psbt::{
     decode_zcash_transaction_meta, ZcashBitGoPsbt, ZcashTransactionMeta,

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

@@ -204,6 +204,17 @@ impl Bip322Namespace {
         Ok(indices.into_iter().map(|i| i as u32).collect())
     }
 
+    /// Get the BIP322 message stored at a PSBT input index.
+    /// Returns null if no message is stored.
+    #[wasm_bindgen]
+    pub fn get_bip322_message(
+        psbt: &super::fixed_script_wallet::BitGoPsbt,
+        input_index: u32,
+    ) -> Result<Option<String>, WasmUtxoError> {
+        bitgo_psbt::get_bip322_message(&psbt.psbt, input_index as usize)
+            .map_err(|e| WasmUtxoError::new(&e))
+    }
+
     /// Verify a single input of a BIP-0322 transaction proof using pubkeys directly
     ///
     /// # Arguments

packages/wasm-utxo/test/bip322/index.ts

@@ -91,6 +91,41 @@ describe("BIP-0322", function () {
         });
       }, /BIP-0322 PSBT must have version 0/);
     });
+
+    it("should store and retrieve BIP322 message via proprietary field", function () {
+      const psbt = fixedScriptWallet.BitGoPsbt.createEmpty("testnet", walletKeys, { version: 0 });
+
+      bip322.addBip322Input(psbt, {
+        message: "Hello, BitGo!",
+        scriptId: { chain: 10, index: 0 },
+        rootWalletKeys: walletKeys,
+      });
+
+      assert.strictEqual(bip322.getBip322Message(psbt, 0), "Hello, BitGo!");
+    });
+
+    it("should store different messages for multiple inputs", function () {
+      const psbt = fixedScriptWallet.BitGoPsbt.createEmpty("testnet", walletKeys, { version: 0 });
+
+      bip322.addBip322Input(psbt, {
+        message: "Message A",
+        scriptId: { chain: 10, index: 0 },
+        rootWalletKeys: walletKeys,
+      });
+      bip322.addBip322Input(psbt, {
+        message: "Message B",
+        scriptId: { chain: 10, index: 1 },
+        rootWalletKeys: walletKeys,
+      });
+
+      assert.strictEqual(bip322.getBip322Message(psbt, 0), "Message A");
+      assert.strictEqual(bip322.getBip322Message(psbt, 1), "Message B");
+    });
+
+    it("should throw for input index out of bounds in getBip322Message", function () {
+      const psbt = fixedScriptWallet.BitGoPsbt.createEmpty("testnet", walletKeys, { version: 0 });
+      assert.throws(() => bip322.getBip322Message(psbt, 0), /out of bounds/);
+    });
   });
 
   describe("sign and verify per-input", function () {