TrueCloudLab/neo-go
8
0
Fork 0
mirror of https://github.com/nspcc-dev/neo-go.git synced 2024-11-22 19:29:39 +00:00
Code Issues Projects Releases Packages Wiki Activity

Merge pull request #2099 from nspcc-dev/wallet-truncate

Browse source 
wallet: truncate file after writing
This commit is contained in:
Roman Khimov 2021-07-29 18:52:18 +03:00 committed by GitHub
commit ebbb9df91e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 58 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
cli/wallet_test.go
View file

@ -19,6 +19,38 @@ import (
"github.com/stretchr/testify/require"
)
func TestWalletAccountRemove(t *testing.T) {
tmpDir, err := ioutil.TempDir("", "neogo.test.walletinit")
require.NoError(t, err)
t.Cleanup(func() {
os.RemoveAll(tmpDir)
})
e := newExecutor(t, false)
walletPath := path.Join(tmpDir, "wallet.json")
e.In.WriteString("acc1\r")
e.In.WriteString("pass\r")
e.In.WriteString("pass\r")
e.Run(t, "neo-go", "wallet", "init", "--wallet", walletPath, "--account")
e.In.WriteString("acc2\r")
e.In.WriteString("pass\r")
e.In.WriteString("pass\r")
e.Run(t, "neo-go", "wallet", "create", "--wallet", walletPath)
w, err := wallet.NewWalletFromFile(walletPath)
require.NoError(t, err)
addr := w.Accounts[0].Address
e.Run(t, "neo-go", "wallet", "remove", "--wallet", walletPath,
"--address", addr, "--force")
rawWallet, err := ioutil.ReadFile(walletPath)
require.NoError(t, err)
require.NoError(t, json.Unmarshal(rawWallet, new(wallet.Wallet)))
}
func TestWalletInit(t *testing.T) {
tmpDir, err := ioutil.TempDir("", "neogo.test.walletinit")
require.NoError(t, err)

32
pkg/wallet/wallet.go
View file

@ -143,25 +143,45 @@ func (w *Wallet) Path() string {
// that is responsible for saving the data. This can
// be a buffer, file, etc..
func (w *Wallet) Save() error {
if err := w.rewind(); err != nil {
data, err := json.Marshal(w)
if err != nil {
return err
}
return json.NewEncoder(w.rw).Encode(w)
return w.writeRaw(data)
}
// savePretty saves wallet in a beautiful JSON.
func (w *Wallet) savePretty() error {
data, err := json.MarshalIndent(w, "", " ")
if err != nil {
return err
}
return w.writeRaw(data)
}
func (w *Wallet) writeRaw(data []byte) error {
if err := w.rewind(); err != nil {
return err
}
enc := json.NewEncoder(w.rw)
enc.SetIndent("", " ")
return enc.Encode(w)
_, err := w.rw.Write(data)
if err != nil {
return err
}
if f, ok := w.rw.(*os.File); ok {
if err := f.Truncate(int64(len(data))); err != nil {
return err
}
}
return nil
}
func (w *Wallet) rewind() error {
if s, ok := w.rw.(io.Seeker); ok {
if _, err := s.Seek(0, 0); err != nil {
if _, err := s.Seek(0, io.SeekStart); err != nil {
return err
}
}