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

wallet: Add new wallet constructors

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

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