From f1420cf48418c01694eaf3110ac411915b217d36 Mon Sep 17 00:00:00 2001
From: Ting-Yu Wang <anivia@google.com>
Date: Fri, 15 Jan 2021 12:48:58 -0800
Subject: Add sanity check on return values from Write

io.Writer.Write requires err to be non-nil if n < len(v).

We could allow this but it will be irreversible if users depend on this
behavior.

Ported the test that discovered this.

PiperOrigin-RevId: 352065946
---
 test/syscalls/linux/BUILD                    |  2 +
 test/syscalls/linux/socket_generic_stress.cc | 71 ++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

(limited to 'test/syscalls/linux')

diff --git a/test/syscalls/linux/BUILD b/test/syscalls/linux/BUILD
index 017f997de..d184712e3 100644
--- a/test/syscalls/linux/BUILD
+++ b/test/syscalls/linux/BUILD
@@ -2304,9 +2304,11 @@ cc_binary(
     deps = [
         ":ip_socket_test_util",
         ":socket_test_util",
+        "@com_google_absl//absl/strings",
         gtest,
         "//test/util:test_main",
         "//test/util:test_util",
+        "//test/util:thread_util",
     ],
 )
 
diff --git a/test/syscalls/linux/socket_generic_stress.cc b/test/syscalls/linux/socket_generic_stress.cc
index 6cd67123d..679586530 100644
--- a/test/syscalls/linux/socket_generic_stress.cc
+++ b/test/syscalls/linux/socket_generic_stress.cc
@@ -18,10 +18,15 @@
 #include <sys/socket.h>
 #include <sys/un.h>
 
+#include <array>
+#include <string>
+
 #include "gtest/gtest.h"
+#include "absl/strings/string_view.h"
 #include "test/syscalls/linux/ip_socket_test_util.h"
 #include "test/syscalls/linux/socket_test_util.h"
 #include "test/util/test_util.h"
+#include "test/util/thread_util.h"
 
 namespace gvisor {
 namespace testing {
@@ -138,5 +143,71 @@ INSTANTIATE_TEST_SUITE_P(
         SetSockOpt(SOL_SOCKET, SO_REUSEADDR, &kSockOptOn)(
             DualStackTCPAcceptBindPersistentListenerSocketPair(0))));
 
+using DataTransferStressTest = SocketPairTest;
+
+TEST_P(DataTransferStressTest, BigDataTransfer) {
+  // TODO(b/165912341): These are too slow on KVM platform with nested virt.
+  SKIP_IF(GvisorPlatform() == Platform::kKVM);
+
+  auto sockets = ASSERT_NO_ERRNO_AND_VALUE(NewSocketPair());
+  int client_fd = sockets->first_fd();
+  int server_fd = sockets->second_fd();
+
+  ScopedThread echo([server_fd]() {
+    std::array<uint8_t, 1024> buf;
+    for (;;) {
+      ssize_t r = read(server_fd, buf.data(), buf.size());
+      ASSERT_THAT(r, SyscallSucceeds());
+      if (r == 0) {
+        break;
+      }
+      for (size_t i = 0; i < r;) {
+        ssize_t w = write(server_fd, buf.data() + i, r - i);
+        ASSERT_GE(w, 0);
+        i += w;
+      }
+    }
+    ASSERT_THAT(shutdown(server_fd, SHUT_WR), SyscallSucceeds());
+  });
+
+  const std::string chunk = "Though this upload be but little, it is fierce.";
+  std::string big_string;
+  while (big_string.size() < 31 << 20) {
+    big_string += chunk;
+  }
+  absl::string_view data = big_string;
+
+  ScopedThread writer([client_fd, data]() {
+    absl::string_view view = data;
+    while (!view.empty()) {
+      ssize_t n = write(client_fd, view.data(), view.size());
+      ASSERT_GE(n, 0);
+      view = view.substr(n);
+    }
+    ASSERT_THAT(shutdown(client_fd, SHUT_WR), SyscallSucceeds());
+  });
+
+  std::string buf;
+  buf.resize(1 << 20);
+  while (!data.empty()) {
+    ssize_t n = read(client_fd, buf.data(), buf.size());
+    ASSERT_GE(n, 0);
+    for (size_t i = 0; i < n; i += chunk.size()) {
+      size_t c = std::min(chunk.size(), n - i);
+      ASSERT_EQ(buf.substr(i, c), data.substr(i, c)) << "offset " << i;
+    }
+    data = data.substr(n);
+  }
+  // Should read EOF now.
+  ASSERT_THAT(read(client_fd, buf.data(), buf.size()),
+              SyscallSucceedsWithValue(0));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    AllConnectedSockets, DataTransferStressTest,
+    ::testing::Values(IPv6TCPAcceptBindPersistentListenerSocketPair(0),
+                      IPv4TCPAcceptBindPersistentListenerSocketPair(0),
+                      DualStackTCPAcceptBindPersistentListenerSocketPair(0)));
+
 }  // namespace testing
 }  // namespace gvisor
-- 
cgit v1.2.3