Add `sendSizeAndMsgWithTimeout()` to send size and data with a single call

diff --git a/pdns/iputils.cc b/pdns/iputils.cc
index e31c2d367374c4f57923b286be17ac468f5c7147..f06e08e0b4e258f579632dbc742371020d2a4e6b 100644 (file)
--- a/pdns/iputils.cc
+++ b/pdns/iputils.cc
@@ -300,3 +300,104 @@ ssize_t sendMsgWithTimeout(int fd, const char* buffer, size_t len, int timeout,
 
 template class NetmaskTree<bool>;
 
+bool sendSizeAndMsgWithTimeout(int sock, uint16_t bufferLen, const char* buffer, int idleTimeout, const ComboAddress* dest, const ComboAddress* local, unsigned int localItf, int totalTimeout, int flags)
+{
+  uint16_t size = htons(bufferLen);
+  char cbuf[256];
+  struct msghdr msgh;
+  struct iovec iov[2];
+  int remainingTime = totalTimeout;
+  time_t start = 0;
+  if (totalTimeout) {
+    start = time(NULL);
+  }
+
+  /* Set up iov and msgh structures. */
+  memset(&msgh, 0, sizeof(struct msghdr));
+  msgh.msg_control = nullptr;
+  msgh.msg_controllen = 0;
+  if (dest) {
+    msgh.msg_name = reinterpret_cast<void*>(const_cast<ComboAddress*>(dest));
+    msgh.msg_namelen = dest->getSocklen();
+  }
+  else {
+    msgh.msg_name = nullptr;
+    msgh.msg_namelen = 0;
+  }
+
+  msgh.msg_flags = 0;
+
+  if (localItf != 0 && local) {
+    addCMsgSrcAddr(&msgh, cbuf, local, localItf);
+  }
+
+  iov[0].iov_base = &size;
+  iov[0].iov_len = sizeof(size);
+  iov[1].iov_base = reinterpret_cast<void*>(const_cast<char*>(buffer));
+  iov[1].iov_len = bufferLen;
+
+  size_t pos = 0;
+  size_t sent = 0;
+  size_t nbElements = sizeof(iov)/sizeof(*iov);
+  while (true) {
+    msgh.msg_iov = &iov[pos];
+    msgh.msg_iovlen = nbElements - pos;
+
+    ssize_t res = sendmsg(sock, &msgh, flags);
+    if (res > 0) {
+      size_t written = static_cast<size_t>(res);
+      sent += written;
+
+      if (sent == (sizeof(size) + bufferLen)) {
+        return true;
+      }
+      /* partial write, we need to keep only the (parts of) elements
+         that have not been written.
+      */
+      do {
+        if (written < iov[pos].iov_len) {
+          iov[pos].iov_len -= written;
+          written = 0;
+        }
+        else {
+          written -= iov[pos].iov_len;
+          iov[pos].iov_len = 0;
+          pos++;
+        }
+      }
+      while (written > 0 && pos < nbElements);
+    }
+    else if (res == -1) {
+      if (errno == EINTR) {
+        continue;
+      }
+      else if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINPROGRESS) {
+        /* EINPROGRESS might happen with non blocking socket,
+           especially with TCP Fast Open */
+        int ret = waitForRWData(sock, false, (totalTimeout == 0 || idleTimeout <= remainingTime) ? idleTimeout : remainingTime, 0);
+        if (ret > 0) {
+          /* there is room available */
+        }
+        else if (ret == 0) {
+          throw runtime_error("Timeout while waiting to send data");
+        } else {
+          throw runtime_error("Error while waiting for room to send data");
+        }
+      }
+      else {
+        unixDie("failed in sendSizeAndMsgWithTimeout");
+      }
+    }
+    if (totalTimeout) {
+      time_t now = time(NULL);
+      int elapsed = now - start;
+      if (elapsed >= remainingTime) {
+        throw runtime_error("Timeout while sending data");
+      }
+      start = now;
+      remainingTime -= elapsed;
+    }
+  }
+
+  return false;
+}

diff --git a/pdns/iputils.hh b/pdns/iputils.hh
index 166dc6516fe9a68b320655833629512a4bf1f86d..4d3a2d4f253900f0e1b98848e5bf813df540907d 100644 (file)
--- a/pdns/iputils.hh
+++ b/pdns/iputils.hh
@@ -901,7 +901,8 @@ bool HarvestTimestamp(struct msghdr* msgh, struct timeval* tv);
 void fillMSGHdr(struct msghdr* msgh, struct iovec* iov, char* cbuf, size_t cbufsize, char* data, size_t datalen, ComboAddress* addr);
 ssize_t sendfromto(int sock, const char* data, size_t len, int flags, const ComboAddress& from, const ComboAddress& to);
 ssize_t sendMsgWithTimeout(int fd, const char* buffer, size_t len, int timeout, ComboAddress& dest, const ComboAddress& local, unsigned int localItf);
-
-#endif
+bool sendSizeAndMsgWithTimeout(int sock, uint16_t bufferLen, const char* buffer, int idleTimeout, const ComboAddress* dest, const ComboAddress* local, unsigned int localItf, int totalTimeout, int flags);
 
 extern template class NetmaskTree<bool>;
+
+#endif