gecko/xputtest/dagwallet/wallet.go

123 lines
3.2 KiB
Go

// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package dagwallet
import (
"fmt"
"math"
"github.com/ava-labs/gecko/ids"
"github.com/ava-labs/gecko/utils/crypto"
"github.com/ava-labs/gecko/utils/timer"
"github.com/ava-labs/gecko/vms/spdagvm"
)
// Wallet is a holder for keys and UTXOs for the Ava DAG.
type Wallet struct {
networkID uint32
chainID ids.ID
clock timer.Clock
keychain *spdagvm.Keychain // Mapping from public address to the SigningKeys
utxoSet *UTXOSet // Mapping from utxoIDs to UTXOs
balance uint64
txFee uint64
}
// NewWallet returns a new Wallet
func NewWallet(networkID uint32, chainID ids.ID, txFee uint64) *Wallet {
return &Wallet{
networkID: networkID,
chainID: chainID,
keychain: spdagvm.NewKeychain(networkID, chainID),
utxoSet: &UTXOSet{},
txFee: txFee,
}
}
// GetAddress returns one of the addresses this wallet manages. If no address
// exists, one will be created.
func (w *Wallet) GetAddress() ids.ShortID {
if w.keychain.Addrs.Len() == 0 {
return w.CreateAddress()
}
return w.keychain.Addrs.CappedList(1)[0]
}
// CreateAddress returns a new address.
// It also saves the address and the private key that controls it
// so the address can be used later
func (w *Wallet) CreateAddress() ids.ShortID {
privKey, _ := w.keychain.New()
return privKey.PublicKey().Address()
}
// ImportKey imports a private key into this wallet
func (w *Wallet) ImportKey(sk *crypto.PrivateKeySECP256K1R) { w.keychain.Add(sk) }
// AddUTXO adds a new UTXO to this wallet if this wallet may spend it
// The UTXO's output must be an OutputPayment
func (w *Wallet) AddUTXO(utxo *spdagvm.UTXO) {
out, ok := utxo.Out().(*spdagvm.OutputPayment)
if !ok {
return
}
if _, _, err := w.keychain.Spend(utxo, math.MaxUint64); err == nil {
w.utxoSet.Put(utxo)
w.balance += out.Amount()
}
}
// Balance returns the amount of the assets in this wallet
func (w *Wallet) Balance() uint64 { return w.balance }
// Send sends [amount] to address [destAddr]
// The output of this transaction may be spent after [locktime]
func (w *Wallet) Send(amount uint64, locktime uint64, destAddr ids.ShortID) *spdagvm.Tx {
builder := spdagvm.Builder{
NetworkID: w.networkID,
ChainID: w.chainID,
}
currentTime := w.clock.Unix()
// Send any change to an address this wallet controls
changeAddr := ids.ShortID{}
if w.keychain.Addrs.Len() < 1000 {
changeAddr = w.CreateAddress()
} else {
changeAddr = w.GetAddress()
}
utxoList := w.utxoSet.UTXOs // List of UTXOs this wallet may spend
destAddrs := []ids.ShortID{destAddr}
// Build the transaction
tx, err := builder.NewTxFromUTXOs(w.keychain, utxoList, amount, w.txFee, locktime, 1, destAddrs, changeAddr, currentTime)
if err != nil {
panic(err)
}
// Remove from [w.utxoSet] any UTXOs used to fund [tx]
for _, in := range tx.Ins() {
if in, ok := in.(*spdagvm.InputPayment); ok {
inUTXOID := in.InputID()
w.utxoSet.Remove(inUTXOID)
w.balance -= in.Amount() // Deduct from [w.balance] the amount sent
}
}
return tx
}
func (w Wallet) String() string {
return fmt.Sprintf(
"Keychain:\n"+
"%s\n"+
"UTXOSet:\n"+
"%s",
w.keychain.PrefixedString(" "),
w.utxoSet.string(" "))
}