SERVER-32631 Deduplicate results in SockAddr::createAll()

mongodb · Jan 12, 2018 · 62c9eed · 62c9eed
1 parent ca539d9
commit 62c9eed
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 4 deletions.
diff --git a/src/mongo/util/net/sockaddr.cpp b/src/mongo/util/net/sockaddr.cpp
@@ -31,6 +31,8 @@
 #include "mongo/platform/basic.h"
 
 #include <iterator>
+#include <set>
+#include <vector>
 
 #include "mongo/util/net/sockaddr.h"
 
@@ -177,7 +179,7 @@ std::vector<SockAddr> SockAddr::createAll(StringData target, int port, sa_family
         return {};
     }
 
-    std::vector<SockAddr> ret;
+    std::set<SockAddr> ret;
     struct sockaddr_storage storage;
     memset(&storage, 0, sizeof(storage));
     for (const auto* addrs = addrErr.second.get(); addrs; addrs = addrs->ai_next) {
@@ -186,9 +188,9 @@ std::vector<SockAddr> SockAddr::createAll(StringData target, int port, sa_family
         // SockAddr constructor below can copy the entire buffer
         // without over-running addrinfo's storage
         memcpy(&storage, addrs->ai_addr, addrs->ai_addrlen);
-        ret.emplace_back(storage, addrs->ai_addrlen);
+        ret.emplace(storage, addrs->ai_addrlen);
     }
-    return ret;
+    return std::vector<SockAddr>(ret.begin(), ret.end());
 }
 
 SockAddr::SockAddr(struct sockaddr_storage& other, socklen_t size)
@@ -323,4 +325,44 @@ bool SockAddr::operator!=(const SockAddr& r) const {
     return !(*this == r);
 }
 
+bool SockAddr::operator<(const SockAddr& r) const {
+    // Address family first
+    if (getType() < r.getType()) {
+        return true;
+    }
+    if (getType() > r.getType()) {
+        return false;
+    }
+
+    // Address second
+    int cmp;
+    switch (getType()) {
+        case AF_INET: {
+            const auto laddr = ntohl(as<sockaddr_in>().sin_addr.s_addr);
+            const auto raddr = ntohl(r.as<sockaddr_in>().sin_addr.s_addr);
+            cmp = (laddr < raddr) ? -1 : (laddr > raddr) ? 1 : 0;
+            break;
+        }
+        case AF_INET6:
+            cmp = memcmp(as<sockaddr_in6>().sin6_addr.s6_addr,
+                         r.as<sockaddr_in6>().sin6_addr.s6_addr,
+                         sizeof(in6_addr));
+            break;
+        case AF_UNIX:
+            cmp = strcmp(as<sockaddr_un>().sun_path, r.as<sockaddr_un>().sun_path);
+            break;
+        default:
+            massert(SOCK_FAMILY_UNKNOWN_ERROR, "unsupported address family", false);
+    }
+    if (cmp < 0) {
+        return true;
+    }
+    if (cmp > 0) {
+        return false;
+    }
+
+    // All else equal, compare port
+    return getPort() < r.getPort();
+}
+
 }  // namespace mongo
diff --git a/src/mongo/util/net/sockaddr.h b/src/mongo/util/net/sockaddr.h
@@ -136,8 +136,8 @@ struct SockAddr {
     bool isAnonymousUNIXSocket() const;
 
     bool operator==(const SockAddr& r) const;
-
     bool operator!=(const SockAddr& r) const;
+    bool operator<(const SockAddr& r) const;
 
     const sockaddr* raw() const {
         return (sockaddr*)&sa;