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