blockchain/wallet/wallet_test.go

   1 package wallet
   2
   3 import (
   4         "context"
   5         "io/ioutil"
   6         "os"
   7         "testing"
   8         "time"
   9
  10         "github.com/tendermint/go-wire/data/base58"
  11         dbm "github.com/tendermint/tmlibs/db"
  12
  13         "github.com/bytom/blockchain/account"
  14         "github.com/bytom/blockchain/asset"
  15         "github.com/bytom/blockchain/pseudohsm"
  16         "github.com/bytom/blockchain/txbuilder"
  17         "github.com/bytom/blockchain/txdb"
  18         cfg "github.com/bytom/config"
  19         "github.com/bytom/consensus"
  20         "github.com/bytom/crypto/ed25519/chainkd"
  21         "github.com/bytom/crypto/sha3pool"
  22         "github.com/bytom/protocol"
  23         "github.com/bytom/protocol/bc"
  24         "github.com/bytom/protocol/bc/legacy"
  25 )
  26
  27 func TestWalletUpdate(t *testing.T) {
  28         dirPath, err := ioutil.TempDir(".", "")
  29         if err != nil {
  30                 t.Fatal(err)
  31         }
  32         defer os.RemoveAll(dirPath)
  33
  34         testDB := dbm.NewDB("testdb", "leveldb", "temp")
  35         defer os.RemoveAll("temp")
  36
  37         store := txdb.NewStore(testDB)
  38         txPool := protocol.NewTxPool()
  39
  40         chain, err := protocol.NewChain(bc.Hash{}, store, txPool)
  41         if err != nil {
  42                 t.Fatal(err)
  43         }
  44
  45         accountManager := account.NewManager(testDB, chain)
  46         hsm, err := pseudohsm.New(dirPath)
  47         if err != nil {
  48                 t.Fatal(err)
  49         }
  50
  51         xpub1, err := hsm.XCreate("test_pub1", "password")
  52         if err != nil {
  53                 t.Fatal(err)
  54         }
  55
  56         testAccount, err := accountManager.Create(nil, []chainkd.XPub{xpub1.XPub}, 1, "testAccount", nil)
  57         if err != nil {
  58                 t.Fatal(err)
  59         }
  60
  61         controlProg, err := accountManager.CreateAddress(nil, testAccount.ID, false)
  62         if err != nil {
  63                 t.Fatal(err)
  64         }
  65
  66         controlProg.KeyIndex = 1
  67
  68         utxo := mockUTXO(controlProg)
  69         _, txData, err := mockTxData(utxo, testAccount)
  70         if err != nil {
  71                 t.Fatal(err)
  72         }
  73
  74         tx := legacy.NewTx(*txData)
  75
  76         reg := asset.NewRegistry(testDB, chain)
  77
  78         w := mockWallet(testDB, accountManager, reg, chain)
  79
  80         block := mockSingleBlock(tx)
  81
  82         err = w.attachBlock(block)
  83         if err != nil {
  84                 t.Fatal(err)
  85         }
  86
  87         want, err := w.GetTransactionsByTxID(tx.ID.String())
  88         if len(want) != 1 {
  89                 t.Fatal(err)
  90         }
  91
  92         wants, err := w.GetTransactionsByTxID("")
  93         if len(wants) != 1 {
  94                 t.Fatal(err)
  95         }
  96 }
  97
  98 func TestExportAndImportPrivKey(t *testing.T) {
  99         dirPath, err := ioutil.TempDir(".", "")
 100         if err != nil {
 101                 t.Fatal(err)
 102         }
 103         defer os.RemoveAll(dirPath)
 104
 105         testDB := dbm.NewDB("testdb", "leveldb", "temp")
 106         defer os.RemoveAll("temp")
 107
 108         store := txdb.NewStore(testDB)
 109         txPool := protocol.NewTxPool()
 110
 111         chain, err := protocol.NewChain(bc.Hash{}, store, txPool)
 112         if err != nil {
 113                 t.Fatal(err)
 114         }
 115
 116         genesisBlock := cfg.GenerateGenesisBlock()
 117
 118         chain.SaveBlock(genesisBlock)
 119         chain.ConnectBlock(genesisBlock)
 120
 121         acntManager := account.NewManager(testDB, chain)
 122         reg := asset.NewRegistry(testDB, chain)
 123
 124         hsm, err := pseudohsm.New(dirPath)
 125         if err != nil {
 126                 t.Fatal(err)
 127         }
 128
 129         pwd := "password"
 130         xpub, err := hsm.XCreate("alias", pwd)
 131         if err != nil {
 132                 t.Fatal(err)
 133         }
 134
 135         w, err := NewWallet(testDB, acntManager, reg, chain, nil)
 136         if err != nil {
 137                 t.Fatal(err)
 138         }
 139
 140         ctx := context.Background()
 141         acnt1, err := w.AccountMgr.Create(ctx, []chainkd.XPub{xpub.XPub}, 1, "account-alias", nil)
 142         if err != nil {
 143                 t.Fatal(err)
 144         }
 145
 146         priv, err := w.ExportAccountPrivKey(hsm, xpub.XPub, pwd)
 147
 148         wantPriv, err := hsm.LoadChainKDKey(xpub.XPub, pwd)
 149         if err != nil {
 150                 t.Fatal(err)
 151         }
 152         var hashed [32]byte
 153         sha3pool.Sum256(hashed[:], wantPriv[:])
 154
 155         tmp := append(wantPriv[:], hashed[:4]...)
 156         res := base58.Encode(tmp)
 157
 158         if res != *priv {
 159                 t.Fatalf("XPrivs should be identical.\nBefore: %v\n After: %v\n", *priv, res)
 160         }
 161
 162         rawPriv, err := base58.Decode(*priv)
 163         if err != nil {
 164                 t.Fatal(err)
 165         }
 166
 167         if len(rawPriv) != 68 {
 168                 t.Fatal("invalid private key hash length")
 169         }
 170
 171         var xprv [64]byte
 172         copy(xprv[:], rawPriv[:64])
 173
 174         _, err = w.ImportAccountPrivKey(hsm, xprv, xpub.Alias, pwd, 0, acnt1.Alias)
 175         if err != pseudohsm.ErrDuplicateKeyAlias {
 176                 t.Fatal(err)
 177         }
 178
 179         hsm.XDelete(xpub.XPub, pwd)
 180
 181         _, err = w.ImportAccountPrivKey(hsm, xprv, xpub.Alias, pwd, 0, acnt1.Alias)
 182         if err != account.ErrDuplicateAlias {
 183                 t.Fatal(err)
 184         }
 185
 186         accountInfo := struct {
 187                 AccountInfo string `json:"account_info"`
 188         }{AccountInfo: acnt1.Alias}
 189
 190         w.AccountMgr.DeleteAccount(accountInfo)
 191
 192         acnt2, err := w.ImportAccountPrivKey(hsm, xprv, xpub.Alias, pwd, 0, acnt1.Alias)
 193         if err != nil {
 194                 t.Fatal(err)
 195         }
 196
 197         if acnt2.XPub != acnt1.XPubs[0] {
 198                 t.Fatalf("XPubs should be identical.\nBefore: %v\n After: %v\n", acnt1.XPubs[0], acnt2.XPub)
 199         }
 200 }
 201
 202 func mockUTXO(controlProg *account.CtrlProgram) *account.UTXO {
 203         utxo := &account.UTXO{}
 204         utxo.OutputID = bc.Hash{V0: 1}
 205         utxo.SourceID = bc.Hash{V0: 2}
 206         utxo.AssetID = *consensus.BTMAssetID
 207         utxo.Amount = 1000000000
 208         utxo.SourcePos = 0
 209         utxo.ControlProgram = controlProg.ControlProgram
 210         utxo.AccountID = controlProg.AccountID
 211         utxo.Address = controlProg.Address
 212         utxo.ControlProgramIndex = controlProg.KeyIndex
 213         return utxo
 214 }
 215
 216 func mockTxData(utxo *account.UTXO, testAccount *account.Account) (*txbuilder.Template, *legacy.TxData, error) {
 217         txInput, sigInst, err := account.UtxoToInputs(testAccount.Signer, utxo, nil)
 218         if err != nil {
 219                 return nil, nil, err
 220         }
 221
 222         b := txbuilder.NewBuilder(time.Now())
 223         b.AddInput(txInput, sigInst)
 224         out := legacy.NewTxOutput(*consensus.BTMAssetID, 100, utxo.ControlProgram, nil)
 225         b.AddOutput(out)
 226         return b.Build()
 227 }
 228
 229 func mockWallet(walletDB dbm.DB, account *account.Manager, asset *asset.Registry, chain *protocol.Chain) *Wallet {
 230         return &Wallet{
 231                 DB:             walletDB,
 232                 AccountMgr:     account,
 233                 AssetReg:       asset,
 234                 chain:          chain,
 235                 rescanProgress: make(chan struct{}, 1),
 236         }
 237 }
 238
 239 func mockSingleBlock(tx *legacy.Tx) *legacy.Block {
 240         return &legacy.Block{
 241                 BlockHeader: legacy.BlockHeader{
 242                         Version: 1,
 243                         Height:  1,
 244                         Bits:    2305843009230471167,
 245                 },
 246                 Transactions: []*legacy.Tx{tx},
 247         }
 248 }