Add balance calculation logic (#419)

* chore: move balance calculation logic

* chore: disable ordinals spending for any address type

sdk/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@gobob/bob-sdk",
-  "version": "3.0.4",
+  "version": "3.1.0",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "scripts": {

sdk/src/wallet/utxo.ts

@@ -98,7 +98,7 @@ export async function createBitcoinPsbt(
         esploraClient.getAddressUtxos(fromAddress),
         feeRate === undefined ? esploraClient.getFeeEstimate(confirmationTarget) : feeRate,
         // cardinal = return UTXOs not containing inscriptions or runes
-        addressInfo.type === AddressType.p2tr ? ordinalsClient.getOutputsFromAddress(fromAddress, 'cardinal') : [],
+        ordinalsClient.getOutputsFromAddress(fromAddress, 'cardinal'),
     ]);
 
     if (confirmedUtxos.length === 0) {
@@ -122,11 +122,7 @@ export async function createBitcoinPsbt(
                 publicKey
             );
             // to support taproot addresses we want to exclude outputs which contain inscriptions
-            if (addressInfo.type === AddressType.p2tr) {
-                if (outpointsSet.has(OutPoint.toString(utxo))) possibleInputs.push(input);
-            } else {
-                possibleInputs.push(input);
-            }
+            if (outpointsSet.has(OutPoint.toString(utxo))) possibleInputs.push(input);
         })
     );
 
@@ -166,13 +162,11 @@ export async function createBitcoinPsbt(
     });
 
     if (!transaction || !transaction.tx) {
+        console.debug('confirmedUtxos', confirmedUtxos);
+        console.debug('outputsFromAddress', outputsFromAddress);
         console.debug(`fromAddress: ${fromAddress}, toAddress: ${toAddress}, amount: ${amount}`);
         console.debug(`publicKey: ${publicKey}, opReturnData: ${opReturnData}`);
         console.debug(`feeRate: ${feeRate}, confirmationTarget: ${confirmationTarget}`);
-        if (addressInfo.type === AddressType.p2tr) {
-            console.debug('confirmedUtxos', confirmedUtxos);
-            console.debug('outputsFromAddress', outputsFromAddress);
-        }
         throw new Error('Failed to create transaction. Do you have enough funds?');
     }
 
@@ -313,7 +307,7 @@ export async function estimateTxFee(
         esploraClient.getAddressUtxos(fromAddress),
         feeRate === undefined ? esploraClient.getFeeEstimate(confirmationTarget) : feeRate,
         // cardinal = return UTXOs not containing inscriptions or runes
-        addressInfo.type === AddressType.p2tr ? ordinalsClient.getOutputsFromAddress(fromAddress, 'cardinal') : [],
+        ordinalsClient.getOutputsFromAddress(fromAddress, 'cardinal'),
     ]);
 
     if (confirmedUtxos.length === 0) {
@@ -336,11 +330,7 @@ export async function estimateTxFee(
             );
 
             // to support taproot addresses we want to exclude outputs which contain inscriptions
-            if (addressInfo.type === AddressType.p2tr) {
-                if (outpointsSet.has(OutPoint.toString(utxo))) possibleInputs.push(input);
-            } else {
-                possibleInputs.push(input);
-            }
+            if (outpointsSet.has(OutPoint.toString(utxo))) possibleInputs.push(input);
         })
     );
 
@@ -387,15 +377,72 @@ export async function estimateTxFee(
     });
 
     if (!transaction || !transaction.tx) {
+        console.debug('confirmedUtxos', confirmedUtxos);
+        console.debug('outputsFromAddress', outputsFromAddress);
         console.debug(`fromAddress: ${fromAddress}, amount: ${amount}`);
         console.debug(`publicKey: ${publicKey}, opReturnData: ${opReturnData}`);
         console.debug(`feeRate: ${feeRate}, confirmationTarget: ${confirmationTarget}`);
-        if (addressInfo.type === AddressType.p2tr) {
-            console.debug('confirmedUtxos', confirmedUtxos);
-            console.debug('outputsFromAddress', outputsFromAddress);
-        }
         throw new Error('Failed to create transaction. Do you have enough funds?');
     }
 
     return transaction.fee;
 }
+
+/**
+ * Get balance of provided address in satoshis.
+ *
+ * @typedef { {confirmed: BigInt, unconfirmed: BigInt, total: bigint} } Balance
+ *
+ * @param {string} [address] The Bitcoin address. If no address specified returning object will contain zeros.
+ * @returns {Promise<Balance>} The balance object of provided address in satoshis.
+ *
+ * @example
+ * ```typescript
+ * const address = 'bc1qar0srrr7xfkvy5l643lydnw9re59gtzzwf5mdq';
+ *
+ * const balance = await getBalance(address);
+ * console.log(balance);
+ * ```
+ *
+ * @dev UTXOs that contain inscriptions or runes will not be used to calculate balance.
+ */
+export async function getBalance(address?: string) {
+    if (!address) {
+        return { confirmed: BigInt(0), unconfirmed: BigInt(0), total: BigInt(0) };
+    }
+
+    const addressInfo = getAddressInfo(address);
+
+    const esploraClient = new EsploraClient(addressInfo.network);
+    const ordinalsClient = new OrdinalsClient(addressInfo.network);
+
+    const [outputs, cardinalOutputs] = await Promise.all([
+        esploraClient.getAddressUtxos(address),
+        // cardinal = return UTXOs not containing inscriptions or runes
+        ordinalsClient.getOutputsFromAddress(address, 'cardinal'),
+    ]);
+
+    const cardinalOutputsSet = new Set(cardinalOutputs.map((output) => output.outpoint));
+
+    const total = outputs.reduce((acc, output) => {
+        if (cardinalOutputsSet.has(OutPoint.toString(output))) {
+            return acc + output.value;
+        }
+
+        return acc;
+    }, 0);
+
+    const confirmed = outputs.reduce((acc, output) => {
+        if (cardinalOutputsSet.has(OutPoint.toString(output)) && output.confirmed) {
+            return acc + output.value;
+        }
+
+        return acc;
+    }, 0);
+
+    return {
+        confirmed: BigInt(confirmed),
+        unconfirmed: BigInt(total - confirmed),
+        total: BigInt(total),
+    };
+}

sdk/test/utxo.test.ts

@@ -1,10 +1,11 @@
-import { vi, describe, it, assert, Mock, expect } from 'vitest';
+import { vi, describe, it, assert, Mock, expect, beforeEach } from 'vitest';
 import { AddressType, getAddressInfo, Network } from 'bitcoin-address-validation';
 import { Address, NETWORK, OutScript, Script, Transaction, p2sh, p2wpkh, selectUTXO } from '@scure/btc-signer';
 import { hex, base64 } from '@scure/base';
-import { createBitcoinPsbt, getInputFromUtxoAndTx, estimateTxFee, Input } from '../src/wallet/utxo';
+import { createBitcoinPsbt, getInputFromUtxoAndTx, estimateTxFee, Input, getBalance } from '../src/wallet/utxo';
 import { TransactionOutput } from '@scure/btc-signer/psbt';
 import { OrdinalsClient, OutPoint } from '../src/ordinal-api';
+import { EsploraClient } from '../src/esplora';
 
 vi.mock(import('@scure/btc-signer'), async (importOriginal) => {
     const actual = await importOriginal();
@@ -15,9 +16,31 @@ vi.mock(import('@scure/btc-signer'), async (importOriginal) => {
     };
 });
 
+vi.mock(import('../src/ordinal-api'), async (importOriginal) => {
+    const actual = await importOriginal();
+
+    actual.OrdinalsClient.prototype.getOutputsFromAddress = vi.fn(
+        actual.OrdinalsClient.prototype.getOutputsFromAddress
+    );
+
+    return actual;
+});
+
+vi.mock(import('../src/esplora'), async (importOriginal) => {
+    const actual = await importOriginal();
+
+    actual.EsploraClient.prototype.getAddressUtxos = vi.fn(actual.EsploraClient.prototype.getAddressUtxos);
+
+    return actual;
+});
+
 // TODO: Add more tests using https://github.com/paulmillr/scure-btc-signer/tree/5ead71ea9a873d8ba1882a9cd6aa561ad410d0d1/test/bitcoinjs-test/fixtures/bitcoinjs
 // TODO: Ensure that the paymentAddresses have sufficient funds to create the transaction
 describe('UTXO Tests', () => {
+    beforeEach(() => {
+        vi.clearAllMocks();
+    });
+
     it('should spend from address to create a transaction with an OP return output', { timeout: 50000 }, async () => {
         // Addresses where randomly picked from blockstream.info
         const paymentAddresses = [
@@ -392,4 +415,56 @@ describe('UTXO Tests', () => {
             'Failed to create transaction. Do you have enough funds?'
         );
     });
+
+    it('should return address balance', async () => {
+        const address = 'bc1peqr5a5kfufvsl66444jm9y8qq0s87ph0zv4lfkcs7h40ew02uvsqkhjav0';
+
+        const balance = await getBalance(address);
+
+        assert(balance.confirmed);
+        assert(balance.total);
+        assert(
+            balance.confirmed === balance.total
+                ? balance.unconfirmed === 0n
+                : balance.unconfirmed === balance.total - balance.confirmed
+        );
+
+        const zeroBalance = await getBalance();
+
+        assert(zeroBalance.confirmed === 0n, 'If no address specified confirmed must be 0');
+        assert(zeroBalance.unconfirmed === 0n, 'If no address specified unconfirmed must be 0');
+        assert(zeroBalance.total === 0n, 'If no address specified total must be 0');
+    });
+
+    it('returns smalled amount if address holds ordinals', async () => {
+        const taprootAddress = 'bc1peqr5a5kfufvsl66444jm9y8qq0s87ph0zv4lfkcs7h40ew02uvsqkhjav0';
+
+        const esploraClient = new EsploraClient('mainnet');
+
+        const outputs = await esploraClient.getAddressUtxos(taprootAddress);
+
+        const total = outputs.reduce((acc, output) => acc + output.value, 0);
+
+        const confirmed = outputs.reduce((acc, output) => {
+            if (output.confirmed) {
+                return acc + output.value;
+            }
+
+            return acc;
+        }, 0);
+
+        // mock half of the UTXOs contain inscriptions or runes
+        (OrdinalsClient.prototype.getOutputsFromAddress as Mock).mockResolvedValueOnce(
+            outputs.slice(Math.ceil(outputs.length / 2)).map((output) => {
+                const outpoint = OutPoint.toString(output);
+
+                return { outpoint };
+            })
+        );
+
+        const balanceData = await getBalance(taprootAddress);
+
+        expect(balanceData.total).toBeLessThan(total);
+        expect(balanceData.confirmed).toBeLessThan(confirmed);
+    });
 });