From f8ad0939a2af2e2090f2b037bf30632bb3b52af3 Mon Sep 17 00:00:00 2001
From: Dave Collins <davec@conformal.com>
Date: Mon, 18 Aug 2014 17:53:37 -0500
Subject: [PATCH] Add new function on extended keys to zero them.

This commit adds a new function named Zero on the hdkeychain.ExtendedKey
which can be used to manually clear the memory used for an extended key.
This is useful for enhanced security by allowing the caller to explicitly
clear the memory when they're done with a key.  Otherwise it might hang
around in memory for a while.

Once a key has been zeroed it is no longer usable.

This commit also contains tests to ensure everything works as expected
after a key has been zeroed.
---
 hdkeychain/extendedkey.go      | 29 +++++++++++++
 hdkeychain/extendedkey_test.go | 79 ++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+)

diff --git a/hdkeychain/extendedkey.go b/hdkeychain/extendedkey.go
index 648ee3e..c0e3e2b 100644
--- a/hdkeychain/extendedkey.go
+++ b/hdkeychain/extendedkey.go
@@ -360,6 +360,10 @@ func (k *ExtendedKey) Address(net *btcnet.Params) (*btcutil.AddressPubKeyHash, e
 
 // String returns the extended key as a human-readable base58-encoded string.
 func (k *ExtendedKey) String() string {
+	if len(k.key) == 0 {
+		return "zeroed extended key"
+	}
+
 	var childNumBytes [4]byte
 	depthByte := byte(k.depth % 256)
 	binary.BigEndian.PutUint32(childNumBytes[:], k.childNum)
@@ -402,6 +406,31 @@ func (k *ExtendedKey) SetNet(net *btcnet.Params) {
 	}
 }
 
+// zero sets all bytes in the passed slice to zero.  This is used to
+// explicitly clear private key material from memory.
+func zero(b []byte) {
+	lenb := len(b)
+	for i := 0; i < lenb; i++ {
+		b[i] = 0
+	}
+}
+
+// Zero manually clears all fields and bytes in the extended key.  This can be
+// used to explicitly clear key material from memory for enhanced security
+// against memory scraping.  This function only clears this particular key and
+// not any children that have already been derived.
+func (k *ExtendedKey) Zero() {
+	zero(k.key)
+	zero(k.pubKey)
+	zero(k.chainCode)
+	zero(k.parentFP)
+	zero(k.version)
+	k.key = nil
+	k.depth = 0
+	k.childNum = 0
+	k.isPrivate = false
+}
+
 // NewMaster creates a new master node for use in creating a hierarchical
 // deterministic key chain.  The seed must be between 128 and 512 bits and
 // should be generated by a cryptographically secure random generation source.
diff --git a/hdkeychain/extendedkey_test.go b/hdkeychain/extendedkey_test.go
index ab0cff4..5527c65 100644
--- a/hdkeychain/extendedkey_test.go
+++ b/hdkeychain/extendedkey_test.go
@@ -674,3 +674,82 @@ func TestErrors(t *testing.T) {
 		}
 	}
 }
+
+// TestZero ensures that zeroing an extended key works as intended.
+func TestZero(t *testing.T) {
+	tests := []struct {
+		name   string
+		extKey string
+	}{
+		// Test vector 1
+		{
+			name:   "test vector 1 chain m",
+			extKey: "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi",
+		},
+	}
+
+	for i, test := range tests {
+		key, err := hdkeychain.NewKeyFromString(test.extKey)
+		if err != nil {
+			t.Errorf("NewKeyFromString #%d (%s): unexpected "+
+				"error: %v", i, test.name, err)
+			continue
+		}
+		key.Zero()
+
+		// Zeroing a key should result in it no longer being private
+		if key.IsPrivate() != false {
+			t.Errorf("IsPrivate #%d (%s): mismatched key type -- "+
+				"want private %v, got private %v", i, test.name,
+				false, key.IsPrivate())
+			continue
+		}
+
+		parentFP := key.ParentFingerprint()
+		if parentFP != 0 {
+			t.Errorf("ParentFingerprint #%d (%s): mismatched "+
+				"parent fingerprint -- want %d, got %d", i,
+				test.name, 0, parentFP)
+			continue
+		}
+
+		wantKey := "zeroed extended key"
+		serializedKey := key.String()
+		if serializedKey != wantKey {
+			t.Errorf("String #%d (%s): mismatched serialized key "+
+				"-- want %s, got %s", i, test.name, wantKey,
+				serializedKey)
+			continue
+		}
+
+		wantErr := hdkeychain.ErrNotPrivExtKey
+		_, err = key.ECPrivKey()
+		if !reflect.DeepEqual(err, wantErr) {
+			t.Errorf("ECPrivKey #%d (%s): mismatched error: want "+
+				"%v, got %v", i, test.name, wantErr, err)
+			continue
+		}
+
+		wantErr = errors.New("pubkey string is empty")
+		_, err = key.ECPubKey()
+		if !reflect.DeepEqual(err, wantErr) {
+			t.Errorf("ECPubKey #%d (%s): mismatched error: want "+
+				"%v, got %v", i, test.name, wantErr, err)
+			continue
+		}
+
+		wantAddr := "1HT7xU2Ngenf7D4yocz2SAcnNLW7rK8d4E"
+		addr, err := key.Address(&btcnet.MainNetParams)
+		if err != nil {
+			t.Errorf("Addres s #%d (%s): unexpected error: %v", i,
+				test.name, err)
+			continue
+		}
+		if addr.EncodeAddress() != wantAddr {
+			t.Errorf("Address #%d (%s): mismatched address -- want "+
+				"%s, got %s", i, test.name, wantAddr,
+				addr.EncodeAddress())
+			continue
+		}
+	}
+}