1 files changed, 93 insertions, 26 deletions

diff --git a/net/socket/socket_test_util.cc b/net/socket/socket_test_util.cc
index 376d542..7730a5b 100644
--- a/net/socket/socket_test_util.cc
+++ b/net/socket/socket_test_util.cc
@@ -832,6 +832,26 @@ int MockTCPClientSocket::Write(net::IOBuffer* buf, int buf_len,
 
   return write_result.result;
 }
+int MockTCPClientSocket::Write(net::IOBuffer* buf, int buf_len,
+                               const net::CompletionCallback& callback) {
+  DCHECK(buf);
+  DCHECK_GT(buf_len, 0);
+
+  if (!connected_)
+    return net::ERR_UNEXPECTED;
+
+  std::string data(buf->data(), buf_len);
+  net::MockWriteResult write_result = data_->OnWrite(data);
+
+  was_used_to_convey_data_ = true;
+
+  if (write_result.async) {
+    RunCallbackAsync(callback, write_result.result);
+    return net::ERR_IO_PENDING;
+  }
+
+  return write_result.result;
+}
 
 int MockTCPClientSocket::Connect(net::OldCompletionCallback* callback) {
   if (connected_)
@@ -972,7 +992,7 @@ DeterministicMockTCPClientSocket::DeterministicMockTCPClientSocket(
     net::NetLog* net_log, net::DeterministicSocketData* data)
     : MockClientSocket(net_log),
       write_pending_(false),
-      write_callback_(NULL),
+      old_write_callback_(NULL),
       write_result_(0),
       read_data_(),
       read_buf_(NULL),
@@ -987,7 +1007,10 @@ DeterministicMockTCPClientSocket::~DeterministicMockTCPClientSocket() {}
 void DeterministicMockTCPClientSocket::CompleteWrite() {
   was_used_to_convey_data_ = true;
   write_pending_ = false;
-  write_callback_->Run(write_result_);
+  if (old_write_callback_)
+    old_write_callback_->Run(write_result_);
+  else
+    write_callback_.Run(write_result_);
 }
 
 int DeterministicMockTCPClientSocket::CompleteRead() {
@@ -1021,30 +1044,6 @@ int DeterministicMockTCPClientSocket::CompleteRead() {
   return result;
 }
 
-int DeterministicMockTCPClientSocket::Write(
-    net::IOBuffer* buf, int buf_len, net::OldCompletionCallback* callback) {
-  DCHECK(buf);
-  DCHECK_GT(buf_len, 0);
-
-  if (!connected_)
-    return net::ERR_UNEXPECTED;
-
-  std::string data(buf->data(), buf_len);
-  net::MockWriteResult write_result = data_->OnWrite(data);
-
-  if (write_result.async) {
-    write_callback_ = callback;
-    write_result_ = write_result.result;
-    DCHECK(write_callback_ != NULL);
-    write_pending_ = true;
-    return net::ERR_IO_PENDING;
-  }
-
-  was_used_to_convey_data_ = true;
-  write_pending_ = false;
-  return write_result.result;
-}
-
 int DeterministicMockTCPClientSocket::Read(
     net::IOBuffer* buf, int buf_len, net::OldCompletionCallback* callback) {
   if (!connected_)
@@ -1092,6 +1091,53 @@ int DeterministicMockTCPClientSocket::Read(
   return CompleteRead();
 }
 
+int DeterministicMockTCPClientSocket::Write(
+    net::IOBuffer* buf, int buf_len, net::OldCompletionCallback* callback) {
+  DCHECK(buf);
+  DCHECK_GT(buf_len, 0);
+
+  if (!connected_)
+    return net::ERR_UNEXPECTED;
+
+  std::string data(buf->data(), buf_len);
+  net::MockWriteResult write_result = data_->OnWrite(data);
+
+  if (write_result.async) {
+    old_write_callback_ = callback;
+    write_result_ = write_result.result;
+    DCHECK(old_write_callback_ != NULL);
+    write_pending_ = true;
+    return net::ERR_IO_PENDING;
+  }
+
+  was_used_to_convey_data_ = true;
+  write_pending_ = false;
+  return write_result.result;
+}
+int DeterministicMockTCPClientSocket::Write(
+    net::IOBuffer* buf, int buf_len, const net::CompletionCallback& callback) {
+  DCHECK(buf);
+  DCHECK_GT(buf_len, 0);
+
+  if (!connected_)
+    return net::ERR_UNEXPECTED;
+
+  std::string data(buf->data(), buf_len);
+  net::MockWriteResult write_result = data_->OnWrite(data);
+
+  if (write_result.async) {
+    write_callback_ = callback;
+    write_result_ = write_result.result;
+    DCHECK(!write_callback_.is_null());
+    write_pending_ = true;
+    return net::ERR_IO_PENDING;
+  }
+
+  was_used_to_convey_data_ = true;
+  write_pending_ = false;
+  return write_result.result;
+}
+
 // TODO(erikchen): Support connect sequencing.
 int DeterministicMockTCPClientSocket::Connect(
     net::OldCompletionCallback* callback) {
@@ -1235,6 +1281,10 @@ int MockSSLClientSocket::Write(net::IOBuffer* buf, int buf_len,
                                net::OldCompletionCallback* callback) {
   return transport_->socket()->Write(buf, buf_len, callback);
 }
+int MockSSLClientSocket::Write(net::IOBuffer* buf, int buf_len,
+                               const net::CompletionCallback& callback) {
+  return transport_->socket()->Write(buf, buf_len, callback);
+}
 
 int MockSSLClientSocket::Connect(net::OldCompletionCallback* callback) {
   OldConnectCallback* connect_callback = new OldConnectCallback(
@@ -1422,6 +1472,23 @@ int MockUDPClientSocket::Write(net::IOBuffer* buf, int buf_len,
   }
   return write_result.result;
 }
+int MockUDPClientSocket::Write(net::IOBuffer* buf, int buf_len,
+                               const net::CompletionCallback& callback) {
+  DCHECK(buf);
+  DCHECK_GT(buf_len, 0);
+
+  if (!connected_)
+    return ERR_UNEXPECTED;
+
+  std::string data(buf->data(), buf_len);
+  MockWriteResult write_result = data_->OnWrite(data);
+
+  if (write_result.async) {
+    RunCallbackAsync(callback, write_result.result);
+    return ERR_IO_PENDING;
+  }
+  return write_result.result;
+}
 
 bool MockUDPClientSocket::SetReceiveBufferSize(int32 size) {
   return true;