Merge pull request #365 from BitGo/veetragjain/cshld-1528-preserve-unified-addresses-in-shielded-output-parsing

fix(wasm-utxo): make isShielded nullable and preserve Unified Address on shielded outputs

Author: veetragjain

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

@@ -15,6 +15,7 @@ pub fn handle_add_shielded_output_command(
     anchor: String,
     ovk: Option<String>,
     memo: Option<String>,
+    unified_address: Option<String>,
 ) -> Result<()> {
     let raw_bytes = read_input_bytes(&path, "PSBT")?;
     let bytes = decode_input(&raw_bytes)?;
@@ -45,9 +46,17 @@ pub fn handle_add_shielded_output_command(
         None => [0u8; 512],
     };
 
-    psbt.add_ironwood_output(&recipient, value, ovk, &anchor, &memo, OsRng)
-        .map_err(|e| anyhow!(e))
-        .context("failed to add shielded output")?;
+    psbt.add_ironwood_output(
+        &recipient,
+        value,
+        ovk,
+        &anchor,
+        &memo,
+        unified_address.as_deref(),
+        OsRng,
+    )
+    .map_err(|e| anyhow!(e))
+    .context("failed to add shielded output")?;
 
     println!("{}", hex::encode(psbt.serialize().map_err(|e| anyhow!(e))?));
     Ok(())

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

@@ -133,6 +133,11 @@ pub enum PsbtCommand {
         /// Memo field, hex-encoded (512 bytes; default: all-zero)
         #[arg(long)]
         memo: Option<String>,
+        /// Full Unified Address this output was addressed to, if known. Its Orchard receiver must
+        /// match --recipient. Stored verbatim so it survives a serialize/deserialize round-trip
+        /// instead of being reconstructed as a single-receiver UA when the PSBT is later parsed.
+        #[arg(long)]
+        unified_address: Option<String>,
     },
     /// Sign one transparent input of a v6 PSBT with a single private key, over the ZIP-244
     /// transparent sighash. Call once per required signature (2-of-3). Prints the updated PSBT
@@ -244,6 +249,7 @@ pub fn handle_command(command: PsbtCommand) -> Result<()> {
             anchor,
             ovk,
             memo,
+            unified_address,
         } => add_shielded_output::handle_add_shielded_output_command(
             path,
             network.into(),
@@ -252,6 +258,7 @@ pub fn handle_command(command: PsbtCommand) -> Result<()> {
             anchor,
             ovk,
             memo,
+            unified_address,
         ),
         PsbtCommand::SignV6Input {
             path,

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

@@ -53,21 +53,11 @@ export type ParsedOutput = {
   paygo: boolean;
   /** Full BIP32 derivation path from the wallet xpub (e.g. "0/1"). Null for external outputs. */
   derivationPath: string | null;
-  /**
-   * True for a shielded (Orchard/Ironwood) output. Such an output has no `unsigned_tx` entry of
-   * its own — it lives in the PSBT's proprietary-map PCZT, read from its plaintext (not
-   * encrypted/decrypted) recipient field. `address` is a single-receiver ZIP-316 unified address
-   * (`u1...`/`utest1...`) encoding that receiver — a real, usable Zcash address, though not
-   * necessarily byte-identical to whatever multi-receiver UA the sender originally pasted in (a
-   * UA with a transparent/Sapling receiver too would round-trip to a different string carrying
-   * only the Orchard one). `script` holds the same receiver as raw 43 bytes.
-   */
-  isShielded: boolean;
 };
 
-export type ParsedTransaction = {
+export type ParsedTransaction<TOutput extends ParsedOutput = ParsedOutput> = {
   inputs: ParsedInput[];
-  outputs: ParsedOutput[];
+  outputs: TOutput[];
   spendAmount: bigint;
   minerFee: bigint;
   virtualSize: number;
@@ -145,7 +135,10 @@ export type HydrationUnspent =
   | { chain: number; index: number; value: bigint } // wallet input
   | { pubkey: Uint8Array; value: bigint }; // P2SH-P2PK replay protection input
 
-export class BitGoPsbt extends PsbtBase<WasmBitGoPsbt> implements IPsbtWithAddress {
+export class BitGoPsbt<TOutput extends ParsedOutput = ParsedOutput>
+  extends PsbtBase<WasmBitGoPsbt>
+  implements IPsbtWithAddress
+{
   protected constructor(wasm: WasmBitGoPsbt) {
     super(wasm);
   }
@@ -626,15 +619,15 @@ export class BitGoPsbt extends PsbtBase<WasmBitGoPsbt> implements IPsbtWithAddre
   parseTransactionWithWalletKeys(
     walletKeys: WalletKeysArg,
     options: ParseTransactionOptions,
-  ): ParsedTransaction {
+  ): ParsedTransaction<TOutput> {
     const keys = RootWalletKeys.from(walletKeys);
     const rp = ReplayProtection.from(options.replayProtection, this._wasm.network());
     const pubkeys = options.payGoPubkeys?.map((arg) => ECPair.from(arg).wasm);
     return this._wasm.parse_transaction_with_wallet_keys(
       keys.wasm,
       rp.wasm,
       pubkeys,
-    ) as ParsedTransaction;
+    ) as ParsedTransaction<TOutput>;
   }
 
   /**
@@ -650,13 +643,10 @@ export class BitGoPsbt extends PsbtBase<WasmBitGoPsbt> implements IPsbtWithAddre
    * @returns Array of parsed outputs
    * @note This method does NOT validate wallet inputs. It only parses outputs.
    */
-  parseOutputsWithWalletKeys(
-    walletKeys: WalletKeysArg,
-    options?: ParseOutputsOptions,
-  ): ParsedOutput[] {
+  parseOutputsWithWalletKeys(walletKeys: WalletKeysArg, options?: ParseOutputsOptions): TOutput[] {
     const keys = RootWalletKeys.from(walletKeys);
     const pubkeys = options?.payGoPubkeys?.map((arg) => ECPair.from(arg).wasm);
-    return this._wasm.parse_outputs_with_wallet_keys(keys.wasm, pubkeys) as ParsedOutput[];
+    return this._wasm.parse_outputs_with_wallet_keys(keys.wasm, pubkeys) as TOutput[];
   }
 
   /**

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

@@ -4,12 +4,30 @@ import {
   zcash_ironwood_version_group_id,
 } from "../wasm/wasm_utxo.js";
 import { type WalletKeysArg, RootWalletKeys } from "./RootWalletKeys.js";
-import { BitGoPsbt, type CreateEmptyOptions, type HydrationUnspent } from "./BitGoPsbt.js";
+import {
+  BitGoPsbt,
+  type CreateEmptyOptions,
+  type HydrationUnspent,
+  type ParsedOutput,
+} from "./BitGoPsbt.js";
 import { ZcashTransaction, type ITransaction } from "../transaction.js";
 
 /** Zcash network names */
 export type ZcashNetworkName = "zcash" | "zcashTest" | "zec" | "tzec";
 
+export type ZcashParsedOutput = ParsedOutput & {
+  /**
+   * True for a shielded (Orchard/Ironwood) output. Such an output has no `unsigned_tx` entry of
+   * its own — it lives in the PSBT's proprietary-map PCZT, read from its plaintext (not
+   * encrypted/decrypted) recipient field. `address` is a single-receiver ZIP-316 unified address
+   * (`u1...`/`utest1...`) encoding that receiver — a real, usable Zcash address, though not
+   * necessarily byte-identical to whatever multi-receiver UA the sender originally pasted in (a
+   * UA with a transparent/Sapling receiver too would round-trip to a different string carrying
+   * only the Orchard one). `script` holds the same receiver as raw 43 bytes.
+   */
+  isShielded: boolean;
+};
+
 /**
  * Zcash v6 (Ironwood) version group id (0xd884b698). Its presence marks a PSBT as v6 — see
  * `ZcashIronwoodBitGoPsbt`.
@@ -62,7 +80,7 @@ export type CreateEmptyZcashWithConsensusBranchIdOptions = CreateEmptyOptions &
  * const psbt = ZcashBitGoPsbt.fromBytes(bytes, "zcash");
  * ```
  */
-export class ZcashBitGoPsbt extends BitGoPsbt {
+export class ZcashBitGoPsbt extends BitGoPsbt<ZcashParsedOutput> {
   /**
    * Create an empty Zcash PSBT with consensus branch ID determined from block height
    *

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

@@ -198,14 +198,27 @@ export class ZcashIronwoodBitGoPsbt extends ZcashBitGoPsbt {
    * @param options.anchor - 32-byte Ironwood note-commitment-tree root
    * @param options.memo - optional 512-byte memo (defaults to the ZIP-302 "no memo" encoding)
    * @param options.ovk - optional 32-byte outgoing viewing key (omit for a keyless build)
+   * @param options.unifiedAddress - optional full Unified Address string this output was addressed
+   *   to. Its Orchard receiver must equal `recipient`. The PCZT itself only carries the raw 43-byte
+   *   receiver — a lossy encoding for a multi-receiver UA, since any transparent/Sapling receiver
+   *   can't be recovered from it — so passing this stores the original UA verbatim, letting a later
+   *   `parseOutputsWithWalletKeys`/`parseTransactionWithWalletKeys` (even after a
+   *   serialize/deserialize round-trip) return it in full instead of a re-encoded single-receiver UA.
    */
   addShieldedOutput(
     recipient: Uint8Array,
     amount: bigint,
-    options: { anchor: Uint8Array; memo?: Uint8Array; ovk?: Uint8Array },
+    options: { anchor: Uint8Array; memo?: Uint8Array; ovk?: Uint8Array; unifiedAddress?: string },
   ): void {
     const memo = options.memo ?? zip302NoMemo();
-    this.wasm.add_ironwood_output(recipient, amount, options.ovk, options.anchor, memo);
+    this.wasm.add_ironwood_output(
+      recipient,
+      amount,
+      options.ovk,
+      options.anchor,
+      memo,
+      options.unifiedAddress,
+    );
   }
 
   /**

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

@@ -43,6 +43,7 @@ export { BitGoKeySubtype, type PsbtKvKey } from "./BitGoKeySubtype.js";
 export {
   ZcashBitGoPsbt,
   type ZcashNetworkName,
+  type ZcashParsedOutput,
   type CreateEmptyZcashOptions,
   IRONWOOD_VERSION_GROUP_ID,
 } from "./ZcashBitGoPsbt.js";

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

@@ -2755,11 +2755,17 @@ impl BitGoPsbt {
         else {
             return Ok(None);
         };
-        let address = crate::zcash::unified_address::encode_orchard_receiver(
-            &recipient,
-            self.network().to_coin_name(),
-        )
-        .map_err(|e| ParseTransactionError::ShieldedOutput(e.to_string()))?;
+        // Prefer the caller's original Unified Address (if `add_ironwood_output` was given one):
+        // it may carry a transparent/Sapling receiver alongside the Orchard one, which a
+        // single-receiver reconstruction from `recipient` alone cannot recover.
+        let address = match propkv::get_ironwood_unified_address(&z.psbt) {
+            Some(ua) => ua,
+            None => crate::zcash::unified_address::encode_orchard_receiver(
+                &recipient,
+                self.network().to_coin_name(),
+            )
+            .map_err(|e| ParseTransactionError::ShieldedOutput(e.to_string()))?,
+        };
         Ok(Some((
             ParsedOutput {
                 address: Some(address),
@@ -2771,7 +2777,7 @@ impl BitGoPsbt {
                 script_id: None,
                 paygo: false,
                 derivation_path: None,
-                is_shielded: true,
+                is_shielded: Some(true),
             },
             amount,
         )))

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

@@ -271,6 +271,13 @@ pub enum ZecV6KeySubtype {
     /// happened, not just "no shielded output was ever added"). Persists even though the PCZT
     /// itself is gone, so a later read can tell the two "no PCZT" states apart.
     IronwoodExtracted = 0x04,
+    /// The full ZIP-316 Unified Address string (UTF-8) the shielded output was addressed to, if
+    /// the caller supplied one to [`crate::fixed_script_wallet::bitgo_psbt::zcash_psbt`]'s
+    /// `add_ironwood_output`. The PCZT itself only carries the raw 43-byte Orchard receiver, which
+    /// is lossy for a multi-receiver UA (transparent/Sapling receivers can't be recovered from it);
+    /// storing the original string here lets output parsing return the exact UA the caller passed,
+    /// receivers and all, after a serialize/deserialize round-trip.
+    UnifiedAddress = 0x05,
 }
 
 fn set_zec_v6(
@@ -292,6 +299,15 @@ fn get_zec_v6(psbt: &miniscript::bitcoin::psbt::Psbt, subtype: ZecV6KeySubtype)
         .map(|(_, v)| v.clone())
 }
 
+fn remove_zec_v6(psbt: &mut miniscript::bitcoin::psbt::Psbt, subtype: ZecV6KeySubtype) {
+    let key = ProprietaryKey {
+        prefix: BITGO_ZEC_V6.to_vec(),
+        subtype: subtype as u8,
+        key: vec![],
+    };
+    psbt.proprietary.remove(&key);
+}
+
 fn set_zec_v6_u32(
     psbt: &mut miniscript::bitcoin::psbt::Psbt,
     subtype: ZecV6KeySubtype,
@@ -377,6 +393,36 @@ pub fn get_zec_v6_params(psbt: &miniscript::bitcoin::psbt::Psbt) -> Option<(u32,
     Some((vgid, expiry))
 }
 
+/// Store the full Unified Address string the Ironwood shielded output was addressed to, so it
+/// survives a serialize/deserialize round-trip verbatim (receivers and all) instead of being
+/// rebuilt from just the raw Orchard receiver. Overwrites any existing value.
+pub fn set_ironwood_unified_address(psbt: &mut miniscript::bitcoin::psbt::Psbt, ua: &str) {
+    set_zec_v6(
+        psbt,
+        ZecV6KeySubtype::UnifiedAddress,
+        ua.as_bytes().to_vec(),
+    );
+}
+
+/// Fetch the Unified Address string stored by [`set_ironwood_unified_address`], if present and
+/// valid UTF-8.
+pub fn get_ironwood_unified_address(psbt: &miniscript::bitcoin::psbt::Psbt) -> Option<String> {
+    let bytes = get_zec_v6(psbt, ZecV6KeySubtype::UnifiedAddress)?;
+    String::from_utf8(bytes).ok()
+}
+
+/// Remove the Unified Address string set by [`set_ironwood_unified_address`], if present.
+///
+/// Callers that build a shielded output without a `unified_address` must call this rather than
+/// simply not calling [`set_ironwood_unified_address`]: `add_ironwood_output` can be called again
+/// on a PSBT whose PCZT was previously extracted (see `take_ironwood_pczt`/
+/// `mark_ironwood_extracted`, which drop only the PCZT key, not this one), and without an explicit
+/// removal a UA stored for an earlier shielded output would otherwise survive and be silently
+/// misattributed to the new one.
+pub fn remove_ironwood_unified_address(psbt: &mut miniscript::bitcoin::psbt::Psbt) {
+    remove_zec_v6(psbt, ZecV6KeySubtype::UnifiedAddress);
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -16,9 +16,10 @@ pub struct ParsedOutput {
     /// `None` for outputs that do not belong to this wallet.
     pub derivation_path: Option<DerivationPath>,
     /// Whether this output is a shielded (Orchard/Ironwood) output rather than a transparent one.
-    /// Always `false` for outputs parsed from `tx_output`/`psbt_output` — set by the caller when
-    /// synthesizing a `ParsedOutput` for the shielded side of a v6 (Ironwood) transaction.
-    pub is_shielded: bool,
+    /// `None` for coins that don't support shielded outputs. Always `Some(false)` for outputs
+    /// parsed from `tx_output`/`psbt_output` on a Zcash PSBT — set to `Some(true)` by the caller
+    /// when synthesizing a `ParsedOutput` for the shielded side of a v6 (Ironwood) transaction.
+    pub is_shielded: Option<bool>,
 }
 
 impl ParsedOutput {
@@ -63,7 +64,7 @@ impl ParsedOutput {
             script_id,
             paygo,
             derivation_path,
-            is_shielded: false,
+            is_shielded: matches!(network, Network::Zcash | Network::ZcashTestnet).then_some(false),
         })
     }