net: use Sock in InterruptibleRecv() and Socks5()

vasild · practicalswift · vasild · commit e82e609c3960 · 2021-01-08T16:22:20.000+01:00
Use the `Sock` class instead of `SOCKET` for `InterruptibleRecv()` and
`Socks5()`.

This way the `Socks5()` function can be tested by giving it a mocked
instance of a socket.

Co-authored-by: practicalswift &lt;practicalswift@users.noreply.github.com&gt;
diff --git a/src/net.cpp b/src/net.cpp
@@ -439,7 +439,7 @@ CNode* CConnman::ConnectNode(CAddress addrConnect, const char *pszDest, bool fCo
                 return nullptr;
             }
             connected = ConnectThroughProxy(proxy, addrConnect.ToStringIP(), addrConnect.GetPort(),
-                                            sock->Get(), nConnectTimeout, proxyConnectionFailed);
+                                            *sock, nConnectTimeout, proxyConnectionFailed);
         } else {
             // no proxy needed (none set for target network)
             sock = CreateSock(addrConnect);
@@ -463,8 +463,8 @@ CNode* CConnman::ConnectNode(CAddress addrConnect, const char *pszDest, bool fCo
         int port = default_port;
         SplitHostPort(std::string(pszDest), port, host);
         bool proxyConnectionFailed;
-        connected = ConnectThroughProxy(proxy, host, port, sock->Get(), nConnectTimeout,
-                                        proxyConnectionFailed);
+        connected =
+            ConnectThroughProxy(proxy, host, port, *sock, nConnectTimeout, proxyConnectionFailed);
     }
     if (!connected) {
         return nullptr;
diff --git a/src/netbase.cpp b/src/netbase.cpp
@@ -331,8 +331,7 @@ enum class IntrRecvError {
  * @param data The buffer where the read bytes should be stored.
  * @param len The number of bytes to read into the specified buffer.
  * @param timeout The total timeout in milliseconds for this read.
- * @param hSocket The socket (has to be in non-blocking mode) from which to read
- *                bytes.
+ * @param sock The socket (has to be in non-blocking mode) from which to read bytes.
  *
  * @returns An IntrRecvError indicating the resulting status of this read.
  *          IntrRecvError::OK only if all of the specified number of bytes were
@@ -342,15 +341,15 @@ enum class IntrRecvError {
  *      Sockets can be made non-blocking with SetSocketNonBlocking(const
  *      SOCKET&, bool).
  */
-static IntrRecvError InterruptibleRecv(uint8_t* data, size_t len, int timeout, const SOCKET& hSocket)
+static IntrRecvError InterruptibleRecv(uint8_t* data, size_t len, int timeout, const Sock& sock)
 {
     int64_t curTime = GetTimeMillis();
     int64_t endTime = curTime + timeout;
     // Maximum time to wait for I/O readiness. It will take up until this time
     // (in millis) to break off in case of an interruption.
     const int64_t maxWait = 1000;
     while (len > 0 && curTime < endTime) {
-        ssize_t ret = recv(hSocket, (char*)data, len, 0); // Optimistically try the recv first
+        ssize_t ret = sock.Recv(data, len, 0); // Optimistically try the recv first
         if (ret > 0) {
             len -= ret;
             data += ret;
@@ -359,25 +358,10 @@ static IntrRecvError InterruptibleRecv(uint8_t* data, size_t len, int timeout, c
         } else { // Other error or blocking
             int nErr = WSAGetLastError();
             if (nErr == WSAEINPROGRESS || nErr == WSAEWOULDBLOCK || nErr == WSAEINVAL) {
-                if (!IsSelectableSocket(hSocket)) {
-                    return IntrRecvError::NetworkError;
-                }
                 // Only wait at most maxWait milliseconds at a time, unless
                 // we're approaching the end of the specified total timeout
                 int timeout_ms = std::min(endTime - curTime, maxWait);
-#ifdef USE_POLL
-                struct pollfd pollfd = {};
-                pollfd.fd = hSocket;
-                pollfd.events = POLLIN;
-                int nRet = poll(&pollfd, 1, timeout_ms);
-#else
-                struct timeval tval = MillisToTimeval(timeout_ms);
-                fd_set fdset;
-                FD_ZERO(&fdset);
-                FD_SET(hSocket, &fdset);
-                int nRet = select(hSocket + 1, &fdset, nullptr, nullptr, &tval);
-#endif
-                if (nRet == SOCKET_ERROR) {
+                if (!sock.Wait(std::chrono::milliseconds{timeout_ms}, Sock::RECV)) {
                     return IntrRecvError::NetworkError;
                 }
             } else {
@@ -431,7 +415,7 @@ static std::string Socks5ErrorString(uint8_t err)
  * @param port The destination port.
  * @param auth The credentials with which to authenticate with the specified
  *             SOCKS5 proxy.
- * @param hSocket The SOCKS5 proxy socket.
+ * @param sock The SOCKS5 proxy socket.
  *
  * @returns Whether or not the operation succeeded.
  *
@@ -441,7 +425,10 @@ static std::string Socks5ErrorString(uint8_t err)
  * @see <a href="https://www.ietf.org/rfc/rfc1928.txt">RFC1928: SOCKS Protocol
  *      Version 5</a>
  */
-static bool Socks5(const std::string& strDest, int port, const ProxyCredentials *auth, const SOCKET& hSocket)
+static bool Socks5(const std::string& strDest,
+                   int port,
+                   const ProxyCredentials* auth,
+                   const Sock& sock)
 {
     IntrRecvError recvr;
     LogPrint(BCLog::NET, "SOCKS5 connecting %s\n", strDest);
@@ -459,12 +446,12 @@ static bool Socks5(const std::string& strDest, int port, const ProxyCredentials
         vSocks5Init.push_back(0x01); // 1 method identifier follows...
         vSocks5Init.push_back(SOCKS5Method::NOAUTH);
     }
-    ssize_t ret = send(hSocket, (const char*)vSocks5Init.data(), vSocks5Init.size(), MSG_NOSIGNAL);
+    ssize_t ret = sock.Send(vSocks5Init.data(), vSocks5Init.size(), MSG_NOSIGNAL);
     if (ret != (ssize_t)vSocks5Init.size()) {
         return error("Error sending to proxy");
     }
     uint8_t pchRet1[2];
-    if ((recvr = InterruptibleRecv(pchRet1, 2, SOCKS5_RECV_TIMEOUT, hSocket)) != IntrRecvError::OK) {
+    if ((recvr = InterruptibleRecv(pchRet1, 2, SOCKS5_RECV_TIMEOUT, sock)) != IntrRecvError::OK) {
         LogPrintf("Socks5() connect to %s:%d failed: InterruptibleRecv() timeout or other failure\n", strDest, port);
         return false;
     }
@@ -481,13 +468,13 @@ static bool Socks5(const std::string& strDest, int port, const ProxyCredentials
         vAuth.insert(vAuth.end(), auth->username.begin(), auth->username.end());
         vAuth.push_back(auth->password.size());
         vAuth.insert(vAuth.end(), auth->password.begin(), auth->password.end());
-        ret = send(hSocket, (const char*)vAuth.data(), vAuth.size(), MSG_NOSIGNAL);
+        ret = sock.Send(vAuth.data(), vAuth.size(), MSG_NOSIGNAL);
         if (ret != (ssize_t)vAuth.size()) {
             return error("Error sending authentication to proxy");
         }
         LogPrint(BCLog::PROXY, "SOCKS5 sending proxy authentication %s:%s\n", auth->username, auth->password);
         uint8_t pchRetA[2];
-        if ((recvr = InterruptibleRecv(pchRetA, 2, SOCKS5_RECV_TIMEOUT, hSocket)) != IntrRecvError::OK) {
+        if ((recvr = InterruptibleRecv(pchRetA, 2, SOCKS5_RECV_TIMEOUT, sock)) != IntrRecvError::OK) {
             return error("Error reading proxy authentication response");
         }
         if (pchRetA[0] != 0x01 || pchRetA[1] != 0x00) {
@@ -507,12 +494,12 @@ static bool Socks5(const std::string& strDest, int port, const ProxyCredentials
     vSocks5.insert(vSocks5.end(), strDest.begin(), strDest.end());
     vSocks5.push_back((port >> 8) & 0xFF);
     vSocks5.push_back((port >> 0) & 0xFF);
-    ret = send(hSocket, (const char*)vSocks5.data(), vSocks5.size(), MSG_NOSIGNAL);
+    ret = sock.Send(vSocks5.data(), vSocks5.size(), MSG_NOSIGNAL);
     if (ret != (ssize_t)vSocks5.size()) {
         return error("Error sending to proxy");
     }
     uint8_t pchRet2[4];
-    if ((recvr = InterruptibleRecv(pchRet2, 4, SOCKS5_RECV_TIMEOUT, hSocket)) != IntrRecvError::OK) {
+    if ((recvr = InterruptibleRecv(pchRet2, 4, SOCKS5_RECV_TIMEOUT, sock)) != IntrRecvError::OK) {
         if (recvr == IntrRecvError::Timeout) {
             /* If a timeout happens here, this effectively means we timed out while connecting
              * to the remote node. This is very common for Tor, so do not print an
@@ -536,24 +523,24 @@ static bool Socks5(const std::string& strDest, int port, const ProxyCredentials
     uint8_t pchRet3[256];
     switch (pchRet2[3])
     {
-        case SOCKS5Atyp::IPV4: recvr = InterruptibleRecv(pchRet3, 4, SOCKS5_RECV_TIMEOUT, hSocket); break;
-        case SOCKS5Atyp::IPV6: recvr = InterruptibleRecv(pchRet3, 16, SOCKS5_RECV_TIMEOUT, hSocket); break;
+        case SOCKS5Atyp::IPV4: recvr = InterruptibleRecv(pchRet3, 4, SOCKS5_RECV_TIMEOUT, sock); break;
+        case SOCKS5Atyp::IPV6: recvr = InterruptibleRecv(pchRet3, 16, SOCKS5_RECV_TIMEOUT, sock); break;
         case SOCKS5Atyp::DOMAINNAME:
         {
-            recvr = InterruptibleRecv(pchRet3, 1, SOCKS5_RECV_TIMEOUT, hSocket);
+            recvr = InterruptibleRecv(pchRet3, 1, SOCKS5_RECV_TIMEOUT, sock);
             if (recvr != IntrRecvError::OK) {
                 return error("Error reading from proxy");
             }
             int nRecv = pchRet3[0];
-            recvr = InterruptibleRecv(pchRet3, nRecv, SOCKS5_RECV_TIMEOUT, hSocket);
+            recvr = InterruptibleRecv(pchRet3, nRecv, SOCKS5_RECV_TIMEOUT, sock);
             break;
         }
         default: return error("Error: malformed proxy response");
     }
     if (recvr != IntrRecvError::OK) {
         return error("Error reading from proxy");
     }
-    if ((recvr = InterruptibleRecv(pchRet3, 2, SOCKS5_RECV_TIMEOUT, hSocket)) != IntrRecvError::OK) {
+    if ((recvr = InterruptibleRecv(pchRet3, 2, SOCKS5_RECV_TIMEOUT, sock)) != IntrRecvError::OK) {
         return error("Error reading from proxy");
     }
     LogPrint(BCLog::NET, "SOCKS5 connected %s\n", strDest);
@@ -778,18 +765,23 @@ bool IsProxy(const CNetAddr &addr) {
  * @param proxy The SOCKS5 proxy.
  * @param strDest The destination service to which to connect.
  * @param port The destination port.
- * @param hSocket The socket on which to connect to the SOCKS5 proxy.
+ * @param sock The socket on which to connect to the SOCKS5 proxy.
  * @param nTimeout Wait this many milliseconds for the connection to the SOCKS5
  *                 proxy to be established.
  * @param[out] outProxyConnectionFailed Whether or not the connection to the
  *                                      SOCKS5 proxy failed.
  *
  * @returns Whether or not the operation succeeded.
  */
-bool ConnectThroughProxy(const proxyType &proxy, const std::string& strDest, int port, const SOCKET& hSocket, int nTimeout, bool& outProxyConnectionFailed)
+bool ConnectThroughProxy(const proxyType& proxy,
+                         const std::string& strDest,
+                         int port,
+                         const Sock& sock,
+                         int nTimeout,
+                         bool& outProxyConnectionFailed)
 {
     // first connect to proxy server
-    if (!ConnectSocketDirectly(proxy.proxy, hSocket, nTimeout, true)) {
+    if (!ConnectSocketDirectly(proxy.proxy, sock.Get(), nTimeout, true)) {
         outProxyConnectionFailed = true;
         return false;
     }
@@ -798,11 +790,11 @@ bool ConnectThroughProxy(const proxyType &proxy, const std::string& strDest, int
         ProxyCredentials random_auth;
         static std::atomic_int counter(0);
         random_auth.username = random_auth.password = strprintf("%i", counter++);
-        if (!Socks5(strDest, (uint16_t)port, &random_auth, hSocket)) {
+        if (!Socks5(strDest, (uint16_t)port, &random_auth, sock)) {
             return false;
         }
     } else {
-        if (!Socks5(strDest, (uint16_t)port, 0, hSocket)) {
+        if (!Socks5(strDest, (uint16_t)port, 0, sock)) {
             return false;
         }
     }
diff --git a/src/netbase.h b/src/netbase.h
@@ -68,7 +68,12 @@ std::unique_ptr<Sock> CreateSockTCP(const CService& address_family);
 extern std::function<std::unique_ptr<Sock>(const CService&)> CreateSock;
 
 bool ConnectSocketDirectly(const CService &addrConnect, const SOCKET& hSocketRet, int nTimeout, bool manual_connection);
-bool ConnectThroughProxy(const proxyType &proxy, const std::string& strDest, int port, const SOCKET& hSocketRet, int nTimeout, bool& outProxyConnectionFailed);
+bool ConnectThroughProxy(const proxyType& proxy,
+                         const std::string& strDest,
+                         int port,
+                         const Sock& sock,
+                         int nTimeout,
+                         bool& outProxyConnectionFailed);
 /** Disable or enable blocking-mode for a socket */
 bool SetSocketNonBlocking(const SOCKET& hSocket, bool fNonBlocking);
 /** Set the TCP_NODELAY flag on a socket */
