update GetControlProgram

[bytom/vapor.git] / wallet / annotated.go

diff --git a/wallet/annotated.go b/wallet/annotated.go
index 6604876..b146e4a 100644 (file)
--- a/wallet/annotated.go
+++ b/wallet/annotated.go
@@ -13,7 +13,6 @@ import (
        "github.com/vapor/consensus"
        "github.com/vapor/consensus/segwit"
        "github.com/vapor/crypto/sha3pool"
-       dbm "github.com/vapor/database/leveldb"
        "github.com/vapor/protocol/bc"
        "github.com/vapor/protocol/bc/types"
 )
@@ -33,28 +32,18 @@ func annotateTxsAsset(w *Wallet, txs []*query.AnnotatedTx) {
 }
 
 func (w *Wallet) getExternalDefinition(assetID *bc.AssetID) json.RawMessage {
-       definitionByte := w.DB.Get(asset.ExtAssetKey(assetID))
-       if definitionByte == nil {
-               return nil
+       externalAsset, err := w.store.GetAssetDefinition(assetID)
+       if err != nil {
+               log.WithFields(log.Fields{"module": logModule, "err": err}).Warning("fail on get asset definition.")
        }
-
-       definitionMap := make(map[string]interface{})
-       if err := json.Unmarshal(definitionByte, &definitionMap); err != nil {
+       if externalAsset == nil {
                return nil
        }
 
-       alias := assetID.String()
-       externalAsset := &asset.Asset{
-               AssetID:           *assetID,
-               Alias:             &alias,
-               DefinitionMap:     definitionMap,
-               RawDefinitionByte: definitionByte,
+       if err := w.AssetReg.SaveAsset(externalAsset, *externalAsset.Alias); err != nil {
+               log.WithFields(log.Fields{"module": logModule, "err": err, "assetAlias": *externalAsset.Alias}).Warning("fail on save external asset to internal asset DB")
        }
-
-       if err := w.AssetReg.SaveAsset(externalAsset, alias); err != nil {
-               log.WithFields(log.Fields{"module": logModule, "err": err, "assetID": alias}).Warning("fail on save external asset to internal asset DB")
-       }
-       return definitionByte
+       return json.RawMessage(externalAsset.RawDefinitionByte)
 }
 
 func (w *Wallet) getAliasDefinition(assetID bc.AssetID) (string, json.RawMessage) {
@@ -82,14 +71,14 @@ func (w *Wallet) getAliasDefinition(assetID bc.AssetID) (string, json.RawMessage
 }
 
 // annotateTxs adds account data to transactions
-func annotateTxsAccount(txs []*query.AnnotatedTx, walletDB dbm.DB) {
+func annotateTxsAccount(txs []*query.AnnotatedTx, store WalletStorer) {
        for i, tx := range txs {
                for j, input := range tx.Inputs {
                        //issue asset tx input SpentOutputID is nil
                        if input.SpentOutputID == nil {
                                continue
                        }
-                       localAccount, err := getAccountFromACP(input.ControlProgram, walletDB)
+                       localAccount, err := getAccountFromACP(input.ControlProgram, store)
                        if localAccount == nil || err != nil {
                                continue
                        }
@@ -97,7 +86,7 @@ func annotateTxsAccount(txs []*query.AnnotatedTx, walletDB dbm.DB) {
                        txs[i].Inputs[j].AccountID = localAccount.ID
                }
                for j, output := range tx.Outputs {
-                       localAccount, err := getAccountFromACP(output.ControlProgram, walletDB)
+                       localAccount, err := getAccountFromACP(output.ControlProgram, store)
                        if localAccount == nil || err != nil {
                                continue
                        }
@@ -107,23 +96,20 @@ func annotateTxsAccount(txs []*query.AnnotatedTx, walletDB dbm.DB) {
        }
 }
 
-func getAccountFromACP(program []byte, walletDB dbm.DB) (*account.Account, error) {
+func getAccountFromACP(program []byte, store WalletStorer) (*account.Account, error) {
        var hash common.Hash
-       accountCP := account.CtrlProgram{}
        localAccount := account.Account{}
 
        sha3pool.Sum256(hash[:], program)
-
-       rawProgram := walletDB.Get(account.ContractKey(hash))
-       if rawProgram == nil {
-               return nil, fmt.Errorf("failed get account control program:%x ", hash)
-       }
-
-       if err := json.Unmarshal(rawProgram, &accountCP); err != nil {
+       accountCP, err := store.GetControlProgram(hash)
+       if err != nil {
                return nil, err
        }
+       if accountCP == nil {
+               return nil, fmt.Errorf("failed get account control program:%x ", hash)
+       }
 
-       accountValue := walletDB.Get(account.Key(accountCP.AccountID))
+       accountValue := store.GetAccountByAccountID(accountCP.AccountID)
        if accountValue == nil {
                return nil, fmt.Errorf("failed get account:%s ", accountCP.AccountID)
        }
@@ -219,18 +205,18 @@ func (w *Wallet) getAddressFromControlProgram(prog []byte, isMainchain bool) str
 
        if segwit.IsP2WPKHScript(prog) {
                if pubHash, err := segwit.GetHashFromStandardProg(prog); err == nil {
-                       return buildP2PKHAddress(pubHash, netParams)
+                       return BuildP2PKHAddress(pubHash, netParams)
                }
        } else if segwit.IsP2WSHScript(prog) {
                if scriptHash, err := segwit.GetHashFromStandardProg(prog); err == nil {
-                       return buildP2SHAddress(scriptHash, netParams)
+                       return BuildP2SHAddress(scriptHash, netParams)
                }
        }
 
        return ""
 }
 
-func buildP2PKHAddress(pubHash []byte, netParams *consensus.Params) string {
+func BuildP2PKHAddress(pubHash []byte, netParams *consensus.Params) string {
        address, err := common.NewAddressWitnessPubKeyHash(pubHash, netParams)
        if err != nil {
                return ""
@@ -239,7 +225,7 @@ func buildP2PKHAddress(pubHash []byte, netParams *consensus.Params) string {
        return address.EncodeAddress()
 }
 
-func buildP2SHAddress(scriptHash []byte, netParams *consensus.Params) string {
+func BuildP2SHAddress(scriptHash []byte, netParams *consensus.Params) string {
        address, err := common.NewAddressWitnessScriptHash(scriptHash, netParams)
        if err != nil {
                return ""