From 2012c56031c902e5c12bfb8d543ffd2276b45b12 Mon Sep 17 00:00:00 2001 From: Evgeny Baydakov Date: Fri, 8 Dec 2023 13:58:48 +0400 Subject: [PATCH] wallet: Add new wallet constructors Closes #3177 Signed-off-by: Evgeny Baydakov --- pkg/wallet/wallet.go | 46 ++++++++++++++-- pkg/wallet/wallet_test.go | 110 ++++++++++++++++++++++++-------------- 2 files changed, 113 insertions(+), 43 deletions(-) diff --git a/pkg/wallet/wallet.go b/pkg/wallet/wallet.go index ade1d3146..b509fd083 100644 --- a/pkg/wallet/wallet.go +++ b/pkg/wallet/wallet.go @@ -1,6 +1,7 @@ package wallet import ( + "bytes" "encoding/json" "errors" "fmt" @@ -18,6 +19,13 @@ const ( walletVersion = "1.0" ) +var ( + // ErrPathIsEmpty appears if wallet was created without linking to file system path, + // for instance with [NewInMemoryWallet] or [NewWalletFromBytes]. + // Despite this, there was an attempt to save it via [Wallet.Save] or [Wallet.SavePretty] without [Wallet.SetPath]. + ErrPathIsEmpty = errors.New("path is empty") +) + // Wallet represents a NEO (NEP-2, NEP-6) compliant wallet. type Wallet struct { // Version of the wallet, used for later upgrades. @@ -53,6 +61,12 @@ func NewWallet(location string) (*Wallet, error) { return newWallet(file), nil } +// NewInMemoryWallet creates a new NEO wallet without linking to the read file on file system. +// If wallet required to be written to the file system, [Wallet.SetPath] should be used to set the path. +func NewInMemoryWallet() *Wallet { + return newWallet(nil) +} + // NewWalletFromFile creates a Wallet from the given wallet file path. func NewWalletFromFile(path string) (*Wallet, error) { file, err := os.Open(path) @@ -70,6 +84,20 @@ func NewWalletFromFile(path string) (*Wallet, error) { return wall, nil } +// NewWalletFromBytes creates a [Wallet] from the given byte slice. +// Parameter wallet contains JSON representation of wallet, see [Wallet.JSON] for details. +// +// NewWalletFromBytes constructor doesn't set wallet's path. If you want to save the wallet to file system, +// use [Wallet.SetPath]. +func NewWalletFromBytes(wallet []byte) (*Wallet, error) { + wall := &Wallet{} + if err := json.NewDecoder(bytes.NewReader(wallet)).Decode(wall); err != nil { + return nil, fmt.Errorf("unmarshal wallet: %w", err) + } + + return wall, nil +} + func newWallet(rw io.ReadWriter) *Wallet { var path string if f, ok := rw.(*os.File); ok { @@ -138,9 +166,15 @@ func (w *Wallet) Path() string { return w.path } -// Save saves the wallet data. It's the internal io.ReadWriter -// that is responsible for saving the data. This can -// be a buffer, file, etc.. +// SetPath sets the location of the wallet on the filesystem. +func (w *Wallet) SetPath(path string) { + w.path = path +} + +// Save saves the wallet data to the file located at the path that was either provided +// via [NewWalletFromFile] constructor or via [Wallet.SetPath]. +// +// Returns [ErrPathIsEmpty] if wallet path is not set. See [Wallet.SetPath]. func (w *Wallet) Save() error { data, err := json.Marshal(w) if err != nil { @@ -151,6 +185,8 @@ func (w *Wallet) Save() error { } // SavePretty saves the wallet in a beautiful JSON. +// +// Returns [ErrPathIsEmpty] if wallet path is not set. See [Wallet.SetPath]. func (w *Wallet) SavePretty() error { data, err := json.MarshalIndent(w, "", " ") if err != nil { @@ -161,6 +197,10 @@ func (w *Wallet) SavePretty() error { } func (w *Wallet) writeRaw(data []byte) error { + if w.path == "" { + return ErrPathIsEmpty + } + return os.WriteFile(w.path, data, 0644) } diff --git a/pkg/wallet/wallet_test.go b/pkg/wallet/wallet_test.go index 9de064319..816269db0 100644 --- a/pkg/wallet/wallet_test.go +++ b/pkg/wallet/wallet_test.go @@ -47,24 +47,29 @@ func TestCreateAccountAndClose(t *testing.T) { } func TestAddAccount(t *testing.T) { - wallet := checkWalletConstructor(t) + wallets := []*Wallet{ + checkWalletConstructor(t), + NewInMemoryWallet(), + } - wallet.AddAccount(&Account{ - privateKey: nil, - Address: "real", - EncryptedWIF: "", - Label: "", - Contract: nil, - Locked: false, - Default: false, - }) - accounts := wallet.Accounts - require.Len(t, accounts, 1) + for _, w := range wallets { + w.AddAccount(&Account{ + privateKey: nil, + Address: "real", + EncryptedWIF: "", + Label: "", + Contract: nil, + Locked: false, + Default: false, + }) + accounts := w.Accounts + require.Len(t, accounts, 1) - require.Error(t, wallet.RemoveAccount("abc")) - require.Len(t, wallet.Accounts, 1) - require.NoError(t, wallet.RemoveAccount("real")) - require.Len(t, wallet.Accounts, 0) + require.Error(t, w.RemoveAccount("abc")) + require.Len(t, w.Accounts, 1) + require.NoError(t, w.RemoveAccount("real")) + require.Len(t, w.Accounts, 0) + } } func TestPath(t *testing.T) { @@ -75,36 +80,47 @@ func TestPath(t *testing.T) { } func TestSave(t *testing.T) { - wallet := checkWalletConstructor(t) + inMemWallet := NewInMemoryWallet() - wallet.AddAccount(&Account{ - privateKey: nil, - Address: "", - EncryptedWIF: "", - Label: "", - Contract: nil, - Locked: false, - Default: false, - }) + tmpDir := t.TempDir() + file := filepath.Join(tmpDir, walletTemplate) + inMemWallet.SetPath(file) - errForSave := wallet.Save() - require.NoError(t, errForSave) + wallets := []*Wallet{ + checkWalletConstructor(t), + inMemWallet, + } - openedWallet, err := NewWalletFromFile(wallet.path) - require.NoError(t, err) - require.Equal(t, wallet.Accounts, openedWallet.Accounts) + for _, w := range wallets { + w.AddAccount(&Account{ + privateKey: nil, + Address: "", + EncryptedWIF: "", + Label: "", + Contract: nil, + Locked: false, + Default: false, + }) - t.Run("change and rewrite", func(t *testing.T) { - err := openedWallet.CreateAccount("test", "pass") + errForSave := w.Save() + require.NoError(t, errForSave) + + openedWallet, err := NewWalletFromFile(w.path) require.NoError(t, err) + require.Equal(t, w.Accounts, openedWallet.Accounts) - w2, err := NewWalletFromFile(openedWallet.path) - require.NoError(t, err) - require.Equal(t, 2, len(w2.Accounts)) - require.NoError(t, w2.Accounts[1].Decrypt("pass", w2.Scrypt)) - _ = w2.Accounts[1].ScriptHash() // openedWallet has it for acc 1. - require.Equal(t, openedWallet.Accounts, w2.Accounts) - }) + t.Run("change and rewrite", func(t *testing.T) { + err := openedWallet.CreateAccount("test", "pass") + require.NoError(t, err) + + w2, err := NewWalletFromFile(openedWallet.path) + require.NoError(t, err) + require.Equal(t, 2, len(w2.Accounts)) + require.NoError(t, w2.Accounts[1].Decrypt("pass", w2.Scrypt)) + _ = w2.Accounts[1].ScriptHash() // openedWallet has it for acc 1. + require.Equal(t, openedWallet.Accounts, w2.Accounts) + }) + } } func TestJSONMarshallUnmarshal(t *testing.T) { @@ -198,3 +214,17 @@ func TestWalletForExamples(t *testing.T) { // we need to keep the owner of the example contracts the same as the wallet account require.Equal(t, "NbrUYaZgyhSkNoRo9ugRyEMdUZxrhkNaWB", w.Accounts[0].Address, "need to change `owner` in the example contracts") } + +func TestFromBytes(t *testing.T) { + wallet := checkWalletConstructor(t) + bts, err := wallet.JSON() + require.NoError(t, err) + + w, err := NewWalletFromBytes(bts) + require.NoError(t, err) + + require.Len(t, w.path, 0) + w.SetPath(wallet.path) + + require.Equal(t, wallet, w) +}