From 62be6f959c789ffaf53f851fc029358525faa0d2 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Thu, 1 Sep 2022 15:28:35 +0300
Subject: [PATCH] keys/wallet: add Destroy/Close/Close

PrivateKey can be destroyed and Account/Wallet can be closed (destroying keys
in the process).
---
 pkg/crypto/keys/private_key.go      |  9 +++++++++
 pkg/crypto/keys/private_key_test.go |  4 ++++
 pkg/wallet/account.go               | 11 +++++++++++
 pkg/wallet/account_test.go          |  4 +++-
 pkg/wallet/wallet.go                | 11 ++++++++---
 pkg/wallet/wallet_test.go           |  5 ++++-
 6 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/pkg/crypto/keys/private_key.go b/pkg/crypto/keys/private_key.go
index 5e5992378..775c8e556 100644
--- a/pkg/crypto/keys/private_key.go
+++ b/pkg/crypto/keys/private_key.go
@@ -117,6 +117,15 @@ func (p *PrivateKey) WIF() string {
 	return w
 }
 
+// Destroy wipes the contents of the private key from memory. Any operations
+// with the key after call to Destroy have undefined behavior.
+func (p *PrivateKey) Destroy() {
+	bits := p.D.Bits()
+	for i := range bits {
+		bits[i] = 0
+	}
+}
+
 // Address derives the public NEO address that is coupled with the private key, and
 // returns it as a string.
 func (p *PrivateKey) Address() string {
diff --git a/pkg/crypto/keys/private_key_test.go b/pkg/crypto/keys/private_key_test.go
index b8b806080..16dc84c7a 100644
--- a/pkg/crypto/keys/private_key_test.go
+++ b/pkg/crypto/keys/private_key_test.go
@@ -2,6 +2,7 @@ package keys
 
 import (
 	"encoding/hex"
+	"math/big"
 	"strings"
 	"testing"
 
@@ -27,6 +28,9 @@ func TestPrivateKey(t *testing.T) {
 		assert.Equal(t, testCase.Wif, wif)
 		pubKey := privKey.PublicKey()
 		assert.Equal(t, hex.EncodeToString(pubKey.Bytes()), testCase.PublicKey)
+		oldD := new(big.Int).Set(privKey.D)
+		privKey.Destroy()
+		assert.NotEqual(t, oldD, privKey.D)
 	}
 }
 
diff --git a/pkg/wallet/account.go b/pkg/wallet/account.go
index ef426c08e..c0bc5910a 100644
--- a/pkg/wallet/account.go
+++ b/pkg/wallet/account.go
@@ -175,6 +175,17 @@ func (a *Account) PrivateKey() *keys.PrivateKey {
 	return a.privateKey
 }
 
+// Close cleans up the private key used by Account and disassociates it from
+// Account. The Account can no longer sign anything after this call, but Decrypt
+// can make it usable again.
+func (a *Account) Close() {
+	if a.privateKey == nil {
+		return
+	}
+	a.privateKey.Destroy()
+	a.privateKey = nil
+}
+
 // NewAccountFromWIF creates a new Account from the given WIF.
 func NewAccountFromWIF(wif string) (*Account, error) {
 	privKey, err := keys.NewPrivateKeyFromWIF(wif)
diff --git a/pkg/wallet/account_test.go b/pkg/wallet/account_test.go
index f424f779c..238430217 100644
--- a/pkg/wallet/account_test.go
+++ b/pkg/wallet/account_test.go
@@ -141,9 +141,11 @@ func TestContractSignTx(t *testing.T) {
 	require.Error(t, acc2.SignTx(0, tx)) // Locked account.
 
 	acc2.Locked = false
-	acc2.privateKey = nil
+	acc2.Close()
 	require.False(t, acc2.CanSign())
 	require.Error(t, acc2.SignTx(0, tx)) // No private key.
+	acc2.Close()                         // No-op.
+	require.False(t, acc2.CanSign())
 
 	tx.Scripts = append(tx.Scripts, transaction.Witness{
 		VerificationScript: acc.Contract.Script,
diff --git a/pkg/wallet/wallet.go b/pkg/wallet/wallet.go
index 024a5679b..a5fa85f96 100644
--- a/pkg/wallet/wallet.go
+++ b/pkg/wallet/wallet.go
@@ -168,9 +168,14 @@ func (w *Wallet) JSON() ([]byte, error) {
 	return json.MarshalIndent(w, " ", "	")
 }
 
-// Deprecated: Close is deprecated.
-// Close performs nothing and is left for backwards compatibility.
-func (w *Wallet) Close() {}
+// Close closes all Wallet accounts making them incapable of signing anything
+// (unless they're decrypted again). It's not doing anything to the underlying
+// wallet file.
+func (w *Wallet) Close() {
+	for _, acc := range w.Accounts {
+		acc.Close()
+	}
+}
 
 // GetAccount returns an account corresponding to the provided scripthash.
 func (w *Wallet) GetAccount(h util.Uint160) *Account {
diff --git a/pkg/wallet/wallet_test.go b/pkg/wallet/wallet_test.go
index 98f940cd6..7d155e9c4 100644
--- a/pkg/wallet/wallet_test.go
+++ b/pkg/wallet/wallet_test.go
@@ -34,13 +34,16 @@ func TestNewWalletFromFile_Negative_NoFile(t *testing.T) {
 	require.Errorf(t, err, "open testWallet: no such file or directory")
 }
 
-func TestCreateAccount(t *testing.T) {
+func TestCreateAccountAndClose(t *testing.T) {
 	wallet := checkWalletConstructor(t)
 
 	errAcc := wallet.CreateAccount("testName", "testPass")
 	require.NoError(t, errAcc)
 	accounts := wallet.Accounts
 	require.Len(t, accounts, 1)
+	require.True(t, wallet.Accounts[0].CanSign())
+	wallet.Close()
+	require.False(t, wallet.Accounts[0].CanSign())
 }
 
 func TestAddAccount(t *testing.T) {