fix(abstract-utxo): refactor backup key recovery psbt creation

Extract PSBT creation logic to separate functions for better maintainability
and testability. Improves error handling for insufficient funds and adds
proper type annotations for fee parameters.

Issue: BTC-2891

Co-authored-by: llm-git <[email protected]>

Author: OttoAllmendinger

modules/abstract-utxo/src/recovery/backupKeyRecovery.ts

@@ -1,6 +1,5 @@
 import _ from 'lodash';
 import * as utxolib from '@bitgo/utxo-lib';
-import { Dimensions } from '@bitgo/unspents';
 import {
   BitGoBase,
   ErrorNoInputToRecover,
@@ -20,6 +19,7 @@ import { generateAddressWithChainAndIndex } from '../address';
 import { forCoin, RecoveryProvider } from './RecoveryProvider';
 import { MempoolApi } from './mempoolApi';
 import { CoingeckoApi } from './coingeckoApi';
+import { createBackupKeyRecoveryPsbt, getRecoveryAmount } from './psbt';
 
 type ScriptType2Of3 = utxolib.bitgo.outputScripts.ScriptType2Of3;
 type ChainCode = utxolib.bitgo.ChainCode;
@@ -100,6 +100,7 @@ export interface RecoverParams {
   apiKey?: string;
   userKeyPath?: string;
   recoveryProvider?: RecoveryProvider;
+  /** Satoshi per byte */
   feeRate?: number;
 }
 
@@ -331,10 +332,6 @@ export async function backupKeyRecovery(
     throw new ErrorNoInputToRecover();
   }
 
-  // Build the psbt
-  const psbt = utxolib.bitgo.createPsbtForNetwork({ network: coin.network });
-  // xpubs can become handy for many things.
-  utxolib.bitgo.addXpubsToPsbt(psbt, walletKeys);
   const txInfo = {} as BackupKeyRecoveryTransansaction;
   const feePerByte: number =
     params.feeRate !== undefined
@@ -346,10 +343,6 @@ export async function backupKeyRecovery(
       ? unspents.map((u) => ({ ...u, value: Number(u.value), valueString: u.value.toString(), prevTx: undefined }))
       : undefined;
 
-  unspents.forEach((unspent) => {
-    utxolib.bitgo.addWalletUnspentToPsbt(psbt, unspent, walletKeys, 'user', 'backup');
-  });
-
   let krsFee = BigInt(0);
   if (isKrsRecovery && params.krsProvider) {
     try {
@@ -360,40 +353,26 @@ export async function backupKeyRecovery(
     }
   }
 
-  const dimensions = Dimensions.fromPsbt(psbt)
-    // Add the recovery output
-    .plus(Dimensions.fromOutput({ script: utxolib.address.toOutputScript(params.recoveryDestination, coin.network) }))
-    // KRS recovery transactions have a 2nd output to pay the recovery fee, like paygo fees.
-    .plus(krsFee > BigInt(0) ? Dimensions.SingleOutput.p2wsh : Dimensions.ZERO);
-  const approximateFee = BigInt(dimensions.getVSize() * feePerByte);
-
-  const recoveryAmount = totalInputAmount - approximateFee - krsFee;
-
-  if (recoveryAmount < BigInt(0)) {
-    throw new Error(`this wallet\'s balance is too low to pay the fees specified by the KRS provider.
-          Existing balance on wallet: ${totalInputAmount.toString()}. Estimated network fee for the recovery transaction
-          : ${approximateFee.toString()}, KRS fee to pay: ${krsFee.toString()}. After deducting fees, your total
-          recoverable balance is ${recoveryAmount.toString()}`);
-  }
-
-  const recoveryOutputScript = utxolib.address.toOutputScript(params.recoveryDestination, coin.network);
-  psbt.addOutput({ script: recoveryOutputScript, value: recoveryAmount });
-
+  let krsFeeAddress: string | undefined;
   if (krsProvider && krsFee > BigInt(0)) {
     if (!krsProvider.feeAddresses) {
       throw new Error(`keyProvider must define feeAddresses`);
     }
 
-    const krsFeeAddress = krsProvider.feeAddresses[coin.getChain()];
+    krsFeeAddress = krsProvider.feeAddresses[coin.getChain()];
 
     if (!krsFeeAddress) {
       throw new Error('this KRS provider has not configured their fee structure yet - recovery cannot be completed');
     }
-
-    const krsFeeOutputScript = utxolib.address.toOutputScript(krsFeeAddress, coin.network);
-    psbt.addOutput({ script: krsFeeOutputScript, value: krsFee });
   }
 
+  const psbt = createBackupKeyRecoveryPsbt(coin.network, walletKeys, unspents, {
+    feeRateSatVB: feePerByte,
+    recoveryDestination: params.recoveryDestination,
+    keyRecoveryServiceFee: krsFee,
+    keyRecoveryServiceFeeAddress: krsFeeAddress,
+  });
+
   if (isUnsignedSweep) {
     return {
       txHex: psbt.toHex(),
@@ -421,6 +400,7 @@ export async function backupKeyRecovery(
   if (isKrsRecovery) {
     txInfo.coin = coin.getChain();
     txInfo.backupKey = params.backupKey;
+    const recoveryAmount = getRecoveryAmount(psbt, params.recoveryDestination);
     txInfo.recoveryAmount = Number(recoveryAmount);
     txInfo.recoveryAmountString = recoveryAmount.toString();
   }

modules/abstract-utxo/src/recovery/psbt.ts

@@ -0,0 +1,92 @@
+import * as utxolib from '@bitgo/utxo-lib';
+import { Dimensions } from '@bitgo/unspents';
+
+type RootWalletKeys = utxolib.bitgo.RootWalletKeys;
+type WalletUnspent<TNumber extends number | bigint> = utxolib.bitgo.WalletUnspent<TNumber>;
+
+class InsufficientFundsError extends Error {
+  constructor(
+    public totalInputAmount: bigint,
+    public approximateFee: bigint,
+    public krsFee: bigint,
+    public recoveryAmount: bigint
+  ) {
+    super(
+      `This wallet's balance is too low to pay the fees specified by the KRS provider.` +
+        `Existing balance on wallet: ${totalInputAmount.toString()}. ` +
+        `Estimated network fee for the recovery transaction: ${approximateFee.toString()}` +
+        `KRS fee to pay: ${krsFee.toString()}. ` +
+        `After deducting fees, your total recoverable balance is ${recoveryAmount.toString()}`
+    );
+  }
+}
+
+export function createBackupKeyRecoveryPsbt(
+  network: utxolib.Network,
+  rootWalletKeys: RootWalletKeys,
+  unspents: WalletUnspent<bigint>[],
+  {
+    feeRateSatVB,
+    recoveryDestination,
+    keyRecoveryServiceFee,
+    keyRecoveryServiceFeeAddress,
+  }: {
+    feeRateSatVB: number;
+    recoveryDestination: string;
+    keyRecoveryServiceFee: bigint;
+    keyRecoveryServiceFeeAddress: string | undefined;
+  }
+): utxolib.bitgo.UtxoPsbt {
+  if (keyRecoveryServiceFee > 0 && !keyRecoveryServiceFeeAddress) {
+    throw new Error('keyRecoveryServiceFeeAddress is required when keyRecoveryServiceFee is provided');
+  }
+
+  const psbt = utxolib.bitgo.createPsbtForNetwork({ network });
+  utxolib.bitgo.addXpubsToPsbt(psbt, rootWalletKeys);
+  unspents.forEach((unspent) => {
+    utxolib.bitgo.addWalletUnspentToPsbt(psbt, unspent, rootWalletKeys, 'user', 'backup');
+  });
+
+  let dimensions = Dimensions.fromPsbt(psbt).plus(
+    Dimensions.fromOutput({ script: utxolib.address.toOutputScript(recoveryDestination, network) })
+  );
+
+  if (keyRecoveryServiceFeeAddress) {
+    dimensions = dimensions.plus(
+      Dimensions.fromOutput({
+        script: utxolib.address.toOutputScript(keyRecoveryServiceFeeAddress, network),
+      })
+    );
+  }
+
+  const approximateFee = BigInt(dimensions.getVSize() * feeRateSatVB);
+
+  const totalInputAmount = utxolib.bitgo.unspentSum(unspents, 'bigint');
+
+  const recoveryAmount = totalInputAmount - approximateFee - keyRecoveryServiceFee;
+
+  // FIXME(BTC-2650): we should check for dust limit here instead
+  if (recoveryAmount < BigInt(0)) {
+    throw new InsufficientFundsError(totalInputAmount, approximateFee, keyRecoveryServiceFee, recoveryAmount);
+  }
+
+  psbt.addOutput({ script: utxolib.address.toOutputScript(recoveryDestination, network), value: recoveryAmount });
+
+  if (keyRecoveryServiceFeeAddress) {
+    psbt.addOutput({
+      script: utxolib.address.toOutputScript(keyRecoveryServiceFeeAddress, network),
+      value: keyRecoveryServiceFee,
+    });
+  }
+
+  return psbt;
+}
+
+export function getRecoveryAmount(psbt: utxolib.bitgo.UtxoPsbt, address: string): bigint {
+  const recoveryOutputScript = utxolib.address.toOutputScript(address, psbt.network);
+  const output = psbt.txOutputs.find((o) => o.script.equals(recoveryOutputScript));
+  if (!output) {
+    throw new Error(`Recovery destination output not found in PSBT: ${address}`);
+  }
+  return output.value;
+}