From 11463e729f5973da856c4d2ce317e83fb0d3bdb4 Mon Sep 17 00:00:00 2001
From: Mike Hearn <mike@plan99.net>
Date: Fri, 9 Jan 2015 15:17:56 +0100
Subject: [PATCH] Add a feature to PeerGroup that lets you find/wait for peers
 that match a certain mask, similar to those methods that exist for peer
 versions.

---
 .../java/org/bitcoinj/core/PeerGroup.java     | 44 ++++++++++++++++++-
 .../java/org/bitcoinj/core/PeerGroupTest.java | 25 +++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/bitcoinj/core/PeerGroup.java b/core/src/main/java/org/bitcoinj/core/PeerGroup.java
index db88f143..e6dd1162 100644
--- a/core/src/main/java/org/bitcoinj/core/PeerGroup.java
+++ b/core/src/main/java/org/bitcoinj/core/PeerGroup.java
@@ -1501,7 +1501,7 @@ public class PeerGroup implements TransactionBroadcaster {
     }
 
     /**
-     * Returns a mutable array list of peers that implement the given protocol version or better.
+     * Returns an array list of peers that implement the given protocol version or better.
      */
     public List<Peer> findPeersOfAtLeastVersion(long protocolVersion) {
         lock.lock();
@@ -1516,6 +1516,48 @@ public class PeerGroup implements TransactionBroadcaster {
         }
     }
 
+    /**
+     * Returns a future that is triggered when there are at least the requested number of connected peers that support
+     * the given protocol version or higher. To block immediately, just call get() on the result.
+     *
+     * @param numPeers How many peers to wait for.
+     * @param mask An integer representing a bit mask that will be ANDed with the peers advertised service masks.
+     * @return a future that will be triggered when the number of connected peers implementing protocolVersion or higher >= numPeers
+     */
+    public ListenableFuture<List<Peer>> waitForPeersWithServiceMask(final int numPeers, final int mask) {
+        List<Peer> foundPeers = findPeersWithServiceMask(mask);
+        if (foundPeers.size() >= numPeers)
+            return Futures.immediateFuture(foundPeers);
+        final SettableFuture<List<Peer>> future = SettableFuture.create();
+        addEventListener(new AbstractPeerEventListener() {
+            @Override
+            public void onPeerConnected(Peer peer, int peerCount) {
+                final List<Peer> peers = findPeersWithServiceMask(mask);
+                if (peers.size() >= numPeers) {
+                    future.set(peers);
+                    removeEventListener(this);
+                }
+            }
+        });
+        return future;
+    }
+
+    /**
+     * Returns an array list of peers that match the requested service bit mask.
+     */
+    public List<Peer> findPeersWithServiceMask(int mask) {
+        lock.lock();
+        try {
+            ArrayList<Peer> results = new ArrayList<Peer>(peers.size());
+            for (Peer peer : peers)
+                if ((peer.getPeerVersionMessage().localServices & mask) == mask)
+                    results.add(peer);
+            return results;
+        } finally {
+            lock.unlock();
+        }
+    }
+
     /**
      * Returns the number of connections that are required before transactions will be broadcast. If there aren't
      * enough, {@link PeerGroup#broadcastTransaction(Transaction)} will wait until the minimum number is reached so
diff --git a/core/src/test/java/org/bitcoinj/core/PeerGroupTest.java b/core/src/test/java/org/bitcoinj/core/PeerGroupTest.java
index 67746bd4..181afb2f 100644
--- a/core/src/test/java/org/bitcoinj/core/PeerGroupTest.java
+++ b/core/src/test/java/org/bitcoinj/core/PeerGroupTest.java
@@ -698,6 +698,31 @@ public class PeerGroupTest extends TestWithPeerGroup {
         assertTrue(future.isDone());
     }
 
+    @Test
+    public void waitForPeersWithServiceFlags() throws Exception {
+        ListenableFuture<List<Peer>> future = peerGroup.waitForPeersWithServiceMask(2, 3);
+
+        VersionMessage ver1 = new VersionMessage(params, 10);
+        ver1.clientVersion = 70000;
+        ver1.localServices = VersionMessage.NODE_NETWORK;
+        VersionMessage ver2 = new VersionMessage(params, 10);
+        ver2.clientVersion = 70000;
+        ver2.localServices = VersionMessage.NODE_NETWORK | 2;
+        peerGroup.start();
+        assertFalse(future.isDone());
+        connectPeer(1, ver1);
+        assertTrue(peerGroup.findPeersWithServiceMask(3).isEmpty());
+        assertFalse(future.isDone());
+        connectPeer(2, ver2);
+        assertFalse(future.isDone());
+        assertEquals(1, peerGroup.findPeersWithServiceMask(3).size());
+        assertTrue(peerGroup.waitForPeersWithServiceMask(1, 0x3).isDone());   // Immediate completion.
+        connectPeer(3, ver2);
+        future.get();
+        assertTrue(future.isDone());
+        peerGroup.stop();
+    }
+
     @Test
     public void preferLocalPeer() throws IOException {
         // Because we are using the same port (8333 or 18333) that is used by Satoshi client