[DFSan] Add custom wrapper for recvmsg.

The wrapper clears shadow for anything written by recvmsg.

Reviewed By: stephan.yichao.zhao

Differential Revision: https://reviews.llvm.org/D92949

compiler-rt/lib/dfsan/dfsan_custom.cpp

@@ -29,6 +29,7 @@
 #include <sys/epoll.h>
 #include <sys/resource.h>
 #include <sys/select.h>
+#include <sys/socket.h>
 #include <sys/stat.h>
 #include <sys/time.h>
 #include <sys/types.h>
@@ -879,6 +880,26 @@ SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_nanosleep(const struct timespec *req,
   return ret;
 }
 
+SANITIZER_INTERFACE_ATTRIBUTE ssize_t __dfsw_recvmsg(
+    int sockfd, struct msghdr *msg, int flags, dfsan_label sockfd_label,
+    dfsan_label msg_label, dfsan_label flags_label, dfsan_label *ret_label) {
+  ssize_t ret = recvmsg(sockfd, msg, flags);
+  if (ret >= 0) {
+    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 remaining = ret, i = 0; remaining > 0; ++i) {
+      assert(i < msg->msg_iovlen);
+      struct iovec *iov = &msg->msg_iov[i];
+      size_t written = remaining < iov->iov_len ? remaining : iov->iov_len;
+      dfsan_set_label(0, iov->iov_base, written);
+      remaining -= written;
+    }
+  }
+  *ret_label = 0;
+  return ret;
+}
+
 SANITIZER_INTERFACE_ATTRIBUTE int
 __dfsw_socketpair(int domain, int type, int protocol, int sv[2],
                   dfsan_label domain_label, dfsan_label type_label,

compiler-rt/lib/dfsan/done_abilist.txt

@@ -116,6 +116,8 @@ fun:connect=discard
 fun:creat=discard
 fun:dladdr=discard
 fun:dlclose=discard
+fun:epoll_create=discard
+fun:epoll_create1=discard
 fun:epoll_ctl=discard
 fun:fclose=discard
 fun:feof=discard
@@ -195,6 +197,7 @@ fun:getrusage=custom
 fun:nanosleep=custom
 fun:pread=custom
 fun:read=custom
+fun:recvmsg=custom
 fun:socketpair=custom
 fun:stat=custom
 fun:time=custom

compiler-rt/test/dfsan/custom.cpp

@@ -25,6 +25,7 @@
 #include <sys/epoll.h>
 #include <sys/resource.h>
 #include <sys/select.h>
+#include <sys/socket.h>
 #include <sys/stat.h>
 #include <sys/time.h>
 #include <sys/types.h>
@@ -336,6 +337,41 @@ void test_calloc() {
   free(crv);
 }
 
+void test_recvmsg() {
+  int sockfds[2];
+  int ret = socketpair(AF_UNIX, SOCK_DGRAM, 0, sockfds);
+  assert(ret != -1);
+
+  char sbuf[] = "abcdefghijkl";
+  struct iovec siovs[2] = {{&sbuf[0], 4}, {&sbuf[4], 4}};
+  struct msghdr smsg = {};
+  smsg.msg_iov = siovs;
+  smsg.msg_iovlen = 2;
+
+  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));
+
+  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);
+
+  close(sockfds[0]);
+  close(sockfds[1]);
+}
+
 void test_read() {
   char buf[16];
   dfsan_set_label(i_label, buf, 1);
@@ -1089,6 +1125,7 @@ int main(void) {
   test_pread();
   test_pthread_create();
   test_read();
+  test_recvmsg();
   test_sched_getaffinity();
   test_select();
   test_sigaction();