dnsdist: Properly handle short reads on backend upgrade discovery

diff --git a/pdns/dnsdistdist/dnsdist-discovery.cc b/pdns/dnsdistdist/dnsdist-discovery.cc
index 66ac1c070a655059a725419b8a8d0c7d661e8436..de76a42ff17f83c1556c0c976d545d5047769a51 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-discovery.cc
+++ b/pdns/dnsdistdist/dnsdist-discovery.cc
@@ -272,8 +272,9 @@ bool ServiceDiscovery::getDiscoveredConfig(const UpgradeableBackend& upgradeable
 
     sock.writenWithTimeout(reinterpret_cast<const char*>(packet.data()), packet.size(), backend->d_config.tcpSendTimeout);
 
+    const struct timeval remainingTime = { .tv_sec = backend->d_config.tcpRecvTimeout, .tv_usec = 0 };
     uint16_t responseSize = 0;
-    auto got = sock.readWithTimeout(reinterpret_cast<char*>(&responseSize), sizeof(responseSize), backend->d_config.tcpRecvTimeout);
+    auto got = readn2WithTimeout(sock.getHandle(), &responseSize, sizeof(responseSize), remainingTime);
     if (got < 0 || static_cast<size_t>(got) != sizeof(responseSize)) {
       if (g_verbose) {
         warnlog("Error while waiting for the ADD upgrade response size from backend %s: %d", addr.toStringWithPort(), got);
@@ -283,7 +284,7 @@ bool ServiceDiscovery::getDiscoveredConfig(const UpgradeableBackend& upgradeable
 
     packet.resize(ntohs(responseSize));
 
-    got = sock.readWithTimeout(reinterpret_cast<char*>(packet.data()), packet.size(), backend->d_config.tcpRecvTimeout);
+    got = readn2WithTimeout(sock.getHandle(), packet.data(), packet.size(), remainingTime);
     if (got < 0 || static_cast<size_t>(got) != packet.size()) {
       if (g_verbose) {
         warnlog("Error while waiting for the ADD upgrade response from backend %s: %d", addr.toStringWithPort(), got);