3 files changed, 35 insertions, 12 deletions

diff --git a/test/syscalls/linux/socket_test_util.cc b/test/syscalls/linux/socket_test_util.cc
index 53b678e94..317db5e0f 100644
--- a/test/syscalls/linux/socket_test_util.cc
+++ b/test/syscalls/linux/socket_test_util.cc
@@ -753,6 +753,19 @@ PosixErrorOr<int> SendMsg(int sock, msghdr* msg, char buf[], int buf_size) {
   return ret;
 }
 
+PosixErrorOr<int> RecvMsgTimeout(int sock, char buf[], int buf_size,
+                                 int timeout) {
+  fd_set rfd;
+  struct timeval to = {.tv_sec = timeout, .tv_usec = 0};
+  FD_ZERO(&rfd);
+  FD_SET(sock, &rfd);
+
+  int ret;
+  RETURN_ERROR_IF_SYSCALL_FAIL(ret = select(1, &rfd, NULL, NULL, &to));
+  RETURN_ERROR_IF_SYSCALL_FAIL(ret = recv(sock, buf, buf_size, MSG_DONTWAIT));
+  return ret;
+}
+
 void RecvNoData(int sock) {
   char data = 0;
   struct iovec iov;
diff --git a/test/syscalls/linux/socket_test_util.h b/test/syscalls/linux/socket_test_util.h
index 734b48b96..468bc96e0 100644
--- a/test/syscalls/linux/socket_test_util.h
+++ b/test/syscalls/linux/socket_test_util.h
@@ -467,6 +467,10 @@ PosixError FreeAvailablePort(int port);
 // SendMsg converts a buffer to an iovec and adds it to msg before sending it.
 PosixErrorOr<int> SendMsg(int sock, msghdr* msg, char buf[], int buf_size);
 
+// RecvMsgTimeout calls select on sock with timeout and then calls recv on sock.
+PosixErrorOr<int> RecvMsgTimeout(int sock, char buf[], int buf_size,
+                                 int timeout);
+
 // RecvNoData checks that no data is receivable on sock.
 void RecvNoData(int sock);
 
diff --git a/test/syscalls/linux/udp_socket.cc b/test/syscalls/linux/udp_socket.cc
index 97db2b321..1a7673317 100644
--- a/test/syscalls/linux/udp_socket.cc
+++ b/test/syscalls/linux/udp_socket.cc
@@ -14,6 +14,9 @@
 
 #include <arpa/inet.h>
 #include <fcntl.h>
+
+#include <ctime>
+
 #ifdef __linux__
 #include <linux/errqueue.h>
 #include <linux/filter.h>
@@ -834,8 +837,9 @@ TEST_P(UdpSocketTest, ReceiveBeforeConnect) {
 
   // Receive the data. It works because it was sent before the connect.
   char received[sizeof(buf)];
-  EXPECT_THAT(recv(bind_.get(), received, sizeof(received), 0),
-              SyscallSucceedsWithValue(sizeof(received)));
+  EXPECT_THAT(
+      RecvMsgTimeout(bind_.get(), received, sizeof(received), 1 /*timeout*/),
+      IsPosixErrorOkAndHolds(sizeof(received)));
   EXPECT_EQ(memcmp(buf, received, sizeof(buf)), 0);
 
   // Send again. This time it should not be received.
@@ -924,7 +928,9 @@ TEST_P(UdpSocketTest, ReadShutdownNonblockPendingData) {
               SyscallSucceedsWithValue(1));
 
   // We should get the data even though read has been shutdown.
-  EXPECT_THAT(recv(bind_.get(), received, 2, 0), SyscallSucceedsWithValue(2));
+  EXPECT_THAT(
+      RecvMsgTimeout(bind_.get(), received, 2 /*buf_size*/, 1 /*timeout*/),
+      IsPosixErrorOkAndHolds(2));
 
   // Because we read less than the entire packet length, since it's a packet
   // based socket any subsequent reads should return EWOULDBLOCK.
@@ -1692,9 +1698,9 @@ TEST_P(UdpSocketTest, RecvBufLimitsEmptyRcvBuf) {
         sendto(sock_.get(), buf.data(), buf.size(), 0, bind_addr_, addrlen_),
         SyscallSucceedsWithValue(buf.size()));
     std::vector<char> received(buf.size());
-    EXPECT_THAT(
-        recv(bind_.get(), received.data(), received.size(), MSG_DONTWAIT),
-        SyscallSucceedsWithValue(received.size()));
+    EXPECT_THAT(RecvMsgTimeout(bind_.get(), received.data(), received.size(),
+                               1 /*timeout*/),
+                IsPosixErrorOkAndHolds(received.size()));
   }
 
   {
@@ -1708,9 +1714,9 @@ TEST_P(UdpSocketTest, RecvBufLimitsEmptyRcvBuf) {
         SyscallSucceedsWithValue(buf.size()));
 
     std::vector<char> received(buf.size());
-    EXPECT_THAT(
-        recv(bind_.get(), received.data(), received.size(), MSG_DONTWAIT),
-        SyscallSucceedsWithValue(received.size()));
+    ASSERT_THAT(RecvMsgTimeout(bind_.get(), received.data(), received.size(),
+                               1 /*timeout*/),
+                IsPosixErrorOkAndHolds(received.size()));
   }
 }
 
@@ -1779,9 +1785,9 @@ TEST_P(UdpSocketTest, RecvBufLimits) {
     for (int i = 0; i < sent - 1; i++) {
       // Receive the data.
       std::vector<char> received(buf.size());
-      EXPECT_THAT(
-          recv(bind_.get(), received.data(), received.size(), MSG_DONTWAIT),
-          SyscallSucceedsWithValue(received.size()));
+      EXPECT_THAT(RecvMsgTimeout(bind_.get(), received.data(), received.size(),
+                                 1 /*timeout*/),
+                  IsPosixErrorOkAndHolds(received.size()));
       EXPECT_EQ(memcmp(buf.data(), received.data(), buf.size()), 0);
     }