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