HD wallet: currentKey is now stable after serialization roundtrip

At the moment currentKeys map of KeyChainGroup is not restored after
deserialization and subsequent call to currentKey produces different
key then expected.
Proposed solution reconstructs currentKey map on deserialization using
stored numbers of issues keys. It is
not future-proof as it assumes only RECEIVE and CHANGE keys are being
used.

core/src/main/java/com/google/bitcoin/wallet/DeterministicKeyChain.java

@@ -826,6 +826,32 @@ public class DeterministicKeyChain implements EncryptableKeyChain {
         return result;
     }
 
+    /**
+     * Returns number of keys used on external path. This may be fewer than the number that have been deserialized
+     * or held in memory, because of the lookahead zone.
+     */
+    public int getIssuedExternalKeys() {
+        lock.lock();
+        try {
+            return issuedExternalKeys;
+        } finally {
+            lock.unlock();
+        }
+    }
+
+    /**
+     * Returns number of keys used on internal path. This may be fewer than the number that have been deserialized
+     * or held in memory, because of the lookahead zone.
+     */
+    public int getIssuedInternalKeys() {
+        lock.lock();
+        try {
+            return issuedInternalKeys;
+        } finally {
+            lock.unlock();
+        }
+    }
+
     /** Returns the seed or null if this chain is encrypted or watching. */
     @Nullable
     public DeterministicSeed getSeed() {

core/src/main/java/com/google/bitcoin/wallet/KeyChainGroup.java

@@ -18,6 +18,7 @@
 package com.google.bitcoin.wallet;
 
 import com.google.bitcoin.core.*;
+import com.google.bitcoin.crypto.ChildNumber;
 import com.google.bitcoin.crypto.DeterministicKey;
 import com.google.bitcoin.crypto.KeyCrypter;
 import com.google.bitcoin.store.UnreadableWalletException;
@@ -62,12 +63,12 @@ public class KeyChainGroup {
 
     /** Creates a keychain group with no basic chain, and a single randomly initialized HD chain. */
     public KeyChainGroup() {
-        this(null, new ArrayList<DeterministicKeyChain>(1), null);
+        this(null, new ArrayList<DeterministicKeyChain>(1), null, null);
     }
 
     /** Creates a keychain group with no basic chain, and an HD chain initialized from the given seed. */
     public KeyChainGroup(DeterministicSeed seed) {
-        this(null, ImmutableList.of(new DeterministicKeyChain(seed)), null);
+        this(null, ImmutableList.of(new DeterministicKeyChain(seed)), null, null);
     }
 
     /**
@@ -75,7 +76,7 @@ public class KeyChainGroup {
      * This HAS to be an account key as returned by {@link DeterministicKeyChain#getWatchingKey()}.
      */
     public KeyChainGroup(DeterministicKey watchKey) {
-        this(null, ImmutableList.of(DeterministicKeyChain.watch(watchKey)), null);
+        this(null, ImmutableList.of(DeterministicKeyChain.watch(watchKey)), null, null);
     }
 
     /**
@@ -84,15 +85,17 @@ public class KeyChainGroup {
      * This HAS to be an account key as returned by {@link DeterministicKeyChain#getWatchingKey()}.
      */
     public KeyChainGroup(DeterministicKey watchKey, long creationTimeSecondsSecs) {
-        this(null, ImmutableList.of(DeterministicKeyChain.watch(watchKey, creationTimeSecondsSecs)), null);
+        this(null, ImmutableList.of(DeterministicKeyChain.watch(watchKey, creationTimeSecondsSecs)), null, null);
     }
 
     // Used for deserialization.
-    private KeyChainGroup(@Nullable BasicKeyChain basicKeyChain, List<DeterministicKeyChain> chains, @Nullable KeyCrypter crypter) {
+    private KeyChainGroup(@Nullable BasicKeyChain basicKeyChain, List<DeterministicKeyChain> chains, @Nullable EnumMap<KeyChain.KeyPurpose, DeterministicKey> currentKeys, @Nullable KeyCrypter crypter) {
         this.basic = basicKeyChain == null ? new BasicKeyChain() : basicKeyChain;
         this.chains = new ArrayList<DeterministicKeyChain>(checkNotNull(chains));
         this.keyCrypter = crypter;
-        this.currentKeys = new EnumMap<KeyChain.KeyPurpose, DeterministicKey>(KeyChain.KeyPurpose.class);
+        this.currentKeys = currentKeys == null
+                ? new EnumMap<KeyChain.KeyPurpose, DeterministicKey>(KeyChain.KeyPurpose.class)
+                : currentKeys;
     }
 
     private void createAndActivateNewHDChain() {
@@ -443,20 +446,42 @@ public class KeyChainGroup {
     public static KeyChainGroup fromProtobufUnencrypted(List<Protos.Key> keys) throws UnreadableWalletException {
         BasicKeyChain basicKeyChain = BasicKeyChain.fromProtobufUnencrypted(keys);
         List<DeterministicKeyChain> chains = DeterministicKeyChain.fromProtobuf(keys, null);
+        EnumMap<KeyChain.KeyPurpose, DeterministicKey> currentKeys = createCurrentKeysMap(chains);
+
         if (chains.isEmpty()) {
             // TODO: Old bag-of-keys style wallet only! Auto-upgrade time!
         }
-        return new KeyChainGroup(basicKeyChain, chains, null);
+        return new KeyChainGroup(basicKeyChain, chains, currentKeys, null);
     }
 
     public static KeyChainGroup fromProtobufEncrypted(List<Protos.Key> keys, KeyCrypter crypter) throws UnreadableWalletException {
         checkNotNull(crypter);
         BasicKeyChain basicKeyChain = BasicKeyChain.fromProtobufEncrypted(keys, crypter);
         List<DeterministicKeyChain> chains = DeterministicKeyChain.fromProtobuf(keys, crypter);
+        EnumMap<KeyChain.KeyPurpose, DeterministicKey> currentKeys = createCurrentKeysMap(chains);
+
         if (chains.isEmpty()) {
             // TODO: Old bag-of-keys style wallet only! Auto-upgrade time!
         }
-        return new KeyChainGroup(basicKeyChain, chains, crypter);
+        return new KeyChainGroup(basicKeyChain, chains, currentKeys, crypter);
+    }
+
+    private static EnumMap<KeyChain.KeyPurpose, DeterministicKey> createCurrentKeysMap(List<DeterministicKeyChain> chains) {
+        DeterministicKeyChain activeChain = chains.get(chains.size() - 1);
+        DeterministicKey currentExternalKey = activeChain.getKeyByPath(
+                ImmutableList.of(ChildNumber.ZERO_HARDENED, ChildNumber.ZERO, new ChildNumber(activeChain.getIssuedExternalKeys() - 1))
+        );
+        DeterministicKey currentInternalKey = activeChain.getKeyByPath(
+                ImmutableList.of(ChildNumber.ZERO_HARDENED, new ChildNumber(1), new ChildNumber(activeChain.getIssuedInternalKeys() - 1))
+        );
+
+        EnumMap<KeyChain.KeyPurpose, DeterministicKey> currentKeys = new EnumMap<KeyChain.KeyPurpose, DeterministicKey>(KeyChain.KeyPurpose.class);
+        // assuming that only RECEIVE and CHANGE keys are being used at the moment, we will treat latest issued external key
+        // as current RECEIVE key and latest issued internal key as CHANGE key. This should be changed as soon as other
+        // kinds of KeyPurpose are introduced.
+        currentKeys.put(KeyChain.KeyPurpose.RECEIVE_FUNDS, currentExternalKey);
+        currentKeys.put(KeyChain.KeyPurpose.CHANGE, currentInternalKey);
+        return currentKeys;
     }
 
     public String toString(@Nullable NetworkParameters params, boolean includePrivateKeys) {

core/src/test/java/com/google/bitcoin/wallet/KeyChainGroupTest.java

@@ -284,20 +284,23 @@ public class KeyChainGroupTest {
     @Test
     public void serialization() throws Exception {
         assertEquals(INITIAL_KEYS + 1 /* for the seed */, group.serializeToProtobuf().size());
-        DeterministicKey key1 = (DeterministicKey) group.freshKey(KeyChain.KeyPurpose.RECEIVE_FUNDS);
+        group.freshKey(KeyChain.KeyPurpose.RECEIVE_FUNDS);
-        DeterministicKey key2 = (DeterministicKey) group.freshKey(KeyChain.KeyPurpose.CHANGE);
+        DeterministicKey key1 = group.freshKey(KeyChain.KeyPurpose.RECEIVE_FUNDS);
+        DeterministicKey key2 = group.freshKey(KeyChain.KeyPurpose.CHANGE);
         List<Protos.Key> protoKeys1 = group.serializeToProtobuf();
-        assertEquals(INITIAL_KEYS + ((LOOKAHEAD_SIZE + 1) * 2) + 1 /* for the seed */, protoKeys1.size());
+        assertEquals(INITIAL_KEYS + ((LOOKAHEAD_SIZE + 1) * 2) + 1 /* for the seed */ + 1, protoKeys1.size());
         group.importKeys(new ECKey());
         List<Protos.Key> protoKeys2 = group.serializeToProtobuf();
-        assertEquals(INITIAL_KEYS + ((LOOKAHEAD_SIZE + 1) * 2) + 1 /* for the seed */ + 1, protoKeys2.size());
+        assertEquals(INITIAL_KEYS + ((LOOKAHEAD_SIZE + 1) * 2) + 1 /* for the seed */ + 2, protoKeys2.size());
 
         group = KeyChainGroup.fromProtobufUnencrypted(protoKeys1);
-        assertEquals(INITIAL_KEYS + ((LOOKAHEAD_SIZE + 1)  * 2)  + 1 /* for the seed */, protoKeys1.size());
+        assertEquals(INITIAL_KEYS + ((LOOKAHEAD_SIZE + 1)  * 2)  + 1 /* for the seed */ + 1, protoKeys1.size());
         assertTrue(group.hasKey(key1));
         assertTrue(group.hasKey(key2));
+        assertEquals(key2, group.currentKey(KeyChain.KeyPurpose.CHANGE));
+        assertEquals(key1, group.currentKey(KeyChain.KeyPurpose.RECEIVE_FUNDS));
         group = KeyChainGroup.fromProtobufUnencrypted(protoKeys2);
-        assertEquals(INITIAL_KEYS + ((LOOKAHEAD_SIZE + 1) * 2) + 1 /* for the seed */ + 1, protoKeys2.size());
+        assertEquals(INITIAL_KEYS + ((LOOKAHEAD_SIZE + 1) * 2) + 1 /* for the seed */ + 2, protoKeys2.size());
         assertTrue(group.hasKey(key1));
         assertTrue(group.hasKey(key2));