Hotfix/353 sign psbt transaction backwards compatible (#372)

babylonlabs-io · Jul 23, 2024 · c3be517 · c3be517
2 parents f3ae459 + 72e22ff
commit c3be517
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 84 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "simple-staking",
-  "version": "0.2.16",
+  "version": "0.2.17",
   "private": true,
   "scripts": {
     "dev": "next dev",

diff --git a/src/app/common/utils/psbt.ts b/src/app/common/utils/psbt.ts
@@ -14,11 +14,16 @@ export const signPsbtTransaction = (wallet: WalletProvider) => {
     const signedHex = await wallet.signPsbt(psbtHex);
     const providerName = await wallet.getWalletProviderName();
     if (SIGN_PSBT_NOT_COMPATIBLE_WALLETS.includes(providerName)) {
-      // The old implementation of signPsbt returns the signed transaction in hex
-      return Transaction.fromHex(signedHex);
+      try {
+        // Try to parse the signedHex as PSBT to see if it follows the new implementation
+        return Psbt.fromHex(signedHex).extractTransaction();
+      } catch {
+        // If parsing fails, it's the old version implementation
+        return Transaction.fromHex(signedHex);
+      }
     }
-    // The new implementation of signPsbt returns the signed PSBT in hex
-    // We need to extract the transaction from the PSBT
+
+    // For compatible wallets, directly extract the transaction from the signed PSBT
     return Psbt.fromHex(signedHex).extractTransaction();
   };
 };
diff --git a/tests/utils/signPsbtTransaction.test.ts b/tests/utils/signPsbtTransaction.test.ts
@@ -1,94 +1,79 @@
-import { Transaction } from "bitcoinjs-lib";
-import { initBTCCurve } from "btc-staking-ts";
+import { Psbt, Transaction } from "bitcoinjs-lib";
 
 import { signPsbtTransaction } from "@/app/common/utils/psbt";
-import { WalletProvider } from "@/utils/wallet/wallet_provider";
 
-import { testingNetworks } from "../helper";
-
-describe("signPsbtTransaction", () => {
-  initBTCCurve();
-  const { dataGenerator } = testingNetworks[0];
-  const btcWallet = {
-    signPsbt: jest.fn(),
-    getWalletProviderName: jest.fn(),
-    pushTx: jest.fn().mockResolvedValue(true),
-  } as unknown as WalletProvider;
-
-  it("should return a signed transaction for compatible wallets", async () => {
-    const randomGlobalParamsVersions =
-      dataGenerator.generateGlobalPramsVersions(
-        dataGenerator.getRandomIntegerBetween(1, 20),
-      );
-    const stakerKeys = dataGenerator.generateRandomKeyPair();
-    const randomStakingTxHeight =
-      randomGlobalParamsVersions[
-        dataGenerator.getRandomIntegerBetween(
-          0,
-          randomGlobalParamsVersions.length - 1,
-        )
-      ].activationHeight + 1;
-
-    const { unsignedPsbt } = dataGenerator.createRandomStakingPsbt(
-      randomGlobalParamsVersions,
-      randomStakingTxHeight,
-      stakerKeys,
-    );
+// Mock the WalletProvider
+const mockWalletProvider = (providerName: string, signedHex: string) => ({
+  signPsbt: jest.fn().mockResolvedValue(signedHex),
+  getWalletProviderName: jest.fn().mockResolvedValue(providerName),
+});
 
-    const unsignedPsbtHex = unsignedPsbt.toHex();
+const mocked_signedTxHexPair = {
+  psbtHex:
+    "70736274ff0100db0200000001103eee601ed9524da7b94953d17ac95cae07540ee4b5fd944669de461ee00c900200000000fdffffff033075000000000000225120b435cd0ce5f2b53a2bb3f0265af2c5240ea604988d7fba3bc4d43e3249641f730000000000000000496a476262643400392fc0de15a9a8cf5488235976bcecf5743ee0ef73f61d6068b2eb91c6449d317259aa77dcb96e09d5ae1204bc0648ce6160d8945db4ca12d883aac5e2042c1f00966cb6000000000000225120e2138f042b525186db6e22526f86de3d29d3ec9bff023f16e4827599f3596570d92003000001012ba72c010000000000225120e2138f042b525186db6e22526f86de3d29d3ec9bff023f16e4827599f3596570011720392fc0de15a9a8cf5488235976bcecf5743ee0ef73f61d6068b2eb91c6449d3100000000",
+  signedTxHex:
+    "02000000000101103eee601ed9524da7b94953d17ac95cae07540ee4b5fd944669de461ee00c900200000000fdffffff033075000000000000225120b435cd0ce5f2b53a2bb3f0265af2c5240ea604988d7fba3bc4d43e3249641f730000000000000000496a476262643400392fc0de15a9a8cf5488235976bcecf5743ee0ef73f61d6068b2eb91c6449d317259aa77dcb96e09d5ae1204bc0648ce6160d8945db4ca12d883aac5e2042c1f00966cb6000000000000225120e2138f042b525186db6e22526f86de3d29d3ec9bff023f16e4827599f359657001406ddb0bf6016660fa315584c32323119ec9788a5ef9a0c1ab6330f8105c9c1c2037ec367f2f376d5e0760bcdc9b74d65b80404074860e37bdd1e983572dff3207d9200300",
+};
+const mocked_signedPsbtHexPair = {
+  psbtHex:
+    "70736274ff0100cf0200000001f5c333aef0e41a646c22953109cb22936386d0b4c6ae4806a65044ffb4333c140200000000fdffffff0330750000000000002251205c3c9613d7c8e1069962e48df5df6b39f5ff799a077f9818e403962951d327940000000000000000496a476262643400634d38e0306a1623db98da1e229eed182b0c61811bdcc26cba5c8373a5df84417259aa77dcb96e09d5ae1204bc0648ce6160d8945db4ca12d883aac5e2042c1f0096a1e63f01000000001600140a70b8068c6c24946a9e41c23e5bb4084f5266a0d92003000001011fe65c4001000000001600140a70b8068c6c24946a9e41c23e5bb4084f5266a000000000",
+  signedPsbtHex:
+    "70736274ff0100cf0200000001f5c333aef0e41a646c22953109cb22936386d0b4c6ae4806a65044ffb4333c140200000000fdffffff0330750000000000002251205c3c9613d7c8e1069962e48df5df6b39f5ff799a077f9818e403962951d327940000000000000000496a476262643400634d38e0306a1623db98da1e229eed182b0c61811bdcc26cba5c8373a5df84417259aa77dcb96e09d5ae1204bc0648ce6160d8945db4ca12d883aac5e2042c1f0096a1e63f01000000001600140a70b8068c6c24946a9e41c23e5bb4084f5266a0d92003000001011fe65c4001000000001600140a70b8068c6c24946a9e41c23e5bb4084f5266a001086b0247304402207b7b53b01ee71512444ebf22af6e0ef45a9c897e3ddb7e895f70815a738fc38502206a658fd7f358c09a46e51f326871191fb27df1e41186bfbff486a441269c659d012103634d38e0306a1623db98da1e229eed182b0c61811bdcc26cba5c8373a5df844100000000",
+};
 
-    (btcWallet.getWalletProviderName as jest.Mock).mockResolvedValue(
-      dataGenerator.generateRandomString(10),
+describe("signPsbtTransaction", () => {
+  it("should handle compatible wallets correctly", async () => {
+    const mockProviderName = "OKX";
+
+    const wallet = mockWalletProvider(
+      mockProviderName,
+      mocked_signedPsbtHexPair.signedPsbtHex,
+    ) as any;
+    const signTransaction = signPsbtTransaction(wallet);
+
+    const result = await signTransaction(mocked_signedPsbtHexPair.psbtHex);
+
+    expect(wallet.getWalletProviderName).toHaveBeenCalled();
+    expect(result.toHex()).toBe(
+      Psbt.fromHex(mocked_signedPsbtHexPair.signedPsbtHex)
+        .extractTransaction()
+        .toHex(),
     );
-
-    (btcWallet.signPsbt as jest.Mock).mockResolvedValue(unsignedPsbtHex);
-
-    const signTransaction = signPsbtTransaction(btcWallet);
-
-    const result = await signTransaction(unsignedPsbtHex);
-
-    expect(btcWallet.signPsbt).toHaveBeenCalledWith(unsignedPsbtHex);
-
-    expect(btcWallet.getWalletProviderName).toHaveBeenCalled();
-
-    expect(result).toBeInstanceOf(Transaction);
   });
 
-  it("should return a signed transaction for non-compatible wallets", async () => {
-    const randomGlobalParamsVersions =
-      dataGenerator.generateGlobalPramsVersions(
-        dataGenerator.getRandomIntegerBetween(1, 20),
-      );
-    const stakerKeys = dataGenerator.generateRandomKeyPair();
-    const randomStakingTxHeight =
-      randomGlobalParamsVersions[
-        dataGenerator.getRandomIntegerBetween(
-          0,
-          randomGlobalParamsVersions.length - 1,
-        )
-      ].activationHeight + 1;
-
-    const { signedPsbt, unsignedPsbt } = dataGenerator.createRandomStakingPsbt(
-      randomGlobalParamsVersions,
-      randomStakingTxHeight,
-      stakerKeys,
-    );
+  it("should handle non-compatible wallets with new implementation", async () => {
+    const mockProviderName = "OneKey";
 
-    const signedPsbtHex = signedPsbt.extractTransaction().toHex();
-    const unsignedPsbtHex = unsignedPsbt.extractTransaction().toHex();
+    const wallet = mockWalletProvider(
+      mockProviderName,
+      mocked_signedPsbtHexPair.signedPsbtHex,
+    ) as any;
+    const signTransaction = signPsbtTransaction(wallet);
 
-    (btcWallet.getWalletProviderName as jest.Mock).mockResolvedValue("OneKey");
+    const result = await signTransaction(mocked_signedPsbtHexPair.psbtHex);
 
-    (btcWallet.signPsbt as jest.Mock).mockResolvedValue(signedPsbtHex);
-
-    const signTransaction = signPsbtTransaction(btcWallet);
+    expect(wallet.getWalletProviderName).toHaveBeenCalled();
+    expect(result.toHex()).toBe(
+      Psbt.fromHex(mocked_signedPsbtHexPair.signedPsbtHex)
+        .extractTransaction()
+        .toHex(),
+    );
+  });
 
-    const result = await signTransaction(unsignedPsbtHex);
+  it("should handle non-compatible wallets with old implementation", async () => {
+    const mockProviderName = "OneKey";
 
-    expect(btcWallet.signPsbt).toHaveBeenCalledWith(unsignedPsbtHex);
+    const wallet = mockWalletProvider(
+      mockProviderName,
+      mocked_signedTxHexPair.signedTxHex,
+    ) as any;
+    const signTransaction = signPsbtTransaction(wallet);
 
-    expect(btcWallet.getWalletProviderName).toHaveBeenCalled();
+    const result = await signTransaction(mocked_signedTxHexPair.psbtHex);
 
-    expect(result).toBeInstanceOf(Transaction);
+    expect(wallet.getWalletProviderName).toHaveBeenCalled();
+    expect(result.toHex()).toBe(
+      Transaction.fromHex(mocked_signedTxHexPair.signedTxHex).toHex(),
+    );
   });
 });