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")
  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         reg := asset.NewRegistry(testDB, chain)
  68         asset, err := reg.Define([]chainkd.XPub{xpub1.XPub}, 1, nil, "TESTASSET")
  69         if err != nil {
  70                 t.Fatal(err)
  71         }
  72
  73         utxos := []*account.UTXO{}
  74         btmUtxo := mockUTXO(controlProg, consensus.BTMAssetID)
  75         utxos = append(utxos, btmUtxo)
  76         OtherUtxo := mockUTXO(controlProg, &asset.AssetID)
  77         utxos = append(utxos, OtherUtxo)
  78
  79         _, txData, err := mockTxData(utxos, testAccount)
  80         if err != nil {
  81                 t.Fatal(err)
  82         }
  83
  84         tx := types.NewTx(*txData)
  85         block := mockSingleBlock(tx)
  86         txStatus := bc.NewTransactionStatus()
  87         store.SaveBlock(block, txStatus)
  88
  89         w := mockWallet(testDB, accountManager, reg, chain)
  90         err = w.AttachBlock(block)
  91         if err != nil {
  92                 t.Fatal(err)
  93         }
  94
  95         want, err := w.GetTransactionsByTxID(tx.ID.String())
  96         if len(want) != 1 {
  97                 t.Fatal(err)
  98         }
  99
 100         wants, err := w.GetTransactionsByTxID("")
 101         if len(wants) != 1 {
 102                 t.Fatal(err)
 103         }
 104 }
 105
 106 func TestExportAndImportPrivKey(t *testing.T) {
 107         dirPath, err := ioutil.TempDir(".", "")
 108         if err != nil {
 109                 t.Fatal(err)
 110         }
 111         defer os.RemoveAll(dirPath)
 112
 113         testDB := dbm.NewDB("testdb", "leveldb", "temp")
 114         defer os.RemoveAll("temp")
 115
 116         store := leveldb.NewStore(testDB)
 117         txPool := protocol.NewTxPool()
 118
 119         chain, err := protocol.NewChain(store, txPool)
 120         if err != nil {
 121                 t.Fatal(err)
 122         }
 123
 124         acntManager := account.NewManager(testDB, chain)
 125         reg := asset.NewRegistry(testDB, chain)
 126
 127         hsm, err := pseudohsm.New(dirPath)
 128         if err != nil {
 129                 t.Fatal(err)
 130         }
 131
 132         pwd := "password"
 133         xpub, err := hsm.XCreate("alias", pwd)
 134         if err != nil {
 135                 t.Fatal(err)
 136         }
 137
 138         w, err := NewWallet(testDB, acntManager, reg, hsm, chain)
 139         if err != nil {
 140                 t.Fatal(err)
 141         }
 142
 143         ctx := context.Background()
 144         acnt1, err := w.AccountMgr.Create(ctx, []chainkd.XPub{xpub.XPub}, 1, "account-alias")
 145         if err != nil {
 146                 t.Fatal(err)
 147         }
 148
 149         priv, err := w.ExportAccountPrivKey(xpub.XPub, pwd)
 150
 151         wantPriv, err := hsm.LoadChainKDKey(xpub.XPub, pwd)
 152         if err != nil {
 153                 t.Fatal(err)
 154         }
 155         var hashed [32]byte
 156         sha3pool.Sum256(hashed[:], wantPriv[:])
 157
 158         tmp := append(wantPriv[:], hashed[:4]...)
 159         res := base58.Encode(tmp)
 160
 161         if res != *priv {
 162                 t.Fatalf("XPrivs should be identical.\nBefore: %v\n After: %v\n", *priv, res)
 163         }
 164
 165         rawPriv, err := base58.Decode(*priv)
 166         if err != nil {
 167                 t.Fatal(err)
 168         }
 169
 170         if len(rawPriv) != 68 {
 171                 t.Fatal("invalid private key hash length")
 172         }
 173
 174         var xprv [64]byte
 175         copy(xprv[:], rawPriv[:64])
 176
 177         _, err = w.ImportAccountPrivKey(xprv, xpub.Alias, pwd, 0, acnt1.Alias)
 178         if err != pseudohsm.ErrDuplicateKeyAlias {
 179                 t.Fatal(err)
 180         }
 181
 182         hsm.XDelete(xpub.XPub, pwd)
 183
 184         _, err = w.ImportAccountPrivKey(xprv, xpub.Alias, pwd, 0, acnt1.Alias)
 185         if err != account.ErrDuplicateAlias {
 186                 t.Fatal(err)
 187         }
 188
 189         w.AccountMgr.DeleteAccount(acnt1.Alias)
 190
 191         acnt2, err := w.ImportAccountPrivKey(xprv, xpub.Alias, pwd, 0, acnt1.Alias)
 192         if err != nil {
 193                 t.Fatal(err)
 194         }
 195
 196         if acnt2.XPub != acnt1.XPubs[0] {
 197                 t.Fatalf("XPubs should be identical.\nBefore: %v\n After: %v\n", acnt1.XPubs[0], acnt2.XPub)
 198         }
 199 }
 200
 201 func mockUTXO(controlProg *account.CtrlProgram, assetID *bc.AssetID) *account.UTXO {
 202         utxo := &account.UTXO{}
 203         utxo.OutputID = bc.Hash{V0: 1}
 204         utxo.SourceID = bc.Hash{V0: 2}
 205         utxo.AssetID = *assetID
 206         utxo.Amount = 1000000000
 207         utxo.SourcePos = 0
 208         utxo.ControlProgram = controlProg.ControlProgram
 209         utxo.AccountID = controlProg.AccountID
 210         utxo.Address = controlProg.Address
 211         utxo.ControlProgramIndex = controlProg.KeyIndex
 212         return utxo
 213 }
 214
 215 func mockTxData(utxos []*account.UTXO, testAccount *account.Account) (*txbuilder.Template, *types.TxData, error) {
 216         tplBuilder := txbuilder.NewBuilder(time.Now())
 217
 218         for _, utxo := range utxos {
 219                 txInput, sigInst, err := account.UtxoToInputs(testAccount.Signer, utxo)
 220                 if err != nil {
 221                         return nil, nil, err
 222                 }
 223                 tplBuilder.AddInput(txInput, sigInst)
 224
 225                 out := &types.TxOutput{}
 226                 if utxo.AssetID == *consensus.BTMAssetID {
 227                         out = types.NewTxOutput(utxo.AssetID, 100, utxo.ControlProgram)
 228                 } else {
 229                         out = types.NewTxOutput(utxo.AssetID, utxo.Amount, utxo.ControlProgram)
 230                 }
 231                 tplBuilder.AddOutput(out)
 232         }
 233
 234         return tplBuilder.Build()
 235 }
 236
 237 func mockWallet(walletDB dbm.DB, account *account.Manager, asset *asset.Registry, chain *protocol.Chain) *Wallet {
 238         wallet := &Wallet{
 239                 DB:                  walletDB,
 240                 AccountMgr:          account,
 241                 AssetReg:            asset,
 242                 chain:               chain,
 243                 rescanProgress:      make(chan struct{}, 1),
 244         }
 245         wallet.status = StatusInfo{
 246                 OnChainAddresses: NewAddressSet(),
 247         }
 248         return wallet
 249 }
 250
 251 func mockSingleBlock(tx *types.Tx) *types.Block {
 252         return &types.Block{
 253                 BlockHeader: types.BlockHeader{
 254                         Version: 1,
 255                         Height:  1,
 256                         Bits:    2305843009230471167,
 257                 },
 258                 Transactions: []*types.Tx{tx},
 259         }
 260 }