diff --git a/core/src/main/java/org/bitcoinj/crypto/HDKeyDerivation.java b/core/src/main/java/org/bitcoinj/crypto/HDKeyDerivation.java index 357340301..b22d174d6 100644 --- a/core/src/main/java/org/bitcoinj/crypto/HDKeyDerivation.java +++ b/core/src/main/java/org/bitcoinj/crypto/HDKeyDerivation.java @@ -23,6 +23,8 @@ import java.math.*; import java.nio.*; import java.security.*; import java.util.*; +import java.util.function.Supplier; +import java.util.stream.Stream; import static com.google.common.base.Preconditions.*; @@ -122,6 +124,18 @@ public final class HDKeyDerivation { } + /** + * Generate an infinite stream of {@link DeterministicKey}s from the given parent and index. + *

+ * Note: Use {@link Stream#limit(long)} to get the desired number of keys. + * @param parent The parent key + * @param childNumber The index of the first child to supply/generate + * @return A stream of {@code DeterministicKey}s + */ + public static Stream generate(DeterministicKey parent, int childNumber) { + return Stream.generate(new KeySupplier(parent, childNumber)); + } + /** * @throws HDDerivationException if private derivation is attempted for a public-only parent key, or * if the resulting derived key is invalid (eg. private key == 0). @@ -228,6 +242,33 @@ public final class HDKeyDerivation { throw new HDDerivationException(errorMessage); } + /** + * A supplier of a sequence of {@code DeterministicKey}s generated from a starting + * parent and child index. + *

Not threadsafe: Should be used on a single thread + */ + private static class KeySupplier implements Supplier { + private final DeterministicKey parent; + private int nextChild; + + /** + * Construct a key supplier with the specified starting point + * @param parent The parent key + * @param nextChild The index of the next child to supply/generate + */ + public KeySupplier(DeterministicKey parent, int nextChild) { + this.parent = parent; + this.nextChild = nextChild; + } + + @Override + public DeterministicKey get() { + DeterministicKey key = HDKeyDerivation.deriveThisOrNextChildKey(parent, nextChild); + nextChild = key.getChildNumber().num() + 1; + return key; + } + } + public static class RawKeyBytes { public final byte[] keyBytes, chainCode; diff --git a/core/src/main/java/org/bitcoinj/wallet/DeterministicKeyChain.java b/core/src/main/java/org/bitcoinj/wallet/DeterministicKeyChain.java index a8edfb9d1..06a496f85 100644 --- a/core/src/main/java/org/bitcoinj/wallet/DeterministicKeyChain.java +++ b/core/src/main/java/org/bitcoinj/wallet/DeterministicKeyChain.java @@ -41,6 +41,7 @@ import java.util.*; import java.util.concurrent.Executor; import java.util.concurrent.locks.ReentrantLock; import java.util.function.Predicate; +import java.util.stream.Collectors; import static com.google.common.base.Preconditions.*; import static java.util.stream.Collectors.collectingAndThen; @@ -1202,23 +1203,19 @@ public class DeterministicKeyChain implements EncryptableKeyChain { checkState(lock.isHeldByCurrentThread()); final int numChildren = hierarchy.getNumChildren(parent.getPath()); final int needed = issued + lookaheadSize + lookaheadThreshold - numChildren; - - if (needed <= lookaheadThreshold) - return new ArrayList<>(); + final int limit = (needed > lookaheadThreshold) ? needed : 0; log.info("{} keys needed for {} = {} issued + {} lookahead size + {} lookahead threshold - {} num children", - needed, parent.getPathAsString(), issued, lookaheadSize, lookaheadThreshold, numChildren); + limit, parent.getPathAsString(), issued, lookaheadSize, lookaheadThreshold, numChildren); - List result = new ArrayList<>(needed); final Stopwatch watch = Stopwatch.createStarted(); - int nextChild = numChildren; - for (int i = 0; i < needed; i++) { - DeterministicKey key = HDKeyDerivation.deriveThisOrNextChildKey(parent, nextChild); - key = key.dropPrivateBytes(); + List result = HDKeyDerivation.generate(parent, numChildren) + .limit(limit) + .map(DeterministicKey::dropPrivateBytes) + .collect(Collectors.toList()); + result.forEach(key -> { hierarchy.putKey(key); - result.add(key); - nextChild = key.getChildNumber().num() + 1; - } + }); watch.stop(); log.info("Took {}", watch); return result; diff --git a/core/src/test/java/org/bitcoinj/crypto/HDKeyDerivationTest.java b/core/src/test/java/org/bitcoinj/crypto/HDKeyDerivationTest.java index 7bef7b09e..e364f4bc7 100644 --- a/core/src/test/java/org/bitcoinj/crypto/HDKeyDerivationTest.java +++ b/core/src/test/java/org/bitcoinj/crypto/HDKeyDerivationTest.java @@ -22,6 +22,9 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import java.math.BigInteger; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; import org.bitcoinj.core.Utils; import org.bitcoinj.crypto.HDKeyDerivation.PublicDeriveMode; @@ -136,4 +139,30 @@ public class HDKeyDerivationTest { assertEquals(EXPECTED_CHILD_PRIVATE_KEY, fromPublicWithInversion.getPrivateKeyAsHex()); assertEquals(EXPECTED_CHILD_PUBLIC_KEY, fromPublicWithInversion.getPublicKeyAsHex()); } + + @Test + public void testGenerate() { + DeterministicKey parent = new DeterministicKey(HDPath.m(), new byte[32], BigInteger.TEN, + null); + assertFalse(parent.isPubKeyOnly()); + assertFalse(parent.isEncrypted()); + + List keys0 = HDKeyDerivation.generate(parent, CHILD_NUMBER.num()) + .limit(0) + .collect(Collectors.toList()); + assertEquals(0, keys0.size()); + + List keys1 = HDKeyDerivation.generate(parent, CHILD_NUMBER.num()) + .limit(1) + .collect(Collectors.toList()); + assertEquals(1, keys1.size()); + assertEquals(HDPath.m(CHILD_NUMBER), keys1.get(0).getPath()); + + List keys2 = HDKeyDerivation.generate(parent, CHILD_NUMBER.num()) + .limit(2) + .collect(Collectors.toList()); + assertEquals(2, keys2.size()); + assertEquals(HDPath.parsePath("m/1"), keys2.get(0).getPath()); + assertEquals(HDPath.parsePath("m/2"), keys2.get(1).getPath()); + } }