aarifullin/neoneo-go
1
0
Fork 0
forked from TrueCloudLab/neoneo-go
Code Pull requests Activity

Merge pull request #3255 from smallhive/3177-wallet-from-bytes

Browse source 
wallet: Add new wallet constructors
This commit is contained in:
Roman Khimov 2023-12-14 18:45:44 +03:00 committed by GitHub
commit 3766206f44
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 113 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
pkg/wallet/wallet.go
View file

@ -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)
}

110
pkg/wallet/wallet_test.go
View file

@ -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)
}