memwallet.go raw

   1  package rpctest
   2  
   3  import (
   4  	"bytes"
   5  	"encoding/binary"
   6  	"fmt"
   7  	"sync"
   8  
   9  	"github.com/p9c/p9/pkg/amt"
  10  	"github.com/p9c/p9/pkg/btcaddr"
  11  	"github.com/p9c/p9/pkg/chaincfg"
  12  
  13  	"github.com/p9c/p9/pkg/qu"
  14  
  15  	"github.com/p9c/p9/pkg/blockchain"
  16  	"github.com/p9c/p9/pkg/chainhash"
  17  	ec "github.com/p9c/p9/pkg/ecc"
  18  	"github.com/p9c/p9/pkg/rpcclient"
  19  	"github.com/p9c/p9/pkg/txscript"
  20  	"github.com/p9c/p9/pkg/util"
  21  	"github.com/p9c/p9/pkg/util/hdkeychain"
  22  	"github.com/p9c/p9/pkg/wire"
  23  )
  24  
  25  var (
  26  	// hdSeed is the BIP 32 seed used by the memWallet to initialize it's HD root key. This value is hard coded in order
  27  	// to ensure deterministic behavior across test runs.
  28  	hdSeed = [chainhash.HashSize]byte{
  29  		0x79, 0xa6, 0x1a, 0xdb, 0xc6, 0xe5, 0xa2, 0xe1,
  30  		0x39, 0xd2, 0x71, 0x3a, 0x54, 0x6e, 0xc7, 0xc8,
  31  		0x75, 0x63, 0x2e, 0x75, 0xf1, 0xdf, 0x9c, 0x3f,
  32  		0xa6, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  33  	}
  34  )
  35  
  36  // utxo represents an unspent output spendable by the memWallet. The maturity height of the transaction is recorded in
  37  // order to properly observe the maturity period of direct coinbase outputs.
  38  type utxo struct {
  39  	pkScript       []byte
  40  	value          amt.Amount
  41  	keyIndex       uint32
  42  	maturityHeight int32
  43  	isLocked       bool
  44  }
  45  
  46  // isMature returns true if the target utxo is considered "mature" at the passed block height. Otherwise, false is
  47  // returned.
  48  func (u *utxo) isMature(height int32) bool {
  49  	return height >= u.maturityHeight
  50  }
  51  
  52  // chainUpdate encapsulates an update to the current main chain. This struct is used to sync up the memWallet each time
  53  // a new block is connected to the main chain.
  54  type chainUpdate struct {
  55  	filteredTxns []*util.Tx
  56  	blockHeight  int32
  57  	isConnect    bool // True if connect, false if disconnect
  58  }
  59  
  60  // undoEntry is functionally the opposite of a chainUpdate. An undoEntry is created for each new block received, then
  61  // stored in a log in order to properly handle block re-orgs.
  62  type undoEntry struct {
  63  	utxosDestroyed map[wire.OutPoint]*utxo
  64  	utxosCreated   []wire.OutPoint
  65  }
  66  
  67  // memWallet is a simple in-memory wallet whose purpose is to provide basic wallet functionality to the harness. The
  68  // wallet uses a hard-coded HD key hierarchy which promotes reproducibility between harness test runs.
  69  type memWallet struct {
  70  	coinbaseKey  *ec.PrivateKey
  71  	coinbaseAddr btcaddr.Address
  72  	// hdRoot is the root master private key for the wallet.
  73  	hdRoot *hdkeychain.ExtendedKey
  74  	// hdIndex is the next available key index offset from the hdRoot.
  75  	hdIndex uint32
  76  	// currentHeight is the latest height the wallet is known to be synced to.
  77  	currentHeight int32
  78  	// addrs tracks all addresses belonging to the wallet.
  79  	// The addresses are indexed by their keypath from the hdRoot.
  80  	addrs map[uint32]btcaddr.Address
  81  	// utxos is the set of utxos spendable by the wallet.
  82  	utxos map[wire.OutPoint]*utxo
  83  	// reorgJournal is a map storing an undo entry for each new block received. Once a block is disconnected, the undo
  84  	// entry for the particular height is evaluated, thereby rewinding the effect of the disconnected block on the
  85  	// wallet's set of spendable utxos.
  86  	reorgJournal      map[int32]*undoEntry
  87  	chainUpdates      []*chainUpdate
  88  	chainUpdateSignal qu.C
  89  	chainMtx          sync.Mutex
  90  	net               *chaincfg.Params
  91  	rpc               *rpcclient.Client
  92  	sync.RWMutex
  93  }
  94  
  95  // newMemWallet creates and returns a fully initialized instance of the memWallet given a particular blockchain's
  96  // parameters.
  97  func newMemWallet(net *chaincfg.Params, harnessID uint32) (*memWallet, error) {
  98  	// The wallet's final HD seed is: hdSeed || harnessID. This method ensures that each harness instance uses a
  99  	// deterministic root seed based on its harness ID.
 100  	var harnessHDSeed [chainhash.HashSize + 4]byte
 101  	copy(harnessHDSeed[:], hdSeed[:])
 102  	binary.BigEndian.PutUint32(harnessHDSeed[:chainhash.HashSize], harnessID)
 103  	hdRoot, e := hdkeychain.NewMaster(harnessHDSeed[:], net)
 104  	if e != nil {
 105  		return nil, nil
 106  	}
 107  	// The first child key from the hd root is reserved as the coinbase generation address.
 108  	coinbaseChild, e := hdRoot.Child(0)
 109  	if e != nil {
 110  		return nil, e
 111  	}
 112  	coinbaseKey, e := coinbaseChild.ECPrivKey()
 113  	if e != nil {
 114  		return nil, e
 115  	}
 116  	coinbaseAddr, e := keyToAddr(coinbaseKey, net)
 117  	if e != nil {
 118  		return nil, e
 119  	}
 120  	// Track the coinbase generation address to ensure we properly track newly generated DUO we can spend.
 121  	addrs := make(map[uint32]btcaddr.Address)
 122  	addrs[0] = coinbaseAddr
 123  	return &memWallet{
 124  		net:               net,
 125  		coinbaseKey:       coinbaseKey,
 126  		coinbaseAddr:      coinbaseAddr,
 127  		hdIndex:           1,
 128  		hdRoot:            hdRoot,
 129  		addrs:             addrs,
 130  		utxos:             make(map[wire.OutPoint]*utxo),
 131  		chainUpdateSignal: qu.T(),
 132  		reorgJournal:      make(map[int32]*undoEntry),
 133  	}, nil
 134  }
 135  
 136  // Start launches all goroutines required for the wallet to function properly.
 137  func (m *memWallet) Start() {
 138  	go m.chainSyncer()
 139  }
 140  
 141  // SyncedHeight returns the height the wallet is known to be synced to. This function is safe for concurrent access.
 142  func (m *memWallet) SyncedHeight() int32 {
 143  	m.RLock()
 144  	defer m.RUnlock()
 145  	return m.currentHeight
 146  }
 147  
 148  // SetRPCClient saves the passed rpc connection to pod as the wallet's personal rpc connection.
 149  func (m *memWallet) SetRPCClient(rpcClient *rpcclient.Client) {
 150  	m.rpc = rpcClient
 151  }
 152  
 153  // IngestBlock is a call-back which is to be triggered each time a new block is connected to the main chain. It queues
 154  // the update for the chain syncer, calling the private version in sequential order.
 155  func (m *memWallet) IngestBlock(height int32, header *wire.BlockHeader, filteredTxns []*util.Tx) {
 156  	// Append this new chain update to the end of the queue of new chain
 157  	// updates.
 158  	m.chainMtx.Lock()
 159  	m.chainUpdates = append(
 160  		m.chainUpdates, &chainUpdate{
 161  			filteredTxns,
 162  			height,
 163  			true,
 164  		},
 165  	)
 166  	m.chainMtx.Unlock()
 167  	// Launch a goroutine to signal the chainSyncer that a new update is available. We do this in a new goroutine in
 168  	// order to avoid blocking the main loop of the rpc client.
 169  	go func() {
 170  		m.chainUpdateSignal <- struct{}{}
 171  	}()
 172  }
 173  
 174  // ingestBlock updates the wallet's internal utxo state based on the outputs created and destroyed within each block.
 175  func (m *memWallet) ingestBlock(update *chainUpdate) {
 176  	// Update the latest synced height, then process each filtered transaction in the block creating and destroying
 177  	// utxos within the wallet as a result.
 178  	m.currentHeight = update.blockHeight
 179  	undo := &undoEntry{
 180  		utxosDestroyed: make(map[wire.OutPoint]*utxo),
 181  	}
 182  	for _, tx := range update.filteredTxns {
 183  		mtx := tx.MsgTx()
 184  		isCoinbase := blockchain.IsCoinBaseTx(mtx)
 185  		txHash := mtx.TxHash()
 186  		m.evalOutputs(mtx.TxOut, &txHash, isCoinbase, undo)
 187  		m.evalInputs(mtx.TxIn, undo)
 188  	}
 189  	// Finally, record the undo entry for this block so we can properly update our internal state in response to the
 190  	// block being re-org'd from the main chain.
 191  	m.reorgJournal[update.blockHeight] = undo
 192  }
 193  
 194  // chainSyncer is a goroutine dedicated to processing new blocks in order to keep the wallet's utxo state up to date.
 195  // NOTE: This MUST be run as a goroutine.
 196  func (m *memWallet) chainSyncer() {
 197  	var update *chainUpdate
 198  	for range m.chainUpdateSignal {
 199  		// A new update is available, so pop the new chain update from the front of the update queue.
 200  		m.chainMtx.Lock()
 201  		update = m.chainUpdates[0]
 202  		m.chainUpdates[0] = nil // Set to nil to prevent GC leak.
 203  		m.chainUpdates = m.chainUpdates[1:]
 204  		m.chainMtx.Unlock()
 205  		m.Lock()
 206  		if update.isConnect {
 207  			m.ingestBlock(update)
 208  		} else {
 209  			m.unwindBlock(update)
 210  		}
 211  		m.Unlock()
 212  	}
 213  }
 214  
 215  // evalOutputs evaluates each of the passed outputs, creating a new matching utxo within the wallet if we're able to
 216  // spend the output.
 217  func (m *memWallet) evalOutputs(
 218  	outputs []*wire.TxOut, txHash *chainhash.Hash,
 219  	isCoinbase bool, undo *undoEntry,
 220  ) {
 221  	for i, output := range outputs {
 222  		pkScript := output.PkScript
 223  		// Scan all the addresses we currently control to see if the output is paying to us.
 224  		for keyIndex, addr := range m.addrs {
 225  			pkHash := addr.ScriptAddress()
 226  			if !bytes.Contains(pkScript, pkHash) {
 227  				continue
 228  			}
 229  			// If this is a coinbase output, then we mark the maturity height at the proper block height in the future.
 230  			var maturityHeight int32
 231  			if isCoinbase {
 232  				maturityHeight = m.currentHeight + int32(m.net.CoinbaseMaturity)
 233  			}
 234  			op := wire.OutPoint{Hash: *txHash, Index: uint32(i)}
 235  			m.utxos[op] = &utxo{
 236  				value:          amt.Amount(output.Value),
 237  				keyIndex:       keyIndex,
 238  				maturityHeight: maturityHeight,
 239  				pkScript:       pkScript,
 240  			}
 241  			undo.utxosCreated = append(undo.utxosCreated, op)
 242  		}
 243  	}
 244  }
 245  
 246  // evalInputs scans all the passed inputs, destroying any utxos within the wallet which are spent by an input.
 247  func (m *memWallet) evalInputs(inputs []*wire.TxIn, undo *undoEntry) {
 248  	for _, txIn := range inputs {
 249  		op := txIn.PreviousOutPoint
 250  		oldUtxo, ok := m.utxos[op]
 251  		if !ok {
 252  			continue
 253  		}
 254  		undo.utxosDestroyed[op] = oldUtxo
 255  		delete(m.utxos, op)
 256  	}
 257  }
 258  
 259  // UnwindBlock is a call-back which is to be executed each time a block is disconnected from the main chain. It queues
 260  // the update for the chain syncer, calling the private version in sequential order.
 261  func (m *memWallet) UnwindBlock(height int32, header *wire.BlockHeader) {
 262  	// Append this new chain update to the end of the queue of new chain
 263  	// updates.
 264  	m.chainMtx.Lock()
 265  	m.chainUpdates = append(
 266  		m.chainUpdates, &chainUpdate{
 267  			nil,
 268  			height,
 269  			false,
 270  		},
 271  	)
 272  	m.chainMtx.Unlock()
 273  	// Launch a goroutine to signal the chainSyncer that a new update is available. We do this in a new goroutine in
 274  	// order to avoid blocking the main loop of the rpc client.
 275  	go func() {
 276  		m.chainUpdateSignal <- struct{}{}
 277  	}()
 278  }
 279  
 280  // unwindBlock undoes the effect that a particular block had on the wallet's internal utxo state.
 281  func (m *memWallet) unwindBlock(update *chainUpdate) {
 282  	undo := m.reorgJournal[update.blockHeight]
 283  	for _, utxo := range undo.utxosCreated {
 284  		delete(m.utxos, utxo)
 285  	}
 286  	for outPoint, utxo := range undo.utxosDestroyed {
 287  		m.utxos[outPoint] = utxo
 288  	}
 289  	delete(m.reorgJournal, update.blockHeight)
 290  }
 291  
 292  // newAddress returns a new address from the wallet's hd key chain. It also loads the address into the RPC client's
 293  // transaction filter to ensure any transactions that involve it are delivered via the notifications.
 294  func (m *memWallet) newAddress() (btcaddr.Address, error) {
 295  	index := m.hdIndex
 296  	childKey, e := m.hdRoot.Child(index)
 297  	if e != nil {
 298  		return nil, e
 299  	}
 300  	privKey, e := childKey.ECPrivKey()
 301  	if e != nil {
 302  		return nil, e
 303  	}
 304  	addr, e := keyToAddr(privKey, m.net)
 305  	if e != nil {
 306  		return nil, e
 307  	}
 308  	e = m.rpc.LoadTxFilter(false, []btcaddr.Address{addr}, nil)
 309  	if e != nil {
 310  		return nil, e
 311  	}
 312  	m.addrs[index] = addr
 313  	m.hdIndex++
 314  	return addr, nil
 315  }
 316  
 317  // NewAddress returns a fresh address spendable by the wallet. This function is safe for concurrent access.
 318  func (m *memWallet) NewAddress() (btcaddr.Address, error) {
 319  	m.Lock()
 320  	defer m.Unlock()
 321  	return m.newAddress()
 322  }
 323  
 324  // fundTx attempts to fund a transaction sending amt bitcoin. The coins are selected such that the final amount spent
 325  // pays enough fees as dictated by the passed fee rate. The passed fee rate should be expressed in satoshis-per-byte.
 326  // The transaction being funded can optionally include a change output indicated by the change boolean. NOTE: The
 327  // memWallet's mutex must be held when this function is called.
 328  func (m *memWallet) fundTx(
 329  	tx *wire.MsgTx, amount amt.Amount,
 330  	feeRate amt.Amount, change bool,
 331  ) (e error) {
 332  	const (
 333  		// spendSize is the largest number of bytes of a sigScript which spends a p2pkh output: OP_DATA_73 <sig>
 334  		// OP_DATA_33 <pubkey>
 335  		spendSize = 1 + 73 + 1 + 33
 336  	)
 337  	var (
 338  		amtSelected amt.Amount
 339  		txSize      int
 340  	)
 341  	for outPoint, utxo := range m.utxos {
 342  		// Skip any outputs that are still currently immature or are currently locked.
 343  		if !utxo.isMature(m.currentHeight) || utxo.isLocked {
 344  			continue
 345  		}
 346  		amtSelected += utxo.value
 347  		// Add the selected output to the transaction, updating the current tx size while accounting for the size of the
 348  		// future sigScript.
 349  		op := outPoint
 350  		tx.AddTxIn(wire.NewTxIn(&op, nil, nil))
 351  		txSize = tx.SerializeSize() + spendSize*len(tx.TxIn)
 352  		// Calculate the fee required for the txn at this point observing the specified fee rate. If we don't have
 353  		// enough coins from he current amount selected to pay the fee, then continue to grab more coins.
 354  		reqFee := amt.Amount(txSize * int(feeRate))
 355  		if amtSelected-reqFee < amount {
 356  			continue
 357  		}
 358  		// If we have any change left over and we should create a change output, then add an additional output to the
 359  		// transaction reserved for it.
 360  		changeVal := amtSelected - amount - reqFee
 361  		if changeVal > 0 && change {
 362  			addr, e := m.newAddress()
 363  			if e != nil {
 364  				return e
 365  			}
 366  			pkScript, e := txscript.PayToAddrScript(addr)
 367  			if e != nil {
 368  				return e
 369  			}
 370  			changeOutput := &wire.TxOut{
 371  				Value:    int64(changeVal),
 372  				PkScript: pkScript,
 373  			}
 374  			tx.AddTxOut(changeOutput)
 375  		}
 376  		return nil
 377  	}
 378  	// If we've reached this point, then coin selection failed due to an insufficient amount of coins.
 379  	return fmt.Errorf("not enough funds for coin selection")
 380  }
 381  
 382  // SendOutputs creates then sends a transaction paying to the specified output while observing the passed fee rate. The
 383  // passed fee rate should be expressed in satoshis-per-byte.
 384  func (m *memWallet) SendOutputs(
 385  	outputs []*wire.TxOut,
 386  	feeRate amt.Amount,
 387  ) (*chainhash.Hash, error) {
 388  	tx, e := m.CreateTransaction(outputs, feeRate, true)
 389  	if e != nil {
 390  		return nil, e
 391  	}
 392  	return m.rpc.SendRawTransaction(tx, true)
 393  }
 394  
 395  // SendOutputsWithoutChange creates and sends a transaction that pays to the specified outputs while observing the
 396  // passed fee rate and ignoring a change output. The passed fee rate should be expressed in sat/b.
 397  func (m *memWallet) SendOutputsWithoutChange(
 398  	outputs []*wire.TxOut,
 399  	feeRate amt.Amount,
 400  ) (*chainhash.Hash, error) {
 401  	tx, e := m.CreateTransaction(outputs, feeRate, false)
 402  	if e != nil {
 403  		return nil, e
 404  	}
 405  	return m.rpc.SendRawTransaction(tx, true)
 406  }
 407  
 408  // CreateTransaction returns a fully signed transaction paying to the specified outputs while observing the desired fee
 409  // rate. The passed fee rate should be expressed in satoshis-per-byte. The transaction being created can optionally
 410  // include a change output indicated by the change boolean. This function is safe for concurrent access.
 411  func (m *memWallet) CreateTransaction(
 412  	outputs []*wire.TxOut,
 413  	feeRate amt.Amount, change bool,
 414  ) (*wire.MsgTx, error) {
 415  	m.Lock()
 416  	defer m.Unlock()
 417  	tx := wire.NewMsgTx(wire.TxVersion)
 418  	// Tally up the total amount to be sent in order to perform coin selection shortly below.
 419  	var outputAmt amt.Amount
 420  	for _, output := range outputs {
 421  		outputAmt += amt.Amount(output.Value)
 422  		tx.AddTxOut(output)
 423  	}
 424  	// Attempt to fund the transaction with spendable utxos.
 425  	if e := m.fundTx(tx, outputAmt, feeRate, change); E.Chk(e) {
 426  		return nil, e
 427  	}
 428  	// Populate all the selected inputs with valid sigScript for spending. Along the way record all outputs being spent
 429  	// in order to avoid a potential double spend.
 430  	spentOutputs := make([]*utxo, 0, len(tx.TxIn))
 431  	for i, txIn := range tx.TxIn {
 432  		outPoint := txIn.PreviousOutPoint
 433  		utxo := m.utxos[outPoint]
 434  		extendedKey, e := m.hdRoot.Child(utxo.keyIndex)
 435  		if e != nil {
 436  			return nil, e
 437  		}
 438  		privKey, e := extendedKey.ECPrivKey()
 439  		if e != nil {
 440  			return nil, e
 441  		}
 442  		sigScript, e := txscript.SignatureScript(
 443  			tx, i, utxo.pkScript,
 444  			txscript.SigHashAll, privKey, true,
 445  		)
 446  		if e != nil {
 447  			return nil, e
 448  		}
 449  		txIn.SignatureScript = sigScript
 450  		spentOutputs = append(spentOutputs, utxo)
 451  	}
 452  	// As these outputs are now being spent by this newly created transaction, mark the outputs are "locked". This
 453  	// action ensures these outputs won't be double spent by any subsequent transactions. These locked outputs can be
 454  	// freed via a call to UnlockOutputs.
 455  	for _, utxo := range spentOutputs {
 456  		utxo.isLocked = true
 457  	}
 458  	return tx, nil
 459  }
 460  
 461  // UnlockOutputs unlocks any outputs which were previously locked due to being selected to fund a transaction via the
 462  // CreateTransaction method. This function is safe for concurrent access.
 463  func (m *memWallet) UnlockOutputs(inputs []*wire.TxIn) {
 464  	m.Lock()
 465  	defer m.Unlock()
 466  	for _, input := range inputs {
 467  		utxo, ok := m.utxos[input.PreviousOutPoint]
 468  		if !ok {
 469  			continue
 470  		}
 471  		utxo.isLocked = false
 472  	}
 473  }
 474  
 475  // ConfirmedBalance returns the confirmed balance of the wallet. This function is safe for concurrent access.
 476  func (m *memWallet) ConfirmedBalance() amt.Amount {
 477  	m.RLock()
 478  	defer m.RUnlock()
 479  	var balance amt.Amount
 480  	for _, utxo := range m.utxos {
 481  		// Prevent any immature or locked outputs from contributing to the wallet's total confirmed balance.
 482  		if !utxo.isMature(m.currentHeight) || utxo.isLocked {
 483  			continue
 484  		}
 485  		balance += utxo.value
 486  	}
 487  	return balance
 488  }
 489  
 490  // keyToAddr maps the passed private to corresponding p2pkh address.
 491  func keyToAddr(key *ec.PrivateKey, net *chaincfg.Params) (btcaddr.Address, error) {
 492  	serializedKey := key.PubKey().SerializeCompressed()
 493  	pubKeyAddr, e := btcaddr.NewPubKey(serializedKey, net)
 494  	if e != nil {
 495  		return nil, e
 496  	}
 497  	return pubKeyAddr.PubKeyHash(), nil
 498  }
 499