wallet: Add new wallet constructors

Closes #3177

Signed-off-by: Evgeny Baydakov <e.bajdakov@gmail.com>
This commit is contained in:
Evgeny Baydakov 2023-12-08 13:58:48 +04:00
parent 441eb8aa86
commit 2012c56031
No known key found for this signature in database
GPG key ID: 8733EE3D72CDB4DE
2 changed files with 113 additions and 43 deletions

View file

@ -1,6 +1,7 @@
package wallet package wallet
import ( import (
"bytes"
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
@ -18,6 +19,13 @@ const (
walletVersion = "1.0" 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. // Wallet represents a NEO (NEP-2, NEP-6) compliant wallet.
type Wallet struct { type Wallet struct {
// Version of the wallet, used for later upgrades. // Version of the wallet, used for later upgrades.
@ -53,6 +61,12 @@ func NewWallet(location string) (*Wallet, error) {
return newWallet(file), nil 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. // NewWalletFromFile creates a Wallet from the given wallet file path.
func NewWalletFromFile(path string) (*Wallet, error) { func NewWalletFromFile(path string) (*Wallet, error) {
file, err := os.Open(path) file, err := os.Open(path)
@ -70,6 +84,20 @@ func NewWalletFromFile(path string) (*Wallet, error) {
return wall, nil 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 { func newWallet(rw io.ReadWriter) *Wallet {
var path string var path string
if f, ok := rw.(*os.File); ok { if f, ok := rw.(*os.File); ok {
@ -138,9 +166,15 @@ func (w *Wallet) Path() string {
return w.path return w.path
} }
// Save saves the wallet data. It's the internal io.ReadWriter // SetPath sets the location of the wallet on the filesystem.
// that is responsible for saving the data. This can func (w *Wallet) SetPath(path string) {
// be a buffer, file, etc.. 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 { func (w *Wallet) Save() error {
data, err := json.Marshal(w) data, err := json.Marshal(w)
if err != nil { if err != nil {
@ -151,6 +185,8 @@ func (w *Wallet) Save() error {
} }
// SavePretty saves the wallet in a beautiful JSON. // SavePretty saves the wallet in a beautiful JSON.
//
// Returns [ErrPathIsEmpty] if wallet path is not set. See [Wallet.SetPath].
func (w *Wallet) SavePretty() error { func (w *Wallet) SavePretty() error {
data, err := json.MarshalIndent(w, "", " ") data, err := json.MarshalIndent(w, "", " ")
if err != nil { if err != nil {
@ -161,6 +197,10 @@ func (w *Wallet) SavePretty() error {
} }
func (w *Wallet) writeRaw(data []byte) error { func (w *Wallet) writeRaw(data []byte) error {
if w.path == "" {
return ErrPathIsEmpty
}
return os.WriteFile(w.path, data, 0644) return os.WriteFile(w.path, data, 0644)
} }

View file

@ -47,24 +47,29 @@ func TestCreateAccountAndClose(t *testing.T) {
} }
func TestAddAccount(t *testing.T) { func TestAddAccount(t *testing.T) {
wallet := checkWalletConstructor(t) wallets := []*Wallet{
checkWalletConstructor(t),
NewInMemoryWallet(),
}
wallet.AddAccount(&Account{ for _, w := range wallets {
privateKey: nil, w.AddAccount(&Account{
Address: "real", privateKey: nil,
EncryptedWIF: "", Address: "real",
Label: "", EncryptedWIF: "",
Contract: nil, Label: "",
Locked: false, Contract: nil,
Default: false, Locked: false,
}) Default: false,
accounts := wallet.Accounts })
require.Len(t, accounts, 1) accounts := w.Accounts
require.Len(t, accounts, 1)
require.Error(t, wallet.RemoveAccount("abc")) require.Error(t, w.RemoveAccount("abc"))
require.Len(t, wallet.Accounts, 1) require.Len(t, w.Accounts, 1)
require.NoError(t, wallet.RemoveAccount("real")) require.NoError(t, w.RemoveAccount("real"))
require.Len(t, wallet.Accounts, 0) require.Len(t, w.Accounts, 0)
}
} }
func TestPath(t *testing.T) { func TestPath(t *testing.T) {
@ -75,36 +80,47 @@ func TestPath(t *testing.T) {
} }
func TestSave(t *testing.T) { func TestSave(t *testing.T) {
wallet := checkWalletConstructor(t) inMemWallet := NewInMemoryWallet()
wallet.AddAccount(&Account{ tmpDir := t.TempDir()
privateKey: nil, file := filepath.Join(tmpDir, walletTemplate)
Address: "", inMemWallet.SetPath(file)
EncryptedWIF: "",
Label: "",
Contract: nil,
Locked: false,
Default: false,
})
errForSave := wallet.Save() wallets := []*Wallet{
require.NoError(t, errForSave) checkWalletConstructor(t),
inMemWallet,
}
openedWallet, err := NewWalletFromFile(wallet.path) for _, w := range wallets {
require.NoError(t, err) w.AddAccount(&Account{
require.Equal(t, wallet.Accounts, openedWallet.Accounts) privateKey: nil,
Address: "",
EncryptedWIF: "",
Label: "",
Contract: nil,
Locked: false,
Default: false,
})
t.Run("change and rewrite", func(t *testing.T) { errForSave := w.Save()
err := openedWallet.CreateAccount("test", "pass") require.NoError(t, errForSave)
openedWallet, err := NewWalletFromFile(w.path)
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, w.Accounts, openedWallet.Accounts)
w2, err := NewWalletFromFile(openedWallet.path) t.Run("change and rewrite", func(t *testing.T) {
require.NoError(t, err) err := openedWallet.CreateAccount("test", "pass")
require.Equal(t, 2, len(w2.Accounts)) require.NoError(t, err)
require.NoError(t, w2.Accounts[1].Decrypt("pass", w2.Scrypt))
_ = w2.Accounts[1].ScriptHash() // openedWallet has it for acc 1. w2, err := NewWalletFromFile(openedWallet.path)
require.Equal(t, openedWallet.Accounts, w2.Accounts) 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) { 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 // 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") 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)
}