addrman: add functionality to select by network

Add an optional parameter to the addrman Select function that allows callers to
specify which network the returned address should be on. Ensure that the proper
table is selected with different cases of whether the new or tried table has
network addresses that match.

Co-authored-by: Martin Zumsande <mzumsande@gmail.com>

3 files changed, 43 insertions, 20 deletions

diff --git a/src/addrman.cpp b/src/addrman.cpp
index fe4a79b2e6..cdfd079fcd 100644
--- a/src/addrman.cpp
+++ b/src/addrman.cpp
@@ -714,28 +714,41 @@ void AddrManImpl::Attempt_(const CService& addr, bool fCountFailure, NodeSeconds
     }
 }
 
-std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only) const
+std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only, std::optional<Network> network) const
 {
     AssertLockHeld(cs);
 
     if (vRandom.empty()) return {};
-    if (new_only && nNew == 0) return {};
+
+    size_t new_count = nNew;
+    size_t tried_count = nTried;
+
+    if (network.has_value()) {
+        auto it = m_network_counts.find(*network);
+        if (it == m_network_counts.end()) return {};
+
+        auto counts = it->second;
+        new_count = counts.n_new;
+        tried_count = counts.n_tried;
+    }
+
+    if (new_only && new_count == 0) return {};
+    if (new_count + tried_count == 0) return {};
 
     // Decide if we are going to search the new or tried table
+    // If either option is viable, use a 50% chance to choose
     bool search_tried;
-    int bucket_count;
-
-    // Use a 50% chance for choosing between tried and new table entries.
-    if (!new_only &&
-       (nTried > 0 &&
-        (nNew == 0 || insecure_rand.randbool() == 0))) {
+    if (new_only || tried_count == 0) {
+        search_tried = false;
+    } else if (new_count == 0) {
         search_tried = true;
-        bucket_count = ADDRMAN_TRIED_BUCKET_COUNT;
     } else {
-        search_tried = false;
-        bucket_count = ADDRMAN_NEW_BUCKET_COUNT;
+        search_tried = insecure_rand.randbool();
     }
 
+    const int bucket_count{search_tried ? ADDRMAN_TRIED_BUCKET_COUNT : ADDRMAN_NEW_BUCKET_COUNT};
+
+    // Loop through the addrman table until we find an appropriate entry
     double chance_factor = 1.0;
     while (1) {
         // Pick a bucket, and an initial position in that bucket.
@@ -748,7 +761,16 @@ std::pair<CAddress, NodeSeconds> AddrManImpl::Select_(bool new_only) const
         for (i = 0; i < ADDRMAN_BUCKET_SIZE; ++i) {
             int position = (initial_position + i) % ADDRMAN_BUCKET_SIZE;
             int node_id = GetEntry(search_tried, bucket, position);
-            if (node_id != -1) break;
+            if (node_id != -1) {
+                if (network.has_value()) {
+                    const auto it{mapInfo.find(node_id)};
+                    assert(it != mapInfo.end());
+                    const auto info{it->second};
+                    if (info.GetNetwork() == *network) break;
+                } else {
+                    break;
+                }
+            }
         }
 
         // If the bucket is entirely empty, start over with a (likely) different one.
@@ -1168,11 +1190,11 @@ std::pair<CAddress, NodeSeconds> AddrManImpl::SelectTriedCollision()
     return ret;
 }
 
-std::pair<CAddress, NodeSeconds> AddrManImpl::Select(bool new_only) const
+std::pair<CAddress, NodeSeconds> AddrManImpl::Select(bool new_only, std::optional<Network> network) const
 {
     LOCK(cs);
     Check();
-    auto addrRet = Select_(new_only);
+    auto addrRet = Select_(new_only, network);
     Check();
     return addrRet;
 }
@@ -1266,9 +1288,9 @@ std::pair<CAddress, NodeSeconds> AddrMan::SelectTriedCollision()
     return m_impl->SelectTriedCollision();
 }
 
-std::pair<CAddress, NodeSeconds> AddrMan::Select(bool new_only) const
+std::pair<CAddress, NodeSeconds> AddrMan::Select(bool new_only, std::optional<Network> network) const
 {
-    return m_impl->Select(new_only);
+    return m_impl->Select(new_only, network);
 }
 
 std::vector<CAddress> AddrMan::GetAddr(size_t max_addresses, size_t max_pct, std::optional<Network> network) const
diff --git a/src/addrman.h b/src/addrman.h
index 374996e501..c8e31fe283 100644
--- a/src/addrman.h
+++ b/src/addrman.h
@@ -146,11 +146,12 @@ public:
     /**
      * Choose an address to connect to.
      *
-     * @param[in] new_only  Whether to only select addresses from the new table.
+     * @param[in] new_only Whether to only select addresses from the new table.
+     * @param[in] network  Select only addresses of this network (nullopt = all)
      * @return    CAddress The record for the selected peer.
      *            seconds  The last time we attempted to connect to that peer.
      */
-    std::pair<CAddress, NodeSeconds> Select(bool new_only = false) const;
+    std::pair<CAddress, NodeSeconds> Select(bool new_only = false, std::optional<Network> network = std::nullopt) const;
 
     /**
      * Return all or many randomly selected addresses, optionally by network.
diff --git a/src/addrman_impl.h b/src/addrman_impl.h
index 3cf3b838d6..7aead2812b 100644
--- a/src/addrman_impl.h
+++ b/src/addrman_impl.h
@@ -127,7 +127,7 @@ public:
 
     std::pair<CAddress, NodeSeconds> SelectTriedCollision() EXCLUSIVE_LOCKS_REQUIRED(!cs);
 
-    std::pair<CAddress, NodeSeconds> Select(bool new_only) const
+    std::pair<CAddress, NodeSeconds> Select(bool new_only, std::optional<Network> network) const
         EXCLUSIVE_LOCKS_REQUIRED(!cs);
 
     std::vector<CAddress> GetAddr(size_t max_addresses, size_t max_pct, std::optional<Network> network) const
@@ -251,7 +251,7 @@ private:
 
     void Attempt_(const CService& addr, bool fCountFailure, NodeSeconds time) EXCLUSIVE_LOCKS_REQUIRED(cs);
 
-    std::pair<CAddress, NodeSeconds> Select_(bool new_only) const EXCLUSIVE_LOCKS_REQUIRED(cs);
+    std::pair<CAddress, NodeSeconds> Select_(bool new_only, std::optional<Network> network) const EXCLUSIVE_LOCKS_REQUIRED(cs);
 
     /** Helper to generalize looking up an addrman entry from either table.
      *