Married HD wallets: Bloom filter adjustments

Pull request: #115
Based on design notes:
https://groups.google.com/forum/#!msg/bitcoinj/Uxl-z40OLuQ/e2m4mEWR6gMJ

core/src/main/java/com/google/bitcoin/core/TransactionOutput.java

@@ -293,6 +293,8 @@ public class TransactionOutput extends ChildMessage implements Serializable {
             if (script.isSentToRawPubKey()) {
                 byte[] pubkey = script.getPubKey();
                 return wallet.isPubKeyMine(pubkey);
+            } if (script.isPayToScriptHash()) {
+                return wallet.isPayToScriptHashMine(script.getPubKeyHash());
             } else {
                 byte[] pubkeyHash = script.getPubKeyHash();
                 return wallet.isPubKeyHashMine(pubkeyHash);

core/src/main/java/com/google/bitcoin/core/Wallet.java

@@ -755,6 +755,27 @@ public class Wallet extends BaseTaggableObject implements Serializable, BlockCha
         return findKeyFromPubKey(pubkey) != null;
     }
 
+    /**
+     * <p>Locates a script from the KeyChainGroup given the script hash. This is needed when finding out which
+     * script we need to use to redeem a transaction output.</p>
+     * Returns null if no such script found
+     */
+    @Nullable
+    public Script findRedeemScriptFromPubHash(byte[] payToScriptHash) {
+        lock.lock();
+        try {
+            return keychain.findRedeemScriptFromPubHash(payToScriptHash);
+        } finally {
+            lock.unlock();
+        }
+    }
+    /**
+     * Returns true if this wallet knows the script corresponding to the given hash
+     */
+    public boolean isPayToScriptHashMine(byte[] payToScriptHash) {
+        return findRedeemScriptFromPubHash(payToScriptHash) != null;
+    }
+
     /**
      * Marks all keys used in the transaction output as used in the wallet.
      * See {@link com.google.bitcoin.wallet.DeterministicKeyChain#markKeyAsUsed(DeterministicKey)} for more info on this.
@@ -3405,7 +3426,8 @@ public class Wallet extends BaseTaggableObject implements Serializable, BlockCha
     }
 
     private boolean isTxOutputBloomFilterable(TransactionOutput out) {
-        return (out.isMine(this) && out.getScriptPubKey().isSentToRawPubKey()) ||
+        boolean isScriptTypeSupported = out.getScriptPubKey().isSentToRawPubKey() || out.getScriptPubKey().isPayToScriptHash();
+        return (out.isMine(this) && isScriptTypeSupported) ||
                 out.isWatched(this);
     }
 

core/src/main/java/com/google/bitcoin/script/ScriptBuilder.java

@@ -24,7 +24,6 @@ import com.google.common.collect.Lists;
 import com.google.common.primitives.UnsignedBytes;
 
 import javax.annotation.Nullable;
-import java.math.BigInteger;
 import java.util.*;
 
 import static com.google.bitcoin.script.ScriptOpCodes.*;
@@ -187,11 +186,28 @@ public class ScriptBuilder {
         return new ScriptBuilder().op(OP_HASH160).data(hash).op(OP_EQUAL).build();
     }
 
+    /**
+     * Creates a scriptPubKey for the given redeem script.
+     */
+    public static Script createP2SHOutputScript(Script redeemScript) {
+        byte[] hash = Utils.sha256hash160(redeemScript.getProgram());
+        return ScriptBuilder.createP2SHOutputScript(hash);
+    }
+
     /**
      * Creates a P2SH output script with given public keys and threshold. Given public keys will be placed in
      * redeem script in the lexicographical sorting order.
      */
     public static Script createP2SHOutputScript(int threshold, List<ECKey> pubkeys) {
+        Script redeemScript = createRedeemScript(threshold, pubkeys);
+        return createP2SHOutputScript(redeemScript);
+    }
+
+    /**
+     * Creates redeem script with given public keys and threshold. Given public keys will be placed in
+     * redeem script in the lexicographical sorting order.
+     */
+    public static Script createRedeemScript(int threshold, List<ECKey> pubkeys) {
         pubkeys = new ArrayList<ECKey>(pubkeys);
         final Comparator comparator = UnsignedBytes.lexicographicalComparator();
         Collections.sort(pubkeys, new Comparator<ECKey>() {
@@ -200,8 +216,7 @@ public class ScriptBuilder {
                 return comparator.compare(k1.getPubKey(), k2.getPubKey());
             }
         });
-        Script redeemScript = ScriptBuilder.createMultiSigOutputScript(threshold, pubkeys);
+
-        byte[] hash = Utils.sha256hash160(redeemScript.getProgram());
+        return ScriptBuilder.createMultiSigOutputScript(threshold, pubkeys);
-        return ScriptBuilder.createP2SHOutputScript(hash);
     }
 }

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

@@ -927,4 +927,19 @@ public class DeterministicKeyChain implements EncryptableKeyChain {
     /* package */ List<ECKey> getKeys() {
         return basicKeyChain.getKeys();
     }
+
+
+    /**
+     * Returns leaf keys issued by this chain (not including lookahead zone)
+     */
+    public List<DeterministicKey> getLeafKeys() {
+        ImmutableList.Builder<DeterministicKey> keys = ImmutableList.builder();
+        for (ECKey key : getKeys()) {
+            DeterministicKey dKey = (DeterministicKey) key;
+            if (dKey.getPath().size() > 2) {
+                keys.add(dKey);
+            }
+        }
+        return keys.build();
+    }
 }

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

@@ -21,6 +21,7 @@ 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.script.Script;
 import com.google.bitcoin.script.ScriptBuilder;
 import com.google.bitcoin.store.UnreadableWalletException;
 import com.google.bitcoin.utils.ListenerRegistration;
@@ -30,6 +31,7 @@ import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Multimap;
+import com.google.protobuf.ByteString;
 import org.bitcoinj.wallet.Protos;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -68,6 +70,9 @@ public class KeyChainGroup {
     // The map keys are the watching keys of the followed chains and values are the following chains
     private Multimap<DeterministicKey, DeterministicKeyChain> followingKeychains;
 
+    // The map holds P2SH redeem scripts issued by this KeyChainGroup (including lookahead) mapped to their scriptPubKey hashes.
+    private LinkedHashMap<ByteString, Script> marriedKeysScripts;
+
     private EnumMap<KeyChain.KeyPurpose, Address> currentAddresses;
     @Nullable private KeyCrypter keyCrypter;
     private int lookaheadSize = -1;
@@ -144,6 +149,39 @@ public class KeyChainGroup {
         if (followingKeychains != null) {
             this.followingKeychains.putAll(followingKeychains);
         }
+        marriedKeysScripts = new LinkedHashMap<ByteString, Script>();
+        maybeLookaheadScripts();
+
+        if (!this.currentKeys.isEmpty()) {
+            DeterministicKey followedWatchKey = getActiveKeyChain().getWatchingKey();
+            for (Map.Entry<KeyChain.KeyPurpose, DeterministicKey> entry : this.currentKeys.entrySet()) {
+                Address address = makeP2SHOutputScript(entry.getValue(), followedWatchKey).getToAddress(params);
+                currentAddresses.put(entry.getKey(), address);
+            }
+        }
+    }
+
+    /**
+     * This keeps {@link #marriedKeysScripts} in sync with the number of keys issued
+     */
+    private void maybeLookaheadScripts() {
+        if (chains.isEmpty())
+            return;
+
+        int numLeafKeys = chains.get(chains.size() - 1).getLeafKeys().size();
+        checkState(marriedKeysScripts.size() <= numLeafKeys, "Number of scripts is greater than number of leaf keys");
+        if (marriedKeysScripts.size() == numLeafKeys)
+            return;
+
+        for (DeterministicKeyChain chain : chains) {
+            if (isMarried(chain)) {
+                for (DeterministicKey followedKey : chain.getLeafKeys()) {
+                    Script redeemScript = makeRedeemScript(followedKey, chain.getWatchingKey());
+                    Script scriptPubKey = ScriptBuilder.createP2SHOutputScript(redeemScript);
+                    marriedKeysScripts.put(ByteString.copyFrom(scriptPubKey.getPubKeyHash()), redeemScript);
+                }
+            }
+        }
     }
 
     private void createAndActivateNewHDChain() {
@@ -239,7 +277,9 @@ public class KeyChainGroup {
         DeterministicKeyChain chain = getActiveKeyChain();
         if (isMarried(chain)) {
             List<ECKey> marriedKeys = freshMarriedKeys(purpose, chain);
-            Address freshAddress = Address.fromP2SHScript(params, ScriptBuilder.createP2SHOutputScript(2, marriedKeys));
+            Script p2shScript = makeP2SHOutputScript(marriedKeys);
+            Address freshAddress = Address.fromP2SHScript(params, p2shScript);
+            maybeLookaheadScripts();
             currentAddresses.put(purpose, freshAddress);
             return freshAddress;
         } else {
@@ -259,6 +299,15 @@ public class KeyChainGroup {
         return keys.build();
     }
 
+    private List<ECKey> getMarriedKeysWithFollowed(DeterministicKey followedKey, Collection<DeterministicKeyChain> followingChains) {
+        ImmutableList.Builder<ECKey> keys = ImmutableList.builder();
+        for (DeterministicKeyChain keyChain : followingChains) {
+            keys.add(keyChain.getKeyByPath(followedKey.getPath()));
+        }
+        keys.add(followedKey);
+        return keys.build();
+    }
+
     /** Returns the key chain that's used for generation of fresh/current keys. This is always the newest HD chain. */
     public DeterministicKeyChain getActiveKeyChain() {
         if (chains.isEmpty()) {
@@ -353,6 +402,15 @@ public class KeyChainGroup {
         return importKeys(encryptedKeys);
     }
 
+    /**
+     * <p>Returns redeem script for the given scriptPubKey hash.
+     * Returns null if no such script found
+     */
+    @Nullable
+    public Script findRedeemScriptFromPubHash(byte[] payToScriptHash) {
+        return marriedKeysScripts.get(ByteString.copyFrom(payToScriptHash));
+    }
+
     @Nullable
     public ECKey findKeyFromPubHash(byte[] pubkeyHash) {
         ECKey result;
@@ -513,8 +571,13 @@ public class KeyChainGroup {
 
     public int getBloomFilterElementCount() {
         int result = basic.numBloomFilterEntries();
-        for (DeterministicKeyChain chain : chains)
+        for (DeterministicKeyChain chain : chains) {
-            result += chain.numBloomFilterEntries();
+            if (isMarried(chain)) {
+                result += chain.getLeafKeys().size() * 2;
+            } else {
+                result += chain.numBloomFilterEntries();
+            }
+        }
         return result;
     }
 
@@ -522,8 +585,16 @@ public class KeyChainGroup {
         BloomFilter filter = new BloomFilter(size, falsePositiveRate, nTweak);
         if (basic.numKeys() > 0)
             filter.merge(basic.getFilter(size, falsePositiveRate, nTweak));
-        for (DeterministicKeyChain chain : chains)
+        for (DeterministicKeyChain chain : chains) {
-            filter.merge(chain.getFilter(size, falsePositiveRate, nTweak));
+            if (isMarried(chain)) {
+                for (Map.Entry<ByteString, Script> entry : marriedKeysScripts.entrySet()) {
+                    filter.insert(entry.getKey().toByteArray());
+                    filter.insert(ScriptBuilder.createP2SHOutputScript(entry.getValue()).getProgram());
+                }
+            } else {
+                filter.merge(chain.getFilter(size, falsePositiveRate, nTweak));
+            }
+        }
         return filter;
     }
 
@@ -532,6 +603,24 @@ public class KeyChainGroup {
         throw new UnsupportedOperationException();   // Unused.
     }
 
+    private Script makeP2SHOutputScript(List<ECKey> marriedKeys) {
+        return ScriptBuilder.createP2SHOutputScript(makeRedeemScript(marriedKeys));
+    }
+
+    private Script makeP2SHOutputScript(DeterministicKey followedKey, DeterministicKey followedAccountKey) {
+        return ScriptBuilder.createP2SHOutputScript(makeRedeemScript(followedKey, followedAccountKey));
+    }
+
+    private Script makeRedeemScript(DeterministicKey followedKey, DeterministicKey followedAccountKey) {
+        Collection<DeterministicKeyChain> followingChains = followingKeychains.get(followedAccountKey);
+        List<ECKey> marriedKeys = getMarriedKeysWithFollowed(followedKey, followingChains);
+        return makeRedeemScript(marriedKeys);
+    }
+
+    private Script makeRedeemScript(List<ECKey> marriedKeys) {
+        return ScriptBuilder.createRedeemScript((marriedKeys.size() / 2) + 1, marriedKeys);
+    }
+
     /** Adds a listener for events that are run when keys are added, on the user thread. */
     public void addEventListener(KeyChainEventListener listener) {
         addEventListener(listener, Threading.USER_THREAD);
@@ -729,12 +818,30 @@ public class KeyChainGroup {
             if (watchingKey.getParent() != null) {
                 builder.append(String.format("Key to watch:  %s%n%n", watchingKey.serializePubB58()));
             }
-            for (ECKey key : chain.getKeys())
+            if (isMarried(chain)) {
-                formatKeyWithAddress(includePrivateKeys, key, builder);
+                Collection<DeterministicKeyChain> followingChains = followingKeychains.get(chain.getWatchingKey());
+                for (DeterministicKeyChain followingChain : followingChains) {
+                    builder.append(String.format("Following chain:  %s%n", followingChain.getWatchingKey().serializePubB58()));
+                }
+                builder.append("\n");
+                for (Script script : marriedKeysScripts.values())
+                    formatScript(ScriptBuilder.createP2SHOutputScript(script), builder);
+            } else {
+                for (ECKey key : chain.getKeys())
+                    formatKeyWithAddress(includePrivateKeys, key, builder);
+            }
         }
         return builder.toString();
     }
 
+    private void formatScript(Script script, StringBuilder builder) {
+        builder.append("  addr:");
+        builder.append(script.getToAddress(params));
+        builder.append("  hash160:");
+        builder.append(Utils.HEX.encode(script.getPubKeyHash()));
+        builder.append("\n");
+    }
+
     private void formatKeyWithAddress(boolean includePrivateKeys, ECKey key, StringBuilder builder) {
         final Address address = key.toAddress(params);
         builder.append("  addr:");

core/src/test/java/com/google/bitcoin/core/WalletTest.java

@@ -20,6 +20,7 @@ package com.google.bitcoin.core;
 import com.google.bitcoin.core.Transaction.SigHash;
 import com.google.bitcoin.core.Wallet.SendRequest;
 import com.google.bitcoin.crypto.*;
+import com.google.bitcoin.store.MemoryBlockStore;
 import com.google.bitcoin.store.WalletProtobufSerializer;
 import com.google.bitcoin.testing.FakeTxBuilder;
 import com.google.bitcoin.testing.MockTransactionBroadcaster;
@@ -1184,6 +1185,29 @@ public class WalletTest extends TestWithWallet {
         assertTrue(wallet.getBloomFilter(1e-12).contains(outPoint.bitcoinSerialize()));
     }
 
+    @Test
+    public void marriedKeychainBloomFilter() throws Exception {
+        wallet = new Wallet(params);
+        blockStore = new MemoryBlockStore(params);
+        chain = new BlockChain(params, wallet, blockStore);
+
+        String XPUB = "xpub68KFnj3bqUx1s7mHejLDBPywCAKdJEu1b49uniEEn2WSbHmZ7xbLqFTjJbtx1LUcAt1DwhoqWHmo2s5WMJp6wi38CiF2hYD49qVViKVvAoi";
+        wallet.addFollowingAccountKeys(ImmutableList.of(DeterministicKey.deserializeB58(null, XPUB)));
+        Address address = wallet.currentReceiveAddress();
+
+        assertTrue(wallet.getBloomFilter(0.001).contains(address.getHash160()));
+
+        Transaction t1 = createFakeTx(params, CENT, address);
+        StoredBlock b1 = createFakeBlock(blockStore, t1).storedBlock;
+
+        TransactionOutPoint outPoint = new TransactionOutPoint(params, 0, t1);
+
+        assertFalse(wallet.getBloomFilter(0.001).contains(outPoint.bitcoinSerialize()));
+
+        wallet.receiveFromBlock(t1, b1, BlockChain.NewBlockType.BEST_CHAIN, 0);
+        assertTrue(wallet.getBloomFilter(0.001).contains(outPoint.bitcoinSerialize()));
+    }
+
     @Test
     public void autosaveImmediate() throws Exception {
         // Test that the wallet will save itself automatically when it changes.

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

@@ -56,7 +56,7 @@ public class KeyChainGroupTest {
     }
 
     private KeyChainGroup createMarriedKeyChainGroup() {
-        byte[] seedBytes = Sha256Hash.create("don't use a string seed like this in real life".getBytes()).getBytes();
+        byte[] seedBytes = Sha256Hash.create("don't use a seed like this in real life".getBytes()).getBytes();
         DeterministicSeed seed = new DeterministicSeed(seedBytes, MnemonicCode.BIP39_STANDARDISATION_TIME_SECS);
         KeyChainGroup group = new KeyChainGroup(params, seed, ImmutableList.of(watchingAccountKey));
         group.setLookaheadSize(LOOKAHEAD_SIZE);
@@ -300,6 +300,44 @@ public class KeyChainGroupTest {
         assertTrue(filter.contains(key2.getPubKey()));
     }
 
+    @Test
+    public void findRedeemScriptFromPubHash() throws Exception {
+        group = createMarriedKeyChainGroup();
+        Address address = group.freshAddress(KeyChain.KeyPurpose.RECEIVE_FUNDS);
+        assertTrue(group.findRedeemScriptFromPubHash(address.getHash160()) != null);
+        KeyChainGroup group2 = createMarriedKeyChainGroup();
+        group2.freshAddress(KeyChain.KeyPurpose.RECEIVE_FUNDS);
+        // test address from lookahead zone
+        for (int i = 0; i < LOOKAHEAD_SIZE; i++) {
+            address = group.freshAddress(KeyChain.KeyPurpose.RECEIVE_FUNDS);
+            assertTrue(group2.findRedeemScriptFromPubHash(address.getHash160()) != null);
+        }
+        assertFalse(group2.findRedeemScriptFromPubHash(group.freshAddress(KeyChain.KeyPurpose.RECEIVE_FUNDS).getHash160()) != null);
+    }
+
+    @Test
+    public void bloomFilterForMarriedChains() throws Exception {
+        group = createMarriedKeyChainGroup();
+        // only leaf keys are used for populating bloom filter, so initial number is zero
+        assertEquals(0, group.getBloomFilterElementCount());
+        Address address1 = group.freshAddress(KeyChain.KeyPurpose.RECEIVE_FUNDS);
+        final int size = (LOOKAHEAD_SIZE + 1 /* for the just created key */) * 2;
+        assertEquals(size, group.getBloomFilterElementCount());
+        BloomFilter filter = group.getBloomFilter(size, 0.001, (long)(Math.random() * Long.MAX_VALUE));
+        assertTrue(filter.contains(address1.getHash160()));
+
+        Address address2 = group.freshAddress(KeyChain.KeyPurpose.CHANGE);
+        assertFalse(filter.contains(address2.getHash160()));
+
+        // Check that the filter contains the lookahead buffer.
+        for (int i = 0; i < LOOKAHEAD_SIZE; i++) {
+            Address address = group.freshAddress(KeyChain.KeyPurpose.RECEIVE_FUNDS);
+            assertTrue(filter.contains(address.getHash160()));
+        }
+        // We ran ahead of the lookahead buffer.
+        assertFalse(filter.contains(group.freshAddress(KeyChain.KeyPurpose.RECEIVE_FUNDS).getHash160()));
+    }
+
     @Test
     public void earliestKeyTime() throws Exception {
         long now = Utils.currentTimeSeconds();   // mock
@@ -391,17 +429,14 @@ public class KeyChainGroupTest {
     @Test
     public void serializeMarried() throws Exception {
         group = createMarriedKeyChainGroup();
-        DeterministicKeyChain keyChain = group.getActiveKeyChain();
+        Address address1 = group.currentAddress(KeyChain.KeyPurpose.RECEIVE_FUNDS);
-        keyChain.getKey(KeyChain.KeyPurpose.RECEIVE_FUNDS);
+        assertTrue(group.isMarried());
-        DeterministicKey key1 = keyChain.getKey(KeyChain.KeyPurpose.RECEIVE_FUNDS);
-        ImmutableList<ChildNumber> path = key1.getPath();
-        assertTrue(group.isMarried(keyChain));
 
-        List<Protos.Key> protoKeys3 = group.serializeToProtobuf();
+        List<Protos.Key> protoKeys = group.serializeToProtobuf();
-        group = KeyChainGroup.fromProtobufUnencrypted(params, protoKeys3);
+        KeyChainGroup group2 = KeyChainGroup.fromProtobufUnencrypted(params, protoKeys);
-        assertTrue(group.isMarried(keyChain));
+        assertTrue(group2.isMarried());
-        DeterministicKey key2 = keyChain.getKeyByPath(path);
+        Address address2 = group2.currentAddress(KeyChain.KeyPurpose.RECEIVE_FUNDS);
-        assertEquals(key1, key2);
+        assertEquals(address1, address2);
     }
 
     @Test

tools/src/main/java/com/google/bitcoin/tools/WalletTool.java

@@ -80,6 +80,7 @@ public class WalletTool {
     private static OptionSpec<Date> dateFlag;
     private static OptionSpec<Integer> unixtimeFlag;
     private static OptionSpec<String> seedFlag, watchFlag;
+    private static OptionSpec<String> xpubkeysFlag;
 
     private static NetworkParameters params;
     private static File walletFile;
@@ -162,6 +163,7 @@ public class WalletTool {
         SEND,
         ENCRYPT,
         DECRYPT,
+        MARRY
     }
 
     public enum WaitForEnum {
@@ -202,6 +204,7 @@ public class WalletTool {
         parser.accepts("privkey").withRequiredArg();
         parser.accepts("addr").withRequiredArg();
         parser.accepts("peers").withRequiredArg();
+        xpubkeysFlag = parser.accepts("xpubkeys").withRequiredArg();
         OptionSpec<String> outputFlag = parser.accepts("output").withRequiredArg();
         parser.accepts("value").withRequiredArg();
         parser.accepts("fee").withRequiredArg();
@@ -352,6 +355,7 @@ public class WalletTool {
                 break;
             case ENCRYPT: encrypt(); break;
             case DECRYPT: decrypt(); break;
+            case MARRY: marry(); break;
         }
 
         if (!wallet.isConsistent()) {
@@ -380,6 +384,19 @@ public class WalletTool {
         shutdown();
     }
 
+    private static void marry() {
+        if (!options.has(xpubkeysFlag)) {
+            throw new IllegalStateException();
+        }
+
+        String[] xpubkeys = options.valueOf(xpubkeysFlag).split(",");
+        ImmutableList.Builder<DeterministicKey> keys = ImmutableList.builder();
+        for (String xpubkey : xpubkeys) {
+            keys.add(DeterministicKey.deserializeB58(null, xpubkey.trim()));
+        }
+        wallet.addFollowingAccountKeys(keys.build());
+    }
+
     private static void encrypt() {
         if (password == null) {
             System.err.println("You must provide a --password");

tools/src/main/resources/com/google/bitcoin/tools/wallet-tool-help.txt

@@ -9,6 +9,8 @@ Usage: wallet-tool --flags action-name
   create               Makes a new wallet in the file specified by --wallet.
                        Will complain and require --force if the wallet already exists.
                        If --seed is present, it should specify either a mnemonic code or hex/base58 raw seed bytes.
+  marry                Makes the wallet married with other parties, requiring multisig to spend funds.
+                       External public keys for other signing parties must be specified with --xpubkeys (comma separated).
   add-key              Adds a new key to the wallet, either specified or freshly generated.
                        If --date is specified, that's the creation date.
                        If --unixtime is specified, that's the creation time and it overrides --date.