From 976ca80056c309b992610dfdcac6f8f20c0de9a9 Mon Sep 17 00:00:00 2001
From: Peter Johnson <johnson.peter@gmail.com>
Date: Sun, 25 Dec 2016 01:30:12 -0800
Subject: [PATCH] Use InetNtop on Win32 rather than WSAAddressToString. (#170)

Similarly, use InetPton rather than WSAStringToAddress.

The WSAAddressToString function is intended to provide a user-readable
string and thus includes the port number.  This breaks some use cases
on Windows which expect to get just the IP address.

Note: The InetPton and InetNtop functions are available only in Vista or above.
---
 src/tcpsockets/TCPAcceptor.cpp  | 10 +++++++---
 src/tcpsockets/TCPConnector.cpp | 10 +++++-----
 src/tcpsockets/TCPStream.cpp    |  4 ++--
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/tcpsockets/TCPAcceptor.cpp b/src/tcpsockets/TCPAcceptor.cpp
index c46cca342a..db5d608722 100644
--- a/src/tcpsockets/TCPAcceptor.cpp
+++ b/src/tcpsockets/TCPAcceptor.cpp
@@ -27,6 +27,7 @@
 #include <cstring>
 #ifdef _WIN32
 #include <WinSock2.h>
+#include <Ws2tcpip.h>
 #pragma comment(lib, "Ws2_32.lib")
 #else
 #include <arpa/inet.h>
@@ -85,11 +86,14 @@ int TCPAcceptor::start() {
 #ifdef _WIN32
     llvm::SmallString<128> addr_copy(m_address);
     addr_copy.push_back('\0');
-    int size = sizeof(address);
-    WSAStringToAddress(addr_copy.data(), PF_INET, nullptr, (struct sockaddr*)&address, &size);
+    int res = InetPton(PF_INET, addr_copy.data(), &(address.sin_addr));
 #else
-    inet_pton(PF_INET, m_address.c_str(), &(address.sin_addr));
+    int res = inet_pton(PF_INET, m_address.c_str(), &(address.sin_addr));
 #endif
+    if (res != 1) {
+      WPI_ERROR(m_logger, "could not resolve " << m_address << " address");
+      return -1;
+    }
   } else {
     address.sin_addr.s_addr = INADDR_ANY;
   }
diff --git a/src/tcpsockets/TCPConnector.cpp b/src/tcpsockets/TCPConnector.cpp
index 9c6d6128af..ac563d17a6 100644
--- a/src/tcpsockets/TCPConnector.cpp
+++ b/src/tcpsockets/TCPConnector.cpp
@@ -89,14 +89,14 @@ std::unique_ptr<NetworkStream> TCPConnector::connect(const char* server,
 #ifdef _WIN32
     llvm::SmallString<128> addr_copy(server);
     addr_copy.push_back('\0');
-    int size = sizeof(address);
-    if (WSAStringToAddress(addr_copy.data(), PF_INET, nullptr, (struct sockaddr*)&address, &size) != 0) {
+    int res = InetPton(PF_INET, addr_copy.data(), &(address.sin_addr));
+#else
+    int res = inet_pton(PF_INET, server, &(address.sin_addr));
+#endif
+    if (res != 1) {
       WPI_ERROR(logger, "could not resolve " << server << " address");
       return nullptr;
     }
-#else
-    inet_pton(PF_INET, server, &(address.sin_addr));
-#endif
   }
   address.sin_port = htons(port);
 
diff --git a/src/tcpsockets/TCPStream.cpp b/src/tcpsockets/TCPStream.cpp
index e7d028e21a..d738abef9d 100644
--- a/src/tcpsockets/TCPStream.cpp
+++ b/src/tcpsockets/TCPStream.cpp
@@ -26,6 +26,7 @@
 #include <fcntl.h>
 #ifdef _WIN32
 #include <WinSock2.h>
+#include <Ws2tcpip.h>
 #else
 #include <arpa/inet.h>
 #include <netinet/tcp.h>
@@ -38,8 +39,7 @@ TCPStream::TCPStream(int sd, sockaddr_in* address)
     : m_sd(sd), m_blocking(true) {
   char ip[50];
 #ifdef _WIN32
-  unsigned long size = sizeof(ip) - 1;
-  WSAAddressToString((sockaddr*)address, sizeof sockaddr_in, nullptr, ip, &size);
+  InetNtop(PF_INET, &(address->sin_addr.s_addr), ip, sizeof(ip) - 1);
 #else
   inet_ntop(PF_INET, (in_addr*)&(address->sin_addr.s_addr), ip,
             sizeof(ip) - 1);