fix: withdraw slashed staking tx (#51)

* feat: withdraw slashed staking tx

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@babylonlabs-io/btc-staking-ts",
-  "version": "0.4.0-canary.5",
+  "version": "0.4.0-canary.6",
   "description": "Library exposing methods for the creation and consumption of Bitcoin transactions pertaining to Babylon's Bitcoin Staking protocol.",
   "module": "dist/index.js",
   "main": "dist/index.cjs",

src/staking/index.ts

@@ -8,6 +8,7 @@ import {
   slashTimelockUnbondedTransaction,
   stakingTransaction, unbondingTransaction,
   withdrawEarlyUnbondedTransaction,
+  withdrawSlashingTransaction,
   withdrawTimelockUnbondedTransaction
 } from "./transactions";
 import {
@@ -16,7 +17,8 @@ import {
 } from "../utils/btc";
 import { 
   deriveStakingOutputAddress,
-  findMatchingStakingTxOutputIndex,
+  deriveSlashingOutputAddress,
+  findMatchingTxOutputIndex,
   validateParams,
   validateStakingTimelock,
   validateStakingTxInputData,
@@ -206,7 +208,7 @@ export class Staking {
     const scripts = this.buildScripts();
 
     // Reconstruct the stakingOutputIndex
-    const stakingOutputIndex = findMatchingStakingTxOutputIndex(
+    const stakingOutputIndex = findMatchingTxOutputIndex(
       stakingTx,
       deriveStakingOutputAddress(scripts, this.network),
       this.network,
@@ -315,7 +317,7 @@ export class Staking {
     const scripts = this.buildScripts();
 
     // Reconstruct the stakingOutputIndex
-    const stakingOutputIndex = findMatchingStakingTxOutputIndex(
+    const stakingOutputIndex = findMatchingTxOutputIndex(
       stakingTx,
       deriveStakingOutputAddress(scripts, this.network),
       this.network,
@@ -421,4 +423,45 @@ export class Staking {
       );
     }
   }
+
+  /**
+   * Create a withdraw transaction that spends a slashing transaction from the
+   * staking output.
+   * 
+   * @param {Transaction} slashingTx - The slashing transaction.
+   * @param {number} feeRate - The fee rate for the transaction in satoshis per byte.
+   * @returns {PsbtResult} - An object containing the unsigned psbt and fee
+   * @throws {StakingError} - If the delegation is invalid or the transaction cannot be built
+   */
+  public createWithdrawSlashingTransaction(
+    slashingTx: Transaction,
+    feeRate: number,
+  ): PsbtResult {
+    // Build scripts
+    const scripts = this.buildScripts();
+
+    // Reconstruct and validate the slashingOutputIndex
+    const slashingOutputIndex = findMatchingTxOutputIndex(
+      slashingTx,
+      deriveSlashingOutputAddress(scripts, this.network),
+      this.network,
+    )
+
+    // Create the withdraw slashed transaction
+    try {
+      return withdrawSlashingTransaction(
+        scripts,
+        slashingTx,
+        this.stakerInfo.address,
+        this.network,
+        feeRate,
+        slashingOutputIndex,
+      );  
+    } catch (error) {
+      throw StakingError.fromUnknown(
+        error, StakingErrorCode.BUILD_TRANSACTION_FAILURE,
+        "Cannot build withdraw slashing transaction",
+      );
+    }
+  }
 }

src/staking/transactions.ts

@@ -236,6 +236,52 @@ export function withdrawTimelockUnbondedTransaction(
   );
 }
 
+/**
+ * Constructs a withdrawal transaction for a slashing transaction.
+ *
+ * This transaction spends the output from the slashing transaction.
+ *
+ * @param {Object} scripts - The unbondingTimelockScript and slashingScript
+ * We use the unbonding timelock script as the timelock of the slashing transaction.
+ * This is due to slashing tx timelock is the same as the unbonding timelock.
+ * @param {Transaction} slashingTx - The slashing transaction.
+ * @param {string} withdrawalAddress - The address to send the withdrawn funds to.
+ * @param {networks.Network} network - The Bitcoin network.
+ * @param {number} feeRate - The fee rate for the transaction in satoshis per byte.
+ * @param {number} outputIndex - The index of the output to be spent in the original transaction.
+ * @returns {PsbtResult} An object containing the partially signed transaction (PSBT).
+ */
+export function withdrawSlashingTransaction(
+  scripts: {
+    unbondingTimelockScript: Buffer;
+    slashingScript: Buffer;
+  },
+  slashingTx: Transaction,
+  withdrawalAddress: string,
+  network: networks.Network,
+  feeRate: number,
+  outputIndex: number,
+): PsbtResult {
+  const scriptTree: Taptree = [
+    {
+      output: scripts.slashingScript,
+    },
+    { output: scripts.unbondingTimelockScript },
+  ];
+
+  return withdrawalTransaction(
+    {
+      timelockScript: scripts.unbondingTimelockScript,
+    },
+    scriptTree,
+    slashingTx,
+    withdrawalAddress,
+    network,
+    feeRate,
+    outputIndex,
+  );
+}
+
 // withdrawalTransaction generates a transaction that
 // spends the staking output of the staking transaction
 function withdrawalTransaction(

src/utils/staking/index.ts

@@ -88,25 +88,57 @@ export const deriveStakingOutputAddress = (
 };
 
 /**
- * Find the matching output index for the given staking transaction.
+ * Derive the slashing output address from the staking scripts.
  * 
- * @param {Transaction} stakingTx - The staking transaction.
- * @param {string} stakingOutputAddress - The staking output address.
+ * @param {StakingScripts} scripts - The unbonding timelock scripts, we use the
+ * unbonding timelock script as the timelock of the slashing transaction.
+ * This is due to slashing tx timelock is the same as the unbonding timelock.
+ * @param {networks.Network} network - The Bitcoin network.
+ * @returns {string} - The slashing output address.
+ * @throws {StakingError} - If the slashing output address cannot be derived.
+ */
+export const deriveSlashingOutputAddress = (
+  scripts: {
+    unbondingTimelockScript: Buffer;
+  },
+  network: networks.Network,
+) => {
+  const slashingOutput = payments.p2tr({
+    internalPubkey,
+    scriptTree: { output: scripts.unbondingTimelockScript },
+    network,
+  });
+
+  if (!slashingOutput.address) {
+    throw new StakingError(
+      StakingErrorCode.INVALID_OUTPUT,
+      "Failed to build slashing output address",
+    );
+  }
+
+  return slashingOutput.address;
+}
+
+/**
+ * Find the matching output index for the given transaction.
+ * 
+ * @param {Transaction} tx - The transaction.
+ * @param {string} outputAddress - The output address.
  * @param {networks.Network} network - The Bitcoin network.
  * @returns {number} - The output index.
  * @throws {Error} - If the matching output is not found.
  */
-export const findMatchingStakingTxOutputIndex = (
-  stakingTx: Transaction,
-  stakingOutputAddress: string,
+export const findMatchingTxOutputIndex = (
+  tx: Transaction,
+  outputAddress: string,
   network: networks.Network,
 ) => {
-  const index = stakingTx.outs.findIndex(output => {
-    return address.fromOutputScript(output.script, network) === stakingOutputAddress;
+  const index = tx.outs.findIndex(output => {
+    return address.fromOutputScript(output.script, network) === outputAddress;
   });
 
   if (index === -1) {
-    throw new Error(`Matching output not found for address: ${stakingOutputAddress}`);
+    throw new Error(`Matching output not found for address: ${outputAddress}`);
   }
 
   return index;

tests/helper/datagen/base.ts

@@ -1,12 +1,16 @@
 import * as ecc from "@bitcoin-js/tiny-secp256k1-asmjs";
 import * as bitcoin from "bitcoinjs-lib";
 import ECPairFactory from "ecpair";
-import { Staking, stakingTransaction } from "../../../src";
+import { slashEarlyUnbondedTransaction, slashTimelockUnbondedTransaction, Staking, stakingTransaction, TransactionResult, unbondingTransaction } from "../../../src";
 import { UTXO } from "../../../src/types/UTXO";
 import { StakingParams } from "../../../src/types/params";
 import { generateRandomAmountSlices } from "../math";
 import { StakingScriptData, StakingScripts } from "../../../src/index";
 import { MIN_UNBONDING_OUTPUT_VALUE } from "../../../src/constants/unbonding";
+import { payments, Psbt, Transaction } from "bitcoinjs-lib";
+import { TRANSACTION_VERSION } from "../../../src/constants/psbt";
+import { NON_RBF_SEQUENCE } from "../../../src/constants/psbt";
+import { internalPubkey } from "../../../src/constants/internalPubkey";
 
 bitcoin.initEccLib(ecc);
 const ECPair = ECPairFactory(ecc);
@@ -20,6 +24,8 @@ export interface KeyPair {
   keyPair: bitcoin.Signer;
 }
 
+export type SlashingType = "earlyUnbonded" | "timelockExpire";
+
 export class StakingDataGenerator {
   network: bitcoin.networks.Network;
 
@@ -287,6 +293,102 @@ export class StakingDataGenerator {
     }
   };
 
+  /**
+   * Generates a random slashing transaction based on the staking transaction 
+   * and staking scripts
+   * @param network - The network to use
+   * @param stakingScripts - The staking scripts to use
+   * @param stakingTx - The staking transaction to use
+   * @param param - The param used in the staking transaction
+   * @param keyPair - The key pair to use. This is used to sign the slashing 
+   * psbt to derive the transaction.
+   * @param type - The type of slashing to use.
+   * @returns {Object} - A random slashing transaction
+   */
+  generateSlashingTransaction = (
+    network: bitcoin.networks.Network,
+    stakingScripts: StakingScripts,
+    stakingTx: Transaction,
+    param: {
+      minSlashingTxFeeSat: number,
+      slashingPkScriptHex: string,
+      slashingRate: number,
+    },
+    keyPair: KeyPair,
+    type: SlashingType = "timelockExpire",
+  ) => {
+    let slashingPsbt: Psbt;
+    let outputValue: number;
+
+    if (type === "earlyUnbonded") {
+      const { transaction: unbondingTx } = unbondingTransaction(
+        stakingScripts,
+        stakingTx,
+        1,
+        network,
+      );
+      const { psbt } = slashEarlyUnbondedTransaction(
+        stakingScripts,
+        unbondingTx,
+        param.slashingPkScriptHex,
+        param.slashingRate,
+        param.minSlashingTxFeeSat,
+        network,
+      );
+      slashingPsbt = psbt;
+      outputValue = unbondingTx.outs[0].value;
+    } else {
+      const { psbt } = slashTimelockUnbondedTransaction(
+        stakingScripts,
+        stakingTx,
+        param.slashingPkScriptHex,
+        param.slashingRate,
+        param.minSlashingTxFeeSat,
+        network,
+      );
+      slashingPsbt = psbt;
+      outputValue = stakingTx.outs[0].value;
+    }
+
+    expect(slashingPsbt).toBeDefined();
+    expect(slashingPsbt.txOutputs.length).toBe(2);
+    // first output shall send slashed amount to the slashing pk script (i.e burn output)
+    expect(Buffer.from(slashingPsbt.txOutputs[0].script).toString("hex")).toBe(
+      param.slashingPkScriptHex,
+    );
+    expect(slashingPsbt.txOutputs[0].value).toBe(
+      Math.floor(outputValue * param.slashingRate),
+    );
+
+     // second output is the change output which send to unbonding timelock script address
+     const changeOutput = payments.p2tr({
+      internalPubkey,
+      scriptTree: { output: stakingScripts.unbondingTimelockScript },
+      network,
+    });
+    expect(slashingPsbt.txOutputs[1].address).toBe(changeOutput.address);
+    const expectedChangeOutputValue =
+      outputValue -
+      Math.floor(outputValue * param.slashingRate) -
+      param.minSlashingTxFeeSat;
+    expect(slashingPsbt.txOutputs[1].value).toBe(expectedChangeOutputValue);
+
+    expect(slashingPsbt.version).toBe(TRANSACTION_VERSION);
+    expect(slashingPsbt.locktime).toBe(0);
+    slashingPsbt.txInputs.forEach((input) => {
+      expect(input.sequence).toBe(NON_RBF_SEQUENCE);
+    });
+
+    const tx = slashingPsbt.signAllInputs(
+      keyPair.keyPair,
+    ).finalizeAllInputs().extractTransaction();
+
+    return {
+      psbt: slashingPsbt,
+      tx,
+    };
+  }
+
   randomBoolean(): boolean {
     return Math.random() >= 0.5;
   };

tests/staking/createSlashingTx.test.ts

@@ -1,7 +1,6 @@
 import * as stakingScript from "../../src/staking/stakingScript";
 import { testingNetworks } from "../helper";
 import * as transaction from "../../src/staking/transactions";
-import { Staking } from "../../src/staking";
 import { opcodes, payments, script } from "bitcoinjs-lib";
 import { internalPubkey } from "../../src/constants/internalPubkey";
 
@@ -89,10 +88,6 @@ describe.each(testingNetworks)("Create slashing transactions", ({
       jest.spyOn(stakingScript, "StakingScriptData").mockImplementation(() => {
         throw new Error("slash timelock unbonded delegation build script error");
       });
-      const staking = new Staking(
-        network, stakerInfo,
-        params, finalityProviderPkNoCoordHex, timelock,
-      );
 
       expect(() => stakingInstance.createStakingOutputSlashingTransaction(
         stakingTx,

tests/staking/createUnbondingtx.test.ts

@@ -5,7 +5,7 @@ import { StakingError, StakingErrorCode } from "../../src/error";
 import { testingNetworks } from "../helper";
 import { NON_RBF_SEQUENCE } from "../../src/constants/psbt";
 import * as stakingScript from "../../src/staking/stakingScript";
-import { deriveStakingOutputAddress, findMatchingStakingTxOutputIndex } from "../../src/utils/staking";
+import { deriveStakingOutputAddress, findMatchingTxOutputIndex } from "../../src/utils/staking";
 
 describe.each(testingNetworks)("Create unbonding transaction", ({
   network, networkName, datagen: { stakingDatagen : dataGenerator }
@@ -66,7 +66,7 @@ describe.each(testingNetworks)("Create unbonding transaction", ({
     expect(psbt.locktime).toBe(0);
 
     // Get staking output index
-    const stakingOutputIndex = findMatchingStakingTxOutputIndex(
+    const stakingOutputIndex = findMatchingTxOutputIndex(
       stakingTx,
       deriveStakingOutputAddress(scripts, network),
       network,