From 9dbc507fca98e94736ed734e57b45469fe8ab61a Mon Sep 17 00:00:00 2001
From: troggy <7r0ggy@gmail.com>
Date: Mon, 23 Jun 2014 15:12:32 +0400
Subject: [PATCH] Improvement of getFollowingKeys

First of all, freshAddress was actually asking for a fresh key twice for non-married keychain. That was fixed by moving first call (needed only for married chain) inside the getFollowingKeys. As the latter now started to return all keys in a marriage and not only following ones, it was renamed to freshMarriedKeys.
Having all the keys in one block allows to do simple derivation path check to make sure keychains are in sync (as per @devrandom suggestion)
---
 .../google/bitcoin/wallet/KeyChainGroup.java  | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/core/src/main/java/com/google/bitcoin/wallet/KeyChainGroup.java b/core/src/main/java/com/google/bitcoin/wallet/KeyChainGroup.java
index 2f7f4fdc..89d1910b 100644
--- a/core/src/main/java/com/google/bitcoin/wallet/KeyChainGroup.java
+++ b/core/src/main/java/com/google/bitcoin/wallet/KeyChainGroup.java
@@ -238,12 +238,9 @@ public class KeyChainGroup {
      */
     public Address freshAddress(KeyChain.KeyPurpose purpose) {
         DeterministicKeyChain chain = getActiveKeyChain();
-        DeterministicKey key = chain.getKey(purpose);
         if (isMarried(chain)) {
-            List<ECKey> keys = ImmutableList.<ECKey>builder()
-                    .addAll(getFollowingKeys(purpose, chain.getWatchingKey()))
-                    .add(key).build();
-            Address freshAddress = Address.fromP2SHScript(params, ScriptBuilder.createP2SHOutputScript(2, keys));
+            List<ECKey> marriedKeys = freshMarriedKeys(purpose, chain);
+            Address freshAddress = Address.fromP2SHScript(params, ScriptBuilder.createP2SHOutputScript(2, marriedKeys));
             currentAddresses.put(purpose, freshAddress);
             return freshAddress;
         } else {
@@ -251,13 +248,16 @@ public class KeyChainGroup {
         }
     }
 
-    private List<ECKey> getFollowingKeys(KeyChain.KeyPurpose purpose, DeterministicKey followedChainWatchKey) {
-        List<ECKey> keys = new ArrayList<ECKey>();
-        Collection<DeterministicKeyChain> keyChains = followingKeychains.get(followedChainWatchKey);
+    private List<ECKey> freshMarriedKeys(KeyChain.KeyPurpose purpose, DeterministicKeyChain followedKeyChain) {
+        DeterministicKey followedKey = followedKeyChain.getKey(purpose);
+        ImmutableList.Builder<ECKey> keys = ImmutableList.<ECKey>builder().add(followedKey);
+        Collection<DeterministicKeyChain> keyChains = followingKeychains.get(followedKeyChain.getWatchingKey());
         for (DeterministicKeyChain keyChain : keyChains) {
-            keys.add(keyChain.getKey(purpose));
+            DeterministicKey followingKey = keyChain.getKey(purpose);
+            checkState(followedKey.getChildNumber().equals(followingKey.getChildNumber()), "Following keychains should be in sync");
+            keys.add(followingKey);
         }
-        return keys;
+        return keys.build();
     }
 
     /** Returns the key chain that's used for generation of fresh/current keys. This is always the newest HD chain. */