From dad5f760211df314d650999e0a76edb0151b4fe1 Mon Sep 17 00:00:00 2001
From: Amiti Uttarwar <amiti@uttarwar.org>
Date: Thu, 9 Dec 2021 18:25:22 +0000
Subject: addrman: Introduce a test-only function to lookup addresses

Co-Authored-By: Martin Zumsande <mzumsande@gmail.com>
---
 src/addrman.cpp    | 37 +++++++++++++++++++++++++++++++++++++
 src/addrman.h      | 34 ++++++++++++++++++++++++++++++++++
 src/addrman_impl.h |  5 +++++
 3 files changed, 76 insertions(+)

diff --git a/src/addrman.cpp b/src/addrman.cpp
index 4d785043b8..f2fe02d06f 100644
--- a/src/addrman.cpp
+++ b/src/addrman.cpp
@@ -930,6 +930,29 @@ std::pair<CAddress, int64_t> AddrManImpl::SelectTriedCollision_()
     return {info_old, info_old.nLastTry};
 }
 
+std::optional<AddressPosition> AddrManImpl::FindAddressEntry_(const CAddress& addr)
+{
+    AssertLockHeld(cs);
+
+    AddrInfo* addr_info = Find(addr);
+
+    if (!addr_info) return std::nullopt;
+
+    if(addr_info->fInTried) {
+        int bucket{addr_info->GetTriedBucket(nKey, m_asmap)};
+        return AddressPosition(/*tried=*/true,
+                               /*multiplicity=*/1,
+                               /*bucket=*/bucket,
+                               /*position=*/addr_info->GetBucketPosition(nKey, false, bucket));
+    } else {
+        int bucket{addr_info->GetNewBucket(nKey, m_asmap)};
+        return AddressPosition(/*tried=*/false,
+                               /*multiplicity=*/addr_info->nRefCount,
+                               /*bucket=*/bucket,
+                               /*position=*/addr_info->GetBucketPosition(nKey, true, bucket));
+    }
+}
+
 void AddrManImpl::Check() const
 {
     AssertLockHeld(cs);
@@ -1116,6 +1139,15 @@ void AddrManImpl::SetServices(const CService& addr, ServiceFlags nServices)
     Check();
 }
 
+std::optional<AddressPosition> AddrManImpl::FindAddressEntry(const CAddress& addr)
+{
+    LOCK(cs);
+    Check();
+    auto entry = FindAddressEntry_(addr);
+    Check();
+    return entry;
+}
+
 const std::vector<bool>& AddrManImpl::GetAsmap() const
 {
     return m_asmap;
@@ -1201,3 +1233,8 @@ const std::vector<bool>& AddrMan::GetAsmap() const
 {
     return m_impl->GetAsmap();
 }
+
+std::optional<AddressPosition> AddrMan::FindAddressEntry(const CAddress& addr)
+{
+    return m_impl->FindAddressEntry(addr);
+}
diff --git a/src/addrman.h b/src/addrman.h
index e9bc372380..1181bdd300 100644
--- a/src/addrman.h
+++ b/src/addrman.h
@@ -22,6 +22,31 @@ class AddrManImpl;
 /** Default for -checkaddrman */
 static constexpr int32_t DEFAULT_ADDRMAN_CONSISTENCY_CHECKS{0};
 
+/** Test-only struct, capturing info about an address in AddrMan */
+struct AddressPosition {
+    // Whether the address is in the new or tried table
+    const bool tried;
+
+    // Addresses in the tried table should always have a multiplicity of 1.
+    // Addresses in the new table can have multiplicity between 1 and
+    // ADDRMAN_NEW_BUCKETS_PER_ADDRESS
+    const int multiplicity;
+
+    // If the address is in the new table, the bucket and position are
+    // populated based on the first source who sent the address.
+    // In certain edge cases, this may not be where the address is currently
+    // located.
+    const int bucket;
+    const int position;
+
+    bool operator==(AddressPosition other) {
+        return std::tie(tried, multiplicity, bucket, position) ==
+               std::tie(other.tried, other.multiplicity, other.bucket, other.position);
+    }
+    explicit AddressPosition(bool tried_in, int multiplicity_in, int bucket_in, int position_in)
+        : tried{tried_in}, multiplicity{multiplicity_in}, bucket{bucket_in}, position{position_in} {}
+};
+
 /** Stochastic address manager
  *
  * Design goals:
@@ -142,6 +167,15 @@ public:
     void SetServices(const CService& addr, ServiceFlags nServices);
 
     const std::vector<bool>& GetAsmap() const;
+
+    /** Test-only function
+     * Find the address record in AddrMan and return information about its
+     * position.
+     * @param[in] addr       The address record to look up.
+     * @return               Information about the address record in AddrMan
+     *                       or nullopt if address is not found.
+     */
+    std::optional<AddressPosition> FindAddressEntry(const CAddress& addr);
 };
 
 #endif // BITCOIN_ADDRMAN_H
diff --git a/src/addrman_impl.h b/src/addrman_impl.h
index bd7caf473b..04d5fe3053 100644
--- a/src/addrman_impl.h
+++ b/src/addrman_impl.h
@@ -137,6 +137,9 @@ public:
     void SetServices(const CService& addr, ServiceFlags nServices)
         EXCLUSIVE_LOCKS_REQUIRED(!cs);
 
+    std::optional<AddressPosition> FindAddressEntry(const CAddress& addr)
+        EXCLUSIVE_LOCKS_REQUIRED(!cs);
+
     const std::vector<bool>& GetAsmap() const;
 
     friend class AddrManTest;
@@ -266,6 +269,8 @@ private:
 
     std::pair<CAddress, int64_t> SelectTriedCollision_() EXCLUSIVE_LOCKS_REQUIRED(cs);
 
+    std::optional<AddressPosition> FindAddressEntry_(const CAddress& addr) EXCLUSIVE_LOCKS_REQUIRED(cs);
+
     //! Consistency check, taking into account m_consistency_check_ratio.
     //! Will std::abort if an inconsistency is detected.
     void Check() const EXCLUSIVE_LOCKS_REQUIRED(cs);
-- 
cgit v1.2.3