TestWallet.ts

import { Signer, SignatureRequest, SignatureResponse, SignTransactionOptions } from "../abstract-signer";
import { Provider, UtxoQueryOptions } from "../abstract-provider";
import { AddressesOption, AddressOption, Network, UTXO } from "../types";
import { bsv } from "scryptlib/dist";
import {parseAddresses} from "../utils";

const DEFAULT_SIGHASH_TYPE = bsv.crypto.Signature.ALL;


/**
 * An implemention of a simple wallet which should just be used in dev/test environments.
 * It can hold multiple private keys and have a feature of cachable in-memory utxo management.
 * 
 * Reminder: DO NOT USE IT IN PRODUCTION ENV.
 */
export class TestWallet extends Signer {


  private readonly _privateKeys: bsv.PrivateKey[];

  private _utxoManagers: Map<string, CacheableUtxoManager>;

  constructor(privateKey: bsv.PrivateKey | bsv.PrivateKey[], provider?: Provider) {
    super(provider);
    if (privateKey instanceof Array) {
      this._privateKeys = privateKey;
    } else {
      this._privateKeys = [privateKey];
    }
    this._utxoManagers = new Map();
  }

  get network(): Network {
    return bsv.Networks.testnet;
  }

  get addresses(): string[] {
    return this._privateKeys.map(p => p.toAddress(this.network).toString());
  }

  addPrivateKey(privateKey: bsv.PrivateKey | bsv.PrivateKey[]): this {
    const keys: bsv.PrivateKey[] = privateKey instanceof Array ? privateKey : [privateKey]
    this._privateKeys.push(...keys)
    return this
  }

  getDefaultAddress(): Promise<bsv.Address> {
    return Promise.resolve(this._defaultPrivateKey.toAddress());
  }

  getDefaultPubKey(): Promise<bsv.PublicKey> {
    return Promise.resolve(this._defaultPrivateKey.toPublicKey());
  }

  getPubKey(address: AddressOption): Promise<bsv.PublicKey> {
    return Promise.resolve(this._getPrivateKeys(address)[0].toPublicKey());
  }

  async signRawTransaction(rawTxHex: string, options: SignTransactionOptions): Promise<string> {
    const sigReqsByInputIndex: Map<number, SignatureRequest> = (options?.sigRequests || []).reduce((m, sigReq) => { m.set(sigReq.inputIndex, sigReq); return m; }, new Map());
    const tx = new bsv.Transaction(rawTxHex);
    tx.inputs.forEach((_, inputIndex) => {
      const sigReq = sigReqsByInputIndex.get(inputIndex);
      if (!sigReq) {
        throw new Error(`\`SignatureRequest\` info should be provided for the input ${inputIndex} to call #signRawTransaction`)
      }
      const script = sigReq.scriptHex ? new bsv.Script(sigReq.scriptHex) : bsv.Script.buildPublicKeyHashOut(sigReq.address.toString());
      // set ref output of the input
      tx.inputs[inputIndex].output = new bsv.Transaction.Output({
        script,
        satoshis: sigReq.satoshis
      })
    });
    const signedTx = await this.signTransaction(tx, options);
    return signedTx.toString();
  }

  async signTransaction(tx: bsv.Transaction, options?: SignTransactionOptions): Promise<bsv.Transaction> {
    const addresses = options?.address;
    this._checkAddressOption(addresses);
    // TODO: take account of SignatureRequests in options.
    return Promise.resolve(tx.sign(this._getPrivateKeys(addresses)));
  }

  signMessage(message: string, address?: AddressOption): Promise<string> {
    throw new Error("Method #signMessage not implemented.");
  }

  getSignatures(rawTxHex: string, sigRequests: SignatureRequest[]): Promise<SignatureResponse[]> {
    this._checkAddressOption(this._getAddressesIn(sigRequests));
    const tx = new bsv.Transaction(rawTxHex);
    const sigResponses: SignatureResponse[] = sigRequests.flatMap(sigReq => {
      tx.inputs[sigReq.inputIndex].output = new bsv.Transaction.Output({
        // TODO: support multiSig?
        script: sigReq.scriptHex ? new bsv.Script(sigReq.scriptHex) : bsv.Script.buildPublicKeyHashOut(parseAddresses(sigReq.address)[0]),
        satoshis: sigReq.satoshis
      });
      const privkeys = this._getPrivateKeys(sigReq.address);
      return privkeys.map(privKey => {
        const sig = tx.getSignature(sigReq.inputIndex, privKey, sigReq.sigHashType);
        return {
          sig: sig as string,
          publicKey: privKey.publicKey.toString(),
          inputIndex: sigReq.inputIndex,
          sigHashType: sigReq.sigHashType || DEFAULT_SIGHASH_TYPE
        }
      })
    })
    return Promise.resolve(sigResponses);
  }


  async connect(provider: Provider): Promise<this> {
    this.provider = provider;
    await this.provider.connect();
    return this;
  }

  override async listUnspent(address: AddressOption, options: UtxoQueryOptions): Promise<UTXO[]> {
    let utxoManager = this._utxoManagers.get(address.toString());
    if (!utxoManager) {
      utxoManager = new CacheableUtxoManager(address, this);
      this._utxoManagers.set(address.toString(), utxoManager);
      await utxoManager.init();
    }
    return utxoManager.fetchUtxos(options?.minSatoshis);
  }

  private _getAddressesIn(sigRequests?: SignatureRequest[]): AddressesOption {
    return (sigRequests || []).flatMap((req) => {
      return req.address instanceof Array ? req.address : [req.address];
    })
  }

  private _checkAddressOption(address?: AddressesOption) {
    if (!address) return;
    if (address instanceof Array) {
      (address as AddressOption[]).forEach(address => this._checkAddressOption(address));
    } else {
      if (!this.addresses.includes(address.toString())) {
        throw new Error(`the address ${address.toString()} is not belong to this SimpleWallet`);
      }
    }
  }

  private get _defaultPrivateKey(): bsv.PrivateKey {
    return this._privateKeys[0];
  }

  private _getPrivateKeys(address?: AddressesOption): bsv.PrivateKey[] {
    if (!address) return [this._defaultPrivateKey];
    this._checkAddressOption(address);
    let addresses = [];
    if (address instanceof Array) {
      (address as AddressOption[]).forEach(addr => addresses.push(addr.toString()));
    } else {
      addresses.push(address.toString())
    }
    return this._privateKeys.filter(priv => addresses.includes(priv.toAddress(this.network).toString()))
  }

}

enum InitState {
  UNINITIALIZED,
  INITIALIZING,
  INITIALIZED
};

class CacheableUtxoManager {

  address: AddressOption;

  private readonly signer: Signer;

  private availableUtxos: UTXO[] = [];
  private initStates: InitState = InitState.UNINITIALIZED;
  private initUtxoCnt: number = 0;

  constructor(address: AddressOption, signer: Signer) {
    this.address = address
    this.signer = signer;
  }

  async init() {
    if (this.initStates === InitState.INITIALIZED) {
      return this;
    }
    if (this.initStates === InitState.UNINITIALIZED) {
      this.initStates = InitState.INITIALIZING;
      this.availableUtxos = await this.signer.connectedProvider.listUnspent(this.address);
      this.initStates = InitState.INITIALIZED;
      this.initUtxoCnt = this.availableUtxos.length;
      console.log(`current balance of address ${this.address} is ${this.availableUtxos.reduce((r, utxo) => r + utxo.satoshis, 0)} satoshis`);
    }
    while (this.initStates === InitState.INITIALIZING) {
      await sleep(1);
    }
    return this;
  }

  async fetchUtxos(targetSatoshis?: number): Promise<UTXO[]> {
    if (this.initStates === InitState.INITIALIZED
      && this.initUtxoCnt > 0
      && this.availableUtxos.length === 0
    ) {
      const timeoutSec = 30;
      for (let i = 0; i < timeoutSec; i++) {
        console.log('waiting for available utxos')
        await sleep(1);
        if (this.availableUtxos.length > 0) {
          break;
        }
      }
    }

    if (targetSatoshis === undefined) {
      const allUtxos = this.availableUtxos;
      this.availableUtxos = [];
      return allUtxos;
    }

    const sortedUtxos = this.availableUtxos.sort((a, b) => a.satoshis - b.satoshis);

    if (targetSatoshis > sortedUtxos.reduce((r, utxo) => r + utxo.satoshis, 0)) {
      throw new Error('no sufficient utxos to pay the fee of ' + targetSatoshis);
    }

    let idx = 0;
    let accAmt = 0;
    for (let i = 0; i < sortedUtxos.length; i++) {
      accAmt += sortedUtxos[i].satoshis;
      if (accAmt >= targetSatoshis) {
        idx = i;
        break;
      }
    }

    const usedUtxos = sortedUtxos.slice(0, idx + 1);

    // update the available utxos, remove used ones
    this.availableUtxos = sortedUtxos.slice(idx + 1);

    const dustLimit = 1;
    if (accAmt > targetSatoshis + dustLimit) {
      // split `accAmt` to `targetSatoshis` + `change` 
      const splitTx =
        new bsv.Transaction().from(usedUtxos)
          .addOutput(new bsv.Transaction.Output({
            script: bsv.Script.buildPublicKeyHashOut(this.address),
            satoshis: targetSatoshis
          }))
          .change(this.address); // here generates a new available utxo for address

      const txId = (await this.signer.signAndsendTransaction(splitTx)).id; // sendTx(splitTx);

      // update the available utxos, add the new created on as the change
      if (splitTx.outputs.length === 2) {
        this.availableUtxos = this.availableUtxos.concat({
          txId,
          outputIndex: 1,
          script: splitTx.outputs[1].script.toHex(),
          satoshis: splitTx.outputs[1].satoshis
        });
      }

      // return the new created utxo which has value of `targetSatoshis`
      return [
        {
          txId,
          outputIndex: 0,
          script: splitTx.outputs[0].script.toHex(),
          satoshis: splitTx.outputs[0].satoshis,
        }
      ];
    } else {
      return usedUtxos;
    }
  }

  collectUtxoFrom(output: bsv.Transaction.Output, txId: string, outputIndex: number) {
    if (output.script.toHex() === this.utxoScriptHex) {
      this.availableUtxos.push({
        txId,
        outputIndex,
        satoshis: output.satoshis,
        script: output.script.toHex()
      });
    }
  }

  private get utxoScriptHex(): string {
    // all managed utxos should have the same P2PKH script for `this.address`
    return bsv.Script.buildPublicKeyHashOut(this.address).toHex();
  }
}

const sleep = async (seconds: number) => {
  return new Promise((resolve) => {
    setTimeout(() => {
      resolve({});
    }, seconds * 1000);
  })
}