[compiler-rt] [DFSan] Fix recvmsg wrapper to support MSG_TRUNC flag. (PR #92599)

[Andrew Browne via llvm-commits]

https://github.com/browneee created https://github.com/llvm/llvm-project/pull/92599

The MSG_TRUNC flag makes recvmsg return the real length of the packet, even if it was too big to fit in the provided buffer. This is commonly used together with MSG_PEEK.

Without this patch, dfsan's clear_msghdr_labels expects the return value of recvmsg (size recieved) to be less than or equal to the iov buffer length where recvmsg writes data, resulting in a crash.

>From 402dc26e81e9e18a2d00c84a532084f088d27cab Mon Sep 17 00:00:00 2001
From: Andrew Browne <browneee at google.com>
Date: Fri, 17 May 2024 19:03:40 +0000
Subject: [PATCH] [DFSan] Fix recvmsg wrapper to support MSG_TRUNC flag.

The MSG_TRUNC flag makes recvmsg return the real length of the packet,
even if it was too big to fit in the provided buffer.
This is commonly used together with MSG_PEEK.

Without this patch, dfsan's clear_msghdr_labels expects the return
value of recvmsg (size recieved) to be less than or equal to the
iov buffer length where recvmsg writes data, resulting in a crash.
---
 compiler-rt/lib/dfsan/dfsan_custom.cpp | 26 +++++++---
 compiler-rt/test/dfsan/custom.cpp      | 67 ++++++++++++++++++--------
 2 files changed, 65 insertions(+), 28 deletions(-)

diff --git a/compiler-rt/lib/dfsan/dfsan_custom.cpp b/compiler-rt/lib/dfsan/dfsan_custom.cpp
index 3af26e9f64c92..af3c1f4d1673c 100644
--- a/compiler-rt/lib/dfsan/dfsan_custom.cpp
+++ b/compiler-rt/lib/dfsan/dfsan_custom.cpp
@@ -1901,17 +1901,27 @@ SANITIZER_INTERFACE_ATTRIBUTE int __dfso_nanosleep(
   return __dfsw_nanosleep(req, rem, req_label, rem_label, ret_label);
 }
 
-static void clear_msghdr_labels(size_t bytes_written, struct msghdr *msg) {
+static void clear_msghdr_labels(size_t bytes_written, struct msghdr *msg,
+                                int flags) {
   dfsan_set_label(0, msg, sizeof(*msg));
   dfsan_set_label(0, msg->msg_name, msg->msg_namelen);
   dfsan_set_label(0, msg->msg_control, msg->msg_controllen);
-  for (size_t i = 0; bytes_written > 0; ++i) {
-    assert(i < msg->msg_iovlen);
+  for (size_t i = 0; i < msg->msg_iovlen; ++i) {
     struct iovec *iov = &msg->msg_iov[i];
-    size_t iov_written =
-        bytes_written < iov->iov_len ? bytes_written : iov->iov_len;
+    size_t iov_written = iov->iov_len;
+
+    // When MSG_TRUNC is not set, we want to avoid setting 0 label on bytes that
+    // may not have changed, using bytes_written to bound the 0 label write.
+    // When MSG_TRUNC flag is set, bytes_written may be larger than the buffer,
+    // and should not be used as a bound.
+    if (!(MSG_TRUNC & flags)) {
+      if (bytes_written < iov->iov_len) {
+        iov_written = bytes_written;
+      }
+      bytes_written -= iov_written;
+    }
+
     dfsan_set_label(0, iov->iov_base, iov_written);
-    bytes_written -= iov_written;
   }
 }
 
@@ -1923,7 +1933,7 @@ SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_recvmmsg(
   int ret = recvmmsg(sockfd, msgvec, vlen, flags, timeout);
   for (int i = 0; i < ret; ++i) {
     dfsan_set_label(0, &msgvec[i].msg_len, sizeof(msgvec[i].msg_len));
-    clear_msghdr_labels(msgvec[i].msg_len, &msgvec[i].msg_hdr);
+    clear_msghdr_labels(msgvec[i].msg_len, &msgvec[i].msg_hdr, flags);
   }
   *ret_label = 0;
   return ret;
@@ -1947,7 +1957,7 @@ SANITIZER_INTERFACE_ATTRIBUTE ssize_t __dfsw_recvmsg(
     dfsan_label msg_label, dfsan_label flags_label, dfsan_label *ret_label) {
   ssize_t ret = recvmsg(sockfd, msg, flags);
   if (ret >= 0)
-    clear_msghdr_labels(ret, msg);
+    clear_msghdr_labels(ret, msg, flags);
   *ret_label = 0;
   return ret;
 }
diff --git a/compiler-rt/test/dfsan/custom.cpp b/compiler-rt/test/dfsan/custom.cpp
index f544e481b7261..cede0d64dbcf2 100644
--- a/compiler-rt/test/dfsan/custom.cpp
+++ b/compiler-rt/test/dfsan/custom.cpp
@@ -768,26 +768,53 @@ void test_recvmsg() {
   ssize_t sent = sendmsg(sockfds[0], &smsg, 0);
   assert(sent > 0);
 
-  char rbuf[128];
-  struct iovec riovs[2] = {{&rbuf[0], 4}, {&rbuf[4], 4}};
-  struct msghdr rmsg = {};
-  rmsg.msg_iov = riovs;
-  rmsg.msg_iovlen = 2;
-
-  dfsan_set_label(i_label, rbuf, sizeof(rbuf));
-  dfsan_set_label(i_label, &rmsg, sizeof(rmsg));
-
-  DEFINE_AND_SAVE_ORIGINS(rmsg)
-
-  ssize_t received = recvmsg(sockfds[1], &rmsg, 0);
-  assert(received == sent);
-  assert(memcmp(sbuf, rbuf, 8) == 0);
-  ASSERT_ZERO_LABEL(received);
-  ASSERT_READ_ZERO_LABEL(&rmsg, sizeof(rmsg));
-  ASSERT_READ_ZERO_LABEL(&rbuf[0], 8);
-  ASSERT_READ_LABEL(&rbuf[8], 1, i_label);
-
-  ASSERT_SAVED_ORIGINS(rmsg)
+  {
+    char rpbuf[2];
+    struct iovec peek_iov;
+    peek_iov.iov_base = rpbuf;
+    peek_iov.iov_len = 2;
+
+    struct msghdr peek_header = {};
+    peek_header.msg_iov = &peek_iov;
+    peek_header.msg_iovlen = 1;
+
+    dfsan_set_label(i_label, rpbuf, sizeof(rpbuf));
+    dfsan_set_label(i_label, &peek_header, sizeof(peek_header));
+
+    DEFINE_AND_SAVE_ORIGINS(peek_header)
+
+    ssize_t received = recvmsg(sockfds[1], &peek_header, MSG_PEEK | MSG_TRUNC);
+    assert(received == sent);
+    assert(memcmp(sbuf, rpbuf, 2) == 0);
+    ASSERT_ZERO_LABEL(received);
+    ASSERT_READ_ZERO_LABEL(&peek_header, sizeof(peek_header));
+    ASSERT_READ_ZERO_LABEL(&rpbuf[0], 0);
+
+    ASSERT_SAVED_ORIGINS(peek_header)
+  }
+
+  {
+    char rbuf[128];
+    struct iovec riovs[2] = {{&rbuf[0], 4}, {&rbuf[4], 4}};
+    struct msghdr rmsg = {};
+    rmsg.msg_iov = riovs;
+    rmsg.msg_iovlen = 2;
+
+    dfsan_set_label(i_label, rbuf, sizeof(rbuf));
+    dfsan_set_label(i_label, &rmsg, sizeof(rmsg));
+
+    DEFINE_AND_SAVE_ORIGINS(rmsg)
+
+    ssize_t received = recvmsg(sockfds[1], &rmsg, 0);
+    assert(received == sent);
+    assert(memcmp(sbuf, rbuf, 8) == 0);
+    ASSERT_ZERO_LABEL(received);
+    ASSERT_READ_ZERO_LABEL(&rmsg, sizeof(rmsg));
+    ASSERT_READ_ZERO_LABEL(&rbuf[0], 8);
+    ASSERT_READ_LABEL(&rbuf[8], 1, i_label);
+
+    ASSERT_SAVED_ORIGINS(rmsg)
+  }
 
   close(sockfds[0]);
   close(sockfds[1]);