Merge pull request #5822 from BitGo/BTC-1933.createDelegationMessageFromPsbt

OttoAllmendinger · web-flow · commit 2ba014510dbf · 2025-03-25T14:46:57.000+01:00
feat(utxo-staking): implement Babylon staking utilities
diff --git a/modules/utxo-staking/src/babylon/delegationMessage.ts b/modules/utxo-staking/src/babylon/delegationMessage.ts
@@ -1,15 +1,20 @@
+/**
+ * https://github.com/babylonlabs-io/babylon/blob/v1.99.0-snapshot.250211/x/btcstaking/types/validate_parsed_message.go
+ */
 import assert from 'assert';
 
 import * as vendor from '@bitgo/babylonlabs-io-btc-staking-ts';
 import * as babylonProtobuf from '@babylonlabs-io/babylon-proto-ts';
 import * as bitcoinjslib from 'bitcoinjs-lib';
 import * as utxolib from '@bitgo/utxo-lib';
 import { Descriptor } from '@bitgo/wasm-miniscript';
+import { toXOnlyPublicKey } from '@bitgo/utxo-core';
 import { toWrappedPsbt } from '@bitgo/utxo-core/descriptor';
 
 import { BabylonDescriptorBuilder } from './descriptor';
 import { createStakingManager } from './stakingManager';
 import { getStakingParams } from './stakingParams';
+import { BabylonNetworkLike, toBitcoinJsNetwork } from './network';
 
 export type ValueWithTypeUrl<T> = { typeUrl: string; value: T };
 
@@ -83,60 +88,106 @@ type Result = {
   stakingTx: bitcoinjslib.Transaction;
 };
 
-/*
- * This is mostly lifted from
- * https://github.com/babylonlabs-io/btc-staking-ts/blob/v0.4.0-rc.2/src/staking/manager.ts#L100-L172
- *
- * The difference is that here we are returning an _unsigned_ delegation message.
+/**
+ * @param stakingKey - this is the single-sig key that is used for co-signing the staking output
+ * @param changeAddress - this is unrelated to the staking key and is used for the change output
  */
-export async function createUnsignedPreStakeRegistrationBabylonTransaction(
-  manager: vendor.BabylonBtcStakingManager,
-  stakingParams: vendor.VersionedStakingParams[],
-  network: bitcoinjslib.Network,
+export function toStakerInfo(
+  stakingKey: utxolib.ECPairInterface | Buffer | string,
+  changeAddress: string
+): vendor.StakerInfo {
+  if (typeof stakingKey === 'object' && 'publicKey' in stakingKey) {
+    stakingKey = stakingKey.publicKey;
+  }
+  if (typeof stakingKey === 'string') {
+    stakingKey = Buffer.from(stakingKey, 'hex');
+  }
+  return {
+    publicKeyNoCoordHex: toXOnlyPublicKey(stakingKey).toString('hex'),
+    address: changeAddress,
+  };
+}
+
+export function createStaking(
+  network: BabylonNetworkLike,
+  blockHeight: number,
   stakerBtcInfo: vendor.StakerInfo,
   stakingInput: vendor.StakingInputs,
-  babylonBtcTipHeight: number,
-  inputUTXOs: vendor.UTXO[],
-  feeRateSatB: number,
-  babylonAddress: string
-): Promise<Result> {
-  if (babylonBtcTipHeight === 0) {
+  versionedParams: vendor.VersionedStakingParams[] = getStakingParams(network)
+): vendor.Staking {
+  if (blockHeight === 0) {
     throw new Error('Babylon BTC tip height cannot be 0');
   }
-  if (inputUTXOs.length === 0) {
-    throw new Error('No input UTXOs provided');
-  }
-  if (!vendor.isValidBabylonAddress(babylonAddress)) {
-    throw new Error('Invalid Babylon address');
-  }
 
   // Get the Babylon params based on the BTC tip height from Babylon chain
-  const params = vendor.getBabylonParamByBtcHeight(babylonBtcTipHeight, stakingParams);
+  const params = vendor.getBabylonParamByBtcHeight(blockHeight, versionedParams);
 
-  const staking = new vendor.Staking(
-    network,
+  return new vendor.Staking(
+    toBitcoinJsNetwork(network),
     stakerBtcInfo,
     params,
     stakingInput.finalityProviderPkNoCoordHex,
     stakingInput.stakingTimelock
   );
+}
 
-  // Create unsigned staking transaction
-  const { transaction } = staking.createStakingTransaction(stakingInput.stakingAmountSat, inputUTXOs, feeRateSatB);
+type TransactionLike =
+  | bitcoinjslib.Psbt
+  | bitcoinjslib.Transaction
+  | utxolib.Transaction
+  | utxolib.bitgo.UtxoTransaction<bigint | number>
+  | utxolib.Psbt
+  | utxolib.bitgo.UtxoPsbt;
 
+function toStakingTransactionFromPsbt(
+  psbt: bitcoinjslib.Psbt | utxolib.Psbt | utxolib.bitgo.UtxoPsbt
+): bitcoinjslib.Transaction {
+  if (!(psbt instanceof utxolib.bitgo.UtxoPsbt)) {
+    psbt = utxolib.bitgo.createPsbtFromBuffer(psbt.toBuffer(), utxolib.networks.bitcoin);
+  }
+  if (psbt instanceof utxolib.bitgo.UtxoPsbt) {
+    // only utxolib.bitgo.UtxoPsbt has the getUnsignedTx method
+    return bitcoinjslib.Transaction.fromHex(psbt.getUnsignedTx().toHex());
+  }
+  throw new Error('illegal state');
+}
+
+export function toStakingTransaction(tx: TransactionLike): bitcoinjslib.Transaction {
+  if (tx instanceof bitcoinjslib.Psbt || tx instanceof utxolib.Psbt) {
+    return toStakingTransactionFromPsbt(tx);
+  }
+  return bitcoinjslib.Transaction.fromHex(tx.toHex());
+}
+
+/*
+ * This is mostly lifted from
+ * https://github.com/babylonlabs-io/btc-staking-ts/blob/v0.4.0-rc.2/src/staking/manager.ts#L100-L172
+ *
+ * The difference is that here we are returning an _unsigned_ delegation message.
+ */
+export async function createDelegationMessageWithTransaction(
+  manager: vendor.BabylonBtcStakingManager,
+  staking: vendor.Staking,
+  stakingAmountSat: number,
+  transaction: TransactionLike,
+  babylonAddress: string
+): Promise<ValueWithTypeUrl<babylonProtobuf.btcstakingtx.MsgCreateBTCDelegation>> {
+  if (!vendor.isValidBabylonAddress(babylonAddress)) {
+    throw new Error('Invalid Babylon address');
+  }
   // Create delegation message without including inclusion proof
-  const msg = await manager.createBtcDelegationMsg(
+  return manager.createBtcDelegationMsg(
     staking,
-    stakingInput,
-    transaction,
+    {
+      stakingTimelock: staking.stakingTimelock,
+      finalityProviderPkNoCoordHex: staking.finalityProviderPkNoCoordHex,
+      stakingAmountSat,
+    },
+    toStakingTransaction(transaction),
     babylonAddress,
-    stakerBtcInfo,
-    params
+    staking.stakerInfo,
+    staking.params
   );
-  return {
-    unsignedDelegationMsg: msg,
-    stakingTx: transaction,
-  };
 }
 
 export async function createUnsignedPreStakeRegistrationBabylonTransactionWithBtcProvider(
@@ -150,16 +201,19 @@ export async function createUnsignedPreStakeRegistrationBabylonTransactionWithBt
   babylonAddress: string,
   stakingParams: vendor.VersionedStakingParams[] = getStakingParams(network)
 ): Promise<Result> {
+  if (inputUTXOs.length === 0) {
+    throw new Error('No input UTXOs provided');
+  }
   const manager = createStakingManager(network, btcProvider, stakingParams);
-  return await createUnsignedPreStakeRegistrationBabylonTransaction(
+  const staking = createStaking(network, babylonBtcTipHeight, stakerBtcInfo, stakingInput, stakingParams);
+  // Create unsigned staking transaction
+  const { transaction } = staking.createStakingTransaction(stakingInput.stakingAmountSat, inputUTXOs, feeRateSatB);
+  const unsignedDelegationMsg = await createDelegationMessageWithTransaction(
     manager,
-    stakingParams,
-    network,
-    stakerBtcInfo,
-    stakingInput,
-    babylonBtcTipHeight,
-    inputUTXOs,
-    feeRateSatB,
+    staking,
+    stakingInput.stakingAmountSat,
+    transaction,
     babylonAddress
   );
+  return { unsignedDelegationMsg, stakingTx: transaction };
 }
diff --git a/modules/utxo-staking/src/babylon/descriptor.ts b/modules/utxo-staking/src/babylon/descriptor.ts
@@ -1,6 +1,7 @@
 /**
  * https://github.com/babylonlabs-io/babylon/tree/main/docs
  * https://github.com/babylonlabs-io/babylon/blob/main/docs/staking-script.md
+ * https://github.com/babylonlabs-io/babylon/blob/v1.99.0-snapshot.250211/btcstaking/staking.go
  */
 
 import { Descriptor, ast } from '@bitgo/wasm-miniscript';
@@ -44,21 +45,6 @@ export class BabylonDescriptorBuilder {
       finalityProviderKeys: Buffer[];
     } & StakingParams
   ): BabylonDescriptorBuilder {
-    /*
-
-  const stakerKey = getECKey('staker');
-  const covenantThreshold = stakingParams.covenantQuorum;
-  const stakingTimelock = stakingParams.minStakingTimeBlocks;
-  const unbondingTimelock = stakingParams.unbondingTime;
-  const vendorBuilder = new vendor.StakingScriptData(
-    getXOnlyPubkey(stakerKey),
-    finalityProviderKeys.map(getXOnlyPubkey),
-    covenantKeys.map(getXOnlyPubkey),
-    covenantThreshold,
-    stakingTimelock,
-    unbondingTimelock
-  );
-     */
     return new BabylonDescriptorBuilder(
       params.stakerKey,
       params.finalityProviderKeys,
diff --git a/modules/utxo-staking/src/babylon/index.ts b/modules/utxo-staking/src/babylon/index.ts
@@ -1,3 +1,10 @@
+import { initBTCCurve } from '@bitgo/babylonlabs-io-btc-staking-ts';
+import * as bitcoinjslib from 'bitcoinjs-lib';
+import * as utxolib from '@bitgo/utxo-lib';
+
+initBTCCurve();
+bitcoinjslib.initEccLib(utxolib.ecc);
+
 export * from './delegationMessage';
 export * from './descriptor';
 export * from './stakingParams';
diff --git a/modules/utxo-staking/src/babylon/keyUtils.ts b/modules/utxo-staking/src/babylon/keyUtils.ts
diff --git a/modules/utxo-staking/src/babylon/network.ts b/modules/utxo-staking/src/babylon/network.ts
@@ -0,0 +1,34 @@
+import * as bitcoinjslib from 'bitcoinjs-lib';
+import * as utxolib from '@bitgo/utxo-lib';
+
+export type BabylonNetwork = 'mainnet' | 'testnet';
+
+export type BabylonNetworkLike = bitcoinjslib.Network | utxolib.Network | BabylonNetwork;
+
+export function toBabylonNetwork(n: BabylonNetworkLike): BabylonNetwork {
+  switch (n) {
+    case bitcoinjslib.networks.bitcoin:
+    case utxolib.networks.bitcoin:
+      return 'mainnet';
+    case bitcoinjslib.networks.testnet:
+    case utxolib.networks.testnet:
+    case utxolib.networks.bitcoinPublicSignet:
+      return 'testnet';
+    case 'mainnet':
+    case 'testnet':
+      return n;
+    default:
+      throw new Error('Unsupported network');
+  }
+}
+
+export function toBitcoinJsNetwork(n: BabylonNetworkLike): bitcoinjslib.Network {
+  switch (toBabylonNetwork(n)) {
+    case 'mainnet':
+      return bitcoinjslib.networks.bitcoin;
+    case 'testnet':
+      return bitcoinjslib.networks.testnet;
+    default:
+      throw new Error('Unsupported network');
+  }
+}
diff --git a/modules/utxo-staking/src/babylon/stakingParams.ts b/modules/utxo-staking/src/babylon/stakingParams.ts
@@ -2,18 +2,20 @@ import * as t from 'io-ts';
 import * as tt from 'io-ts-types';
 import { isLeft } from 'fp-ts/Either';
 import { PathReporter } from 'io-ts/lib/PathReporter';
-import * as bitcoinjslib from 'bitcoinjs-lib';
 import * as utxolib from '@bitgo/utxo-lib';
 import {
   getBabylonParamByVersion,
   StakerInfo,
   StakingInputs,
+  StakingParams,
   VersionedStakingParams,
 } from '@bitgo/babylonlabs-io-btc-staking-ts';
+export { getBabylonParamByVersion, getBabylonParamByBtcHeight } from '@bitgo/babylonlabs-io-btc-staking-ts';
 
 import { BabylonDescriptorBuilder } from './descriptor';
 import jsonMainnetParams from './params.mainnet.json';
 import jsonTestnetParams from './params.testnet.json';
+import { BabylonNetworkLike, toBabylonNetwork } from './network';
 
 const BabylonParamsJSON = t.type({
   covenant_pks: t.array(t.string),
@@ -68,27 +70,6 @@ function toVersionedParamsFromJson(jsonParams: unknown[]): VersionedStakingParam
   );
 }
 
-type BabylonNetwork = 'mainnet' | 'testnet';
-
-type BabylonNetworkLike = bitcoinjslib.Network | utxolib.Network | BabylonNetwork;
-
-function toBabylonNetwork(n: BabylonNetworkLike): BabylonNetwork {
-  switch (n) {
-    case bitcoinjslib.networks.bitcoin:
-    case utxolib.networks.bitcoin:
-      return 'mainnet';
-    case bitcoinjslib.networks.testnet:
-    case utxolib.networks.testnet:
-    case utxolib.networks.bitcoinPublicSignet:
-      return 'testnet';
-    case 'mainnet':
-    case 'testnet':
-      return n;
-    default:
-      throw new Error('Unsupported network');
-  }
-}
-
 export const mainnetStakingParams: readonly VersionedStakingParams[] = Object.freeze(
   toVersionedParamsFromJson(jsonMainnetParams)
 );
@@ -114,13 +95,16 @@ export const testnetFinalityProvider0 = Buffer.from(
   'hex'
 );
 
+type DescriptorStakingParams = Pick<
+  StakingParams,
+  'covenantNoCoordPks' | 'covenantQuorum' | 'minStakingTimeBlocks' | 'unbondingTime'
+>;
+
 export function getDescriptorBuilderForParams(
   userKey: utxolib.BIP32Interface | utxolib.ECPairInterface | Buffer,
   finalityProviderKeys: Buffer[],
-  params: Pick<
-    VersionedStakingParams,
-    'covenantNoCoordPks' | 'covenantQuorum' | 'minStakingTimeBlocks' | 'unbondingTime'
-  >
+  stakingTimelock: number,
+  params: DescriptorStakingParams
 ): BabylonDescriptorBuilder {
   if (!Buffer.isBuffer(userKey)) {
     userKey = userKey.publicKey;
@@ -130,32 +114,32 @@ export function getDescriptorBuilderForParams(
     finalityProviderKeys,
     params.covenantNoCoordPks.map((pk) => Buffer.from(pk, 'hex')),
     params.covenantQuorum,
-    params.minStakingTimeBlocks,
+    stakingTimelock,
     params.unbondingTime
   );
 }
 
 export function getDescriptorProviderForStakingParams(
-  stakerBtcInfo: StakerInfo,
+  stakerBtcInfo: Pick<StakerInfo, 'publicKeyNoCoordHex'>,
   stakingInput: StakingInputs,
-  stakingParams: VersionedStakingParams
+  stakingParams: DescriptorStakingParams
 ): BabylonDescriptorBuilder {
   const userKey = Buffer.from(stakerBtcInfo.publicKeyNoCoordHex, 'hex');
   const finalityProviderKey = Buffer.from(stakingInput.finalityProviderPkNoCoordHex, 'hex');
-  return getDescriptorBuilderForParams(userKey, [finalityProviderKey], stakingParams);
+  return getDescriptorBuilderForParams(userKey, [finalityProviderKey], stakingInput.stakingTimelock, stakingParams);
 }
 
 export function getTestnetDescriptorBuilder(
   userKey: utxolib.BIP32Interface | utxolib.ECPairInterface | Buffer,
   {
     finalityProviderKeys = [testnetFinalityProvider0],
+    params = getBabylonParamByVersion(5, getStakingParams('testnet')),
+    stakingTimelock = params.minStakingTimeBlocks,
   }: {
     finalityProviderKeys?: Buffer[];
+    params?: StakingParams;
+    stakingTimelock?: number;
   } = {}
 ): BabylonDescriptorBuilder {
-  return getDescriptorBuilderForParams(
-    userKey,
-    finalityProviderKeys,
-    getBabylonParamByVersion(5, getStakingParams('testnet'))
-  );
+  return getDescriptorBuilderForParams(userKey, finalityProviderKeys, stakingTimelock, params);
 }
diff --git a/modules/utxo-staking/src/transaction.ts b/modules/utxo-staking/src/transaction.ts
@@ -1,5 +1,8 @@
 import * as utxolib from '@bitgo/utxo-lib';
 import { Dimensions } from '@bitgo/unspents';
+import * as bitcoinjslib from 'bitcoinjs-lib';
+
+bitcoinjslib.initEccLib(utxolib.ecc);
 
 /**
  * Build a staking transaction for a wallet that assumes 2-of-3 multisig for the inputs
diff --git a/modules/utxo-staking/test/unit/babylon/transactions.ts b/modules/utxo-staking/test/unit/babylon/transactions.ts
