feat: wrap updateInputWithDescriptor

packages/wasm-miniscript/src/descriptor.rs

@@ -7,14 +7,14 @@ use std::str::FromStr;
 use wasm_bindgen::prelude::wasm_bindgen;
 use wasm_bindgen::{JsError, JsValue};
 
-enum WrapDescriptorEnum {
+pub(crate) enum WrapDescriptorEnum {
     Derivable(Descriptor<DescriptorPublicKey>, KeyMap),
     Definite(Descriptor<DefiniteDescriptorKey>),
     String(Descriptor<String>),
 }
 
 #[wasm_bindgen]
-pub struct WrapDescriptor(WrapDescriptorEnum);
+pub struct WrapDescriptor(pub(crate) WrapDescriptorEnum);
 
 #[wasm_bindgen]
 impl WrapDescriptor {

packages/wasm-miniscript/src/psbt.rs

@@ -1,6 +1,9 @@
 use miniscript::bitcoin::Psbt;
+use miniscript::psbt::PsbtExt;
 use wasm_bindgen::prelude::wasm_bindgen;
 use wasm_bindgen::{JsError};
+use crate::descriptor::WrapDescriptorEnum;
+use crate::WrapDescriptor;
 
 #[wasm_bindgen]
 pub struct WrapPsbt(Psbt);
@@ -14,4 +17,19 @@ impl WrapPsbt {
     pub fn serialize(&self) -> Vec<u8> {
         self.0.serialize()
     }
+
+    #[wasm_bindgen(js_name = updateInputWithDescriptor)]
+    pub fn update_input_with_descriptor(&mut self, input_index: usize, descriptor: WrapDescriptor) -> Result<(), JsError> {
+        match descriptor.0 {
+            WrapDescriptorEnum::Definite(d) => {
+                self.0.update_input_with_descriptor(input_index, &d).map_err(JsError::from)
+            }
+            WrapDescriptorEnum::Derivable(_, _) => {
+                Err(JsError::new("Cannot update input with a derivable descriptor"))
+            }
+            WrapDescriptorEnum::String(_) => {
+                Err(JsError::new("Cannot update input with a string descriptor"))
+            }
+        }
+    }
 }

packages/wasm-miniscript/test/descriptorUtil.ts

@@ -0,0 +1,40 @@
+import * as utxolib from "@bitgo/utxo-lib";
+
+/** Expand a template with the given root wallet keys and chain code */
+function expand(template: string, rootWalletKeys: utxolib.bitgo.RootWalletKeys, chainCode: number) {
+  return template.replace(/\$([0-9])/g, (_, i) => {
+    const keyIndex = parseInt(i, 10);
+    if (keyIndex !== 0 && keyIndex !== 1 && keyIndex !== 2) {
+      throw new Error("Invalid key index");
+    }
+    const xpub = rootWalletKeys.triple[keyIndex].neutered().toBase58();
+    const prefix = rootWalletKeys.derivationPrefixes[keyIndex];
+    return xpub + "/" + prefix + "/" + chainCode + "/*";
+  });
+}
+
+/**
+ * Get a standard output descriptor that corresponds to the proprietary HD wallet setup
+ * used in BitGo wallets.
+ * Only supports a subset of script types.
+ */
+export function getDescriptorForScriptType(
+  rootWalletKeys: utxolib.bitgo.RootWalletKeys,
+  scriptType: utxolib.bitgo.outputScripts.ScriptType2Of3,
+  scope: "internal" | "external",
+): string {
+  const chain =
+    scope === "external"
+      ? utxolib.bitgo.getExternalChainCode(scriptType)
+      : utxolib.bitgo.getInternalChainCode(scriptType);
+  switch (scriptType) {
+    case "p2sh":
+      return expand("sh(multi(2,$0,$1,$2))", rootWalletKeys, chain);
+    case "p2shP2wsh":
+      return expand("sh(wsh(multi(2,$0,$1,$2)))", rootWalletKeys, chain);
+    case "p2wsh":
+      return expand("wsh(multi(2,$0,$1,$2))", rootWalletKeys, chain);
+    default:
+      throw new Error(`Unsupported script type ${scriptType}`);
+  }
+}

packages/wasm-miniscript/test/fixedScriptToDescriptor.ts

@@ -1,45 +1,7 @@
 import * as assert from "assert";
 import * as utxolib from "@bitgo/utxo-lib";
 import { Descriptor } from "../js";
-
-/** Expand a template with the given root wallet keys and chain code */
-function expand(template: string, rootWalletKeys: utxolib.bitgo.RootWalletKeys, chainCode: number) {
-  return template.replace(/\$([0-9])/g, (_, i) => {
-    const keyIndex = parseInt(i, 10);
-    if (keyIndex !== 0 && keyIndex !== 1 && keyIndex !== 2) {
-      throw new Error("Invalid key index");
-    }
-    const xpub = rootWalletKeys.triple[keyIndex].neutered().toBase58();
-    const prefix = rootWalletKeys.derivationPrefixes[keyIndex];
-    return xpub + "/" + prefix + "/" + chainCode + "/*";
-  });
-}
-
-/**
- * Get a standard output descriptor that corresponds to the proprietary HD wallet setup
- * used in BitGo wallets.
- * Only supports a subset of script types.
- */
-function getDescriptorForScriptType(
-  rootWalletKeys: utxolib.bitgo.RootWalletKeys,
-  scriptType: utxolib.bitgo.outputScripts.ScriptType2Of3,
-  scope: "internal" | "external",
-): string {
-  const chain =
-    scope === "external"
-      ? utxolib.bitgo.getExternalChainCode(scriptType)
-      : utxolib.bitgo.getInternalChainCode(scriptType);
-  switch (scriptType) {
-    case "p2sh":
-      return expand("sh(multi(2,$0,$1,$2))", rootWalletKeys, chain);
-    case "p2shP2wsh":
-      return expand("sh(wsh(multi(2,$0,$1,$2)))", rootWalletKeys, chain);
-    case "p2wsh":
-      return expand("wsh(multi(2,$0,$1,$2))", rootWalletKeys, chain);
-    default:
-      throw new Error(`Unsupported script type ${scriptType}`);
-  }
-}
+import { getDescriptorForScriptType } from "./descriptorUtil";
 
 const rootWalletKeys = new utxolib.bitgo.RootWalletKeys(utxolib.testutil.getKeyTriple("wasm"));
 const scriptTypes = ["p2sh", "p2shP2wsh", "p2wsh"] as const;

packages/wasm-miniscript/test/psbt.ts

@@ -1,36 +1,91 @@
 import * as utxolib from "@bitgo/utxo-lib";
 import * as assert from "node:assert";
-import { getPsbtFixtures } from "./psbtFixtures";
-import { Psbt } from "../js";
+import { getPsbtFixtures, toPsbtWithPrevOutOnly } from "./psbtFixtures";
+import { Descriptor, Psbt } from "../js";
 
-getPsbtFixtures().forEach(({ psbt, name }) => {
-  describe(`PSBT fixture ${name}`, function () {
+import { getDescriptorForScriptType } from "./descriptorUtil";
+
+const rootWalletKeys = new utxolib.bitgo.RootWalletKeys(utxolib.testutil.getKeyTriple("wasm"));
+
+function toWrappedPsbt(psbt: utxolib.bitgo.UtxoPsbt | Buffer | Uint8Array) {
+  if (psbt instanceof utxolib.bitgo.UtxoPsbt) {
+    psbt = psbt.toBuffer();
+  }
+  if (psbt instanceof Buffer || psbt instanceof Uint8Array) {
+    return Psbt.deserialize(psbt);
+  }
+  throw new Error("Invalid input");
+}
+
+function toUtxoPsbt(psbt: Psbt | Buffer | Uint8Array) {
+  if (psbt instanceof Psbt) {
+    psbt = psbt.serialize();
+  }
+  if (psbt instanceof Buffer || psbt instanceof Uint8Array) {
+    return utxolib.bitgo.UtxoPsbt.fromBuffer(Buffer.from(psbt), {
+      network: utxolib.networks.bitcoin,
+    });
+  }
+  throw new Error("Invalid input");
+}
+
+const fixtures = getPsbtFixtures(rootWalletKeys);
+
+function describeUpdateInputWithDescriptor(
+  psbt: utxolib.bitgo.UtxoPsbt,
+  scriptType: utxolib.bitgo.outputScripts.ScriptType2Of3,
+) {
+  const fullSignedFixture = fixtures.find(
+    (f) => f.scriptType === scriptType && f.stage === "fullsigned",
+  );
+  if (!fullSignedFixture) {
+    throw new Error("Could not find fullsigned fixture");
+  }
+
+  describe("updateInputWithDescriptor", function () {
+    it("should update the input with the descriptor", function () {
+      const descriptorStr = getDescriptorForScriptType(rootWalletKeys, scriptType, "internal");
+      const index = 0;
+      const descriptor = Descriptor.fromString(descriptorStr, "derivable");
+      const wrappedPsbt = toWrappedPsbt(toPsbtWithPrevOutOnly(psbt));
+      wrappedPsbt.updateInputWithDescriptor(0, descriptor.atDerivationIndex(index));
+      const updatedPsbt = toUtxoPsbt(wrappedPsbt);
+      updatedPsbt.signAllInputsHD(rootWalletKeys.triple[0]);
+      updatedPsbt.signAllInputsHD(rootWalletKeys.triple[2]);
+      updatedPsbt.finalizeAllInputs();
+      assert.deepStrictEqual(
+        fullSignedFixture.psbt
+          .clone()
+          .finalizeAllInputs()
+          .extractTransaction()
+          .toBuffer()
+          .toString("hex"),
+        updatedPsbt.extractTransaction().toBuffer().toString("hex"),
+      );
+    });
+  });
+}
+
+fixtures.forEach(({ psbt, scriptType, stage }) => {
+  describe(`PSBT fixture ${scriptType} ${stage}`, function () {
     let buf: Buffer;
     let wrappedPsbt: Psbt;
 
     before(function () {
       buf = psbt.toBuffer();
-      wrappedPsbt = Psbt.deserialize(buf);
+      wrappedPsbt = toWrappedPsbt(buf);
     });
 
     it("should map to same hex", function () {
-      assert.strictEqual(
-        buf.toString("hex"),
-        // it seems that the utxolib impl sometimes adds two extra bytes zero bytes at the end
-        // they probably are insignificant so we just add them here
-        Buffer.from(wrappedPsbt.serialize()).toString("hex") + (name === "empty" ? "0000" : ""),
-      );
+      assert.strictEqual(buf.toString("hex"), Buffer.from(wrappedPsbt.serialize()).toString("hex"));
     });
 
     it("should round-trip utxolib -> ms -> utxolib", function () {
-      assert.strictEqual(
-        buf.toString("hex"),
-        utxolib.bitgo.UtxoPsbt.fromBuffer(Buffer.from(wrappedPsbt.serialize()), {
-          network: utxolib.networks.bitcoin,
-        })
-          .toBuffer()
-          .toString("hex"),
-      );
+      assert.strictEqual(buf.toString("hex"), toUtxoPsbt(wrappedPsbt).toBuffer().toString("hex"));
     });
+
+    if (stage === "bare") {
+      describeUpdateInputWithDescriptor(psbt, scriptType);
+    }
   });
 });

packages/wasm-miniscript/test/psbtFixtures.ts

@@ -1,15 +1,38 @@
 import * as utxolib from "@bitgo/utxo-lib";
+import { RootWalletKeys } from "@bitgo/utxo-lib/dist/src/bitgo";
 
-function getEmptyPsbt() {
-  return new utxolib.bitgo.UtxoPsbt();
+type PsbtStage = "bare" | "unsigned" | "halfsigned" | "fullsigned";
+
+export function toPsbtWithPrevOutOnly(psbt: utxolib.bitgo.UtxoPsbt) {
+  const psbtCopy = utxolib.bitgo.UtxoPsbt.createPsbt({
+    network: utxolib.networks.bitcoin,
+  });
+  psbtCopy.setVersion(psbt.version);
+  psbtCopy.setLocktime(psbt.locktime);
+  psbt.txInputs.forEach((input, vin) => {
+    const { witnessUtxo, nonWitnessUtxo } = psbt.data.inputs[vin];
+    psbtCopy.addInput({
+      hash: input.hash,
+      index: input.index,
+      sequence: input.sequence,
+      ...(witnessUtxo ? { witnessUtxo } : { nonWitnessUtxo }),
+    });
+  });
+  psbt.txOutputs.forEach((output, vout) => {
+    psbtCopy.addOutput(output);
+  });
+  return psbtCopy;
 }
 
 function getPsbtWithScriptTypeAndStage(
-  seed: string,
+  keys: RootWalletKeys,
   scriptType: utxolib.bitgo.outputScripts.ScriptType2Of3,
-  stage: "unsigned" | "halfsigned" | "fullsigned",
+  stage: PsbtStage,
 ) {
-  const keys = new utxolib.bitgo.RootWalletKeys(utxolib.testutil.getKeyTriple(seed));
+  if (stage === "bare") {
+    const psbt = getPsbtWithScriptTypeAndStage(keys, scriptType, "unsigned");
+    return toPsbtWithPrevOutOnly(psbt);
+  }
   return utxolib.testutil.constructPsbt(
     [
       {
@@ -25,27 +48,28 @@ function getPsbtWithScriptTypeAndStage(
     ],
     utxolib.networks.bitcoin,
     keys,
-    "unsigned",
+    stage,
   );
 }
 
 export type PsbtFixture = {
   psbt: utxolib.bitgo.UtxoPsbt;
-  name: string;
+  scriptType: utxolib.bitgo.outputScripts.ScriptType2Of3;
+  stage: PsbtStage;
 };
 
-export function getPsbtFixtures(): PsbtFixture[] {
+export function getPsbtFixtures(keys: RootWalletKeys): PsbtFixture[] {
   const testMatrixScriptTypes = ["p2sh", "p2shP2wsh", "p2wsh"] as const;
-  const testMatrixStages = ["unsigned", "halfsigned", "fullsigned"] as const;
+  const testMatrixStages = ["bare", "unsigned", "halfsigned", "fullsigned"] as const;
 
-  const fixturesBitGo2Of3 = testMatrixStages.flatMap((stage) => {
+  return testMatrixStages.flatMap((stage) => {
     return testMatrixScriptTypes.map((scriptType) => {
       return {
-        psbt: getPsbtWithScriptTypeAndStage("wasm", scriptType, stage),
+        psbt: getPsbtWithScriptTypeAndStage(keys, scriptType, stage),
         name: `${scriptType}-${stage}`,
+        scriptType,
+        stage,
       };
     });
   });
-
-  return [{ psbt: getEmptyPsbt(), name: "empty" }, ...fixturesBitGo2Of3];
 }