Fix WAITALL and RCVTIMEO interaction

PiperOrigin-RevId: 225424296
Change-Id: I60fcc2b859339dca9963cb32227a287e719ab765

4 files changed, 39 insertions, 2 deletions

diff --git a/pkg/sentry/socket/epsocket/epsocket.go b/pkg/sentry/socket/epsocket/epsocket.go
index b49ef21ad..19af7bc45 100644
--- a/pkg/sentry/socket/epsocket/epsocket.go
+++ b/pkg/sentry/socket/epsocket/epsocket.go
@@ -1352,6 +1352,9 @@ func (s *SocketOperations) RecvMsg(t *kernel.Task, dst usermem.IOSequence, flags
 		dst = dst.DropFirst(rn)
 
 		if err := t.BlockWithDeadline(ch, haveDeadline, deadline); err != nil {
+			if n > 0 {
+				return n, senderAddr, senderAddrLen, controlMessages, nil
+			}
 			if err == syserror.ETIMEDOUT {
 				return 0, nil, 0, socket.ControlMessages{}, syserr.ErrTryAgain
 			}
diff --git a/pkg/sentry/socket/socket.go b/pkg/sentry/socket/socket.go
index f31729819..f73127ea6 100644
--- a/pkg/sentry/socket/socket.go
+++ b/pkg/sentry/socket/socket.go
@@ -86,6 +86,8 @@ type Socket interface {
 	//
 	// senderAddrLen is the address length to be returned to the application,
 	// not necessarily the actual length of the address.
+	//
+	// If err != nil, the recv was not successful.
 	RecvMsg(t *kernel.Task, dst usermem.IOSequence, flags int, haveDeadline bool, deadline ktime.Time, senderRequested bool, controlDataLen uint64) (n int, senderAddr interface{}, senderAddrLen uint32, controlMessages ControlMessages, err *syserr.Error)
 
 	// SendMsg implements the sendmsg(2) linux syscall. SendMsg does not take
diff --git a/pkg/sentry/socket/unix/unix.go b/pkg/sentry/socket/unix/unix.go
index 11cad411d..4c9dcbd61 100644
--- a/pkg/sentry/socket/unix/unix.go
+++ b/pkg/sentry/socket/unix/unix.go
@@ -538,6 +538,9 @@ func (s *SocketOperations) RecvMsg(t *kernel.Task, dst usermem.IOSequence, flags
 			}
 			total += n
 			if err != nil || !waitAll || s.isPacket || n >= dst.NumBytes() {
+				if total > 0 {
+					err = nil
+				}
 				return int(total), from, fromLen, socket.ControlMessages{Unix: r.Control}, syserr.FromError(err)
 			}
 
@@ -546,10 +549,13 @@ func (s *SocketOperations) RecvMsg(t *kernel.Task, dst usermem.IOSequence, flags
 		}
 
 		if err := t.BlockWithDeadline(ch, haveDeadline, deadline); err != nil {
+			if total > 0 {
+				err = nil
+			}
 			if err == syserror.ETIMEDOUT {
-				return 0, nil, 0, socket.ControlMessages{}, syserr.ErrTryAgain
+				return int(total), nil, 0, socket.ControlMessages{}, syserr.ErrTryAgain
 			}
-			return 0, nil, 0, socket.ControlMessages{}, syserr.FromError(err)
+			return int(total), nil, 0, socket.ControlMessages{}, syserr.FromError(err)
 		}
 	}
 }
diff --git a/test/syscalls/linux/socket_generic.cc b/test/syscalls/linux/socket_generic.cc
index fdc346d4d..a9edbb950 100644
--- a/test/syscalls/linux/socket_generic.cc
+++ b/test/syscalls/linux/socket_generic.cc
@@ -395,6 +395,8 @@ TEST_P(AllSocketPairTest, RecvWaitAll) {
   ASSERT_THAT(RetryEINTR(recv)(sockets->second_fd(), received_data,
                                sizeof(received_data), MSG_WAITALL),
               SyscallSucceedsWithValue(sizeof(sent_data)));
+
+  EXPECT_EQ(0, memcmp(sent_data, received_data, sizeof(sent_data)));
 }
 
 TEST_P(AllSocketPairTest, RecvWaitAllDontWait) {
@@ -406,5 +408,29 @@ TEST_P(AllSocketPairTest, RecvWaitAllDontWait) {
               SyscallFailsWithErrno(EAGAIN));
 }
 
+TEST_P(AllSocketPairTest, RecvTimeoutWaitAll) {
+  auto sockets = ASSERT_NO_ERRNO_AND_VALUE(NewSocketPair());
+
+  struct timeval tv {
+    .tv_sec = 0, .tv_usec = 200000  // 200ms
+  };
+  EXPECT_THAT(setsockopt(sockets->second_fd(), SOL_SOCKET, SO_RCVTIMEO, &tv,
+                         sizeof(tv)),
+              SyscallSucceeds());
+
+  char sent_data[100];
+  RandomizeBuffer(sent_data, sizeof(sent_data));
+
+  ASSERT_THAT(write(sockets->first_fd(), sent_data, sizeof(sent_data)),
+              SyscallSucceedsWithValue(sizeof(sent_data)));
+
+  char received_data[sizeof(sent_data) * 2] = {};
+  ASSERT_THAT(RetryEINTR(recv)(sockets->second_fd(), received_data,
+                               sizeof(received_data), MSG_WAITALL),
+              SyscallSucceedsWithValue(sizeof(sent_data)));
+
+  EXPECT_EQ(0, memcmp(sent_data, received_data, sizeof(sent_data)));
+}
+
 }  // namespace testing
 }  // namespace gvisor