[msan] Choose in runtime if IPv4 or IPv6 are supported.

Summary: This reverts commit cab5051c691ce27a7ffac41e8e76ceb222ad9549.

Reviewers: eugenis

Subscribers: mgorny, llvm-commits

Differential Revision: https://reviews.llvm.org/D31894

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@299884 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 54 insertions, 9 deletions

diff --git a/lib/msan/tests/CMakeLists.txt b/lib/msan/tests/CMakeLists.txt
index 8e911dc10..65fbc732d 100644
--- a/lib/msan/tests/CMakeLists.txt
+++ b/lib/msan/tests/CMakeLists.txt
@@ -35,6 +35,7 @@ set(MSAN_UNITTEST_COMMON_CFLAGS
   -Wno-zero-length-array
   -Wno-uninitialized
   -Werror=sign-compare
+  -Wno-gnu-zero-variadic-macro-arguments
 )
 set(MSAN_UNITTEST_INSTRUMENTED_CFLAGS
   ${MSAN_UNITTEST_COMMON_CFLAGS}
diff --git a/lib/msan/tests/msan_test.cc b/lib/msan/tests/msan_test.cc
index e85db097d..12d6092ef 100644
--- a/lib/msan/tests/msan_test.cc
+++ b/lib/msan/tests/msan_test.cc
@@ -902,9 +902,7 @@ class SocketAddr4 : public SocketAddr {
     sai_.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
   }
 
-  sockaddr *ptr() override {
-    return reinterpret_cast<sockaddr *>(&sai_);
-  }
+  sockaddr *ptr() override { return reinterpret_cast<sockaddr *>(&sai_); }
 
   size_t size() const override { return sizeof(sai_); }
 
@@ -912,14 +910,60 @@ class SocketAddr4 : public SocketAddr {
   sockaddr_in sai_;
 };
 
-template <class... Args>
-std::unique_ptr<SocketAddr> CreateSockAddr(Args... args) {
-  return std::unique_ptr<SocketAddr>(new SocketAddr4(args...));
+class SocketAddr6 : public SocketAddr {
+ public:
+  SocketAddr6() { EXPECT_POISONED(sai_); }
+  explicit SocketAddr6(uint16_t port) {
+    memset(&sai_, 0, sizeof(sai_));
+    sai_.sin6_family = AF_INET6;
+    sai_.sin6_port = port;
+    sai_.sin6_addr = in6addr_loopback;
+  }
+
+  sockaddr *ptr() override { return reinterpret_cast<sockaddr *>(&sai_); }
+
+  size_t size() const override { return sizeof(sai_); }
+
+ private:
+  sockaddr_in6 sai_;
+};
+
+class MemorySanitizerIpTest : public ::testing::TestWithParam<int> {
+ public:
+  void SetUp() override {
+    ASSERT_TRUE(GetParam() == AF_INET || GetParam() == AF_INET6);
+  }
+
+  template <class... Args>
+  std::unique_ptr<SocketAddr> CreateSockAddr(Args... args) const {
+    if (GetParam() == AF_INET)
+      return std::unique_ptr<SocketAddr>(new SocketAddr4(args...));
+    return std::unique_ptr<SocketAddr>(new SocketAddr6(args...));
+  }
+
+  int CreateSocket(int socket_type) const {
+    return socket(GetParam(), socket_type, 0);
+  }
+};
+
+std::vector<int> GetAvailableIpSocketFamilies() {
+  std::vector<int> result;
+
+  for (int i : std::vector<int>(AF_INET, AF_INET6)) {
+    int s = socket(i, SOCK_STREAM, 0);
+    if (s > 0) {
+      result.push_back(i);
+      close(s);
+    }
+  }
+
+  return result;
 }
 
-int CreateSocket(int socket_type) { return socket(AF_INET, socket_type, 0); }
+INSTANTIATE_TEST_CASE_P(IpTests, MemorySanitizerIpTest,
+                        ::testing::ValuesIn(GetAvailableIpSocketFamilies()));
 
-TEST(MemorySanitizer, accept) {
+TEST_P(MemorySanitizerIpTest, accept) {
   int listen_socket = CreateSocket(SOCK_STREAM);
   ASSERT_LT(0, listen_socket);
 
@@ -963,7 +1007,7 @@ TEST(MemorySanitizer, accept) {
   close(listen_socket);
 }
 
-TEST(MemorySanitizer, recvmsg) {
+TEST_P(MemorySanitizerIpTest, recvmsg) {
   int server_socket = CreateSocket(SOCK_DGRAM);
   ASSERT_LT(0, server_socket);