@atomiqlabs/sdk/src/bitcoin/wallet/BitcoinWallet.ts

atomiq labs SDK for cross-chain swaps between smart chains and bitcoin

import {coinSelect, maxSendable, CoinselectAddressTypes, CoinselectTxInput} from "../coinselect2";
import {BTC_NETWORK, NETWORK, TEST_NETWORK} from "@scure/btc-signer/utils"
import {p2wpkh, OutScript, Transaction, p2tr, Address} from "@scure/btc-signer";
import {BitcoinWalletUtxo, BitcoinWalletUtxoBase, IBitcoinWallet} from "./IBitcoinWallet";
import {Buffer} from "buffer";
import {randomBytes} from "../../utils/Utils";
import {getDummyOutputScript, toCoinselectAddressType, toOutputScript} from "../../utils/BitcoinUtils";
import {TransactionInputUpdate} from "@scure/btc-signer/psbt";
import {getLogger} from "../../utils/Logger";
import {BitcoinNetwork, BitcoinRpcWithAddressIndex} from "@atomiqlabs/base";
import {utils} from "../coinselect2/utils";

/**
 * Identifies the address type of a Bitcoin address
 *
 * @category Bitcoin
 */
export function identifyAddressType(address: string, network: BTC_NETWORK): CoinselectAddressTypes {
    switch(Address(network).decode(address).type) {
        case "pkh":
            return "p2pkh";
        case "wpkh":
            return "p2wpkh";
        case "tr":
            return "p2tr";
        case "sh":
            return "p2sh-p2wpkh";
        case "wsh":
            return "p2wsh";
        default:
            throw new Error("Unknown address type of "+address);
    }
}

const btcNetworkMapping = {
    [BitcoinNetwork.MAINNET]: NETWORK,
    [BitcoinNetwork.TESTNET]: TEST_NETWORK,
    [BitcoinNetwork.TESTNET4]: TEST_NETWORK,
    [BitcoinNetwork.REGTEST]: {
        ...TEST_NETWORK,
        bech32: "bcrt"
    }
}

const logger = getLogger("BitcoinWallet: ");

/**
 * Abstract base class for Bitcoin wallet implementations, using bitcoin rpc with address index
 *  as a backend for fetching balances, UTXOs, etc.
 *
 * @category Bitcoin
 */
export abstract class BitcoinWallet implements IBitcoinWallet {

    protected readonly rpc: BitcoinRpcWithAddressIndex<any>;
    protected readonly network: BTC_NETWORK;
    protected feeMultiplier: number;
    protected feeOverride?: number;

    constructor(
        mempoolApi: BitcoinRpcWithAddressIndex<any>,
        network: BitcoinNetwork | BTC_NETWORK,
        feeMultiplier: number = 1.25, feeOverride?: number
    ) {
        this.rpc = mempoolApi;
        this.network = typeof(network)==="object" ? network : BitcoinWallet.bitcoinNetworkToObject(network);
        this.feeMultiplier = feeMultiplier;
        this.feeOverride = feeOverride;
    }

    /**
     * @inheritDoc
     */
    async getFeeRate(): Promise<number> {
        if(this.feeOverride!=null) {
            return this.feeOverride;
        }
        return Math.floor((await this.rpc.getFeeRate())*this.feeMultiplier);
    }

    /**
     * Internal helper function for sending a raw transaction through the underlying RPC
     *
     * @param rawHex Serialized bitcoin transaction in hexadecimal format
     * @returns txId Transaction ID of the submitted bitcoin transaction
     *
     * @protected
     */
    protected _sendTransaction(rawHex: string): Promise<string> {
        return this.rpc.sendRawTransaction(rawHex);
    }

    /**
     * Internal helper function for fetching the balance of the wallet given a specific bitcoin wallet address
     *
     * @param address
     * @protected
     */
    protected _getBalance(address: string): Promise<{ confirmedBalance: bigint; unconfirmedBalance: bigint }> {
        return this.rpc.getAddressBalances(address);
    }

    /**
     * Internal helper function for fetching the UTXO set of a given wallet address
     *
     * @param sendingAddress
     * @param sendingAddressType
     * @protected
     */
    protected async _getUtxoPool(
        sendingAddress: string,
        sendingAddressType: CoinselectAddressTypes
    ): Promise<BitcoinWalletUtxo[]> {
        const utxos = await this.rpc.getAddressUTXOs(sendingAddress);

        let totalSpendable = 0;

        const outputScript = toOutputScript(this.network, sendingAddress);

        const utxoPool: BitcoinWalletUtxo[] = [];

        for(let utxo of utxos) {
            const value = Number(utxo.value);
            totalSpendable += value;
            utxoPool.push({
                vout: utxo.vout,
                txId: utxo.txid,
                value: value,
                type: sendingAddressType,
                outputScript: outputScript,
                address: sendingAddress,
                cpfp: !utxo.confirmed ? await this.rpc.getCPFPData(utxo.txid).then((result) => {
                    if(result==null) return;
                    return {
                        txVsize: result.adjustedVsize,
                        txEffectiveFeeRate: result.effectiveFeePerVsize
                    }
                }) : undefined,
                confirmed: utxo.confirmed
            })
        }

        logger.debug("_getUtxoPool(): Total spendable value: "+totalSpendable+" num utxos: "+utxoPool.length);

        return utxoPool;
    }

    /**
     *
     * @param sendingAccounts
     * @param recipient
     * @param amount
     * @param feeRate
     * @protected
     */
    protected async _getPsbt(
        sendingAccounts: {
            pubkey: string,
            address: string,
            addressType: CoinselectAddressTypes,
        }[],
        recipient: string,
        amount: number,
        feeRate?: number
    ): Promise<{
        fee: number,
        psbt?: Transaction,
        inputAddressIndexes?: {[address: string]: number[]}
    }> {
        const psbt = new Transaction({PSBTVersion: 0});
        psbt.addOutput({
            amount: BigInt(amount),
            script: toOutputScript(this.network, recipient)
        });
        return this._fundPsbt(sendingAccounts, psbt, feeRate);
    }

    protected async _fundPsbt(
        sendingAccounts: {
            pubkey: string,
            address: string,
            addressType: CoinselectAddressTypes,
        }[],
        psbt: Transaction,
        _feeRate?: number,
        utxos?: BitcoinWalletUtxo[],
        spendFully?: boolean
    ): Promise<{
        fee: number,
        psbt?: Transaction,
        inputAddressIndexes?: {[address: string]: number[]}
    }> {
        const feeRate = _feeRate ?? await this.getFeeRate();
        const utxoPool: BitcoinWalletUtxo[] = utxos ?? (await Promise.all(sendingAccounts.map(acc => this._getUtxoPool(acc.address, acc.addressType)))).flat();

        if(spendFully && utxoPool==null) throw new Error("Cannot fully spend when no utxos are passed!");

        logger.debug("_fundPsbt(): fee rate: "+feeRate+" utxo pool: ", utxoPool);

        const accountPubkeys: Record<string, string> = {};
        sendingAccounts.forEach(acc => accountPubkeys[acc.address] = acc.pubkey);

        const requiredInputs: CoinselectTxInput[] = [];
        for(let i=0;i<psbt.inputsLength;i++) {
            const input = psbt.getInput(i);
            if(input.index==null || input.txid==null) throw new Error("Inputs need txid & index!");
            let amount: bigint;
            let script: Uint8Array;
            if(input.witnessUtxo!=null) {
                amount = input.witnessUtxo.amount as bigint;
                script = input.witnessUtxo.script as Uint8Array;
            } else if(input.nonWitnessUtxo!=null) {
                amount = input.nonWitnessUtxo.outputs[input.index].amount;
                script = input.nonWitnessUtxo.outputs[input.index].script;
            } else throw new Error("Either witnessUtxo or nonWitnessUtxo has to be defined!");
            requiredInputs.push({
                txId: Buffer.from(input.txid).toString('hex'),
                vout: input.index,
                value: Number(amount),
                type: toCoinselectAddressType(script)
            })
        }

        const targets: {value: number, script: Buffer}[] = [];
        for(let i=0;i<psbt.outputsLength;i++) {
            const output = psbt.getOutput(i);
            if(output.amount==null || output.script==null) throw new Error("Outputs need amount & script defined!");
            targets.push({
                value: Number(output.amount),
                script: Buffer.from(output.script)
            })
        }
        logger.debug("_fundPsbt(): Coinselect targets: ", targets);

        let coinselectResult = spendFully
            ? utils.finalize(requiredInputs.concat(utxoPool.filter(utxo => !utils.isDetrimentalInput(feeRate, utxo))), targets, feeRate, null)
            : coinSelect(utxoPool, targets, feeRate, sendingAccounts[0].addressType, requiredInputs);
        logger.debug("_fundPsbt(): Coinselect result: ", coinselectResult);

        if(coinselectResult.inputs==null || coinselectResult.outputs==null || coinselectResult.effectiveFeeRate==null) {
            return {
                fee: coinselectResult.fee
            };
        }

        if(spendFully && feeRate!=null) {
            const maximumAllowedFeeRate = (1.5*feeRate) + 10;
            if(coinselectResult.effectiveFeeRate > maximumAllowedFeeRate)
                throw new Error(`Effective fee rate too high, feeRate: ${coinselectResult.effectiveFeeRate} sats/vB, maximum: ${maximumAllowedFeeRate} sats/vB!`);
            const minimumAllowedFeeRate = 0.9*feeRate;
            if(coinselectResult.effectiveFeeRate < minimumAllowedFeeRate)
                throw new Error(`Effective fee rate too low, feeRate: ${coinselectResult.effectiveFeeRate} sats/vB, minimum: ${minimumAllowedFeeRate} sats/vB!`);
        }

        // Remove in/outs that are already in the PSBT
        coinselectResult.inputs.splice(0, psbt.inputsLength);
        coinselectResult.outputs.splice(0, psbt.outputsLength);

        const inputAddressIndexes: {[address: string]: number[]} = {};
        coinselectResult.inputs.forEach((input, index) => {
            inputAddressIndexes[input.address!] ??= [];
            inputAddressIndexes[input.address!].push(index);
        });

        const formattedInputs: TransactionInputUpdate[] = await Promise.all<TransactionInputUpdate>(coinselectResult.inputs.map(async (input) => {
            switch(input.type) {
                case "p2tr":
                    const parsed = p2tr(Buffer.from(accountPubkeys[input.address!], "hex"));
                    return {
                        txid: input.txId,
                        index: input.vout,
                        witnessUtxo: {
                            script: input.outputScript!,
                            amount: BigInt(input.value)
                        },
                        tapInternalKey: parsed.tapInternalKey,
                        tapMerkleRoot: parsed.tapMerkleRoot,
                        tapLeafScript: parsed.tapLeafScript
                    };
                case "p2wpkh":
                    return {
                        txid: input.txId,
                        index: input.vout,
                        witnessUtxo: {
                            script: input.outputScript!,
                            amount: BigInt(input.value)
                        },
                        sighashType: 0x01
                    };
                case "p2sh-p2wpkh":
                    return {
                        txid: input.txId,
                        index: input.vout,
                        witnessUtxo: {
                            script: input.outputScript!,
                            amount: BigInt(input.value)
                        },
                        redeemScript: p2wpkh(Buffer.from(accountPubkeys[input.address!], "hex"), this.network).script,
                        sighashType: 0x01
                    };
                case "p2pkh":
                    const tx = await this.rpc.getTransaction(input.txId);
                    if(tx==null) throw new Error("Cannot fetch existing tx "+input.txId);
                    return {
                        txid: input.txId,
                        index: input.vout,
                        nonWitnessUtxo: tx.raw,
                        sighashType: 0x01
                    };
                default:
                    throw new Error("Invalid input type: "+input.type);
            }
        }));

        formattedInputs.forEach(input => psbt.addInput(input));

        coinselectResult.outputs.forEach(output => {
            if(output.script==null && output.address==null) {
                //Change output
                psbt.addOutput({
                    script: toOutputScript(this.network, sendingAccounts[0].address),
                    amount: BigInt(Math.floor(output.value))
                });
            } else {
                psbt.addOutput({
                    script: output.script ?? toOutputScript(this.network, output.address!),
                    amount: BigInt(output.value)
                });
            }
        });

        return {
            psbt,
            fee: coinselectResult.fee,
            inputAddressIndexes
        };
    }

    protected async _getSpendableBalance(
        sendingAccounts: {
            address: string,
            addressType: CoinselectAddressTypes,
        }[],
        psbt?: Transaction,
        feeRate?: number,
        outputAddressType?: CoinselectAddressTypes,
        utxoPool?: BitcoinWalletUtxoBase[]
    ): Promise<{
        balance: bigint,
        feeRate: number,
        totalFee: number
    }> {
        feeRate ??= await this.getFeeRate();
        utxoPool ??= (await Promise.all(sendingAccounts.map(acc => this._getUtxoPool(acc.address, acc.addressType)))).flat();

        return {
            ...BitcoinWallet.getSpendableBalance(
                utxoPool ?? (await Promise.all(sendingAccounts.map(acc => this._getUtxoPool(acc.address, acc.addressType)))).flat(),
                feeRate ?? await this.getFeeRate(),
                psbt,
                outputAddressType
            ),
            feeRate
        };
    }

    abstract sendTransaction(address: string, amount: bigint, feeRate?: number): Promise<string>;
    abstract fundPsbt(psbt: Transaction, feeRate?: number, utxos?: BitcoinWalletUtxo[], spendFully?: boolean): Promise<Transaction>;
    abstract signPsbt(psbt: Transaction, signInputs: number[]): Promise<Transaction>;

    abstract getTransactionFee(address: string, amount: bigint, feeRate?: number): Promise<number>;
    abstract getFundedPsbtFee(psbt: Transaction, feeRate?: number): Promise<number>;

    abstract getReceiveAddress(): string;
    abstract getBalance(): Promise<{
        confirmedBalance: bigint,
        unconfirmedBalance: bigint
    }>;
    abstract getSpendableBalance(psbt?: Transaction, feeRate?: number): Promise<{
        balance: bigint,
        feeRate: number,
        totalFee: number
    }>;

    static bitcoinNetworkToObject(network: BitcoinNetwork): BTC_NETWORK {
        return btcNetworkMapping[network];
    }

    static getSpendableBalance(
        utxoPool: BitcoinWalletUtxoBase[],
        feeRate: number,
        psbt?: Transaction,
        outputAddressType?: CoinselectAddressTypes
    ): {
        balance: bigint,
        totalFee: number
    } {
        const requiredInputs: CoinselectTxInput[] = [];
        if(psbt!=null) for(let i=0;i<psbt.inputsLength;i++) {
            const input = psbt.getInput(i);
            if(input.index==null || input.txid==null) throw new Error("Inputs need txid & index!");
            let amount: bigint;
            let script: Uint8Array;
            if(input.witnessUtxo!=null) {
                amount = input.witnessUtxo.amount as bigint;
                script = input.witnessUtxo.script as Uint8Array;
            } else if(input.nonWitnessUtxo!=null) {
                amount = input.nonWitnessUtxo.outputs[input.index].amount;
                script = input.nonWitnessUtxo.outputs[input.index].script;
            } else throw new Error("Either witnessUtxo or nonWitnessUtxo has to be defined!");
            requiredInputs.push({
                txId: Buffer.from(input.txid).toString('hex'),
                vout: input.index,
                value: Number(amount),
                type: toCoinselectAddressType(script)
            })
        }

        const additionalOutputs: {value: number, script: Buffer}[] = [];
        if(psbt!=null) for(let i=0;i<psbt.outputsLength;i++) {
            const output = psbt.getOutput(i);
            if(output.amount==null || output.script==null) throw new Error("Outputs need amount & script!");
            additionalOutputs.push({
                value: Number(output.amount),
                script: Buffer.from(output.script)
            })
        }

        const target: Uint8Array = getDummyOutputScript(outputAddressType ?? "p2wsh");
        let coinselectResult = maxSendable(utxoPool, {script: Buffer.from(target), type: outputAddressType ?? "p2wsh"}, feeRate, requiredInputs, additionalOutputs);

        logger.debug("_getSpendableBalance(): Max spendable result: ", coinselectResult);

        return {
            balance: BigInt(Math.floor(coinselectResult.value)),
            totalFee: coinselectResult.fee
        }
    }

}