wallet/utxo.go

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