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>

1 files changed, 38 insertions, 16 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