wallet/utxo.go

   1 package wallet
   2
   3 import (
   4         "encoding/json"
   5         log "github.com/sirupsen/logrus"
   6
   7         "github.com/bytom/bytom/account"
   8         "github.com/bytom/bytom/consensus"
   9         "github.com/bytom/bytom/consensus/segwit"
  10         "github.com/bytom/bytom/crypto/sha3pool"
  11         dbm "github.com/bytom/bytom/database/leveldb"
  12         "github.com/bytom/bytom/errors"
  13         "github.com/bytom/bytom/protocol/bc/types"
  14 )
  15
  16 // GetAccountUtxos return all account unspent outputs
  17 func (w *Wallet) GetAccountUtxos(accountID string, id string, unconfirmed, isSmartContract bool) []*account.UTXO {
  18         prefix := account.UTXOPreFix
  19         if isSmartContract {
  20                 prefix = account.SUTXOPrefix
  21         }
  22
  23         accountUtxos := []*account.UTXO{}
  24         if unconfirmed {
  25                 accountUtxos = w.AccountMgr.ListUnconfirmedUtxo(accountID, isSmartContract)
  26         }
  27
  28         accountUtxoIter := w.DB.IteratorPrefix([]byte(prefix + id))
  29         defer accountUtxoIter.Release()
  30
  31         for accountUtxoIter.Next() {
  32                 accountUtxo := &account.UTXO{}
  33                 if err := json.Unmarshal(accountUtxoIter.Value(), accountUtxo); err != nil {
  34                         log.WithFields(log.Fields{"module": logModule, "err": err}).Warn("GetAccountUtxos fail on unmarshal utxo")
  35                         continue
  36                 }
  37
  38                 if accountID == accountUtxo.AccountID || accountID == "" {
  39                         accountUtxos = append(accountUtxos, accountUtxo)
  40                 }
  41         }
  42         return accountUtxos
  43 }
  44
  45 func (w *Wallet) attachUtxos(batch dbm.Batch, b *types.Block) {
  46         for txIndex, tx := range b.Transactions {
  47                 //hand update the transaction input utxos
  48                 inputUtxos := txInToUtxos(tx)
  49                 for _, inputUtxo := range inputUtxos {
  50                         if segwit.IsP2WScript(inputUtxo.ControlProgram) {
  51                                 batch.Delete(account.StandardUTXOKey(inputUtxo.OutputID))
  52                         } else {
  53                                 batch.Delete(account.ContractUTXOKey(inputUtxo.OutputID))
  54                         }
  55                 }
  56
  57                 //hand update the transaction output utxos
  58                 validHeight := uint64(0)
  59                 if txIndex == 0 {
  60                         validHeight = b.Height + consensus.CoinbasePendingBlockNumber
  61                 }
  62                 outputUtxos := txOutToUtxos(tx, validHeight)
  63                 utxos := w.filterAccountUtxo(outputUtxos)
  64                 if err := batchSaveUtxos(utxos, batch); err != nil {
  65                         log.WithFields(log.Fields{"module": logModule, "err": err}).Error("attachUtxos fail on batchSaveUtxos")
  66                 }
  67         }
  68 }
  69
  70 func (w *Wallet) detachUtxos(batch dbm.Batch, b *types.Block) {
  71         for txIndex := len(b.Transactions) - 1; txIndex >= 0; txIndex-- {
  72                 tx := b.Transactions[txIndex]
  73                 for j := range tx.Outputs {
  74                         resOut, err := tx.Output(*tx.ResultIds[j])
  75                         if err != nil {
  76                                 continue
  77                         }
  78
  79                         if segwit.IsP2WScript(resOut.ControlProgram.Code) {
  80                                 batch.Delete(account.StandardUTXOKey(*tx.ResultIds[j]))
  81                         } else {
  82                                 batch.Delete(account.ContractUTXOKey(*tx.ResultIds[j]))
  83                         }
  84                 }
  85
  86                 inputUtxos := txInToUtxos(tx)
  87                 utxos := w.filterAccountUtxo(inputUtxos)
  88                 if err := batchSaveUtxos(utxos, batch); err != nil {
  89                         log.WithFields(log.Fields{"module": logModule, "err": err}).Error("detachUtxos fail on batchSaveUtxos")
  90                         return
  91                 }
  92         }
  93 }
  94
  95 func (w *Wallet) filterAccountUtxo(utxos []*account.UTXO) []*account.UTXO {
  96         outsByScript := make(map[string][]*account.UTXO, len(utxos))
  97         for _, utxo := range utxos {
  98                 scriptStr := string(utxo.ControlProgram)
  99                 outsByScript[scriptStr] = append(outsByScript[scriptStr], utxo)
 100         }
 101
 102         result := make([]*account.UTXO, 0, len(utxos))
 103         for s := range outsByScript {
 104                 if !segwit.IsP2WScript([]byte(s)) {
 105                         continue
 106                 }
 107
 108                 var hash [32]byte
 109                 sha3pool.Sum256(hash[:], []byte(s))
 110                 data := w.DB.Get(account.ContractKey(hash))
 111                 if data == nil {
 112                         continue
 113                 }
 114
 115                 cp := &account.CtrlProgram{}
 116                 if err := json.Unmarshal(data, cp); err != nil {
 117                         log.WithFields(log.Fields{"module": logModule, "err": err}).Error("filterAccountUtxo fail on unmarshal control program")
 118                         continue
 119                 }
 120
 121                 for _, utxo := range outsByScript[s] {
 122                         utxo.AccountID = cp.AccountID
 123                         utxo.Address = cp.Address
 124                         utxo.ControlProgramIndex = cp.KeyIndex
 125                         utxo.Change = cp.Change
 126                         result = append(result, utxo)
 127                 }
 128         }
 129         return result
 130 }
 131
 132 func batchSaveUtxos(utxos []*account.UTXO, batch dbm.Batch) error {
 133         for _, utxo := range utxos {
 134                 data, err := json.Marshal(utxo)
 135                 if err != nil {
 136                         return errors.Wrap(err, "failed marshal accountutxo")
 137                 }
 138
 139                 if segwit.IsP2WScript(utxo.ControlProgram) {
 140                         batch.Set(account.StandardUTXOKey(utxo.OutputID), data)
 141                 } else {
 142                         batch.Set(account.ContractUTXOKey(utxo.OutputID), data)
 143                 }
 144         }
 145         return nil
 146 }
 147
 148 func txInToUtxos(tx *types.Tx) []*account.UTXO {
 149         utxos := []*account.UTXO{}
 150         for _, inpID := range tx.Tx.InputIDs {
 151                 sp, err := tx.Spend(inpID)
 152                 if err != nil {
 153                         continue
 154                 }
 155
 156                 resOut, err := tx.Output(*sp.SpentOutputId)
 157                 if err != nil {
 158                         log.WithFields(log.Fields{"module": logModule, "err": err}).Error("txInToUtxos fail on get resOut")
 159                         continue
 160                 }
 161
 162                 utxos = append(utxos, &account.UTXO{
 163                         OutputID:       *sp.SpentOutputId,
 164                         AssetID:        *resOut.Source.Value.AssetId,
 165                         Amount:         resOut.Source.Value.Amount,
 166                         ControlProgram: resOut.ControlProgram.Code,
 167                         SourceID:       *resOut.Source.Ref,
 168                         SourcePos:      resOut.Source.Position,
 169                 })
 170         }
 171         return utxos
 172 }
 173
 174 func txOutToUtxos(tx *types.Tx, vaildHeight uint64) []*account.UTXO {
 175         utxos := []*account.UTXO{}
 176         for i, out := range tx.Outputs {
 177                 bcOut, err := tx.Output(*tx.ResultIds[i])
 178                 if err != nil {
 179                         continue
 180                 }
 181
 182                 utxo := &account.UTXO{
 183                         OutputID:       *tx.OutputID(i),
 184                         AssetID:        *out.AssetAmount.AssetId,
 185                         Amount:         out.Amount,
 186                         ControlProgram: out.ControlProgram,
 187                         SourceID:       *bcOut.Source.Ref,
 188                         SourcePos:      bcOut.Source.Position,
 189                         ValidHeight:    vaildHeight,
 190                 }
 191                 utxos = append(utxos, utxo)
 192         }
 193         return utxos
 194 }