1 files changed, 48 insertions, 11 deletions

diff --git a/src/netaddress.cpp b/src/netaddress.cpp
index 2ee4ae3ee3..ce3e17197e 100644
--- a/src/netaddress.cpp
+++ b/src/netaddress.cpp
@@ -6,6 +6,7 @@
 #include <netaddress.h>
 #include <hash.h>
 #include <util/strencodings.h>
+#include <util/asmap.h>
 #include <tinyformat.h>
 
 static const unsigned char pchIPv4[12] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff };
@@ -400,6 +401,39 @@ bool CNetAddr::GetIn6Addr(struct in6_addr* pipv6Addr) const
     return true;
 }
 
+uint32_t CNetAddr::GetNetClass() const {
+    uint32_t net_class = NET_IPV6;
+    if (IsLocal()) {
+        net_class = 255;
+    }
+    if (IsInternal()) {
+        net_class = NET_INTERNAL;
+    } else if (!IsRoutable()) {
+        net_class = NET_UNROUTABLE;
+    } else if (IsIPv4() || IsRFC6145() || IsRFC6052() || IsRFC3964() || IsRFC4380()) {
+        net_class = NET_IPV4;
+    } else if (IsTor()) {
+        net_class = NET_ONION;
+    }
+    return net_class;
+}
+
+uint32_t CNetAddr::GetMappedAS(const std::vector<bool> &asmap) const {
+    uint32_t net_class = GetNetClass();
+    if (asmap.size() == 0 || (net_class != NET_IPV4 && net_class != NET_IPV6)) {
+        return 0; // Indicates not found, safe because AS0 is reserved per RFC7607.
+    }
+    std::vector<bool> ip_bits(128);
+    for (int8_t byte_i = 0; byte_i < 16; ++byte_i) {
+        uint8_t cur_byte = GetByte(15 - byte_i);
+        for (uint8_t bit_i = 0; bit_i < 8; ++bit_i) {
+            ip_bits[byte_i * 8 + bit_i] = (cur_byte >> (7 - bit_i)) & 1;
+        }
+    }
+    uint32_t mapped_as = Interpret(asmap, ip_bits);
+    return mapped_as;
+}
+
 /**
  * Get the canonical identifier of our network group
  *
@@ -410,56 +444,61 @@ bool CNetAddr::GetIn6Addr(struct in6_addr* pipv6Addr) const
  * @note No two connections will be attempted to addresses with the same network
  *       group.
  */
-std::vector<unsigned char> CNetAddr::GetGroup() const
+std::vector<unsigned char> CNetAddr::GetGroup(const std::vector<bool> &asmap) const
 {
     std::vector<unsigned char> vchRet;
-    int nClass = NET_IPV6;
+    uint32_t net_class = GetNetClass();
+    // If non-empty asmap is supplied and the address is IPv4/IPv6,
+    // return ASN to be used for bucketing.
+    uint32_t asn = GetMappedAS(asmap);
+    if (asn != 0) { // Either asmap was empty, or address has non-asmappable net class (e.g. TOR).
+        vchRet.push_back(NET_IPV6); // IPv4 and IPv6 with same ASN should be in the same bucket
+        for (int i = 0; i < 4; i++) {
+            vchRet.push_back((asn >> (8 * i)) & 0xFF);
+        }
+        return vchRet;
+    }
+
+    vchRet.push_back(net_class);
     int nStartByte = 0;
     int nBits = 16;
 
     // all local addresses belong to the same group
     if (IsLocal())
     {
-        nClass = 255;
         nBits = 0;
     }
     // all internal-usage addresses get their own group
     if (IsInternal())
     {
-        nClass = NET_INTERNAL;
         nStartByte = sizeof(g_internal_prefix);
         nBits = (sizeof(ip) - sizeof(g_internal_prefix)) * 8;
     }
     // all other unroutable addresses belong to the same group
     else if (!IsRoutable())
     {
-        nClass = NET_UNROUTABLE;
         nBits = 0;
     }
     // for IPv4 addresses, '1' + the 16 higher-order bits of the IP
     // includes mapped IPv4, SIIT translated IPv4, and the well-known prefix
     else if (IsIPv4() || IsRFC6145() || IsRFC6052())
     {
-        nClass = NET_IPV4;
         nStartByte = 12;
     }
     // for 6to4 tunnelled addresses, use the encapsulated IPv4 address
     else if (IsRFC3964())
     {
-        nClass = NET_IPV4;
         nStartByte = 2;
     }
     // for Teredo-tunnelled IPv6 addresses, use the encapsulated IPv4 address
     else if (IsRFC4380())
     {
-        vchRet.push_back(NET_IPV4);
         vchRet.push_back(GetByte(3) ^ 0xFF);
         vchRet.push_back(GetByte(2) ^ 0xFF);
         return vchRet;
     }
     else if (IsTor())
     {
-        nClass = NET_ONION;
         nStartByte = 6;
         nBits = 4;
     }
@@ -470,8 +509,6 @@ std::vector<unsigned char> CNetAddr::GetGroup() const
     else
         nBits = 32;
 
-    vchRet.push_back(nClass);
-
     // push our ip onto vchRet byte by byte...
     while (nBits >= 8)
     {