Rearrange extensions MojoDataPump::ReceiveMore and SendMore

Rearrange ReceiveMore and SendMore, so that we check MOJO_RESULT_SHOULD_WAIT
before !MOJO_RESULT_OK. This CL adds three unit tests that pass
with or without the change.

Bug: 721401
Change-Id: I6b132b2399f9091374d037d890d7ca98ba358445
Reviewed-on: https://chromium-review.googlesource.com/1150428Reviewed-by: Reilly Grant <reillyg@chromium.org>
Commit-Queue: Helen Li <xunjieli@chromium.org>
Cr-Commit-Position: refs/heads/master@{#578597}

extensions/browser/BUILD.gn

@@ -538,6 +538,7 @@ source_set("unit_tests") {
     "api/networking_private/networking_private_chromeos_unittest.cc",
     "api/power/power_api_unittest.cc",
     "api/runtime/restart_after_delay_api_unittest.cc",
+    "api/socket/mojo_data_pump_unittest.cc",
     "api/sockets_tcp/sockets_tcp_api_unittest.cc",
     "api/sockets_udp/sockets_udp_api_unittest.cc",
     "api/storage/settings_quota_unittest.cc",

extensions/browser/api/socket/mojo_data_pump.cc

@@ -69,28 +69,24 @@ void MojoDataPump::ReceiveMore(MojoResult result,
   DCHECK(read_callback_);
   DCHECK_NE(0u, read_size_);
 
-  if (result != MOJO_RESULT_OK) {
-    read_size_ = 0;
-    std::move(read_callback_).Run(net::ERR_FAILED, nullptr);
-    return;
-  }
-
   uint32_t num_bytes = read_size_;
-  auto io_buffer =
-      base::MakeRefCounted<net::IOBuffer>(static_cast<size_t>(num_bytes));
-  result = receive_stream_->ReadData(io_buffer->data(), &num_bytes,
-                                     MOJO_READ_DATA_FLAG_NONE);
-  if (result != MOJO_RESULT_OK) {
-    read_size_ = 0;
-    std::move(read_callback_).Run(net::ERR_FAILED, nullptr);
-    return;
-  }
 
+  scoped_refptr<net::IOBuffer> io_buffer;
+  if (result == MOJO_RESULT_OK) {
+    io_buffer =
+        base::MakeRefCounted<net::IOBuffer>(static_cast<size_t>(num_bytes));
+    result = receive_stream_->ReadData(io_buffer->data(), &num_bytes,
+                                       MOJO_READ_DATA_FLAG_NONE);
+  }
   if (result == MOJO_RESULT_SHOULD_WAIT) {
     receive_stream_watcher_.ArmOrNotify();
     return;
   }
   read_size_ = 0;
+  if (result != MOJO_RESULT_OK) {
+    std::move(read_callback_).Run(net::ERR_FAILED, nullptr);
+    return;
+  }
   std::move(read_callback_).Run(num_bytes, io_buffer);
 }
 
@@ -99,8 +95,10 @@ void MojoDataPump::SendMore(MojoResult result,
   DCHECK(write_callback_);
 
   uint32_t num_bytes = pending_write_buffer_size_;
-  result = send_stream_->WriteData(pending_write_buffer_->data(), &num_bytes,
-                                   MOJO_WRITE_DATA_FLAG_NONE);
+  if (result == MOJO_RESULT_OK) {
+    result = send_stream_->WriteData(pending_write_buffer_->data(), &num_bytes,
+                                     MOJO_WRITE_DATA_FLAG_NONE);
+  }
   if (result == MOJO_RESULT_SHOULD_WAIT) {
     send_stream_watcher_.ArmOrNotify();
     return;

extensions/browser/api/socket/mojo_data_pump_unittest.cc

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "extensions/browser/api/socket/mojo_data_pump.h"
+
+#include <memory>
+#include <utility>
+#include <string>
+
+#include "base/run_loop.h"
+#include "base/test/bind_test_util.h"
+#include "base/test/scoped_task_environment.h"
+#include "mojo/public/cpp/system/data_pipe.h"
+#include "net/base/test_completion_callback.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace extensions {
+
+// Tests that if |MojoDataPump::receive_stream_| is not ready, MojoDataPump will
+// wait and not error out.
+TEST(MojoDataPumpTest, ReceiveStreamNotReady) {
+  base::test::ScopedTaskEnvironment scoped_task_environment;
+  mojo::DataPipe receive_pipe;
+  mojo::DataPipe send_pipe;
+  auto pump =
+      std::make_unique<MojoDataPump>(std::move(receive_pipe.consumer_handle),
+                                     std::move(send_pipe.producer_handle));
+  std::string data("dummy");
+  base::RunLoop run_loop;
+  bool callback_called = false;
+  pump->Read(10 /*count*/,
+             base::BindLambdaForTesting(
+                 [&](int result, scoped_refptr<net::IOBuffer> io_buffer) {
+                   callback_called = true;
+                   ASSERT_EQ(static_cast<int>(data.size()), result);
+                   EXPECT_EQ(data,
+                             std::string(io_buffer->data(), result));
+                   run_loop.Quit();
+                 }));
+
+  // Spin the message loop so that MojoDataPump::ReceiveMore() is called but the
+  // callback will not be executed yet because there is no data to read.
+  base::RunLoop().RunUntilIdle();
+  EXPECT_FALSE(callback_called);
+
+  mojo::ScopedDataPipeProducerHandle receive_pipe_producer =
+      std::move(receive_pipe.producer_handle);
+
+  uint32_t num_bytes = data.size();
+  // WriteData() completes synchronously because |data| is much smaller than
+  // data pipe's internal buffer.
+  MojoResult r = receive_pipe_producer->WriteData(data.data(), &num_bytes,
+                                                  MOJO_WRITE_DATA_FLAG_NONE);
+  ASSERT_EQ(MOJO_RESULT_OK, r);
+  ASSERT_EQ(data.size(), num_bytes);
+
+  // Now pump->Read() should complete.
+  run_loop.Run();
+  EXPECT_TRUE(callback_called);
+}
+
+// Tests that if |MojoDataPump::receive_stream_| is closed, an error is
+// propagated.
+TEST(MojoDataPumpTest, ReceiveStreamClosed) {
+  base::test::ScopedTaskEnvironment scoped_task_environment;
+  mojo::DataPipe receive_pipe;
+  mojo::DataPipe send_pipe;
+  auto pump =
+      std::make_unique<MojoDataPump>(std::move(receive_pipe.consumer_handle),
+                                     std::move(send_pipe.producer_handle));
+  base::RunLoop run_loop;
+  pump->Read(10 /*count*/,
+             base::BindLambdaForTesting(
+                 [&](int result, scoped_refptr<net::IOBuffer> io_buffer) {
+                   EXPECT_EQ(net::ERR_FAILED, result);
+                   EXPECT_EQ(nullptr, io_buffer);
+                   run_loop.Quit();
+                 }));
+
+  mojo::ScopedDataPipeProducerHandle receive_pipe_producer =
+      std::move(receive_pipe.producer_handle);
+  receive_pipe_producer.reset();
+
+  // Now pump->Read() should complete.
+  run_loop.Run();
+}
+
+// Tests that if |MojoDataPump::send_stream_| is closed, Write() will fail.
+TEST(MojoDataPumpTest, SendStreamClosed) {
+  base::test::ScopedTaskEnvironment scoped_task_environment;
+  mojo::DataPipe receive_pipe;
+  mojo::DataPipe send_pipe;
+  auto pump =
+      std::make_unique<MojoDataPump>(std::move(receive_pipe.consumer_handle),
+                                     std::move(send_pipe.producer_handle));
+  scoped_refptr<net::StringIOBuffer> write_buffer =
+      base::MakeRefCounted<net::StringIOBuffer>("dummy");
+  net::TestCompletionCallback callback;
+  mojo::ScopedDataPipeConsumerHandle send_pipe_consumer =
+      std::move(send_pipe.consumer_handle);
+  send_pipe_consumer.reset();
+  pump->Write(write_buffer.get(), write_buffer->size(), callback.callback());
+  EXPECT_EQ(net::ERR_FAILED, callback.WaitForResult());
+}
+
+}  // namespace extensions