Update from https://crrev.com/312600

mojo/services/network/udp_socket_apptest.cc

Index: mojo/services/network/udp_socket_apptest.cc
diff --git a/mojo/services/network/udp_socket_apptest.cc b/mojo/services/network/udp_socket_apptest.cc
index 079ed8fd40c3e5ba72dc8e2cddebde88c55de5ac..86f3387a6232f30fda270f1b17020a5d36a42773 100644
--- a/mojo/services/network/udp_socket_apptest.cc
+++ b/mojo/services/network/udp_socket_apptest.cc
@@ -222,82 +222,82 @@ class TestReceiveCallback
   Array<uint8_t> data_;
 };
 
-class UDPSocketAppTest : public test::ApplicationTestBase {
- public:
-  UDPSocketAppTest() {}
-  ~UDPSocketAppTest() override {}
-
-  void SetUp() override {
-    ApplicationTestBase::SetUp();
+struct ReceiveResult {
+  NetworkErrorPtr result;
+  NetAddressPtr addr;
+  Array<uint8_t> data;
+};
 
-    ApplicationConnection* connection =
-        application_impl()->ConnectToApplication("mojo:network_service");
-    connection->ConnectToService(&network_service_);
+class UDPSocketClientImpl : public UDPSocketClient {
+ public:
+  UDPSocketClientImpl() : run_loop_(nullptr), expected_receive_count_(0) {}
 
-    network_service_->CreateUDPSocket(GetProxy(&udp_socket_));
-    udp_socket_.set_client(&udp_socket_client_);
+  ~UDPSocketClientImpl() override {
+    while (!results_.empty()) {
+      delete results_.front();
+      results_.pop();
+    }
   }
 
- protected:
-  struct ReceiveResult {
-    NetworkErrorPtr result;
-    NetAddressPtr addr;
-    Array<uint8_t> data;
-  };
+  std::queue<ReceiveResult*>* results() {
+    return &results_;
+  }
 
-  class UDPSocketClientImpl : public UDPSocketClient {
-   public:
+  void WaitForReceiveResults(size_t count) {
+    if (results_.size() == count)
+      return;
 
-    UDPSocketClientImpl() : run_loop_(nullptr), expected_receive_count_(0) {}
+    expected_receive_count_ = count;
+    base::RunLoop run_loop;
+    run_loop_ = &run_loop;
+    run_loop.Run();
+    run_loop_ = nullptr;
+  }
 
-    ~UDPSocketClientImpl() override {
-      while (!results_.empty()) {
-        delete results_.front();
-        results_.pop();
-      }
+ private:
+  void OnReceived(NetworkErrorPtr result,
+                  NetAddressPtr src_addr,
+                  Array<uint8_t> data) override {
+    ReceiveResult* entry = new ReceiveResult();
+    entry->result = result.Pass();
+    entry->addr = src_addr.Pass();
+    entry->data = data.Pass();
+
+    results_.push(entry);
+
+    if (results_.size() == expected_receive_count_ && run_loop_) {
+      expected_receive_count_ = 0;
+      run_loop_->Quit();
     }
+  }
 
-    void OnReceived(NetworkErrorPtr result,
-                    NetAddressPtr src_addr,
-                    Array<uint8_t> data) override {
-      ReceiveResult* entry = new ReceiveResult();
-      entry->result = result.Pass();
-      entry->addr = src_addr.Pass();
-      entry->data = data.Pass();
-
-      results_.push(entry);
-
-      if (results_.size() == expected_receive_count_ && run_loop_) {
-        expected_receive_count_ = 0;
-        run_loop_->Quit();
-      }
-    }
+  base::RunLoop* run_loop_;
+  std::queue<ReceiveResult*> results_;
+  size_t expected_receive_count_;
 
-    base::RunLoop* run_loop_;
-    std::queue<ReceiveResult*> results_;
-    size_t expected_receive_count_;
+  DISALLOW_COPY_AND_ASSIGN(UDPSocketClientImpl);
+};
 
-    DISALLOW_COPY_AND_ASSIGN(UDPSocketClientImpl);
-  };
+class UDPSocketAppTest : public test::ApplicationTestBase {
+ public:
+  UDPSocketAppTest() {}
+  ~UDPSocketAppTest() override {}
 
-  std::queue<ReceiveResult*>* GetReceiveResults() {
-    return &udp_socket_client_.results_;
-  }
+  void SetUp() override {
+    ApplicationTestBase::SetUp();
 
-  void WaitForReceiveResults(size_t count) {
-    if (GetReceiveResults()->size() == count)
-      return;
+    ApplicationConnection* connection =
+        application_impl()->ConnectToApplication("mojo:network_service");
+    connection->ConnectToService(&network_service_);
 
-    udp_socket_client_.expected_receive_count_ = count;
-    base::RunLoop run_loop;
-    udp_socket_client_.run_loop_ = &run_loop;
-    run_loop.Run();
-    udp_socket_client_.run_loop_ = nullptr;
+    network_service_->CreateUDPSocket(GetProxy(&socket_));
+    socket_.set_client(&receiver_);
   }
 
+ protected:
   NetworkServicePtr network_service_;
-  UDPSocketPtr udp_socket_;
-  UDPSocketClientImpl udp_socket_client_;
+  UDPSocketPtr socket_;
+  UDPSocketClientImpl receiver_;
 
   DISALLOW_COPY_AND_ASSIGN(UDPSocketAppTest);
 };
@@ -306,58 +306,58 @@ class UDPSocketAppTest : public test::ApplicationTestBase {
 
 TEST_F(UDPSocketAppTest, Settings) {
   TestCallback callback1;
-  udp_socket_->AllowAddressReuse(callback1.callback());
+  socket_->AllowAddressReuse(callback1.callback());
   callback1.WaitForResult();
   EXPECT_EQ(net::OK, callback1.result()->code);
 
   // Should fail because the socket hasn't been bound.
   TestCallback callback2;
-  udp_socket_->SetSendBufferSize(1024, callback2.callback());
+  socket_->SetSendBufferSize(1024, callback2.callback());
   callback2.WaitForResult();
   EXPECT_NE(net::OK, callback2.result()->code);
 
   // Should fail because the socket hasn't been bound.
   TestCallback callback3;
-  udp_socket_->SetReceiveBufferSize(2048, callback3.callback());
+  socket_->SetReceiveBufferSize(2048, callback3.callback());
   callback3.WaitForResult();
   EXPECT_NE(net::OK, callback3.result()->code);
 
   TestCallbackWithAddress callback4;
-  udp_socket_->Bind(GetLocalHostWithAnyPort(), callback4.callback());
+  socket_->Bind(GetLocalHostWithAnyPort(), callback4.callback());
   callback4.WaitForResult();
   EXPECT_EQ(net::OK, callback4.result()->code);
   EXPECT_NE(0u, callback4.net_address()->ipv4->port);
 
   // Should fail because the socket has been bound.
   TestCallback callback5;
-  udp_socket_->AllowAddressReuse(callback5.callback());
+  socket_->AllowAddressReuse(callback5.callback());
   callback5.WaitForResult();
   EXPECT_NE(net::OK, callback5.result()->code);
 
   TestCallback callback6;
-  udp_socket_->SetSendBufferSize(1024, callback6.callback());
+  socket_->SetSendBufferSize(1024, callback6.callback());
   callback6.WaitForResult();
   EXPECT_EQ(net::OK, callback6.result()->code);
 
   TestCallback callback7;
-  udp_socket_->SetReceiveBufferSize(2048, callback7.callback());
+  socket_->SetReceiveBufferSize(2048, callback7.callback());
   callback7.WaitForResult();
   EXPECT_EQ(net::OK, callback7.result()->code);
 
   TestCallbackWithUint32 callback8;
-  udp_socket_->NegotiateMaxPendingSendRequests(0, callback8.callback());
+  socket_->NegotiateMaxPendingSendRequests(0, callback8.callback());
   callback8.WaitForResult();
   EXPECT_GT(callback8.result(), 0u);
 
   TestCallbackWithUint32 callback9;
-  udp_socket_->NegotiateMaxPendingSendRequests(16, callback9.callback());
+  socket_->NegotiateMaxPendingSendRequests(16, callback9.callback());
   callback9.WaitForResult();
   EXPECT_GT(callback9.result(), 0u);
 }
 
 TEST_F(UDPSocketAppTest, TestReadWrite) {
   TestCallbackWithAddress callback1;
-  udp_socket_->Bind(GetLocalHostWithAnyPort(), callback1.callback());
+  socket_->Bind(GetLocalHostWithAnyPort(), callback1.callback());
   callback1.WaitForResult();
   ASSERT_EQ(net::OK, callback1.result()->code);
   ASSERT_NE(0u, callback1.net_address()->ipv4->port);
@@ -377,7 +377,7 @@ TEST_F(UDPSocketAppTest, TestReadWrite) {
 
   const size_t kDatagramCount = 6;
   const size_t kDatagramSize = 255;
-  udp_socket_->ReceiveMore(kDatagramCount);
+  socket_->ReceiveMore(kDatagramCount);
 
   for (size_t i = 0; i < kDatagramCount; ++i) {
     TestCallback callback;
@@ -389,23 +389,97 @@ TEST_F(UDPSocketAppTest, TestReadWrite) {
     EXPECT_EQ(255, callback.result()->code);
   }
 
-  WaitForReceiveResults(kDatagramCount);
+  receiver_.WaitForReceiveResults(kDatagramCount);
+  for (size_t i = 0; i < kDatagramCount; ++i) {
+    scoped_ptr<ReceiveResult> result(receiver_.results()->front());
+    receiver_.results()->pop();
+
+    EXPECT_EQ(static_cast<int>(kDatagramSize), result->result->code);
+    EXPECT_TRUE(result->addr.Equals(client_addr));
+    EXPECT_TRUE(result->data.Equals(
+        CreateTestMessage(static_cast<uint8_t>(i), kDatagramSize)));
+  }
+}
+
+TEST_F(UDPSocketAppTest, TestConnectedReadWrite) {
+  TestCallbackWithAddress callback1;
+  socket_->Bind(GetLocalHostWithAnyPort(), callback1.callback());
+  callback1.WaitForResult();
+  ASSERT_EQ(net::OK, callback1.result()->code);
+  ASSERT_NE(0u, callback1.net_address()->ipv4->port);
+
+  NetAddressPtr server_addr = callback1.net_address().Clone();
+
+  UDPSocketPtr client_socket;
+  network_service_->CreateUDPSocket(GetProxy(&client_socket));
+  UDPSocketClientImpl client_socket_receiver;
+  client_socket.set_client(&client_socket_receiver);
+
+  TestCallbackWithAddress callback2;
+  client_socket->Connect(server_addr.Clone(), callback2.callback());
+  callback2.WaitForResult();
+  ASSERT_EQ(net::OK, callback2.result()->code);
+  ASSERT_NE(0u, callback2.net_address()->ipv4->port);
+
+  NetAddressPtr client_addr = callback2.net_address().Clone();
+
+  const size_t kDatagramCount = 6;
+  const size_t kDatagramSize = 255;
+
+  // Test send using a connected socket.
+  socket_->ReceiveMore(kDatagramCount);
+
+  for (size_t i = 0; i < kDatagramCount; ++i) {
+    TestCallback callback;
+    client_socket->SendTo(
+        nullptr,
+        CreateTestMessage(static_cast<uint8_t>(i), kDatagramSize),
+        callback.callback());
+    callback.WaitForResult();
+    EXPECT_EQ(255, callback.result()->code);
+  }
+
+  receiver_.WaitForReceiveResults(kDatagramCount);
   for (size_t i = 0; i < kDatagramCount; ++i) {
-    scoped_ptr<ReceiveResult> result(GetReceiveResults()->front());
-    GetReceiveResults()->pop();
+    scoped_ptr<ReceiveResult> result(receiver_.results()->front());
+    receiver_.results()->pop();
 
     EXPECT_EQ(static_cast<int>(kDatagramSize), result->result->code);
     EXPECT_TRUE(result->addr.Equals(client_addr));
     EXPECT_TRUE(result->data.Equals(
         CreateTestMessage(static_cast<uint8_t>(i), kDatagramSize)));
   }
+
+  // Test receive using a connected socket.
+  client_socket->ReceiveMore(kDatagramCount);
+
+  for (size_t i = 0; i < kDatagramCount; ++i) {
+    TestCallback callback;
+    socket_->SendTo(
+        client_addr.Clone(),
+        CreateTestMessage(static_cast<uint8_t>(i), kDatagramSize),
+        callback.callback());
+    callback.WaitForResult();
+    EXPECT_EQ(255, callback.result()->code);
+  }
+
+  client_socket_receiver.WaitForReceiveResults(kDatagramCount);
+  for (size_t i = 0; i < kDatagramCount; ++i) {
+    scoped_ptr<ReceiveResult> result(client_socket_receiver.results()->front());
+    client_socket_receiver.results()->pop();
+
+    EXPECT_EQ(static_cast<int>(kDatagramSize), result->result->code);
+    EXPECT_FALSE(result->addr);
+    EXPECT_TRUE(result->data.Equals(
+        CreateTestMessage(static_cast<uint8_t>(i), kDatagramSize)));
+  }
 }
 
-TEST_F(UDPSocketAppTest, TestUDPSocketWrapper) {
-  UDPSocketWrapper udp_socket(udp_socket_.Pass(), 4, 4);
+TEST_F(UDPSocketAppTest, TestWrapperReadWrite) {
+  UDPSocketWrapper socket(socket_.Pass(), 4, 4);
 
   TestCallbackWithAddress callback1;
-  udp_socket.Bind(GetLocalHostWithAnyPort(), callback1.callback());
+  socket.Bind(GetLocalHostWithAnyPort(), callback1.callback());
   callback1.WaitForResult();
   ASSERT_EQ(net::OK, callback1.result()->code);
   ASSERT_NE(0u, callback1.net_address()->ipv4->port);
@@ -438,7 +512,7 @@ TEST_F(UDPSocketAppTest, TestUDPSocketWrapper) {
           CreateTestMessage(static_cast<uint8_t>(j), kDatagramSize),
           send_callbacks[j].callback());
 
-      udp_socket.ReceiveFrom(receive_callbacks[j].callback());
+      socket.ReceiveFrom(receive_callbacks[j].callback());
     }
 
     receive_callbacks[i - 1].WaitForResult();
@@ -453,5 +527,84 @@ TEST_F(UDPSocketAppTest, TestUDPSocketWrapper) {
   }
 }
 
+TEST_F(UDPSocketAppTest, TestWrapperConnectedReadWrite) {
+  UDPSocketWrapper socket(socket_.Pass(), 4, 4);
+
+  TestCallbackWithAddress callback1;
+  socket.Bind(GetLocalHostWithAnyPort(), callback1.callback());
+  callback1.WaitForResult();
+  ASSERT_EQ(net::OK, callback1.result()->code);
+  ASSERT_NE(0u, callback1.net_address()->ipv4->port);
+
+  NetAddressPtr server_addr = callback1.net_address().Clone();
+
+  UDPSocketPtr raw_client_socket;
+  network_service_->CreateUDPSocket(GetProxy(&raw_client_socket));
+  UDPSocketWrapper client_socket(raw_client_socket.Pass(), 4, 4);
+
+  TestCallbackWithAddress callback2;
+  client_socket.Connect(server_addr.Pass(), callback2.callback());
+  callback2.WaitForResult();
+  ASSERT_EQ(net::OK, callback2.result()->code);
+  ASSERT_NE(0u, callback2.net_address()->ipv4->port);
+
+  NetAddressPtr client_addr = callback2.net_address().Clone();
+
+  const size_t kDatagramCount = 16;
+  const size_t kDatagramSize = 255;
+
+  // Test send using a connected socket.
+  for (size_t i = 1; i < kDatagramCount; ++i) {
+    scoped_ptr<TestCallback[]> send_callbacks(new TestCallback[i]);
+    scoped_ptr<TestReceiveCallback[]> receive_callbacks(
+        new TestReceiveCallback[i]);
+
+    for (size_t j = 0; j < i; ++j) {
+      client_socket.SendTo(
+          nullptr,
+          CreateTestMessage(static_cast<uint8_t>(j), kDatagramSize),
+          send_callbacks[j].callback());
+
+      socket.ReceiveFrom(receive_callbacks[j].callback());
+    }
+
+    receive_callbacks[i - 1].WaitForResult();
+
+    for (size_t j = 0; j < i; ++j) {
+      EXPECT_EQ(static_cast<int>(kDatagramSize),
+                receive_callbacks[j].result()->code);
+      EXPECT_TRUE(receive_callbacks[j].src_addr().Equals(client_addr));
+      EXPECT_TRUE(receive_callbacks[j].data().Equals(
+          CreateTestMessage(static_cast<uint8_t>(j), kDatagramSize)));
+    }
+  }
+
+  // Test receive using a connected socket.
+  for (size_t i = 1; i < kDatagramCount; ++i) {
+    scoped_ptr<TestCallback[]> send_callbacks(new TestCallback[i]);
+    scoped_ptr<TestReceiveCallback[]> receive_callbacks(
+        new TestReceiveCallback[i]);
+
+    for (size_t j = 0; j < i; ++j) {
+      socket.SendTo(
+          client_addr.Clone(),
+          CreateTestMessage(static_cast<uint8_t>(j), kDatagramSize),
+          send_callbacks[j].callback());
+
+      client_socket.ReceiveFrom(receive_callbacks[j].callback());
+    }
+
+    receive_callbacks[i - 1].WaitForResult();
+
+    for (size_t j = 0; j < i; ++j) {
+      EXPECT_EQ(static_cast<int>(kDatagramSize),
+                receive_callbacks[j].result()->code);
+      EXPECT_FALSE(receive_callbacks[j].src_addr());
+      EXPECT_TRUE(receive_callbacks[j].data().Equals(
+          CreateTestMessage(static_cast<uint8_t>(j), kDatagramSize)));
+    }
+  }
+}
+
 }  // namespace service
 }  // namespace mojo